load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target")
oncall("executorch")
# Any targets that should be shared between fbcode and xplat must be defined in
# targets.bzl. This file can contain xplat-only targets.

load(":targets.bzl", "define_common_targets")


non_fbcode_target(_kind = define_common_targets,)

# !!!! fbcode/executorch/extension/llm/custom_ops/TARGETS was merged into this file, see https://fburl.com/workplace/xl8l9yuo for more info !!!!

# Any targets that should be shared between fbcode and xplat must be defined in
# targets.bzl. This file can contain fbcode-only targets.

load("@fbsource//tools/target_determinator/macros:ci.bzl", "ci")
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
load(":targets.bzl", "define_common_targets")


fbcode_target(_kind = define_common_targets,)

fbcode_target(_kind = runtime.python_test,
    name = "test_sdpa_with_kv_cache",
    srcs = [
        "test_sdpa_with_kv_cache.py",
    ],
    preload_deps = [
        ":custom_ops_aot_lib_mkl_noomp",
        ":custom_ops_aot_py",
    ],
    deps = [
        "//caffe2:torch",
        "//executorch/extension/pybindings:portable_lib",
    ],
)

fbcode_target(_kind = runtime.python_test,
    name = "test_update_cache",
    srcs = [
        "test_update_cache.py",
    ],
    preload_deps = [
        ":custom_ops_aot_lib",
        ":custom_ops_aot_py",
    ],
    deps = [
        "//caffe2:torch",
    ],
)

fbcode_target(_kind = runtime.python_test,
    name = "test_gated_delta",
    srcs = [
        "test_gated_delta.py",
    ],
    # Builds the runtime, kernels and custom_ops with -DET_EVENT_TRACER_ENABLED so
    # the .pte-execution test can assert the op's per-op native_call_ ETDump event.
    labels = [] if runtime.is_oss else ci.labels(
        ci.buckconfig("executorch.event_tracer_enabled", "true"),
    ),
    preload_deps = [
        ":custom_ops_aot_lib",
        ":custom_ops_aot_py",
    ],
    deps = [
        "//caffe2:torch",
        "//executorch/devtools:lib",
        "//executorch/exir:lib",
        "//executorch/runtime:runtime",
    ],
)

fbcode_target(_kind = runtime.python_test,
    name = "test_preprocess_custom_ops",
    srcs = [
        "test_preprocess_custom_ops.py",
    ],
    preload_deps = [
        ":preprocess_custom_ops_py",
    ],
    deps = [
        "//caffe2:torch",
    ],
)

fbcode_target(_kind = runtime.python_test,
    name = "test_quantized_sdpa",
    srcs = [
        "test_quantized_sdpa.py",
    ],
    preload_deps = [
        ":custom_ops_aot_lib_mkl_noomp",
        ":custom_ops_aot_py",
    ],
    deps = [
        "//caffe2:torch",
        "//executorch/extension/pybindings:portable_lib",
    ],
)

fbcode_target(_kind = runtime.python_test,
    name = "test_update_cross_attn_cache",
    srcs = [
        "test_update_cross_attn_cache.py",
    ],
    preload_deps = [
        ":custom_ops_aot_lib",
        ":custom_ops_aot_py",
    ],
    deps = [
        "//caffe2:torch",
    ],
)

fbcode_target(_kind = runtime.python_library,
    name = "bench_gated_delta_rule_lib",
    srcs = ["bench_gated_delta_rule.py"],
    base_module = "executorch.extension.llm.custom_ops",
    deps = ["//caffe2:torch"],
)

fbcode_target(_kind = runtime.python_binary,
    name = "bench_gated_delta_rule",
    main_module = "executorch.extension.llm.custom_ops.bench_gated_delta_rule",
    preload_deps = [
        ":custom_ops_aot_lib_mkl_noomp",
        ":custom_ops_aot_py",
    ],
    deps = [":bench_gated_delta_rule_lib"],
)

fbcode_target(_kind = runtime.python_test,
    name = "test_quantized_moe",
    srcs = [
        "test_quantized_moe.py",
    ],
    preload_deps = [
        ":custom_ops_aot_lib_mkl_noomp",
        ":custom_ops_aot_py",
        # The Python source transform calls
        # `torch.ops.torchao._pack_8bit_act_4bit_weight` to pack each
        # expert's INT4 weights, so preload the torchao AOT library too.
        "//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten",
    ],
    deps = [
        "//caffe2:torch",
        "//executorch/extension/pybindings:portable_lib",
        "//executorch/examples/models/llama:llama_transformer",
        "//executorch/examples/models/llama:transformer_modules",
        "//executorch/examples/models/llama:source_transformation",
        "//executorch/examples/models/llama:export_library",
        "//executorch/extension/llm/export/config:llm_config",
    ],
)
