load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")

oncall("executorch")

fbcode_target(
    _kind = runtime.python_library,
    name = "coalesced_int4_tensor",
    srcs = [
        "coalesced_int4_tensor.py",
    ],
    visibility = [
        "//executorch/...",
    ],
    deps = [
        "//caffe2:torch",
        "//pytorch/ao:torchao",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "dp4a_planar_int5_tensor",
    srcs = [
        "dp4a_planar_int5_tensor.py",
    ],
    visibility = [
        "//executorch/...",
    ],
    deps = [
        "//caffe2:torch",
        "//pytorch/ao:torchao",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "dp4a_planar_int6_tensor",
    srcs = [
        "dp4a_planar_int6_tensor.py",
    ],
    visibility = [
        "//executorch/...",
    ],
    deps = [
        "//caffe2:torch",
        "//pytorch/ao:torchao",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "quantize_op_dispatch",
    srcs = [
        "quantize_op_dispatch/__init__.py",
        "quantize_op_dispatch/_gemm_family_dispatch.py",
        "quantize_op_dispatch/int4_dispatch.py",
        "quantize_op_dispatch/int5_dispatch.py",
        "quantize_op_dispatch/int6_dispatch.py",
        "quantize_op_dispatch/int8_dispatch.py",
    ],
    visibility = [
        "//executorch/...",
    ],
    deps = [
        ":coalesced_int4_tensor",
        ":dp4a_planar_int5_tensor",
        ":dp4a_planar_int6_tensor",
        ":triton_kernels",
        "//caffe2:torch",
        "//pytorch/ao:torchao",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "cuda_passes",
    srcs = [
        "passes/__init__.py",
        "passes/move_cond_predicate_to_cpu.py",
        "passes/lower_offgraph_kv.py",
        "passes/remove_assert_async.py",
        "passes/replace_fused_sdpa.py",
        "passes/replace_int64_floordiv.py",
        "passes/replace_masked_copy_with_where.py",
    ],
    visibility = [
        "//executorch/backends/cuda/...",
    ],
    deps = [
        ":triton_kernels",
        "//caffe2:torch",
        "//executorch/backends/transforms:utils",
        "//executorch/exir/dialects:lib",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "autotune",
    srcs = [
        "autotune/__init__.py",
        "autotune/cuda_graph_timing.py",
        "autotune/inputs.py",
        "autotune/launch_params.py",
    ],
    visibility = [
        "//executorch/...",
    ],
    deps = [
        "//caffe2:torch",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "cuda_backend",
    srcs = [
        "cuda_backend.py",
        "target_smem.py",
        "cuda_weight_collector.py",
    ],
    visibility = ["PUBLIC"],
    deps = [
        ":autotune",
        ":cuda_passes",
        ":triton_replacement_pass",
        "//caffe2:torch",
        "//executorch/backends/aoti/passes:passes",
        "//executorch/exir/_serialize:lib",
        "//executorch/exir/backend:backend_details",
        "//executorch/exir/backend:compile_spec_schema",
        "//executorch/backends/aoti:aoti_backend",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "merge_ptes",
    srcs = [
        "merge_ptes.py",
    ],
    visibility = ["PUBLIC"],
    deps = [
        ":cuda_backend",
        "//caffe2:torch",
        "//executorch/exir/_serialize:lib",
        "//executorch/exir:schema",
        "//executorch/extension/flat_tensor/serialize:serialize",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "cuda_partitioner",
    srcs = [
        "cuda_partitioner.py",
    ],
    visibility = ["PUBLIC"],
    deps = [
        "//caffe2:torch",
        "//executorch/backends/aoti:aoti_partitioner",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "triton_kernels",
    srcs = [
        "triton/kernels/__init__.py",
        "triton/kernels/fused_moe.py",
        "triton/kernels/int4_quantized_gemm.py",
        "triton/kernels/int5_quantized_gemm.py",
        "triton/kernels/int6_quantized_gemm.py",
        "triton/kernels/int8_quantized_gemm.py",
        "triton/kernels/int4_matmul.py",
        "triton/kernels/quantized_gemm_family.py",
        "triton/kernels/quantized_gemm_utils.py",
        "triton/kernels/sdpa.py",
        "triton/kernels/topk.py",
        "triton/kernels/tq4_sdpa.py",
        "triton/kernels/triton_mutation_cache.py",
    ],
    visibility = [
        "//executorch/backends/cuda/...",
    ],
    deps = [
        ":autotune",
        "//caffe2:torch",
    ],
)

fbcode_target(
    _kind = runtime.python_library,
    name = "triton_replacement_pass",
    srcs = [
        "triton/__init__.py",
        "triton/replacement_pass.py",
    ],
    visibility = [
        "//executorch/...",
    ],
    deps = [
        ":triton_kernels",
        "//caffe2:torch",
        "//executorch/exir/dialects:lib",
    ],
)
