From 8fd147e44b203b43adcd5a8aaa0491491d0da41f Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 4 Aug 2026 14:54:41 +0000 Subject: [PATCH] revert ops changes --- test/unit/test_function.py | 9 --------- tinygrad/llm/kernels/__init__.py | 7 ++++++- tinygrad/uop/ops.py | 6 +++--- 3 files changed, 9 insertions(+), 13 deletions(-) diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 9714f64ebb..e333d8f9a6 100644 --- a/test/unit/test_function.py +++ b/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: diff --git a/tinygrad/llm/kernels/__init__.py b/tinygrad/llm/kernels/__init__.py index e469e6a2c1..d6406d2cb8 100644 --- a/tinygrad/llm/kernels/__init__.py +++ b/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)) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 6c7847cccb..8c708e2bdf 100644 --- a/tinygrad/uop/ops.py +++ b/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]