revert ops changes

test/unit/test_function.py

@@ -557,15 +557,6 @@ class TestFunctionTuple(unittest.TestCase):
     def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
     with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
 
-  def test_custom_kernel_bound_scalar(self):
-    def add_scalar(out:UOp, x:UOp, value:UOp):
-      i = UOp.range(x.shape[0], 0)
-      return out[i].store(x[i] + value).end(i).sink(arg=KernelInfo(name="add_scalar"))
-    value = Tensor(UOp.variable("value", 0, 10).bind(3))
-    x = Tensor([0, 1, 2, 3])
-    out = Tensor.custom_kernel(Tensor.empty_like(x), x, value, fxn=add_scalar)[0]
-    np.testing.assert_equal(out.numpy(), [3, 4, 5, 6])
-
   def test_custom_kernel_write_only_persistent_output_is_implicit(self):
     # a write-only custom_kernel output that is a realized buffer must be captured
     def write(C:UOp, A:UOp) -> UOp:

tinygrad/llm/kernels/__init__.py

@@ -73,4 +73,9 @@ def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor,
     from tinygrad.llm.kernels.amd import _gated_delta_prefill_kernel as kernel
   core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
   srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
-  return Tensor.custom_kernel(*srcs, *((start_pos,) if start_pos is not None else ()), fxn=kernel)[0]
+  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=kernel)[0]
+  contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
+  params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
+  assert start_pos.uop.op is Ops.BIND
+  call = kernel(*params, start_pos.uop.src[0]).call(*contig, start_pos.uop)
+  return Tensor(contig[0].after(call))

tinygrad/uop/ops.py

@@ -1157,7 +1157,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
   def param_like(self, slot:int):
     addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot))
+    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=addrspace))
     return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
 
   @staticmethod
@@ -1177,8 +1177,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     body = self if self.op is Ops.TUPLE else UOp.maketuple(self)
     return UOp(Ops.FUNCTION, src=(body,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    contig_srcs = tuple(x if x.op in (Ops.AFTER, Ops.BIND) else x.contiguous() for x in srcs)
-    placeholders = [s.param_like(i) if s.op is Ops.BIND else UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
+    contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
+    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
     kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
     return [s.after(kernel) for s in contig_srcs]