metrale_gpu_runtime/cuda_backend/
gpu_impl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `impl GpuBackend for MetraleCudaBackend`: allocation, the
4//! async copies, kernel launch and lookup, and one-line delegations to the
5//! inherent methods in `gpu_copy.rs`, `gpu_impl_graph.rs` and `redzone.rs`.
6//!
7//! Each `unsafe` block is one call: a driver function declared in
8//! `cuda_backend.rs`, `cudaMemcpy2DAsync`, or the `unsafe fn`
9//! `MetraleRegistry::launch_on_stream`.
10//! The caller provides what they need and nothing here checks it: the
11//! backend's context current on the calling thread (`bind_to_thread`),
12//! device pointers into live allocations, and byte counts that fit them.
13//!
14//! Owner: gpu-runtime (CUDA backend).
15//! Invariants:
16//! - `alloc` and `alloc_managed` record every pointer they return `Ok` in the
17//!   allocation ledger; `free` removes it before calling `cuMemFree_v2`.
18//! - `alloc_arena` and `free_arena` never touch the ledger.
19
20use std::ffi::c_void;
21use std::sync::OnceLock;
22
23use super::n_tile::record_published;
24use crate::registry::{RawCudaFunc, cuda_error_text};
25use anyhow::{Result, bail};
26use cudarc::driver::LaunchConfig;
27
28use super::{
29    MetraleCudaBackend, cuMemAlloc_v2, cuMemAllocManaged, cuMemFree_v2, cuMemGetInfo_v2,
30    cuMemcpyDtoDAsync_v2, cuMemcpyDtoHAsync_v2, cuStreamSynchronize,
31};
32use crate::gpu::{DevicePtr, GpuBackend, GraphHandle, KernelHandle};
33
34mod d2h_trace;
35mod h2d;
36mod raw;
37
38use d2h_trace::d2h_trace_tick;
39use h2d::{h2d_enqueue, warn_pinned_transient_source};
40
41impl GpuBackend for MetraleCudaBackend {
42    #[track_caller]
43    fn alloc(&self, bytes: usize) -> Result<DevicePtr> {
44        let site = std::panic::Location::caller();
45        let mut dptr: u64 = 0;
46        // 2026-09-25: Red zone (`METRALE_REDZONE=<bytes>`, unset or 0 = off,
47        // and only for creation indices >= `METRALE_REDZONE_MIN_IDX`):
48        // allocate `pad` extra bytes after the buffer and return the base, so
49        // the caller's pointer is what an unpadded allocation would give. The
50        // pad is filled with `redzone_fill()` here and checked by
51        // `scan_redzones`.
52        let seq = super::ALLOC_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
53        let pad = if seq >= super::redzone_min_idx() {
54            super::redzone_bytes()
55        } else {
56            0
57        };
58        let status = unsafe { cuMemAlloc_v2(&mut dptr, bytes + pad) };
59        if status != 0 {
60            let mut free: usize = 0;
61            let mut total: usize = 0;
62            unsafe { cuMemGetInfo_v2(&mut free, &mut total) };
63            bail!(
64                "cuMemAlloc_v2 failed: status {status}, requested {bytes} bytes \
65                 (device reports {:.1} MB free / {:.1} GB total)",
66                free as f64 / (1024.0 * 1024.0),
67                total as f64 / (1024.0 * 1024.0 * 1024.0),
68            );
69        }
70        if pad > 0 {
71            self.alloc_poison_redzone(dptr, bytes, pad, seq)?;
72        }
73        self.record_alloc(DevicePtr(dptr), bytes, site);
74        if bytes >= 32 * 1024 * 1024 {
75            tracing::debug!(
76                "alloc {:.1} MB (device ptr {dptr:#x})",
77                bytes as f64 / (1024.0 * 1024.0)
78            );
79        }
80        Ok(DevicePtr(dptr))
81    }
82
83    fn scan_redzones(&self) -> Result<usize> {
84        if super::redzone_bytes() == 0 {
85            return Ok(0);
86        }
87        MetraleCudaBackend::scan_redzones(self)
88    }
89
90    fn poison_redzones(&self, lo: usize, hi: usize) -> Result<()> {
91        if super::redzone_bytes() == 0 {
92            return Ok(());
93        }
94        MetraleCudaBackend::poison_redzones(self, lo, hi)
95    }
96
97    #[track_caller]
98    fn alloc_managed(&self, bytes: usize) -> Result<DevicePtr> {
99        let site = std::panic::Location::caller();
100        let mut dptr: u64 = 0;
101        const CU_MEM_ATTACH_GLOBAL: u32 = 0x1;
102        let status = unsafe { cuMemAllocManaged(&mut dptr, bytes, CU_MEM_ATTACH_GLOBAL) };
103        if status != 0 {
104            bail!(
105                "cuMemAllocManaged failed: status {status}, requested {bytes} bytes. \
106                 Check system swap space: swapon --show"
107            );
108        }
109        self.record_alloc(DevicePtr(dptr), bytes, site);
110        Ok(DevicePtr(dptr))
111    }
112
113    fn free(&self, ptr: DevicePtr) -> Result<()> {
114        if ptr.is_null() {
115            return Ok(());
116        }
117        // 2026-09-25: Off the ledger before the free: an entry left behind
118        // would be freed again by `sweep_unreleased`.
119        self.forget_alloc(ptr);
120        if super::redzone_bytes() > 0 {
121            self.forget_redzone(ptr.0);
122        }
123        let status = unsafe { cuMemFree_v2(ptr.0) };
124        // 2026-09-25: A status meaning the context is already gone
125        // (`registry::is_teardown_noop`: 4, 201, 709) is not an error.
126        if status != 0 && !crate::registry::is_teardown_noop(status) {
127            bail!("cuMemFree_v2 failed: status {status}, ptr {ptr}");
128        }
129        Ok(())
130    }
131
132    fn live_bytes(&self) -> Option<usize> {
133        Some(MetraleCudaBackend::live_bytes(self))
134    }
135
136    fn alloc_report(&self, top_n: usize, min_mb: usize) -> Option<String> {
137        Some(MetraleCudaBackend::alloc_report(self, top_n, min_mb))
138    }
139
140    fn chunk_slack_bytes(&self) -> Option<usize> {
141        Some(MetraleCudaBackend::chunk_slack_bytes(self))
142    }
143
144    fn live_allocations(&self) -> Option<Vec<(DevicePtr, usize, String)>> {
145        Some(MetraleCudaBackend::live_allocations(self))
146    }
147
148    fn sweep_unreleased(&self) -> usize {
149        MetraleCudaBackend::sweep_unreleased(self)
150    }
151
152    fn copy_h2d(&self, src: &[u8], dst: DevicePtr) -> Result<()> {
153        MetraleCudaBackend::copy_h2d_impl(self, src, dst)
154    }
155
156    fn copy_d2h(&self, src: DevicePtr, dst: &mut [u8]) -> Result<()> {
157        d2h_trace_tick();
158        MetraleCudaBackend::copy_d2h_impl(self, src, dst)
159    }
160
161    fn copy_d2h_on_stream(&self, src: DevicePtr, dst: &mut [u8], stream: u64) -> Result<()> {
162        d2h_trace_tick();
163        MetraleCudaBackend::copy_d2h_on_stream_impl(self, src, dst, stream)
164    }
165
166    fn copy_d2h_async(&self, src: DevicePtr, dst: &mut [u8], stream: u64) -> Result<()> {
167        // 2026-09-25: No wait here, unlike `copy_d2h` and
168        // `copy_d2h_on_stream`. The caller must synchronise `stream` before
169        // reading `dst`.
170        d2h_trace_tick();
171        let status = unsafe {
172            cuMemcpyDtoHAsync_v2(dst.as_mut_ptr() as *mut c_void, src.0, dst.len(), stream)
173        };
174        if status != 0 {
175            bail!("cuMemcpyDtoHAsync_v2 (async) failed: status {status}");
176        }
177        Ok(())
178    }
179
180    fn copy_d2d(&self, src: DevicePtr, dst: DevicePtr, bytes: usize) -> Result<()> {
181        if metrale_telemetry::launch_trace::on() {
182            metrale_telemetry::launch_trace::record(metrale_telemetry::launch_trace::Entry {
183                kind: "d2d",
184                func: 0,
185                grid: [0, 0, 0],
186                block: [0, 0, 0],
187                smem: 0,
188                args: vec![src.0, dst.0, bytes as u64],
189            });
190        }
191        MetraleCudaBackend::copy_d2d_impl(self, src, dst, bytes)
192    }
193
194    fn launch(
195        &self,
196        func: KernelHandle,
197        grid: [u32; 3],
198        block: [u32; 3],
199        shared_mem: u32,
200        stream: u64,
201        params: &mut [*mut c_void],
202    ) -> Result<()> {
203        let raw_func = RawCudaFunc(func.0 as *mut c_void);
204        let cfg = LaunchConfig {
205            grid_dim: (grid[0], grid[1], grid[2]),
206            block_dim: (block[0], block[1], block[2]),
207            shared_mem_bytes: shared_mem,
208        };
209        let registry = self.registry();
210        let span = crate::timing::kernel_begin(self, func.0, stream);
211        let launched =
212            unsafe { registry.launch_on_stream(raw_func, cfg, stream, params) }.map_err(|e| {
213                // 2026-09-25: Probe the context and latch the fault if it is
214                // gone (`fault_probe`); the error is returned either way.
215                super::fault_probe::note_failure("kernel launch", &e.to_string());
216                anyhow::anyhow!("Kernel launch failed: {e}")
217            });
218        crate::timing::kernel_end(span, stream);
219        launched
220    }
221
222    fn stream_is_capturing(&self, stream: u64) -> bool {
223        // 2026-09-25: Under `metrale_scale` `cuStreamIsCapturing` is not
224        // declared (`cuda_backend.rs`), so report not capturing.
225        #[cfg(metrale_scale)]
226        {
227            let _ = stream;
228            false
229        }
230        #[cfg(not(metrale_scale))]
231        {
232            let mut status: u32 = 0;
233            // 2026-09-25: CU_STREAM_CAPTURE_STATUS_NONE = 0. A failed query
234            // counts as capturing.
235            let rc = unsafe { super::cuStreamIsCapturing(stream, &mut status) };
236            rc != 0 || status != 0
237        }
238    }
239
240    fn synchronize(&self, stream: u64) -> Result<()> {
241        metrale_telemetry::global().stream_sync();
242        let status = unsafe { cuStreamSynchronize(stream) };
243        if status != 0 {
244            bail!("cuStreamSynchronize failed: {}", cuda_error_text(status));
245        }
246        Ok(())
247    }
248
249    fn default_stream(&self) -> u64 {
250        self.default_stream
251    }
252
253    fn op_cache(&self) -> &crate::op_cache::OpCache {
254        &self.op_cache
255    }
256
257    fn debug_sync_kernels(&self) -> bool {
258        MetraleCudaBackend::debug_sync_kernels(self)
259    }
260
261    fn kernel_registry(&self) -> Option<std::sync::Arc<crate::registry::MetraleRegistry>> {
262        Some(self.registry().clone())
263    }
264
265    #[track_caller]
266    fn kernel(&self, module: &str, func_name: &str) -> Result<KernelHandle> {
267        // 2026-09-25: The caller's `file:line`, carried through by
268        // `#[track_caller]` here and on the trait declaration.
269        let site = std::panic::Location::caller();
270        // 2026-09-25: A fresh `OnceLock` per call: nothing is cached across
271        // calls.
272        let cache: OnceLock<RawCudaFunc> = OnceLock::new();
273        let registry = self.registry();
274        match registry.raw_function_cached(&cache, module, func_name) {
275            Ok(raw) => {
276                metrale_telemetry::kernel_audit::record(module, func_name, true, site);
277                metrale_telemetry::launch_trace::name_kernel(raw.0 as u64, module, func_name);
278                let handle = KernelHandle(raw.0 as u64);
279                record_published(registry, &self.op_cache, module, func_name, handle)?;
280                Ok(handle)
281            }
282            Err(e) => {
283                // 2026-09-25: A failed lookup is recorded in the kernel audit
284                // before the error is returned.
285                metrale_telemetry::kernel_audit::record(module, func_name, false, site);
286                Err(anyhow::anyhow!("Kernel lookup {module}::{func_name}: {e}"))
287            }
288        }
289    }
290
291    fn has_module(&self, module: &str) -> bool {
292        self.registry().has_module(module)
293    }
294
295    fn copy_h2d_async(&self, src: &[u8], dst: DevicePtr, stream: u64) -> Result<()> {
296        h2d_enqueue(src, dst, stream)?;
297        // 2026-09-25: The trait lets the caller drop `src` on return. For a
298        // page-locked source (`pinned_hosts::is_pinned`) the copy reads `src`
299        // after the enqueue, so wait for `stream` before returning.
300        if crate::pinned_hosts::is_pinned(src) {
301            warn_pinned_transient_source();
302            let sync = unsafe { cuStreamSynchronize(stream) };
303            if sync != 0 {
304                bail!(
305                    "cuStreamSynchronize after pinned-source H2D failed: {}",
306                    cuda_error_text(sync)
307                );
308            }
309        }
310        Ok(())
311    }
312
313    fn copy_h2d_async_retained(&self, src: &[u8], dst: DevicePtr, stream: u64) -> Result<()> {
314        // 2026-09-25: The caller keeps `src` alive until its next sync on
315        // `stream`, so no wait is added.
316        h2d_enqueue(src, dst, stream)
317    }
318
319    fn copy_d2d_async(
320        &self,
321        src: DevicePtr,
322        dst: DevicePtr,
323        bytes: usize,
324        stream: u64,
325    ) -> Result<()> {
326        let status = unsafe { cuMemcpyDtoDAsync_v2(dst.0, src.0, bytes, stream) };
327        if status != 0 {
328            // 2026-09-25: Backtrace for the same reason as in `copy_d2d_impl`.
329            tracing::error!(
330                "copy_d2d_async failed (status {status}) at:\n{}",
331                std::backtrace::Backtrace::force_capture()
332            );
333            bail!("cuMemcpyDtoDAsync_v2 (copy_d2d_async) failed: status {status}");
334        }
335        Ok(())
336    }
337
338    fn copy_d2d_2d_async(
339        &self,
340        src: DevicePtr,
341        src_pitch: usize,
342        dst: DevicePtr,
343        dst_pitch: usize,
344        width_bytes: usize,
345        height: usize,
346        stream: u64,
347    ) -> Result<()> {
348        raw::memcpy_2d_async(src, src_pitch, dst, dst_pitch, width_bytes, height, stream)
349    }
350
351    fn begin_capture(&self, stream: u64) -> Result<()> {
352        self.begin_capture_cu(stream)
353    }
354    fn end_capture(&self, stream: u64) -> Result<GraphHandle> {
355        self.end_capture_cu(stream)
356    }
357
358    fn abort_capture_if_active(&self, stream: u64) {
359        self.abort_capture_if_active_cu(stream)
360    }
361
362    fn launch_graph(&self, graph: GraphHandle, stream: u64) -> Result<()> {
363        self.launch_graph_cu(graph, stream)
364    }
365    fn destroy_graph(&self, graph: GraphHandle) -> Result<()> {
366        self.destroy_graph_cu(graph)
367    }
368    fn memset(&self, ptr: DevicePtr, value: u8, bytes: usize) -> Result<()> {
369        self.memset_cu(ptr, value, bytes)
370    }
371    fn memset_async(&self, ptr: DevicePtr, value: u8, bytes: usize, stream: u64) -> Result<()> {
372        if metrale_telemetry::launch_trace::on() {
373            metrale_telemetry::launch_trace::record(metrale_telemetry::launch_trace::Entry {
374                kind: "memset",
375                func: 0,
376                grid: [0, 0, 0],
377                block: [0, 0, 0],
378                smem: 0,
379                args: vec![ptr.0, value as u64, bytes as u64],
380            });
381        }
382        self.memset_async_cu(ptr, value, bytes, stream)
383    }
384    fn total_memory(&self) -> Result<usize> {
385        self.total_memory_cu()
386    }
387    fn free_memory(&self) -> Result<usize> {
388        self.free_memory_cu()
389    }
390    fn device_free_memory(&self) -> Result<usize> {
391        self.device_free_memory_cu()
392    }
393    fn live_alloc_count(&self) -> usize {
394        self.live_alloc_len()
395    }
396    fn sm_count(&self) -> Result<u32> {
397        self.sm_count_cu()
398    }
399    fn create_stream(&self) -> Result<u64> {
400        self.create_stream_cu()
401    }
402    fn bind_to_thread(&self) -> Result<()> {
403        self.bind_to_thread_cu()
404    }
405    fn create_event(&self) -> Result<u64> {
406        self.create_event_cu()
407    }
408    fn record_event(&self, event: u64, stream: u64) -> Result<()> {
409        self.record_event_cu(event, stream)
410    }
411    fn stream_wait_event(&self, stream: u64, event: u64) -> Result<()> {
412        self.stream_wait_event_cu(stream, event)
413    }
414    fn event_synchronize(&self, event: u64) -> Result<()> {
415        self.event_synchronize_cu(event)
416    }
417    fn event_query(&self, event: u64) -> Result<bool> {
418        self.event_query_cu(event)
419    }
420    fn destroy_event(&self, event: u64) -> Result<()> {
421        self.destroy_event_cu(event)
422    }
423    fn host_ptr_to_device(&self, host: *mut u8) -> Result<DevicePtr> {
424        let mut dptr: u64 = 0;
425        let status =
426            unsafe { super::cuMemHostGetDevicePointer_v2(&mut dptr, host as *mut c_void, 0) };
427        if status != 0 {
428            bail!("cuMemHostGetDevicePointer_v2 failed: status {status}");
429        }
430        Ok(DevicePtr(dptr))
431    }
432
433    fn alloc_host_pinned(&self, bytes: usize) -> Result<*mut u8> {
434        if bytes >= 32 * 1024 * 1024 {
435            tracing::debug!(
436                "alloc_host_pinned {:.1} MB",
437                bytes as f64 / (1024.0 * 1024.0)
438            );
439        }
440        self.alloc_host_pinned_cu(bytes)
441    }
442    fn free_host_pinned(&self, ptr: *mut u8, _bytes: usize) -> Result<()> {
443        self.free_host_pinned_cu(ptr, _bytes)
444    }
445
446    fn alloc_arena(&self, bytes: usize) -> Result<DevicePtr> {
447        raw::alloc_arena(bytes)
448    }
449    fn free_arena(&self, ptr: DevicePtr) -> Result<()> {
450        raw::free_arena(ptr)
451    }
452}
metrale_gpu_runtime/cuda_backend/gpu_impl.rs

metrale_gpu_runtime/cuda_backend/
gpu_impl.rs
