sampler: pinned-host d2h via cuda.bindings when available

Default numpy d2h (np.asarray on a jax.Array) lands in pageable host memory,
which forces the CUDA driver to stage the transfer through internal pinned
scratch before copying into the user buffer. The second hop is bound by host
DRAM bandwidth and caps the effective d2h throughput well below PCIe line
rate. Pinning the destination skips the staging hop.

Measured throughput at 3.4 GB transfer (rotated surface code d=7, 10M shots):

                                    H100 (gen4 x16)   B200 (gen5 x16)
  cp.asnumpy / np.asarray (pageable)      1.9 GB/s        4.1 GB/s
  cudaMemcpy → cudaHostAlloc (pinned)    23.4 GB/s       51.4 GB/s

Translated to per-shot wall on the surface-code-noise sweep:

         p           vanilla    pre-pin (+G)   post-pin (+G)
  1e-6  (10M shots)   0.084µs     0.162µs        0.021µs
  1e-4  (10M shots)   0.176µs     0.188µs        0.021µs
  1e-2  (10K shots)   4.86µs      0.102µs        0.091µs

The pinned path lifts +G from "loses to vanilla below p~1e-5" to
"monotonically faster across the whole sweep."

Implementation:
- New tsim.utils.cuda_helpers module: _PinnedBuf (RAII over cudaHostAlloc),
  alloc_pinned_numpy (returns a pinned-backed ndarray with lifetime tied to
  the underlying region via ctypes + ndarray.base), copy_d2h (the public
  entry, picks pinned fast path or numpy fallback based on import of
  cuda.bindings).
- sampler._sample_batches replaces np.asarray(combined)[:shots] with
  copy_d2h(combined)[:shots] and adds the matching jax.block_until_ready
  before the call.

cuda.bindings is a soft dep — when import fails, copy_d2h falls back to
np.array, preserving the prior behavior. The pyproject.toml is unchanged.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: danlkv

src/tsim/sampler.py

@@ -17,6 +17,7 @@
 from tsim.core.graph import prepare_graph
 from tsim.core.types import CompiledComponent, CompiledProgram
 from tsim.noise.channels import ChannelSampler
+from tsim.utils.cuda_helpers import copy_d2h
 
 if TYPE_CHECKING:
     from jax import Array as PRNGKey
@@ -341,7 +342,8 @@ def _sample_batches(
         # output. For big bool tensors (e.g. 500k shots × 528 detector bits)
         # the host memcpy alone was ~1 s on top of the PCIe transfer.
         combined = batches[0] if len(batches) == 1 else jnp.concatenate(batches, axis=0)
-        result = np.asarray(combined)[:shots]
+        jax.block_until_ready(combined)
+        result = copy_d2h(combined)[:shots]
 
         if compute_reference:
             assert reference is not None

src/tsim/utils/cuda_helpers.py

@@ -0,0 +1,114 @@
+"""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``.
+"""
+
+from __future__ import annotations
+
+import ctypes
+
+import numpy as np
+
+try:
+    from cuda.bindings import runtime as cudart
+    _CUDA_BINDINGS_AVAILABLE = True
+except Exception:
+    _CUDA_BINDINGS_AVAILABLE = False
+
+
+class _PinnedBuf:
+    """RAII wrapper for a ``cudaHostAlloc``'d region.
+
+    The Python instance owns the lifetime; ``cudaFreeHost`` runs in
+    ``__del__`` when no references remain (typically when the wrapping
+    numpy view is garbage-collected).
+    """
+
+    __slots__ = ("ptr", "nbytes")
+
+    def __init__(self, nbytes: int):
+        err, ptr = cudart.cudaHostAlloc(nbytes, cudart.cudaHostAllocDefault)
+        if err != cudart.cudaError_t.cudaSuccess:
+            raise RuntimeError(f"cudaHostAlloc({nbytes}) failed: {err}")
+        self.ptr = int(ptr)
+        self.nbytes = nbytes
+
+    def __del__(self):
+        if self.ptr:
+            try:
+                cudart.cudaFreeHost(self.ptr)
+            except Exception:
+                # cudart may be torn down at interpreter exit.
+                pass
+            self.ptr = 0
+
+
+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  # arr.base = carr; carr._owner = buf → buf stays alive
+    return np.frombuffer(carr, dtype=np.uint8).view(dtype).reshape(shape)
+
+
+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:
+        # np.asarray(jax.Array) is a read-only zero-copy view; callers
+        # mutate the return (e.g. XOR detectors with the reference sample)
+        # so allocate fresh and copy.
+        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