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 lamina_platform::PointerWidth;
6
7use crate::{
8    block::Block,
9    instruction::Instruction,
10    register::Register,
11    register::VirtualReg,
12    types::{Aggregate, MirType, Signedness},
13    verify::verify_function,
14};
15use std::collections::HashMap;
16use std::fmt;
17
18/// Function parameter
19#[derive(Debug, Clone, PartialEq)]
20pub struct Parameter {
21    /// Virtual register assigned to this parameter
22    pub reg: Register,
23
24    /// Type of the parameter
25    pub ty: MirType,
26
27    /// Signedness of an integer parameter
28    pub signedness: Signedness,
29
30    /// Set when the parameter is an aggregate passed by value. `reg` then names the
31    /// frame slot holding the callee's copy, and `ty` is the pointer to it.
32    pub aggregate: Option<Aggregate>,
33}
34
35impl Parameter {
36    /// Creates a signed parameter that is not an aggregate; chain `with_signedness` or
37    /// `with_aggregate` for others.
38    pub fn new(reg: Register, ty: MirType) -> Self {
39        Self {
40            reg,
41            ty,
42            signedness: Signedness::Signed,
43            aggregate: None,
44        }
45    }
46
47    pub fn with_signedness(mut self, signedness: Signedness) -> Self {
48        self.signedness = signedness;
49        self
50    }
51
52    pub fn with_aggregate(mut self, aggregate: Aggregate) -> Self {
53        self.aggregate = Some(aggregate);
54        self
55    }
56}
57
58/// Launch limits for a GPU kernel. GPU codegen rejects a zero workgroup dimension, a
59/// workgroup larger than 1024 threads or than `max_threads`, and a `max_threads` outside
60/// 1..=1024.
61#[derive(Debug, Clone, Default, PartialEq, Eq)]
62pub struct KernelMetadata {
63    pub required_workgroup_size: Option<[u32; 3]>,
64    pub max_threads: Option<u32>,
65}
66
67/// ELF-style symbol visibility requested by a function annotation.
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum SymbolVisibility {
70    Internal,
71    Private,
72    Hidden,
73    Protected,
74}
75
76/// Section a function is placed in, so the linker groups it with other cold or hot code.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum CodePlacement {
79    Cold,
80    Hot,
81}
82
83/// Function signature
84#[derive(Debug, Clone, PartialEq)]
85pub struct Signature {
86    /// Function name
87    pub name: String,
88
89    /// Parameters (in order)
90    pub params: Vec<Parameter>,
91
92    /// Return type (None for void)
93    pub ret_ty: Option<MirType>,
94
95    /// Signedness of an integer return value
96    pub ret_signedness: Signedness,
97
98    /// Set when the function returns an aggregate by value. `ret_ty` is then the pointer
99    /// to it, and a call's result register names the caller's slot for the copy.
100    pub ret_aggregate: Option<Aggregate>,
101
102    pub kernel: Option<KernelMetadata>,
103
104    /// Set by `@noinline`; the inliner must leave calls to this function alone.
105    pub no_inline: bool,
106
107    /// Set by `@inline`; the inliner accepts a larger body than its default size limit.
108    pub inline_hint: bool,
109
110    /// Set by `@noreturn`; no code after a call to this function is reachable.
111    pub noreturn: bool,
112
113    /// Set by `@pure` or `@const`; a call whose result is unused can be deleted.
114    pub pure: bool,
115
116    /// Set by `@internal`, `@private`, `@hidden` or `@protected`.
117    pub visibility: Option<SymbolVisibility>,
118
119    /// Set by `@cold` or `@hot`.
120    pub placement: Option<CodePlacement>,
121
122    /// Set by `@weak`; the symbol is emitted as a weak definition.
123    pub weak: bool,
124
125    /// Set by `@section("name")`; the function is emitted into that ELF section.
126    pub section: Option<String>,
127
128    /// Set by `@align(n)`; the function entry is aligned to `n` bytes.
129    pub align: Option<u32>,
130
131    /// Set by `@naked`; the function is emitted with no prologue or epilogue.
132    pub naked: bool,
133
134    /// Set by a trailing `...` on an `@extern` declaration. A call to one follows the
135    /// variadic side of the platform convention.
136    pub variadic: bool,
137}
138
139impl Signature {
140    /// Creates a void, parameterless, non-kernel signature that the inliner may inline.
141    pub fn new(name: impl Into<String>) -> Self {
142        Self {
143            name: name.into(),
144            params: Vec::new(),
145            ret_ty: None,
146            ret_signedness: Signedness::Signed,
147            ret_aggregate: None,
148            kernel: None,
149            no_inline: false,
150            inline_hint: false,
151            noreturn: false,
152            pure: false,
153            visibility: None,
154            placement: None,
155            weak: false,
156            section: None,
157            align: None,
158            naked: false,
159            variadic: false,
160        }
161    }
162
163    /// Marks the function as a GPU kernel, which GPU codegen requires to be void with a
164    /// single block.
165    pub fn with_kernel(mut self, metadata: KernelMetadata) -> Self {
166        self.kernel = Some(metadata);
167        self
168    }
169
170    pub fn with_params(mut self, params: Vec<Parameter>) -> Self {
171        self.params = params;
172        self
173    }
174
175    pub fn with_return(mut self, ty: MirType) -> Self {
176        self.ret_ty = Some(ty);
177        self
178    }
179
180    pub fn with_return_signedness(mut self, signedness: Signedness) -> Self {
181        self.ret_signedness = signedness;
182        self
183    }
184
185    pub fn with_return_aggregate(mut self, aggregate: Aggregate) -> Self {
186        self.ret_aggregate = Some(aggregate);
187        self
188    }
189
190    pub fn param_count(&self) -> usize {
191        self.params.len()
192    }
193
194    /// True when there is no return value. An aggregate return is not void, since `ret_ty`
195    /// then holds the pointer.
196    pub fn is_void(&self) -> bool {
197        self.ret_ty.is_none()
198    }
199}
200
201/// LUMIR function
202#[derive(Debug, Clone, PartialEq)]
203pub struct Function {
204    /// Function signature
205    pub sig: Signature,
206
207    /// Basic blocks (ordered map for deterministic iteration)
208    pub blocks: Vec<Block>,
209
210    /// Entry block label
211    pub entry: String,
212
213    /// Extra 8-byte stack slots a virtual register needs beyond its own, for
214    /// stack allocations larger than one slot. The register addresses the
215    /// lowest byte of the reservation.
216    pub stack_reservations: HashMap<VirtualReg, usize>,
217
218    /// Required alignment in bytes for a virtual register's stack slot, set by
219    /// `stackalloc ... align n` when n exceeds the 8-byte slot size.
220    pub stack_alignments: HashMap<VirtualReg, usize>,
221
222    /// Width of data pointers on the target this function is lowered for. Sizes of
223    /// pointer-typed values derive from it.
224    pub pointer_width: PointerWidth,
225
226    /// Whether integer registers on the target this function is lowered for are 64 bits.
227    /// Sets the register type of an I/O intrinsic's result, which is an i64 wherever
228    /// registers are 64 bits.
229    pub int64_regs: bool,
230}
231
232impl Function {
233    /// Create a new function with the given signature. Pointers are 64-bit until
234    /// [`Function::with_pointer_width`] says otherwise.
235    pub fn new(sig: Signature) -> Self {
236        Self {
237            sig,
238            blocks: Vec::new(),
239            entry: "entry".to_string(),
240            stack_reservations: HashMap::new(),
241            stack_alignments: HashMap::new(),
242            pointer_width: PointerWidth::Bits64,
243            int64_regs: true,
244        }
245    }
246
247    pub fn with_pointer_width(mut self, pointer_width: PointerWidth) -> Self {
248        self.pointer_width = pointer_width;
249        self
250    }
251
252    pub fn with_int64_regs(mut self, int64_regs: bool) -> Self {
253        self.int64_regs = int64_regs;
254        self
255    }
256
257    /// Set the entry block label
258    pub fn with_entry(mut self, entry: impl Into<String>) -> Self {
259        self.entry = entry.into();
260        self
261    }
262
263    /// Add a basic block to this function
264    pub fn add_block(&mut self, block: Block) {
265        self.blocks.push(block);
266    }
267
268    /// Get a basic block by label
269    pub fn get_block(&self, label: &str) -> Option<&Block> {
270        self.blocks.iter().find(|b| b.label == label)
271    }
272
273    /// Get a mutable reference to a basic block by label
274    pub fn get_block_mut(&mut self, label: &str) -> Option<&mut Block> {
275        self.blocks.iter_mut().find(|b| b.label == label)
276    }
277
278    /// Get the entry block
279    pub fn entry_block(&self) -> Option<&Block> {
280        self.get_block(&self.entry)
281    }
282
283    /// Get a mutable reference to the entry block
284    pub fn entry_block_mut(&mut self) -> Option<&mut Block> {
285        let entry = self.entry.clone();
286        self.get_block_mut(&entry)
287    }
288
289    /// Get all block labels
290    pub fn block_labels(&self) -> Vec<&str> {
291        self.blocks.iter().map(|b| b.label.as_str()).collect()
292    }
293
294    /// Total number of instructions across all blocks
295    pub fn instruction_count(&self) -> usize {
296        self.blocks.iter().map(Block::len).sum()
297    }
298
299    /// Check if this function is well-formed
300    pub fn validate(&self) -> Result<(), String> {
301        verify_function(self).map_err(|errors| errors.join("\n"))
302    }
303}
304
305impl fmt::Display for Function {
306    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
307        // Signature header
308        write!(f, "fn {}(", self.sig.name)?;
309        for (i, p) in self.sig.params.iter().enumerate() {
310            if i > 0 {
311                write!(f, ", ")?;
312            }
313            write!(f, "{} {}", p.reg, p.ty)?;
314        }
315        write!(f, ")")?;
316        if let Some(ret) = &self.sig.ret_ty {
317            write!(f, " -> {ret}")?;
318        }
319        writeln!(f, " {{")?;
320
321        // Emit blocks in order
322        for block in &self.blocks {
323            writeln!(f, "{block}")?;
324        }
325
326        write!(f, "}}")
327    }
328}
329
330/// Builder for constructing LUMIR functions
331pub struct FunctionBuilder {
332    function: Function,
333    current_block: Option<String>,
334}
335
336impl FunctionBuilder {
337    /// Create a new function builder
338    pub fn new(name: impl Into<String>) -> Self {
339        Self {
340            function: Function::new(Signature::new(name)),
341            current_block: None,
342        }
343    }
344
345    /// Add a parameter to the function
346    pub fn param(mut self, reg: Register, ty: MirType) -> Self {
347        self.function.sig.params.push(Parameter::new(reg, ty));
348        self
349    }
350
351    /// Set the return type
352    pub fn returns(mut self, ty: MirType) -> Self {
353        self.function.sig.ret_ty = Some(ty);
354        self
355    }
356
357    /// Create a new basic block and make it the current block
358    pub fn block(mut self, label: impl Into<String>) -> Self {
359        let label = label.into();
360        self.function.add_block(Block::new(label.clone()));
361        self.current_block = Some(label);
362        self
363    }
364
365    /// Add an instruction to the current block
366    pub fn instr(mut self, instr: Instruction) -> Self {
367        if let Some(ref label) = self.current_block
368            && let Some(block) = self.function.get_block_mut(label)
369        {
370            block.push(instr);
371        }
372        self
373    }
374
375    /// Build the function
376    pub fn build(self) -> Function {
377        self.function
378    }
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384    use crate::{
385        instruction::{Instruction, IntBinOp, Operand},
386        register::VirtualReg,
387        types::ScalarType,
388    };
389
390    #[test]
391    fn test_signature_creation() {
392        let sig = Signature::new("test_func").with_return(MirType::Scalar(ScalarType::I64));
393
394        assert_eq!(sig.name, "test_func");
395        assert!(sig.ret_ty.is_some());
396        assert_eq!(sig.param_count(), 0);
397    }
398
399    #[test]
400    fn test_function_builder() {
401        let func = FunctionBuilder::new("add")
402            .param(
403                Register::Virtual(VirtualReg::gpr(0)),
404                MirType::Scalar(ScalarType::I64),
405            )
406            .param(
407                Register::Virtual(VirtualReg::gpr(1)),
408                MirType::Scalar(ScalarType::I64),
409            )
410            .returns(MirType::Scalar(ScalarType::I64))
411            .block("entry")
412            .instr(Instruction::IntBinary {
413                op: IntBinOp::Add,
414                ty: MirType::Scalar(ScalarType::I64),
415                dst: Register::Virtual(VirtualReg::gpr(2)),
416                lhs: Operand::Register(Register::Virtual(VirtualReg::gpr(0))),
417                rhs: Operand::Register(Register::Virtual(VirtualReg::gpr(1))),
418            })
419            .instr(Instruction::Ret {
420                value: Some(Operand::Register(Register::Virtual(VirtualReg::gpr(2)))),
421            })
422            .build();
423
424        assert_eq!(func.sig.name, "add");
425        assert_eq!(func.sig.param_count(), 2);
426        assert_eq!(func.blocks.len(), 1);
427        assert_eq!(func.instruction_count(), 2);
428    }
429
430    #[test]
431    fn test_function_validation() {
432        let mut func = Function::new(Signature::new("test"));
433
434        // No blocks at all is a declaration, which carries only its signature.
435        assert!(func.validate().is_ok());
436
437        // A block that the entry does not name leaves the entry dangling.
438        let mut other = Block::new("other");
439        other.push(Instruction::Ret { value: None });
440        func.add_block(other);
441        assert!(func.validate().is_err());
442
443        // Add entry block with terminator
444        let mut entry = Block::new("entry");
445        entry.push(Instruction::Ret { value: None });
446        func.add_block(entry);
447
448        assert!(func.validate().is_ok());
449    }
450}
lamina_mir/function.rs

lamina_mir/
function.rs
