lift_core/
model_builder.rs

1// ============================================================================
2// model_builder.rs — High-Level Programmatic Model Builder
3// ============================================================================
4//
5// Provides a fluent API for constructing LIFT IR models from Rust code.
6// Automatically generates valid `.lif` text output and handles all SSA wiring.
7//
8// Usage:
9//   let lif = ModelBuilder::new("gpt2")
10//       .function("forward")
11//           .param("x", tensor(1, 768, DataType::FP32))
12//           .param("w", tensor_2d(768, 768, DataType::FP32))
13//           .op("tensor.matmul", &["x", "w"], "h", tensor(1, 768, DataType::FP32))
14//           .op("tensor.relu", &["h"], "out", tensor(1, 768, DataType::FP32))
15//           .returns("out")
16//           .done()
17//       .build_lif();
18//
19// ============================================================================
20
21use crate::attributes::{Attribute, Attributes};
22use crate::context::Context;
23use crate::functions::FunctionData;
24use crate::location::Location;
25use crate::types::{Dimension, MemoryLayout};
26
27/// Re-exported so users can write
28/// `use lift_core::model_builder::{ModelBuilder, DataType, tensor}`.
29pub use crate::types::DataType;
30
31/// Shape descriptor for tensor types.
32#[derive(Debug, Clone)]
33pub struct Shape(pub Vec<usize>);
34
35impl Shape {
36    pub fn new(dims: &[usize]) -> Self {
37        Self(dims.to_vec())
38    }
39}
40
41/// Describes a value type in the model.
42#[derive(Debug, Clone)]
43pub enum ModelType {
44    Tensor { shape: Vec<usize>, dtype: DataType },
45    Qubit,
46    Bit,
47    Integer { bits: u32 },
48}
49
50/// Convenience: create a tensor type descriptor.
51pub fn tensor(shape: &[usize], dtype: DataType) -> ModelType {
52    ModelType::Tensor {
53        shape: shape.to_vec(),
54        dtype,
55    }
56}
57
58/// Convenience: create a 2D tensor type descriptor.
59pub fn tensor_2d(m: usize, n: usize, dtype: DataType) -> ModelType {
60    ModelType::Tensor {
61        shape: vec![m, n],
62        dtype,
63    }
64}
65
66/// Convenience: create a 3D tensor type descriptor.
67pub fn tensor_3d(b: usize, s: usize, d: usize, dtype: DataType) -> ModelType {
68    ModelType::Tensor {
69        shape: vec![b, s, d],
70        dtype,
71    }
72}
73
74/// Convenience: create a 4D tensor type descriptor.
75pub fn tensor_4d(b: usize, c: usize, h: usize, w: usize, dtype: DataType) -> ModelType {
76    ModelType::Tensor {
77        shape: vec![b, c, h, w],
78        dtype,
79    }
80}
81
82/// Convenience: create a 1D tensor type descriptor.
83pub fn tensor_1d(n: usize, dtype: DataType) -> ModelType {
84    ModelType::Tensor {
85        shape: vec![n],
86        dtype,
87    }
88}
89
90/// Represents a single operation in the builder.
91#[derive(Debug, Clone)]
92struct OpDef {
93    op_name: String,
94    inputs: Vec<String>,
95    result_name: String,
96    result_type: ModelType,
97    attrs: Vec<(String, Attribute)>,
98}
99
100/// Represents a function being built.
101#[derive(Debug)]
102pub struct FunctionBuilder {
103    name: String,
104    params: Vec<(String, ModelType)>,
105    ops: Vec<OpDef>,
106    return_name: Option<String>,
107    dialect: String,
108}
109
110impl FunctionBuilder {
111    fn new(name: &str) -> Self {
112        Self {
113            name: name.to_string(),
114            params: Vec::new(),
115            ops: Vec::new(),
116            return_name: None,
117            dialect: "tensor".to_string(),
118        }
119    }
120
121    /// Add a parameter to the function.
122    pub fn param(mut self, name: &str, ty: ModelType) -> Self {
123        self.params.push((name.to_string(), ty));
124        self
125    }
126
127    /// Add an operation.
128    pub fn op(
129        mut self,
130        op_name: &str,
131        inputs: &[&str],
132        result: &str,
133        result_type: ModelType,
134    ) -> Self {
135        self.ops.push(OpDef {
136            op_name: op_name.to_string(),
137            inputs: inputs.iter().map(|s| s.to_string()).collect(),
138            result_name: result.to_string(),
139            result_type,
140            attrs: Vec::new(),
141        });
142        self
143    }
144
145    /// Add an operation with attributes.
146    pub fn op_with_attrs(
147        mut self,
148        op_name: &str,
149        inputs: &[&str],
150        result: &str,
151        result_type: ModelType,
152        attrs: Vec<(&str, Attribute)>,
153    ) -> Self {
154        self.ops.push(OpDef {
155            op_name: op_name.to_string(),
156            inputs: inputs.iter().map(|s| s.to_string()).collect(),
157            result_name: result.to_string(),
158            result_type,
159            attrs: attrs.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
160        });
161        self
162    }
163
164    /// Set the return value.
165    pub fn returns(mut self, name: &str) -> Self {
166        self.return_name = Some(name.to_string());
167        self
168    }
169
170    /// Set the dialect prefix.
171    pub fn dialect(mut self, d: &str) -> Self {
172        self.dialect = d.to_string();
173        self
174    }
175}
176
177/// Main model builder — constructs a complete LIFT module.
178#[derive(Debug)]
179pub struct ModelBuilder {
180    module_name: String,
181    functions: Vec<FunctionBuilder>,
182    dialect_directive: String,
183}
184
185impl ModelBuilder {
186    /// Create a new model builder with a module name.
187    pub fn new(name: &str) -> Self {
188        Self {
189            module_name: name.to_string(),
190            functions: Vec::new(),
191            dialect_directive: "tensor".to_string(),
192        }
193    }
194
195    /// Set the dialect directive.
196    pub fn dialect(mut self, d: &str) -> Self {
197        self.dialect_directive = d.to_string();
198        self
199    }
200
201    /// Add a function definition.
202    pub fn function(mut self, name: &str) -> FunctionBuilderHandle {
203        let fb = FunctionBuilder::new(name);
204        let idx = self.functions.len();
205        self.functions.push(fb);
206        FunctionBuilderHandle {
207            model: self,
208            func_idx: idx,
209        }
210    }
211
212    /// Build the model into a LIFT Context and return it.
213    pub fn build_context(&self) -> Context {
214        let mut ctx = Context::new();
215        let mod_idx = ctx.create_module(&self.module_name);
216
217        for fb in &self.functions {
218            let func_data = self.build_function_data(&mut ctx, fb);
219            ctx.add_function_to_module(mod_idx, func_data);
220        }
221
222        ctx
223    }
224
225    /// Build the model and return the `.lif` source text (parseable format).
226    pub fn build_lif(&self) -> String {
227        let mut out = format!("#dialect {}\n\n", self.dialect_directive);
228        out.push_str(&format!("module @{} {{\n\n", self.module_name));
229
230        for fb in &self.functions {
231            out.push_str(&self.emit_function_source(fb));
232            out.push('\n');
233        }
234
235        out.push_str("}\n");
236        out
237    }
238
239    /// Build and write `.lif` file to disk.
240    pub fn write_lif(&self, path: &str) -> std::io::Result<()> {
241        let lif = self.build_lif();
242        std::fs::write(path, lif)
243    }
244
245    fn emit_function_source(&self, fb: &FunctionBuilder) -> String {
246        let mut out = String::new();
247
248        // func @name(%p0: type, %p1: type) -> ret_type {
249        out.push_str(&format!("    func @{}(", fb.name));
250        for (i, (pname, pty)) in fb.params.iter().enumerate() {
251            if i > 0 {
252                out.push_str(", ");
253            }
254            out.push_str(&format!("%{}: {}", pname, format_model_type(pty)));
255        }
256        out.push(')');
257
258        // Return type
259        if let Some(ref ret_name) = fb.return_name {
260            if let Some(op) = fb.ops.iter().find(|o| o.result_name == *ret_name) {
261                out.push_str(&format!(" -> {}", format_model_type(&op.result_type)));
262            } else if let Some((_, ty)) = fb.params.iter().find(|(n, _)| n == ret_name) {
263                out.push_str(&format!(" -> {}", format_model_type(ty)));
264            }
265        }
266
267        out.push_str(" {\n");
268
269        // Operations
270        for op_def in &fb.ops {
271            let inputs: String = op_def
272                .inputs
273                .iter()
274                .map(|n| format!("%{}", n))
275                .collect::<Vec<_>>()
276                .join(", ");
277
278            let input_types: String = op_def
279                .inputs
280                .iter()
281                .filter_map(|n| {
282                    // Look up type: first in params, then in previous op results
283                    if let Some((_, ty)) = fb.params.iter().find(|(pn, _)| pn == n) {
284                        Some(format_model_type(ty))
285                    } else {
286                        fb.ops
287                            .iter()
288                            .find(|o| o.result_name == *n)
289                            .map(|prev_op| format_model_type(&prev_op.result_type))
290                    }
291                })
292                .collect::<Vec<_>>()
293                .join(", ");
294
295            let result_type = format_model_type(&op_def.result_type);
296
297            out.push_str(&format!(
298                "        %{} = \"{}\"({}) : ({}) -> {}\n",
299                op_def.result_name, op_def.op_name, inputs, input_types, result_type
300            ));
301        }
302
303        // Return
304        if let Some(ref ret_name) = fb.return_name {
305            out.push_str(&format!("        return %{}\n", ret_name));
306        }
307
308        out.push_str("    }\n");
309        out
310    }
311
312    fn build_function_data(&self, ctx: &mut Context, fb: &FunctionBuilder) -> FunctionData {
313        let name_id = ctx.intern_string(&fb.name);
314
315        // Resolve parameter types
316        let param_types: Vec<_> = fb
317            .params
318            .iter()
319            .map(|(_, ty)| self.resolve_model_type(ctx, ty))
320            .collect();
321
322        // Determine return type
323        let return_types = if let Some(ref ret_name) = fb.return_name {
324            // Find the op that produces the return value
325            if let Some(op) = fb.ops.iter().find(|o| o.result_name == *ret_name) {
326                vec![self.resolve_model_type(ctx, &op.result_type)]
327            } else if let Some((_, ty)) = fb.params.iter().find(|(n, _)| n == ret_name) {
328                vec![self.resolve_model_type(ctx, ty)]
329            } else {
330                vec![]
331            }
332        } else {
333            vec![]
334        };
335
336        let mut func_data = FunctionData::new(name_id, param_types.clone(), return_types);
337
338        // Create function body
339        let region = ctx.create_region();
340        let block = ctx.create_block();
341        ctx.add_block_to_region(region, block);
342
343        // Create block args for parameters and build name map
344        let mut name_map = std::collections::HashMap::new();
345        for (i, (pname, _)) in fb.params.iter().enumerate() {
346            let val = ctx.create_block_arg(block, param_types[i]);
347            name_map.insert(pname.clone(), val);
348        }
349
350        // Build operations
351        for op_def in &fb.ops {
352            let inputs: Vec<_> = op_def
353                .inputs
354                .iter()
355                .filter_map(|name| name_map.get(name).copied())
356                .collect();
357
358            let result_ty = self.resolve_model_type(ctx, &op_def.result_type);
359
360            let (dialect, _) = op_def
361                .op_name
362                .split_once('.')
363                .unwrap_or(("core", &op_def.op_name));
364
365            let mut attrs = Attributes::new();
366            for (key, val) in &op_def.attrs {
367                attrs.set(key.clone(), val.clone());
368            }
369
370            let (op_key, results) = ctx.create_op(
371                &op_def.op_name,
372                dialect,
373                inputs,
374                vec![result_ty],
375                attrs,
376                Location::unknown(),
377            );
378            ctx.add_op_to_block(block, op_key);
379
380            if !results.is_empty() {
381                name_map.insert(op_def.result_name.clone(), results[0]);
382            }
383        }
384
385        // Add return operation
386        if let Some(ref ret_name) = fb.return_name {
387            if let Some(&ret_val) = name_map.get(ret_name) {
388                let (ret_op, _) = ctx.create_op(
389                    "core.return",
390                    "core",
391                    vec![ret_val],
392                    vec![],
393                    Attributes::new(),
394                    Location::unknown(),
395                );
396                ctx.add_op_to_block(block, ret_op);
397            }
398        }
399
400        func_data.body = Some(region);
401        func_data
402    }
403
404    fn resolve_model_type(&self, ctx: &mut Context, mt: &ModelType) -> crate::types::TypeId {
405        match mt {
406            ModelType::Tensor { shape, dtype } => {
407                let dims: Vec<Dimension> = shape.iter().map(|&d| Dimension::Constant(d)).collect();
408                ctx.make_tensor_type(dims, *dtype, MemoryLayout::Contiguous)
409            }
410            ModelType::Qubit => ctx.make_qubit_type(),
411            ModelType::Bit => ctx.make_bit_type(),
412            ModelType::Integer { bits } => ctx.make_integer_type(*bits, true),
413        }
414    }
415}
416
417/// Handle that chains function-building methods back to the model builder.
418pub struct FunctionBuilderHandle {
419    model: ModelBuilder,
420    func_idx: usize,
421}
422
423impl FunctionBuilderHandle {
424    pub fn param(mut self, name: &str, ty: ModelType) -> Self {
425        let fb = std::mem::replace(
426            &mut self.model.functions[self.func_idx],
427            FunctionBuilder::new("__tmp__"),
428        );
429        self.model.functions[self.func_idx] = fb.param(name, ty);
430        self
431    }
432
433    pub fn op(
434        mut self,
435        op_name: &str,
436        inputs: &[&str],
437        result: &str,
438        result_type: ModelType,
439    ) -> Self {
440        let fb = std::mem::replace(
441            &mut self.model.functions[self.func_idx],
442            FunctionBuilder::new("__tmp__"),
443        );
444        self.model.functions[self.func_idx] = fb.op(op_name, inputs, result, result_type);
445        self
446    }
447
448    pub fn op_with_attrs(
449        mut self,
450        op_name: &str,
451        inputs: &[&str],
452        result: &str,
453        result_type: ModelType,
454        attrs: Vec<(&str, Attribute)>,
455    ) -> Self {
456        let fb = std::mem::replace(
457            &mut self.model.functions[self.func_idx],
458            FunctionBuilder::new("__tmp__"),
459        );
460        self.model.functions[self.func_idx] =
461            fb.op_with_attrs(op_name, inputs, result, result_type, attrs);
462        self
463    }
464
465    pub fn returns(mut self, name: &str) -> Self {
466        let fb = std::mem::replace(
467            &mut self.model.functions[self.func_idx],
468            FunctionBuilder::new("__tmp__"),
469        );
470        self.model.functions[self.func_idx] = fb.returns(name);
471        self
472    }
473
474    pub fn dialect(mut self, d: &str) -> Self {
475        let fb = std::mem::replace(
476            &mut self.model.functions[self.func_idx],
477            FunctionBuilder::new("__tmp__"),
478        );
479        self.model.functions[self.func_idx] = fb.dialect(d);
480        self
481    }
482
483    /// Finish the function and return to the model builder.
484    pub fn done(self) -> ModelBuilder {
485        self.model
486    }
487}
488
489/// Format a ModelType as a `.lif` type string.
490fn format_model_type(mt: &ModelType) -> String {
491    match mt {
492        ModelType::Tensor { shape, dtype } => {
493            let dims: String = shape
494                .iter()
495                .map(|d| d.to_string())
496                .collect::<Vec<_>>()
497                .join("x");
498            let dt = match dtype {
499                DataType::FP64 => "f64",
500                DataType::FP32 => "f32",
501                DataType::FP16 => "f16",
502                DataType::BF16 => "bf16",
503                DataType::FP8E4M3 => "f8e4m3",
504                DataType::FP8E5M2 => "f8e5m2",
505                DataType::INT64 => "i64",
506                DataType::INT32 => "i32",
507                DataType::INT16 => "i16",
508                DataType::INT8 => "i8",
509                DataType::INT4 => "i4",
510                DataType::INT2 => "i2",
511                DataType::UINT8 => "ui8",
512                DataType::Bool => "i1",
513                DataType::Index => "index",
514            };
515            format!("tensor<{}x{}>", dims, dt)
516        }
517        ModelType::Qubit => "qubit".to_string(),
518        ModelType::Bit => "bit".to_string(),
519        ModelType::Integer { bits } => format!("i{}", bits),
520    }
521}
522
523/// Generate a `.lith` configuration string.
524pub fn build_lith_config(
525    backend: &str,
526    device: &str,
527    precision: &str,
528    passes: &[&str],
529    max_flops: Option<u64>,
530    max_memory: Option<u64>,
531) -> String {
532    let mut out = String::new();
533    out.push_str(&format!(
534        "[target]\nbackend = \"{}\"\ndevice = \"{}\"\nprecision = \"{}\"\n\n",
535        backend, device, precision
536    ));
537
538    if max_flops.is_some() || max_memory.is_some() {
539        out.push_str("[budget]\n");
540        if let Some(f) = max_flops {
541            out.push_str(&format!("max_flops = {}\n", f));
542        }
543        if let Some(m) = max_memory {
544            out.push_str(&format!("max_memory_bytes = {}\n", m));
545        }
546        out.push('\n');
547    }
548
549    out.push_str("[optimisation]\nlevel = O3\n");
550    out.push_str(&format!("passes = {}\n", passes.join(", ")));
551    out.push_str("max_iterations = 10\n\n");
552
553    out.push_str("[simulation]\nshape_propagation = true\nflop_counting = true\nmemory_analysis = true\nnoise_simulation = false\n");
554
555    out
556}
557
558#[cfg(test)]
559mod tests {
560    use super::*;
561
562    #[test]
563    fn test_simple_model() {
564        let lif = ModelBuilder::new("test_mlp")
565            .function("forward")
566            .param("x", tensor(&[1, 784], DataType::FP32))
567            .param("w", tensor_2d(784, 256, DataType::FP32))
568            .op(
569                "tensor.matmul",
570                &["x", "w"],
571                "h",
572                tensor(&[1, 256], DataType::FP32),
573            )
574            .op(
575                "tensor.relu",
576                &["h"],
577                "out",
578                tensor(&[1, 256], DataType::FP32),
579            )
580            .returns("out")
581            .done()
582            .build_lif();
583
584        assert!(lif.contains("#dialect tensor"));
585        assert!(lif.contains("module @test_mlp"));
586        assert!(lif.contains("tensor.matmul"));
587        assert!(lif.contains("tensor.relu"));
588        assert!(lif.contains("return %out"));
589    }
590
591    #[test]
592    fn test_build_context() {
593        let ctx = ModelBuilder::new("ctx_test")
594            .function("f")
595            .param("a", tensor(&[4, 4], DataType::FP32))
596            .param("b", tensor(&[4, 4], DataType::FP32))
597            .op(
598                "tensor.add",
599                &["a", "b"],
600                "c",
601                tensor(&[4, 4], DataType::FP32),
602            )
603            .returns("c")
604            .done()
605            .build_context();
606
607        assert_eq!(ctx.modules.len(), 1);
608        assert_eq!(ctx.modules[0].functions.len(), 1);
609        assert!(ctx.ops.len() >= 2); // add + return
610    }
611
612    #[test]
613    fn test_write_lif() {
614        let path = "/tmp/lift_builder_test.lif";
615        ModelBuilder::new("write_test")
616            .function("f")
617            .param("x", tensor(&[8], DataType::FP32))
618            .op("tensor.relu", &["x"], "y", tensor(&[8], DataType::FP32))
619            .returns("y")
620            .done()
621            .write_lif(path)
622            .unwrap();
623
624        let content = std::fs::read_to_string(path).unwrap();
625        assert!(content.contains("#dialect tensor"));
626        assert!(content.contains("tensor.relu"));
627    }
628
629    #[test]
630    fn test_lith_config() {
631        let config = build_lith_config(
632            "llvm",
633            "h100",
634            "fp16",
635            &["canonicalize", "dce", "tensor-fusion"],
636            Some(1_000_000_000),
637            Some(80_000_000_000),
638        );
639        assert!(config.contains("backend = \"llvm\""));
640        assert!(config.contains("canonicalize, dce, tensor-fusion"));
641        assert!(config.contains("max_flops = 1000000000"));
642    }
643}
lift_core/model_builder.rs

lift_core/
model_builder.rs
