#!/usr/bin/env python3

"""Recon-tree tool of a sealed binary reference (GENERATED by `forge init --binary-ref`,

copied in as binref.py; see forgekernel.binref).

The reference is a SOL-ExecBench / FlashInfer-Bench solution compiled into

ref_bin/<module>.so the way SOL-ExecBench builds a cuda_cpp solution (torch

cpp_extension; sm_100a unless the solution names an arch; CUTLASS headers; only the

solution's top-level C/C++/CUDA files are compiled, the rest are there to #include). Its

sources live under src/ and never leave this tree: the card only gets the .so.

python3 binref.py build compile src/ into ref_bin/<module>.so; prints SO_* lines

python3 binref.py check load the .so, run it once on the card's inputs

python3 binref.py crosscheck judge the .so with the card's own judge: must PASS

python3 binref.py bench time it on the card's clock (the card's --mode ref JSON)

python3 binref.py ncu NCU window: warm calls, then calls inside profiler start/stop

"""

from __future__ import annotations

import hashlib

import json

import os

import shutil

import sys

from pathlib import Path

HERE = Path(__file__).resolve().parent

META = json.loads((HERE / "binref.json").read_text())

SO = HERE / "ref_bin" / f"{META['module']}.so"

BUILD = Path(os.environ.get("BINREF_BUILD_DIR", "/workspace/binref_build")) / META["module"]

CPP_SUFFIXES = (".cu", ".cpp", ".cc", ".cxx", ".c")

BLACKWELL_GENCODE = "-gencode=arch=compute_100a,code=sm_100a"

# SOL-ExecBench CompileOptions defaults (core/data/solution.py): a solution that gives

# compile_options but leaves a key out gets these. -lcuda is what lets driver-API calls

# (cuTensorMapEncodeTiled for TMA descriptors) resolve; without it the .so fails to import.

SOL_DEFAULT_CUDA_CFLAGS = ["-O3", "--use_fast_math"]

SOL_DEFAULT_LD_FLAGS = ["-lcuda"]

# SOL-ExecBench's build_ext: CUTLASS_DIR, else the image's /usr/local/cutlass; include/ and

# tools/util/include/ go on the include path whatever the solution's dependencies say.

SOL_DEFAULT_CUTLASS_DIR = "/usr/local/cutlass"

def compiled_sources(sources: list[str]) -> list[str]:

"""The sources SOL-ExecBench's build_ext compiles: the C/C++/CUDA files at the top of the

solution directory (``HERE.iterdir()`` + ``is_file()``). A file in a subdirectory is only

there to be #included -- a spliced solution wraps each component's ``cN/kernel.cu`` in a

top-level ``cN__kernel.cu``, and two components' ``kernel.cu`` would otherwise be two

ninja rules for one ``kernel.cuda.o``."""

return [r for r in sources if "/" not in r.replace("\\", "/") and Path(r).suffix in CPP_SUFFIXES]

def cutlass_dir(env=None) -> tuple[str, str | None]:

"""(the CUTLASS root the build puts on the include path, why its headers are missing or

None). The root is SOL-ExecBench's: $CUTLASS_DIR, else /usr/local/cutlass."""

env = os.environ if env is None else env

root = env.get("CUTLASS_DIR") or SOL_DEFAULT_CUTLASS_DIR

if os.path.isfile(os.path.join(root, "include", "cutlass", "cutlass.h")):

return root, None

return root, (f"CUTLASS headers not found under {root}/include (CUTLASS_DIR "

+ (f"= {env['CUTLASS_DIR']}" if env.get("CUTLASS_DIR") else

f"unset, the default {SOL_DEFAULT_CUTLASS_DIR}") + ")")

def compile_flags(opts: dict | None) -> tuple[list[str], list[str], list[str]]:

"""(cuda_cflags, cflags, ld_flags) as SOL-ExecBench's build_ext resolves them: no

compile_options at all -> no extra flags; otherwise a missing key takes its default."""

if opts is None:

return [], [], []

def pick(key: str, default: list[str]) -> list[str]:

v = opts.get(key)

return list(default) if v is None else list(v)

return (pick("cuda_cflags", SOL_DEFAULT_CUDA_CFLAGS), pick("cflags", []),

pick("ld_flags", SOL_DEFAULT_LD_FLAGS))

def _lb():

sys.path.insert(0, str(HERE / "eval"))

import leaderboard # the card's judge, inputs and clock (a verbatim copy)

return leaderboard

def build() -> int:

"""Stage the sealed sources in BUILD, compile, copy the .so into ref_bin/ -- and remove

BUILD on every way out: it sits outside the tree's clone dir (a payload sync does not

wipe it) and holds the sources and the compile options in the clear, on a machine other

trees' jobs may run on (a group card's pricing shares this tree's GPU machine)."""

try:

return _build()

finally:

shutil.rmtree(BUILD, ignore_errors=True)

def _build() -> int:

import torch

import torch.utils.cpp_extension as ext

if BUILD.exists():

shutil.rmtree(BUILD)

BUILD.mkdir(parents=True)

for rel in META["sources"]:

dst = BUILD / rel

dst.parent.mkdir(parents=True, exist_ok=True)

shutil.copy2(HERE / "src" / rel, dst)

cuda, cflags, ldflags = compile_flags(META.get("compile_options"))

if not any("-gencode" in f or "-arch" in f for f in cuda):

cuda = [BLACKWELL_GENCODE] + cuda

cutlass, missing = cutlass_dir()

if missing: # a source that includes CUTLASS / CuTe fails below on this node: say why up front

print(f"[binref] WARNING {missing}: a source that includes cutlass/ or cute/ "

f"headers cannot compile on this node", flush=True)

top = compiled_sources(META["sources"])

if not top:

print("[binref] no C/C++/CUDA source at the top of the solution (SOL-ExecBench compiles "

"only those; files in subdirectories are #included)")

return 1

ext.load(name=META["module"],

sources=[str(BUILD / r) for r in top],

extra_cuda_cflags=list(cuda), extra_cflags=list(cflags), extra_ldflags=list(ldflags), # load appends in place

extra_include_paths=[str(BUILD), f"{cutlass}/include", f"{cutlass}/tools/util/include"],

build_directory=str(BUILD), verbose=True, is_python_module=True)

built = sorted(BUILD.glob(f"{META['module']}*.so"))

if not built:

print(f"[binref] no {META['module']}*.so produced in {BUILD}")

return 1

SO.parent.mkdir(parents=True, exist_ok=True)

shutil.copy2(built[0], SO)

data = SO.read_bytes()

print(f"SO_PATH {SO}")

print(f"SO_SIZE {len(data)}")

print(f"SO_SHA256 {hashlib.sha256(data).hexdigest()}")

print(f"SO_TORCH {torch.__version__} CUDA {torch.version.cuda}")

print(f"SO_FLAGS {json.dumps(cuda)}")

print(f"SO_LDFLAGS {json.dumps(ldflags)}")

return 0

def _ensure_built() -> None:

if not SO.exists():

rc = build()

if rc:

raise SystemExit(rc)

def check() -> int:

_ensure_built()

lb = _lb()

problem = lb.load_problem()

fn, dps = lb._impl(problem, "ref")

_, custom = lb.reference_callables(problem)

inp = lb.gen_inputs(problem, "cuda", custom_inputs_fn=custom)

outs = lb.call_impl(problem, fn, inp, dps, "cuda", no_grad=True)

print("[binref] check ok:", [tuple(o.shape) for o in outs])

return 0

def crosscheck() -> int:

_ensure_built()

lb = _lb()

problem = lb.load_problem()

fn, dps = lb._impl(problem, "ref")

verdict, _ = lb.judge_workloads(problem, fn, dps, "cuda") # every workload of a group card

print(f"[binref] crosscheck {'PASS' if verdict['pass'] else 'FAIL'}: {verdict['status']}"

+ (f" — {verdict.get('error')}" if not verdict["pass"] else ""))

return 0 if verdict["pass"] else 1

def bench() -> int:

_ensure_built()

return _lb().run_leaderboard_bench(["--mode", "ref"])

def ncu() -> int:

_ensure_built()

_lb().ncu_window(which="ref")

return 0

if __name__ == "__main__":

cmd = sys.argv[1] if len(sys.argv) > 1 else ""

fns = {"build": build, "check": check, "crosscheck": crosscheck, "bench": bench, "ncu": ncu}

if cmd not in fns:

print(__doc__)

raise SystemExit(2)

raise SystemExit(fns[cmd]())