plan(ml-alpha): Phase A API addendum — cudarc 0.19 canonical patterns
Plan 1 was written against an older cudarc device-centric API. cudarc 0.19 moved alloc/launch ownership to the stream (partly for CUDA Graph capture hygiene). This addendum pins: - MlDevice -> CudaContext -> CudaStream construction - Cubin load + module + function caching - MappedF32Buffer staging -> DtoD-async -> CudaSlice (canonical CPU->GPU) - Slow-path readback via DtoD into a staging MappedF32Buffer - launch_builder(&func).arg(...).launch(cfg) idiom - IsvBus and MappedPinnedSnapshotSlot/FillSlot using the real MappedF32Buffer API (host_slice_mut, read_all, dev_ptr field) - CUDA Graph A capture via stream.begin_capture / end_capture / instantiate Plan 1 kernel .cu source, CPU oracles, finite-diff thresholds, smoke criteria, and gate logic are unchanged. Only Rust binding code uses these patterns. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,348 @@
|
||||
# Plan 1 API Addendum — cudarc 0.19 + foxhunt canonical patterns
|
||||
|
||||
> Companion to `2026-05-16-ml-alpha-phase-a.md`. When the plan's binding-code
|
||||
> blocks reference `dev.htod_copy`, `dev.alloc_zeros`, `dev.load_cubin_function`,
|
||||
> `func.launch_async`, or `dev.default_stream`, **use the patterns in this
|
||||
> addendum instead.** The plan's kernel `.cu` source, CPU oracles, tests, and
|
||||
> assertion thresholds are unchanged.
|
||||
|
||||
This addendum exists because Plan 1 was written against an older cudarc
|
||||
device-centric API. The vendored cudarc 0.19 (`vendor/cudarc/`) moved
|
||||
allocation, copy, and kernel-launch ownership from the device to the
|
||||
stream — partly to make CUDA Graph capture clean (everything that becomes a
|
||||
graph node has to happen on the captured stream).
|
||||
|
||||
---
|
||||
|
||||
## Imports
|
||||
|
||||
```rust
|
||||
use std::sync::Arc;
|
||||
use cudarc::driver::{
|
||||
CudaContext, CudaStream, CudaModule, CudaFunction, CudaSlice,
|
||||
DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg, DriverError,
|
||||
};
|
||||
use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer;
|
||||
use ml_core::device::MlDevice;
|
||||
```
|
||||
|
||||
`cudarc::driver::result` is the raw-call namespace (used for
|
||||
`memcpy_dtod_async`).
|
||||
|
||||
---
|
||||
|
||||
## Pattern 1 — Device + stream construction
|
||||
|
||||
Production / runtime:
|
||||
|
||||
```rust
|
||||
let dev = MlDevice::cuda(0)?; // panics if no GPU
|
||||
let ctx: &Arc<CudaContext> = dev.cuda_context()?;
|
||||
let stream: &Arc<CudaStream> = dev.cuda_stream()?;
|
||||
```
|
||||
|
||||
Tests:
|
||||
|
||||
```rust
|
||||
fn test_device() -> MlDevice {
|
||||
MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
|
||||
}
|
||||
```
|
||||
|
||||
The plan's references to `MlDevice::default_for_test()` should be replaced
|
||||
with a local `test_device()` helper inside the test module (don't add a
|
||||
public shim to `ml-core` for this).
|
||||
|
||||
---
|
||||
|
||||
## Pattern 2 — Cubin load + function handle
|
||||
|
||||
The build script emits cubins to `OUT_DIR`. At runtime:
|
||||
|
||||
```rust
|
||||
const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin"));
|
||||
|
||||
fn load_fn(ctx: &Arc<CudaContext>, name: &str) -> Result<CudaFunction, DriverError> {
|
||||
let module: Arc<CudaModule> = ctx.load_cubin(CUBIN.to_vec())?;
|
||||
module.load_function(name)
|
||||
}
|
||||
```
|
||||
|
||||
**Cache the `CudaFunction` and the `Arc<CudaModule>`** as struct fields where
|
||||
the kernel is used repeatedly (CfcTrunk, MultiHorizonHeads, BCE, AdamW
|
||||
holders). Don't reload on every forward call — that defeats the whole point
|
||||
of the pre-compiled cubin discipline.
|
||||
|
||||
---
|
||||
|
||||
## Pattern 3 — Host data → mapped-pinned → device (the canonical CPU→GPU path)
|
||||
|
||||
For Phase A bit-equiv tests and the training inner loop, the standard pattern
|
||||
is **mapped-pinned staging → DtoD-async into a `CudaSlice<f32>` → pass the
|
||||
slice to launch**. This keeps kernel signatures uniform (every kernel takes
|
||||
`CudaSlice` args; mapped-pinned is just an upload optimization that is invisible
|
||||
to the kernel).
|
||||
|
||||
```rust
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>, anyhow::Error> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("mapped-pinned alloc: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
|
||||
let mut dst = stream.alloc_zeros::<f32>(n)?;
|
||||
if n > 0 {
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
```
|
||||
|
||||
Use `upload()` everywhere the plan's bindings say `dev.htod_copy(vec)`. Same
|
||||
shape for `i32` etc. via `MappedI32Buffer`.
|
||||
|
||||
**No plain `cuMemcpyHtoD` anywhere** per `feedback_no_htod_htoh_only_mapped_pinned.md`.
|
||||
|
||||
---
|
||||
|
||||
## Pattern 4 — Device → host (slow-path readback only)
|
||||
|
||||
```rust
|
||||
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>, anyhow::Error> {
|
||||
let n = src.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("readback alloc: {e}"))?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (src_ptr, _g) = src.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
Ok(staging.read_all())
|
||||
}
|
||||
```
|
||||
|
||||
Use `download()` only for slow-path test readback and ISV snapshots. Never
|
||||
on the hot path.
|
||||
|
||||
---
|
||||
|
||||
## Pattern 5 — Kernel launch (the builder)
|
||||
|
||||
```rust
|
||||
fn launch<P: PushKernelArg>(stream: &Arc<CudaStream>, func: &CudaFunction, cfg: LaunchConfig, args: impl FnOnce(&mut cudarc::driver::LaunchArgs)) -> Result<(), DriverError> {
|
||||
let mut builder = stream.launch_builder(func);
|
||||
args(&mut builder);
|
||||
unsafe { builder.launch(cfg) }
|
||||
}
|
||||
```
|
||||
|
||||
But idiomatic usage is inline — no need for the helper:
|
||||
|
||||
```rust
|
||||
let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (128,1,1), shared_mem_bytes: 0 };
|
||||
let mut launch = stream.launch_builder(&func);
|
||||
launch
|
||||
.arg(&w_in_d) // CudaSlice<f32>
|
||||
.arg(&w_rec_d)
|
||||
.arg(&b_d)
|
||||
.arg(&x_d)
|
||||
.arg(&h_old_d)
|
||||
.arg(&dt_s) // scalar f32
|
||||
.arg(&(n_in as i32))
|
||||
.arg(&(n_hid as i32))
|
||||
.arg(&mut h_new_d); // mutable arg uses .arg() too — &mut CudaSlice
|
||||
unsafe { launch.launch(cfg)?; }
|
||||
```
|
||||
|
||||
Scalar args (`f32`, `i32`, `u64`) pass by `&value` — they implement
|
||||
`DeviceRepr`. Device slices pass by `&slice` or `&mut slice`.
|
||||
|
||||
The plan's `func.launch_async(cfg, (args...))` should be replaced with the
|
||||
above builder.
|
||||
|
||||
---
|
||||
|
||||
## Pattern 6 — IsvBus (Task 3)
|
||||
|
||||
The plan's IsvBus uses `dev.alloc_zeros` and `dev.default_stream` — both
|
||||
nonexistent. The corrected impl:
|
||||
|
||||
```rust
|
||||
pub struct IsvBus {
|
||||
stream: Arc<CudaStream>,
|
||||
buffer: CudaSlice<f32>,
|
||||
}
|
||||
|
||||
impl IsvBus {
|
||||
pub fn new(dev: &MlDevice) -> Result<Self> {
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let buffer = stream.alloc_zeros::<f32>(SLOTS)?;
|
||||
Ok(Self { stream, buffer })
|
||||
}
|
||||
|
||||
pub fn write_slot(&mut self, idx: usize, value: f32) -> Result<()> {
|
||||
assert!(idx < SLOTS);
|
||||
let staging = unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow!("isv staging: {e}"))?;
|
||||
staging.write_from_slice(&[value]);
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = self.buffer.device_ptr_mut(&self.stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr + (idx * 4) as u64,
|
||||
staging.dev_ptr,
|
||||
4,
|
||||
self.stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
self.stream.synchronize()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> Result<[f32; SLOTS]> {
|
||||
let v = download(&self.stream, &self.buffer)?;
|
||||
let mut out = [0f32; SLOTS];
|
||||
out.copy_from_slice(&v);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub fn buffer(&self) -> &CudaSlice<f32> { &self.buffer }
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Pattern 7 — `MappedPinnedSnapshotSlot` (Task 4)
|
||||
|
||||
The plan invented `as_slice_mut()`, `as_slice()`, `device_ptr()` accessors
|
||||
that don't exist on `MappedF32Buffer`. The real accessors are
|
||||
`host_slice_mut() -> &mut [f32]`, `read_all() -> Vec<f32>`, and the public
|
||||
`dev_ptr: CUdeviceptr` field.
|
||||
|
||||
Corrected `MappedPinnedSnapshotSlot`:
|
||||
|
||||
```rust
|
||||
pub struct MappedPinnedSnapshotSlot { backing: MappedF32Buffer }
|
||||
|
||||
impl MappedPinnedSnapshotSlot {
|
||||
pub fn new() -> Result<Self> {
|
||||
let n = (std::mem::size_of::<SnapshotPayload>() + 3) / 4;
|
||||
let backing = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow!("snapshot slot alloc: {e}"))?;
|
||||
Ok(Self { backing })
|
||||
}
|
||||
|
||||
pub fn write_volatile(&mut self, payload: &SnapshotPayload) {
|
||||
let host = self.backing.host_slice_mut();
|
||||
let dst = host.as_mut_ptr() as *mut SnapshotPayload;
|
||||
unsafe { std::ptr::write_volatile(dst, *payload); }
|
||||
std::sync::atomic::compiler_fence(std::sync::atomic::Ordering::Release);
|
||||
}
|
||||
|
||||
pub fn peek(&self) -> SnapshotPayload {
|
||||
let v = self.backing.read_all();
|
||||
let src = v.as_ptr() as *const SnapshotPayload;
|
||||
unsafe { std::ptr::read_volatile(src) }
|
||||
}
|
||||
|
||||
pub fn device_ptr(&self) -> cudarc::driver::sys::CUdeviceptr { self.backing.dev_ptr }
|
||||
}
|
||||
```
|
||||
|
||||
Tests asserting `slot.device_ptr() != 0` work unchanged (the field is
|
||||
`CUdeviceptr`, a `u64`, so `!= 0` is meaningful).
|
||||
|
||||
`MappedPinnedFillSlot` follows the same shape.
|
||||
|
||||
---
|
||||
|
||||
## Pattern 8 — Kernel argument: passing a `MappedF32Buffer`'s pinned region
|
||||
|
||||
If a kernel must read **directly** from a mapped-pinned slot (hot-path
|
||||
zero-copy), the kernel arg type on the Rust side is `&u64` carrying the
|
||||
`CUdeviceptr`:
|
||||
|
||||
```rust
|
||||
let dev_ptr_arg: u64 = slot.device_ptr();
|
||||
let mut launch = stream.launch_builder(&func);
|
||||
launch.arg(&dev_ptr_arg); // passed as u64 -> CUdeviceptr in the kernel signature
|
||||
unsafe { launch.launch(cfg)?; }
|
||||
```
|
||||
|
||||
The CUDA kernel signature for that arg is `const float* __restrict__` as
|
||||
usual — no kernel-side change.
|
||||
|
||||
For Phase A, **prefer the DtoD-into-CudaSlice pattern** (Pattern 3) — it
|
||||
keeps every kernel binding uniform. Save zero-copy pinned reads for the
|
||||
live hot path (Plan 3 work).
|
||||
|
||||
---
|
||||
|
||||
## Pattern 9 — CUDA Graph A capture (Task 11)
|
||||
|
||||
cudarc 0.19's safe graph API:
|
||||
|
||||
```rust
|
||||
use cudarc::driver::{CudaGraph, CudaGraphExec};
|
||||
use cudarc::driver::sys::CUstreamCaptureMode;
|
||||
|
||||
impl CfcTrunk {
|
||||
pub fn capture_graph_a(&mut self) -> Result<()> {
|
||||
// Warm-up (kernels loaded, allocations stable)
|
||||
self.forward_snapshot(&Mbp10RawInput::default())?;
|
||||
|
||||
let stream = self.stream.clone();
|
||||
unsafe {
|
||||
stream.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
|
||||
self.dispatch_perception_kernels()?;
|
||||
let graph = stream.end_capture()?;
|
||||
let exec = graph.instantiate(cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)?;
|
||||
self.graph_a_exec = Some(exec);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn perception_forward_captured(&mut self, input: &Mbp10RawInput) -> Result<([f32; 5], [f32; 8])> {
|
||||
// Stage scalars + raw input device buffers (outside the graph for v1).
|
||||
self.upload_input(input)?;
|
||||
let exec = self.graph_a_exec.as_ref().expect("Graph A not captured");
|
||||
let stream = &self.stream;
|
||||
unsafe { exec.launch(stream)?; }
|
||||
stream.synchronize()?;
|
||||
Ok((download_array5(&stream, &self.probs_d)?, download_array8(&stream, &self.proj_out_d)?))
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Verify `stream.begin_capture` / `stream.end_capture` / `graph.instantiate` /
|
||||
`exec.launch` signatures from `vendor/cudarc/src/driver/safe/graph.rs` when
|
||||
landing the task — minor `flags` shape may differ.
|
||||
|
||||
---
|
||||
|
||||
## Per-task summary of plan-vs-real changes
|
||||
|
||||
| Plan 1 Task | Change |
|
||||
|---|---|
|
||||
| 3 (IsvBus) | Use Pattern 6 |
|
||||
| 4 (pinned slots) | Use Pattern 7 |
|
||||
| 5 (snap_feature_assemble) | Pattern 1 + 2 + 3 + 5; replace `func.launch_async((args...))` with `launch_builder` |
|
||||
| 6 (cfc_step) | Same as Task 5 |
|
||||
| 7 (multi_horizon_heads) | Same |
|
||||
| 8 (projection) | Same |
|
||||
| 9 (CfcTrunk) | Cache `Arc<CudaStream>`, `CudaFunction`s, and `CudaSlice`s in struct; remove device-method calls |
|
||||
| 10 (BCE + AdamW) | Same |
|
||||
| 11 (Graph A capture) | Use Pattern 9 |
|
||||
| 13 (Phase A trainer) | Inherits CfcTrunk's stream/handles |
|
||||
| 14 (Backprop kernels) | Same as Task 5 for binding; kernels themselves are the real work |
|
||||
|
||||
The plan's kernel `.cu` source, CPU oracles, finite-diff thresholds, smoke
|
||||
criteria, and gate verdict logic are **unchanged**. Only the Rust binding
|
||||
code uses these patterns.
|
||||
Reference in New Issue
Block a user