diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 3a942c5522..909385f073 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -246,6 +246,24 @@ class TestFunction(unittest.TestCase): r0 = f(buf, x, v.bind(0)).numpy() np.testing.assert_equal(r0, [[1.,0.,0.,0.,0.,0.,0.,0.], [2.,0.,0.,0.,0.,0.,0.,0.]]) + def test_single_after_store_precompile(self): + """precompiled AFTER(buf, STORE(view, data)) should return buf after the store.""" + @function(precompile=True) + def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor: + slice_uop = buf[:, start_pos:start_pos+1].uop + assigned = Tensor(buf.uop.after(slice_uop.store(x.uop))) + return assigned + + x = Tensor([[1.], [2.]]).realize() + v = UOp.variable("sp", 0, 7) + for sp in (0, 2): + with self.subTest(sp=sp): + buf = Tensor.zeros(2, 8).clone().realize() + expected = np.zeros((2, 8), dtype=np.float32) + expected[:, sp] = [1., 2.] + np.testing.assert_equal(f(buf, x, v.bind(sp)).numpy(), expected) + np.testing.assert_equal(buf.numpy(), expected) + @unittest.expectedFailure def test_assign_slice(self): @function diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 45a8ee10e9..8dfad59647 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -128,7 +128,7 @@ def transform_precompiled_call(c:UOp) -> UOp|None: subs[s] = placed items.append(s.after(*after_deps) if after_deps else s) else: - items.append(t.after(t.store(s), *after_deps)) + items.append(t.after(t.store(s.after(*after_deps)))) fxn = UOp.sink(*(x.substitute(subs) for x in items)) # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)