geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

hcq2: rewrite patch split (#16907)

Browse Source 
* x

* new patch splitter

* style

* x

* x
This commit is contained in:
nimlgen
2026-07-07 20:43:52 +03:00
committed by GitHub
parent 0035bb6fa8
commit 91cefdb52a
1 changed files with 90 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/hcq2/hcq2.py
+90 -76
View File
@@ -152,8 +152,7 @@ def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp:
return UOp.param(next(UOp.unique_num) if unique else 0, dtype.ptr(size), device=devs).rtag(name or "buf")
def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
base = buf.dtype.base
return UOp(Ops.SHRINK, base, (buf, UOp.const(dtypes.int, off//base.itemsize), UOp.const(dtypes.int, 1))).store(val.cast(dtype or base))
return buf.index(UOp.const(dtypes.int, off//buf.dtype.base.itemsize)).store(val.cast(dtype or buf.dtype.base))
def make_cmdbuf(lin, devs):
blob, patches = b'', []
@@ -161,12 +160,7 @@ def make_cmdbuf(lin, devs):
if s.op is not Ops.CONST: patches.append((len(blob), s))
blob += struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)
buf = make_placeholder(devs, len(blob) // 4, dtypes.uint32)
# pull patches to cmdbuf
afters = dedup(u for _, s in patches for u in s.toposort() if u.op is Ops.AFTER)
deps = tuple(d for p in afters for d in p.src[1:])
cmdbuf = buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *[make_patch(buf, off, s) for off, s in patches], *deps)
return cmdbuf.substitute({p: p.src[0] for p in afters}) if afters else cmdbuf
return buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *[make_patch(buf, off, s) for off, s in patches])
def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))
@@ -383,6 +377,72 @@ pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbu
# *****************
def unwrap_mstack(u): return u.src if u.op is Ops.MSTACK else (u,)
def _is_link_patch(p:UOp, buf:UOp, jit=False) -> bool:
if p.op is not Ops.STORE or p.buf_uop is not buf: return False # this is not a patch :(
assert all(x.op is Ops.PARAM for x in unwrap_mstack(p.buf_uop))
has_loads = any(u.op in (Ops.LOAD, Ops.INDEX) for u in p.src[1].backward_slice)
param_is_input = all(x.tag is None and x.op is Ops.PARAM for x in unwrap_mstack(p.src[1].buf_uop))
return not has_loads and not param_is_input if True else (p.buf_uop.tag in {"program"})
def trim_link_patches(ctx:tuple[bool, list[UOp]], a:UOp) -> UOp|None:
links, kept = partition(a.src[1:], lambda p: _is_link_patch(p, a.src[0], jit=ctx[0]))
# keep all patches from the link-time patches' subtrees in the C code
afters = [u for u in UOp.sink(*links).toposort() if u.op is Ops.AFTER]
ctx[1].extend(UOp.sink(*links).substitute({p: p.src[0] for p in afters}).src)
return a.src[0].after(*kept, *[d for p in afters for d in p.src[1:]]) if links else None
pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])
def split_patches(ctx:bool, call:UOp) -> UOp|None:
body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(ctx, lt_patches:=[]), name=f"trim link-time patches ({call.arg.aux.name})")
lt_srcs = collections.defaultdict(list)
for p in lt_patches: lt_srcs[p.buf_uop].append(p)
return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()]))
pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
# *****************
def _make_getaddrs_sub(call:UOp, gaddrs:list[UOp], name:str):
bare = {g: g.replace(src=(unwrap_after(g.src[0]),)) for g in gaddrs}
order = sorted(dedup(bare.values()), key=lambda g: (g.buf_uop.arg.slot, to_tuple(g.buf_uop.tag)))
b = make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name)
sub = {g: b.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, order.index(gr))).load() for g,gr in bare.items()}
return sub, (b.after(*[make_patch(b, i * b.dtype.base.itemsize, gr) for i,gr in enumerate(order)]),) if order else ()
def rm_rt_getaddrs(call:UOp) -> UOp|None:
if not (gaddrs:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR]): return None
inputs, systems = partition(gaddrs, lambda g: all(x.tag is None for x in unwrap_mstack(g.buf_uop)))
(inpsub, _), (syssub, sysarg) = _make_getaddrs_sub(call, inputs, "inputs"), _make_getaddrs_sub(call, systems, "systems")
return call.replace(src=(call.src[0].substitute(inpsub | syssub), *call.src[1:], *sysarg),
arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(g.buf_uop.arg.slot for g in inputs))))))
pm_rm_rt_getaddrs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_getaddrs)])
# *****************
def replace_params(call:UOp) -> UOp|None:
body, variables, param_ops = call.src[0], call.src[0].variables(), {Ops.PARAM, Ops.MSTACK}
args = dedup([s for u in body.toposort(gate=lambda u: u.op not in param_ops) for s in u.src if s.op in param_ops and s not in variables])
patched, refhold = partition(call.src[1:], lambda x: x.src[0] in args)
by_root = {p.src[0]: p for p in patched}
c_args = [by_root.get(a, a) for a in args]
sub = {unwrap_after(u): UOp.param(i, u.dtype, device=u.device) for i,u in enumerate(c_args)} | \
{v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM}
info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None))
return call.replace(src=(body.substitute(sub), *c_args, *refhold), arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
# *****************
def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
itemsize = bv.src[0].dtype.itemsize if unwrap_after(bv.src[0]).op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
return UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0],), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)
@@ -392,75 +452,27 @@ pm_early_simplify = PatternMatcher([
(UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"),), name="g"), resolve_getaddr_slice),
])
# *****************
def replace_params(call:UOp) -> UOp|None:
gaddrs = [u for u in call.src[0].toposort(enter_calls=False) if u.op is Ops.GETADDR and u.src[0].op is Ops.PARAM and u.src[0].tag is None]
if not gaddrs: return None
idxs:dict[int, int] = {}
for g in gaddrs: idxs.setdefault(g.src[0].arg.slot, len(idxs))
inputs = make_placeholder(call.arg.aux.device, len(idxs), dtypes.uint64, "inputs")
body = call.src[0].substitute({g: inputs.index(UOp.const(dtypes.int, idxs[g.src[0].arg.slot])).load() for g in gaddrs})
return call.replace(src=(body, *call.src[1:], inputs), arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(idxs))))
pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
# *****************
def changes_per_submit(u:UOp) -> bool: return u.op in (Ops.LOAD, Ops.INDEX) or (u.op is Ops.PARAM and u.tag is None)
def is_placeholder(b:UOp) -> bool: return (b.op is Ops.PARAM and b.tag is not None) or (b.op is Ops.MSTACK and all(is_placeholder(x) for x in b.src))
def is_link_patch(s:UOp) -> bool: return is_placeholder(s.buf_uop) and not any(changes_per_submit(u) for u in s.backward_slice)
def trim_link_patches(ctx:list[UOp], a:UOp) -> UOp|None:
links, kept = partition(a.src[1:], is_link_patch)
ctx += links
return a.src[0].after(*kept) if links else None
pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])
def split_patches(call:UOp) -> UOp|None:
# trim link-time patches
body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(lt_patches:=[]), name=f"trim link-time patches ({call.arg.aux.name})")
units:dict[UOp, None] = {}
def unit_gate(u:UOp) -> bool:
if (is_plc:=is_placeholder(u)): units[u] = None
return not is_plc
body.toposort(gate=unit_gate)
srcs = dedup(list(call.src[1:]) + list(units) + [s.buf_uop for s in lt_patches])
param_sub = {u: UOp.param(i, u.dtype, device=u.device) for i,u in enumerate(srcs)}
for b in dedup(s.buf_uop for s in lt_patches):
idx = param_sub[b].arg.slot
srcs[idx] = srcs[idx].after(*dedup(s for s in lt_patches if s.buf_uop is b))
param_sub |= {v: v.replace(arg=replace(v.arg, slot=-1)) for v in body.variables() if v.op is Ops.PARAM}
info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(srcs) if unwrap_after(u).tag == "inputs"), None))
return call.replace(src=(body.substitute(param_sub), *srcs), arg=replace(call.arg, aux=info))
pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
# *****************
# 5.3. pack placeholders buffers
def pack_hcq_placeholders(call:UOp) -> UOp|None:
bufs = [b for b in call.src[0].toposort() if b.op is Ops.PARAM and b.tag in (maxtags:={"scratch"}) | (sumtags:={"program", "kernargs"})]
# def pack_hcq_placeholders(call:UOp) -> UOp|None:
# bufs = [b for b in call.src[0].toposort() if b.op is Ops.PARAM and b.tag in (maxtags:={"scratch"}) | (sumtags:={"program", "kernargs"})]
off_per_buf:dict[UOp, int] = {}
size_per_tag:dict[str, int] = {}
for b in bufs:
bsz = b.max_numel()
if b.tag in maxtags: size_per_tag[b.tag] = max(size_per_tag.get(b.tag, 0), bsz)
elif b.tag in sumtags:
off_per_buf[b] = round_up(size_per_tag.get(b.tag, 0), {"program": 0x1000}.get(b.tag, 128))
size_per_tag[b.tag] = off_per_buf[b] + bsz
# off_per_buf:dict[UOp, int] = {}
# size_per_tag:dict[str, int] = {}
# for b in bufs:
# bsz = b.max_numel()
# if b.tag in maxtags: size_per_tag[b.tag] = max(size_per_tag.get(b.tag, 0), bsz)
# elif b.tag in sumtags:
# off_per_buf[b] = round_up(size_per_tag.get(b.tag, 0), {"program": 0x1000}.get(b.tag, 128))
# size_per_tag[b.tag] = off_per_buf[b] + bsz
count_per_tag = collections.Counter(b.tag for b in bufs)
ref_bufs = {b.tag:b for b in bufs if count_per_tag[b.tag] > 1}
bases = {tag:UOp.new_buffer(b.device, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()}
subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, off_per_buf.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
# count_per_tag = collections.Counter(b.tag for b in bufs)
# ref_bufs = {b.tag:b for b in bufs if count_per_tag[b.tag] > 1}
# bases = {tag:UOp.new_buffer(b.device, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()}
# subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, off_per_buf.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
# return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
# pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
# *****************
# 8. callify hcq programs
@@ -491,8 +503,10 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp:
linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
# pie
linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace params")
linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split rt/lt patches")
linear = graph_rewrite(linear, pm_rm_rt_getaddrs, walk=True, name="replace rt getaddrs")
linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace with args")
linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="early simplify patches", enter_calls=True)
# and compile it
@@ -524,7 +538,7 @@ def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
return UOp(Ops.NOOP)
def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
assert buf.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT)
assert buf.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT), f"{buf.op}"
devs, b = g.arg, buf.buffer
bufs = tuple(cast(Buffer, x.buffer) for x in buf.src) if buf.op is Ops.MSTACK else tuple(b.bufs if isinstance(b, MultiBuffer) else (b,)*len(devs))
@@ -542,7 +556,7 @@ pm_resolve_patches = PatternMatcher([
# folders
(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
(UPat(Ops.SHRINK, src=(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf"), UPat.cvar("off"), UPat(Ops.CONST)))
(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat.cvar("off"))
.store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), fold_const_store),
])