Skip to main content
Lazy-imports cuda.bindings. When the import succeeds, host-pinned allocations + a direct cudaMemcpy replace numpy’s default pageable d2h: the pageable path forces the driver to stage the transfer through internal pinned scratch before copying into the user buffer (host-DRAM-bandwidth bound), while a pinned destination skips the staging hop and lets d2h reach PCIe line rate. When the import fails, copy_d2h transparently falls back to numpy.array.

​
alloc_pinned_numpy

Allocate a pinned host region and return it as an ndarray view. The returned array’s base chain pins the underlying _PinnedBuf alive until the array and all derived views are dropped; only then does cudaFreeHost run. Parameters:
  • nbytes (int) — Size of the underlying allocation in bytes. Must be at least prod(shape) * dtype.itemsize.
  • dtype — numpy-compatible dtype for the returned view.
  • shape — Shape of the returned view.
Returns:
  • np.ndarray — ndarray of the requested shape and dtype, backed by pinned memory.
Raises:
  • RuntimeError — if cuda.bindings is unavailable, or the underlying cudaHostAlloc fails.

​
copy_d2h

Device-to-host copy, pinned-destination fast path when available. Parameters:
  • src — Single-device contiguous array-like exposing unsafe_buffer_pointer(), nbytes, shape, and dtype. The caller must sync to the source’s stream before invocation (jax.block_until_ready(src) for a jax.Array).
  • dst (np.ndarray | None) — Optional pre-allocated pinned ndarray to write into. Must have at least src.nbytes bytes.
Returns:
  • np.ndarray — ndarray with the same shape and dtype as src.