pt_defs.oss.bzl

load("@bazel_skylib//lib:paths.bzl", "paths")
load("//tools/build_defs:fb_xplat_genrule.bzl", "fb_xplat_genrule")
load("//tools/build_defs:type_defs.bzl", "is_list", "is_string")
load(":build_variables.bzl", "aten_native_source_list")
load(
    ":ufunc_defs.bzl",
    "aten_ufunc_generated_cpu_kernel_sources",
    "aten_ufunc_generated_cpu_sources",
    "aten_ufunc_generated_cuda_sources",
)

USED_PT_BACKENDS = [
    "CPU",
    "QuantizedCPU",
    "SparseCPU",  # brings ~20 kb size regression
]

# This needs to be kept in sync with https://github.com/pytorch/pytorch/blob/release/1.9/torchgen/gen.py#L892
PT_BACKEND_HEADERS = [
    "CPU",
    "CUDA",
    "CompositeExplicitAutograd",
    "CompositeImplicitAutograd",
    "Meta",
]

PT_BASE_OPS = [
    "aten::_coalesced_",
    "aten::_copy_from",
    "aten::_empty_affine_quantized",
    "aten::_empty_per_channel_affine_quantized",
    "aten::_indices",
    "aten::_nnz",
    "aten::_values",
    "aten::add",
    "aten::add_",
    "aten::arange",
    "aten::as_strided",
    "aten::as_strided_",
    "aten::cat",
    "aten::clone",
    "aten::coalesce",
    "aten::contiguous",
    "aten::copy_",
    "aten::copy_sparse_to_sparse_",
    "aten::dense_dim",
    "aten::dequantize",
    "aten::div",
    "aten::div_",
    "aten::empty",
    "aten::empty_like",
    "aten::empty_strided",
    "aten::empty.memory_format",
    "aten::eq",
    "aten::equal",
    "aten::expand",
    "aten::fill_",
    "aten::is_coalesced",
    "aten::is_complex",
    "aten::is_floating_point",
    "aten::is_leaf",
    "aten::is_nonzero",
    "aten::item",
    "aten::max",
    "aten::min",
    "aten::mul",
    "aten::mul_",
    "aten::narrow",
    "aten::ne",
    "aten::permute",
    "aten::q_per_channel_axis",
    "aten::q_per_channel_scales",
    "aten::q_per_channel_zero_points",
    "aten::q_scale",
    "aten::q_zero_point",
    "aten::qscheme",
    "aten::quantize_per_tensor",
    "aten::reshape",
    "aten::_reshape_alias",
    "aten::resize_",
    "aten::resize_as_",
    "aten::scalar_tensor",
    "aten::select",
    "aten::set_",
    "aten::size",
    "aten::slice",
    "aten::sparse_dim",
    "aten::sparse_resize_and_clear_",
    "aten::squeeze",
    "aten::squeeze_",
    "aten::stride",
    "aten::sub",
    "aten::sub_",
    "aten::sum",
    "aten::t",
    "aten::to",
    "aten::_to_copy",
    "aten::unsqueeze",
    "aten::view",
    "aten::zero_",
    "aten::zeros",
    "aten::zeros_like",
]

def get_aten_compiler_flags():
    return ATEN_COMPILER_FLAGS

def get_generate_code_bin_outs():
    return {
        "autograd/generated/ADInplaceOrViewTypeEverything.cpp": ["autograd/generated/ADInplaceOrViewTypeEverything.cpp"],
        "autograd/generated/ADInplaceOrViewType_0.cpp": ["autograd/generated/ADInplaceOrViewType_0.cpp"],
        "autograd/generated/ADInplaceOrViewType_1.cpp": ["autograd/generated/ADInplaceOrViewType_1.cpp"],
        "autograd/generated/Functions.cpp": ["autograd/generated/Functions.cpp"],
        "autograd/generated/Functions.h": ["autograd/generated/Functions.h"],
        "autograd/generated/TraceTypeEverything.cpp": ["autograd/generated/TraceTypeEverything.cpp"],
        "autograd/generated/TraceType_0.cpp": ["autograd/generated/TraceType_0.cpp"],
        "autograd/generated/TraceType_1.cpp": ["autograd/generated/TraceType_1.cpp"],
        "autograd/generated/TraceType_2.cpp": ["autograd/generated/TraceType_2.cpp"],
        "autograd/generated/TraceType_3.cpp": ["autograd/generated/TraceType_3.cpp"],
        "autograd/generated/TraceType_4.cpp": ["autograd/generated/TraceType_4.cpp"],
        "autograd/generated/VariableType.h": ["autograd/generated/VariableType.h"],
        "autograd/generated/VariableTypeEverything.cpp": ["autograd/generated/VariableTypeEverything.cpp"],
        "autograd/generated/VariableType_0.cpp": ["autograd/generated/VariableType_0.cpp"],
        "autograd/generated/VariableType_1.cpp": ["autograd/generated/VariableType_1.cpp"],
        "autograd/generated/VariableType_2.cpp": ["autograd/generated/VariableType_2.cpp"],
        "autograd/generated/VariableType_3.cpp": ["autograd/generated/VariableType_3.cpp"],
        "autograd/generated/VariableType_4.cpp": ["autograd/generated/VariableType_4.cpp"],
        "autograd/generated/variable_factories.h": ["autograd/generated/variable_factories.h"],
    }

ATEN_COMPILER_FLAGS = [
    "-fexceptions",
    "-frtti",
    "-fPIC",
    "-Os",
    "-Wno-absolute-value",
    "-Wno-deprecated-declarations",
    "-Wno-macro-redefined",
    "-Wno-tautological-constant-out-of-range-compare",
    "-Wno-unknown-pragmas",
    "-Wno-unknown-warning-option",
    "-Wno-unused-function",
    "-Wno-unused-variable",
    "-Wno-pass-failed",
    "-Wno-shadow",
]

PT_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",  # to accommodate Eigen
]

def get_template_source_dict():
    ret = {}
    for file_path in TEMPLATE_SOURCE_LIST:
        path_prefix = paths.dirname(file_path)
        if path_prefix not in ret:
            ret[path_prefix] = []
        ret[path_prefix].append(file_path)
    return ret

def get_gen_oplist_outs():
    return {
        #"SupportedMobileModelsRegistration.cpp": [
        #    "SupportedMobileModelsRegistration.cpp",
        #],
        "selected_mobile_ops.h": [
            "selected_mobile_ops.h",
        ],
        "selected_operators.yaml": [
            "selected_operators.yaml",
        ],
    }

def get_pt_compiler_flags():
    return PT_COMPILER_FLAGS

def get_aten_preprocessor_flags():
    # read_config is not allowed outside of function in Starlark
    ATEN_PREPROCESSOR_FLAGS = [
        "-DC10_MOBILE",
        "-DCPU_CAPABILITY_DEFAULT",
        "-DCPU_CAPABILITY=DEFAULT",
        "-DCAFFE2_USE_LITE_PROTO",
        "-DATEN_CUDNN_ENABLED_FBXPLAT=0",
        "-DATEN_MKLDNN_ENABLED_FBXPLAT=0",
        "-DATEN_NNPACK_ENABLED_FBXPLAT=0",
        "-DATEN_MKL_ENABLED_FBXPLAT=0",
        "-DATEN_MKL_SEQUENTIAL_FBXPLAT=0",
        "-DUSE_PYTORCH_METAL",
        "-DUSE_PYTORCH_QNNPACK",
        "-DUSE_XNNPACK",
        "-DNO_EXPORT",
        "-DPYTORCH_QNNPACK_RUNTIME_QUANTIZATION",
        "-DAT_PARALLEL_OPENMP_FBXPLAT=0",
        "-DAT_PARALLEL_NATIVE_FBXPLAT=1",
        "-DAT_PARALLEL_NATIVE_TBB_FBXPLAT=0",
        "-DUSE_LAPACK_FBXPLAT=0",
        "-DAT_BLAS_F2C_FBXPLAT=0",
        "-DAT_BLAS_USE_CBLAS_DOT_FBXPLAT=0",
        "-DUSE_RUY_QMATMUL",  # need third_party:ruy
    ]

    # if get_disable_per_op_profiling():
    ATEN_PREPROCESSOR_FLAGS.append("-DPYTORCH_DISABLE_PER_OP_PROFILING")
    return ATEN_PREPROCESSOR_FLAGS

TEMPLATE_SOURCE_LIST = [
    "torch/csrc/jit/runtime/register_prim_ops.cpp",
    "torch/csrc/jit/runtime/register_special_ops.cpp",
] + aten_native_source_list

# For selective build, we can lump the CPU and CPU kernel sources altogether
# because there is only ever one vectorization variant that is compiled
def aten_ufunc_generated_all_cpu_sources(gencode_pattern = "{}"):
    return (
        aten_ufunc_generated_cpu_sources(gencode_pattern) +
        aten_ufunc_generated_cpu_kernel_sources(gencode_pattern)
    )

def get_template_registration_files_outs():
    outs = {}

    for file_path in TEMPLATE_SOURCE_LIST:
        outs[file_path] = [file_path]

    for base_name in aten_ufunc_generated_all_cpu_sources():
        file_path = "aten/src/ATen/{}".format(base_name)
        outs[file_path] = [file_path]

    return outs

def get_pt_preprocessor_flags():
    # read_config is not allowed outside of function in Starlark
    PT_PREPROCESSOR_FLAGS = [
        "-D_THP_CORE",
        "-DC10_MOBILE",
        "-DUSE_SCALARS",
        "-DNO_CUDNN_DESTROY_HANDLE",
        "-DNO_EXPORT",
        "-DBUILD_CAFFE2",
    ]
    return PT_PREPROCESSOR_FLAGS

def is_arvr_mode():
    return False

def get_build_from_deps_query():
    build_from_query = native.read_config("pt", "build_from_deps_query", "1")
    return bool(int(build_from_query))

def get_enable_lightweight_dispatch():
    enable_lightweight_dispatch = native.read_config("pt", "enable_lightweight_dispatch", "0")
    return bool(int(enable_lightweight_dispatch))

def get_static_dispatch_backend():
    static_dispatch_backend = native.read_config("pt", "static_dispatch_backend", None)
    if static_dispatch_backend == None:
        return []
    return static_dispatch_backend.split(";")

def get_aten_codegen_extra_params(backends):
    if get_build_from_deps_query():
        extra_params = {
            "force_schema_registration": True,
        }
        static_backends = get_static_dispatch_backend()
        if static_backends:
            extra_params["static_dispatch_backend"] = static_backends
            extra_params["enabled_backends"] = static_backends
        else:
            extra_params["enabled_backends"] = backends
        return extra_params
    else:
        return {}

def gen_aten_files(
        name,
        extra_flags = {},
        visibility = [],
        compatible_with = []):
    extra_params = []
    force_schema_registration = extra_flags.get("force_schema_registration", False)
    op_registration_allowlist = extra_flags.get("op_registration_allowlist", None)
    op_selection_yaml_path = extra_flags.get("op_selection_yaml_path", None)
    enabled_backends = extra_flags.get("enabled_backends", None)
    static_dispatch_backend = extra_flags.get("static_dispatch_backend", None)

    if force_schema_registration:
        extra_params.append("--force_schema_registration")
    if op_registration_allowlist != None and is_string(op_registration_allowlist):
        extra_params.append("--op_registration_whitelist")
        extra_params.append(op_registration_allowlist)
    if op_selection_yaml_path != None and is_string(op_selection_yaml_path):
        extra_params.append("--op_selection_yaml_path")
        extra_params.append(op_selection_yaml_path)
    if enabled_backends != None and is_list(enabled_backends):
        extra_params.append("--backend_whitelist")
        extra_params.extend(enabled_backends)
    if get_enable_lightweight_dispatch():
        extra_params.append("--skip_dispatcher_op_registration")
    if static_dispatch_backend:
        extra_params.append("--static_dispatch_backend")
        extra_params.extend(static_dispatch_backend)
        backends = static_dispatch_backend
    else:
        backends = enabled_backends
    fb_xplat_genrule(
        name = name,
        default_outs = ["."],
        outs = get_aten_generated_files(backends),
        cmd = "$(exe //torchgen:gen) " + " ".join([
            "--source-path $(location //:aten_src_path)/aten/src/ATen",
            "--install_dir $OUT",
        ] + extra_params),
        visibility = visibility,
        compatible_with = compatible_with,
    )

def get_aten_generated_files(enabled_backends):
    # NB: RegisterMeta counts as an optionally enabled backend,
    # and is intentionally omitted from here
    src_files = [
        "RegisterBackendSelect.cpp",
        "RegisterCompositeImplicitAutograd.cpp",
        "RegisterCompositeExplicitAutograd.cpp",
        "CompositeViewCopyKernels.cpp",
        "RegisterSchema.cpp",
        "Declarations.yaml",
        "Functions.cpp",
        "Functions.h",
        "RedispatchFunctions.h",
        "NativeFunctions.h",
        "NativeMetaFunctions.h",
        "MethodOperators.h",
        "FunctionalInverses.h",
        "Operators.h",
        "Operators_0.cpp",
        "Operators_1.cpp",
        "Operators_2.cpp",
        "Operators_3.cpp",
        "Operators_4.cpp",
        "CompositeImplicitAutogradFunctions.h",
        "CompositeImplicitAutogradFunctions_inl.h",
        "CompositeExplicitAutogradFunctions.h",
        "CompositeExplicitAutogradFunctions_inl.h",
        "core/ATenOpList.cpp",
        "core/TensorBody.h",
        "core/TensorMethods.cpp",
        "core/aten_interned_strings.h",
        "core/enum_tag.h",
    ] + get_aten_derived_type_srcs(enabled_backends)

    # This is tiresome.  A better strategy would be to unconditionally
    # generate these files, and then only actually COMPILE them depended
    # on the generated set.  C'est la vie...
    if "CPU" in enabled_backends:
        src_files.extend(aten_ufunc_generated_cpu_sources())
        src_files.extend(aten_ufunc_generated_cpu_kernel_sources())
    if "CUDA" in enabled_backends:
        # Cannot unconditionally include this, because in the Edge selective
        # build CUDA is not enabled and thus the ufunc codegen for CUDA gets
        # skipped
        src_files.extend(aten_ufunc_generated_cuda_sources())

    res = {}
    for file_name in src_files:
        res[file_name] = [file_name]
    return res

def get_template_registration_file_rules(rule_name):
    rules = []
    for file_path in TEMPLATE_SOURCE_LIST:
        rules.append(":{}[{}]".format(rule_name, file_path))
    for file_path in aten_ufunc_generated_all_cpu_sources():
        rules.append(":{}[aten/src/ATen/{}]".format(rule_name, file_path))

    return rules

# Originally, there were two sets of codes in caffe2:aten_cpu, native codes and non-native.
# Now we have only non-naitve sources in aten_cpu. However, there are some aten related
# tests that may require both native and non-native codes. This rule is used to generate
# both aten_cpu and aten_native_cpu. They are using the same compilation setups.
def build_aten_cpu(name, srcs, deps = []):
    cxx_library(
        name = name,
        srcs = srcs,
        header_namespace = "",
        compiler_flags = get_pt_compiler_flags(),
        exported_preprocessor_flags = get_aten_preprocessor_flags(),
        link_whole = True,
        linker_flags = ["-Wl,--no-as-needed", "-ldl"],
        visibility = ["PUBLIC"],
        deps = [
            "//third_party:cpuinfo",
            "//third_party:glog",
            "//third_party:XNNPACK",
            #"//third_party/linker_lib:omp",
        ],
        exported_deps = [
            "//third_party:fmt",
            "//aten/src/ATen/native/quantized/cpu/qnnpack:pytorch_qnnpack",
            "//c10:c10",
            ":aten_header",
            ":caffe2_headers",
            ":common_core",
            ":generated_aten_config_header",
            ":generated_aten_headers_cpu",
            ":jit_core_headers",
            ":pthreadpool",
            "//third_party:ruy_lib",
        ],
    )

######### selective build #########

def get_pt_ops_deps(name, deps, train = False, enforce_traced_op_list = False, enable_flatbuffer = False, **kwargs):
    if not get_build_from_deps_query():
        return deps
    pt_operator_registry(
        name,
        deps,
        train = train,
        enforce_traced_op_list = enforce_traced_op_list,
        enable_flatbuffer = enable_flatbuffer,
        **kwargs
    )
    return deps + [":" + name]

# pt_operator_registry is the method that defines the fb_xplat_cxx_library that contains
# code for all selected PyTorch Operators and kernel functions. This also includes
# operator registration into the dispatcher.
#
# template_select: bool: Indicates if template based selective build is enabled.
#
# enforce_traced_op_list: bool: Enforces that only new-style operator
# lists based on the all_mobile_model_configs.yaml file and tracing based selective
# build are used in this library.
#
# train: bool: Build this library for training (True) or inference only (False).
# If built for training, codegen for VariableType is also included.
#
# pt_allow_forced_schema_registration: Manually disables forced schema registration when set to false, Default is true.
# Only does anything when train=True and the app requires full jit then force_schema_registration needs to occur.
# As Federated Learning migrates to lite interpreter
# we can slowly turn off forced schema registration as it is useless space and floods the compatibility api
#
def pt_operator_registry(
        name,
        deps = [],
        train = False,
        labels = [],
        env = [],
        template_select = True,
        enforce_traced_op_list = False,
        pt_allow_forced_schema_registration = True,
        enable_flatbuffer = False,
        **kwargs):
    compatible_with = kwargs.get("compatible_with", [])
    code_gen_files = pt_operator_query_codegen(name, deps = deps, train = train, enforce_traced_op_list = enforce_traced_op_list, pt_allow_forced_schema_registration = pt_allow_forced_schema_registration, compatible_with = compatible_with)
    code_gen_srcs = code_gen_files["srcs"]

    lib_deps = [
        ":aten_cpu",
        ":torch_mobile_core",
        "//c10:c10",
        "//third_party:glog",
    ]

    #if train:
    #    lib_deps = lib_deps + ["fbsource//xplat/caffe2:torch_mobile_train"]

    exported_preprocessor_flags = get_aten_preprocessor_flags()
    exported_preprocessor_flags += kwargs.pop("exported_preprocessor_flags", [])
    if template_select:
        # In addition to the
        # original code-gen select, this option further filter more operators based on
        # compile-time calculation. Examples include prim ops and any other ops that were
        # not filtered out before. The purpose of this option is to reduce the production
        # size further. However, it may have less flexibility, especially for tests from
        # python, where the used operator list is not explicitly generated. If the tests
        # are for functionality but not for size, and it's difficult to maintain an explicit
        # operator list, it's suggested to turn this option off.
        exported_preprocessor_flags.append("-DTEMPLATE_SELECTIVE_BUILD")
    kwargs.pop("exported_headers", [])
    cxx_library(
        name = name,
        srcs = code_gen_srcs,
        linker_flags = [
            "-Wl,--no-as-needed",
            "-ldl",
        ],
        link_whole = True,
        soname = "libtorch-code-gen.$(ext)",
        compiler_flags = get_aten_compiler_flags(),
        platform_compiler_flags = get_cpukernel_avx2_flags(),
        platform_deps = get_cpukernel_avx2_deps(),
        header_namespace = "ATen",
        exported_headers = code_gen_files["headers"],
        exported_preprocessor_flags = exported_preprocessor_flags,
        headers = kwargs.pop("headers", []),
        deps = lib_deps + [
            "//third_party:XNNPACK",
        ],
        **kwargs
    )

def get_aten_derived_type_src_rules(aten_rule_name, enabled_backends):
    return [
        ":{}[{}]".format(aten_rule_name, "Register" + backend + ".cpp")
        for backend in enabled_backends
    ]

def get_aten_selective_cpp_rules(aten_rule_name, enabled_backends):
    return [
        ":{}[{}]".format(aten_rule_name, f)
        for f in ["RegisterCompositeImplicitAutograd.cpp", "RegisterCompositeExplicitAutograd.cpp", "RegisterSchema.cpp", "RegisterBackendSelect.cpp", "CompositeViewCopyKernels.cpp"]
    ] + get_aten_derived_type_src_rules(aten_rule_name, enabled_backends)

def get_aten_derived_type_srcs(enabled_backends):
    return [
        "Register" + derived_type + ".cpp"
        for derived_type in enabled_backends
    ] + [
        derived_type + "Functions.h"
        for derived_type in enabled_backends
        if derived_type in PT_BACKEND_HEADERS or derived_type in get_static_dispatch_backend()
    ] + [
        derived_type + "Functions_inl.h"
        for derived_type in enabled_backends
        if derived_type in PT_BACKEND_HEADERS or derived_type in get_static_dispatch_backend()
    ]

def pt_operator_query_codegen(name, deps = [], train = False, enforce_traced_op_list = False, pt_allow_forced_schema_registration = True, compatible_with = []):
    oplist_dir_name = name + "_pt_oplist"

    # @lint-ignore BUCKLINT
    fb_xplat_genrule(
        name = oplist_dir_name,
        cmd = ("$(exe //:gen_oplist) " +
               "--model_file_list_path $(@query_outputs 'attrfilter(labels, pt_operator_library, deps(set({deps})))') " +
               ("" if enforce_traced_op_list else "--allow_include_all_overloads ") +
               "--output_dir $OUT ").format(deps = " ".join(["\"{}\"".format(d) for d in deps])),
        outs = get_gen_oplist_outs(),
        default_outs = ["."],
        compatible_with = compatible_with,
    )

    # Aten files
    aten_genrule = name + "_aten"
    extra_flags = {
        "enabled_backends": USED_PT_BACKENDS,
        "op_selection_yaml_path": "$(location :{}[selected_operators.yaml])".format(oplist_dir_name),
    }

    if train and pt_allow_forced_schema_registration:
        extra_flags["force_schema_registration"] = True

    # if get_enable_lightweight_dispatch():
    #    unboxing_genrule = name + "_unboxing"
    #    gen_aten_unboxing_files(
    #        unboxing_genrule,
    #        extra_flags = extra_flags,
    #    )

    static_dispatch_backend = get_static_dispatch_backend()
    if static_dispatch_backend:
        extra_flags["static_dispatch_backend"] = static_dispatch_backend

    gen_aten_files(
        aten_genrule,
        extra_flags = extra_flags,
        compatible_with = compatible_with,
    )

    # unboxing_wrappers files
    extra_params = [
        "--operators_yaml_path",
        "$(location :" + oplist_dir_name + "[selected_operators.yaml])",
    ]
    unboxing_and_autograd_genrule = name + "_unboxing_and_autograd"
    gen_aten_libtorch_files(unboxing_and_autograd_genrule, extra_params, compatible_with)

    # Template runtime files (prim ops, etc)
    template_registration_genrule = name + "_template_registration"
    copy_template_registration_files(template_registration_genrule)

    srcs = get_aten_selective_cpp_rules(
        aten_genrule,
        static_dispatch_backend if static_dispatch_backend else USED_PT_BACKENDS,
    ) + get_template_registration_file_rules(
        template_registration_genrule,
    ) + ([
        ":{}[autograd/generated/VariableType_0.cpp]".format(unboxing_and_autograd_genrule),
        ":{}[autograd/generated/VariableType_1.cpp]".format(unboxing_and_autograd_genrule),
        ":{}[autograd/generated/VariableType_2.cpp]".format(unboxing_and_autograd_genrule),
        ":{}[autograd/generated/VariableType_3.cpp]".format(unboxing_and_autograd_genrule),
        ":{}[autograd/generated/VariableType_4.cpp]".format(unboxing_and_autograd_genrule),
        ":{}[autograd/generated/ADInplaceOrViewType_0.cpp]".format(unboxing_and_autograd_genrule),
        ":{}[autograd/generated/ADInplaceOrViewType_1.cpp]".format(unboxing_and_autograd_genrule),
    ] if train else []) + ([
        #":{}[SupportedMobileModelsRegistration.cpp]".format(oplist_dir_name),
    ])

    headers = {
        "selected_mobile_ops.h": ":{}[selected_mobile_ops.h]".format(oplist_dir_name),
    }

    # if get_enable_lightweight_dispatch():
    #     srcs.extend([
    #         ":{}[UnboxingFunctions_0.cpp]".format(unboxing_genrule),
    #         ":{}[UnboxingFunctions_1.cpp]".format(unboxing_genrule),
    #         ":{}[UnboxingFunctions_2.cpp]".format(unboxing_genrule),
    #         ":{}[UnboxingFunctions_3.cpp]".format(unboxing_genrule),
    #         ":{}[UnboxingFunctions_4.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_0.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_1.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_2.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_3.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_4.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_5.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_6.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_7.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_8.cpp]".format(unboxing_genrule),
    #         ":{}[RegisterCodegenUnboxedKernels_9.cpp]".format(unboxing_genrule),
    #     ])
    #     headers["UnboxingFunctions.h"] = ":{}[UnboxingFunctions.h]".format(unboxing_genrule)
    return {"headers": headers, "srcs": srcs}

def gen_aten_libtorch_files(name, extra_params = [], compatible_with = []):
    fb_xplat_genrule(
        name = name,
        outs = get_generate_code_bin_outs(),
        default_outs = ["."],
        cmd = "mkdir -p tools && " +
              "$(exe //tools/setup_helpers:generate_code_bin) " + " ".join(
            # Mobile build only needs libtorch - skip python bindings for now, except
            # for ovrsource, which needs Python bindings.
            (["--subset libtorch"] if not is_arvr_mode() else []) + [
                "--native-functions-path $(location :aten_src_path)/aten/src/ATen/native/native_functions.yaml",
                "--tags-path $(location :aten_src_path)/aten/src/ATen/native/tags.yaml",  # todo D35992309
                "--install_dir $OUT",
            ] + extra_params,
        ),
        cmd_exe = "@powershell -Command New-Item -Path tools -ItemType Directory -Force; " +
                  "$(exe //tools/setup_helpers:generate_code_bin) " + " ".join(
            # Mobile build only needs libtorch - skip python bindings for now, except
            # for ovrsource, which needs Python bindings.
            (["--subset libtorch"] if not is_arvr_mode() else []) + [
                "--native-functions-path $(location :aten_src_path)/aten/src/ATen/native/native_functions.yaml",
                "--tags-path $(location :aten_src_path)/aten/src/ATen/native/tags.yaml",
                "--install_dir $OUT",
            ] + extra_params,
        ),
        compatible_with = compatible_with,
    )

def copy_template_registration_files(name):
    cmd = []
    cmd_exe = []

    template_source_dict = get_template_source_dict()

    # Ideally, we would run one copy command for a single source directory along
    # with all its child directories, but it's somewhat hard to know if a directory
    # is a child of another just bu looking at the metadata (directory relative
    # path) that we currently have since 1 directory could look like a parent of
    # another and yet come from a different filegroup() rule.
    #
    for (path_prefix, file_paths) in template_source_dict.items():
        cmd.append("mkdir -p $OUT/{}".format(path_prefix))
        cmd_exe.append("md $OUT/{}".format(path_prefix))

        # Adding *.cpp is a workaround to prevent cp from thrown an error when it
        # encounters a directory (since -r was not specified). If files with an
        # extension other than .cpp need to be copied, then the command below
        # will not work and will need to be updated.
        #
        cmd.append("cp -f {0}/{1}/*.cpp $OUT/{1}/".format("$(location :templated_selective_build_srcs)", path_prefix))
        cmd_exe.append("robocopy /E {0}/{1} $OUT/{1}".format("$(location :templated_selective_build_srcs)", path_prefix))

    cmd.append("mkdir -p $OUT/aten/src/ATen")
    cmd_exe.append("md $OUT/aten/src/ATen")

    # NB: CUDA is skipped here because this is selective build and CUDA is not
    # supported for selective build
    for ufunc_file in aten_ufunc_generated_all_cpu_sources("$(location :gen_aten[{}])"):
        cmd.append("cp -f " + ufunc_file + " $OUT/aten/src/ATen")
        cmd_exe.append("copy " + ufunc_file + " $OUT/aten/src/ATen")

    fb_xplat_genrule(
        name = name,
        cmd = " && ".join(cmd),
        cmd_exe = "@powershell -Command " + ("; ".join(cmd_exe)),
        outs = get_template_registration_files_outs(),
        default_outs = ["."],
    )

def pt_operator_library(
        name,
        ops = [],
        exported_deps = [],
        check_decl = True,
        train = False,
        model = None,
        include_all_operators = False,
        **kwargs):
    model_name = name

    if get_build_from_deps_query():
        ops = [op.strip() for op in ops]

        # If ops are specified, then we are in static selective build mode, so we append
        # base ops to this list to avoid additional special case logic in subsequent code.
        if len(ops) > 0:
            ops.extend(PT_BASE_OPS)

        visibility = kwargs.pop("visibility", ["PUBLIC"])

        fb_xplat_genrule(
            name = name,
            out = "model_operators.yaml",
            cmd = (
                "$(exe :gen_operators_yaml) " +
                "{optionally_root_ops} " +
                "{optionally_training_root_ops} " +
                "--rule_name {rule_name} " +
                "--output_path \"${{OUT}}\" " +
                "--model_name {model_name} " +
                "--dep_graph_yaml_path pytorch_op_deps.yaml " +
                "--models_yaml_path all_mobile_model_configs.yaml " +
                #"{optionally_model_versions} " +
                #"{optionally_model_assets} " +
                #"{optionally_model_traced_backends} " +
                "{optionally_include_all_operators}"
            ).format(
                rule_name = name,
                model_name = model_name,
                optionally_root_ops = "--root_ops " + (",".join(ops)) if len(ops) > 0 else "",
                optionally_training_root_ops = "--training_root_ops " + (",".join(ops)) if len(ops) > 0 and train else "",
                #optionally_model_versions = "--model_versions " + (",".join(model_versions)) if model_versions != None else "",
                #optionally_model_assets = "--model_assets " + (",".join(model_assets)) if model_assets != None else "",
                #optionally_model_traced_backends = "--model_traced_backends " + (",".join(model_traced_backends)) if model_traced_backends != None else "",
                optionally_include_all_operators = "--include_all_operators " if include_all_operators else "",
            ),
            labels = ["pt_operator_library"],  # for pt_operator_query_codegen query
            visibility = visibility,
            **kwargs
        )
    else:
        if check_decl:
            pass
            # ensure_ops_are_declared(ops)

        cxx_library(
            name = name,
            compiler_flags = get_pt_compiler_flags(),
            cxx_platform_compiler_flags = get_cpukernel_avx2_flags(),
            exported_deps = exported_deps,
            **kwargs
        )

def compose_platform_setting_list(settings):
    """Settings object:
    os/cpu pair: should be valid key, or at most one part can be wildcard.
    flags: the values added to the compiler flags
    """
    result = []
    for setting in settings:
        result = result.append([
            "^{}-{}$".format(setting["os"], setting["cpu"]),
            setting["flags"],
        ])
    return result

def get_cpukernel_avx2_flags():
    # flags = compose_platform_setting_list([
    #     {
    #         "cpu": "x86_64",
    #         "flags": ["-DHAVE_AVX2_CPU_DEFINITION"],
    #         "os": "macosx",
    #     },
    # ]) if build_cpukernel_avx2() else []
    return []

def build_cpukernel_avx2():
    return not is_arvr_mode()

def get_cpukernel_avx2_deps():
    # flags = compose_platform_setting_list([
    #     {
    #         "cpu": "x86_64",
    #         "flags": ["fbsource//xplat/caffe2:cpukernel_avx2"],
    #         "os": "macosx",
    #     },
    # ]) if build_cpukernel_avx2() else []
    return []