From 6ece327cf3dd6cc55b3d2bf81b557ce3de17b0a9 Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 25 Aug 2026 20:15:06 -0400 Subject: [PATCH] CUSTOM arg is (str, dtype) [PR] (#17737) --- extra/gemm/moe_routing.py | 2 +- .../quantize_fp8_delayed/__init__.py | 2 +- .../llama_kernels/quantize_mxfp4/__init__.py | 2 +- test/null/test_viz.py | 12 ++--- tinygrad/llm/kernels/amd.py | 16 +++---- tinygrad/nn/__init__.py | 2 +- tinygrad/renderer/cstyle.py | 2 +- tinygrad/runtime/ops_qcom.py | 5 +- tinygrad/uop/ops.py | 3 +- tinygrad/uop/spec.py | 5 +- tinygrad/uop/upat.py | 46 ++++++++++--------- 11 files changed, 52 insertions(+), 45 deletions(-) diff --git a/extra/gemm/moe_routing.py b/extra/gemm/moe_routing.py index eb4982c6e6..321f241aa0 100644 --- a/extra/gemm/moe_routing.py +++ b/extra/gemm/moe_routing.py @@ -53,7 +53,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple: g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk)) row = idx.index(g, m).cast(dtypes.weakint) val = gout.index(g, m, j).load().cast(dtypes.float32) - atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=atomic_str) + atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void)) return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=())) grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0] return (None, grad_table.cast(table_u.dtype).uop, None) diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index db9bb5fd68..f7d42b44e1 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -50,7 +50,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state: else: raise NotImplementedError(f"no atomic max for device {device}") amax_idx = amax_out.reshape((1,)).index(UOp.const(0)) max_val = lds[0].load() - atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg) + atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void)) return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=())) @functools.cache diff --git a/extra/llama_kernels/quantize_mxfp4/__init__.py b/extra/llama_kernels/quantize_mxfp4/__init__.py index e68fbb146b..48682b8a10 100644 --- a/extra/llama_kernels/quantize_mxfp4/__init__.py +++ b/extra/llama_kernels/quantize_mxfp4/__init__.py @@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO mem = M*N*2 + M*N + M*N//16 # read bf16, write row+col fp4 + e8m0 outputs = (row_fp4, row_scale, col_fp4, col_scale) sink = UOp.sink(*(o.base for o in outputs), x.base, - *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg="") for o in outputs), + *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg=("", dtypes.void)) for o in outputs), UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"), arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem))) src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text() diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 9812cefda2..06ce496137 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -185,18 +185,18 @@ class TestViz(unittest.TestCase): @dataclass(frozen=True) class TestStruct: colored_field: str - a = UOp(Ops.CUSTOM, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue"))) + a = UOp(Ops.PYLITERAL, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue"))) a2 = uop_to_json(VizData(), a)[id(a)] - self.assertEqual(ansistrip(a2["label"]), f"CUSTOM\n{TestStruct.__qualname__}(colored_field='xyz12345')") + self.assertEqual(ansistrip(a2["label"]), f"PYLITERAL\n{TestStruct.__qualname__}(colored_field='xyz12345')") def test_colored_label_multiline(self): with save_viz() as viz: arg = colored("x", "green")+"\n"+colored("y", "red")+colored("z", "yellow")+colored("ww\nw", "magenta") src = [Tensor.empty(1).uop for _ in range(10)] - a = UOp(Ops.CUSTOM, src=tuple(src), arg=arg) + a = UOp(Ops.PYLITERAL, src=tuple(src), arg=arg) exec_rewrite(a, [PatternMatcher([])]) a2 = next(viz.get_details(0, 0))["graph"][id(a)] - self.assertEqual(ansistrip(a2["label"]), "CUSTOM\nx\nyzww\nw") + self.assertEqual(ansistrip(a2["label"]), "PYLITERAL\nx\nyzww\nw") def test_inf_loop(self): a = UOp.const(3) @@ -347,7 +347,7 @@ class TestVizGC(unittest.TestCase): init = bufs_allocated() a = UOp.new_buffer("NULL", 10, dtypes.char) a.buffer.allocate() - exec_rewrite(UOp(Ops.CUSTOM, src=(a,), arg=a), [PatternMatcher([])]) + exec_rewrite(UOp(Ops.PYLITERAL, src=(a,), arg=a), [PatternMatcher([])]) del a self.assertEqual(bufs_allocated()-init, 0) lst = viz.list_items() @@ -474,7 +474,7 @@ class TestVizIntegration(unittest.TestCase): def custom_fn(X:UOp): X = X.flatten() i = UOp.range(X.numel(), 0) - custom_op = UOp(Ops.CUSTOMI, src=(X[i],), arg="{} + undeclared_name") + custom_op = UOp(Ops.CUSTOMI, src=(X[i],), arg=("{} + undeclared_name", X.dtype)) return X[i].store(custom_op).end(i).sink(arg=KernelInfo(name=f"custom_fn_{X.numel()}")) x = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), fxn=custom_fn)[0] with save_viz() as viz: diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index b206c3415f..ecb4d5a79e 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -35,8 +35,8 @@ def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool: def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp: for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)): if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load() - other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg= - f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))") + other = UOp(Ops.CUSTOM, src=(val,), arg= + (f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))", dtypes.float)) val = val.maximum(other) if maximum else val + other return val @@ -77,14 +77,14 @@ class Linear(nn.Linear): return super().__call__(x) def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp: - return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c), arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)") + return UOp(Ops.CUSTOMI, src=(a.int(), b.int(), c), arg=("__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)", dtypes.int32)) def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp: - return UOp(Ops.CUSTOMI, dtypes.uint32, tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg="__builtin_amdgcn_perm({}, {}, {})") + return UOp(Ops.CUSTOMI, src=tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg=("__builtin_amdgcn_perm({}, {}, {})", dtypes.uint32)) def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp: assert ptr.op is Ops.INDEX - if lanes is None: return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})") + if lanes is None: return UOp(Ops.CUSTOMI, src=(ptr,), arg=("__builtin_nontemporal_load({0})", ptr.dtype)) buf, coords = ptr.src[0], ptr.src[1:] idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0)) return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(dtype=ptr.dtype) @@ -191,7 +191,7 @@ def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int): output_waves = 2 if out_features % (32*output_tiles) == 0 else 1 token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1) lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL) - hw_lane = UOp(Ops.CUSTOM, dtypes.int32, (lane.int(),), arg="__builtin_amdgcn_mbcnt_lo(-1, 0)").cast(dtypes.weakint) + hw_lane = UOp(Ops.CUSTOM, src=(lane.int(),), arg=("__builtin_amdgcn_mbcnt_lo(-1, 0)", dtypes.int32)).cast(dtypes.weakint) col, half = hw_lane % 16, hw_lane // 16 outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles)) inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16)) @@ -201,8 +201,8 @@ def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int): def _wmma_stores(out, outputs, tokens, accs, update, half): def values(acc:UOp) -> tuple[UOp, ...]: vals = tuple(acc.after(update)[i].load() for i in range(8)) - swapped = tuple(UOp(Ops.CUSTOM, dtypes.float32, (value,), - arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))") for value in vals) + swapped = tuple(UOp(Ops.CUSTOM, src=(value,), + arg=("__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))", dtypes.float32)) for value in vals) low = half.eq(0) return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2)) return [out[token, output].store(value) for output,output_accs in zip(outputs, accs) diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index e98186111e..9b663935a4 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -359,7 +359,7 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple: if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);" elif device == "AMD": atomic_arg = "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" else: raise NotImplementedError(f"no atomics for device {device}") - atomic = UOp(Ops.CUSTOM, src=(grad_weight.index(local_token_id, j_idx), grad_val), arg = atomic_arg) + atomic = UOp(Ops.CUSTOM, src=(grad_weight.index(local_token_id, j_idx), grad_val), arg=(atomic_arg, dtypes.void)) return atomic.end(i, j_outer, j_inner).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=())) grad_weight_uop = grad_weight_uop.custom_kernel(grad_emb, idx, fxn=_embedding_bwd_kernel)[0] diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index ac9f2eda85..f139ce9dc2 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -71,7 +71,7 @@ base_rewrite = PatternMatcher([ f"({', '.join(f'({ctx.render_type(y)})({ctx[y]})' for y in x.src[1:])}))" + (";" if x.dtype is dtypes.void else "")), # custom passes through with format - (UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg.format(*[ctx[y] for y in x.src])), + (UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg[0].format(*[ctx[y] for y in x.src])), ]) def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True): diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index f69a14c554..7b12d8cc82 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -22,8 +22,9 @@ def dcache_flush(): from tinygrad.codegen import to_program buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU) i = UOp.range(n, 0, dtype=dtypes.int) - flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");') - sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"), tag=1) + flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg=('__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");', dtypes.void)) + sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg=('__asm__ volatile("dsb sy" ::: "memory");', dtypes.void)), + arg=KernelInfo(name="dcache_flush"), tag=1) prg = to_program(sink, Device["CPU"].renderer) return Device["CPU"].runtime(prg.to_elf()) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 2ff57a458e..359c55c15e 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -123,7 +123,8 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None: # a CALL of an opaque body is void, a CALL of an address can return a value return dtypes.void if src[0].dtype is dtypes.void else None case Ops.CUSTOM | Ops.CUSTOMI: - return None + assert isinstance(arg, tuple) and len(arg) == 2 and isinstance(arg[1], DType), f"CUSTOM/CUSTOMI arg must be (str, DType), got {arg}" + return arg[1] case Ops.INS: return None case Ops.NOOP: diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 09b9c17c90..3fa820c480 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -100,8 +100,9 @@ spec_shared = PatternMatcher([ Ops.AFTER, Ops.UNSHARD, Ops.BITCAST, Ops.INS})),), allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)), - # CUSTOM (inline and non inline) - (UPat((Ops.CUSTOMI, Ops.CUSTOM)), lambda: True), + # CUSTOM (inline and non inline): the arg is the source string and the dtype it produces, void for a bare statement + (UPat((Ops.CUSTOMI, Ops.CUSTOM), name="x"), + lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and isinstance(x.arg[0], str) and isinstance(x.arg[1], DType)), # CALL of an external function (UPat(Ops.CALL, src=(UPat(),), allow_any_len=True, name="x"), diff --git a/tinygrad/uop/upat.py b/tinygrad/uop/upat.py index 9bbcab452f..55ab01eaeb 100644 --- a/tinygrad/uop/upat.py +++ b/tinygrad/uop/upat.py @@ -2,6 +2,7 @@ from typing import Any, Callable import itertools, inspect, functools, types from tinygrad.helpers import partition, dedup, Context from tinygrad.uop.ops import UPat, UOp, Ops, PatternMatcher, graph_rewrite, deconstruct_function +from tinygrad.dtype import dtypes class UPatCompileError(Exception): pass @@ -18,40 +19,42 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp: # build the and_clause for acceptance and_clause:list[UOp] = [] if self.op is not None: - if len(self.op) > 1: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(int(x) for x in self.op))), arg="{0}.op in {1}")) - else: and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.op == "+str(self.op[0].value))) + if len(self.op) > 1: + and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(int(x) for x in self.op))), arg=("{0}.op in {1}", dtypes.void))) + else: and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("{0}.op == "+str(self.op[0].value), dtypes.void))) if self.arg is not None: - if isinstance(self.arg, int): and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.arg == "+str(int(self.arg)))) - else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.arg)), arg="{0}.arg == {1}")) + if isinstance(self.arg, int): and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("{0}.arg == "+str(int(self.arg)), dtypes.void))) + else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.arg)), arg=("{0}.arg == {1}", dtypes.void))) if self.strict_length or self.required_len > 0: - and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len)))) - if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.CUSTOMI, arg=self.name), base))) + and_clause.append(UOp(Ops.CUSTOM, src=(base,), + arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len), dtypes.void))) + if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.CUSTOMI, arg=(self.name, dtypes.void)), base))) if self.match_dtype is not None: if len(self.match_dtype) > 1: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(self.match_dtype))), - arg="{0}.dtype in {1}")) + arg=("{0}.dtype in {1}", dtypes.void))) else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.match_dtype[0])), - arg="{0}.dtype == {1}")) + arg=("{0}.dtype == {1}", dtypes.void))) if self.match_tag is not None: if len(self.match_tag) > 1: - and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(self.match_tag))), arg="{0}.tag in {1}")) - else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.match_tag[0])), arg="{0}.tag == {1}")) + and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(self.match_tag))), arg=("{0}.tag in {1}", dtypes.void))) + else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.match_tag[0])), arg=("{0}.tag == {1}", dtypes.void))) if self.src is not None: # single match if len(self.src) == 1 and isinstance(self.src[0], tuple): and_clause += [_get_clause(s, base.index(i), depth) for i,s in enumerate(self.src[0])] # repeat match elif len(self.src) == 1 and isinstance(self.src[0], itertools.repeat): - it = UOp(Ops.CUSTOMI, arg=f"ituop{depth}") + it = UOp(Ops.CUSTOMI, arg=(f"ituop{depth}", dtypes.void)) match = _get_clause(next(self.src[0]), it, depth+1) - and_clause.append(UOp(Ops.CUSTOM, src=(match, it, base), arg="all([{0} for {1} in {2}.src])")) + and_clause.append(UOp(Ops.CUSTOM, src=(match, it, base), arg=("all([{0} for {1} in {2}.src])", dtypes.void))) # multi match (fork) elif len(self.src) > 1 and all(isinstance(x, tuple) for x in self.src): fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.index(i), depth) for i,s in enumerate(ss)])) for ss in self.src] and_clause.append(UOp(Ops.OR, src=tuple(fork_cond))) else: raise RuntimeError("broken") - return UOp(Ops.AND, src=tuple(and_clause)) if and_clause else UOp(Ops.CUSTOMI, arg="True") + return UOp(Ops.AND, src=tuple(and_clause)) if and_clause else UOp(Ops.CUSTOMI, arg=("True", dtypes.void)) # *** pattern matcher *** @@ -91,7 +94,7 @@ def do_process_and(a:UOp) -> UOp|None: for store in stores: if store.src[0] in dict_stores: # duplicate store is an identity compare - new_src.append(UOp(Ops.CUSTOM, src=(dict_stores[store.src[0]], store.src[1]), arg="{0} is {1}")) + new_src.append(UOp(Ops.CUSTOM, src=(dict_stores[store.src[0]], store.src[1]), arg=("{0} is {1}", dtypes.void))) found = True else: dict_stores[store.src[0]] = store.src[1] @@ -108,17 +111,18 @@ pm_proc = PatternMatcher([(UPat(Ops.AND, name="a"), do_process_and)], compiled=F # renderer def wrap(ctx, x) -> UOp: ctx[ret:=f"a{len(ctx)}"] = x.arg - return UOp(Ops.CUSTOMI, arg=ret) + return UOp(Ops.CUSTOMI, arg=(ret, dtypes.void)) pm_renderer = PatternMatcher([ (UPat(Ops.PYLITERAL, name="x"), wrap), # AND of CUSTOMI fragments inside a CUSTOM becomes a single CUSTOMI (joined with " and ") (UPat(Ops.CUSTOM, src=(UPat(Ops.AND, src=UPat(Ops.CUSTOMI), name="x"), UPat(), UPat()), name="r"), - lambda r,x: r.replace(src=(UOp(Ops.CUSTOMI, arg="(" + ' and '.join(y.arg for y in x.src) + ")"),)+r.src[1:])), + lambda r,x: r.replace(src=(UOp(Ops.CUSTOMI, arg=("(" + ' and '.join(y.arg[0] for y in x.src) + ")", dtypes.void)),)+r.src[1:])), - (UPat(Ops.CUSTOM, src=UPat(Ops.CUSTOMI), name="x"), lambda x: UOp(Ops.CUSTOMI, arg=x.arg.format(*[y.arg for y in x.src]))), - (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"), lambda x,c,g: x.replace(arg=x.arg+f".src[{c.val}]")) + (UPat(Ops.CUSTOM, src=UPat(Ops.CUSTOMI), name="x"), lambda x: UOp(Ops.CUSTOMI, arg=(x.arg[0].format(*[y.arg[0] for y in x.src]), dtypes.void))), + (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"), + lambda x,c,g: x.replace(arg=(x.arg[0]+f".src[{c.val}]", dtypes.void))) ], compiled=False) def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]: @@ -131,8 +135,8 @@ def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]: for ss in s.src: or_pieces.extend(_final_render(ss, has_ctx, depth+1)) elif s.op is Ops.STORE: assert s.src[0].op is Ops.CUSTOMI and s.src[1].op is Ops.CUSTOMI - store_pieces.append(f"{s.src[0].arg}={s.src[1].arg}") - elif s.op is Ops.CUSTOMI: and_pieces.append(s.arg) + store_pieces.append(f"{s.src[0].arg[0]}={s.src[1].arg[0]}") + elif s.op is Ops.CUSTOMI: and_pieces.append(s.arg[0]) else: raise UPatCompileError(f"can't compile this {s}") # if we have an or, render it if len(or_pieces): @@ -145,7 +149,7 @@ def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]: return [f"{' '*depth}if {and_clause}: return _ret"] def _get_code(self:UPat, has_ctx:bool): - ret = _get_clause(self, UOp(Ops.CUSTOMI, arg="uop")) + ret = _get_clause(self, UOp(Ops.CUSTOMI, arg=("uop", dtypes.void))) try: # TODO: this should be tracked in a "system" rewrite, not untracked or tracked with kernel with Context(TRACK_MATCH_STATS=0):