geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions 7 Activity

remove vec from const (#16889)

Browse Source 
* remove vec from const

* reject in spec

* remove all vector dtypes

* not needed

* remove that

* clean up invalid
This commit is contained in:
George Hotz
2026-07-06 12:23:07 -07:00
committed by GitHub
parent 47f9677f67
commit 9512dc30f4
13 changed files with 50 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/mi350x_uop_matmul.py
+1 -1
View File
@@ -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]:
test/null/test_const_folding.py
+1 -1
View File
@@ -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))
test/null/test_graph_rewrite.py
+8 -8
View File
@@ -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<rhs
opt = apply_rewrite(unopt)
print(unopt)
@@ -182,28 +182,28 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
class TestGEPAndVectorizeRewrite(unittest.TestCase):
def test_gep_single_element_extraction(self):
# GEP on a vector dtype to extract a single element
base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)
def test_gep_tuple_extraction(self):
# GEP on a vector dtype to extract multiple elements as a vector
base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
def test_gep_on_const_stack(self):
# GEP on a const STACK to extract a single element
const_stack = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
const_stack = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0)
def test_gep_tuple_on_const_stack(self):
# GEP on a const STACK using a tuple to extract multiple elements
const_stack = UOp.const(dtypes.float32.vec(4), (7.0, 8.0, 9.0, 10.0))
const_stack = UOp.const(dtypes.float32, (7.0, 8.0, 9.0, 10.0))
self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
def test_vectorize_multiple_elements(self):
# Vectorizing multiple elements using GEP
base_vector = UOp.const(dtypes.float32.vec(4), (5.0, 10.0, 15.0, 20.0))
vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), src=tuple(base_vector.index(i) for i in range(4)))
base_vector = UOp.const(dtypes.float32, (5.0, 10.0, 15.0, 20.0))
vectorized_uop = UOp(Ops.STACK, dtypes.float32, src=tuple(base_vector.index(i) for i in range(4)))
self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
test/null/test_simplify_valid_idx.py
+1 -1
View File
@@ -19,7 +19,7 @@ def get_gated_load_uop(valid:UOp, idx:UOp):
))
def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
return UOp(Ops.LOAD, dtypes.float.vec(4), (
return UOp(Ops.LOAD, dtypes.float, (
UOp.param(0, dtypes.imagef(image_shape)).index(idx[1].valid(valid), idx[0].valid(valid), ptr=True),
))
test/null/test_uop_graph.py
+2 -2
View File
@@ -20,7 +20,7 @@ def const_values(u:UOp):
class TestGraphRewriteConst(unittest.TestCase):
def test_gep_const(self):
v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
v1 = UOp.const(dtypes.int, (0,1,2))
v2 = v1.index(1)
ret = graph_rewrite(v2, sym)
self.assertEqual(ret.dtype, dtypes.int)
@@ -303,7 +303,7 @@ class TestUOpGraph(unittest.TestCase):
def test_gep_vec_const_fold(self):
for vec_size in [2, 4, 8]:
consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
vec = UOp(Ops.STACK, dtypes.float.vec(vec_size), tuple(consts))
vec = UOp(Ops.STACK, dtypes.float, tuple(consts))
with Context(SPEC=0):
uops = to_uops_list([vec.index(i) for i in range(vec_size)])
for uop, const in zip(uops, consts):
test/null/test_uop_symbolic.py
+3 -3
View File
@@ -1289,9 +1289,9 @@ class TestInvalidIndex(unittest.TestCase):
self.assertIs((UOp.invalid()<Variable("a",0,10)).simplify().dtype, dtypes.bool)
def test_alu_invalid_vconst(self):
c1 = UOp.const(dtypes.weakint.vec(4), (1, 1, Invalid, Invalid))
c2 = UOp.const(dtypes.weakint.vec(4), (1, Invalid, 1, 1))
self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint.vec(4), (2, Invalid, Invalid, Invalid)))
c1 = UOp.const(dtypes.weakint, (1, 1, Invalid, Invalid))
c2 = UOp.const(dtypes.weakint, (1, Invalid, 1, 1))
self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint, (2, Invalid, Invalid, Invalid)))
class TestStoreLoadFolding(unittest.TestCase):
"""Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
test/null/test_uop_vmin_vmax.py
+9 -9
View File
@@ -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)
test/null/test_uops.py
+7 -7
View File
@@ -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__':
tinygrad/codegen/__init__.py
+2 -5
View File
@@ -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):
tinygrad/codegen/opt/postrange.py
+1 -1
View File
@@ -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]
tinygrad/schedule/rangeify.py
+3 -3
View File
@@ -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([
tinygrad/uop/ops.py
+7 -8
View File
@@ -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))
tinygrad/uop/spec.py
+5 -1
View File
@@ -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),