lamina_mir/
function.rs

1//! Function representation in LUMIR.
2//!
3//! This module defines the structures for representing functions in LUMIR,
4//! including function signatures, parameters, and basic blocks.
5use crate::{
6    block::Block,
7    instruction::Instruction,
8    register::Register,
9    register::VirtualReg,
10    types::{Aggregate, MirType, Signedness},
11    verify::verify_function,
12};
13use std::collections::HashMap;
14use std::fmt;
15
16/// Function parameter
17#[derive(Debug, Clone, PartialEq)]
18pub struct Parameter {
19    /// Virtual register assigned to this parameter
20    pub reg: Register,
21
22    /// Type of the parameter
23    pub ty: MirType,
24
25    /// Signedness of an integer parameter
26    pub signedness: Signedness,
27
28    /// Set when the parameter is an aggregate passed by value. `reg` then names the
29    /// frame slot holding the callee's copy, and `ty` is the pointer to it.
30    pub aggregate: Option<Aggregate>,
31}
32
33impl Parameter {
34    /// Creates a signed parameter that is not an aggregate; chain `with_signedness` or
35    /// `with_aggregate` for others.
36    pub fn new(reg: Register, ty: MirType) -> Self {
37        Self {
38            reg,
39            ty,
40            signedness: Signedness::Signed,
41            aggregate: None,
42        }
43    }
44
45    pub fn with_signedness(mut self, signedness: Signedness) -> Self {
46        self.signedness = signedness;
47        self
48    }
49
50    pub fn with_aggregate(mut self, aggregate: Aggregate) -> Self {
51        self.aggregate = Some(aggregate);
52        self
53    }
54}
55
56/// Launch limits for a GPU kernel. GPU codegen rejects a zero workgroup dimension, a
57/// workgroup larger than 1024 threads or than `max_threads`, and a `max_threads` outside
58/// 1..=1024.
59#[derive(Debug, Clone, Default, PartialEq, Eq)]
60pub struct KernelMetadata {
61    pub required_workgroup_size: Option<[u32; 3]>,
62    pub max_threads: Option<u32>,
63}
64
65/// Function signature
66#[derive(Debug, Clone, PartialEq)]
67pub struct Signature {
68    /// Function name
69    pub name: String,
70
71    /// Parameters (in order)
72    pub params: Vec<Parameter>,
73
74    /// Return type (None for void)
75    pub ret_ty: Option<MirType>,
76
77    /// Signedness of an integer return value
78    pub ret_signedness: Signedness,
79
80    /// Set when the function returns an aggregate by value. `ret_ty` is then the pointer
81    /// to it, and a call's result register names the caller's slot for the copy.
82    pub ret_aggregate: Option<Aggregate>,
83
84    pub kernel: Option<KernelMetadata>,
85
86    /// Set by `@noinline`; the inliner must leave calls to this function alone.
87    pub no_inline: bool,
88}
89
90impl Signature {
91    /// Creates a void, parameterless, non-kernel signature that the inliner may inline.
92    pub fn new(name: impl Into<String>) -> Self {
93        Self {
94            name: name.into(),
95            params: Vec::new(),
96            ret_ty: None,
97            ret_signedness: Signedness::Signed,
98            ret_aggregate: None,
99            kernel: None,
100            no_inline: false,
101        }
102    }
103
104    /// Marks the function as a GPU kernel, which GPU codegen requires to be void with a
105    /// single block.
106    pub fn with_kernel(mut self, metadata: KernelMetadata) -> Self {
107        self.kernel = Some(metadata);
108        self
109    }
110
111    pub fn with_params(mut self, params: Vec<Parameter>) -> Self {
112        self.params = params;
113        self
114    }
115
116    pub fn with_return(mut self, ty: MirType) -> Self {
117        self.ret_ty = Some(ty);
118        self
119    }
120
121    pub fn with_return_signedness(mut self, signedness: Signedness) -> Self {
122        self.ret_signedness = signedness;
123        self
124    }
125
126    pub fn with_return_aggregate(mut self, aggregate: Aggregate) -> Self {
127        self.ret_aggregate = Some(aggregate);
128        self
129    }
130
131    pub fn param_count(&self) -> usize {
132        self.params.len()
133    }
134
135    /// True when there is no return value. An aggregate return is not void, since `ret_ty`
136    /// then holds the pointer.
137    pub fn is_void(&self) -> bool {
138        self.ret_ty.is_none()
139    }
140}
141
142/// LUMIR function
143#[derive(Debug, Clone, PartialEq)]
144pub struct Function {
145    /// Function signature
146    pub sig: Signature,
147
148    /// Basic blocks (ordered map for deterministic iteration)
149    pub blocks: Vec<Block>,
150
151    /// Entry block label
152    pub entry: String,
153
154    /// Extra 8-byte stack slots a virtual register needs beyond its own, for
155    /// stack allocations larger than one slot. The register addresses the
156    /// lowest byte of the reservation.
157    pub stack_reservations: HashMap<VirtualReg, usize>,
158}
159
160impl Function {
161    /// Create a new function with the given signature
162    pub fn new(sig: Signature) -> Self {
163        Self {
164            sig,
165            blocks: Vec::new(),
166            entry: "entry".to_string(),
167            stack_reservations: HashMap::new(),
168        }
169    }
170
171    /// Set the entry block label
172    pub fn with_entry(mut self, entry: impl Into<String>) -> Self {
173        self.entry = entry.into();
174        self
175    }
176
177    /// Add a basic block to this function
178    pub fn add_block(&mut self, block: Block) {
179        self.blocks.push(block);
180    }
181
182    /// Get a basic block by label
183    pub fn get_block(&self, label: &str) -> Option<&Block> {
184        self.blocks.iter().find(|b| b.label == label)
185    }
186
187    /// Get a mutable reference to a basic block by label
188    pub fn get_block_mut(&mut self, label: &str) -> Option<&mut Block> {
189        self.blocks.iter_mut().find(|b| b.label == label)
190    }
191
192    /// Get the entry block
193    pub fn entry_block(&self) -> Option<&Block> {
194        self.get_block(&self.entry)
195    }
196
197    /// Get a mutable reference to the entry block
198    pub fn entry_block_mut(&mut self) -> Option<&mut Block> {
199        let entry = self.entry.clone();
200        self.get_block_mut(&entry)
201    }
202
203    /// Get all block labels
204    pub fn block_labels(&self) -> Vec<&str> {
205        self.blocks.iter().map(|b| b.label.as_str()).collect()
206    }
207
208    /// Total number of instructions across all blocks
209    pub fn instruction_count(&self) -> usize {
210        self.blocks.iter().map(Block::len).sum()
211    }
212
213    /// Check if this function is well-formed
214    pub fn validate(&self) -> Result<(), String> {
215        verify_function(self).map_err(|errors| errors.join("\n"))
216    }
217}
218
219impl fmt::Display for Function {
220    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
221        // Signature header
222        write!(f, "fn {}(", self.sig.name)?;
223        for (i, p) in self.sig.params.iter().enumerate() {
224            if i > 0 {
225                write!(f, ", ")?;
226            }
227            write!(f, "{} {}", p.reg, p.ty)?;
228        }
229        write!(f, ")")?;
230        if let Some(ret) = &self.sig.ret_ty {
231            write!(f, " -> {ret}")?;
232        }
233        writeln!(f, " {{")?;
234
235        // Emit blocks in order
236        for block in &self.blocks {
237            writeln!(f, "{block}")?;
238        }
239
240        write!(f, "}}")
241    }
242}
243
244/// Builder for constructing LUMIR functions
245pub struct FunctionBuilder {
246    function: Function,
247    current_block: Option<String>,
248}
249
250impl FunctionBuilder {
251    /// Create a new function builder
252    pub fn new(name: impl Into<String>) -> Self {
253        Self {
254            function: Function::new(Signature::new(name)),
255            current_block: None,
256        }
257    }
258
259    /// Add a parameter to the function
260    pub fn param(mut self, reg: Register, ty: MirType) -> Self {
261        self.function.sig.params.push(Parameter::new(reg, ty));
262        self
263    }
264
265    /// Set the return type
266    pub fn returns(mut self, ty: MirType) -> Self {
267        self.function.sig.ret_ty = Some(ty);
268        self
269    }
270
271    /// Create a new basic block and make it the current block
272    pub fn block(mut self, label: impl Into<String>) -> Self {
273        let label = label.into();
274        self.function.add_block(Block::new(label.clone()));
275        self.current_block = Some(label);
276        self
277    }
278
279    /// Add an instruction to the current block
280    pub fn instr(mut self, instr: Instruction) -> Self {
281        if let Some(ref label) = self.current_block
282            && let Some(block) = self.function.get_block_mut(label)
283        {
284            block.push(instr);
285        }
286        self
287    }
288
289    /// Build the function
290    pub fn build(self) -> Function {
291        self.function
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298    use crate::{
299        instruction::{Instruction, IntBinOp, Operand},
300        register::VirtualReg,
301        types::ScalarType,
302    };
303
304    #[test]
305    fn test_signature_creation() {
306        let sig = Signature::new("test_func").with_return(MirType::Scalar(ScalarType::I64));
307
308        assert_eq!(sig.name, "test_func");
309        assert!(sig.ret_ty.is_some());
310        assert_eq!(sig.param_count(), 0);
311    }
312
313    #[test]
314    fn test_function_builder() {
315        let func = FunctionBuilder::new("add")
316            .param(
317                Register::Virtual(VirtualReg::gpr(0)),
318                MirType::Scalar(ScalarType::I64),
319            )
320            .param(
321                Register::Virtual(VirtualReg::gpr(1)),
322                MirType::Scalar(ScalarType::I64),
323            )
324            .returns(MirType::Scalar(ScalarType::I64))
325            .block("entry")
326            .instr(Instruction::IntBinary {
327                op: IntBinOp::Add,
328                ty: MirType::Scalar(ScalarType::I64),
329                dst: Register::Virtual(VirtualReg::gpr(2)),
330                lhs: Operand::Register(Register::Virtual(VirtualReg::gpr(0))),
331                rhs: Operand::Register(Register::Virtual(VirtualReg::gpr(1))),
332            })
333            .instr(Instruction::Ret {
334                value: Some(Operand::Register(Register::Virtual(VirtualReg::gpr(2)))),
335            })
336            .build();
337
338        assert_eq!(func.sig.name, "add");
339        assert_eq!(func.sig.param_count(), 2);
340        assert_eq!(func.blocks.len(), 1);
341        assert_eq!(func.instruction_count(), 2);
342    }
343
344    #[test]
345    fn test_function_validation() {
346        let mut func = Function::new(Signature::new("test"));
347
348        // No blocks at all is a declaration, which carries only its signature.
349        assert!(func.validate().is_ok());
350
351        // A block that the entry does not name leaves the entry dangling.
352        let mut other = Block::new("other");
353        other.push(Instruction::Ret { value: None });
354        func.add_block(other);
355        assert!(func.validate().is_err());
356
357        // Add entry block with terminator
358        let mut entry = Block::new("entry");
359        entry.push(Instruction::Ret { value: None });
360        func.add_block(entry);
361
362        assert!(func.validate().is_ok());
363    }
364}
lamina_mir/function.rs

lamina_mir/
function.rs
