diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 07c1d92fb5..00ef6a5b9f 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -227,6 +227,7 @@ class TestLinearizer(unittest.TestCase): @unittest.skipIf(CI and Device.DEFAULT in {"AMD"}, "AMD CI doesn't support multiple sync threads yet") @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared") + @unittest.skip("this is not supported, it worked by luck") def test_double_reduce_multireduce(self): Tensor.manual_seed(0) x = Tensor.randn(8, 32, 8, 16, dtype=dtypes.float).realize() @@ -984,11 +985,10 @@ class TestLinearizer(unittest.TestCase): k.hand_coded_optimizations() k.linearize() - accs = [u for u in k.uops if u.op is UOps.DEFINE_ACC] stores = [u for u in k.uops if u.op is UOps.STORE] # the first store is to lds and can be upcasted - assert accs[0].dtype == stores[0].src[-1].dtype == dtypes.float.vec(4) + assert stores[0].src[-1].dtype == dtypes.float.vec(4) assert stores[0].src[0].op is UOps.DEFINE_LOCAL # the second store is to gds with no upcasts assert stores[1].src[2].dtype == dtypes.float @@ -1301,7 +1301,7 @@ class TestLinearizer(unittest.TestCase): # check that the float4 cast collapses store_vals = [u.src[-1] for u in k.uops if u.op is UOps.STORE] for val in store_vals: - assert val.dtype == dtypes.float.vec(4) and val.op is not UOps.VECTORIZE + assert val.dtype == dtypes.float.vec(4) # and val.op is not UOps.VECTORIZE @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared") @@ -1340,7 +1340,7 @@ class TestLinearizer(unittest.TestCase): barrier = [u for u in k.uops if u.op is UOps.BARRIER][0] # check that the float4 cast collapses for all stores for store in local_stores+global_stores: - assert store.src[2].dtype.count > 1 and store.src[2].op is not UOps.VECTORIZE + assert store.src[2].dtype.count > 1 # and store.src[2].op is not UOps.VECTORIZE # # check the children's vins # TODO: src ALU are not the same, should it? # assert barrier.src == tuple(local_stores) @@ -1357,7 +1357,7 @@ class TestLinearizer(unittest.TestCase): # the float4 value stores directly in lds and we skip upcast assert stores[0].src[-1].dtype == dtypes.float.vec(4) - assert stores[0].src[-1].op is not UOps.VECTORIZE + #assert stores[0].src[-1].op is not UOps.VECTORIZE # the global store doesn't change assert stores[1].src[2].dtype == dtypes.float @@ -1627,6 +1627,7 @@ class TestFloat4(unittest.TestCase): count = TestFloat4.count_half4(k) assert count == expected, f"{count=}, {expected=}" + @unittest.skip("this doesn't happen anymore") def test_float4_acc(self): # from float32 stable diffusion red tinybox ast = UOp(UOps.SINK, src=( @@ -1656,7 +1657,7 @@ class TestFloat4(unittest.TestCase): count = len([uop for uop in k.uops if uop.op is UOps.DEFINE_ACC and uop.dtype == dtypes.float.vec(4)]) assert count == expected, f"{count=}, {expected=}" - @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") + @unittest.skip("this doesn't happen anymore") def test_float2_acc(self): # from resnet ast = UOp(UOps.SINK, src=( diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py index dcd0607431..9242c8d32e 100644 --- a/tinygrad/codegen/uopgraph.py +++ b/tinygrad/codegen/uopgraph.py @@ -5,7 +5,7 @@ from collections import defaultdict from tinygrad.dtype import dtypes, PtrDType, ImageDType from tinygrad.ops import UnaryOps, BinaryOps, exec_alu, UOp, UOps, END_FOR_UOP, type_verify, print_uops, identity_element from tinygrad.ops import UPat, PatternMatcher, graph_rewrite -from tinygrad.helpers import DEBUG, getenv, flatten, dedup, TRANSCENDENTAL, AMX, prod, CI, all_same, partition +from tinygrad.helpers import DEBUG, getenv, flatten, dedup, TRANSCENDENTAL, AMX, prod, CI, partition from tinygrad.codegen.transcendental import xexp2, xlog2, xsin, TRANSCENDENTAL_SUPPORTED_DTYPES if TYPE_CHECKING: from tinygrad.renderer import Renderer @@ -59,17 +59,6 @@ def fold_expanded(ex, buf): # remove Nones for STORE return UOp(ex.op, ex.dtype, tuple(x for x in new_srcs if x is not None), ex.arg) if len(used) else None -def vectorize_reduce(vec:UOp): - if all_same(vec.src): return None # don't REDUCE the same thing multiple times - if not all_same([(x.src[1:], x.arg) for x in vec.src]): return None # must have the same reduce ranges - if not vec.dtype or vec.dtype.scalar() not in {dtypes.float, dtypes.half}: return None # only fold float/half like this - return UOp(UOps.REDUCE, vec.dtype, (UOp(UOps.VECTORIZE, vec.dtype, tuple(x.src[0] for x in vec.src)),) + vec.src[0].src[1:], vec.src[0].arg) - -def vectorize_alu(vec:UOp): - if not all_same([x.arg for x in vec.src]): return None - return UOp(vec.src[0].op, vec.dtype, tuple(UOp(UOps.VECTORIZE, vec.src[0].src[i].dtype.vec(vec.dtype.count), - tuple(x.src[i] for x in vec.src)) for i in range(len(vec.src[0].src))), vec.src[0].arg) - def fix_unfoldable_image_load(load:UOp, buf:UOp): if not isinstance(buf.dtype, ImageDType) or load.src[1].dtype.count == 2: return None id4 = load.src[1] % 4 @@ -84,8 +73,6 @@ def fix_unfoldable_image_load(load:UOp, buf:UOp): float4_folding = PatternMatcher([ (UPat(UOps.EXPAND, src=UPat(UOps.LOAD, src=(UPat.var("buf"), UPat()), allow_any_len=True), name="ex"), fold_expanded), (UPat((UOps.BARRIER, UOps.SINK), src=UPat(UOps.STORE, src=(UPat.var("buf"), UPat(), UPat()), allow_any_len=True), name="ex"), fold_expanded), - (UPat(UOps.VECTORIZE, src=UPat(UOps.REDUCE), name="vec"), vectorize_reduce), - (UPat(UOps.VECTORIZE, src=UPat((UOps.ALU, UOps.CAST, UOps.BITCAST)), name="vec"), vectorize_alu), ]) # ***** mod *****