from __future__ import annotations

from .base import DeviceBackend


class MusaBackend(DeviceBackend):
    name = "musa"
    visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
    torch_namespace = "musa"
    flaggems_vendor = "musa"
    # torch_musa supplies Stream, Event, synchronize, graph and MUSAGraph, so
    # capture works once the timing path takes them from Triton's device
    # interface instead of torch.cuda. Measured on an S5000 over 15 operators:
    # 11 captured, 3 fell back cleanly with a marker (dropout/rand on the RNG
    # seed, sort on a capture-illegal stream op), 1 poisoned the context (see
    # capture_deny_ops). Capture is worth having because it removes launch
    # overhead the direct timer cannot: batch_norm 16x16x8x48 went 0.0758ms ->
    # 0.0089ms, tanh/gelu 4096x4096 about 3.3x, while GEMMs stayed at 1.00x.
    supports_cuda_graph = True
    supports_mm_cluster_fix = False
    supports_triton_hooks = True
    # MUSA reports its own wording; none of the CUDA markers match it.
    fatal_error_markers = (
        "musa error: an illegal memory access",
        "out of memory",
        "device-side assert triggered",
        "misaligned address",
        "uncorrectable ecc",
    )
    # "unknown error" during capture is not recoverable on this runtime: the
    # fallback runs, and the process then dies in teardown with zero rows. Treat
    # it as poisoning so the operator is reported as capture-incompatible rather
    # than silently losing its measurement.
    poisoning_error_markers = ("musa error: unknown error",)
    # Keyed by FLAGGEMS_PERF_CURRENT_OP, which is the pytest FUNCTION name minus
    # `test_` -- `unique2`, `perf_sort` -- not the benchmark file name. Getting
    # that wrong is silent: the entry simply never matches.
    #
    # Observed on an S5000: capturing test_unique2 takes two recoverable
    # fallbacks and then raises "MUSA error: unknown error", ending the run at
    # rc=1 with no measured rows. Skipping the attempt costs that operator's
    # timing precision instead of its measurement. Revisit if a newer torch_musa
    # changes the behaviour.
    capture_deny_ops = ("unique2",)
    # Several FlagGems benchmarks skip on `not torch.cuda.is_available()` while
    # the code under test dispatches through flag_gems.device. On MThreads that
    # turns a working operator into "not measured"; the plugin drops those
    # guards, pinned to one file and one exact skip reason each.
    compat_plugins = ("_musa_benchmark_compat_plugin",)
    # Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
    # a vendor CLI here; confirm the real listing tool before adding it.
    list_command = None