Spaces:
Running
Running
File size: 5,983 Bytes
9d98328 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """Resident bounded delay ring with ordered emissions and explicit snapshot/restore."""
from __future__ import annotations
import ctypes
import math
from collections.abc import Sequence
from typing import Any
from .runtime import CudaDriver, CudaDriverError, DeviceAllocation, DriverModule
def _is_number(value: object) -> bool:
return not isinstance(value, bool) and isinstance(value, (int, float))
class DeviceDelayQueue:
def __init__(
self, driver: CudaDriver, loaded: DriverModule, n_neurons: int, capacity: int
) -> None:
if (
type(n_neurons) is not int
or not 1 <= n_neurons <= 1024
or type(capacity) is not int
or not 1 <= capacity <= 20000
):
raise ValueError("invalid CUDA delay queue capacity")
self.driver, self.n, self.capacity = driver, n_neurons, capacity
self.consume_kernel = driver.kernel(loaded, "pan_delay_consume")
self.enqueue_kernel = driver.kernel(loaded, "pan_delay_enqueue")
self.device: dict[str, DeviceAllocation] = {}
self.host: dict[str, Any] = {}
self.consumed_ticks = 0
self.enqueued_events = 0
try:
for name, count, kind in (
("queue", 65 * self.n, ctypes.c_double),
("currents", self.n, ctypes.c_double),
("targets", capacity, ctypes.c_uint32),
("delays", capacity, ctypes.c_uint32),
("amplitudes", capacity, ctypes.c_double),
("invalid", 1, ctypes.c_uint32),
):
self.host[name] = (kind * count)()
self.device[name] = driver.alloc_device(ctypes.sizeof(self.host[name]))
driver.copy_host_to_device(self.device["queue"], self.host["queue"])
except BaseException:
self.close()
raise
def close(self) -> None:
for allocation in reversed(list(self.device.values())):
try:
self.driver.free_device(allocation)
except CudaDriverError:
pass
self.device.clear()
@staticmethod
def _tick(tick: int) -> None:
if type(tick) is not int or not 0 <= tick < 2**32:
raise ValueError("CUDA queue tick must be uint32")
def consume(self, tick: int) -> list[float]:
self._tick(tick)
self.driver.launch_kernel(
self.consume_kernel,
grid=((self.n + 127) // 128, 1, 1),
block=(128, 1, 1),
arguments=[
ctypes.c_uint32(self.n),
ctypes.c_uint32(tick % 65),
ctypes.c_uint64(self.device["queue"].ptr),
ctypes.c_uint64(self.device["currents"].ptr),
],
)
self.driver.synchronize()
self.driver.copy_device_to_host(self.host["currents"], self.device["currents"])
values = list(self.host["currents"])
if any(not math.isfinite(x) for x in values):
raise ValueError("CUDA delay current is non-finite")
self.consumed_ticks += 1
return values
def enqueue(self, tick: int, events: Sequence[tuple[int, int, float]]) -> None:
self._tick(tick)
count = len(events)
if count > self.capacity or any(
type(target) is not int
or not 0 <= target < self.n
or type(delay) is not int
or not 1 <= delay <= 64
or not _is_number(amplitude)
or not math.isfinite(amplitude)
for target, delay, amplitude in events
):
raise ValueError("invalid CUDA delay event")
if not count:
return
self.host["targets"][:count] = [e[0] for e in events]
self.host["delays"][:count] = [e[1] for e in events]
self.host["amplitudes"][:count] = [e[2] for e in events]
self.host["invalid"][0] = 0
for name, size in (
("targets", count * 4),
("delays", count * 4),
("amplitudes", count * 8),
("invalid", 4),
):
self.driver.copy_host_to_device(
self.device[name], self.host[name], size_bytes=size
)
args: list[object] = [
ctypes.c_uint32(self.n),
ctypes.c_uint32(count),
ctypes.c_uint32(tick),
]
args.extend(
ctypes.c_uint64(self.device[name].ptr)
for name in ("targets", "delays", "amplitudes", "queue", "invalid")
)
self.driver.launch_kernel(
self.enqueue_kernel,
grid=((self.n + 127) // 128, 1, 1),
block=(128, 1, 1),
arguments=args,
)
self.driver.synchronize()
self.driver.copy_device_to_host(self.host["invalid"], self.device["invalid"])
if self.host["invalid"][0]:
raise ValueError("CUDA delay accumulation is non-finite")
self.enqueued_events += count
def snapshot(self) -> list[list[float]]:
self.driver.copy_device_to_host(self.host["queue"], self.device["queue"])
values = list(self.host["queue"])
if any(not math.isfinite(x) for x in values):
raise ValueError("CUDA queue snapshot is non-finite")
return [values[i * self.n : (i + 1) * self.n] for i in range(65)]
def restore(self, rows: Sequence[Sequence[float]]) -> None:
if len(rows) != 65 or any(len(row) != self.n for row in rows):
raise ValueError("CUDA queue snapshot shape mismatch")
if any(not _is_number(x) for row in rows for x in row):
raise ValueError("CUDA queue snapshot values must be numbers")
values = [float(x) for row in rows for x in row]
if any(not math.isfinite(x) for x in values):
raise ValueError("CUDA queue snapshot is non-finite")
self.host["queue"][:] = values
self.driver.copy_host_to_device(self.device["queue"], self.host["queue"])
|