Expand description
Which host buffers are page-locked, and therefore which copy_h2d_async
calls are genuinely asynchronous.
cuMemcpyHtoDAsync_v2 behaves in two completely different ways depending on
the SOURCE, and nothing in its signature says which one you get:
- Pageable source — the driver copies into its own staging buffer before returning. The call is async with respect to the GPU but synchronous with respect to the caller’s memory, so dropping the source immediately is safe.
- Page-locked source — the DMA engine reads the caller’s pages directly, after the call returns. Dropping or rewriting the source is a use-after-free / torn transfer.
Almost every copy_h2d_async call site in Atlas hands over a stack array or
a local Vec that dies on the next line. Those are sound today purely
because they are pageable. Nothing recorded that dependency, so pinning any
one of those buffers — a normal, desirable optimisation, and one this tree
has already started doing for the SSM spill staging — would have turned a
whole class of call sites unsound at once, with no compile error and no
runtime complaint.
This registry is what makes that fail loudly instead. Every page-locked
allocation Atlas makes is recorded here; the CUDA backend consults it on the
copy_h2d_async path and, for a pinned source, adds the synchronisation the
pageable path was getting from the driver for free. The cost of pinning a
buffer that a call site then drops is a stalled stream and a one-time
warning, not corruption.
Scope and honesty about it: this tracks what goes through
crate::gpu::GpuBackend::alloc_host_pinned, which is the only door to
page-locked memory in this workspace today. Memory page-locked by some other
route — a raw cuMemHostRegister on an existing arena, which
model/ssm_snapshot_spill.rs explicitly contemplates — would not be seen.
Any such call must register here too.
Functions§
- is_
pinned - Does
srclie inside a live page-locked region? - live_
count - Live region count, for tests and teardown assertions.
- register
- Record a page-locked region. Idempotent for a repeated identical base.
- unregister
- Forget a region. Called from
free_host_pinned, so a freed-then-reused address is not misreported as still pinned.