diff --git a/test/external/external_test_hcq.py b/test/external/external_test_hcq.py index c992560209..a7be0047ea 100644 --- a/test/external/external_test_hcq.py +++ b/test/external/external_test_hcq.py @@ -51,7 +51,7 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0.synchronize() # wait for copyins to complete def test_run_1000_times_one_submit(self): - temp_signal, temp_value = TestHCQ.d0._get_signal(value=0), 0 + temp_signal, temp_value = TestHCQ.d0._alloc_signal(value=0), 0 q = TestHCQ.compute_queue() for _ in range(1000): q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) @@ -69,7 +69,7 @@ class TestHCQ(unittest.TestCase): assert (val:=TestHCQ.a.lazydata.buffer.as_buffer().cast("f")[0]) == 2000.0, f"got val {val}" def test_run_1000_times(self): - temp_signal = TestHCQ.d0._get_signal(value=0) + temp_signal = TestHCQ.d0._alloc_signal(value=0) q = TestHCQ.compute_queue() q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) q.signal(temp_signal, 2).wait(temp_signal, 2) @@ -84,7 +84,7 @@ class TestHCQ(unittest.TestCase): assert (val:=TestHCQ.a.lazydata.buffer.as_buffer().cast("f")[0]) == 2000.0, f"got val {val}" def test_run_to_3(self): - temp_signal = TestHCQ.d0._get_signal(value=0) + temp_signal = TestHCQ.d0._alloc_signal(value=0) q = TestHCQ.compute_queue() q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) q.signal(temp_signal, 1).wait(temp_signal, 1) @@ -109,7 +109,7 @@ class TestHCQ(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind") def test_bind_run(self): - temp_signal = TestHCQ.d0._get_signal(value=0) + temp_signal = TestHCQ.d0._alloc_signal(value=0) q = TestHCQ.compute_queue() q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) q.signal(temp_signal, 2).wait(temp_signal, 2) @@ -141,7 +141,7 @@ class TestHCQ(unittest.TestCase): @unittest.skipIf(CI, "Can't handle async update on CPU") def test_wait_signal(self): - temp_signal = TestHCQ.d0._get_signal(value=0) + temp_signal = TestHCQ.d0._alloc_signal(value=0) TestHCQ.compute_queue().wait(temp_signal, value=1).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0) with self.assertRaises(RuntimeError): TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50) @@ -152,7 +152,7 @@ class TestHCQ(unittest.TestCase): @unittest.skipIf(CI, "Can't handle async update on CPU") def test_wait_copy_signal(self): - temp_signal = TestHCQ.d0._get_signal(value=0) + temp_signal = TestHCQ.d0._alloc_signal(value=0) TestHCQ.copy_queue().wait(temp_signal, value=1).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0) with self.assertRaises(RuntimeError): TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50) @@ -267,7 +267,7 @@ class TestHCQ(unittest.TestCase): q = TestHCQ.compute_queue() qc = TestHCQ.copy_queue() q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) # b = [1, 2] - q.signal(sig:=TestHCQ.d0._get_signal(value=0), value=1) + q.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=1) qc.wait(sig, value=1) qc.copy(TestHCQ.a.lazydata.buffer._buf.va_addr, TestHCQ.b.lazydata.buffer._buf.va_addr, 8) qc.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) @@ -282,7 +282,7 @@ class TestHCQ(unittest.TestCase): d1 = Device[f"{Device.DEFAULT}:1"] q1 = TestHCQ.compute_queue() q2 = TestHCQ.compute_queue() - q1.signal(sig:=TestHCQ.d0._get_signal(value=0), value=0xfff) + q1.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=0xfff) q2.wait(sig, value=0xfff) q2.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) q2.submit(TestHCQ.d0) diff --git a/test/external/fuzz_kfd.py b/test/external/fuzz_kfd.py index d7bc396fda..527ef20d0e 100644 --- a/test/external/fuzz_kfd.py +++ b/test/external/fuzz_kfd.py @@ -16,7 +16,7 @@ if __name__ == "__main__": d2, b2 = random.choice(buffers) d1._gpu_map(b2) q = HWComputeQueue() - q.signal(sig:=AMDDevice._get_signal(10)) + q.signal(sig:=AMDDevice._alloc_signal(10)) qc = HWCopyQueue() qc.wait(sig) qc.copy(b1.va_addr, b2.va_addr, min(b1.size, b2.size)) diff --git a/test/test_hcq.py b/test/test_hcq.py index 035b33ea65..a5c13acbc5 100644 --- a/test/test_hcq.py +++ b/test/test_hcq.py @@ -48,7 +48,7 @@ class TestHCQ(unittest.TestCase): def test_signal_update(self): for queue_type in [TestHCQ.d0.hw_compute_queue_t]: with self.subTest(name=str(queue_type)): - q = queue_type().signal(fake_signal := TestHCQ.d0._get_signal(), 0x1000) + q = queue_type().signal(fake_signal := TestHCQ.d0._alloc_signal(), 0x1000) q.update_signal(0, signal=TestHCQ.d0.timeline_signal, value=TestHCQ.d0.timeline_value).submit(TestHCQ.d0) TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) @@ -58,26 +58,26 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) TestHCQ.d0.timeline_value += 1 - TestHCQ.d0.signals_pool.append(fake_signal) + TestHCQ.d0._free_signal(fake_signal) # Test wait def test_wait(self): for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]: with self.subTest(name=str(queue_type)): - fake_signal = TestHCQ.d0._get_signal() + fake_signal = TestHCQ.d0._alloc_signal() TestHCQ.d0._set_signal(fake_signal, 1) queue_type().wait(fake_signal, 1) \ .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0) TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) TestHCQ.d0.timeline_value += 1 - TestHCQ.d0.signals_pool.append(fake_signal) + TestHCQ.d0._free_signal(fake_signal) @unittest.skipIf(MOCKGPU, "Can't handle async update on MOCKGPU for now") def test_wait_late_set(self): for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]: with self.subTest(name=str(queue_type)): - fake_signal = TestHCQ.d0._get_signal() + fake_signal = TestHCQ.d0._alloc_signal() queue_type().wait(fake_signal, 1) \ .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0) @@ -89,12 +89,12 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0.timeline_value += 1 - TestHCQ.d0.signals_pool.append(fake_signal) + TestHCQ.d0._free_signal(fake_signal) def test_wait_update(self): for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]: with self.subTest(name=str(queue_type)): - fake_signal = TestHCQ.d0._get_signal() + fake_signal = TestHCQ.d0._alloc_signal() q = queue_type().wait(TestHCQ.d0.timeline_signal, 0xffffffff).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) TestHCQ.d0._set_signal(fake_signal, 0x30) @@ -103,7 +103,7 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) TestHCQ.d0.timeline_value += 1 - TestHCQ.d0.signals_pool.append(fake_signal) + TestHCQ.d0._free_signal(fake_signal) # Test exec def test_exec_one_kernel(self): @@ -167,7 +167,7 @@ class TestHCQ(unittest.TestCase): with self.subTest(name=str(queue_type)): if not hasattr(queue_type(), 'bind'): self.skipTest("queue does not support bind api") - fake_signal = TestHCQ.d0._get_signal() + fake_signal = TestHCQ.d0._alloc_signal() q = queue_type().wait(TestHCQ.d0.timeline_signal, 0xffffffff).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) q.bind(TestHCQ.d0) @@ -177,13 +177,13 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value) TestHCQ.d0.timeline_value += 1 - TestHCQ.d0.signals_pool.append(fake_signal) + TestHCQ.d0._free_signal(fake_signal) # Test multidevice def test_multidevice_signal_wait(self): d1 = Device[f"{Device.DEFAULT}:1"] - TestHCQ.d0.hw_copy_queue_t().signal(sig:=TestHCQ.d0._get_signal(value=0), value=0xfff) \ + TestHCQ.d0.hw_copy_queue_t().signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=0xfff) \ .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0) d1.hw_copy_queue_t().wait(sig, value=0xfff) \ @@ -195,13 +195,13 @@ class TestHCQ(unittest.TestCase): d1._wait_signal(d1.timeline_signal, d1.timeline_value) d1.timeline_value += 1 - TestHCQ.d0.signals_pool.append(sig) + TestHCQ.d0._free_signal(sig) # Test profile api def test_speed_exec_time(self): TestHCQ.d0._prof_setup() - sig_st, sig_en = TestHCQ.d0._get_signal(), TestHCQ.d0._get_signal() + sig_st, sig_en = TestHCQ.d0._alloc_signal(), TestHCQ.d0._alloc_signal() TestHCQ.d0.hw_compute_queue_t().timestamp(sig_st) \ .exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) \ .timestamp(sig_en) \ @@ -212,7 +212,8 @@ class TestHCQ(unittest.TestCase): et = TestHCQ.d0._gpu2cpu_time(TestHCQ.d0._read_timestamp(sig_en), True) - TestHCQ.d0._gpu2cpu_time(TestHCQ.d0._read_timestamp(sig_st), True) - TestHCQ.d0.signals_pool += [sig_st, sig_en] + TestHCQ.d0._free_signal(sig_st) + TestHCQ.d0._free_signal(sig_en) print(f"exec kernel time: {et:.2f} us") assert 1 <= et <= (2000 if CI else 20) @@ -225,7 +226,7 @@ class TestHCQ(unittest.TestCase): a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferOptions(nolru=True)).allocate() b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferOptions(nolru=True)).allocate() - sig_st, sig_en = TestHCQ.d0._get_signal(), TestHCQ.d0._get_signal() + sig_st, sig_en = TestHCQ.d0._alloc_signal(), TestHCQ.d0._alloc_signal() TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \ .copy(a._buf.va_addr, b._buf.va_addr, SZ) \ .timestamp(sig_en) \ @@ -237,7 +238,8 @@ class TestHCQ(unittest.TestCase): et = TestHCQ.d0._gpu2cpu_time(TestHCQ.d0._read_timestamp(sig_en), True) - TestHCQ.d0._gpu2cpu_time(TestHCQ.d0._read_timestamp(sig_st), True) et_ms = et / 1e3 - TestHCQ.d0.signals_pool += [sig_st, sig_en] + TestHCQ.d0._free_signal(sig_st) + TestHCQ.d0._free_signal(sig_en) gb_s = ((SZ / 1e9) / et_ms) * 1e3 print(f"same device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s") @@ -251,7 +253,7 @@ class TestHCQ(unittest.TestCase): a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferOptions(nolru=True)).allocate() TestHCQ.d0._gpu_map(b._buf) - sig_st, sig_en = TestHCQ.d0._get_signal(), TestHCQ.d0._get_signal() + sig_st, sig_en = TestHCQ.d0._alloc_signal(), TestHCQ.d0._alloc_signal() TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \ .copy(a._buf.va_addr, b._buf.va_addr, SZ) \ .timestamp(sig_en) \ @@ -263,7 +265,8 @@ class TestHCQ(unittest.TestCase): et = TestHCQ.d0._gpu2cpu_time(TestHCQ.d0._read_timestamp(sig_en), True) - TestHCQ.d0._gpu2cpu_time(TestHCQ.d0._read_timestamp(sig_st), True) et_ms = et / 1e3 - TestHCQ.d0.signals_pool += [sig_st, sig_en] + TestHCQ.d0._free_signal(sig_st) + TestHCQ.d0._free_signal(sig_en) gb_s = ((SZ / 1e9) / et_ms) * 1e3 print(f"cross device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s") diff --git a/tinygrad/device.py b/tinygrad/device.py index 0cbf6b33af..b270de12d5 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -188,7 +188,7 @@ class Compiled: @contextlib.contextmanager def hcq_profile(dev, queue_type, enabled, desc): - st, en = (dev._get_signal(), dev._get_signal()) if enabled else (None, None) + st, en = (dev._alloc_signal(), dev._alloc_signal()) if enabled else (None, None) if enabled: queue_type().timestamp(st).submit(dev) try: yield (st, en) finally: @@ -217,7 +217,10 @@ class HCQCompatCompiled(Compiled): def _set_signal(self, sig, value): raise NotImplementedError("need _set_signal") # sets a value for a signal @classmethod - def _get_signal(self, value=0, **kwargs): raise NotImplementedError("need _get_signal") # allocates a new signal + def _alloc_signal(self, value=0, **kwargs): raise NotImplementedError("need _alloc_signal") # allocates a new signal + + @classmethod + def _free_signal(self, sig): raise NotImplementedError("need _free_signal") # frees a signal @classmethod def _wait_signal(self, signal, value=0, timeout=10000): raise NotImplementedError("need _wait_signal") # waits for a signal value @@ -239,7 +242,7 @@ class HCQCompatCompiled(Compiled): def _prof_process_events(self): self.raw_prof_records += [(self._read_timestamp(st), self._read_timestamp(en), name, is_cp) for st, en, name, is_cp in self.sig_prof_records] - for st, en, _, _ in self.sig_prof_records: self.signals_pool += [st, en] # type: ignore + for st, en, _, _ in self.sig_prof_records: map(self._alloc_signal, [st, en]) self.sig_prof_records = [] def _prof_finalize(self): diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index 0db56be3a6..5be28f634a 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -45,8 +45,8 @@ class HCQGraph(MultiGraphRunner): self.copy_queues: Dict[Compiled, Any] = {dev: dev.hw_copy_queue_t() for dev in self.devices} self.signal_sched: Dict[int, Tuple[List, Optional[int], Optional[List]]] = {} # Dict[ji_idx, (deps, sigval, prof_info)] - self.signals: Dict[Any, Any] = {q: self.devices[0]._get_signal(value=0) for q in list(self.comp_queues.values())+list(self.copy_queues.values())} - self.dev_kickoff_signal = {dev: self.devices[0]._get_signal(value=0) for dev in self.devices + ['CPU']} # Dict[dev, signal] + self.signals = {q: self.devices[0]._alloc_signal(value=0) for q in list(self.comp_queues.values())+list(self.copy_queues.values())} + self.dev_kickoff_signal = {dev: self.devices[0]._alloc_signal(value=0) for dev in self.devices + ['CPU']} # Dict[dev, signal] self.kickoff_value = 0 self.save_devs: Dict[Any, Set] = {q: set() for q in list(self.comp_queues.values()) + list(self.copy_queues.values())} @@ -77,7 +77,7 @@ class HCQGraph(MultiGraphRunner): self.signal_sched[val - 1] = self.signal_sched[val - 1][:1] + (val,) + self.signal_sched[val - 1][2:] prof_ji_desc = ji.prg.clprg.name if isinstance(ji.prg, CompiledRunner) else f"{ji.bufs[1].device} -> {ji.bufs[0].device}" # type: ignore - prof_info = ([enqueue_dev._get_signal() for _ in range(2)] + [enqueue_dev, prof_ji_desc, isinstance(ji.prg, BufferXfer)]) if PROFILE else None + prof_info = ([enqueue_dev._alloc_signal() for _ in range(2)] + [enqueue_dev, prof_ji_desc, isinstance(ji.prg, BufferXfer)]) if PROFILE else None self.signal_sched[j] = (deps, None if isinstance(ji.prg, CompiledRunner) else (j + 1), prof_info) self.last_ji[enqueue_queue] = j @@ -183,5 +183,5 @@ class HCQGraph(MultiGraphRunner): if PROFILE and self.kickoff_value > 1: for _,_,(st,en,dev,desc,is_cp) in self.signal_sched.values(): dev.sig_prof_records += [(st, en, desc, is_cp)] #type: ignore - self.devices[0].signals_pool += list(self.dev_kickoff_signal.values()) + list(self.signals.values()) # type: ignore + map(self.devices[0]._free_signal, list(self.dev_kickoff_signal.values()) + list(self.signals.values())) for dev, buf in self.kernargs_bufs.items(): dev.allocator._free(buf, BufferOptions(cpu_access=True)) diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index e9923372d7..646e967cbc 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -366,7 +366,7 @@ class AMDProgram: for i in range(len(args)): args_st.__setattr__(f'f{i}', args[i].va_addr) for i in range(len(vals)): args_st.__setattr__(f'v{i}', vals[i]) - sig_st, sig_en = (self.device._get_signal(), self.device._get_signal()) if PROFILE else (self.device.time_event_st, self.device.time_event_en) + sig_st, sig_en = (self.device._alloc_signal(), self.device._alloc_signal()) if PROFILE else (self.device.time_event_st, self.device.time_event_en) q = HWPM4Queue() q.wait(self.device.timeline_signal, self.device.timeline_value - 1).memory_barrier() @@ -453,7 +453,7 @@ class AMDDevice(HCQCompatCompiled): def _set_signal(self, sig, value): sig.value = value @classmethod - def _get_signal(self, value=0, **kwargs) -> hsa.amd_signal_t: + def _alloc_signal(self, value=0, **kwargs) -> hsa.amd_signal_t: self._set_signal(ret := self.signals_pool.pop(), value) if (sync_event:=kwargs.get('sync_event')) is not None: ret.event_mailbox_ptr = AMDDevice.event_page.va_addr + sync_event.event_slot_index*8 @@ -461,6 +461,9 @@ class AMDDevice(HCQCompatCompiled): else: ret.event_mailbox_ptr = ret.event_id = 0 return ret + @classmethod + def _free_signal(self, sig): self.signals_pool.append(sig) + @classmethod def _wait_signal(self, signal:hsa.amd_signal_t, value=0, timeout=10000): assert signal.event_id != 0, "can't wait on this signal" @@ -496,7 +499,7 @@ class AMDDevice(HCQCompatCompiled): self._gpu_map(AMDDevice.event_page) sync_event = kio.create_event(AMDDevice.kfd, auto_reset=1) - self.time_event_st, self.time_event_en = AMDDevice._get_signal(), AMDDevice._get_signal() + self.time_event_st, self.time_event_en = AMDDevice._alloc_signal(), AMDDevice._alloc_signal() self.kernargs = self._gpu_alloc(0x1000000, kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM) self.kernargs_ptr = self.kernargs.va_addr @@ -515,7 +518,7 @@ class AMDDevice(HCQCompatCompiled): self.sdma_queue = self._alloc_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x100000) super().__init__(device, AMDAllocator(self), AMDRenderer(), AMDCompiler(self.arch), functools.partial(AMDProgram, self), HWPM4Queue, HWCopyQueue, - timeline_signals=[self._get_signal(sync_event=sync_event), self._get_signal(sync_event=kio.create_event(AMDDevice.kfd, auto_reset=1))]) + timeline_signals=[self._alloc_signal(sync_event=sync_event), self._alloc_signal(sync_event=kio.create_event(AMDDevice.kfd, auto_reset=1))]) def _gpu2cpu_time(self, gpu_time, is_copy): if is_copy: return self.copy_cpu_start_time + (gpu_time - self.copy_gpu_start_time) / 1e2 diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index bf219a6c9f..d468df029b 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -328,7 +328,7 @@ class NVProgram: if MOCKGPU: self.constbuffer_0[0:2] = [len(args), len(vals)] kernargs = [arg_half for arg in args for arg_half in nvdata64_le(arg.base)] + list(vals) - sig_st, sig_en = (self.device._get_signal(), self.device._get_signal()) if PROFILE else (self.device.time_event_st, self.device.time_event_en) + sig_st, sig_en = (self.device._alloc_signal(), self.device._alloc_signal()) if PROFILE else (self.device.time_event_st, self.device.time_event_en) queue = HWComputeQueue() queue.wait(self.device.timeline_signal, self.device.timeline_value - 1) @@ -530,7 +530,7 @@ class NVDevice(HCQCompatCompiled): en_fifo_params = nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1) rm_control(self.fd_ctl, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, self.root, channel_group, en_fifo_params) - self.time_event_st, self.time_event_en = NVDevice._get_signal(), NVDevice._get_signal() + self.time_event_st, self.time_event_en = NVDevice._alloc_signal(), NVDevice._alloc_signal() self.cmdq_page: nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS = self._gpu_alloc(0x200000, map_to_cpu=True, huge_page=True) self.cmdq: memoryview = to_mv(self.cmdq_page.base, 0x200000).cast("I") @@ -543,7 +543,7 @@ class NVDevice(HCQCompatCompiled): compiler_t = (PTXCompiler if PTX else CUDACompiler) if MOCKGPU else (NVPTXCompiler if PTX else NVCompiler) super().__init__(device, NVAllocator(self), PTXRenderer(self.arch, device="NV") if PTX else NVRenderer(self.arch), compiler_t(self.arch), - functools.partial(NVProgram, self), HWComputeQueue, HWCopyQueue, timeline_signals=[self._get_signal(), self._get_signal()]) + functools.partial(NVProgram, self), HWComputeQueue, HWCopyQueue, timeline_signals=[self._alloc_signal(), self._alloc_signal()]) self._cmdq_setup_compute_gpfifo() self._cmdq_setup_dma_gpfifo() @@ -560,10 +560,13 @@ class NVDevice(HCQCompatCompiled): def _set_signal(self, sig, value): sig[0] = value @classmethod - def _get_signal(self, value=0, **kwargs) -> memoryview: + def _alloc_signal(self, value=0, **kwargs) -> memoryview: self._set_signal(sig := self.signals_pool.pop(), value) return sig + @classmethod + def _free_signal(self, sig): self.signals_pool.append(sig) + @classmethod def _wait_signal(self, signal, value=0, timeout=10000): start_time = time.time() * 1000