fix(cuda-async): require unsafe kernel argument setup

cuda-async/src/launch.rs

@@ -23,11 +23,12 @@ use std::vec::Vec;
 ///
 /// Arguments are copied into an inline arena by [`push_arg`](Self::push_arg) /
 /// [`push_device_ptr`](Self::push_device_ptr) and handed to the driver as a
-/// `*mut c_void` array at launch; the driver copies the parameter values out
-/// before `cuLaunchKernel` returns, so nothing outlives the launch.
+/// `*mut c_void` array at launch. The driver copies the parameter values out
+/// before `cuLaunchKernel` returns, but device addresses in those values must
+/// remain valid until the kernel has finished.
 #[derive(Debug)]
 pub struct AsyncKernelLaunch {
-    pub func: Arc<Function>,
+    func: Arc<Function>,
     args: KernelArgStorage,
     cfg: Option<LaunchConfig>,
     programmatic_dependent_launch: bool,
@@ -196,7 +197,22 @@ impl KernelArgStorage {
 
 impl AsyncKernelLaunch {
     /// Creates a new kernel launch builder for the given CUDA function.
-    pub fn new(func: Arc<Function>) -> AsyncKernelLaunch {
+    ///
+    /// ```compile_fail,E0133
+    /// use cuda_async::launch::AsyncKernelLaunch;
+    /// use cuda_core::Function;
+    /// use std::sync::Arc;
+    ///
+    /// fn function() -> Arc<Function> { panic!("compile-only example") }
+    /// AsyncKernelLaunch::new(function());
+    /// ```
+    ///
+    /// # Safety
+    /// Before the builder is executed, its complete argument list must match
+    /// the function's parameter ABI. The function must be valid to execute on
+    /// the selected device, and any device allocations it accesses must remain
+    /// valid until the kernel completes.
+    pub unsafe fn new(func: Arc<Function>) -> AsyncKernelLaunch {
         AsyncKernelLaunch {
             func,
             args: KernelArgStorage::default(),
@@ -205,17 +221,84 @@ impl AsyncKernelLaunch {
         }
     }
 
+    /// Returns the CUDA function this builder will launch.
+    pub fn function(&self) -> &Arc<Function> {
+        &self.func
+    }
+
     /// Pushes a kernel argument by value.
+    ///
+    /// ```compile_fail,E0133
+    /// use cuda_async::launch::AsyncKernelLaunch;
+    /// use cuda_core::Function;
+    /// use std::sync::Arc;
+    ///
+    /// fn function() -> Arc<Function> { panic!("compile-only example") }
+    /// let mut launch = unsafe { AsyncKernelLaunch::new(function()) };
+    /// launch.push_arg(0xdead_beef_u64);
+    /// ```
+    ///
+    /// Calling the trait method directly is unsafe as well:
+    ///
+    /// ```compile_fail,E0133
+    /// use cuda_async::launch::{AsyncKernelLaunch, KernelArgument};
+    /// use cuda_core::Function;
+    /// use std::sync::Arc;
+    ///
+    /// fn function() -> Arc<Function> { panic!("compile-only example") }
+    /// let mut launch = unsafe { AsyncKernelLaunch::new(function()) };
+    /// KernelArgument::push_arg(0xdead_beef_u64, &mut launch);
+    /// ```
+    ///
+    /// The function cannot be replaced through this builder:
+    ///
+    /// ```compile_fail,E0616
+    /// use cuda_async::launch::AsyncKernelLaunch;
+    /// use cuda_core::Function;
+    /// use std::sync::Arc;
+    ///
+    /// fn function() -> Arc<Function> { panic!("compile-only example") }
+    /// let mut launch = unsafe { AsyncKernelLaunch::new(function()) };
+    /// launch.func = function();
+    /// ```
+    ///
+    /// # Safety
+    /// The value must match the ABI and meaning of the next formal parameter
+    /// in `func`. If it encodes a device address, the allocation must be valid
+    /// on the launch's device and remain alive until the kernel completes.
     #[inline(always)]
-    pub fn push_arg<T: KernelArgument>(&mut self, arg: T) -> &mut Self {
-        arg.push_arg(self);
+    pub unsafe fn push_arg<T: KernelArgument>(&mut self, arg: T) -> &mut Self {
+        // SAFETY: the caller guarantees that the argument matches `func`.
+        unsafe { arg.push_arg(self) };
         self
     }
 
     /// Pushes a kernel argument from an `Arc` reference.
+    ///
+    /// ```compile_fail,E0133
+    /// use cuda_async::launch::{ArcKernelArgument, AsyncKernelLaunch};
+    /// use cuda_core::Function;
+    /// use std::sync::Arc;
+    ///
+    /// struct Arg;
+    /// impl ArcKernelArgument for Arg {
+    ///     unsafe fn push_arg_arc(self: &Arc<Self>, _: &mut AsyncKernelLaunch) {}
+    /// }
+    /// fn function() -> Arc<Function> { panic!("compile-only example") }
+    /// let mut launch = unsafe { AsyncKernelLaunch::new(function()) };
+    /// launch.push_arg_arc(&Arc::new(Arg));
+    /// ```
+    ///
+    /// # Safety
+    /// The arguments appended by `arg` must match the next formal parameters
+    /// in `func`. Any device allocations they reference must remain valid on
+    /// the launch's device until the kernel completes. This method borrows the
+    /// `Arc`; it does not retain it for the duration of the kernel.
     #[inline(always)]
-    pub fn push_arg_arc<T: ArcKernelArgument>(&mut self, arg: &Arc<T>) -> &mut Self {
-        arg.push_arg_arc(self);
+    pub unsafe fn push_arg_arc<T: ArcKernelArgument>(&mut self, arg: &Arc<T>) -> &mut Self {
+        // SAFETY: the caller guarantees that the arguments match `func` and
+        // their referenced allocations remain valid through kernel completion.
+        unsafe { arg.push_arg_arc(self) };
         self
     }
 
@@ -310,27 +393,46 @@ impl AsyncKernelLaunch {
 /// A kernel argument that can be pushed from an `Arc` reference.
 pub trait ArcKernelArgument {
     // #[inline(always)] Dont think this is necessary. This will be deprecated for required trait methods
-    fn push_arg_arc(self: &Arc<Self>, launcher: &mut AsyncKernelLaunch);
+    /// ```compile_fail,E0133
+    /// use cuda_async::launch::{ArcKernelArgument, AsyncKernelLaunch};
+    /// use cuda_core::Function;
+    /// use std::sync::Arc;
+    ///
+    /// struct Arg;
+    /// impl ArcKernelArgument for Arg {
+    ///     unsafe fn push_arg_arc(self: &Arc<Self>, _: &mut AsyncKernelLaunch) {}
+    /// }
+    /// fn function() -> Arc<Function> { panic!("compile-only example") }
+    /// let mut launch = unsafe { AsyncKernelLaunch::new(function()) };
+    /// <Arg as ArcKernelArgument>::push_arg_arc(&Arc::new(Arg), &mut launch);
+    /// ```
+    ///
+    /// # Safety
+    /// The arguments appended by this method must match the next formal
+    /// parameters of the launch's function, and any referenced device
+    /// allocations must remain valid until the kernel completes.
+    unsafe fn push_arg_arc(self: &Arc<Self>, launcher: &mut AsyncKernelLaunch);
 }
 
 /// A kernel argument that can be pushed by value into an `AsyncKernelLaunch`.
 pub trait KernelArgument {
     // #[inline(always)] Dont think this is necessary. This will be deprecated for required trait methods
-    fn push_arg(self, launcher: &mut AsyncKernelLaunch);
+    ///
+    /// # Safety
+    /// The arguments appended by this method must match the next formal
+    /// parameters of the launch's function, including their ABI and meaning.
+    /// Any device allocations they reference must remain valid until the
+    /// kernel completes.
+    unsafe fn push_arg(self, launcher: &mut AsyncKernelLaunch);
 }
 
-/// Safe implementation for scalar types. Values implementing `DType` are copied
-/// into the kernel's parameter space during launch — the kernel reads the value,
-/// not a device pointer, so no `unsafe` is required.
+/// Implementation for scalar types. `DType` guarantees the host representation
+/// of the value, but the caller must still ensure it matches the kernel ABI.
 impl<T: DType> KernelArgument for T {
-    fn push_arg(self, launcher: &mut AsyncKernelLaunch) {
-        // SAFETY: a `DType` scalar is a plain `Copy` value with the layout the
-        // compiled kernel declares for that scalar type; the launcher's
-        // signature validation (in cutile) checks the position, and the value
-        // is copied out by the driver at launch, so nothing outlives the box.
-        unsafe {
-            launcher.push_arg_raw(self);
-        }
+    unsafe fn push_arg(self, launcher: &mut AsyncKernelLaunch) {
+        // SAFETY: the caller guarantees that this `DType` scalar matches the
+        // kernel's parameter at this position.
+        unsafe { launcher.push_arg_raw(self) };
     }
 }
 

cuda-async/tests/device_fault.rs

@@ -100,7 +100,9 @@ impl DeviceOp for FaultingOp {
                 )));
             }
         }
-        let mut launch = AsyncKernelLaunch::new(self.trap);
+        // SAFETY: `trap` is the known zero-parameter kernel compiled from
+        // TRAP_PTX and is launched on the context's device and stream.
+        let mut launch = unsafe { AsyncKernelLaunch::new(self.trap) };
         launch.set_launch_config(LaunchConfig {
             grid_dim: (1, 1, 1),
             block_dim: (1, 1, 1),

cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs

@@ -66,10 +66,17 @@ impl DeviceOp for ThresholdScaleKernel {
         self,
         ctx: &ExecutionContext,
     ) -> Result<<Self as DeviceOp>::Output, DeviceError> {
-        let mut launcher = AsyncKernelLaunch::new(self.function);
-        launcher.push_arg(self.n);
-        launcher.push_arg(self.threshold);
-        launcher.push_arg(self.scale);
+        // SAFETY: this is the compiled threshold/scale kernel, and the full
+        // parameter list below matches its PTX signature.
+        let mut launcher = unsafe { AsyncKernelLaunch::new(self.function) };
+        // SAFETY: the scalar values and device pointers follow the kernel's
+        // declared parameter order; both tensors remain live through kernel
+        // completion.
+        unsafe {
+            launcher.push_arg(self.n);
+            launcher.push_arg(self.threshold);
+            launcher.push_arg(self.scale);
+        }
         unsafe {
             launcher
                 .push_device_ptr(self.input.device_pointer().cu_deviceptr())

cutile-rs/cutile-book/guide/interoperability.md

@@ -56,14 +56,17 @@ Launch the kernel via `AsyncKernelLaunch`, which is a `DeviceOp` wrapping the CU
 use cuda_async::launch::AsyncKernelLaunch;
 use cuda_core::LaunchConfig;
 
-let mut launcher = AsyncKernelLaunch::new(function.clone());
-launcher.push_arg(num_elements as u32);
-launcher.push_arg(scale);
-// SAFETY: input and output are valid device allocations with at least
-// num_elements f32 elements. output is exclusively written; input is
-// read-only. Both remain allocated until this operation completes.
+// SAFETY: the function is a valid kernel, and the complete parameter list
+// below matches its PTX ABI and uses live device allocations.
+let mut launcher = unsafe { AsyncKernelLaunch::new(function.clone()) };
+// SAFETY: the argument order and types match the PTX signature. Input and
+// output are valid device allocations with at least num_elements f32 elements;
+// output is exclusively written and input is read-only. Both remain allocated
+// until this operation completes.
 unsafe {
     launcher
+        .push_arg(num_elements as u32)
+        .push_arg(scale)
         .push_device_ptr(input.device_pointer().cu_deviceptr())
         .push_device_ptr(output.device_pointer().cu_deviceptr());
 }
@@ -77,16 +80,16 @@ launcher.set_launch_config(LaunchConfig {
 launcher.await?;
 ```
 
-Scalar arguments implementing `DType` use the safe `push_arg` method. Device
-pointers require `unsafe { push_device_ptr() }`. The caller must check that
-allocations are valid on the target GPU, large enough for the kernel, and free
-of conflicting accesses. Argument order and types must match the kernel's
-signature. Rust and the CUDA driver cannot validate these conditions; getting
-them wrong can cause undefined behavior or GPU faults.
+Both `push_arg` and `push_device_ptr` are unsafe for raw kernel launches. The
+caller must check that allocations are valid on the target GPU, large enough
+for the kernel, and free of conflicting accesses. Argument order and types
+must match the kernel's signature. Rust and the CUDA driver cannot validate
+these conditions; getting them wrong can cause undefined behavior or GPU
+faults.
 
 To prevent data races, use stream ordering: operations chained with `.then()` on the same stream execute in order and see each other's writes. Operations on different streams require explicit synchronization.
 
-> **Why generated cuTile kernels don't require `unsafe`:** the `#[cutile::entry]` macro generates launchers that call `push_device_ptr` internally, but they do so safely because the framework controls both sides — device pointers come from framework-managed allocations, and the ownership model (`Partition` for exclusive access, `Arc<Tensor>` for shared reads) prevents aliasing at the type level. Custom kernels bypass this: you are pushing pointers the framework didn't allocate and can't track, so the safety burden falls on you.
+> **Why generated cuTile kernels don't require `unsafe`:** the `#[cutile::entry]` macro generates launchers that push arguments inside `unsafe` blocks. The generated argument sequence follows the compiled kernel signature, and the framework retains tensor storage and checks specialization constraints before enqueueing. Custom kernels bypass those checks, so callers must establish the argument ABI and device-memory validity themselves.
 
 ---
 
@@ -121,13 +124,16 @@ impl DeviceOp for ScaleKernel {
         self,
         ctx: &ExecutionContext,
     ) -> Result<<Self as DeviceOp>::Output, DeviceError> {
-        let mut launcher = AsyncKernelLaunch::new(self.function);
-        launcher.push_arg(self.n);
-        launcher.push_arg(self.scale);
-        // SAFETY: input and output are framework-managed Tensor allocations.
-        // input is shared (Arc, read-only); output is exclusively written.
+        // SAFETY: this wrapper supplies the complete signature declared by
+        // its PTX kernel and retains both buffers through completion.
+        let mut launcher = unsafe { AsyncKernelLaunch::new(self.function) };
+        // SAFETY: these values match the `n` and `scale` parameters in the
+        // PTX kernel, and the pointers below match its input and output params.
+        // Input is shared (Arc, read-only); output is exclusively written.
         unsafe {
             launcher
+                .push_arg(self.n)
+                .push_arg(self.scale)
                 .push_device_ptr(self.input.device_pointer().cu_deviceptr())
                 .push_device_ptr(self.output.device_pointer().cu_deviceptr());
         }

cutile-rs/cutile-book/reference/host-api.md

@@ -303,18 +303,22 @@ LaunchConfig {
 }
 ```
 
-**`AsyncKernelLaunch`** — wraps a CUDA driver kernel launch as a `DeviceOp`. Build the argument list with `push_arg` (safe, for `DType` scalars) or `push_device_ptr` (`unsafe`, for raw device pointers), set the launch config, then `.await` or `.sync_on()`:
+**`AsyncKernelLaunch`** — wraps a CUDA driver kernel launch as a `DeviceOp`. Push arguments with `push_arg` or `push_device_ptr` inside an `unsafe` block, set the launch config, then `.await` or `.sync_on()`:
 
 ```rust
 use cuda_async::launch::AsyncKernelLaunch;
 
-let mut launcher = AsyncKernelLaunch::new(function.clone());
-launcher.push_arg(num_elements as u32);
-launcher.push_arg(scale);
+// SAFETY: the function is a valid kernel and the complete parameter list below
+// matches its PTX ABI; the device allocations remain valid through completion.
+let mut launcher = unsafe { AsyncKernelLaunch::new(function.clone()) };
 let input_ptr = input.device_pointer();
 let output_ptr = output.device_pointer();
+// SAFETY: the values and pointers match the PTX signature, and both buffers
+// are valid on the launch device until the kernel completes.
 unsafe {
     launcher
+        .push_arg(num_elements as u32)
+        .push_arg(scale)
         .push_device_ptr(input_ptr.cu_deviceptr())
         .push_device_ptr(output_ptr.cu_deviceptr());
 }
@@ -326,7 +330,7 @@ launcher.set_launch_config(LaunchConfig {
 launcher.await?;  // Executes as a DeviceOp
 ```
 
-See [Interoperability](../guide/interoperability.md) for the full walkthrough and the wrapper pattern that hides `unsafe` at the call site.
+See [Interoperability](../guide/interoperability.md) for the full walkthrough and the wrapper pattern that keeps `unsafe` inside a typed `DeviceOp`.
 
 **`.generics(Vec<String>)`** — `#[cutile::entry]`-generated launchers accept this method to bind const generics and type parameters at runtime: