load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
load("@fbsource//third-party/llama.cpp:ggml_ci.bzl", "ggml_ci_labels")
load("@fbsource//xplat/executorch/backends/native/runtime/engine:targets.bzl", "CPU_COMPAT", "CUDA_COMPAT", "HIP_COMPAT")
load("@fbcode_macros//build_defs/lib:re_test_utils.bzl", "re_test_utils")
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")

oncall("executorch")

fbcode_target(
    _kind = runtime.cxx_test,
    name = "ggml_engine_test",
    srcs = ["test_ggml_engine.cpp"],
    deps = [
        "//executorch/backends/native/runtime:native_graph_schema",
        "//executorch/backends/native/runtime:runtime",
        "//executorch/backends/native/runtime/deserialize:owned_bytes",
        "//executorch/backends/native/runtime/deserialize:package",
        "//executorch/backends/native/runtime/deserialize:package_test_data",
        "//executorch/backends/native/runtime/engine:engine",
        "//executorch/backends/native/runtime/engine:ggml_engine_cpu",
    ],
    **CPU_COMPAT
)

# The op suite runs once per ggml backend set; only the engine variant, the
# runner's --device, link deps and resources, the test host, and the CI
# configuration differ.
_VARIANTS = [
    {"suffix": "", "engine": "ggml_engine_cpu", "compat": CPU_COMPAT, "device": "cpu"},
    {
        "suffix": "_cuda",
        "engine": "ggml_engine_cuda",
        "compat": CUDA_COMPAT,
        "labels": ggml_ci_labels("cuda"),
        "device": "gpu",
        "remote_execution": re_test_utils.remote_execution(
            platform = "gpu-remote-execution",
        ),
        # The CUDA backend is a shared library that needs the ggml core shared
        # beside it, and its VMM pool calls the CUDA driver API.
        "deps": ["fbsource//third-party/llama.cpp:ggml_base[prefer-shared]"],
        "external_deps": [("cuda", None, "cuda_driver_api-lazy")],
        # Buck does not put the CUDA driver stub's library on the runner's
        # loader path, so it ships beside the runner.
        "resources": {
            "//tools/lazify:lazy_stub_binder[shared]": "libtools_lazify_lazy_stub_binder.so",
        },
    },
    {
        "suffix": "_hip",
        "engine": "ggml_engine_hip",
        "compat": HIP_COMPAT,
        "labels": ggml_ci_labels("hip"),
        "device": "gpu",
        "remote_execution": re_test_utils.remote_execution(
            platform = "amd-gpu",
            subplatform = "MI300",
        ),
    },
]

[
    fbcode_target(
        _kind = runtime.cxx_binary,
        name = "ggml_ptn_runner" + v["suffix"],
        srcs = ["ggml_ptn_runner.cpp"],
        visibility = ["//executorch/backends/native/test/..."],
        deps = [
            "//executorch/backends/native/runtime/engine:" + v["engine"],
            "//executorch/backends/native/test/runtime:ptn_runner_main",
        ] + v.get("deps", []),
        external_deps = ["gflags"] + v.get("external_deps", []),
        labels = v.get("labels", []),
        **v.get("compat", {})
    )
    for v in _VARIANTS
]

[
    fbcode_target(
        _kind = runtime.python_test,
        name = "test_ggml_ops" + v["suffix"],
        srcs = ["test_ggml_ops.py"],
        env = {"GGML_PTN_DEVICE": v["device"]},
        labels = v.get("labels", []),
        remote_execution = v.get("remote_execution"),
        resources = {":ggml_ptn_runner" + v["suffix"]: "ggml_ptn_runner"} |
                    v.get("resources", {}),
        deps = [
            "//caffe2:torch",
            "//executorch/exir/native:native",
            "//executorch/extension/llm/export:gguf",
            "//pytorch/ao:torchao",
            "fbsource//third-party/pypi/transformers:transformers",
        ],
        **v.get("compat", {})
    )
    for v in _VARIANTS
]
