From 20242fdf1dda40d89db8f58dcda7e74e15066ae8 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 1 Jun 2026 19:24:43 -0700 Subject: [PATCH] update test + spec from shrink_in_render (#16467) * update test + spec from shrink_in_render * cast --- test/backend/test_linearizer.py | 12 ++++++------ test/mockgpu/amd/emu.py | 17 +++++++++++------ test/null/test_simplify_valid_idx.py | 3 +-- test/null/test_tensor.py | 2 +- test/null/test_uop_graph.py | 10 +++------- test/opt/test_gen_float4.py | 8 ++++---- tinygrad/uop/__init__.py | 2 +- tinygrad/uop/ops.py | 15 +++++++++++---- tinygrad/uop/spec.py | 18 ++++++++++-------- tinygrad/uop/validate.py | 4 +++- 10 files changed, 51 insertions(+), 40 deletions(-) diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index b89c1f94da..3ebc911536 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -70,9 +70,9 @@ class TestLinearizer(unittest.TestCase): uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1] # only valid test if outermost range is the reduce if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE: - load_types = [u.src[0].dtype for u in uops[uslice+1:] if u.op == Ops.LOAD] + load_idxs = [u.src[0] for u in uops[uslice+1:] if u.op == Ops.LOAD] # assert that there is a global load after the reduce ends - assert any(dt.addrspace == AddrSpace.GLOBAL for dt in load_types) + assert any(u.addrspace == AddrSpace.GLOBAL for u in load_idxs) def _test_no_nested_ranges(self, lins, skip=None): for l in lins: @@ -165,7 +165,6 @@ class TestLinearizer(unittest.TestCase): stores = [u for u in uops if u.op is Ops.STORE] assert len(accs) == 0 # it's removed now assert len(stores) == 1 - assert stores[0].src[1].dtype == dtypes.float.vec(4) # NOTE: can reenable, it does work. it just makes BEAM slow @unittest.expectedFailure @@ -186,12 +185,13 @@ class TestLinearizer(unittest.TestCase): opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)] program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer) - stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG] + stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG] # the first store is to lds and can be upcasted - assert stores[0].src[1].dtype == dtypes.float.vec(4) + assert stores[0].src[1].max_numel() == 4 assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort()) # the second store is to gds with no upcasts + assert stores[1].src[1].max_numel() == 1 assert stores[1].src[1].dtype == dtypes.float assert any(x.op is Ops.PARAM for x in stores[1].toposort()) @@ -367,7 +367,7 @@ class TestLinearizer(unittest.TestCase): assert len(barrier) == 1 # check that the float4 cast collapses for all stores for store in local_stores+global_stores: - assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE + assert store.src[1].max_numel() > 1 # and store.src[2].op is not Ops.VECTORIZE # # check the children's vins # TODO: src ALU are not the same, should it? # assert barrier.src == tuple(local_stores) diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index b23c8021e8..cfa2bdd65a 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -1344,6 +1344,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: src0_r = src0_off - _c(256) # VGPR-relative index (only valid when src is VGPR) src1_r = src1_off - _c(256) src2_off = ctx.inst_field(type(inst).src2) + use_acc = bool(getattr(inst, 'acc_cd', 0)) # Check if sources are VGPRs (offset >= 256) vs inline constants/SGPRs src0_is_vgpr = src0_off >= _c(256) src1_is_vgpr = src1_off >= _c(256) @@ -1504,7 +1505,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: # So: m_base = half * 16 + (out_reg // 4) * 4 + (out_reg % 4) m_base = c_half * UOp.const(dtypes.int, 16) + UOp.const(dtypes.int, (out_reg // 4) * 4 + (out_reg % 4)) - acc_v = ctx.raccvgpr_dyn(src2_r + _c(out_reg), compute_lane, src2_is_vgpr) + acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr) if is_int_out: acc_v = acc_v.cast(dtypes.int32) else: acc_v = acc_v.bitcast(dtypes.float32) acc = src2_is_vgpr.where(acc_v, acc_scalar) @@ -1515,16 +1516,18 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: acc = acc + a_val * b_val if is_int_out: - compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask)) + compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)( + vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask)) else: - compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask)) + compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)( + vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask)) else: # 16x16 and 4x4: each lane computes out_per_lane outputs n_idx = compute_lane % UOp.const(dtypes.int, grp_sub) c_grp = compute_lane // UOp.const(dtypes.int, grp_sub) for out_reg in range(out_per_lane): - acc_v = ctx.raccvgpr_dyn(src2_r + _c(out_reg), compute_lane, src2_is_vgpr) + acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr) if is_int_out: acc_v = acc_v.cast(dtypes.int32) else: acc_v = acc_v.bitcast(dtypes.float32) acc = src2_is_vgpr.where(acc_v, acc_scalar) @@ -1545,9 +1548,11 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: acc = acc + a_val * b_val if is_int_out: - compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask)) + compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)( + vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask)) else: - compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask)) + compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)( + vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask)) compute_phase = UOp.group(*compute_stores).end(compute_lane) return UOp.sink(read_phase, compute_phase, *ctx.inc_pc()) diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index a537a51c72..c1ff696b31 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -292,8 +292,7 @@ class TestImageSimplification(unittest.TestCase): load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx) with Context(NOOPT=1, SPEC=0): load = full_rewrite(load.sink()).src[0] - self.assertEqual(load.op, Ops.STACK) - self.assertEqual(load.dtype.count, 4) + self.assertFalse(load.op_in_backward_slice_with_self(Ops.LOAD)) def test_openpilot_conv1(self): # first conv in openpilot diff --git a/test/null/test_tensor.py b/test/null/test_tensor.py index cf6a2001d0..aca886d4e6 100644 --- a/test/null/test_tensor.py +++ b/test/null/test_tensor.py @@ -78,7 +78,7 @@ class TestIdxUpcast(unittest.TestCase): if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)): assert idx.op is Ops.INDEX idx_val = idx.src[1] - self.assertIs(idx_val.dtype, dtype) + self.assertFalse(idx_val.overflows(idx_val.dtype.base.scalar())) # use expand to generate kernel that uses large idx def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3): diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 6f38064ca7..018eaf9d60 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -536,9 +536,8 @@ class TestUOpGraph(unittest.TestCase): ld0 = glbl1.index(UOp.invalid()) ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True))) uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(idx, ptr=True), ld1+ld0))]) - ld0 = uops[-2].src[-1] # -2 to skip SINK # the gate and invalid value are deleted from ld1 - self.assertEqual(ld0, UOp.load(glbl2.index(idx, ptr=True), dtype=dtypes.int)) + self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1) def test_fold_gated_load_local(self): glbl0 = UOp.param(0, dtypes.int.ptr(16)) @@ -550,21 +549,18 @@ class TestUOpGraph(unittest.TestCase): ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(dtypes.bool, True))) uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(lidx, ptr=True), ld1+ld0))]) - ld0 = uops[-2].src[-1] # -2 to skip SINK # the gate and invalid value are deleted from ld1 - self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2, ptr=True)) + self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2) def test_fold_gated_store(self): glbl = UOp.param(0, dtypes.int.ptr(1)) idx0 = UOp.const(dtypes.int, 0) - idx1 = UOp.const(dtypes.int, 0) val = UOp.const(dtypes.int, 42) st0 = glbl.index(UOp.invalid(), ptr=True).store(val) st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True)), ptr=True).store(val) uops = to_uops_list([st0, st1]) # only the second store happens - self.assertEqual(len(uops), 7) # +1 for SINK, +1 for PARAM shape sentinel - self.assertEqual(uops[-2], glbl.index(idx1, ptr=True).store(val)) # -2 to skip SINK + self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1) @unittest.skip("this is a uop type error") def test_asserts_bad_gate(self): diff --git a/test/opt/test_gen_float4.py b/test/opt/test_gen_float4.py index 66c3c848ad..4fcaa6a093 100644 --- a/test/opt/test_gen_float4.py +++ b/test/opt/test_gen_float4.py @@ -13,12 +13,12 @@ AMX = "AMX" in DEV.arch class TestFloat4(unittest.TestCase): @staticmethod def count_float4(uops: list[UOp], n=4): - return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float.vec(n)]), - len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float.vec(n)])) + return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.float and uop.shape == (4,)]), + len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.float and uop.shape == (4,)])) @staticmethod def count_half4(uops: list[UOp]): - return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half.vec(4)]), - len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half.vec(4)])) + return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.half and uop.shape == (4,)]), + len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.half and uop.shape == (4,)])) def test_float4_basic(self): a = Tensor.empty(2, 8).realize() diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 98ab6490e9..b05756930c 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -117,7 +117,7 @@ class GroupOp: Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW, Ops.FLOORDIV, Ops.FLOORMOD} Ternary = {Ops.WHERE, Ops.MULACC} ALU = set.union(Unary, Binary, Ternary) - Broadcastable = set.union(Binary, Ternary, {Ops.GROUP, Ops.STORE}) + Broadcastable = set.union(Binary, Ternary, {Ops.GROUP}) # TODO: is BITCAST always Elementwise if it's shape changing? Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST}) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b5aa11b486..b66def9594 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -288,6 +288,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass): # STAGE adds the existing shape to the front, opposite of INDEX return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape case Ops.DEFINE_LOCAL | Ops.DEFINE_REG: + if len(self.src) >= 1: + # NOTE: this is the same as PARAM + return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count)) if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size, self.dtype.count) if self.dtype.count > 1 else (self.ptrdtype.size,) return (self.dtype.count,) if self.dtype.count > 1 else () @@ -301,7 +304,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): # passthrough ops case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \ - Ops.COPY | Ops.ALLREDUCE: + Ops.COPY | Ops.ALLREDUCE | Ops.STORE: return self.src[0]._shape # REDUCE with empty axis is passthrough (lowered form) case Ops.REDUCE if len(self.arg[1]) == 0: @@ -689,7 +692,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass): def marg(self): match self.op: case Ops.RESHAPE | Ops.EXPAND: return tuple(ssimplify(self.src[1].sgep(i)) for i in range(self.src[1].dtype.count)) - case Ops.PAD | Ops.SHRINK: return tuple((self.src[1].sgep(i), self.src[2].sgep(i)) for i in range(self.src[1].dtype.count)) + case Ops.PAD | Ops.SHRINK: + # this is like broadcasting for shapes + return tuple(((ssimplify(self.src[1]) if self.src[1].shape == () else self.src[1].sgep(i)), + (ssimplify(self.src[2]) if self.src[2].shape == () else self.src[2].sgep(i))) + for i in range(max(self.src[1].dtype.count, self.src[2].dtype.count))) case Ops.PERMUTE | Ops.FLIP: return self.arg case _: raise RuntimeError(f"{self.op} is not a MovementOp") @@ -762,7 +769,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.GEP}: return self.src[0].addrspace if self.op in GroupOp.Movement: return self.src[0].addrspace - if self.op is Ops.STACK or self.op in GroupOp.Elementwise: + if self.op in {Ops.STACK, Ops.WMMA} or self.op in GroupOp.Elementwise: ad = [x.addrspace for x in self.src if x.addrspace is not None] if not len(ad) or not all_same(ad): return None return ad[0] @@ -1109,7 +1116,7 @@ class ProgramInfo: if u.op is Ops.DEFINE_VAR: _vars.append(u) if u.op is Ops.PARAM: _globals.append(u.arg.slot) if u.op in (Ops.STORE, Ops.LOAD): - if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX): + if (idx:=u.src[0]).op in (Ops.INDEX, Ops.SHRINK) or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX): if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg.slot) if u.op is Ops.SPECIAL: if u.arg[0] == 'i': local_size = None diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 87f37a0169..580a2e0828 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -81,8 +81,7 @@ spec_shared = PatternMatcher([ (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.UNROLL, Ops.INS))), lambda: True), # TOOD: these should be buffer with different addrspace - (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL), - (UPat(Ops.DEFINE_REG, src=()), lambda: True), + (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)), lambda: True), # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.DEFINE_REG, Ops.DEFINE_LOCAL, Ops.AFTER, Ops.MULTI, @@ -102,11 +101,11 @@ spec_shared = PatternMatcher([ (UPat(Ops.INS), lambda: True), # LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE - (UPat(Ops.INDEX, name="uidx").or_casted().load(), validate_index), - (UPat(Ops.INDEX, name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"), + (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(), validate_index), + (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"), lambda uidx,gate,alt,load: validate_index(uidx, gate) if alt.dtype == load.dtype else False), - (UPat(Ops.INDEX, name="uidx").or_casted().store(UPat()), validate_index), - (UPat(Ops.INDEX, name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index), + (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat()), validate_index), + (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index), # STORE in tensor graph: store a value into a target (UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x: True), @@ -197,6 +196,9 @@ spec_program = PatternMatcher([ # weakint is not allowed in programs (UPat(GroupOp.All, dtypes.weakint), lambda: False), + # allow special SHRINK + (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True), + # movement ops are not allowed in programs (UPat(GroupOp.Movement), lambda: False), @@ -208,11 +210,11 @@ spec_program = PatternMatcher([ lambda x: False if x.dtype.count > 1 and (x.dtype.count,) != x.shape else None), # STACK/GEP in program. TODO: this should match Tensor - (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)), + (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1), (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()), # if has a - (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(dtype=dtypes.bool), UPat((Ops.CAST, Ops.INDEX)))), lambda: True), + (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(dtype=dtypes.bool), UPat((Ops.CAST, Ops.INDEX, Ops.SHRINK)))), lambda: True), (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True), ])+spec_shared diff --git a/tinygrad/uop/validate.py b/tinygrad/uop/validate.py index c24af4135f..6c58be3d3a 100644 --- a/tinygrad/uop/validate.py +++ b/tinygrad/uop/validate.py @@ -51,7 +51,9 @@ z3_renderer = PatternMatcher([ ]) def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]: - lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.dtype.scalar() in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK))[:-1] + # 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 \ + (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: z3_rewritten = z3_renderer.rewrite(u, ctx=(solver.ctx, z3map))