Moving data to the device
A spike sorter reads hours of recording, window by window. Whatever the kernels do, every window must cross from disk to host memory to the device, so how it crosses sets a floor under the run time. Two rules shape the code:
- A window crosses each way at most once. It is uploaded, every stage (filters, whitening, delays, detection) runs on the device copy, and only results come back: spikes, or the clips template learning needs.
- It crosses in its most compact form. Recordings are usually stored as 16-bit integers with a
gain and an offset per channel. Converting them to
f32on the host doubles the bytes on the bus and spends a host thread on arithmetic the device does for free.
How the Kilosort4 / EMUsort runner does it
Section titled “How the Kilosort4 / EMUsort runner does it”DeviceWindows (in the dspk-sorter-kilosort runner) is the one path every pass uses (preprocessing
fit, clip collection, detection):
- A background thread reads the next window while the device works on the current one
(
dspk_core::WindowLoader, two recycled buffers: host memory stays at two windows whatever the recording length). - At start it asks the source once whether it can return stored values
(
RecordingSource::read_stored; the NWB, raw binary, Zarr and mtscomp readers can). If so, the window is uploaded as raw integer bytes andPipelineWorkspace::process_stored_chunk_in_vramunpacks and scales them on the device (WGSL has no 16-bit integers, so the kernel reads 32-bit words and sign-extends). If not, it uploads scaledf32. - The window is preprocessed in device buffers the workspace keeps (no allocation per window), aligned for channel delays when EMUsort estimated them, and handed to the pass.
The test stored_upload_matches_f32_upload runs Kilosort4 on an int16 recording both ways (the
second through a wrapper that hides stored reads) and requires the same spikes.
Owned buffers, persistent device buffers
Section titled “Owned buffers, persistent device buffers”Two details decide whether an upload costs a transfer or much more (measured in
Kilosort4 detection): CubeCL’s create_from_slice copies through a slow general
path (~0.5 GB/s for a 92 MB window), and a fresh host allocation per window pays its page faults on
the calling thread. So uploads hand over owned bytes (client.create / client.write with
Bytes), the workspace writes a persistent device buffer in place, and the read-ahead thread
allocates each window itself (WindowLoader::stream_owned_while): the main thread copies nothing.
What it saves
Section titled “What it saves”For 384 channels in 60 000-sample windows, an f32 window is 92 MB and an int16 window 46 MB.
Over PCIe 3.0 ×16 (~12 GB/s in practice) that is about 4 ms saved per window, plus the host-side
conversion it no longer does, and half the pinned staging memory. It matters most where the
device work per window is small (detection on a fast GPU), because then the transfer is the run
time.