llzk 0.4.0

Rust bindings to the LLZK C API.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use llzk::{
    builder::OpBuilder,
    dialect::array::ArrayCtor,
    prelude::melior_dialects::arith,
    prelude::*,
    value_ext::{OwningValueRange, ValueRange},
};
use melior::ir::{Location, Type, r#type::FunctionType};

mod common;

#[test]
fn array_new_empty() {
    common::setup();
    let context = LlzkContext::new();
    let location = Location::unknown(&context);
    let module = llzk_module(location);
    let index_type = Type::index(&context);
    let f = dialect::function::def(
        location,
        "array_new",
        FunctionType::new(&context, &[], &[]),
        &[],
        None,
    )
    .unwrap();
    {
        let block = Block::new(&[]);
        let builder = OpBuilder::new(&context);
        let array_type = ArrayType::new(index_type, &[IntegerAttribute::new(index_type, 2).into()]);
        let _array = block.append_operation(dialect::array::new(
            &builder,
            location,
            array_type,
            ArrayCtor::Empty,
        ));
        block.append_operation(dialect::function::r#return(location, &[]));
        f.region(0)
            .expect("function.def must have at least 1 region")
            .append_block(block);
    }

    assert_eq!(f.region_count(), 1);
    let f = module.body().append_operation(f.into());
    assert!(f.verify());
    log::info!("Op passed verification");
    let ir = format!("{f}");
    let expected = concat!(
        "function.def @array_new() {\n",
        "  %array = array.new  : <2 x index> \n",
        "  function.return\n",
        "}"
    );
    assert_eq!(ir, expected);
}

#[test]
fn array_new_affine_map() {
    common::setup();
    let context = LlzkContext::new();
    let location = Location::unknown(&context);
    let module = llzk_module(location);
    let index_type = Type::index(&context);
    let f = dialect::function::def(
        location,
        "array_new",
        FunctionType::new(&context, &[index_type, index_type], &[]),
        &[],
        None,
    )
    .unwrap();
    {
        let block_arg = (index_type, location);
        let block = Block::new(&[block_arg, block_arg]);
        let arg0: Value = block.argument(0).unwrap().into();
        let arg1: Value = block.argument(1).unwrap().into();
        let builder = OpBuilder::new(&context);
        let affine_map = Attribute::parse(&context, "affine_map<()[s0, s1] -> (s0 + s1)>")
            .expect("failed to parse affine_map");
        let array_type = ArrayType::new(index_type, &[affine_map]);
        let owning_value_range = OwningValueRange::from([arg0, arg1].as_slice());
        let value_range = ValueRange::try_from(&owning_value_range).unwrap();
        let _array = block.append_operation(dialect::array::new(
            &builder,
            location,
            array_type,
            ArrayCtor::MapDimSlice(&[value_range], &[0]),
        ));
        block.append_operation(dialect::function::r#return(location, &[]));
        f.region(0)
            .expect("function.def must have at least 1 region")
            .append_block(block);
    }

    assert_eq!(f.region_count(), 1);
    let f = module.body().append_operation(f.into());
    assert!(f.verify());
    log::info!("Op passed verification");
    let ir = format!("{f}");
    let expected = concat!(
        "function.def @array_new(%arg0: index, %arg1: index) {\n",
        "  %array = array.new{()[%arg0, %arg1]} : <affine_map<()[s0, s1] -> (s0 + s1)> x index> \n",
        "  function.return\n",
        "}"
    );
    assert_eq!(ir, expected);
}

#[test]
fn array_len() {
    common::setup();
    let dim = 77;

    let ctx = LlzkContext::new();
    let unknown = Location::unknown(&ctx);
    let index_ty = Type::index(&ctx);
    let ty = ArrayType::new_with_dims(index_ty, &[dim]);
    let op = dialect::array::new(&OpBuilder::new(&ctx), unknown, ty, ArrayCtor::Values(&[]));
    assert_eq!(1, op.result_count(), "op {op} must only have one result");
    let arr_ref = op.result(0).unwrap();
    let arr_dim_op = arith::constant(&ctx, IntegerAttribute::new(index_ty, 0).into(), unknown);
    assert_eq!(
        1,
        arr_dim_op.result_count(),
        "op {arr_dim_op} must only have one result"
    );
    let arr_dim = arr_dim_op.result(0).unwrap();
    let len = dialect::array::len(unknown, arr_ref.into(), arr_dim.into());
    assert!(len.verify(), "op {len} failed to verify");
    assert!(dialect::array::is_array_len(&len));
}