test fixes

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:
@@ -185,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())
 
@@ -366,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)

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

tinygrad/codegen/__init__.py

@@ -136,7 +136,7 @@ pm_linearize_cleanups = PatternMatcher([
   # if statements are not allowed in the graph
   (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
   # gated STORE becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
-  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
+  (UPat(Ops.STORE, name="u", src=(UPat((Ops.INDEX, Ops.SHRINK)).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
    lambda u, gate: ((st:=u.replace(src=u.src[0:2])), [mif:=UOp(Ops.IF, src=(gate, u.src[0])), st, UOp(Ops.ENDIF, src=(mif,))]))
 ])
 

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})

tinygrad/uop/ops.py

@@ -304,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:

tinygrad/uop/spec.py

@@ -214,7 +214,7 @@ spec_program = PatternMatcher([
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
 
   # if has a <gate, index_for_dedup>
-  (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