Module pinned_hosts

Module pinned_hosts 

Source
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 src lie 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.