Skip to main content

melior/dialect/
arith.rs

1//! `arith` dialect.
2
3use crate::{
4    ir::{
5        attribute::IntegerAttribute, operation::OperationBuilder, r#type::IntegerType, Attribute,
6        Identifier, Location, Operation, Value, ValueLike,
7    },
8    Context,
9};
10
11// spell-checker: disable
12
13/// Creates an `arith.constant` operation.
14pub fn constant<'c>(
15    context: &'c Context,
16    value: Attribute<'c>,
17    location: Location<'c>,
18) -> Operation<'c> {
19    OperationBuilder::new("arith.constant", location)
20        .add_attributes(&[(Identifier::new(context, "value"), value)])
21        .enable_result_type_inference()
22        .build()
23        .expect("valid operation")
24}
25
26/// `arith.cmpf` predicate
27pub enum CmpfPredicate {
28    False,
29    Oeq,
30    Ogt,
31    Oge,
32    Olt,
33    Ole,
34    One,
35    Ord,
36    Ueq,
37    Ugt,
38    Uge,
39    Ult,
40    Ule,
41    Une,
42    Uno,
43    True,
44}
45
46/// Creates an `arith.cmpf` operation.
47pub fn cmpf<'c>(
48    context: &'c Context,
49    predicate: CmpfPredicate,
50    lhs: Value<'c, '_>,
51    rhs: Value<'c, '_>,
52    location: Location<'c>,
53) -> Operation<'c> {
54    cmp(context, "arith.cmpf", predicate as i64, lhs, rhs, location)
55}
56
57/// `arith.cmpi` predicate
58pub enum CmpiPredicate {
59    Eq,
60    Ne,
61    Slt,
62    Sle,
63    Sgt,
64    Sge,
65    Ult,
66    Ule,
67    Ugt,
68    Uge,
69}
70
71/// Creates an `arith.cmpi` operation.
72pub fn cmpi<'c>(
73    context: &'c Context,
74    predicate: CmpiPredicate,
75    lhs: Value<'c, '_>,
76    rhs: Value<'c, '_>,
77    location: Location<'c>,
78) -> Operation<'c> {
79    cmp(context, "arith.cmpi", predicate as i64, lhs, rhs, location)
80}
81
82fn cmp<'c>(
83    context: &'c Context,
84    name: &str,
85    predicate: i64,
86    lhs: Value<'c, '_>,
87    rhs: Value<'c, '_>,
88    location: Location<'c>,
89) -> Operation<'c> {
90    OperationBuilder::new(name, location)
91        .add_attributes(&[(
92            Identifier::new(context, "predicate"),
93            IntegerAttribute::new(IntegerType::new(context, 64).into(), predicate).into(),
94        )])
95        .add_operands(&[lhs, rhs])
96        .enable_result_type_inference()
97        .build()
98        .expect("valid operation")
99}
100
101/// Creates an `arith.select` operation.
102pub fn select<'c>(
103    condition: Value<'c, '_>,
104    true_value: Value<'c, '_>,
105    false_value: Value<'c, '_>,
106    location: Location<'c>,
107) -> Operation<'c> {
108    OperationBuilder::new("arith.select", location)
109        .add_operands(&[condition, true_value, false_value])
110        .add_results(&[true_value.r#type()])
111        .build()
112        .expect("valid operation")
113}
114
115melior_macro::binary_operations!(
116    arith,
117    [
118        addf,
119        addi,
120        addui_extended,
121        andi,
122        ceildivsi,
123        ceildivui,
124        divf,
125        divsi,
126        divui,
127        floordivsi,
128        maxf,
129        maxsi,
130        maxui,
131        minf,
132        minsi,
133        minui,
134        mulf,
135        muli,
136        mulsi_extended,
137        mului_extended,
138        ori,
139        remf,
140        remsi,
141        remui,
142        shli,
143        shrsi,
144        shrui,
145        subf,
146        subi,
147        xori,
148    ]
149);
150
151melior_macro::unary_operations!(arith, [negf, truncf]);
152
153melior_macro::typed_unary_operations!(
154    arith,
155    [
156        bitcast,
157        extf,
158        extsi,
159        extui,
160        fptosi,
161        fptoui,
162        index_cast,
163        index_castui,
164        sitofp,
165        trunci,
166        uitofp
167    ]
168);
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use crate::{
174        dialect::func,
175        ir::{
176            attribute::{StringAttribute, TypeAttribute},
177            r#type::FunctionType,
178            Attribute, Block, Location, Module, Region, Type,
179        },
180        test::load_all_dialects,
181        Context,
182    };
183
184    fn create_context() -> Context {
185        let context = Context::new();
186        load_all_dialects(&context);
187        context
188    }
189
190    fn compile_operation<'c>(
191        context: &'c Context,
192        operation: impl Fn(&Block<'c>) -> Operation<'c>,
193        block_argument_types: &[Type<'c>],
194        function_type: FunctionType<'c>,
195    ) {
196        let location = Location::unknown(context);
197        let module = Module::new(location);
198
199        let block = Block::new(
200            &block_argument_types
201                .iter()
202                .map(|&r#type| (r#type, location))
203                .collect::<Vec<_>>(),
204        );
205
206        let operation = operation(&block);
207        let name = operation.name();
208        let name = name.as_string_ref().as_str().unwrap();
209
210        block.append_operation(func::r#return(
211            &[block.append_operation(operation).result(0).unwrap().into()],
212            location,
213        ));
214
215        let region = Region::new();
216        region.append_block(block);
217
218        let function = func::func(
219            context,
220            StringAttribute::new(context, "foo"),
221            TypeAttribute::new(function_type.into()),
222            region,
223            &[],
224            Location::unknown(context),
225        );
226
227        module.body().append_operation(function);
228
229        assert!(module.as_operation().verify());
230        insta::assert_snapshot!(name, module.as_operation());
231    }
232
233    #[test]
234    fn compile_constant() {
235        let context = create_context();
236        let integer_type = IntegerType::new(&context, 64).into();
237
238        compile_operation(
239            &context,
240            |_| {
241                constant(
242                    &context,
243                    Attribute::parse(&context, "42 : i64").unwrap(),
244                    Location::unknown(&context),
245                )
246            },
247            &[integer_type],
248            FunctionType::new(&context, &[integer_type], &[integer_type]),
249        );
250    }
251
252    #[test]
253    fn compile_negf() {
254        let context = create_context();
255        let f64_type = Type::float64(&context);
256
257        compile_operation(
258            &context,
259            |block| {
260                negf(
261                    block.argument(0).unwrap().into(),
262                    Location::unknown(&context),
263                )
264            },
265            &[Type::float64(&context)],
266            FunctionType::new(&context, &[f64_type], &[f64_type]),
267        );
268    }
269
270    mod cmp {
271        use super::*;
272
273        #[test]
274        fn compile_cmpf() {
275            let context = create_context();
276            let float_type = Type::float64(&context);
277
278            compile_operation(
279                &context,
280                |block| {
281                    cmpf(
282                        &context,
283                        CmpfPredicate::Oeq,
284                        block.argument(0).unwrap().into(),
285                        block.argument(1).unwrap().into(),
286                        Location::unknown(&context),
287                    )
288                },
289                &[float_type, float_type],
290                FunctionType::new(
291                    &context,
292                    &[float_type, float_type],
293                    &[IntegerType::new(&context, 1).into()],
294                ),
295            );
296        }
297
298        #[test]
299        fn compile_cmpi() {
300            let context = create_context();
301            let integer_type = IntegerType::new(&context, 64).into();
302
303            compile_operation(
304                &context,
305                |block| {
306                    cmpi(
307                        &context,
308                        CmpiPredicate::Eq,
309                        block.argument(0).unwrap().into(),
310                        block.argument(1).unwrap().into(),
311                        Location::unknown(&context),
312                    )
313                },
314                &[integer_type, integer_type],
315                FunctionType::new(
316                    &context,
317                    &[integer_type, integer_type],
318                    &[IntegerType::new(&context, 1).into()],
319                ),
320            );
321        }
322    }
323
324    mod typed_unary {
325        use super::*;
326
327        #[test]
328        fn compile_bitcast() {
329            let context = create_context();
330            let integer_type = IntegerType::new(&context, 64).into();
331            let float_type = Type::float64(&context);
332
333            compile_operation(
334                &context,
335                |block| {
336                    bitcast(
337                        block.argument(0).unwrap().into(),
338                        float_type,
339                        Location::unknown(&context),
340                    )
341                },
342                &[integer_type],
343                FunctionType::new(&context, &[integer_type], &[float_type]),
344            );
345        }
346
347        #[test]
348        fn compile_extf() {
349            let context = create_context();
350
351            compile_operation(
352                &context,
353                |block| {
354                    extf(
355                        block.argument(0).unwrap().into(),
356                        Type::float64(&context),
357                        Location::unknown(&context),
358                    )
359                },
360                &[Type::float32(&context)],
361                FunctionType::new(
362                    &context,
363                    &[Type::float32(&context)],
364                    &[Type::float64(&context)],
365                ),
366            );
367        }
368
369        #[test]
370        fn compile_extsi() {
371            let context = create_context();
372
373            compile_operation(
374                &context,
375                |block| {
376                    extsi(
377                        block.argument(0).unwrap().into(),
378                        IntegerType::new(&context, 64).into(),
379                        Location::unknown(&context),
380                    )
381                },
382                &[IntegerType::new(&context, 32).into()],
383                FunctionType::new(
384                    &context,
385                    &[IntegerType::new(&context, 32).into()],
386                    &[IntegerType::new(&context, 64).into()],
387                ),
388            );
389        }
390
391        #[test]
392        fn compile_extui() {
393            let context = create_context();
394
395            compile_operation(
396                &context,
397                |block| {
398                    extui(
399                        block.argument(0).unwrap().into(),
400                        IntegerType::new(&context, 64).into(),
401                        Location::unknown(&context),
402                    )
403                },
404                &[IntegerType::new(&context, 32).into()],
405                FunctionType::new(
406                    &context,
407                    &[IntegerType::new(&context, 32).into()],
408                    &[IntegerType::new(&context, 64).into()],
409                ),
410            );
411        }
412
413        #[test]
414        fn compile_fptosi() {
415            let context = create_context();
416
417            compile_operation(
418                &context,
419                |block| {
420                    fptosi(
421                        block.argument(0).unwrap().into(),
422                        IntegerType::new(&context, 64).into(),
423                        Location::unknown(&context),
424                    )
425                },
426                &[Type::float32(&context)],
427                FunctionType::new(
428                    &context,
429                    &[Type::float32(&context)],
430                    &[IntegerType::new(&context, 64).into()],
431                ),
432            );
433        }
434
435        #[test]
436        fn compile_fptoui() {
437            let context = create_context();
438
439            compile_operation(
440                &context,
441                |block| {
442                    fptoui(
443                        block.argument(0).unwrap().into(),
444                        IntegerType::new(&context, 64).into(),
445                        Location::unknown(&context),
446                    )
447                },
448                &[Type::float32(&context)],
449                FunctionType::new(
450                    &context,
451                    &[Type::float32(&context)],
452                    &[IntegerType::new(&context, 64).into()],
453                ),
454            );
455        }
456
457        #[test]
458        fn compile_index_cast() {
459            let context = create_context();
460
461            compile_operation(
462                &context,
463                |block| {
464                    index_cast(
465                        block.argument(0).unwrap().into(),
466                        IntegerType::new(&context, 64).into(),
467                        Location::unknown(&context),
468                    )
469                },
470                &[Type::index(&context)],
471                FunctionType::new(
472                    &context,
473                    &[Type::index(&context)],
474                    &[IntegerType::new(&context, 64).into()],
475                ),
476            );
477        }
478
479        #[test]
480        fn compile_index_castui() {
481            let context = create_context();
482
483            compile_operation(
484                &context,
485                |block| {
486                    index_castui(
487                        block.argument(0).unwrap().into(),
488                        IntegerType::new(&context, 64).into(),
489                        Location::unknown(&context),
490                    )
491                },
492                &[Type::index(&context)],
493                FunctionType::new(
494                    &context,
495                    &[Type::index(&context)],
496                    &[IntegerType::new(&context, 64).into()],
497                ),
498            );
499        }
500
501        #[test]
502        fn compile_sitofp() {
503            let context = create_context();
504
505            compile_operation(
506                &context,
507                |block| {
508                    sitofp(
509                        block.argument(0).unwrap().into(),
510                        Type::float64(&context),
511                        Location::unknown(&context),
512                    )
513                },
514                &[IntegerType::new(&context, 32).into()],
515                FunctionType::new(
516                    &context,
517                    &[IntegerType::new(&context, 32).into()],
518                    &[Type::float64(&context)],
519                ),
520            );
521        }
522
523        #[test]
524        fn compile_trunci() {
525            let context = create_context();
526
527            compile_operation(
528                &context,
529                |block| {
530                    trunci(
531                        block.argument(0).unwrap().into(),
532                        IntegerType::new(&context, 32).into(),
533                        Location::unknown(&context),
534                    )
535                },
536                &[IntegerType::new(&context, 64).into()],
537                FunctionType::new(
538                    &context,
539                    &[IntegerType::new(&context, 64).into()],
540                    &[IntegerType::new(&context, 32).into()],
541                ),
542            );
543        }
544
545        #[test]
546        fn compile_uitofp() {
547            let context = create_context();
548
549            compile_operation(
550                &context,
551                |block| {
552                    uitofp(
553                        block.argument(0).unwrap().into(),
554                        Type::float64(&context),
555                        Location::unknown(&context),
556                    )
557                },
558                &[IntegerType::new(&context, 32).into()],
559                FunctionType::new(
560                    &context,
561                    &[IntegerType::new(&context, 32).into()],
562                    &[Type::float64(&context)],
563                ),
564            );
565        }
566    }
567
568    #[test]
569    fn compile_addi() {
570        let context = Context::new();
571        load_all_dialects(&context);
572
573        let location = Location::unknown(&context);
574        let module = Module::new(location);
575
576        let integer_type = IntegerType::new(&context, 64).into();
577
578        let function = {
579            let block = Block::new(&[(integer_type, location), (integer_type, location)]);
580
581            let sum = block.append_operation(addi(
582                block.argument(0).unwrap().into(),
583                block.argument(1).unwrap().into(),
584                location,
585            ));
586
587            block.append_operation(func::r#return(&[sum.result(0).unwrap().into()], location));
588
589            let region = Region::new();
590            region.append_block(block);
591
592            func::func(
593                &context,
594                StringAttribute::new(&context, "foo"),
595                TypeAttribute::new(
596                    FunctionType::new(&context, &[integer_type, integer_type], &[integer_type])
597                        .into(),
598                ),
599                region,
600                &[],
601                Location::unknown(&context),
602            )
603        };
604
605        module.body().append_operation(function);
606
607        assert!(module.as_operation().verify());
608        insta::assert_snapshot!(module.as_operation());
609    }
610
611    #[test]
612    fn compile_select() {
613        let context = Context::new();
614        load_all_dialects(&context);
615
616        let location = Location::unknown(&context);
617        let module = Module::new(location);
618
619        let integer_type = IntegerType::new(&context, 64).into();
620        let bool_type = IntegerType::new(&context, 1).into();
621
622        let function = {
623            let block = Block::new(&[
624                (bool_type, location),
625                (integer_type, location),
626                (integer_type, location),
627            ]);
628
629            let val = block.append_operation(select(
630                block.argument(0).unwrap().into(),
631                block.argument(1).unwrap().into(),
632                block.argument(2).unwrap().into(),
633                location,
634            ));
635
636            block.append_operation(func::r#return(&[val.result(0).unwrap().into()], location));
637
638            let region = Region::new();
639            region.append_block(block);
640
641            func::func(
642                &context,
643                StringAttribute::new(&context, "foo"),
644                TypeAttribute::new(
645                    FunctionType::new(
646                        &context,
647                        &[bool_type, integer_type, integer_type],
648                        &[integer_type],
649                    )
650                    .into(),
651                ),
652                region,
653                &[],
654                Location::unknown(&context),
655            )
656        };
657
658        module.body().append_operation(function);
659
660        assert!(module.as_operation().verify());
661        insta::assert_snapshot!(module.as_operation());
662    }
663}