File size: 4,709 Bytes
52a1e9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """Kernel canonicalization + dedup for the search loop.
The advisor's correction (recorded in CLAUDE.md): in `discovery_specialist`, diversity IS
the objective and `mech_of` powers a novelty-or-die REWARD. Here the objective is the
single FASTEST CORRECT kernel — diversity is just search exploration. So the analog of
`mech_of` lives HERE as **dedup**, NOT as a reward:
- `khash(src)` : a hash that is invariant to comments, whitespace, docstrings and
local-variable renaming, so the search never spends a (costly) harness
evaluation on a kernel it already measured. This is the
"AST-normalize then hash" move from the LLVM-peephole plan.
- `knobs(src)` : extract the performance-relevant features (BLOCK size, num_warps,
num_stages, autotune present, ops used) — useful for steering mutation
and for the open-trace, never as a correctness/speed signal.
Nothing in this file ever decides correctness or speed — that is the harness's sole job.
"""
from __future__ import annotations
import ast
import hashlib
import re
class _Canon(ast.NodeTransformer):
"""Alpha-rename local names to positional ids and strip docstrings, so two kernels that
differ only cosmetically canonicalize identically."""
def __init__(self):
self._names: dict[str, str] = {}
def _id(self, name: str) -> str:
# keep dunders / library names (torch, triton, tl, range, tl.* attrs) stable;
# only canonicalize ordinary identifiers the author chose.
if name in _KEEP or name.startswith("__"):
return name
if name not in self._names:
self._names[name] = f"v{len(self._names)}"
return self._names[name]
def visit_Name(self, node):
node.id = self._id(node.id)
return node
def visit_arg(self, node):
node.arg = self._id(node.arg)
return node
def visit_FunctionDef(self, node):
# drop a leading string-literal docstring
if (node.body and isinstance(node.body[0], ast.Expr)
and isinstance(getattr(node.body[0], "value", None), ast.Constant)
and isinstance(node.body[0].value.value, str)):
node.body = node.body[1:]
self.generic_visit(node)
return node
# identifiers that carry meaning across kernels — do NOT rename these
_KEEP = {"torch", "triton", "tl", "run", "range", "len", "float", "int", "tuple",
"constexpr", "self", "True", "False", "None"}
def canonical_src(src: str) -> str:
"""Normalized source: parse -> strip docstrings + alpha-rename locals -> unparse.
Falls back to comment/whitespace stripping if the source doesn't parse standalone."""
try:
tree = ast.parse(src)
tree = _Canon().visit(tree)
ast.fix_missing_locations(tree)
return ast.unparse(tree)
except Exception:
# not parseable in isolation (rare) — coarse normalize: drop comments + collapse ws
no_comments = re.sub(r"#.*", "", src)
return re.sub(r"\s+", " ", no_comments).strip()
def khash(src: str) -> str:
return hashlib.sha1(canonical_src(src).encode()).hexdigest()[:16]
_INT = r"(\d+)"
def knobs(src: str) -> dict:
"""Performance-relevant features (for steering mutation + the open-trace, NOT scoring)."""
def _find(pat, default=None, cast=int):
m = re.search(pat, src)
return cast(m.group(1)) if m else default
return {
"block": _find(rf"BLOCK\s*=\s*{_INT}"),
"block_n": _find(rf"BLOCK_N\s*=\s*{_INT}"),
"num_warps": _find(rf"num_warps\s*=\s*{_INT}"),
"num_stages": _find(rf"num_stages\s*=\s*{_INT}"),
"autotune": "autotune" in src,
"uses_rsqrt": "rsqrt" in src,
"uses_max_sub": bool(re.search(r"-\s*\w*max", src)) or "maximum" in src,
"n_loads": src.count("tl.load"),
"n_stores": src.count("tl.store"),
}
if __name__ == "__main__":
# demo: two cosmetically-different RMSNorm kernels canonicalize to the SAME hash;
# a kernel with a different BLOCK is a genuinely distinct candidate (different hash).
a = open(__file__.rsplit("/", 1)[0] + "/seed_kernels/rmsnorm.py").read()
b = "# totally different comment\n" + a.replace("acc", "accum").replace("row", "r")
c = a.replace("BLOCK = 1024", "BLOCK = 2048")
print("rmsnorm.py khash:", khash(a))
print("renamed+recommented khash:", khash(b), "(== a)" if khash(a) == khash(b) else "(MISMATCH!)")
print("BLOCK 1024->2048 khash:", khash(c), "(distinct)" if khash(a) != khash(c) else "(collision!)")
print("knobs(a):", knobs(a))
|