fix precompiled STORE+AFTER (#16822)

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

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)