From bc91fffc5df3217745302b270db5c54f8eb6e83e Mon Sep 17 00:00:00 2001 From: Ignacio Sica Date: Thu, 3 Apr 2025 12:48:28 +0800 Subject: [PATCH] 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 | 2 ++ test/test_renderer_failures.py | 32 +++++++++++++++++++++++--------- tinygrad/runtime/ops_python.py | 11 +++++------ 3 files changed, 30 insertions(+), 15 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e081e0eb90..4cd1b34a77 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/test/test_renderer_failures.py b/test/test_renderer_failures.py index d37c0a5f65..e428271f7f 100644 --- a/test/test_renderer_failures.py +++ b/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) diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index c30b3fd2f6..3d66448c0c 100644 --- a/tinygrad/runtime/ops_python.py +++ b/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: