From fd76ac992e9a0271815e4f1123584a3347685048 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 9 Jun 2026 18:36:01 -0700 Subject: [PATCH] cstyle renderer is new style [pr] (#16484) * cstyle new style * switch cstyle renderer to new style * fix hip * fixes * fix webgpu * correct webgpu is_packed * fix dsp * fixes * fix Ops.RANGE must be CONST * old style render access * this is correct * fix cstyle to good * dl/dr * as array * fix spec * remove define_local/define_reg * buffer in shrink * fix test_tiny * all tests fix * param args aren't realized * wgsl fix * work * new gate * fix opencl qcom * process replay * sort order * fix render index --- .github/workflows/benchmark.yml | 2 +- tinygrad/codegen/__init__.py | 2 +- tinygrad/renderer/cstyle.py | 105 ++++++++++++++++++++++---------- tinygrad/renderer/nir.py | 1 + tinygrad/renderer/wgsl.py | 44 +++++++------ tinygrad/uop/__init__.py | 9 ++- tinygrad/uop/validate.py | 4 +- 7 files changed, 105 insertions(+), 62 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index cfb5117aa0..2f674ea463 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -356,7 +356,7 @@ jobs: - name: Train MNIST run: time PYTHONPATH=. DEV=NV TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py - name: Run 10 CIFAR training steps - run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py + run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=130 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py - name: Run 10 CIFAR training steps w HALF run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py - name: Run 10 CIFAR training steps w BF16 diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index ee7c7be49f..b8e2cb42b3 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -41,7 +41,7 @@ pm_remove_vec_dtypes = PatternMatcher([ # no LOADs on register dtypes (UPat(Ops.LOAD, name="x"), lambda x: x.src[0] if x.src[0].addrspace == AddrSpace.REG else None), # remove all vec dtypes - (UPat(GroupOp.All-{Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}, name="x"), + (UPat(GroupOp.All-{Ops.PARAM, Ops.BUFFER, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}, name="x"), lambda x: x.replace(dtype=x.dtype.base.scalar().base)), # replace DEFINE_LOCAL/DEFINE_REG with BUFFER (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="x"), lambda x: diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index da18b9c78b..a1993acf98 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -10,9 +10,8 @@ from tinygrad.codegen.late.devectorizer import no_vectorized_alu base_rewrite = PatternMatcher([ - # defines - (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.max_numel()}];"), - (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.max_numel()}];"), + # local/reg buffers + (UPat(Ops.BUFFER, name="x"), lambda ctx,x: ctx.render_buffer(x)), # range/if/endif (UPat(Ops.RANGE, name="x"), @@ -21,11 +20,10 @@ base_rewrite = PatternMatcher([ (UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"), # casting - (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_dtype(x.dtype)})" \ - if x.max_numel() > 1 and not isinstance(x.dtype, PtrDType) else None), + (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \ + if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None), (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"), - (UPat(Ops.BITCAST, name="x"), lambda ctx,x: - f"__builtin_bit_cast({ctx.render_dtype(x.dtype)}, ({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"), + (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"), # GPU stuff (UPat(Ops.BARRIER), lambda ctx: ctx.barrier), @@ -47,18 +45,18 @@ base_rewrite = PatternMatcher([ # default const render (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)), - # movement ops - (UPat.var("buf").index(UPat.var('idx')), lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"), + # SHRINK/INDEX + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), + (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), (UPat(Ops.STACK, name="x"), - lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_dtype(x.dtype))}" + \ - f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"), - (UPat(Ops.GEP, name="x"), lambda ctx,x: ctx[x.src[0]] + \ - (f"[{x.arg[0]}]" if x.src[0].max_numel() > ctx.gep_arr_threshold else f".{'xyzwabcd'[x.arg[0]]}")), + lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_type(x))}" + \ + f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"), # load/store - (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"), - (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"), - (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"), + (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"({ctx.render_access(bidx)})"), + (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), + lambda ctx,bidx,var,gate: f"({ctx[gate]}?{ctx.render_access(bidx)}:{ctx[var]})"), + (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"{ctx.render_access(bidx)} = {ctx[var]};"), # alu/gep (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"), @@ -106,13 +104,22 @@ pm_manual_bf16_cast = PatternMatcher([ (UPat(Ops.CAST, dtype=dtypes.bfloat16, src=(UPat.var("x", dtype=dtypes.float),)), cast_float_to_bf16), ]) -def uops_to_dtypes(uops:list[UOp]) -> list[DType]: return dedup(u.dtype for u in uops if not isinstance(u.dtype, (ImageDType, PtrDType))) +def uops_to_dtypes(uops:list[UOp]) -> list[DType]: + ret = [] + seen = set() + for u in uops: + if u.addrspace in (AddrSpace.REG, None) and u.dtype != dtypes.void and u._shape is not None and (key:=(u.dtype, u.max_numel())) not in seen: + # TODO: this eventually needs to be removed + ret.append(u.dtype.vec(u.max_numel())) + seen.add(key) + return ret # (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes) def wmma_args(uops:list[UOp]): return dedup((uop.arg[0], uop.arg[1], uop.arg[2], uop.dtype.scalar(), *(uop.arg[4:8])) for uop in uops if uop.op is Ops.WMMA) class CStyleLanguage(Renderer): + new_style = True kernel_typedef: str = "void" buffer_prefix: str = "" buffer_suffix: str = "" @@ -146,8 +153,8 @@ class CStyleLanguage(Renderer): tmp = "" if any(isinstance(u.dtype, ImageDType) for _,(u,_) in bufs): tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n" - buftypes = [(name, self.render_dtype(u.dtype, mutable)+self.buffer_suffix if isinstance(u.dtype, (ImageDType, PtrDType)) else - self.arg_int_prefix if u.dtype == dtypes.int else None) for name,(u,mutable) in bufs] + buftypes = [(name, self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable)+self.buffer_suffix \ + if u.addrspace == AddrSpace.GLOBAL else self.arg_int_prefix if u.dtype == dtypes.int else None) for name,(u,mutable) in bufs] local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"] launch_bounds = prod([d.vmax for d in local_dims]) prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] + @@ -155,16 +162,49 @@ class CStyleLanguage(Renderer): [") {\n" + tmp] + ['\n'.join(kernel), "\n}"]) return prg if prefix is None else "\n".join(prefix)+f"\n{prg}" - def render_cast(self, u:UOp, val:str) -> str: return f"({self.render_dtype(u.dtype)})({val})" + def render_index(self, x:UOp, buf:UOp, idx:UOp): + if buf.addrspace == AddrSpace.REG and buf.op not in {Ops.AFTER, Ops.BUFFER}: + # this is lane access in C + assert idx.op is Ops.CONST, f"{idx.op} must be CONST" + return self[buf]+(f"[{idx.arg}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.arg]}") + ptr = f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})" + if buf.addrspace != AddrSpace.REG: return ptr + # REG buffers have no LOAD, so the access is rendered at the INDEX. the cast handles vector access, same as render_access + return f"(*(({self.render_type(x)}*)({ptr})))" if x.max_numel() > 1 else f"(*{ptr})" + + def render_buffer(self, x:UOp): + shp = x.src[0].as_shape + lanes = 1 + prefix = f"{self.smem_align}{self.smem_prefix}" if x.addrspace == AddrSpace.LOCAL else "" + suffix = f"[{shp[0]}]" if len(shp) else "" + if len(shp) > 1: + # for DEFINE_REG, if it's a 2-D shape it's the number of lanes + assert isinstance(shp[1], int) + lanes = shp[1] + return f"{prefix}{self._render_dtype(x.dtype, sz=lanes)} {self[x]}{suffix};" + + def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.REG, mutable=True): + if isinstance(dtype, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t" + prefix, suffix = "", "" + if addrspace in (AddrSpace.LOCAL, AddrSpace.GLOBAL): + if addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast: prefix = self.smem_prefix + if addrspace == AddrSpace.GLOBAL: prefix = self.buffer_prefix + suffix = "*" + if sz > 1: + return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name).replace(" ", "_") + str(sz) + suffix + return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name) + suffix + + def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace) + def render_access(self, u:UOp): + if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL): + if u.max_numel() > 1: return f"*(({self.render_type(u)})({self[u]}))" + else: return f"*{self[u]}" + return self[u] + def render_cast(self, u:UOp, val:str) -> str: return f"({self.render_type(u)})({val})" + + # LEGACY def render_dtype(self, dt:DType, mutable=True) -> str: - if isinstance(dt, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t" - if isinstance(dt, PtrDType): - prefix = "" - if dt.addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast: prefix = self.smem_prefix - if dt.addrspace == AddrSpace.GLOBAL: prefix = self.buffer_prefix - return prefix + self.render_dtype(dt.base) + "*" - if dt.count > 1: return self.type_map.get(scalar:=dt.scalar(), scalar.name).replace(" ", "_") + str(dt.count) - return self.type_map.get(scalar:=dt.scalar(), scalar.name) + return self._render_dtype(dt, dt.count, dt.addrspace if isinstance(dt, PtrDType) else AddrSpace.REG) def __getitem__(self, key): return self.r[key] # hacky helper def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[UOp,bool]]]]: @@ -181,6 +221,7 @@ class CStyleLanguage(Renderer): name = "test" for u in uops: if u.op in {Ops.NOOP, Ops.GROUP}: continue + if u.op == Ops.STACK and len(u.src) == 0: continue if u.op is Ops.AFTER: r[u] = r[u.src[0]] continue @@ -199,7 +240,7 @@ class CStyleLanguage(Renderer): if u.op is Ops.SPECIAL: r[u] = u.arg elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.arg[-1]]}idx"+range_str(u) else: - prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const", + prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const", Ops.BUFFER: "buf", Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.STACK: "cast", Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu") r[u] = f"{prefix}{c[prefix]}" @@ -208,14 +249,14 @@ class CStyleLanguage(Renderer): assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}" if u.op in {Ops.ENDIF, Ops.END}: depth -= 1 - if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \ + if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \ (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG) or \ (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \ (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))): r[u] = l else: - if u.op not in {Ops.RANGE, Ops.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG} and u.dtype != dtypes.void: - l = f"{self.render_dtype(u.dtype)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "") + if u.op not in {Ops.RANGE, Ops.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG, Ops.BUFFER} and u.dtype != dtypes.void: + l = f"{self.render_type(u)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "") kernel.append(" "*depth + l) if prefix: c[prefix] += 1 # if it was used, increment if u.op in {Ops.IF, Ops.RANGE}: depth += 1 diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 16cbad6265..443ee4876b 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -269,6 +269,7 @@ _nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2 lambda b,img,idx_y,idx_x,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load"))) class IR3Renderer(NIRRenderer, OpenCLRenderer): + new_style = False has_aux = True def nload_img(ctx,img,idx_y,idx_x): diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 3117803355..e7efeb4008 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -1,4 +1,4 @@ -from tinygrad.dtype import DType, PtrDType, dtypes, truncate, AddrSpace +from tinygrad.dtype import DType, dtypes, truncate, AddrSpace from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite, extra_pm from tinygrad.helpers import strip_parens @@ -27,13 +27,12 @@ def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None, gate:UOp|Non val = (load.cast(dtypes.uint32) >> shift_am) & mask return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype) -def is_packed(dt:DType, odt:DType|None = None) -> bool: - if odt is None: odt = dt - # registers aren't packed - if isinstance(odt, PtrDType) and odt.addrspace == AddrSpace.REG: return False - return dt.itemsize < 4 and dt.base != dtypes.half -def _packed_size(dt:PtrDType): return dt.size // (4//dt.itemsize) if is_packed(dt) else dt.size - +def is_packed(x:UOp): + if x.op is Ops.LOAD: dt, addrspace = x.dtype, x.src[0].addrspace + elif x.op is Ops.STORE: dt, addrspace = x.src[1].dtype, x.src[0].addrspace + else: dt, addrspace = x.dtype, x.addrspace + return dt.itemsize < 4 and dt != dtypes.half and addrspace != AddrSpace.REG +def _packed_size(u:UOp): return u.max_numel() // (4//u.dtype.itemsize) if is_packed(u) else u.max_numel() def is_nan(a): bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype) return (a.bitcast(getattr(dtypes, f"uint{bs}")) & ((1 << (bs - 1)) - 1)) > (((1 << exp) - 1) << mant) @@ -43,12 +42,11 @@ wgsl_matcher = PatternMatcher([ lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)), # TODO: load alt value doesnt have to be a const (UPat.load(UPat.var("b"), UPat.cvar("c"), UPat.var("gate"), name="l"), - lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l.dtype, b.dtype) else None), - (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None), - (UPat.store(UPat.var("bidx"), UPat.var("var"), UPat.var("gate")), - lambda bidx,var,gate: packed_store(bidx,var,gate) if is_packed(var.dtype, bidx.dtype) else None), - (UPat.store(UPat.var("bidx"), UPat.var("var")), - lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype, bidx.dtype) else None), + lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l) else None), + (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l,b,l.dtype) if is_packed(l) else None), + (UPat.store(UPat.var("b"), UPat.var("var"), UPat.var("gate"), name="s"), + lambda b,var,gate,s: packed_store(b,var,gate) if is_packed(s) else None), + (UPat.store(UPat.var("b"), UPat.var("var"), name="s"), lambda b,var,s: packed_store(b,var) if is_packed(s) else None), (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None), # fix nan check: 'a != a -> is_nan()' @@ -73,8 +71,8 @@ class WGSLRenderer(CStyleLanguage): (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"), lambda x: f"bitcast({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"), (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}"), - (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{ctx.buf_map(x.dtype.base)},{_packed_size(x.dtype)}>;"), - (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{_packed_size(x.dtype)}>;"), + (UPat(Ops.BUFFER, name="x"), lambda ctx,x: + f"var{'' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"), (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)), lambda ctx,x: f"bitcast>({ctx[x.src[0]]})[0]"), (UPat(Ops.BITCAST, dtype=dtypes.uchar, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}&0xFF)"), @@ -86,20 +84,20 @@ class WGSLRenderer(CStyleLanguage): (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]})"), # TODO: load alt value doesnt have to be a const (UPat.load(UPat.var("b"), UPat.cvar("v"), UPat.var("gate")), - lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[gate]})"), - (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)), + lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b], b.src[0])}, {ctx[gate]})"), + (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b)), (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\ # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1] - f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \ + f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b) \ else f"{ctx[b]} = {ctx[v]};"), (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))), lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"), ]) + base_rewrite def render_cast(self, u:UOp, val: str) -> str: return f"{self.type_map[u.dtype]}({val})" - def render_dtype(self, dt:DType, mutable=True) -> str: return "var" - def render_load(self, x:str, dt:DType) -> str: return f"atomicLoad(&{x})" if is_packed(dt) else x - def buf_map(self, dt:DType) -> str: return "atomic" if is_packed(dt) else self.type_map[dt.base] + def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.REG, mutable=True): return "var" + def render_load(self, x:str, u:UOp) -> str: return f"atomicLoad(&{x})" if is_packed(u) else x + def buf_map(self, u:UOp) -> str: return "atomic" if is_packed(u) else self.type_map[u.dtype.base] def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str: local_size = [u.src[0].ssimplify() for u in sorted([u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == 'l'], key=lambda u: u.arg)] if not local_size: local_size = [1] @@ -111,7 +109,7 @@ class WGSLRenderer(CStyleLanguage): prg += "@group(0) @binding(0)\nvar INFINITY : f32;\n" prg += "\n".join((external_local_bufs or [])+[f"@group(0) @binding({next(bind_it)+1})" + f"{'var' if u.addrspace == AddrSpace.GLOBAL else 'var'}" + - f"{name}:{f'array<{self.buf_map(u.dtype.base)}>' if u.addrspace == AddrSpace.GLOBAL else self.buf_map(u.dtype)};" for name,(u,_) in bufs]) + f"{name}:{f'array<{self.buf_map(u)}>' if u.addrspace == AddrSpace.GLOBAL else self.buf_map(u)};" for name,(u,_) in bufs]) prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3," return prg + "@builtin(local_invocation_id) lindex: vec3) {\n" + "\n".join(kernel) + "\n}" diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 3a6c34b518..3311218e7e 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -22,6 +22,9 @@ class Ops(FastEnum): # define LOCAL/REG allocate things DEFINE_LOCAL = auto(); DEFINE_REG = auto() + # BUFFER is the new LOCAL/REG + BUFFER = auto() + # ** 2 -- non op uops ** # uops that aren't rendered @@ -49,7 +52,7 @@ class Ops(FastEnum): # ** 3 -- load/store ** # INDEX is a BinaryOp similar to ADD, but it operates on pointers - INDEX = auto() + INDEX = auto(); SHRINK = auto() # load/store before math LOAD = auto(); STORE = auto() @@ -99,10 +102,10 @@ class Ops(FastEnum): CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto() # buffer ops - STAGE = auto(); COPY = auto(); BUFFER = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto() + STAGE = auto(); COPY = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto() # the core 6 movement ops! these only exist in the tensor graph - RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto() + RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); FLIP = auto() MULTI = auto() # MULTI is really a movement op # reduce diff --git a/tinygrad/uop/validate.py b/tinygrad/uop/validate.py index 6c58be3d3a..51c638c578 100644 --- a/tinygrad/uop/validate.py +++ b/tinygrad/uop/validate.py @@ -51,8 +51,8 @@ z3_renderer = PatternMatcher([ ]) def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]: - # gate on any upstream INDEX as a replacement for PtrDType - lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op is not Ops.INDEX and \ + # gate on upstream AFTER/BUFFER as a replacement for PtrDType, but keep INDEX as an unknown LOAD + lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM} and \ (x.dtype.scalar() in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK)))[:-1] z3map: dict[UOp, z3.ExprRef] = {} for u in lst: