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

lamina_mir/
function.rs
