diff --git a/test/unit/test_hcq_graph.py b/test/unit/test_hcq_graph.py index d8f0322a4b..270aa988d4 100644 --- a/test/unit/test_hcq_graph.py +++ b/test/unit/test_hcq_graph.py @@ -11,7 +11,7 @@ from test.mockgpu.usb import MockUSB @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "HCQ device required to run") class TestHCQUnit(unittest.TestCase): @unittest.skipIf(Device.DEFAULT == "CPU", "requires non-CPU HCQ device") - def test_supports_exec_item(self): + def test_supports_uop(self): d0, cpu_dev = Device[Device.DEFAULT], Device["CPU"] @TinyJit @@ -20,23 +20,23 @@ class TestHCQUnit(unittest.TestCase): inp, inp_cpu = Tensor.randn(10, 10, device=Device.DEFAULT).realize(), Tensor.randn(10, 10, device="CPU").realize() for _ in range(5): f(inp, inp_cpu) - # construct minimal CALL UOps for supports_exec_item (graphs only see PROGRAMs after compile_linear) + # construct minimal CALL UOps for supports_uop (graphs only see PROGRAMs after compile_linear) gpu_call = UOp(Ops.PROGRAM).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float)) cpu_call = UOp(Ops.PROGRAM).call(UOp.new_buffer("CPU", 1, dtypes.float)) gpu_devs = [d0] # local MMIO: GPU works alone and with CPU in batch (cpu_support=True) - assert HCQGraph.supports_exec_item(gpu_devs, gpu_call) is True - assert HCQGraph.supports_exec_item(gpu_devs, cpu_call) is True - assert HCQGraph.supports_exec_item(gpu_devs + [cpu_dev], gpu_call) is True + assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True + assert HCQGraph.supports_uop(gpu_devs, cpu_call) is True + assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is True # USB MMIO: GPU-only still works, but CPU batching must be rejected (cpu_support=False) orig_view = d0.timeline_signal.base_buf.view try: d0.timeline_signal.base_buf.view = USBMMIOInterface(MockUSB(bytearray(256)), 0, 16, fmt='B') - assert HCQGraph.supports_exec_item(gpu_devs, gpu_call) is True - assert HCQGraph.supports_exec_item(gpu_devs, cpu_call) is False - assert HCQGraph.supports_exec_item(gpu_devs + [cpu_dev], gpu_call) is False + assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True + assert HCQGraph.supports_uop(gpu_devs, cpu_call) is False + assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is False finally: d0.timeline_signal.base_buf.view = orig_view diff --git a/test/unit/test_metal_graph.py b/test/unit/test_metal_graph.py index 74c733ba90..3038dcfb34 100644 --- a/test/unit/test_metal_graph.py +++ b/test/unit/test_metal_graph.py @@ -27,18 +27,18 @@ class TestMetalGraph(unittest.TestCase): c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs) return c - def test_supports_exec_item_normal_offset(self): - assert self.MetalGraph.supports_exec_item([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True + def test_supports_uop_normal_offset(self): + assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True - def test_supports_exec_item_overflow_offset(self): - assert self.MetalGraph.supports_exec_item([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False + def test_supports_uop_overflow_offset(self): + assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False - def test_supports_exec_item_nonmetal_buf(self): + def test_supports_uop_nonmetal_buf(self): # non-BUFFER_VIEW ops should not be checked for offset buf = MagicMock() buf.op = Ops.BUFFER buf.device = Device.DEFAULT - self.MetalGraph.supports_exec_item([self.dev], self.call(buf)) + self.MetalGraph.supports_uop([self.dev], self.call(buf)) if __name__ == "__main__": unittest.main() diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 6f7fe24763..6c8d3fff3f 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -48,8 +48,8 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp: devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))]) graph_t = graph_class(devs[0]) if devs[0].graph is not None else None - can_graph = graph_t is not None and graph_t.supports_exec_item(devs, si) - can_extend = can_graph and graph_t is not None and (not current_batch_devs or graph_t.supports_exec_item(current_batch_devs, si)) \ + can_graph = graph_t is not None and graph_t.supports_uop(devs, si) + can_extend = can_graph and graph_t is not None and (not current_batch_devs or graph_t.supports_uop(current_batch_devs, si)) \ and (max_batch_size == 0 or len(current_batch) < max_batch_size) if not can_extend and current_batch: flush_batch() @@ -166,13 +166,13 @@ class GraphRunner: for x in (b.device if isinstance(b.device, tuple) else (b.device,))]) @staticmethod - def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool: + def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool: return new_call.src[0].op is Ops.PROGRAM and len(GraphRunner._all_devs(batch_devs, new_call)) == 1 # a marker for your graph supporting multiple devices of the same type class MultiGraphRunner(GraphRunner): @staticmethod - def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool: + def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool: # Devices must be the same type return new_call.src[0].op in (Ops.PROGRAM, Ops.COPY) and len(dedup([type(d) for d in GraphRunner._all_devs(batch_devs, new_call)])) == 1 diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index 2cc769f887..e4da62d609 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -316,7 +316,7 @@ class HCQGraph(MultiGraphRunner): for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferSpec(cpu_access=True)) @staticmethod - def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool: + def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool: # Check if all devices are HCQ all_devs = cast(list[HCQCompiled], GraphRunner._all_devs(batch_devs, new_call)) if not all(issubclass(type(d), HCQCompiled) for d in all_devs): return False diff --git a/tinygrad/runtime/graph/metal.py b/tinygrad/runtime/graph/metal.py index b33305c607..25c89257e5 100644 --- a/tinygrad/runtime/graph/metal.py +++ b/tinygrad/runtime/graph/metal.py @@ -107,7 +107,7 @@ class MetalGraph(GraphRunner): self.collect_timestamps() @staticmethod - def supports_exec_item(batch_devs, new_call:UOp) -> bool: + def supports_uop(batch_devs, new_call:UOp) -> bool: # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range. if any(b.op is Ops.BUFFER_VIEW and b.arg[1] * b.dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False - return GraphRunner.supports_exec_item(batch_devs, new_call) + return GraphRunner.supports_uop(batch_devs, new_call)