diff --git a/cuda-async/src/launch.rs b/cuda-async/src/launch.rs index d948cb0a85..2e2860d985 100644 --- a/cuda-async/src/launch.rs +++ b/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, + func: Arc, args: KernelArgStorage, cfg: Option, 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) -> AsyncKernelLaunch { + /// + /// ```compile_fail,E0133 + /// use cuda_async::launch::AsyncKernelLaunch; + /// use cuda_core::Function; + /// use std::sync::Arc; + /// + /// fn function() -> Arc { 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) -> 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 { + &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 { 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 { 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 { 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(&mut self, arg: T) -> &mut Self { - arg.push_arg(self); + pub unsafe fn push_arg(&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, _: &mut AsyncKernelLaunch) {} + /// } + /// fn function() -> Arc { 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(&mut self, arg: &Arc) -> &mut Self { - arg.push_arg_arc(self); + pub unsafe fn push_arg_arc(&mut self, arg: &Arc) -> &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, 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, _: &mut AsyncKernelLaunch) {} + /// } + /// fn function() -> Arc { panic!("compile-only example") } + /// let mut launch = unsafe { AsyncKernelLaunch::new(function()) }; + /// ::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, 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 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) }; } } diff --git a/cuda-async/tests/device_fault.rs b/cuda-async/tests/device_fault.rs index 13e1f20277..56501c52d2 100644 --- a/cuda-async/tests/device_fault.rs +++ b/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), diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs index 6d931332da..71d0e7f832 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs +++ b/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<::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()) diff --git a/cutile-rs/cutile-book/guide/interoperability.md b/cutile-rs/cutile-book/guide/interoperability.md index c804768401..ba4cd09bb7 100644 --- a/cutile-rs/cutile-book/guide/interoperability.md +++ b/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` 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<::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()); } diff --git a/cutile-rs/cutile-book/reference/host-api.md b/cutile-rs/cutile-book/reference/host-api.md index d56cfada89..bbf1ae9ef4 100644 --- a/cutile-rs/cutile-book/reference/host-api.md +++ b/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)`** — `#[cutile::entry]`-generated launchers accept this method to bind const generics and type parameters at runtime: diff --git a/cutile-rs/cutile-examples/examples/async_and_then_example.rs b/cutile-rs/cutile-examples/examples/async_and_then_example.rs index 2560aa1684..6ea226a226 100644 --- a/cutile-rs/cutile-examples/examples/async_and_then_example.rs +++ b/cutile-rs/cutile-examples/examples/async_and_then_example.rs @@ -80,41 +80,53 @@ async fn main() -> Result<(), DeviceError> { // Chain a bunch of kernel launches. let cuda_async_ops = tokio::spawn(async move { let y_part = y.partition(y_partition_shape); - let mut launcher = AsyncKernelLaunch::new(function.clone()); - launcher - .push_arg(&y_part) - .push_arg(a) - .push_arg_arc(&x) - .set_launch_config(LaunchConfig { - grid_dim: y_part.grid().expect("Invalid grid."), - block_dim: (1, 1, 1), - shared_mem_bytes: 0, - }); + // SAFETY: `function` is the compiled SAXPY kernel and the arguments + // below provide its complete parameter list. + let mut launcher = unsafe { AsyncKernelLaunch::new(function.clone()) }; + // SAFETY: the arguments match the compiled SAXPY kernel's + // partitioned output, scalar, and tensor input parameters; their + // storage stays alive until the awaited launch completes. + unsafe { + launcher.push_arg(&y_part).push_arg(a).push_arg_arc(&x); + } + launcher.set_launch_config(LaunchConfig { + grid_dim: y_part.grid().expect("Invalid grid."), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }); launcher .then(|_| { - let mut launcher = AsyncKernelLaunch::new(function.clone()); - launcher - .push_arg(&y_part) - .push_arg(a) - .push_arg_arc(&x) - .set_launch_config(LaunchConfig { - grid_dim: y_part.grid().expect("Invalid grid."), - block_dim: (1, 1, 1), - shared_mem_bytes: 0, - }); + // SAFETY: `function` is the compiled SAXPY kernel and the + // arguments below provide its complete parameter list. + let mut launcher = unsafe { AsyncKernelLaunch::new(function.clone()) }; + // SAFETY: the arguments match the compiled SAXPY kernel's + // partitioned output, scalar, and tensor input parameters; + // their storage stays alive until the awaited launch completes. + unsafe { + launcher.push_arg(&y_part).push_arg(a).push_arg_arc(&x); + } + launcher.set_launch_config(LaunchConfig { + grid_dim: y_part.grid().expect("Invalid grid."), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }); launcher }) .then(|_| { - let mut launcher = AsyncKernelLaunch::new(function.clone()); - launcher - .push_arg(&y_part) - .push_arg(a) - .push_arg_arc(&x) - .set_launch_config(LaunchConfig { - grid_dim: y_part.grid().expect("Invalid grid."), - block_dim: (1, 1, 1), - shared_mem_bytes: 0, - }); + // SAFETY: `function` is the compiled SAXPY kernel and the + // arguments below provide its complete parameter list. + let mut launcher = unsafe { AsyncKernelLaunch::new(function.clone()) }; + // SAFETY: the arguments match the compiled SAXPY kernel's + // partitioned output, scalar, and tensor input parameters; + // their storage stays alive until the awaited launch completes. + unsafe { + launcher.push_arg(&y_part).push_arg(a).push_arg_arc(&x); + } + launcher.set_launch_config(LaunchConfig { + grid_dim: y_part.grid().expect("Invalid grid."), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }); launcher }) .await diff --git a/cutile-rs/cutile-examples/examples/interop.rs b/cutile-rs/cutile-examples/examples/interop.rs index 2f2ebe04b6..3d8d7554e4 100644 --- a/cutile-rs/cutile-examples/examples/interop.rs +++ b/cutile-rs/cutile-examples/examples/interop.rs @@ -126,9 +126,15 @@ impl DeviceOp for ScaleKernel { self, ctx: &ExecutionContext, ) -> Result<::Output, DeviceError> { - let mut launcher = AsyncKernelLaunch::new(self.function); - launcher.push_arg(self.n); - launcher.push_arg(self.scale); + // SAFETY: the PTX function is a kernel for this device, and this + // DeviceOp supplies its complete parameter list below. + let mut launcher = unsafe { AsyncKernelLaunch::new(self.function) }; + // SAFETY: these values follow the `n` and `scale` scalar parameters + // declared by the PTX kernel above. + unsafe { + launcher.push_arg(self.n); + launcher.push_arg(self.scale); + } // Safety: // Input okay since it's behind an Arc. // Output okay since this struct owns the tensor we will be writing to. diff --git a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs index 27dac6ab6c..0dd7e4d9be 100644 --- a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs +++ b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs @@ -590,7 +590,9 @@ pub fn generate_kernel_launcher( Some(format!("vec![{type_name}::DTYPE.as_str().to_string()]")), ); } - builder_statements.push(parse_stmt(format!("kernel_launch.push_arg({var_name});"))); + builder_statements.push(parse_stmt(format!( + "// SAFETY: the generated argument order follows this kernel signature.\nunsafe {{ kernel_launch.push_arg({var_name}); }}" + ))); // For integer scalar params, auto-compute DivHint at launch time. if cutile_compiler::specialization::is_integer_scalar(&type_name) { scalar_hint_exprs.push(format!( @@ -1053,7 +1055,7 @@ pub fn generate_kernel_launcher( )); let mut launch_only_stmts: Vec = vec![parse_stmt( - "let mut kernel_launch = AsyncKernelLaunch::new(function.clone());".to_string(), + "// SAFETY: this generated launcher supplies the complete ABI derived from the compiled Tile function.\nlet mut kernel_launch = unsafe { AsyncKernelLaunch::new(function.clone()) };".to_string(), )]; launch_only_stmts.extend(builder_statements); launch_only_stmts.extend( @@ -1838,7 +1840,7 @@ fn get_tensor_code( "KernelOutputStored::retain(&{var_name}, ctx)?;" ))); builder_statements.push(parse_stmt(format!( - "KernelOutputStored::push_kernel_args(&{var_name}, &mut kernel_launch);" + "// SAFETY: the generated Tile binding and this stored output use the same parameter ABI.\nunsafe {{ KernelOutputStored::push_kernel_args(&{var_name}, &mut kernel_launch); }}" ))); launch_grid_expr_strs.push(format!("KernelOutputStored::grid_bound(&{var_name})?")); syn::parse2::(quote! {{ @@ -1871,7 +1873,7 @@ fn get_tensor_code( "KernelInputStored::retain(&{var_name}, ctx)?;" ))); builder_statements.push(parse_stmt(format!( - "KernelInputStored::push_kernel_args(&{var_name}, &mut kernel_launch);" + "// SAFETY: the generated Tile binding and this stored input use the same parameter ABI.\nunsafe {{ KernelInputStored::push_kernel_args(&{var_name}, &mut kernel_launch); }}" ))); syn::parse2::(quote! {{ diff --git a/cutile-rs/cutile/src/tensor.rs b/cutile-rs/cutile/src/tensor.rs index 68cf53b8ce..2f6b02a5eb 100644 --- a/cutile-rs/cutile/src/tensor.rs +++ b/cutile-rs/cutile/src/tensor.rs @@ -1866,7 +1866,12 @@ pub enum GridBound { pub trait KernelOutputStored: Send { /// Retain storage and acquire exclusive device access before enqueueing. fn retain(&self, ctx: &ExecutionContext) -> Result<(), DeviceError>; - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch); + /// + /// # Safety + /// The launch function must expect the parameter sequence emitted by this + /// binding, and the referenced device storage must remain valid until the + /// kernel completes. + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch); fn grid(&self) -> Result<(u32, u32, u32), Error>; /// This binding's launch-grid constraint. Defaults to exact coverage; /// only bindings that explicitly opted into partial coverage return @@ -1913,21 +1918,23 @@ impl KernelOutputStored for Partition> { }) } - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: callers must provide a Tile function with the generated + // partitioned-output ABI and retain this tensor through completion. unsafe { launcher.push_device_ptr(self.object.cu_deviceptr()); - } - for dim in self.object.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.object.strides.iter() { - launcher.push_arg(*stride); - } - for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); - } - for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + for dim in self.object.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.object.strides.iter() { + launcher.push_arg(*stride); + } + for dim in self.partition_shape.iter() { + launcher.push_arg(*dim as i32); + } + for stride in self.partition_strides.iter() { + launcher.push_arg(*stride as i32); + } } } fn grid(&self) -> Result<(u32, u32, u32), Error> { @@ -1971,21 +1978,23 @@ impl KernelOutputStored for Partition<&mut Tensor> { }) } - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: callers must provide a Tile function with the generated + // partitioned-output ABI and retain this tensor through completion. unsafe { launcher.push_device_ptr(self.object.cu_deviceptr()); - } - for dim in self.object.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.object.strides.iter() { - launcher.push_arg(*stride); - } - for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); - } - for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + for dim in self.object.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.object.strides.iter() { + launcher.push_arg(*stride); + } + for dim in self.partition_shape.iter() { + launcher.push_arg(*dim as i32); + } + for stride in self.partition_strides.iter() { + launcher.push_arg(*stride as i32); + } } } fn grid(&self) -> Result<(u32, u32, u32), Error> { @@ -2020,8 +2029,9 @@ impl KernelOutputStored for MappedLaunchPartition Result<(), DeviceError> { self.partition.retain(ctx) } - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { - self.partition.push_kernel_args(launcher); + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: the wrapped partition emits the declared Tile ABI. + unsafe { self.partition.push_kernel_args(launcher) }; } fn grid(&self) -> Result<(u32, u32, u32), Error> { @@ -2060,8 +2070,9 @@ impl KernelOutputStored for MappedLaunchPartition Result<(), DeviceError> { self.partition.retain(ctx) } - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { - self.partition.push_kernel_args(launcher); + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: the wrapped partition emits the declared Tile ABI. + unsafe { self.partition.push_kernel_args(launcher) }; } fn grid(&self) -> Result<(u32, u32, u32), Error> { @@ -2156,7 +2167,12 @@ impl<'a, T: DType> KernelOutput for MappedLaunchPartition Result<(), DeviceError>; - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch); + /// + /// # Safety + /// The launch function must expect the parameter sequence emitted by this + /// binding, and the referenced device storage must remain valid until the + /// kernel completes. + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch); fn shape(&self) -> &[i32]; fn strides(&self) -> &[i32]; fn spec(&self) -> &SpecializationBits; @@ -2184,15 +2200,17 @@ impl KernelInputStored for Arc> { fn retain(&self, ctx: &ExecutionContext) -> Result<(), DeviceError> { self.storage.retain(ctx, false) } - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: callers must provide a Tile function with the generated + // tensor-input ABI and retain this storage through completion. unsafe { launcher.push_device_ptr(self.cu_deviceptr()); - } - for dim in self.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.strides.iter() { - launcher.push_arg(*stride); + for dim in self.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.strides.iter() { + launcher.push_arg(*stride); + } } } fn shape(&self) -> &[i32] { @@ -2213,15 +2231,17 @@ impl KernelInputStored for &Tensor { fn retain(&self, ctx: &ExecutionContext) -> Result<(), DeviceError> { self.storage.retain(ctx, false) } - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: callers must provide a Tile function with the generated + // tensor-input ABI and retain this storage through completion. unsafe { launcher.push_device_ptr(self.cu_deviceptr()); - } - for dim in self.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.strides.iter() { - launcher.push_arg(*stride); + for dim in self.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.strides.iter() { + launcher.push_arg(*stride); + } } } fn shape(&self) -> &[i32] { @@ -2279,17 +2299,19 @@ impl<'a, T: DType + Sync> KernelInputStored for &'a TensorView<'a, T> { fn retain(&self, ctx: &ExecutionContext) -> Result<(), DeviceError> { self.base.storage.retain(ctx, false) } - fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + unsafe fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { // Push the already-offset device pointer. The offset is applied // host-side so the kernel sees the correct base address directly. + // SAFETY: callers must provide a Tile function with the generated + // tensor-view ABI and retain the base storage through completion. unsafe { launcher.push_device_ptr(self.base.cu_deviceptr() + self.offset_bytes as u64); - } - for dim in self.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.strides.iter() { - launcher.push_arg(*stride); + for dim in self.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.strides.iter() { + launcher.push_arg(*stride); + } } } fn shape(&self) -> &[i32] { diff --git a/cutile-rs/cutile/src/tile_kernel.rs b/cutile-rs/cutile/src/tile_kernel.rs index 503c826e95..131c89486d 100644 --- a/cutile-rs/cutile/src/tile_kernel.rs +++ b/cutile-rs/cutile/src/tile_kernel.rs @@ -1429,16 +1429,17 @@ where /// Pushes the device pointer, shape, and stride information to the kernel launcher /// in the order expected by compiled tile functions. impl ArcKernelArgument for Tensor { - fn push_arg_arc(self: &Arc, launcher: &mut AsyncKernelLaunch) { - // TODO (hme): document safety + unsafe fn push_arg_arc(self: &Arc, launcher: &mut AsyncKernelLaunch) { + // SAFETY: callers must use a Tile function with the generated + // tensor-input ABI and retain this storage through kernel completion. unsafe { launcher.push_device_ptr(self.cu_deviceptr()); - } - for dim in self.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.strides.iter() { - launcher.push_arg(*stride); + for dim in self.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.strides.iter() { + launcher.push_arg(*stride); + } } } } @@ -1449,43 +1450,46 @@ impl ArcKernelArgument for Tensor { /// and strides. This allows kernels to access both the full tensor and the partition /// information for block-level indexing. impl KernelArgument for &Partition> { - fn push_arg(self, launcher: &mut AsyncKernelLaunch) { - // TODO (hme): document safety + unsafe fn push_arg(self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: callers must use a Tile function with the generated + // partitioned-input ABI and retain the tensor through completion. unsafe { launcher.push_device_ptr(self.object.cu_deviceptr()); - } - for dim in self.object.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.object.strides.iter() { - launcher.push_arg(*stride); - } - for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); - } - for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + for dim in self.object.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.object.strides.iter() { + launcher.push_arg(*stride); + } + for dim in self.partition_shape.iter() { + launcher.push_arg(*dim as i32); + } + for stride in self.partition_strides.iter() { + launcher.push_arg(*stride as i32); + } } } } /// Same as above but for borrowed mutable tensor partitions. impl KernelArgument for &Partition<&mut Tensor> { - fn push_arg(self, launcher: &mut AsyncKernelLaunch) { + unsafe fn push_arg(self, launcher: &mut AsyncKernelLaunch) { + // SAFETY: callers must use a Tile function with the generated + // partitioned-output ABI and retain the tensor through completion. unsafe { launcher.push_device_ptr(self.object.cu_deviceptr()); - } - for dim in self.object.shape.iter() { - launcher.push_arg(*dim); - } - for stride in self.object.strides.iter() { - launcher.push_arg(*stride); - } - for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); - } - for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + for dim in self.object.shape.iter() { + launcher.push_arg(*dim); + } + for stride in self.object.strides.iter() { + launcher.push_arg(*stride); + } + for dim in self.partition_shape.iter() { + launcher.push_arg(*dim as i32); + } + for stride in self.partition_strides.iter() { + launcher.push_arg(*stride as i32); + } } } } diff --git a/cutile-rs/cutile/tests/gpu/async_saxpy_unsafe.rs b/cutile-rs/cutile/tests/gpu/async_saxpy_unsafe.rs index c4a20f6ca0..de59c1b2c8 100644 --- a/cutile-rs/cutile/tests/gpu/async_saxpy_unsafe.rs +++ b/cutile-rs/cutile/tests/gpu/async_saxpy_unsafe.rs @@ -98,17 +98,21 @@ fn smoke_async_saxpy_unsafe() { let cuda_async_op = async move { let y_part = y.partition(y_partition_shape); - let mut launcher = AsyncKernelLaunch::new(function.clone()); + // SAFETY: `function` is the compiled SAXPY kernel and the + // arguments below provide its complete parameter list. + let mut launcher = unsafe { AsyncKernelLaunch::new(function.clone()) }; let x_arc: Arc<_> = x.into(); - launcher - .push_arg(&y_part) - .push_arg(a) - .push_arg_arc(&x_arc) - .set_launch_config(LaunchConfig { - grid_dim: y_part.grid().expect("Invalid grid."), - block_dim: (1, 1, 1), - shared_mem_bytes: 0, - }); + // SAFETY: the arguments match the compiled SAXPY kernel's + // partitioned output, scalar, and tensor input parameters; + // their storage stays alive until the awaited launch completes. + unsafe { + launcher.push_arg(&y_part).push_arg(a).push_arg_arc(&x_arc); + } + launcher.set_launch_config(LaunchConfig { + grid_dim: y_part.grid().expect("Invalid grid."), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }); launcher.await.expect("Kernel launch failed."); y_part.unpartition() };