BUCK.oss

load("//tools/build_defs:glob_defs.bzl", "subdir_glob")
load("//tools/build_defs:fb_xplat_genrule.bzl", "fb_xplat_genrule")
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",
    "TEMPLATE_SOURCE_LIST",
)

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'],
)

cxx_library(
    name = "common_core",
    srcs = ['caffe2/core/common.cc'],
    deps = [':caffe2_headers', '//c10:c10'],
    exported_deps = [],
    compiler_flags = ['-frtti', '-Os', '-Wno-unknown-pragmas', '-Wno-write-strings', '-Wno-unused-variable', '-Wno-unused-function', '-Wno-deprecated-declarations', '-Wno-shadow', '-Wno-global-constructors', '-Wno-missing-prototypes', '-std=gnu++17'],
    preferred_linkage = "static",
    header_namespace = "caffe2",
    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 = "aten_header",
    header_namespace = "",
    exported_headers = subdir_glob([
        # ATen Core
        ("aten/src", "ATen/core/**/*.h"),
        ("aten/src", "ATen/ops/*.h"),
        # ATen Base
        ("aten/src", "ATen/*.h"),
        ("aten/src", "ATen/cpu/**/*.h"),
        ("aten/src", "ATen/detail/*.h"),
        ("aten/src", "ATen/quantized/*.h"),
        ("aten/src", "ATen/vulkan/*.h"),
        ("aten/src", "ATen/metal/*.h"),
        ("aten/src", "ATen/mps/*.h"),
        ("aten/src", "ATen/nnapi/*.h"),
        # ATen Native
        ("aten/src", "ATen/native/*.h"),
        ("aten/src", "ATen/native/ao_sparse/quantized/cpu/*.h"),
        ("aten/src", "ATen/native/cpu/**/*.h"),
        ("aten/src", "ATen/native/sparse/*.h"),
        ("aten/src", "ATen/native/mps/*.h"),
        ("aten/src", "ATen/native/nested/*.h"),
        ("aten/src", "ATen/native/quantized/*.h"),
        ("aten/src", "ATen/native/quantized/cpu/*.h"),
        ("aten/src", "ATen/native/transformers/*.h"),
        ("aten/src", "ATen/native/ufunc/*.h"),
        ("aten/src", "ATen/native/utils/*.h"),
        ("aten/src", "ATen/native/vulkan/ops/*.h"),
        ("aten/src", "ATen/native/xnnpack/*.h"),
        # Remove the following after modifying codegen for mobile.
        ("aten/src", "ATen/mkl/*.h"),
        ("aten/src", "ATen/native/mkl/*.h"),
        ("aten/src", "ATen/native/mkldnn/*.h"),
    ], exclude = ["aten/src/ATen/Config.h"]),
    visibility = ["PUBLIC"],
)

cxx_library(
    name = "jit_core_headers",
    header_namespace = "",
    exported_headers = subdir_glob([("", x) for x in jit_core_headers]),
)

cxx_library(
    name = "generated_aten_config_header",
    header_namespace = "ATen",
    exported_headers = {
        "Config.h": ":generate_aten_config[Config.h]",
    },
)

cxx_library(
    name = "torch_mobile_headers",
    header_namespace = "",
    exported_headers = subdir_glob(
        [
            ("", "torch/csrc/jit/mobile/*.h"),
        ],
    ),
    visibility = ["PUBLIC"],
)

fb_xplat_genrule(
    name = "generate_aten_config",
    srcs = [
        "aten/src/ATen/Config.h.in",
    ],
    cmd = " ".join([
        "sed",
        "-e 's/@AT_MKLDNN_ENABLED@/ATEN_MKLDNN_ENABLED_FBXPLAT/g'",
        "-e 's/@AT_MKL_ENABLED@/ATEN_MKL_ENABLED_FBXPLAT/g'",
        "-e 's/@AT_MKL_SEQUENTIAL@/ATEN_MKL_SEQUENTIAL_FBXPLAT/g'",
        "-e 's/@AT_FFTW_ENABLED@/0/g'",
        "-e 's/@AT_POCKETFFT_ENABLED@/0/g'",
        "-e 's/@AT_NNPACK_ENABLED@/ATEN_NNPACK_ENABLED_FBXPLAT/g'",
        "-e 's/@CAFFE2_STATIC_LINK_CUDA_INT@/CAFFE2_STATIC_LINK_CUDA_FBXPLAT/g'",
        "-e 's/@AT_BUILD_WITH_BLAS@/USE_BLAS_FBXPLAT/g'",
        "-e 's/@AT_PARALLEL_OPENMP@/AT_PARALLEL_OPENMP_FBXPLAT/g'",
        "-e 's/@AT_PARALLEL_NATIVE@/AT_PARALLEL_NATIVE_FBXPLAT/g'",
        "-e 's/@AT_PARALLEL_NATIVE_TBB@/AT_PARALLEL_NATIVE_TBB_FBXPLAT/g'",
        "-e 's/@AT_BUILD_WITH_LAPACK@/USE_LAPACK_FBXPLAT/g'",
        "-e 's/@AT_BLAS_F2C@/AT_BLAS_F2C_FBXPLAT/g'",
        "-e 's/@AT_BLAS_USE_CBLAS_DOT@/AT_BLAS_USE_CBLAS_DOT_FBXPLAT/g'",
        "aten/src/ATen/Config.h.in > $OUT/Config.h"
    ]),
    outs = {
        "Config.h": ["Config.h"],
    },
    default_outs = ["."],
)

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]",
    "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]",
}

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

filegroup(
    name = "aten_src_path",
    srcs = [
        "aten/src/ATen/native/native_functions.yaml",
        "aten/src/ATen/native/tags.yaml",
    ] + glob(["aten/src/ATen/templates/*"]),
    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_libtorch_files(name = "gen_aten_libtorch")


GENERATED_AUTOGRAD_H = {
    "Functions.h": ":gen_aten_libtorch[autograd/generated/Functions.h]",
    "VariableType.h": ":gen_aten_libtorch[autograd/generated/VariableType.h]",
    "variable_factories.h": ":gen_aten_libtorch[autograd/generated/variable_factories.h]",

    # Don't build python bindings on mobile.
    #"python_functions.h",
}

cxx_library(
    name = "generated-autograd-headers",
    header_namespace = "torch/csrc/autograd/generated",
    exported_headers = GENERATED_AUTOGRAD_H,
    visibility = ["PUBLIC"],
)

cxx_library(
    name = "torch_mobile_observer",
    srcs = [
        "torch/csrc/jit/mobile/observer.cpp",
        #"torch/fb/observers/MobileObserverUtil.cpp",
    ],
    header_namespace = "",
    exported_headers = subdir_glob(
        [
            ("", "torch/csrc/jit/mobile/observer.h"),
            #("", "torch/fb/observers/ObserverUtil.h"),
            #("", "torch/fb/observers/MobileObserverUtil.h"),
        ],
    ),
    visibility = ["PUBLIC"],
    deps = [
        "//c10:c10",
    ],
)

fb_xplat_genrule(
    name = "generate-version-header",
    srcs = [
        "torch/csrc/api/include/torch/version.h.in",
        "version.txt",
    ],
    cmd = "$(exe //tools/setup_helpers:gen-version-header) " + " ".join([
        "--template-path",
        "torch/csrc/api/include/torch/version.h.in",
        "--version-path",
        "version.txt",
        "--output-path",
        "$OUT/version.h",
    ]),
    outs = {
        "version.h": ["version.h"],
    },
    default_outs = ["."],
)

cxx_library(
    name = "generated-version-header",
    header_namespace = "torch",
    exported_headers = {
        "version.h": ":generate-version-header[version.h]",
    },
)

cxx_library(
    name = "torch_headers",
    header_namespace = "",
    exported_headers = subdir_glob(
        [
            ("torch/csrc/api/include", "torch/**/*.h"),
            ("", "torch/csrc/**/*.h"),
            ("", "torch/csrc/generic/*.cpp"),
            ("", "torch/script.h"),
            ("", "torch/library.h"),
            ("", "torch/custom_class.h"),
            ("", "torch/custom_class_detail.h"),
            # Add again due to namespace difference from aten_header.
            ("", "aten/src/ATen/*.h"),
            ("", "aten/src/ATen/quantized/*.h"),
        ],
        exclude = [
            # Don't need on mobile.
            "torch/csrc/Exceptions.h",
            "torch/csrc/python_headers.h",
            "torch/csrc/utils/auto_gil.h",
            "torch/csrc/jit/serialization/mobile_bytecode_generated.h",
            "torch/csrc/api/include/torch/version.h",
        ],
    ),
    visibility = ["PUBLIC"],
    deps = [
        ":generated-version-header",
    ],
)


cxx_library(
    name = "torch_common",
    srcs = core_sources_common,
    compiler_flags = get_pt_compiler_flags(),
    exported_preprocessor_flags = get_pt_preprocessor_flags(),
    link_whole = True,
    visibility = ["PUBLIC"],
    deps = [
        ":aten_cpu",
        ":generated-autograd-headers",
        ":torch_headers",
        "//third_party:glog",
        "//c10:c10",
    ],
)


cxx_library(
    name = "torch_mobile_deserialize_common",
    srcs = [
        "torch/csrc/jit/mobile/parse_bytecode.cpp",
        "torch/csrc/jit/mobile/parse_operators.cpp",
        "torch/csrc/jit/mobile/upgrader_mobile.cpp",
        "torch/csrc/jit/serialization/import_read.cpp",
        "torch/csrc/jit/serialization/unpickler.cpp",
    ],
    header_namespace = "",
    exported_headers = [
        "torch/csrc/jit/serialization/import_read.h",
        "torch/csrc/jit/serialization/unpickler.h",
    ],
    compiler_flags = get_pt_compiler_flags(),
    link_whole = True,
    linker_flags = [
        "-Wl,--no-as-needed",
    ],
    visibility = ["PUBLIC"],
    exported_deps = [
        ":aten_cpu",
        ":caffe2_headers",
        ":caffe2_serialize",
        ":torch_common",
        ":torch_headers",
        ":torch_mobile_headers",
        ":torch_mobile_module",
        ":torch_mobile_observer",
        "//third_party:glog",
        "//c10:c10",
    ],
)

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",
    ],
)

cxx_library(
    name = "torch_mobile_deserialize",
    srcs = [
        "torch/csrc/jit/mobile/import.cpp",
    ],
    header_namespace = "",
    exported_headers = [
        "torch/csrc/jit/mobile/import.h",
    ],
    compiler_flags = get_pt_compiler_flags(),
    link_whole = True,
    linker_flags = [
        "-Wl,--no-as-needed",
    ],
    visibility = ["PUBLIC"],
    exported_deps = [
        ":aten_cpu",
        ":caffe2_headers",
        ":caffe2_serialize",
        ":torch_common",
        ":torch_headers",
        ":torch_mobile_headers",
        ":torch_mobile_module",
        ":torch_mobile_observer",
        "//third_party:glog",
        "//c10:c10",
        ":torch_mobile_deserialize_common",
    ],
)

cxx_library(
    name = "torch_mobile_module",
    srcs = [
        "torch/csrc/jit/mobile/function.cpp",
        "torch/csrc/jit/mobile/interpreter.cpp",
        "torch/csrc/jit/mobile/module.cpp",
    ],
    header_namespace = "",
    exported_headers = [],
    compiler_flags = get_pt_compiler_flags(),
    link_whole = True,
    linker_flags = [
        "-Wl,--no-as-needed",
    ],
    visibility = ["PUBLIC"],
    exported_deps = [
        ":aten_cpu",
        ":caffe2_headers",
        ":torch_common",
        ":torch_headers",
        ":torch_mobile_headers",
        ":torch_mobile_observer",
        "//third_party:glog",
        "//c10:c10",
    ],
)

cxx_library(
    name = "torch_mobile_core",
    srcs = [],
    header_namespace = "",
    exported_headers = [],
    compiler_flags = get_pt_compiler_flags(),
    exported_preprocessor_flags = get_pt_preprocessor_flags(),
    link_whole = True,
    linker_flags = [
        "-Wl,--no-as-needed",
        # "-ldl",
    ],
    visibility = ["PUBLIC"],
    deps = [
        ":generated-autograd-headers",
        ":torch_mobile_observer",
        ":torch_mobile_headers",
    ],
    exported_deps = [
        ":aten_cpu",
        ":torch_common",
        ":torch_mobile_deserialize",
    ],
)

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,
    ),
)

python_library(
    name = "gen_oplist_lib",
    srcs = subdir_glob([
        ("tools/code_analyzer", "gen_oplist.py"),
        ("tools/code_analyzer", "gen_op_registration_allowlist.py"),
    ]),
    base_module = "",
    deps = [
        "//third_party:pyyaml",
        "//tools/lite_interpreter:gen_selected_mobile_ops_header",
        "//torchgen:torchgen",
    ],
)

python_binary(
    name = "gen_oplist",
    main_module = "gen_oplist",
    visibility = ["PUBLIC"],
    deps = [
        ":gen_oplist_lib",
    ],
)

python_library(
    name = "gen_operators_yaml_lib",
    srcs = subdir_glob([
        ("tools/code_analyzer", "gen_operators_yaml.py"),
        ("tools/code_analyzer", "gen_op_registration_allowlist.py"),
    ]),
    base_module = "",
    deps = [
        "//third_party:pyyaml",
        "//torchgen:torchgen",
    ],
)

python_binary(
    name = "gen_operators_yaml",
    main_module = "gen_operators_yaml",
    visibility = ["PUBLIC"],
    deps = [
        ":gen_operators_yaml_lib",
    ],
)

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",
    ],
)