melior/dialect/
memref.rs

1//! `memref` dialect.
2
3use crate::{
4    ir::{
5        attribute::{
6            DenseI32ArrayAttribute, DenseI64ArrayAttribute, FlatSymbolRefAttribute,
7            IntegerAttribute, StringAttribute, TypeAttribute,
8        },
9        operation::OperationBuilder,
10        r#type::MemRefType,
11        Attribute, Identifier, Location, Operation, Value,
12    },
13    Context,
14};
15
16/// Create a `memref.alloc` operation.
17pub fn alloc<'c>(
18    context: &'c Context,
19    r#type: MemRefType<'c>,
20    dynamic_sizes: &[Value<'c, '_>],
21    symbols: &[Value<'c, '_>],
22    alignment: Option<IntegerAttribute<'c>>,
23    location: Location<'c>,
24) -> Operation<'c> {
25    allocate(
26        context,
27        "memref.alloc",
28        r#type,
29        dynamic_sizes,
30        symbols,
31        alignment,
32        location,
33    )
34}
35
36/// Create a `memref.alloca` operation.
37pub fn alloca<'c>(
38    context: &'c Context,
39    r#type: MemRefType<'c>,
40    dynamic_sizes: &[Value<'c, '_>],
41    symbols: &[Value<'c, '_>],
42    alignment: Option<IntegerAttribute<'c>>,
43    location: Location<'c>,
44) -> Operation<'c> {
45    allocate(
46        context,
47        "memref.alloca",
48        r#type,
49        dynamic_sizes,
50        symbols,
51        alignment,
52        location,
53    )
54}
55
56fn allocate<'c>(
57    context: &'c Context,
58    name: &str,
59    r#type: MemRefType<'c>,
60    dynamic_sizes: &[Value<'c, '_>],
61    symbols: &[Value<'c, '_>],
62    alignment: Option<IntegerAttribute<'c>>,
63    location: Location<'c>,
64) -> Operation<'c> {
65    let mut builder = OperationBuilder::new(name, location);
66
67    builder = builder.add_attributes(&[(
68        Identifier::new(context, "operand_segment_sizes"),
69        DenseI32ArrayAttribute::new(context, &[dynamic_sizes.len() as i32, symbols.len() as i32])
70            .into(),
71    )]);
72    builder = builder.add_operands(dynamic_sizes).add_operands(symbols);
73
74    if let Some(alignment) = alignment {
75        builder =
76            builder.add_attributes(&[(Identifier::new(context, "alignment"), alignment.into())]);
77    }
78
79    builder
80        .add_results(&[r#type.into()])
81        .build()
82        .expect("valid operation")
83}
84
85/// Create a `memref.cast` operation.
86pub fn cast<'c>(
87    value: Value<'c, '_>,
88    r#type: MemRefType<'c>,
89    location: Location<'c>,
90) -> Operation<'c> {
91    OperationBuilder::new("memref.cast", location)
92        .add_operands(&[value])
93        .add_results(&[r#type.into()])
94        .build()
95        .expect("valid operation")
96}
97
98/// Create a `memref.dealloc` operation.
99pub fn dealloc<'c>(value: Value<'c, '_>, location: Location<'c>) -> Operation<'c> {
100    OperationBuilder::new("memref.dealloc", location)
101        .add_operands(&[value])
102        .build()
103        .expect("valid operation")
104}
105
106/// Create a `memref.dim` operation.
107pub fn dim<'c>(
108    value: Value<'c, '_>,
109    index: Value<'c, '_>,
110    location: Location<'c>,
111) -> Operation<'c> {
112    OperationBuilder::new("memref.dim", location)
113        .add_operands(&[value, index])
114        .enable_result_type_inference()
115        .build()
116        .expect("valid operation")
117}
118
119/// Create a `memref.get_global` operation.
120pub fn get_global<'c>(
121    context: &'c Context,
122    name: &str,
123    r#type: MemRefType<'c>,
124    location: Location<'c>,
125) -> Operation<'c> {
126    OperationBuilder::new("memref.get_global", location)
127        .add_attributes(&[(
128            Identifier::new(context, "name"),
129            FlatSymbolRefAttribute::new(context, name).into(),
130        )])
131        .add_results(&[r#type.into()])
132        .build()
133        .expect("valid operation")
134}
135
136/// Create a `memref.view` operation.
137pub fn view<'c>(
138    context: &'c Context,
139    source: Value<'c, '_>,
140    byte_shift: Value<'c, '_>,
141    sizes: &[Value<'c, '_>],
142    result_type: MemRefType<'c>,
143    location: Location<'c>,
144) -> Operation<'c> {
145    OperationBuilder::new("memref.view", location)
146        .add_operands(&[source])
147        .add_operands(&[byte_shift])
148        .add_operands(sizes)
149        .add_results(&[result_type.into()])
150        .add_attributes(&[(
151            Identifier::new(context, "operand_segment_sizes"),
152            DenseI32ArrayAttribute::new(context, &[1, 1, sizes.len() as i32]).into(),
153        )])
154        .build()
155        .expect("valid operation")
156}
157
158/// Create a `memref.subview` operation.
159#[allow(clippy::too_many_arguments)]
160pub fn subview<'c>(
161    context: &'c Context,
162    source: Value<'c, '_>,
163    offsets: &[Value<'c, '_>],
164    sizes: &[Value<'c, '_>],
165    strides: &[Value<'c, '_>],
166    static_offsets: &[i64],
167    static_sizes: &[i64],
168    static_strides: &[i64],
169    result_type: MemRefType<'c>,
170    location: Location<'c>,
171) -> Operation<'c> {
172    OperationBuilder::new("memref.subview", location)
173        .add_operands(&[source])
174        .add_operands(offsets)
175        .add_operands(sizes)
176        .add_operands(strides)
177        .add_results(&[result_type.into()])
178        .add_attributes(&[
179            (
180                Identifier::new(context, "operand_segment_sizes"),
181                DenseI32ArrayAttribute::new(
182                    context,
183                    &[
184                        1,
185                        offsets.len() as i32,
186                        sizes.len() as i32,
187                        strides.len() as i32,
188                    ],
189                )
190                .into(),
191            ),
192            (
193                Identifier::new(context, "static_offsets"),
194                DenseI64ArrayAttribute::new(context, static_offsets).into(),
195            ),
196            (
197                Identifier::new(context, "static_sizes"),
198                DenseI64ArrayAttribute::new(context, static_sizes).into(),
199            ),
200            (
201                Identifier::new(context, "static_strides"),
202                DenseI64ArrayAttribute::new(context, static_strides).into(),
203            ),
204        ])
205        .build()
206        .expect("valid operation")
207}
208
209/// Create a `memref.global` operation.
210#[allow(clippy::too_many_arguments)]
211pub fn global<'c>(
212    context: &'c Context,
213    name: &str,
214    visibility: Option<&str>,
215    r#type: MemRefType<'c>,
216    value: Option<Attribute<'c>>,
217    constant: bool,
218    alignment: Option<IntegerAttribute<'c>>,
219    location: Location<'c>,
220) -> Operation<'c> {
221    let mut builder = OperationBuilder::new("memref.global", location).add_attributes(&[
222        (
223            Identifier::new(context, "sym_name"),
224            StringAttribute::new(context, name).into(),
225        ),
226        (
227            Identifier::new(context, "type"),
228            TypeAttribute::new(r#type.into()).into(),
229        ),
230        (
231            Identifier::new(context, "initial_value"),
232            value.unwrap_or_else(|| Attribute::unit(context)),
233        ),
234    ]);
235
236    if let Some(visibility) = visibility {
237        builder = builder.add_attributes(&[(
238            Identifier::new(context, "sym_visibility"),
239            StringAttribute::new(context, visibility).into(),
240        )]);
241    }
242
243    if constant {
244        builder = builder.add_attributes(&[(
245            Identifier::new(context, "constant"),
246            Attribute::unit(context),
247        )]);
248    }
249
250    if let Some(alignment) = alignment {
251        builder =
252            builder.add_attributes(&[(Identifier::new(context, "alignment"), alignment.into())]);
253    }
254
255    builder.build().expect("valid operation")
256}
257
258/// Create a `memref.load` operation.
259pub fn load<'c>(
260    memref: Value<'c, '_>,
261    indices: &[Value<'c, '_>],
262    location: Location<'c>,
263) -> Operation<'c> {
264    OperationBuilder::new("memref.load", location)
265        .add_operands(&[memref])
266        .add_operands(indices)
267        .enable_result_type_inference()
268        .build()
269        .expect("valid operation")
270}
271
272/// Create a `memref.rank` operation.
273pub fn rank<'c>(value: Value<'c, '_>, location: Location<'c>) -> Operation<'c> {
274    OperationBuilder::new("memref.rank", location)
275        .add_operands(&[value])
276        .enable_result_type_inference()
277        .build()
278        .expect("valid operation")
279}
280
281/// Create a `memref.store` operation.
282pub fn store<'c>(
283    value: Value<'c, '_>,
284    memref: Value<'c, '_>,
285    indices: &[Value<'c, '_>],
286    location: Location<'c>,
287) -> Operation<'c> {
288    OperationBuilder::new("memref.store", location)
289        .add_operands(&[value, memref])
290        .add_operands(indices)
291        .build()
292        .expect("valid operation")
293}
294
295/// Create a `memref.realloc` operation.
296pub fn realloc<'c>(
297    context: &'c Context,
298    value: Value<'c, '_>,
299    size: Option<Value<'c, '_>>,
300    r#type: MemRefType<'c>,
301    alignment: Option<IntegerAttribute<'c>>,
302    location: Location<'c>,
303) -> Operation<'c> {
304    let mut builder = OperationBuilder::new("memref.realloc", location)
305        .add_operands(&[value])
306        .add_results(&[r#type.into()]);
307
308    if let Some(size) = size {
309        builder = builder.add_operands(&[size]);
310    }
311
312    if let Some(alignment) = alignment {
313        builder =
314            builder.add_attributes(&[(Identifier::new(context, "alignment"), alignment.into())]);
315    }
316
317    builder.build().expect("valid operation")
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use crate::{
324        dialect::{func, index},
325        ir::{
326            attribute::{DenseElementsAttribute, StringAttribute, TypeAttribute},
327            r#type::{FunctionType, IntegerType, RankedTensorType},
328            Block, Module, Region, Type,
329        },
330        test::create_test_context,
331    };
332
333    fn compile_operation(name: &str, context: &Context, build_block: impl Fn(&Block)) {
334        let location = Location::unknown(context);
335        let module = Module::new(location);
336
337        let function = {
338            let block = Block::new(&[]);
339
340            build_block(&block);
341            block.append_operation(func::r#return(&[], location));
342
343            let region = Region::new();
344            region.append_block(block);
345
346            func::func(
347                context,
348                StringAttribute::new(context, "foo"),
349                TypeAttribute::new(FunctionType::new(context, &[], &[]).into()),
350                region,
351                &[],
352                Location::unknown(context),
353            )
354        };
355
356        module.body().append_operation(function);
357        assert!(module.as_operation().verify());
358        insta::assert_snapshot!(name, module.as_operation());
359    }
360
361    #[test]
362    fn compile_alloc_and_dealloc() {
363        let context = create_test_context();
364        let location = Location::unknown(&context);
365
366        compile_operation("alloc", &context, |block| {
367            let memref = block.append_operation(alloc(
368                &context,
369                MemRefType::new(Type::index(&context), &[], None, None),
370                &[],
371                &[],
372                None,
373                location,
374            ));
375            block.append_operation(dealloc(memref.result(0).unwrap().into(), location));
376        })
377    }
378
379    #[test]
380    fn compile_alloc_and_realloc() {
381        let context = create_test_context();
382        let location = Location::unknown(&context);
383
384        compile_operation("realloc", &context, |block| {
385            let memref = block.append_operation(alloc(
386                &context,
387                MemRefType::new(Type::index(&context), &[8], None, None),
388                &[],
389                &[],
390                None,
391                location,
392            ));
393            block.append_operation(realloc(
394                &context,
395                memref.result(0).unwrap().into(),
396                None,
397                MemRefType::new(Type::index(&context), &[42], None, None),
398                None,
399                location,
400            ));
401        })
402    }
403
404    #[test]
405    fn compile_alloca() {
406        let context = create_test_context();
407        let location = Location::unknown(&context);
408
409        compile_operation("alloca", &context, |block| {
410            block.append_operation(alloca(
411                &context,
412                MemRefType::new(Type::index(&context), &[], None, None),
413                &[],
414                &[],
415                None,
416                location,
417            ));
418        })
419    }
420
421    #[test]
422    fn compile_cast() {
423        let context = create_test_context();
424        let location = Location::unknown(&context);
425
426        compile_operation("cast", &context, |block| {
427            let memref = block.append_operation(alloca(
428                &context,
429                MemRefType::new(Type::float64(&context), &[42], None, None),
430                &[],
431                &[],
432                None,
433                location,
434            ));
435
436            block.append_operation(cast(
437                memref.result(0).unwrap().into(),
438                Type::parse(&context, "memref<?xf64>")
439                    .unwrap()
440                    .try_into()
441                    .unwrap(),
442                location,
443            ));
444        })
445    }
446
447    #[test]
448    fn compile_dim() {
449        let context = create_test_context();
450        let location = Location::unknown(&context);
451
452        compile_operation("dim", &context, |block| {
453            let memref = block.append_operation(alloca(
454                &context,
455                MemRefType::new(Type::index(&context), &[1], None, None),
456                &[],
457                &[],
458                None,
459                location,
460            ));
461
462            let index = block.append_operation(index::constant(
463                &context,
464                IntegerAttribute::new(Type::index(&context), 0),
465                location,
466            ));
467
468            block.append_operation(dim(
469                memref.result(0).unwrap().into(),
470                index.result(0).unwrap().into(),
471                location,
472            ));
473        })
474    }
475
476    #[test]
477    fn compile_get_global() {
478        let context = create_test_context();
479        let location = Location::unknown(&context);
480        let module = Module::new(location);
481        let mem_ref_type = MemRefType::new(Type::index(&context), &[], None, None);
482
483        module.body().append_operation(global(
484            &context,
485            "foo",
486            None,
487            mem_ref_type,
488            None,
489            false,
490            None,
491            location,
492        ));
493
494        module.body().append_operation(func::func(
495            &context,
496            StringAttribute::new(&context, "bar"),
497            TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
498            {
499                let block = Block::new(&[]);
500
501                block.append_operation(get_global(&context, "foo", mem_ref_type, location));
502                block.append_operation(func::r#return(&[], location));
503
504                let region = Region::new();
505                region.append_block(block);
506                region
507            },
508            &[],
509            location,
510        ));
511
512        assert!(module.as_operation().verify());
513        insta::assert_snapshot!(module.as_operation());
514    }
515
516    #[test]
517    fn compile_global() {
518        let context = create_test_context();
519        let location = Location::unknown(&context);
520        let module = Module::new(location);
521
522        module.body().append_operation(global(
523            &context,
524            "foo",
525            None,
526            MemRefType::new(Type::index(&context), &[], None, None),
527            None,
528            false,
529            None,
530            location,
531        ));
532
533        assert!(module.as_operation().verify());
534        insta::assert_snapshot!(module.as_operation());
535    }
536
537    #[test]
538    fn compile_global_with_options() {
539        let context = create_test_context();
540        let location = Location::unknown(&context);
541        let module = Module::new(location);
542        let r#type = IntegerType::new(&context, 64).into();
543
544        module.body().append_operation(global(
545            &context,
546            "foo",
547            Some("private"),
548            MemRefType::new(r#type, &[], None, None),
549            Some(
550                DenseElementsAttribute::new(
551                    RankedTensorType::new(&[], r#type, None).into(),
552                    &[IntegerAttribute::new(r#type, 42).into()],
553                )
554                .unwrap()
555                .into(),
556            ),
557            true,
558            Some(IntegerAttribute::new(
559                IntegerType::new(&context, 64).into(),
560                8,
561            )),
562            location,
563        ));
564
565        assert!(module.as_operation().verify());
566        insta::assert_snapshot!(module.as_operation());
567    }
568
569    #[test]
570    fn compile_load() {
571        let context = create_test_context();
572        let location = Location::unknown(&context);
573
574        compile_operation("load", &context, |block| {
575            let memref = block.append_operation(alloca(
576                &context,
577                MemRefType::new(Type::index(&context), &[], None, None),
578                &[],
579                &[],
580                None,
581                location,
582            ));
583            block.append_operation(load(memref.result(0).unwrap().into(), &[], location));
584        })
585    }
586
587    #[test]
588    fn compile_load_with_index() {
589        let context = create_test_context();
590        let location = Location::unknown(&context);
591
592        compile_operation("load_with_index", &context, |block| {
593            let memref = block.append_operation(alloca(
594                &context,
595                MemRefType::new(Type::index(&context), &[1], None, None),
596                &[],
597                &[],
598                None,
599                location,
600            ));
601
602            let index = block.append_operation(index::constant(
603                &context,
604                IntegerAttribute::new(Type::index(&context), 0),
605                location,
606            ));
607
608            block.append_operation(load(
609                memref.result(0).unwrap().into(),
610                &[index.result(0).unwrap().into()],
611                location,
612            ));
613        })
614    }
615
616    #[test]
617    fn compile_rank() {
618        let context = create_test_context();
619        let location = Location::unknown(&context);
620
621        compile_operation("rank", &context, |block| {
622            let memref = block.append_operation(alloca(
623                &context,
624                MemRefType::new(Type::index(&context), &[1], None, None),
625                &[],
626                &[],
627                None,
628                location,
629            ));
630            block.append_operation(rank(memref.result(0).unwrap().into(), location));
631        })
632    }
633
634    #[test]
635    fn compile_store() {
636        let context = create_test_context();
637        let location = Location::unknown(&context);
638
639        compile_operation("store", &context, |block| {
640            let memref = block.append_operation(alloca(
641                &context,
642                MemRefType::new(Type::index(&context), &[], None, None),
643                &[],
644                &[],
645                None,
646                location,
647            ));
648
649            let value = block.append_operation(index::constant(
650                &context,
651                IntegerAttribute::new(Type::index(&context), 42),
652                location,
653            ));
654
655            block.append_operation(store(
656                value.result(0).unwrap().into(),
657                memref.result(0).unwrap().into(),
658                &[],
659                location,
660            ));
661        })
662    }
663
664    #[test]
665    fn compile_store_with_index() {
666        let context = create_test_context();
667        let location = Location::unknown(&context);
668
669        compile_operation("store_with_index", &context, |block| {
670            let memref = block.append_operation(alloca(
671                &context,
672                MemRefType::new(Type::index(&context), &[1], None, None),
673                &[],
674                &[],
675                None,
676                location,
677            ));
678
679            let value = block.append_operation(index::constant(
680                &context,
681                IntegerAttribute::new(Type::index(&context), 42),
682                location,
683            ));
684
685            let index = block.append_operation(index::constant(
686                &context,
687                IntegerAttribute::new(Type::index(&context), 0),
688                location,
689            ));
690
691            block.append_operation(store(
692                value.result(0).unwrap().into(),
693                memref.result(0).unwrap().into(),
694                &[index.result(0).unwrap().into()],
695                location,
696            ));
697        })
698    }
699
700    #[test]
701    fn compile_view() {
702        let context = create_test_context();
703        let location = Location::unknown(&context);
704
705        compile_operation("view", &context, |block| {
706            let byte_type = IntegerType::new(&context, 8).into();
707
708            let memref = block.append_operation(alloc(
709                &context,
710                MemRefType::new(byte_type, &[8], None, None),
711                &[],
712                &[],
713                None,
714                location,
715            ));
716
717            let byte_shift = block.append_operation(index::constant(
718                &context,
719                IntegerAttribute::new(Type::index(&context), 0),
720                location,
721            ));
722
723            block.append_operation(view(
724                &context,
725                memref.result(0).unwrap().into(),
726                byte_shift.result(0).unwrap().into(),
727                &[],
728                MemRefType::new(byte_type, &[1], None, None),
729                location,
730            ));
731        });
732    }
733
734    #[test]
735    fn compile_subview() {
736        let context = create_test_context();
737        let location = Location::unknown(&context);
738
739        compile_operation("subview", &context, |block| {
740            let memref = block.append_operation(alloc(
741                &context,
742                MemRefType::new(Type::index(&context), &[8, 8], None, None),
743                &[],
744                &[],
745                None,
746                location,
747            ));
748
749            block.append_operation(subview(
750                &context,
751                memref.result(0).unwrap().into(),
752                &[],
753                &[],
754                &[],
755                &[0, 0],
756                &[4, 4],
757                &[1, 1],
758                MemRefType::new(
759                    Type::index(&context),
760                    &[4, 4],
761                    // TODO Construct a strided attribute using the C API without parsing.
762                    Some(Attribute::parse(&context, "strided<[8, 1]>").unwrap()),
763                    None,
764                ),
765                location,
766            ));
767        });
768    }
769}
melior/dialect/memref.rs

melior/dialect/
memref.rs
