# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Shared fixtures for the mamba test suite.

Pre-compiles the ``checkpointing_ssu`` JIT variants the tests exercise in a
single batched ninja invocation. ``checkpointing_ssu`` is not part of the AOT
``flashinfer-jit-cache`` package, so a fresh CI container otherwise compiles
each variant serially on first touch (~60s each x ~40 variants dominates the
suite's wall time). Batch-building lets ninja schedule all of them in
parallel across cores, and is a near-free no-op when the modules are already
present in the JIT disk cache.
"""

import torch

# Every _CHECKPOINTING_SSU_VARIANT_FIELDS combination requested by
# tests/mamba/test_checkpointing_ssu.py, captured by instrumenting
# flashinfer.mamba.checkpointing_ssu._get_module over a full run of the file.
# Keep in sync when adding parametrizations: a missing entry is harmless (the
# test falls back to a serial first-touch compile), just slower.
_CHECKPOINTING_SSU_VARIANT_FIELDS = (
    "state_dtype",
    "input_dtype",
    "dt_dtype",
    "weight_dtype",
    "matrixA_dtype",
    "stateIndex_dtype",
    "state_scale_dtype",
    "dim",
    "dstate",
    "npredicted",
    "max_window",
    "heads_per_group",
    "num_groups",
    "philox_rounds",
    "enable_pdl",
)
_CHECKPOINTING_SSU_VARIANTS = [
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        4,
        8,
        16,
        1,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        4,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        6,
        16,
        1,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        6,
        16,
        1,
        0,
        True,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        8,
        16,
        1,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        8,
        16,
        1,
        0,
        True,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        8,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        10,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        16,
        16,
        8,
        2,
        0,
        False,
    ),
    (
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        16,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        3,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        4,
        8,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        4,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        6,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        6,
        16,
        1,
        10,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        6,
        16,
        2,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        6,
        16,
        2,
        10,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        8,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        8,
        16,
        64,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        10,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        10,
        16,
        16,
        2,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        16,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        16,
        16,
        16,
        1,
        10,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        16,
        16,
        16,
        2,
        0,
        False,
    ),
    (
        torch.float16,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        16,
        16,
        16,
        2,
        10,
        False,
    ),
    (
        torch.float32,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        4,
        8,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float32,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        4,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float32,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        6,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float32,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        6,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float32,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        8,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float32,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        10,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float32,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        None,
        64,
        128,
        16,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        4,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        1,
        10,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        2,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        2,
        10,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        8,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        8,
        16,
        64,
        1,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        1,
        10,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        2,
        0,
        False,
    ),
    (
        torch.float8_e4m3fn,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        2,
        10,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        1,
        1,
        16,
        1,
        0,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        1,
        1,
        16,
        1,
        5,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        4,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        1,
        0,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        1,
        10,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        2,
        0,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        6,
        6,
        16,
        2,
        10,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        8,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        1,
        0,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        1,
        10,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        2,
        0,
        False,
    ),
    (
        torch.int8,
        torch.bfloat16,
        torch.bfloat16,
        torch.bfloat16,
        torch.float32,
        torch.int32,
        torch.float32,
        64,
        128,
        16,
        16,
        16,
        2,
        10,
        False,
    ),
]

# Only prewarm when a substantial chunk of the checkpointing suite is about to
# run; a developer running one or two tests should not wait for 43 modules.
_PREWARM_MIN_TESTS = 8


def _triton_supports_current_arch() -> bool:
    """Whether the installed Triton's PTX toolchain can target this GPU.

    The mamba references are Triton kernels; when Triton's bundled
    nvptxcompiler predates the device (e.g. torch-pinned Triton on sm_107),
    ``make_ptx`` calls ``abort()`` and takes the whole pytest process down —
    unskippable at test time, so probe the matching ``ptxas`` binary in a
    subprocess instead. Defaults to True on any probe uncertainty.
    """
    import glob
    import os
    import subprocess

    try:
        if not torch.cuda.is_available():
            return True
        major, minor = torch.cuda.get_device_capability()

        import triton

        troot = os.path.dirname(triton.__file__)
        cands = glob.glob(os.path.join(troot, "backends", "nvidia", "bin", "ptxas"))
        if not cands:
            return True
        for gpu_name in (f"sm_{major}{minor}a", f"sm_{major}{minor}"):
            r = subprocess.run(
                [cands[0], "--gpu-name", gpu_name, os.devnull, "-o", os.devnull],
                capture_output=True,
                text=True,
                timeout=30,
            )
            err = (r.stderr or "") + (r.stdout or "")
            if "not defined for option 'gpu-name'" not in err:
                return True
        return False
    except Exception:
        return True


def pytest_collection_modifyitems(config, items):
    if items and not _triton_supports_current_arch():
        import pytest

        skip_triton = pytest.mark.skip(
            reason="installed Triton cannot target this GPU architecture "
            "(the mamba reference implementations are Triton kernels and "
            "abort the process at compile time)"
        )
        for item in items:
            item.add_marker(skip_triton)
        return

    n_checkpointing = sum(
        1 for item in items if "test_checkpointing_ssu" in item.nodeid
    )
    if n_checkpointing < _PREWARM_MIN_TESTS:
        return

    import concurrent.futures
    import contextlib
    import os
    import warnings

    from flashinfer.jit.mamba.checkpointing_ssu import gen_checkpointing_ssu_module

    # The prewarm is purely an optimization: any failure below degrades to
    # the normal serial first-touch JIT compile at test time (where a real
    # build error will surface attributed to the test that needs the module),
    # and must never abort collection.
    try:
        # Skip variants whose .so already exists (warm JIT cache / AOT package).
        specs = [
            spec
            for spec in (
                gen_checkpointing_ssu_module(
                    **dict(zip(_CHECKPOINTING_SSU_VARIANT_FIELDS, variant, strict=True))
                )
                for variant in _CHECKPOINTING_SSU_VARIANTS
            )
            if not spec.get_library_path().exists()
        ]
    except Exception as exc:  # noqa: BLE001 — best-effort prewarm
        warnings.warn(
            f"checkpointing_ssu prewarm skipped (spec generation failed): {exc}",
            stacklevel=2,
        )
        return
    if not specs:
        return

    # Build missing modules through the standard per-module path (spec.build),
    # parallelized with a thread pool. Do NOT use build_jit_specs() here: its
    # single batched ninja invocation records commands in the cache root's
    # .ninja_log, while each module's build.ninja sets builddir to the module
    # directory — so the per-test build_and_load() would find no log entries
    # for the batch-built outputs and recompile every module serially again.
    #
    # Concurrency: total concurrent compiler processes ~= pool workers x
    # per-module ninja jobs (MAX_JOBS). Scale both from the core count so
    # small CI runners are not oversubscribed, and treat a pre-set MAX_JOBS
    # (the documented total ninja-job budget) as a cap.
    prev_max_jobs = os.environ.get("MAX_JOBS")
    total_jobs = os.cpu_count() or 4
    if prev_max_jobs:
        with contextlib.suppress(ValueError):
            total_jobs = min(total_jobs, max(1, int(prev_max_jobs)))
    max_workers = max(1, total_jobs // 2)
    os.environ["MAX_JOBS"] = str(max(1, total_jobs // max_workers))
    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
            futures = {pool.submit(spec.build, verbose=False): spec for spec in specs}
            for fut in concurrent.futures.as_completed(futures):
                try:
                    fut.result()
                except Exception as exc:  # noqa: BLE001 — best-effort prewarm
                    warnings.warn(
                        f"checkpointing_ssu prewarm build failed for "
                        f"{futures[fut].name}; affected tests fall back to "
                        f"serial first-touch JIT: {exc}",
                        stacklevel=2,
                    )
    finally:
        if prev_max_jobs is None:
            os.environ.pop("MAX_JOBS", None)
        else:
            os.environ["MAX_JOBS"] = prev_max_jobs
