xlog_cuda/provider/
launch_safe.rs

1//! Recorded asynchronous memset paths through the launch recorder.
2//!
3//! Provides two narrow helpers. Public
4//! [`CudaKernelProvider::memset_recorded`] and the public column helper
5//! [`CudaKernelProvider::memset_column_recorded`] use the consuming recorder
6//! enqueue path, so unsupported tracking is rejected during preparation before
7//! the asynchronous memset.
8//! Both helpers record the destination write on a caller-supplied launch stream,
9//! which may differ from the allocation stream.
10
11use cudarc::driver::sys;
12use xlog_core::{Result, XlogError};
13
14use crate::device_runtime::StreamId;
15use crate::launch::{LaunchEnqueueError, LaunchRecorder};
16use crate::memory::{CudaColumn, TrackedCudaSlice};
17
18impl super::CudaKernelProvider {
19    /// Async memset of `value` into every byte of `dst` on
20    /// `launch_stream`, then record the use against the
21    /// runtime.
22    ///
23    /// Requires the provider's `GpuMemoryManager` to carry the runtime attached
24    /// by [`crate::CudaProviderBuilder`] (so `dst.runtime_block()` is `Some`).
25    /// On a legacy/no-runtime manager, returns [`XlogError::Kernel`].
26    ///
27    /// # Errors
28    ///   * `XlogError::Kernel("memset_recorded requires
29    ///     runtime-backed manager")` if the manager has no
30    ///     runtime attached.
31    ///   * `XlogError::Kernel` from preparation, `cuMemsetD8Async`, cleanup,
32    ///     or commit failure.
33    pub fn memset_recorded(
34        &self,
35        dst: &mut TrackedCudaSlice<u8>,
36        value: u8,
37        launch_stream: StreamId,
38    ) -> Result<()> {
39        let runtime = self.memory().runtime().ok_or_else(|| {
40            XlogError::Kernel(
41                "memset_recorded requires a runtime-backed GpuMemoryManager \
42                 (constructed by CudaProviderBuilder)"
43                    .to_string(),
44            )
45        })?;
46        let dst_ptr = dst.device_ptr_value();
47        let dst_len = dst.len();
48
49        let mut rec = LaunchRecorder::new_strict(launch_stream);
50        rec.write(dst);
51
52        // SAFETY: `dst_ptr` remains live for this synchronous enqueue call,
53        // spans exactly `dst_len` bytes because the element type is `u8`, and
54        // was registered above as a write. The closure enqueues exactly one
55        // operation, immediately, and only on the supplied stream.
56        let enqueued = unsafe {
57            rec.enqueue(runtime, |stream| {
58                let result =
59                    sys::cuMemsetD8Async(dst_ptr, value, dst_len, stream.stream().cu_stream());
60                if result == sys::cudaError_enum::CUDA_SUCCESS {
61                    Ok(())
62                } else {
63                    Err(XlogError::Kernel(format!(
64                        "cuMemsetD8Async failed: {result:?}"
65                    )))
66                }
67            })
68        }
69        .map_err(|error| match error {
70            LaunchEnqueueError::Preparation(preparation) => XlogError::Kernel(format!(
71                "memset_recorded: launch recorder preparation failed: {preparation}"
72            )),
73            LaunchEnqueueError::PreparationAndCleanup {
74                preparation,
75                cleanup,
76            } => XlogError::Kernel(format!(
77                "memset_recorded: launch recorder preparation failed: {preparation}; \
78                 launch reservation cleanup also failed: {cleanup}"
79            )),
80            LaunchEnqueueError::Operation(operation) => operation,
81            LaunchEnqueueError::OperationAndCleanup { operation, cleanup } => XlogError::Kernel(
82                format!("{operation}; launch reservation cleanup also failed: {cleanup}"),
83            ),
84        })?;
85
86        enqueued.commit().map_err(|error| {
87            XlogError::Kernel(format!(
88                "memset_recorded: launch recorder commit failed: {error}"
89            ))
90        })
91    }
92
93    /// Column-level variant of [`Self::memset_recorded`] —
94    /// exercises the `LaunchRecorder::write_column` path. Used
95    /// by tests that prove `CudaColumn::Owned` records its
96    /// runtime block automatically; strict mode rejects
97    /// `CudaColumn::Dlpack` / `CudaColumn::ArrowDevice` at
98    /// preflight (no CUDA work queued).
99    pub fn memset_column_recorded(
100        &self,
101        dst: &mut CudaColumn,
102        value: u8,
103        launch_stream: StreamId,
104    ) -> Result<()> {
105        let runtime = self.memory().runtime().ok_or_else(|| {
106            XlogError::Kernel(
107                "memset_column_recorded requires a runtime-backed GpuMemoryManager".to_string(),
108            )
109        })?;
110        let dst_ptr = *dst.device_ptr();
111        let dst_len = <CudaColumn as cudarc::driver::DeviceSlice<u8>>::len(dst);
112
113        let mut rec = LaunchRecorder::new_strict(launch_stream);
114        rec.write_column(dst);
115        // SAFETY: `dst_ptr` remains live for this synchronous enqueue call and
116        // spans exactly `dst_len` bytes. The destination was registered above
117        // as a write, and the closure uses only the supplied stream.
118        let enqueued = unsafe {
119            rec.enqueue(runtime, |stream| {
120                let result =
121                    sys::cuMemsetD8Async(dst_ptr, value, dst_len, stream.stream().cu_stream());
122                if result == sys::cudaError_enum::CUDA_SUCCESS {
123                    Ok(())
124                } else {
125                    Err(XlogError::Kernel(format!(
126                        "cuMemsetD8Async (column) failed: {result:?}"
127                    )))
128                }
129            })
130        }
131        .map_err(|error| match error {
132            LaunchEnqueueError::Preparation(preparation) => XlogError::Kernel(format!(
133                "memset_column_recorded: launch recorder preparation failed: {preparation}"
134            )),
135            LaunchEnqueueError::PreparationAndCleanup {
136                preparation,
137                cleanup,
138            } => XlogError::Kernel(format!(
139                "memset_column_recorded: launch recorder preparation failed: {preparation}; \
140                 launch reservation cleanup also failed: {cleanup}"
141            )),
142            LaunchEnqueueError::Operation(operation) => operation,
143            LaunchEnqueueError::OperationAndCleanup { operation, cleanup } => XlogError::Kernel(
144                format!("{operation}; launch reservation cleanup also failed: {cleanup}"),
145            ),
146        })?;
147
148        enqueued.commit().map_err(|e| {
149            XlogError::Kernel(format!(
150                "memset_column_recorded: launch recorder commit failed: {}",
151                e
152            ))
153        })
154    }
155}
xlog_cuda/provider/launch_safe.rs

xlog_cuda/provider/
launch_safe.rs
