From 4a2e1f10766a2f856ed5222acbfc53b80bbf8e73 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 1 May 2026 15:00:27 -0700 Subject: [PATCH] STORE doesn't have ranges anymore (#16019) * STORE doesn't have ranges anymore * fix --- tinygrad/codegen/late/devectorizer.py | 14 +++++++------- tinygrad/codegen/late/expander.py | 2 +- tinygrad/codegen/simplify.py | 2 +- tinygrad/renderer/cstyle.py | 2 +- tinygrad/renderer/nir.py | 10 +++++----- tinygrad/renderer/ptx.py | 4 ++-- tinygrad/schedule/rangeify.py | 2 +- tinygrad/uop/ops.py | 2 +- tinygrad/uop/symbolic.py | 12 ++++++------ 9 files changed, 25 insertions(+), 25 deletions(-) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 0d40e84356..dd15089e13 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -116,22 +116,22 @@ def fold_expanded_index(midx:UOp): post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret)) return post_cat.gep(tuple(cast(list[int], idxs))) -def cat_after_store(cat:UOp, data:UOp, sto:UOp): +def cat_after_store(cat:UOp, data:UOp): # TODO: this is written in many places offset = 0 ret: list[UOp] = [] for s in cat.src: - ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))), *sto.src[2:])) + ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))))) offset += s.dtype.count return UOp.group(*ret) -def gep_on_store(gep:UOp, st:UOp, sto:UOp): +def gep_on_store(gep:UOp, st:UOp): # NOTE: we need to invert the gep here, but it may be an expanding gep # fake argsort. TODO: handle duplicates a = {} for i,x in enumerate(gep.arg): a[x] = i new_arg = tuple(x[1] for x in sorted(a.items())) - return gep.src[0].store(st.gep(new_arg), *sto.src[2:]) + return gep.src[0].store(st.gep(new_arg)) load_store_folding = PatternMatcher([ (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index), @@ -140,12 +140,12 @@ load_store_folding = PatternMatcher([ (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True), lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)), # GEP on data of STORE - (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), name="sto"), gep_on_store), + (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st"))), gep_on_store), # put PTRCAT after LOAD (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True), lambda cat,ld: UOp(Ops.VCAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))), # put PTRCAT after STORE - (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), name="sto"), cat_after_store), + (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data"))), cat_after_store), ]) # *** correct load/store *** @@ -187,7 +187,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp): if global_offset+fold_length > sz: continue lidx = buf.index((offset + global_offset).valid(mask), ptr=True) if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace)) - if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))+ls.src[2:])) + if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length)))))) else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length))) global_offset += fold_length break diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index 2b57e9768a..d09dcd9b77 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -45,7 +45,7 @@ def do_expand(root:UOp): else: # non-UNROLL input if root.op in range_start and i >= range_start[root.op]: - # for any range args of STORE/REDUCE, pass them through + # for any range args of REDUCE/WMMA/END/etc., pass them through new_srcs.append(src) elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType): new_srcs.append(src) diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 3b37f465f3..145071c312 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -14,7 +14,7 @@ def flatten_range(r:UOp) -> UOp|None: pm_flatten_range = PatternMatcher([ # real ranges only - (UPat((Ops.REDUCE, Ops.STORE, Ops.END), name="r"), flatten_range), + (UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range), ]) def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.IDIV, Ops.MOD} for u in x.backward_slice) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 45a6f7d0c6..9e64a07c92 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -307,7 +307,7 @@ class OpenCLRenderer(CStyleLanguage): (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))),)), lambda ctx,buf,idx: f"read_imagef({ctx[buf]}, smp, {ctx[idx]})"), (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), allow_any_len=True), - UPat.var("var", dtypes.float.vec(4))), allow_any_len=True), + UPat.var("var", dtypes.float.vec(4)))), lambda ctx,buf,idx,var: f"write_imagef({ctx[buf]}, {ctx[idx]}, {ctx[var]});"), ]) + base_rewrite diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 5bc9894a23..8ee9eef396 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -128,8 +128,8 @@ class NIRRenderer(Renderer): # load/store bool -> uint8 (UPat(Ops.LOAD, dtypes.bool, name="x"), lambda x: x.replace(dtype=dtypes.uint8, src=x.src[0:1]+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())+x.src[2:]).cast(dtypes.bool)), - (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True), - lambda x: x.replace(src=x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])), + (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x"), + lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint8)))), # NIR requires shift amount to be 32 bit: https://docs.mesa3d.org/nir/alu.html#nir-alu-op-ishl (UPat((Ops.SHL, Ops.SHR), name="x"), lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint))) if x.src[1].dtype.bitsize != 32 else None), # OpConvertFToU is undefined if Result Type is not wide enough, cast through int32 @@ -146,7 +146,7 @@ class NIRRenderer(Renderer): (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)), (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))), - (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val")), allow_any_len=True), + (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))), lambda ctx,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)), (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True, name="x"), lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate], @@ -268,8 +268,8 @@ class IR3Renderer(NIRRenderer, OpenCLRenderer): return _nload_img(ctx.b, ctx.r[img], ctx.r[coord], img.dtype) def_rewrite = PatternMatcher([ - (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), allow_any_len=True), UPat.var("val")), - allow_any_len=True), lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)), + (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), allow_any_len=True), UPat.var("val"))), + lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)), (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("alt"))), lambda ctx,img,coord,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, coord), lambda: ctx.r[alt])), (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))),)), nload_img), diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index a1564770ab..7a6fb53825 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -48,8 +48,8 @@ ptx_matcher = PatternMatcher([ # load/store bool -> uint8 (UPat(Ops.LOAD, dtypes.bool, src=(UPat(dtype=dtypes.int64),), name="x", allow_any_len=True), lambda x: UOp(x.op, dtypes.uint8, x.src[0:1] + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ()) + x.src[2:]).cast(dtypes.bool)), - (UPat(Ops.STORE, src=(UPat(dtype=dtypes.int64), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True), - lambda x: UOp(x.op, dtypes.void, x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])), + (UPat(Ops.STORE, src=(UPat(dtype=dtypes.int64), UPat(dtype=dtypes.bool)), name="x"), + lambda x: UOp(x.op, dtypes.void, (x.src[0], x.src[1].cast(dtypes.uint8)))), # indexing on PTX is in uint64, we do the math while it's still in the graph (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op", allow_any_len=True), lambda buf,idx,op: UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)+op.src[2:]) \ diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 3499eba493..279a0adcb9 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -442,7 +442,7 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([ (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.AFTER, name="y"))), lambda x,y: x.after(*y.src[1:])), # remove invalid writes - (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),))), allow_any_len=True), lambda: UOp(Ops.NOOP)), + (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),)))), lambda: UOp(Ops.NOOP)), (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.NOOP, src=()))), lambda x: x), (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.END, src=(UPat(Ops.NOOP, src=()),), allow_any_len=True))), lambda x: x), ]) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 082b78d84a..7e7bab8451 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -26,7 +26,7 @@ axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5} -range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1, +range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1, Ops.COPY: 2, Ops.BUFFER_VIEW: 1, Ops.LINEAR: 0} # https://en.wikipedia.org/wiki/Identity_element diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 9fb519dfbb..beb8346695 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -442,16 +442,16 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([ # ** load/store folding ** (UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)), (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"), - UPat.load(UPat(Ops.INDEX, name="index"))), allow_any_len=True, name="store"), - lambda index, gate, alt, store: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt, *store.src[2:])), + UPat.load(UPat(Ops.INDEX, name="index")))), + lambda index, gate, alt: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt)), # fold gated LOAD/STORE - (UPat(Ops.STORE, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"), lambda x: UOp(Ops.NOOP)), + (UPat(Ops.STORE, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(), UPat())), lambda: UOp(Ops.NOOP)), (UPat(Ops.LOAD, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"), lambda x: x.src[1] if len(x.src) > 1 else x.const_like(0)), # invalid load produces 0, or the alt value if we have one - (UPat(Ops.STORE, src=(UPat(), invalid_pat), allow_any_len=True), lambda i: UOp(Ops.NOOP)), + (UPat(Ops.STORE, src=(UPat(), invalid_pat)), lambda i: UOp(Ops.NOOP)), # store of where with invalid -> gated store - (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat)), allow_any_len=True, name="store"), - lambda index, cond, val, store, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val, *store.src[2:])), + (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat))), + lambda index, cond, val, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val)), ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()), # 1/(x^c) -> (1/x)^c ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()), ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)