CUDA-runtime helpers used by the sampler hot path.
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(nbytes: int, dtype, shape) -> np.ndarray
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:
| Name | Type | Description | Default |
nbytes | int | Size of the underlying allocation in bytes. Must be at least prod(shape) * dtype.itemsize. | required |
dtype | | numpy-compatible dtype for the returned view. | required |
shape | | Shape of the returned view. | required |
Returns:
| Type | Description |
ndarray | ndarray of the requested shape and dtype, backed by pinned memory. |
Raises:
| Type | Description |
RuntimeError | if cuda.bindings is unavailable, or the underlying cudaHostAlloc fails. |
Source code in src/tsim/utils/cuda_helpers.py
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 | def alloc_pinned_numpy(nbytes: int, dtype, shape) -> np.ndarray:
"""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.
Args:
nbytes: 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:
ndarray of the requested shape and dtype, backed by pinned memory.
Raises:
RuntimeError: if cuda.bindings is unavailable, or the underlying
``cudaHostAlloc`` fails.
"""
if not _CUDA_BINDINGS_AVAILABLE:
raise RuntimeError(
"cuda.bindings not importable; install 'cuda-bindings' or use "
"copy_d2h() for a transparent fallback."
)
buf = _PinnedBuf(nbytes)
carr = (ctypes.c_uint8 * nbytes).from_address(buf.ptr)
carr._owner = buf # type: ignore[attr-defined] # arr.base = carr; carr._owner = buf → buf stays alive
return np.frombuffer(carr, dtype=np.uint8).view(dtype).reshape(shape)
|
copy_d2h(src, *, dst: ndarray | None = None) -> np.ndarray
Device-to-host copy, pinned-destination fast path when available.
Parameters:
| Name | Type | Description | Default |
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). | required |
dst | ndarray | None | Optional pre-allocated pinned ndarray to write into. Must have at least src.nbytes bytes. | None |
Returns:
| Type | Description |
ndarray | ndarray with the same shape and dtype as src. |
Source code in src/tsim/utils/cuda_helpers.py
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 | def copy_d2h(src, *, dst: np.ndarray | None = None) -> np.ndarray:
"""Device-to-host copy, pinned-destination fast path when available.
Args:
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: Optional pre-allocated pinned ndarray to write into. Must have
at least ``src.nbytes`` bytes.
Returns:
ndarray with the same shape and dtype as ``src``.
"""
if not _CUDA_BINDINGS_AVAILABLE or _src_on_host(src):
# Fallback path. Two cases:
# 1. cuda.bindings isn't importable
# 2. src lives on a CPU device (JAX on a CPU jaxlib build, or any
# array marked CPU) — its ``unsafe_buffer_pointer`` is a host
# pointer, so ``cudaMemcpy(..., DeviceToHost)`` would fail.
# ``np.asarray(jax.Array)`` returns a read-only zero-copy view in
# newer JAX, so allocate fresh and copy in to a writable buffer.
out = np.empty(src.shape, dtype=src.dtype)
out[:] = src
return out
if dst is None:
dst = alloc_pinned_numpy(src.nbytes, src.dtype, src.shape)
err = cudart.cudaMemcpy(
dst.ctypes.data,
src.unsafe_buffer_pointer(),
src.nbytes,
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
)[0]
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"cudaMemcpy d2h failed: {err}")
return dst
|