DeepInfer/scripts/flaggems_port/screen_ops.py

"""Screen FlagGems ops for the defect class that broke `mul`, so the exclusion list is evidence-based.

`mul` returns the *first operand's shape* whenever broadcasting would expand it:

    (1,64,8) * (4,1,8)  ->  FlagGems (1,64,8),   PyTorch (4,64,8)
    (64,1)   * (1,128)  ->  FlagGems (64,1),     PyTorch (64,128)

Silent, no error. It broke `engram_hash_ref` inside the engine, and it is not dtype-specific.

The traceback went through `flag_gems/utils/pointwise_dynamic.py` and `shape_utils.broadcast_shapes`,
which is shared by every pointwise op -- so the question is which *other* ops inherit it. This script
answers that: each op is run twice, same-shape and broadcasting, and the reference is always taken
before `enable()` so the comparison cannot be FlagGems-against-itself.

Output is a table, which is what the exclusion list in `deepinfer/flaggems_bootstrap.py` is built
from.

Run: PYTHONPATH=. python scripts/flaggems_port/screen_ops.py
"""

from __future__ import annotations

import sys
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

import torch

DEV = "cuda"


@dataclass
class Case:
    op: str
    label: str
    run: Callable[[], Any]
    ref: Any


def ref_of(fn: Callable[[], Any]) -> Any:
    out = fn()
    if isinstance(out, torch.Tensor):
        return out.clone()
    if isinstance(out, (tuple, list)):
        return type(out)(t.clone() if isinstance(t, torch.Tensor) else t for t in out)
    return out


def same(ref: Any, got: Any) -> tuple[bool, str]:
    if isinstance(ref, torch.Tensor):
        if not isinstance(got, torch.Tensor):
            return False, f"type {type(got).__name__}"
        if ref.shape != got.shape:
            return False, f"shape {tuple(ref.shape)} -> {tuple(got.shape)}"
        if ref.dtype != got.dtype:
            return False, f"dtype {ref.dtype} -> {got.dtype}"
        if ref.dtype.is_floating_point:
            d = (ref.float() - got.float()).abs().max().item()
            s = ref.float().abs().max().item()
            tol = 2.0**-8 * max(s, 1e-9) * 2
            return (d <= tol, f"max diff {d:.2e}")
        eq = torch.equal(ref, got)
        if eq:
            return True, "exact"
        return False, f"{(ref != got).sum().item()}/{ref.numel()} wrong"
    if isinstance(ref, (tuple, list)):
        for r, g in zip(ref, got, strict=False):
            ok, msg = same(r, g)
            if not ok:
                return False, msg
        return True, "exact"
    return ref == got, ""


def build_cases() -> list[Case]:
    torch.manual_seed(0)
    b = 64
    n = 8

    # (1, b, n) vs (4, 1, n): the exact shape pattern that broke engram_hash.
    x3 = torch.randn((1, b, n), device=DEV, dtype=torch.bfloat16)
    y3 = torch.randn((4, 1, n), device=DEV, dtype=torch.bfloat16)
    i3 = torch.randint(0, 1000, (1, b, n), device=DEV, dtype=torch.int64)
    j3 = torch.randint(1, 100, (4, 1, n), device=DEV, dtype=torch.int64)

    # (b, 1) vs (1, m): the most common broadcast in transformer code (masks, scales).
    x2 = torch.randn((b, 1), device=DEV, dtype=torch.bfloat16)
    y2 = torch.randn((1, 128), device=DEV, dtype=torch.bfloat16)

    xs = torch.randn((256, 512), device=DEV, dtype=torch.bfloat16)
    ys = torch.randn((256, 512), device=DEV, dtype=torch.bfloat16)

    cases: list[Case] = []

    def add(name: str, label: str, fn: Callable[[], Any]) -> None:
        cases.append(Case(name, label, fn, ref_of(fn)))

    # binary pointwise, same shape (control: isolates broadcast as the cause)
    add("add", "same-shape", lambda: torch.add(xs, ys))
    add("mul", "same-shape", lambda: torch.mul(xs, ys))
    add("sub", "same-shape", lambda: torch.sub(xs, ys))
    add("div", "same-shape", lambda: torch.div(xs, ys))
    add("maximum", "same-shape", lambda: torch.maximum(xs, ys))
    add("minimum", "same-shape", lambda: torch.minimum(xs, ys))

    # binary pointwise, broadcasting (the suspect class)
    add("add", "broadcast 3d", lambda: torch.add(x3, y3))
    add("mul", "broadcast 3d", lambda: torch.mul(x3, y3))
    add("sub", "broadcast 3d", lambda: torch.sub(x3, y3))
    add("div", "broadcast 3d", lambda: torch.div(x3, y3))
    add("maximum", "broadcast 3d", lambda: torch.maximum(x3, y3))
    add("minimum", "broadcast 3d", lambda: torch.minimum(x3, y3))
    add("add", "broadcast 2d", lambda: torch.add(x2, y2))
    add("mul", "broadcast 2d", lambda: torch.mul(x2, y2))
    add("sub", "broadcast 2d", lambda: torch.sub(x2, y2))

    # integer bitwise / remainder: what engram_hash also used
    add("bitwise_xor_", "broadcast int64", lambda: i3.bitwise_xor(j3))
    add("remainder", "broadcast int64", lambda: torch.remainder(i3, j3))

    # scalar variants: same code path, different signature
    add("mul", "scalar", lambda: torch.mul(xs, 3.0))
    add("add", "scalar", lambda: torch.add(xs, 3.0))

    # unary: should not involve broadcasting at all
    add("silu", "unary", lambda: torch.nn.functional.silu(xs))
    add("gelu", "unary", lambda: torch.nn.functional.gelu(xs))
    add("abs", "unary", lambda: torch.abs(xs))

    # reductions
    add("softmax", "last dim", lambda: torch.softmax(xs, dim=-1))
    add("amax", "last dim", lambda: torch.amax(xs, dim=-1))
    add("sum", "last dim", lambda: torch.sum(xs, dim=-1))

    # normalisation; guarded because rms_norm is only in newer torch
    if hasattr(torch.nn.functional, "rms_norm"):
        add("rms_norm", "1x512", lambda: torch.nn.functional.rms_norm(xs, (512,)))

    # more pointwise / reduction candidates for an allow-list
    add("exp", "unary", lambda: torch.exp(xs))
    add("sigmoid", "unary", lambda: torch.sigmoid(xs))
    add("rsqrt", "unary", lambda: torch.rsqrt(xs.abs() + 1))
    add("neg", "unary", lambda: torch.neg(xs))
    add("relu", "unary", lambda: torch.relu(xs))
    add("sqr", "unary", lambda: torch.square(xs))
    add("clamp", "min/max", lambda: torch.clamp(xs, -0.5, 0.5))
    add("pow", "scalar", lambda: torch.pow(xs, 2))
    add("reciprocal", "unary", lambda: torch.reciprocal(xs.abs() + 1))
    add("mean", "last dim", lambda: torch.mean(xs, dim=-1))
    add("cumsum", "last dim", lambda: torch.cumsum(xs, dim=-1))
    add("masked_fill", "scalar", lambda: xs.masked_fill(xs > 0, -1.0))
    add("where", "self", lambda: torch.where(xs > 0, xs, -xs))
    add("eq", "scalar", lambda: torch.eq(xs, 0.0))
    add("lt", "scalar", lambda: torch.lt(xs, 0.0))
    add("ge", "scalar", lambda: torch.ge(xs, 0.0))
    add("nonzero", "unary", lambda: torch.nonzero(xs > 0.5))
    add("gather", "dim0", lambda: torch.gather(xs, 0, torch.zeros((8, 512), device=DEV, dtype=torch.int64)))
    add("index_select", "dim0", lambda: torch.index_select(xs, 0, torch.arange(8, device=DEV)))
    add("sort", "last dim", lambda: torch.sort(xs, dim=-1).values)
    add("sin", "unary", lambda: torch.sin(xs))
    add("cos", "unary", lambda: torch.cos(xs))
    add("sqrt", "unary", lambda: torch.sqrt(xs.abs()))
    add("bitwise_and", "scalar", lambda: torch.bitwise_and(i3, 255))
    add("bincount", "int", lambda: torch.bincount(torch.zeros(64, device=DEV, dtype=torch.int64), minlength=8))

    # shape ops
    add("cat", "2 tensors", lambda: torch.cat([xs, xs], dim=0))
    add("stack", "2 tensors", lambda: torch.stack([xs, xs], dim=0))
    add("repeat_interleave", "self_int 2", lambda: torch.repeat_interleave(xs, 2, dim=0))

    return cases


def main() -> int:
    cases = build_cases()  # references captured here, before enable()
    import flag_gems

    flag_gems.enable()
    registered = sorted(flag_gems.all_registered_ops())
    print(f"flag_gems {flag_gems.__version__}, {len(registered)} ops registered\n")

    def served(op: str) -> bool:
        """Did FlagGems actually take this op?

        Without this the screen lies by omission: `torch.bitwise_xor` has no FlagGems registration at
        all, so a "match" there is PyTorch compared with PyTorch and says nothing about FlagGems. Any
        case that is not served is reported as such and cannot enter an allow-list.
        """
        return any(n == op or n.startswith(f"{op}.") or n.startswith(f"{op}_") for n in registered)

    bad: dict[str, list[str]] = {}
    unserved: set[str] = set()
    rows = []
    for c in cases:
        if not served(c.op):
            unserved.add(c.op)
            rows.append((c.op, c.label, None, "not registered in flag_gems (PyTorch answered)"))
            continue
        try:
            got = c.run()
            torch.cuda.synchronize()
            ok, msg = same(c.ref, got)
        except Exception as exc:  # noqa: BLE001 - an exception is itself a finding
            ok, msg = False, f"raised {type(exc).__name__}: {str(exc)[:70]}"
        rows.append((c.op, c.label, ok, msg))
        if not ok:
            bad.setdefault(c.op, []).append(c.label)

    width = max(len(r[0]) for r in rows) + 2
    for op, label, ok, msg in rows:
        mark = "skip" if ok is None else ("ok  " if ok else "BAD ")
        print(f"  {mark} {op:<{width}} {label:<16} {msg}")

    if unserved:
        print(f"\nnot served by flag_gems (cannot be judged here): {sorted(unserved)}")
    print()
    if bad:
        print("UNSAFE OPS (exclude these):")
        for op, labels in sorted(bad.items()):
            print(f"  {op}: {', '.join(labels)}")
        return 1
    print("all served ops matched")
    return 0


if __name__ == "__main__":
    sys.exit(main())