Skip to main content

melior/
lib.rs

1#![doc = include_str!("../README.md")]
2
3extern crate self as melior;
4
5#[macro_use]
6mod r#macro;
7mod context;
8pub mod diagnostic;
9pub mod dialect;
10mod error;
11mod execution_engine;
12#[cfg(feature = "helpers")]
13pub mod helpers;
14pub mod ir;
15mod logical_result;
16pub mod pass;
17mod string_ref;
18
19#[cfg(test)]
20mod test;
21pub mod utility;
22
23pub use self::{
24    context::{Context, ContextRef},
25    error::Error,
26    execution_engine::ExecutionEngine,
27    string_ref::StringRef,
28};
29
30pub use melior_macro::dialect;
31
32#[cfg(test)]
33mod tests {
34    use crate::{
35        context::Context,
36        dialect::{self, arith, func, scf},
37        ir::{
38            attribute::{IntegerAttribute, StringAttribute, TypeAttribute},
39            operation::OperationBuilder,
40            r#type::{FunctionType, IntegerType},
41            Block, Location, Module, Region, Type, Value,
42        },
43        test::load_all_dialects,
44    };
45
46    #[test]
47    fn build_module() {
48        let context = Context::new();
49        let module = Module::new(Location::unknown(&context));
50
51        assert!(module.as_operation().verify());
52        insta::assert_snapshot!(module.as_operation());
53    }
54
55    #[test]
56    fn build_module_with_dialect() {
57        let registry = dialect::DialectRegistry::new();
58        let context = Context::new();
59        context.append_dialect_registry(&registry);
60        let module = Module::new(Location::unknown(&context));
61
62        assert!(module.as_operation().verify());
63        insta::assert_snapshot!(module.as_operation());
64    }
65
66    #[test]
67    fn build_add() {
68        let context = Context::new();
69        load_all_dialects(&context);
70
71        let location = Location::unknown(&context);
72        let module = Module::new(location);
73
74        let integer_type = IntegerType::new(&context, 64).into();
75
76        let function = {
77            let block = Block::new(&[(integer_type, location), (integer_type, location)]);
78
79            let sum = block.append_operation(arith::addi(
80                block.argument(0).unwrap().into(),
81                block.argument(1).unwrap().into(),
82                location,
83            ));
84
85            block.append_operation(func::r#return(&[sum.result(0).unwrap().into()], location));
86
87            let region = Region::new();
88            region.append_block(block);
89
90            func::func(
91                &context,
92                StringAttribute::new(&context, "add"),
93                TypeAttribute::new(
94                    FunctionType::new(&context, &[integer_type, integer_type], &[integer_type])
95                        .into(),
96                ),
97                region,
98                &[],
99                Location::unknown(&context),
100            )
101        };
102
103        module.body().append_operation(function);
104
105        assert!(module.as_operation().verify());
106        insta::assert_snapshot!(module.as_operation());
107    }
108
109    #[test]
110    fn build_sum() {
111        let context = Context::new();
112        load_all_dialects(&context);
113
114        let location = Location::unknown(&context);
115        let module = Module::new(location);
116
117        let memref_type = Type::parse(&context, "memref<?xf32>").unwrap();
118
119        let function = {
120            let function_block = Block::new(&[(memref_type, location), (memref_type, location)]);
121            let index_type = Type::parse(&context, "index").unwrap();
122
123            let zero = function_block.append_operation(arith::constant(
124                &context,
125                IntegerAttribute::new(Type::index(&context), 0).into(),
126                location,
127            ));
128
129            let dim = function_block.append_operation(
130                OperationBuilder::new("memref.dim", location)
131                    .add_operands(&[
132                        function_block.argument(0).unwrap().into(),
133                        zero.result(0).unwrap().into(),
134                    ])
135                    .add_results(&[index_type])
136                    .build()
137                    .unwrap(),
138            );
139
140            let loop_block = Block::new(&[(index_type, location)]);
141
142            let one = function_block.append_operation(arith::constant(
143                &context,
144                IntegerAttribute::new(Type::index(&context), 1).into(),
145                location,
146            ));
147
148            {
149                let f32_type = Type::float32(&context);
150
151                let lhs = loop_block.append_operation(
152                    OperationBuilder::new("memref.load", location)
153                        .add_operands(&[
154                            function_block.argument(0).unwrap().into(),
155                            loop_block.argument(0).unwrap().into(),
156                        ])
157                        .add_results(&[f32_type])
158                        .build()
159                        .unwrap(),
160                );
161
162                let rhs = loop_block.append_operation(
163                    OperationBuilder::new("memref.load", location)
164                        .add_operands(&[
165                            function_block.argument(1).unwrap().into(),
166                            loop_block.argument(0).unwrap().into(),
167                        ])
168                        .add_results(&[f32_type])
169                        .build()
170                        .unwrap(),
171                );
172
173                let add = loop_block.append_operation(arith::addf(
174                    lhs.result(0).unwrap().into(),
175                    rhs.result(0).unwrap().into(),
176                    location,
177                ));
178
179                loop_block.append_operation(
180                    OperationBuilder::new("memref.store", location)
181                        .add_operands(&[
182                            add.result(0).unwrap().into(),
183                            function_block.argument(0).unwrap().into(),
184                            loop_block.argument(0).unwrap().into(),
185                        ])
186                        .build()
187                        .unwrap(),
188                );
189
190                loop_block.append_operation(scf::r#yield(&[], location));
191            }
192
193            function_block.append_operation(scf::r#for(
194                zero.result(0).unwrap().into(),
195                dim.result(0).unwrap().into(),
196                one.result(0).unwrap().into(),
197                {
198                    let loop_region = Region::new();
199                    loop_region.append_block(loop_block);
200                    loop_region
201                },
202                location,
203            ));
204
205            function_block.append_operation(func::r#return(&[], location));
206
207            let function_region = Region::new();
208            function_region.append_block(function_block);
209
210            func::func(
211                &context,
212                StringAttribute::new(&context, "sum"),
213                TypeAttribute::new(
214                    FunctionType::new(&context, &[memref_type, memref_type], &[]).into(),
215                ),
216                function_region,
217                &[],
218                Location::unknown(&context),
219            )
220        };
221
222        module.body().append_operation(function);
223
224        assert!(module.as_operation().verify());
225        insta::assert_snapshot!(module.as_operation());
226    }
227
228    #[test]
229    fn return_value_from_function() {
230        let context = Context::new();
231        load_all_dialects(&context);
232
233        let location = Location::unknown(&context);
234        let module = Module::new(location);
235
236        let integer_type = IntegerType::new(&context, 64).into();
237
238        fn compile_add<'c, 'a>(
239            context: &'c Context,
240            block: &'a Block<'c>,
241            lhs: Value<'c, '_>,
242            rhs: Value<'c, '_>,
243        ) -> Value<'c, 'a> {
244            block
245                .append_operation(arith::addi(lhs, rhs, Location::unknown(context)))
246                .result(0)
247                .unwrap()
248                .into()
249        }
250
251        module.body().append_operation(func::func(
252            &context,
253            StringAttribute::new(&context, "add"),
254            TypeAttribute::new(
255                FunctionType::new(&context, &[integer_type, integer_type], &[integer_type]).into(),
256            ),
257            {
258                let block = Block::new(&[(integer_type, location), (integer_type, location)]);
259
260                block.append_operation(func::r#return(
261                    &[compile_add(
262                        &context,
263                        &block,
264                        block.argument(0).unwrap().into(),
265                        block.argument(1).unwrap().into(),
266                    )],
267                    location,
268                ));
269
270                let region = Region::new();
271                region.append_block(block);
272                region
273            },
274            &[],
275            Location::unknown(&context),
276        ));
277
278        assert!(module.as_operation().verify());
279        insta::assert_snapshot!(module.as_operation());
280    }
281}