zl_bench/_bespoke_shape_plugin.py

"""pytest plugin: make the shape yaml authoritative for ops whose inputs are
*bespoke* (built inline / from TestParam objects / a sampled kit) and so are
unreachable by `_shape_inject_plugin` and `_shape_iter_inject_plugin`.

A per-op registry; each handler REUSES the class's own input construction and
only swaps the shape source for our yaml (so we never measure a different
kernel). Covered: upsample_bicubic2d_aa_backward, flash_mla_sparse_fwd,
fused_deepseek_..._quant_insert, cutlass_scaled_mm. Ops with no yaml entry fall
through unchanged. Couples to upstream internals; a mismatch surfaces as a loud
benchmark FAIL, not silent corruption.
"""

from __future__ import annotations

import os
import sys

import yaml as _yaml

_YAML_CACHE: dict | None = None


def _load_yaml(path: str) -> dict:
    if not path or not os.path.isfile(path):
        return {}
    try:
        with open(path, "r") as f:
            return _yaml.safe_load(f) or {}
    except Exception as exc:
        print(f"[bespoke-shape-plugin] failed to load {path}: {exc}", file=sys.stderr)
        return {}


def _yaml_shapes(op_name):
    entry = (_YAML_CACHE or {}).get(op_name)
    if isinstance(entry, dict):
        shapes = entry.get("shapes")
        if shapes:
            return shapes
    return None


# --- per-op builders: reuse the class's own construction, swap the shape source ---

def _upsample_builder(self, dtype, shapes, original):
    # _cfgs entries are (N,C,Hi,Wi,Ho,Wo,ac,label); yaml carries the 7 numerics.
    self._cfgs = [tuple(shape) + ("yaml",) for shape in shapes]
    yield from original(self, dtype)


def _fused_deepseek_builder(self, dtype, shapes, original):
    TestParam = original.__globals__["TestParam"]
    for shape in shapes:
        num_tokens, num_heads, num_tokens_insert, block_size, max_pos, eps = shape
        param = TestParam(
            num_tokens=num_tokens,
            num_heads=num_heads,
            num_tokens_insert=num_tokens_insert,
            block_size=block_size,
            max_pos=max_pos,
            eps=eps,
        )
        yield from type(self).make_input(param)


def _flash_mla_sparse_builder(self, dtype, shapes, original):
    TestParam = original.__globals__["TestParam"]
    for shape in shapes:
        s_q, s_kv, topk, h_q, d_qk = shape
        param = TestParam(
            s_q=s_q, s_kv=s_kv, topk=topk, h_q=h_q, d_qk=d_qk,
            have_attn_sink=True,  # every hardcoded case uses True
        )
        yield from type(self).make_input_flashmla(param)


_WRAP_BUILDERS = {
    "UpsampleBicubic2dAaBackwardBenchmark": _upsample_builder,
    "FusedDeepseekV4QnormRopeKVRopeQuantInsertBenchmark": _fused_deepseek_builder,
    "FlashmlaSparseBenchmark": _flash_mla_sparse_builder,
}


def _wrap_get_input_iter(cls, builder) -> None:
    original = cls.get_input_iter

    def patched(self, dtype):
        shapes = _yaml_shapes(getattr(self, "op_name", None))
        if shapes:
            yield from builder(self, dtype, shapes, original)
        else:
            yield from original(self, dtype)

    cls.get_input_iter = patched


def _patch_cutlass(cls) -> bool:
    """Redirect CutlassScaledMMPerfKit's hardcoded mnk to the yaml M,N,K, reusing
    the kit's own combination/sampling pipeline (the quant-mode sweep stays
    intrinsic)."""
    shapes = _yaml_shapes("cutlass_scaled_mm")
    if not shapes:
        return False
    kit = getattr(sys.modules.get(cls.__module__), "CutlassScaledMMPerfKit", None)
    if kit is None:
        return False
    import torch
    from itertools import product

    mnk = [tuple(shape) for shape in shapes]

    def _get_all_combinations():
        # Mirror upstream verbatim except mnk (from yaml); keep in sync with
        # CutlassScaledMMPerfKit._get_all_combinations.
        scale_shape_types = ["scalar", "vector", "matrix"]
        if_use_bias = [True, False]
        dtypes = [(torch.int8, torch.float16), (torch.float8_e4m3fn, torch.bfloat16)]
        return product(mnk, scale_shape_types, scale_shape_types, if_use_bias, dtypes)

    kit._get_all_combinations = staticmethod(_get_all_combinations)
    return True


def pytest_collection_finish(session):
    global _YAML_CACHE
    from benchmark import base as fg_base
    from benchmark.conftest import Config as fg_Config

    _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")

    def _subclasses(cls):
        for sub in cls.__subclasses__():
            yield sub
            yield from _subclasses(sub)

    covered = []
    seen = set()
    for cls in _subclasses(fg_base.Benchmark):
        if cls in seen:
            continue
        seen.add(cls)
        builder = _WRAP_BUILDERS.get(cls.__name__)
        if builder is not None and "get_input_iter" in cls.__dict__:
            _wrap_get_input_iter(cls, builder)
            covered.append(cls.__name__)
        if cls.__name__ == "CutlassScaledMMBenchmark" and _patch_cutlass(cls):
            covered.append("CutlassScaledMMBenchmark(mnk)")

    print(f"[bespoke-shape-plugin] yaml-driven inputs for: {', '.join(covered) or 'none'}",
          file=sys.stderr, flush=True)