pxl_rust/
module.rs

1//! Represents a compiled device kernel.
2//!
3//! A `Module` contains one or more device-side functions, and is typically loaded from a file containing compiled code (e.g., `.mubin` format).//!
4//! This module provides a Rust abstraction around the C++ `pxl::Module` class.
5
6use crate::Function;
7use crate::ffi;
8use std::ffi::CString;
9use std::pin::Pin;
10
11/// Represents a compiled device kernel.
12///
13/// This struct wraps a FFI pointer to the C++ `pxl::Module` class.
14/// It provides methods for creating functions.
15pub struct Module {
16    /// Pointer to the underlying FFI `Module` object.
17    inner: *mut ffi::Module,
18}
19
20/// Implementation of the `Module` struct.
21impl Module {
22    /// Creates a new `Module` wrapper from a raw FFI pointer.
23    /// This is called in runtime wrapper APIs requiring an actual `Module` object.
24    /// # Arguments
25    /// * `ptr` - A raw pointer to the underlying FFI `Module` object.
26    /// # Safety
27    /// Caller must ensure that `ptr` is valid and remains valid for the lifetime of `Module`.
28    pub fn new(filename: &str) -> Self {
29        let filename_cstr = CString::new(filename).expect("Failed to create CString");
30        Module {
31            inner: unsafe { ffi::createModule(filename_cstr.as_ptr() as *const i8) },
32        }
33    }
34
35    /// Returns a raw pointer to the underlying FFI `Module` object.
36    /// This is called in runtime wrapper APIs requiring a raw pointer.
37    /// # Safety
38    /// Caller must ensure that FFI object is valid.
39    pub fn get(&self) -> *mut ffi::Module {
40        self.inner
41    }
42
43    /// Creates a new `Function` object.
44    /// # Arguments
45    /// * `name` - Name of the `Function` to be created.
46    /// # Returns
47    /// A new `Function` object.
48    /// # Safety
49    /// Caller must ensure that `name` is valid and corresponds to an actual function in the module.
50    /// # Example
51    /// ```
52    /// let module = pxl::create_module("tests/mu_kernel.mubin");
53    /// let function = module.create_function("sort_with_ptr");
54    /// ```
55    pub fn create_function(&self, name: &str) -> Function {
56        let name_cstr: CString = CString::new(name).expect("Failed to create CString");
57        let function_ptr = unsafe {
58            Pin::new_unchecked(&mut *self.inner).createFunction(name_cstr.as_ptr() as *const i8)
59        };
60        Function::new(function_ptr)
61    }
62}