morok_codegen/
types.rs

1//! Types for code generation.
2
3use morok_dtype::DType;
4
5// Re-export new unified types from device crate
6pub use morok_device::device::{ProgramSpec, Variable};
7pub use morok_dtype::DeviceSpec;
8
9/// A rendered kernel ready for compilation and execution.
10#[derive(Debug, Clone)]
11pub struct RenderedKernel {
12    /// The generated code (LLVM IR, CUDA C, etc.)
13    pub code: String,
14
15    /// Kernel name (used as entry point and for debugging/caching).
16    pub name: String,
17
18    /// Buffer argument information.
19    pub buffer_args: Vec<BufferArg>,
20
21    /// Variable names in order (for populating vars array at runtime).
22    /// Includes thread_id at the end if threading is enabled.
23    pub var_names: Vec<String>,
24
25    /// Global work size (for GPU backends).
26    pub global_size: Option<[usize; 3]>,
27
28    /// Local work size (for GPU backends).
29    pub local_size: Option<[usize; 3]>,
30}
31
32/// Information about a buffer argument to the kernel.
33#[derive(Debug, Clone)]
34pub struct BufferArg {
35    /// Argument index.
36    pub index: usize,
37
38    /// Buffer name.
39    pub name: String,
40
41    /// Data type.
42    pub dtype: DType,
43
44    /// Whether this is an output buffer.
45    pub is_output: bool,
46}
47
48impl RenderedKernel {
49    /// Create a new rendered kernel.
50    pub fn new(code: String, name: String) -> Self {
51        Self { code, name, buffer_args: Vec::new(), var_names: Vec::new(), global_size: None, local_size: None }
52    }
53
54    /// Add a buffer argument.
55    pub fn add_buffer_arg(&mut self, arg: BufferArg) {
56        self.buffer_args.push(arg);
57    }
58
59    /// Set work sizes for GPU execution.
60    pub fn set_work_sizes(&mut self, global: [usize; 3], local: [usize; 3]) {
61        self.global_size = Some(global);
62        self.local_size = Some(local);
63    }
64}
morok_codegen/types.rs

morok_codegen/
types.rs
