sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

CUSTOM arg is (str, dtype) [PR] (#17737)

Browse Source
This commit is contained in:
chenyu
2026-08-25 20:15:06 -04:00
committed by GitHub
parent 07268b724f
commit 6ece327cf3
11 changed files with 52 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/moe_routing.py
+1 -1
View File
@@ -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)
extra/llama_kernels/quantize_fp8_delayed/__init__.py
+1 -1
View File
@@ -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
extra/llama_kernels/quantize_mxfp4/__init__.py
+1 -1
View File
@@ -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()
test/null/test_viz.py
+6 -6
View File
@@ -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:
tinygrad/llm/kernels/amd.py
+8 -8
View File
@@ -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)
tinygrad/nn/__init__.py
+1 -1
View File
@@ -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]
tinygrad/renderer/cstyle.py
+1 -1
View File
@@ -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):
tinygrad/runtime/ops_qcom.py
+3 -2
View File
@@ -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())
tinygrad/uop/ops.py
+2 -1
View File
@@ -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:
tinygrad/uop/spec.py
+3 -2
View File
@@ -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"),
tinygrad/uop/upat.py
+25 -21
View File
@@ -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):