fix gated store with index in python backend (#9703)

* add default gate in index

* assert store

* add TestRendererFailures

- move test_gated_store_with_alu to new TestRenderFailures class for
tests that fail on multiple renderers
- add test_renderer_failures.py run on python CI

* add test for gated index in 2d

* test TestRenderFailures

.github/workflows/test.yml

@@ -291,6 +291,8 @@ jobs:
       run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
     - name: test_linearizer_failures with Python emulator
       run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_linearizer_failures.py::TestLinearizerFailures::test_failure_1
+    - name: test_renderer_failures with Python emulator
+      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
 
   linter:
     name: Linters

test/test_renderer_failures.py

@@ -9,6 +9,7 @@ from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import dedup, flatten, prod
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.ops import UOp, Ops
 from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor, _to_np_dtype
@@ -27,6 +28,28 @@ def _test_uop_result(inputs:List[Tensor], stores:List[UOp], local_size=None):
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 
+class TestRendererFailures(unittest.TestCase):
+  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
+  def test_gated_store_with_alu(self):
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx0', 4))).ne(0)
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, gate_alu), UOp.const(dtypes.int, 1)))
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
+    uops = linearize_uop(full_graph_rewrite(sink, Device[Device.DEFAULT].renderer))
+    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
+    np.testing.assert_equal(ret, [0, 1, 1, 1])
+
+  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
+  def test_gated_store_with_alu_2d(self):
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx0', 4))).ne(0)
+    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx1', 2))).ne(0)
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0+lidx1*4, gate_alu_0&gate_alu_1), UOp.const(dtypes.int, 1)))
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
+    uops = linearize_uop(full_graph_rewrite(sink, Device[Device.DEFAULT].renderer))
+    ret = _test_uop_result([], uops, local_size=[4, 2, 1])[0]
+    np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
+
 @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
 class TestCStyleFailures(unittest.TestCase):
   def test_inline_const_alu(self):
@@ -44,15 +67,6 @@ class TestCStyleFailures(unittest.TestCase):
 
 @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "tests for ptx renderer")
 class TestPTXFailures(unittest.TestCase):
-  def test_gated_store_with_alu(self):
-    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
-    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (), ('lidx0', 4))).ne(0)
-    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, gate_alu), UOp.const(dtypes.int, 1)))
-    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
-    uops = linearize_uop(full_graph_rewrite(sink, Device[Device.DEFAULT].renderer))
-    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
-    np.testing.assert_equal(ret, [0, 1, 1, 1])
-
   @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
   def test_gated_store_with_if(self):
     a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)

tinygrad/runtime/ops_python.py

@@ -18,7 +18,7 @@ def _load(m, i):
 
 def load(inp, j=0):
   if len(inp) == 2: return [_load(m, x+j if x is not None else None) if gate else default for (m,x,gate),default in zip(*inp)]
-  return [_load(m, x+j if x is not None else None) for m,x in inp[0]]
+  return [_load(m, x+j if x is not None else None) for m,x,_ in inp[0]]
 
 def _store(m, i, v):
   if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
@@ -46,13 +46,13 @@ class PythonProgram:
         dtp = [dl[v] for v in idp if self.uops[v][0] not in void_ops]
         if getenv("TRACE"): print(i, uop, dtype, arg, inp, dtp)
         if uop is Ops.STORE:
-          if len(inp) == 2: inp.append([True] * len(inp[0]))  # set the gate to True
+          assert len(inp) == 2, "expected store is ([(memory, offset, gate)], [value])"
           if dtp[1].count > 1:
             for j,val in enumerate(inp[1]):
-              for (m,o),v,g in zip(inp[0], val, inp[2]):
+              for (m,o,g),v in zip(inp[0], val):
                 if g: _store(m, o+j, v)
           else:
-            for (m,o),v,g in zip(*inp):
+            for (m,o,g),v in zip(*inp):
               if g: _store(m, o, v)
           i += 1
           continue
@@ -87,8 +87,7 @@ class PythonProgram:
               else: ret.append((m, ox*4 + oy*dtp[0].shape[1]*4))
           else:
             for m,o in zip(inp[0], inp[1]): ret.append((m,o))
-          if len(inp) == 3: ret = [(m,o,g) for (m,o),g in zip(ret, inp[2])] # set the gate last
-          ul[i] = ret
+          ul[i] = [(m,o,g) for (m,o),g in zip(ret, inp[2] if len(inp) == 3 else [True]*len(ret))] # set the gate last
         elif uop is Ops.CAST and isinstance(dtype, PtrDType):
           ul[i] = inp[0]
         elif uop is Ops.RANGE: