From a9fbc7db7bef3ddb2843abc82d05b9754eb0783e Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Mon, 13 Jul 2026 18:53:32 -0400 Subject: [PATCH] expect _offset support, CL and WEBGPU are outliers (#17014) --- test/backend/test_graph.py | 4 ++-- test/backend/test_multitensor.py | 4 ++-- test/backend/test_subbuffer.py | 2 +- tinygrad/callify.py | 5 ----- tinygrad/device.py | 7 ++----- tinygrad/runtime/ops_hip.py | 1 + tinygrad/runtime/ops_null.py | 2 +- tinygrad/schedule/memory.py | 4 ++-- tinygrad/uop/ops.py | 8 ++++++++ 9 files changed, 19 insertions(+), 18 deletions(-) diff --git a/test/backend/test_graph.py b/test/backend/test_graph.py index a10ba30934..20fefe99c3 100644 --- a/test/backend/test_graph.py +++ b/test/backend/test_graph.py @@ -63,7 +63,7 @@ def zero_bufs(bufs): @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required") class TestGraph(unittest.TestCase): def skip_if_no_offset(self): - if not hasattr(Device[Device.DEFAULT].allocator, "_offset"): self.skipTest("device does not support _offset") + if Device.DEFAULT in {"WEBGPU", "CL"}: self.skipTest("device does not support _offset") def skip_if_not_multigraph(self): graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g @@ -213,8 +213,8 @@ class TestGraph(unittest.TestCase): def test_graph_offset_bufs(self): self.skip_if_not_multigraph() + self.skip_if_no_offset() d0 = Device.DEFAULT - if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset") b0 = make_buffer(d0, fill=True) b1 = make_view(b0, 0, b0.size) diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index a0a65e1021..caa127dfc7 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -385,7 +385,7 @@ class TestMultiBufferView(unittest.TestCase): b_ref = view_fn(a_ref) b_multi = view_fn(a_multi).contiguous() linear, var_vals = b_multi.linear_with_vars() - if all(hasattr(Device[d].allocator, "_offset") for d in b_multi.device): + if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device): compiled = [call for call in linear.src if call.src[0].op is Ops.SINK] self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}") run_linear(linear, var_vals) @@ -417,7 +417,7 @@ class TestMultiBufferView(unittest.TestCase): a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize() out = a[5].contiguous() linear, var_vals = out.linear_with_vars() - if all(hasattr(Device[d].allocator, "_offset") for d in out.device): + if all(not d.startswith(("WEBGPU", "CL")) for d in out.device): compiled = [call for call in linear.src if call.src[0].op is Ops.SINK] self.assertEqual(len(compiled), 0) run_linear(linear, var_vals) diff --git a/test/backend/test_subbuffer.py b/test/backend/test_subbuffer.py index d6d866597a..8b1b3aef3c 100644 --- a/test/backend/test_subbuffer.py +++ b/test/backend/test_subbuffer.py @@ -4,7 +4,7 @@ from tinygrad.device import Buffer from tinygrad.helpers import Context, DEV from test.helpers import needs_second_gpu -@unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported") +@unittest.skipIf(Device.DEFAULT in {"WEBGPU", "CL"}, "subbuffer not supported") class TestSubBuffer(unittest.TestCase): def setUp(self): self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated() diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 176894000b..4299e19c48 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -73,11 +73,6 @@ def contiguous_mops_to_view(c:UOp, src:UOp): # no symbolic shape if not all_int(c.shape): return None - # check if view is supported - from tinygrad.device import Device - devs = (src.device,) if isinstance(src.device, str) else src.device - if not all(hasattr(Device[d].allocator, "_offset") for d in devs): return None - if buf.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None: view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape) return c.replace(src=(view,)) if c.op is Ops.COPY else view diff --git a/tinygrad/device.py b/tinygrad/device.py index 80b65a78e5..b53e65faa1 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -135,9 +135,7 @@ class Buffer: if device not in self._bufs: allocator = Device[device].allocator if device == self.device: self.ensure_allocated() - elif self._base is not None: - assert hasattr(allocator, "_offset"), "offset function required for view" - self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset) + elif self._base is not None: self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset) else: self._bufs[device] = allocator.map(self.ensure_allocated()) return self._bufs[device] def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_initialized() else self @@ -152,7 +150,6 @@ class Buffer: if self._base is not None: self._base.ensure_allocated() self._base.allocated_views += 1 - assert hasattr(self.allocator, "_offset"), "offset function required for view" self._bufs[self.device] = self.allocator._offset(self.base._buf, self.nbytes, self.offset) else: self._bufs[self.device] = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options) @@ -246,7 +243,7 @@ class Allocator(Generic[DeviceType]): def _map(self, buf): raise NotImplementedError("need map") def _unmap(self, mb): pass # default no-op; override if _map allocates iface-side state # def _as_buffer(self, src) -> memoryview: - # def _offset(self, buf, size:int, offset:int): + def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset") # def _transfer(self, dest, src, sz:int, src_dev, dest_dev): def _encode_decode(self, bufout, bufin, desc, hist:list, shape:tuple[int,...], frame_pos:int): raise NotImplementedError("need encdec") # optional diff --git a/tinygrad/runtime/ops_hip.py b/tinygrad/runtime/ops_hip.py index 073f55ea26..3224789669 100644 --- a/tinygrad/runtime/ops_hip.py +++ b/tinygrad/runtime/ops_hip.py @@ -66,3 +66,4 @@ class HIPAllocator(LRUAllocator[HIPDevice]): def _copyout(self, dest:memoryview, src): self.dev.synchronize() check(hip.hipMemcpy(mv_address(dest), src, len(dest), hip.hipMemcpyDeviceToHost)) + def _offset(self, buf, size:int, offset:int): return hip.hipDeviceptr_t(buf.value + offset) diff --git a/tinygrad/runtime/ops_null.py b/tinygrad/runtime/ops_null.py index fc355c2a80..c00509eb01 100644 --- a/tinygrad/runtime/ops_null.py +++ b/tinygrad/runtime/ops_null.py @@ -28,7 +28,7 @@ class NullAllocator(Allocator['NullDevice']): if not NULL_ALLOW_COPYOUT: raise RuntimeError("no copyout on NULL") def _transfer(self, dest, src, sz:int, src_dev, dest_dev): with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", f"{src_dev.device}:SDMA:0"): pass - def _offset(self, buf, offset:int, size:int): pass + def _offset(self, buf, size:int, offset:int): pass class NullGraph(MultiGraphRunner): def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None: diff --git a/tinygrad/schedule/memory.py b/tinygrad/schedule/memory.py index 8308ba5224..4a77f17bb2 100644 --- a/tinygrad/schedule/memory.py +++ b/tinygrad/schedule/memory.py @@ -1,5 +1,4 @@ from collections import defaultdict -from tinygrad.device import Device from tinygrad.helpers import NO_MEMORY_PLANNER, DEBUG, round_up from tinygrad.uop.ops import UOp, Ops from tinygrad.dtype import dtypes @@ -13,7 +12,8 @@ def _collect_bufs(u:UOp) -> list[UOp]: def _can_plan(b:UOp, held_bufs:set[UOp]) -> bool: if b in held_bufs: return False devs = (b.device,) if isinstance(b.device, str) else b.device - return all(not d.startswith(("DISK", "TINYFS")) and hasattr(Device[d].allocator, "_offset") for d in devs) + # CL and WEBGPU do not support views, see explanation in contiguous_view_offset + return all(not d.startswith(("DISK", "TINYFS", "CL", "WEBGPU")) for d in devs) LaneKey = tuple[str, int] diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 39c556f835..56458c0e25 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -822,6 +822,14 @@ class UOp(RandMixin, metaclass=UOpMetaClass): """If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None.""" from tinygrad.schedule.rangeify import pm_mops from tinygrad.uop.symbolic import symbolic + + # WEBGPU and CL do not support views. + # WEBGPU requires that minUniformBufferOffsetAlignment be at least 32 bytes: https://gpuweb.github.io/gpuweb/#adapter-capability-guarantees + # CL 1.1 provides the clCreateSubBuffer API, but at the time of writing, relevant CL runtimes (rusticl, adreno, nvidia, amd) do not provide + # reasonable values for CL_DEVICE_MEM_BASE_ADDR_ALIGN. cl_ext_buffer_device_address could potentially help, but this extension is not provided + # by relevant CL runtimes at time of writing. + if any(d.startswith(("WEBGPU", "CL")) for d in ((self.device,) if isinstance(self.device, str) else self.device)): return None + numel = self.numel() out = graph_rewrite(self.flatten().index(UOp.range(numel, 0)), pm_mops+symbolic, name="contiguous_view_offset") if out.op is not Ops.INDEX: return None