From 91cefdb52aa87265a31af263eaed989e793842e9 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Tue, 7 Jul 2026 20:43:52 +0300 Subject: [PATCH] hcq2: rewrite patch split (#16907) * x * new patch splitter * style * x * x --- extra/hcq2/hcq2.py | 166 ++++++++++++++++++++++++--------------------- 1 file changed, 90 insertions(+), 76 deletions(-) diff --git a/extra/hcq2/hcq2.py b/extra/hcq2/hcq2.py index af6c23a3a1..be672dc587 100644 --- a/extra/hcq2/hcq2.py +++ b/extra/hcq2/hcq2.py @@ -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), ])