pytorch/BUCK.oss

load("//tools/build_defs:glob_defs.bzl", "subdir_glob")
load(
    ":build_variables.bzl",
    "aten_cpu_source_list",
    "aten_native_source_list",
    "core_sources_common",
    "jit_core_headers",
    "jit_core_sources",
    "libtorch_profiler_sources",
)
load(
    ":pt_defs.oss.bzl",
    "USED_PT_BACKENDS",
    "build_aten_cpu",
    "gen_aten_files",
    "gen_aten_libtorch_files",
    "get_aten_codegen_extra_params",
    "get_pt_compiler_flags",
    "get_pt_preprocessor_flags",
    "pt_operator_library",
    "get_pt_ops_deps",
    "aten_ufunc_generated_all_cpu_sources",
    "pt_xplat_cxx_library",
    "TEMPLATE_SOURCE_LIST",
)
load(":buckbuild.bzl",
    "define_buck_targets",
)

define_buck_targets(
    pt_xplat_cxx_library = pt_xplat_cxx_library,
)

cxx_library(
    name = "pthreadpool",
    srcs = ['caffe2/utils/threadpool/pthreadpool.cc', 'caffe2/utils/threadpool/pthreadpool_impl.cc', 'caffe2/utils/threadpool/pthreadpool-cpp.cc', 'caffe2/utils/threadpool/thread_pool_guard.cpp', 'caffe2/utils/threadpool/ThreadPool.cc'],
    deps = [':caffe2_headers', '//third_party:cpuinfo', '//third_party:glog', '//c10:c10', '//third_party:FXdiv'],
    exported_deps = ['//third_party:pthreadpool'],
    compiler_flags = ['-Wno-unused-function'],
    preferred_linkage = "static",
    exported_headers = subdir_glob([("", "caffe2/utils/threadpool/*.h")]),
    exported_preprocessor_flags = ['-DUSE_PTHREADPOOL'],
    header_namespace = "",
    headers = [],
    link_whole = True,
    platform_preprocessor_flags = [['windows', ['-D_WINDOWS', '-D_WIN32', '-DWIN32', '-DNOMINMAX', '-D_CRT_SECURE_NO_WARNINGS', '-D_USE_MATH_DEFINES']], ['windows.*64$', ['-D_WIN64']]],
    visibility = ['PUBLIC'],
)

cxx_library(
    name = "caffe2_headers",
    deps = ['//c10:c10'],
    exported_headers = subdir_glob(
        [
            ("", "caffe2/**/*.h"),
            ("", "binaries/**/*.h"),
            ("modules", "**/*.h"),
            ("aten/src", "ATen/core/**/*.h"),
        ],
        exclude = [
            "caffe2/fb/**/*.h",
            "caffe2/mobile/contrib/libopencl-stub/**/*.h",
            "caffe2/mobile/contrib/libvulkan-stub/**/*.h",
            "caffe2/mobile/contrib/nnapi/**/*.h",
            "caffe2/mobile/fb/binary/**/*.h",
            "caffe2/mobile/fb/snpe_so/**/*.h",
            "caffe2/mobile/fb/boltnn/bolt_lib/include/**/*.h",
            "caffe2/mobile/contrib/snpe/**/*.h",
            "caffe2/mobile/fb/qpl/jni/QuickPerformanceLogger.h",
            "caffe2/share/fb/x3d/ldi/*.h",
            "**/*.pb.h",
        ],
    ),
    compiler_flags = ['-Os', '-fexceptions', '-frtti', '-Wno-shadow', '-Wno-unknown-pragmas', '-Wno-unused-variable', '-Wno-sign-compare', '-Icaffe2', '-Imodules', '-DEIGEN_NO_DEBUG', '-DCAFFE2_USE_LITE_PROTO', '-DCAFFE2_USE_GOOGLE_GLOG', '-DCAFFE2_RNN_NO_TEXT_FORMAT', '-DGEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK=1', '-DCAFFE2_IS_XPLAT_BUILD', '-DSTRIP_ERROR_MESSAGES', '-DUSE_INTERNAL_PTHREADPOOL_IMPL', '-DCAFFE2_USE_HPTT'],
    preferred_linkage = "static",
    platform_preprocessor_flags = [['windows', ['-D_WINDOWS', '-D_WIN32', '-DWIN32', '-DNOMINMAX', '-D_CRT_SECURE_NO_WARNINGS', '-D_USE_MATH_DEFINES']], ['windows.*64$', ['-D_WIN64']]],
    preprocessor_flags = ['-DUSE_INTERNAL_PTHREADPOOL_IMPL'],
    visibility = ['PUBLIC'],
)

build_aten_cpu(
    name = "aten_cpu",
    srcs = jit_core_sources +
           aten_cpu_source_list + [
        # Generated
        ":gen_aten[Functions.cpp]",
        ":gen_aten[Operators_0.cpp]",
        ":gen_aten[Operators_1.cpp]",
        ":gen_aten[Operators_2.cpp]",
        ":gen_aten[Operators_3.cpp]",
        ":gen_aten[Operators_4.cpp]",
        ":gen_aten[core/ATenOpList.cpp]",
        ":gen_aten[core/TensorMethods.cpp]",
    ] + [
        # Needed by ATen/native/EmbeddingBag.cpp
        "caffe2/perfkernels/embedding_lookup_idx.cc",
    ],
)

gen_aten_files(
    name = "gen_aten",
    extra_flags = get_aten_codegen_extra_params(USED_PT_BACKENDS),
    visibility = ["PUBLIC"],
)

ATEN_EXPORTED_HEADERS = {
    "CPUFunctions.h": ":gen_aten[CPUFunctions.h]",
    "CPUFunctions_inl.h": ":gen_aten[CPUFunctions_inl.h]",
    "CompositeExplicitAutogradFunctions.h": ":gen_aten[CompositeExplicitAutogradFunctions.h]",
    "CompositeExplicitAutogradFunctions_inl.h": ":gen_aten[CompositeExplicitAutogradFunctions_inl.h]",
    "CompositeExplicitAutogradNonFunctionalFunctions.h": ":gen_aten[CompositeExplicitAutogradNonFunctionalFunctions.h]",
    "CompositeExplicitAutogradNonFunctionalFunctions_inl.h": ":gen_aten[CompositeExplicitAutogradNonFunctionalFunctions_inl.h]",
    "CompositeImplicitAutogradFunctions.h": ":gen_aten[CompositeImplicitAutogradFunctions.h]",
    "CompositeImplicitAutogradFunctions_inl.h": ":gen_aten[CompositeImplicitAutogradFunctions_inl.h]",
    "FunctionalInverses.h": ":gen_aten[FunctionalInverses.h]",
    "Functions.h": ":gen_aten[Functions.h]",
    "MethodOperators.h": ":gen_aten[MethodOperators.h]",
    "NativeFunctions.h": ":gen_aten[NativeFunctions.h]",
    "NativeMetaFunctions.h": ":gen_aten[NativeMetaFunctions.h]",
    "Operators.h": ":gen_aten[Operators.h]",
    "RedispatchFunctions.h": ":gen_aten[RedispatchFunctions.h]",
    "core/TensorBody.h": ":gen_aten[core/TensorBody.h]",
    "core/aten_interned_strings.h": ":gen_aten[core/aten_interned_strings.h]",
    "core/enum_tag.h": ":gen_aten[core/enum_tag.h]",
}

cxx_library(
    name = "generated_aten_headers_cpu",
    header_namespace = "ATen",
    exported_headers = ATEN_EXPORTED_HEADERS,
)
gen_aten_libtorch_files(name = "gen_aten_libtorch")

cxx_library(
    name = "caffe2_serialize",
    srcs = [
        "caffe2/serialize/file_adapter.cc",
        "caffe2/serialize/inline_container.cc",
        "caffe2/serialize/istream_adapter.cc",
        "caffe2/serialize/read_adapter_interface.cc",
    ],
    visibility = ["PUBLIC"],
    deps = [
        ":caffe2_headers",
        "//third_party:glog",
        "//c10:c10",
        "//third_party:miniz",
    ],
)

pt_operator_library(
    name = "torch_mobile_ops_full_dev",
    check_decl = False,
    include_all_operators = True,
)

cxx_library(
    name = "torch_mobile_all_ops",
    visibility = ["PUBLIC"],
    deps = get_pt_ops_deps(
        name = "pt_ops_full",
        train = False,
        deps = [
            ":torch_mobile_ops_full_dev",
        ],
        enable_flatbuffer = False,
    ),
)

cxx_binary(
    name = 'ptmobile_benchmark',
    srcs = [
        'binaries/speed_benchmark_torch.cc',
    ],
    compiler_flags = [
        "-fexceptions",
        "-frtti",
        "-Wno-deprecated-declarations",
    ],
    preprocessor_flags = [
        "-DBUILD_LITE_INTERPRETER",
    ],
    platform_linker_flags = [
        (
            "^linux.*$",
            [
                "-Wl,--no-as-needed",
                "-ldl",
                "-pthread",
            ],
        ),
    ],
    deps = [
        ":torch_mobile_core",
        ":torch_mobile_all_ops",
        "//c10:c10",
    ],
)

filegroup(
    name = "templated_selective_build_srcs",
    # NB: no glob here, there are generated targets in this list!
    srcs = glob(TEMPLATE_SOURCE_LIST) + aten_ufunc_generated_all_cpu_sources(":gen_aten[{}]"),
    visibility = [
        "PUBLIC",
    ],
)