mlx_c_rust/metal/
cuda_kernel.rs

1//! The MLX-C CUDA fast-kernel family.
2//!
3//! The CUDA backend's symbols exist in the linked native image but the
4//! backend reports itself unavailable on Apple Silicon; the wrappers keep the
5//! surface callable and typed so builds that target CUDA-capable hosts need
6//! no bridge work.
7
8use std::ffi::CString;
9
10use crate::error::check_status;
11use crate::raw;
12use crate::{MlxArrayVector, MlxBindingsContext, MlxCError, MlxDtype, MlxStream};
13
14/// Owned MLX-C CUDA kernel configuration handle.
15#[derive(Debug)]
16pub struct MlxCudaKernelConfig(raw::mlx_fast_cuda_kernel_config);
17
18impl MlxCudaKernelConfig {
19    /// Creates an empty CUDA kernel configuration.
20    pub fn new() -> Result<Self, MlxCError> {
21        const OPERATION: &str = "configure an MLX custom CUDA kernel";
22        // SAFETY: MLX returns one owned config handle that enters RAII
23        // ownership immediately.
24        let raw_config = unsafe { raw::mlx_fast_cuda_kernel_config_new() };
25        if raw_config.ctx.is_null() {
26            return Err(MlxCError {
27                operation: OPERATION,
28                description: "MLX returned an empty custom CUDA kernel config".to_owned(),
29            });
30        }
31        Ok(Self(raw_config))
32    }
33
34    /// Adds one typed output with a static shape.
35    pub fn add_output(&mut self, shape: &[i32], dtype: MlxDtype) -> Result<(), MlxCError> {
36        const OPERATION: &str = "configure an MLX custom CUDA kernel output";
37        // SAFETY: The shape slice remains live for this copying call.
38        let status = unsafe {
39            raw::mlx_fast_cuda_kernel_config_add_output_arg(
40                self.raw(),
41                shape.as_ptr(),
42                shape.len(),
43                dtype.to_raw(),
44            )
45        };
46        check_status(status, OPERATION)
47    }
48
49    /// Sets the launch grid dimensions.
50    pub fn set_grid(&mut self, grid: [i32; 3]) -> Result<(), MlxCError> {
51        const OPERATION: &str = "configure an MLX custom CUDA kernel grid";
52        // SAFETY: The config handle is live and the values are plain.
53        let status = unsafe {
54            raw::mlx_fast_cuda_kernel_config_set_grid(self.raw(), grid[0], grid[1], grid[2])
55        };
56        check_status(status, OPERATION)
57    }
58
59    /// Sets the thread-group dimensions.
60    pub fn set_thread_group(&mut self, thread_group: [i32; 3]) -> Result<(), MlxCError> {
61        const OPERATION: &str = "configure an MLX custom CUDA kernel thread group";
62        // SAFETY: The config handle is live and the values are plain.
63        let status = unsafe {
64            raw::mlx_fast_cuda_kernel_config_set_thread_group(
65                self.raw(),
66                thread_group[0],
67                thread_group[1],
68                thread_group[2],
69            )
70        };
71        check_status(status, OPERATION)
72    }
73
74    /// Sets the scalar initialization value for output buffers.
75    pub fn set_init_value(&mut self, initial_value: f32) -> Result<(), MlxCError> {
76        const OPERATION: &str = "initialize an MLX custom CUDA kernel output";
77        // SAFETY: The config handle is live and the value is copied by MLX.
78        let status =
79            unsafe { raw::mlx_fast_cuda_kernel_config_set_init_value(self.raw(), initial_value) };
80        check_status(status, OPERATION)
81    }
82
83    /// Toggles MLX's verbose logging for this kernel launch.
84    pub fn set_verbose(&mut self, verbose: bool) -> Result<(), MlxCError> {
85        const OPERATION: &str = "configure MLX custom CUDA kernel logging";
86        // SAFETY: The config handle is live and the flag is a plain value.
87        let status = unsafe { raw::mlx_fast_cuda_kernel_config_set_verbose(self.raw(), verbose) };
88        check_status(status, OPERATION)
89    }
90
91    /// Adds a dtype template argument.
92    pub fn add_dtype_template_argument(
93        &mut self,
94        name: &str,
95        dtype: MlxDtype,
96    ) -> Result<(), MlxCError> {
97        const OPERATION: &str = "add an MLX custom CUDA kernel dtype template argument";
98        let name_argument = CString::new(name).map_err(|_| MlxCError {
99            operation: OPERATION,
100            description: "template argument name contains an interior null byte".to_owned(),
101        })?;
102        // SAFETY: The config handle is live and the name outlives the call.
103        let status = unsafe {
104            raw::mlx_fast_cuda_kernel_config_add_template_arg_dtype(
105                self.raw(),
106                name_argument.as_ptr(),
107                dtype.to_raw(),
108            )
109        };
110        check_status(status, OPERATION)
111    }
112
113    /// Adds an integer template argument.
114    pub fn add_int_template_argument(
115        &mut self,
116        name: &str,
117        argument: i32,
118    ) -> Result<(), MlxCError> {
119        const OPERATION: &str = "add an MLX custom CUDA kernel integer template argument";
120        let name_argument = CString::new(name).map_err(|_| MlxCError {
121            operation: OPERATION,
122            description: "template argument name contains an interior null byte".to_owned(),
123        })?;
124        // SAFETY: The config handle is live and the name outlives the call.
125        let status = unsafe {
126            raw::mlx_fast_cuda_kernel_config_add_template_arg_int(
127                self.raw(),
128                name_argument.as_ptr(),
129                argument,
130            )
131        };
132        check_status(status, OPERATION)
133    }
134
135    /// Adds a boolean template argument.
136    pub fn add_bool_template_argument(
137        &mut self,
138        name: &str,
139        argument: bool,
140    ) -> Result<(), MlxCError> {
141        const OPERATION: &str = "add an MLX custom CUDA kernel boolean template argument";
142        let name_argument = CString::new(name).map_err(|_| MlxCError {
143            operation: OPERATION,
144            description: "template argument name contains an interior null byte".to_owned(),
145        })?;
146        // SAFETY: The config handle is live and the name outlives the call.
147        let status = unsafe {
148            raw::mlx_fast_cuda_kernel_config_add_template_arg_bool(
149                self.raw(),
150                name_argument.as_ptr(),
151                argument,
152            )
153        };
154        check_status(status, OPERATION)
155    }
156
157    pub(crate) const fn raw(&self) -> raw::mlx_fast_cuda_kernel_config {
158        self.0
159    }
160}
161
162impl Drop for MlxCudaKernelConfig {
163    fn drop(&mut self) {
164        // SAFETY: This owner releases its live config exactly once and never
165        // accesses the handle afterward.
166        unsafe {
167            raw::mlx_fast_cuda_kernel_config_free(self.0);
168        }
169    }
170}
171
172/// Owned MLX-C CUDA kernel handle.
173#[derive(Debug)]
174pub struct MlxCudaKernel(raw::mlx_fast_cuda_kernel);
175
176impl MlxCudaKernel {
177    /// Compiles a CUDA kernel from source with named inputs and outputs.
178    ///
179    /// # Errors
180    /// Returns the captured MLX-C description when compilation fails.
181    pub fn new(
182        kernel_name: &str,
183        input_names: &[&str],
184        output_names: &[&str],
185        kernel_header: &str,
186        kernel_source: &str,
187        ensure_row_contiguous: bool,
188        shared_memory_bytes: i32,
189    ) -> Result<Self, MlxCError> {
190        const OPERATION: &str = "compile an MLX custom CUDA kernel";
191        let kernel_name_argument = CString::new(kernel_name).map_err(|_| MlxCError {
192            operation: OPERATION,
193            description: "kernel name contains an interior null byte".to_owned(),
194        })?;
195        let input_names_vector = crate::types::vectors::MlxVectorString::from_strings(input_names)?;
196        let output_names_vector =
197            crate::types::vectors::MlxVectorString::from_strings(output_names)?;
198        let header_argument = CString::new(kernel_header).map_err(|_| MlxCError {
199            operation: OPERATION,
200            description: "kernel header contains an interior null byte".to_owned(),
201        })?;
202        let source_argument = CString::new(kernel_source).map_err(|_| MlxCError {
203            operation: OPERATION,
204            description: "kernel source contains an interior null byte".to_owned(),
205        })?;
206        // SAFETY: Every C string and vector handle remains live for this
207        // compiling constructor and the returned handle enters RAII
208        // ownership.
209        let raw_kernel = unsafe {
210            raw::mlx_fast_cuda_kernel_new(
211                kernel_name_argument.as_ptr(),
212                input_names_vector.raw(),
213                output_names_vector.raw(),
214                source_argument.as_ptr(),
215                header_argument.as_ptr(),
216                ensure_row_contiguous,
217                shared_memory_bytes,
218            )
219        };
220        if raw_kernel.ctx.is_null() {
221            return Err(MlxCError {
222                operation: OPERATION,
223                description: "MLX returned an empty custom CUDA kernel handle".to_owned(),
224            });
225        }
226        Ok(Self(raw_kernel))
227    }
228
229    /// Applies the kernel to the inputs on the given stream.
230    ///
231    /// # Errors
232    /// Returns the captured MLX-C description when the launch fails.
233    pub fn apply(
234        &self,
235        context: &MlxBindingsContext,
236        inputs: &MlxArrayVector,
237        config: &MlxCudaKernelConfig,
238    ) -> Result<MlxArrayVector, MlxCError> {
239        const OPERATION: &str = "apply an MLX custom CUDA kernel";
240        let mut outputs = MlxArrayVector::empty(OPERATION)?;
241        // SAFETY: All handles are live and the output pointer is valid
242        // writable storage.
243        let status = unsafe {
244            raw::mlx_fast_cuda_kernel_apply(
245                outputs.raw_mut(),
246                self.0,
247                inputs.raw(),
248                config.raw(),
249                context.gpu_stream().raw(),
250            )
251        };
252        check_status(status, OPERATION)?;
253        Ok(outputs)
254    }
255}
256
257impl Drop for MlxCudaKernel {
258    fn drop(&mut self) {
259        // SAFETY: This owner releases its live kernel exactly once and never
260        // accesses the handle afterward.
261        unsafe {
262            raw::mlx_fast_cuda_kernel_free(self.0);
263        }
264    }
265}
266
267impl MlxStream {
268    /// Applies the CUDA kernel with an explicit stream argument.
269    ///
270    /// # Errors
271    /// Returns the captured MLX-C description when the launch fails.
272    pub fn apply_cuda_kernel(
273        &self,
274        kernel: &MlxCudaKernel,
275        inputs: &MlxArrayVector,
276        config: &MlxCudaKernelConfig,
277    ) -> Result<MlxArrayVector, MlxCError> {
278        const OPERATION: &str = "apply an MLX custom CUDA kernel on a stream";
279        let mut outputs = MlxArrayVector::empty(OPERATION)?;
280        // SAFETY: All handles are live and the output pointer is valid
281        // writable storage.
282        let status = unsafe {
283            raw::mlx_fast_cuda_kernel_apply(
284                outputs.raw_mut(),
285                kernel.0,
286                inputs.raw(),
287                config.raw(),
288                self.raw(),
289            )
290        };
291        check_status(status, OPERATION)?;
292        Ok(outputs)
293    }
294}
mlx_c_rust/metal/cuda_kernel.rs

mlx_c_rust/metal/
cuda_kernel.rs
