github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

add _alloc_signal/_free_signal to hcq (#5264)

Browse Source 
* add _alloc_signal/_free_signal api

* oops, revert this

* linter
This commit is contained in:
nimlgen
2024-07-02 23:35:39 +03:00
committed by GitHub
parent 9a25ee0b9a
commit 7be776f9af
7 changed files with 54 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
test/external/external_test_hcq.py vendored
View File

@@ -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)

2
test/external/fuzz_kfd.py vendored
View File

@@ -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))

39
test/test_hcq.py
View File

@@ -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")

9
tinygrad/device.py
View File

@@ -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):

8
tinygrad/runtime/graph/hcq.py
View File

@@ -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))

11
tinygrad/runtime/ops_amd.py
View File

@@ -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

11
tinygrad/runtime/ops_nv.py
View File

@@ -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