From 9512dc30f4577060c08fcada23cacb22202b644a Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 6 Jul 2026 12:23:07 -0700 Subject: [PATCH] remove vec from const (#16889) * remove vec from const * reject in spec * remove all vector dtypes * not needed * remove that * clean up invalid --- extra/gemm/mi350x_uop_matmul.py | 2 +- test/null/test_const_folding.py | 2 +- test/null/test_graph_rewrite.py | 16 ++++++++-------- test/null/test_simplify_valid_idx.py | 2 +- test/null/test_uop_graph.py | 4 ++-- test/null/test_uop_symbolic.py | 6 +++--- test/null/test_uop_vmin_vmax.py | 18 +++++++++--------- test/null/test_uops.py | 14 +++++++------- tinygrad/codegen/__init__.py | 7 ++----- tinygrad/codegen/opt/postrange.py | 2 +- tinygrad/schedule/rangeify.py | 6 +++--- tinygrad/uop/ops.py | 15 +++++++-------- tinygrad/uop/spec.py | 6 +++++- 13 files changed, 50 insertions(+), 50 deletions(-) diff --git a/extra/gemm/mi350x_uop_matmul.py b/extra/gemm/mi350x_uop_matmul.py index 8d0725d609..6d10c59852 100644 --- a/extra/gemm/mi350x_uop_matmul.py +++ b/extra/gemm/mi350x_uop_matmul.py @@ -79,7 +79,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: # this is the big accumulator acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG) assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N - acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float.vec(4), 0.0), end=init_l) + acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l) # create locals (note A is permuted, and the stride is changed to avoid bank conflicts) def make_locals(slot) -> tuple[UOp, UOp]: diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index a6557e0dad..9e595ff9f4 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -127,7 +127,7 @@ class TestBitcastConstFolding(unittest.TestCase): def test_vec_bitcast(self): with Context(SPEC=0): - srcs = full_rewrite(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src + srcs = full_rewrite(UOp.const(dtypes.int32, (-1, -2**31, 75)).bitcast(dtypes.uint32).sink()).src self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs)) self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75)) diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index 782e6c770a..f7b496074c 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -129,8 +129,8 @@ class TestModuloAndDivisionFolding(unittest.TestCase): def test_graph_rewrite_div_folding_bug(self): lhs = UOp(Ops.ADD, dtypes.int.vec(4), src=( UOp(Ops.STACK, dtypes.int.vec(4), arg=None, src=(UOp(Ops.SPECIAL, dtypes.int, arg='lidx0', src=(UOp.const(dtypes.int, 32),)),)*4), - UOp.const(dtypes.int.vec(4), (0, 256, 512, 768)))) - rhs = UOp.const(dtypes.int.vec(4), 2) + UOp.const(dtypes.int, (0, 256, 512, 768)))) + rhs = UOp.const(dtypes.int, (2,)*4) unopt = lhs NOOP rule. This rule matches patterns that EMERGE during simplification.""" diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index f6c939a8f6..00073e473e 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -160,7 +160,7 @@ class TestVminVmaxProperties(unittest.TestCase): self.assertNotEqual(i.vmin, i.vmax) def test_vmin_vmax_invalid_vconst(self): - x = UOp.const(dtypes.weakint.vec(4), (0, 4, Invalid, Invalid)) + x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid)) self.assertLess(x.vmin, 0) self.assertGreater(x.vmax, 4) @@ -280,37 +280,37 @@ class TestVminVmaxDivMod(unittest.TestCase): class TestVminVmaxVConst(unittest.TestCase): def test_vmin_vmax_vconst_single_element(self): # vmin and vmax for a single-element vector constant - uop = UOp.const(dtypes.int32.vec(1), (42,)) + uop = UOp.const(dtypes.int32, (42,)) self.assertEqual(uop.vmin, 42) self.assertEqual(uop.vmax, 42) def test_vmin_vmax_vconst_multiple_elements(self): # vmin and vmax for a multi-element vector constant - uop = UOp.const(dtypes.int32.vec(4), (10, 20, -5, 7)) + uop = UOp.const(dtypes.int32, (10, 20, -5, 7)) self.assertEqual(uop.vmin, -5) self.assertEqual(uop.vmax, 20) def test_vmin_vmax_vconst_all_equal(self): # vmin and vmax for a vector where all elements are equal - uop = UOp.const(dtypes.int32.vec(3), (7, 7, 7)) + uop = UOp.const(dtypes.int32, (7, 7, 7)) self.assertEqual(uop.vmin, 7) self.assertEqual(uop.vmax, 7) def test_vmin_vmax_vconst_with_negative_values(self): # vmin and vmax for a vector constant containing negative values - uop = UOp.const(dtypes.int32.vec(4), (-10, -20, -5, -15)) + uop = UOp.const(dtypes.int32, (-10, -20, -5, -15)) self.assertEqual(uop.vmin, -20) self.assertEqual(uop.vmax, -5) def test_vmin_vmax_vconst_with_floats(self): # vmin and vmax for a vector constant of float values - uop = UOp.const(dtypes.float32.vec(3), (1.5, -3.2, 0.0)) + uop = UOp.const(dtypes.float32, (1.5, -3.2, 0.0)) self.assertEqual(uop.vmin, -3.2) self.assertEqual(uop.vmax, 1.5) def test_vmin_vmax_vconst_with_bools(self): # vmin and vmax for a vector constant of bool values - uop = UOp.const(dtypes.bool.vec(3), (True, False, False)) + uop = UOp.const(dtypes.bool, (True, False, False)) self.assertIs(uop.vmin, False) self.assertIs(uop.vmax, True) @@ -318,8 +318,8 @@ class TestVminVmaxVConst(unittest.TestCase): # vmin and vmax for a vector constant of bool values d1 = UOp.param(1, dtypes.int.ptr()) idx = UOp.const(dtypes.int, 0) - val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),)) - uop = (val // 32).index(0) + val = UOp(Ops.LOAD, dtypes.int, (d1.index(idx).cast(dtypes.int.ptr()),)) + uop = (val // 32) self.assertEqual(uop.vmin, -67108864) self.assertEqual(uop.vmax, 67108863) diff --git a/test/null/test_uops.py b/test/null/test_uops.py index c11e86f9a2..15f4bfc226 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -315,7 +315,7 @@ class TestUOpStr(unittest.TestCase): assert str(eval(str(a))) == str(a) def test_vectorized_str(self): - vec = UOp(Ops.STACK, dtypes.int.vec(4), tuple(UOp.const(dtypes.int, x) for x in range(4))) + vec = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4))) assert str(eval(str(vec))) == str(vec) def test_reduceop_arg(self): @@ -344,22 +344,22 @@ class TestUopsObject(unittest.TestCase): class TestUOpRender(unittest.TestCase): def test_render_vectorize_empty(self): - u = UOp(Ops.STACK, dtype=dtypes.int.vec(0), src=()) + u = UOp(Ops.STACK, dtype=dtypes.int, src=()) self.assertEqual(u.render(simplify=False), "{}") def test_render_vectorize_empty_simplified(self): - u = UOp(Ops.STACK, dtype=dtypes.int.vec(0), src=()) + u = UOp(Ops.STACK, dtype=dtypes.int, src=()) self.assertEqual(u.render(), "{}") def test_render_vectorize_same(self): - u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0))) + u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0))) self.assertEqual(u.render(simplify=False), "{0,0,0}") def test_render_vectorize_different(self): - u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2))) + u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2))) self.assertEqual(u.render(simplify=False), "{0,1,2}") def test_render_vectorize_same_simplified(self): - u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0))) + u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0))) self.assertEqual(u.render(), "{0,0,0}") def test_render_vectorize_different_simplified(self): - u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2))) + u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2))) self.assertEqual(u.render(), "{0,1,2}") if __name__ == '__main__': diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 92896ba81c..489db49ece 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -29,16 +29,13 @@ from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element from tinygrad.schedule.rangeify import BufferizeOpts pm_remove_vec_dtypes = PatternMatcher([ - # CONST must be stacked CONST - (UPat(Ops.CONST, name='c'), - lambda c: UOp(Ops.STACK, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None), # rewrite PARAM to non pointer (UPat((Ops.PARAM, Ops.BUFFER), name="buf"), lambda buf: buf.replace(dtype=buf.dtype.base, src=(UOp.const(dtypes.int, buf.ptrdtype.size),)) \ if isinstance(buf.dtype, PtrDType) and not isinstance(buf.dtype, ImageDType) else None), - # remove all vec dtypes + # remove pointer dtypes from non-PARAM/BUFFER ops (UPat(GroupOp.All-{Ops.PARAM, Ops.BUFFER}, name="x"), - lambda x: x.replace(dtype=x.dtype.base.scalar().base)), + lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, PtrDType) else None), ])+pm_clean_up_group_sink def do_number_param(ctx:list[int], x:UOp): diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 926454fcaf..65df66489f 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -303,7 +303,7 @@ class Scheduler: # they need to be moved into the WMMA srcs wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.ren.target.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes) tc_uop = UOp(Ops.WMMA, dtype=tc.dtype_out, src=( - srcs[0], srcs[1], UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1) + srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2])), arg=wmma_arg, tag=1) # preserve extra reduces reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes] diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index d2d223b7e9..c436bd532a 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -30,7 +30,7 @@ def lower_shaped_wmma(ctx, x): tc_upcast_axes = tuple(((u.arg[0], s.shape[-1]),) for s, u in upcasts) name = f"WMMA_{'_'.join(map(str, dims))}_{dtype_in.name}_{dtype_out.name}" wmma_arg = (name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ()) - wmma = UOp(Ops.WMMA, dtype_out.vec(x.src[2].shape[-1]), tuple(s[u].contract(u) for s, u in upcasts), arg=wmma_arg) + wmma = UOp(Ops.WMMA, dtype_out, tuple(s[u].contract(u) for s, u in upcasts), arg=wmma_arg) tmp = UOp.placeholder((x.src[2].shape[-1],), dtype_out, slot=next(ctx), addrspace=AddrSpace.REG) return tmp.after(UOp.group(*[tmp[e].store(wmma.index(e)) for e in range(x.src[2].shape[-1])])) @@ -563,10 +563,10 @@ rangeify_codegen = PatternMatcher([ (UPat(Ops.BUFFER).f(Ops.AFTER, allow_any_len=True).broadcast(name="dg").f(Ops.INDEX, name="idx", allow_any_len=True), lambda dg,idx: None if dg.addrspace is not AddrSpace.LOCAL or isinstance(idx.dtype, PtrDType) else - idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar().vec(dg.dtype.vcount))), + idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar())), (UPat(Ops.BUFFER).f(Ops.AFTER, allow_any_len=True).gep(name="dg").f(Ops.INDEX, name="idx", allow_any_len=True), lambda dg,idx: None if dg.addrspace is not AddrSpace.LOCAL or isinstance(idx.dtype, PtrDType) else - idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar().vec(dg.dtype.vcount))), + idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar())), ]) pm_add_range_tags = PatternMatcher([ diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 6091014851..819ab147d6 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -494,15 +494,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if isinstance(x, UOp): return x # float self keeps its dtype for any scalar, int self only for int/Invalid scalars if dtypes.is_float(self.dtype) or (dtypes.is_int(self.dtype) and isinstance(x, (int, InvalidType))): return self.const_like(x) - return self.const_like(x, dtypes.from_py(x).vec(self.dtype.vcount)) + return self.const_like(x, dtypes.from_py(x)) def broadcast(self, count:int): assert self.dtype.vcount == 1 if count == 1: return self - return UOp(Ops.STACK, self.dtype.vec(count), (self,)*count) + return UOp(Ops.STACK, self.dtype, (self,)*count) def cast(self, dtype:DTypeLike): dtype = to_dtype(dtype) - # TODO: we shouldn't have to check for dtype.count == 1 here, but CAST is misused in AMD LLVM - if dtype.count == 1 and dtype.count != self.dtype.count: dtype = dtype.vec(self.dtype.count) if self.dtype == dtype: return self return UOp(Ops.CAST, dtype, (self,)) def bitcast(self, dtype:DTypeLike): @@ -529,14 +527,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass): out_shape = _broadcast_shape(*shapes) all_srcs = tuple(x._broadcast_to(out_shape) if x._shape else x for x in all_srcs) out_dtype = all_srcs[-1].dtype - if op in {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool + if op in {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}: out_dtype = dtypes.bool return UOp(op, out_dtype, all_srcs, **kwargs) @staticmethod def const(dtype:DType, b:ConstLike, shape:tuple[sint, ...]|None=None): if isinstance(b, UOp): return b.cast(dtype) + assert dtype.vcount == 1, f"const dtype must be scalar, got {dtype}" # NOTE: it always has to be STACK now, even if they are all the same if isinstance(b, tuple): - stk = [UOp(Ops.CONST, dtype.scalar(), arg=dtype.const(c), src=()) for c in b] + stk = [UOp(Ops.CONST, dtype, arg=dtype.const(c), src=()) for c in b] ret = UOp.vectorize(*stk) else: ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=()) @@ -553,9 +552,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): ret = UOp(Ops.REDUCE, self.dtype, (self,), (op, reduce_axis)) if len(reduce_axis) else self return ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis)) if axis != reduce_axis else ret @staticmethod - def invalid(count=1): return UOp(Ops.CONST, dtypes.weakint.vec(count), src=(), arg=Invalid) + def invalid(): return UOp.const(dtypes.weakint, Invalid) def valid(self, cond): - return cond.where(self.cast(dtypes.weakint), UOp.invalid(self.dtype.count)) + return cond.where(self.cast(dtypes.weakint), UOp.invalid()) def get_idx(self) -> UOp: assert self.dtype.scalar() is dtypes.weakint, "Can only call get_idx on index dtype" if self.op is Ops.STACK: return UOp.vectorize(*(x.get_idx() for x in self.src)) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 44e4c46d0a..e2b50dd91b 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -47,7 +47,11 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): # these ops can be used in the tensor graph and programs spec_shared = PatternMatcher([ - (UPat(Ops.SINK, dtypes.void), lambda: True), # NOTE: for testing, we let sinks be anything + # no vec dtypes allowed + (UPat(GroupOp.All, name="x"), lambda x: False if x.dtype.vcount > 1 else None), + + # NOTE: for testing, we let sinks be anything + (UPat(Ops.SINK, dtypes.void), lambda: True), # NOOP. TODO: remove this (UPat(Ops.NOOP), lambda: True),