lift-opt 0.4.4

LIFT compiler pass manager: 13 optimisation passes — constant folding, DCE, CSE, tensor fusion, gate cancellation, qubit routing
Documentation 
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
132
133
134
135
136
137
138
139
140

use lift_core::attributes::Attribute;
use lift_core::context::Context;
use lift_core::pass::{AnalysisCache, Pass, PassResult};

#[derive(Debug)]
pub struct Canonicalize;

impl Pass for Canonicalize {
    fn name(&self) -> &str {
        "canonicalize"
    }

    fn run(&self, ctx: &mut Context, _cache: &mut AnalysisCache) -> PassResult {
        let mut changed = false;

        let op_keys: Vec<_> = ctx.ops.keys().collect();

        for op_key in op_keys {
            let op = match ctx.ops.get(op_key) {
                Some(op) => op,
                None => continue,
            };

            let op_name = ctx.strings.resolve(op.name).to_string();

            match op_name.as_str() {
                // x + 0 -> x
                "tensor.add" => {
                    if op.inputs.len() == 2 {
                        if is_zero_constant(ctx, op.inputs[1]) {
                            let input = op.inputs[0];
                            let result = op.results[0];
                            rewire_value(ctx, result, input, op_key);
                            changed = true;
                        } else if is_zero_constant(ctx, op.inputs[0]) {
                            let input = op.inputs[1];
                            let result = op.results[0];
                            rewire_value(ctx, result, input, op_key);
                            changed = true;
                        }
                    }
                }
                // x * 1 -> x
                "tensor.mul" => {
                    if op.inputs.len() == 2 {
                        if is_one_constant(ctx, op.inputs[1]) {
                            let input = op.inputs[0];
                            let result = op.results[0];
                            rewire_value(ctx, result, input, op_key);
                            changed = true;
                        } else if is_one_constant(ctx, op.inputs[0]) {
                            let input = op.inputs[1];
                            let result = op.results[0];
                            rewire_value(ctx, result, input, op_key);
                            changed = true;
                        }
                    }
                }
                // reshape(reshape(x)) -> reshape(x)
                "tensor.reshape" if op.inputs.len() == 1 => {
                    if let Some(val) = ctx.get_value(op.inputs[0]) {
                        if let lift_core::values::DefSite::OpResult { op: prev_op, .. } = &val.def {
                            if let Some(prev) = ctx.get_op(*prev_op) {
                                let prev_name = ctx.strings.resolve(prev.name);
                                if prev_name == "tensor.reshape" && !prev.inputs.is_empty() {
                                    let original_input = prev.inputs[0];
                                    if let Some(op_mut) = ctx.ops.get_mut(op_key) {
                                        op_mut.inputs = vec![original_input];
                                        changed = true;
                                    }
                                }
                            }
                        }
                    }
                }
                _ => {}
            }
        }

        if changed {
            PassResult::Changed
        } else {
            PassResult::Unchanged
        }
    }
}

fn is_zero_constant(ctx: &Context, val_key: lift_core::values::ValueKey) -> bool {
    if let Some(val) = ctx.get_value(val_key) {
        if let lift_core::values::DefSite::OpResult { op, .. } = &val.def {
            if let Some(op_data) = ctx.get_op(*op) {
                let name = ctx.strings.resolve(op_data.name);
                if name == "core.constant" {
                    return match op_data.attrs.get("value") {
                        Some(Attribute::Integer(0)) => true,
                        Some(Attribute::Float(v)) => *v == 0.0,
                        _ => false,
                    };
                }
            }
        }
    }
    false
}

fn is_one_constant(ctx: &Context, val_key: lift_core::values::ValueKey) -> bool {
    if let Some(val) = ctx.get_value(val_key) {
        if let lift_core::values::DefSite::OpResult { op, .. } = &val.def {
            if let Some(op_data) = ctx.get_op(*op) {
                let name = ctx.strings.resolve(op_data.name);
                if name == "core.constant" {
                    return match op_data.attrs.get("value") {
                        Some(Attribute::Integer(1)) => true,
                        Some(Attribute::Float(v)) => *v == 1.0,
                        _ => false,
                    };
                }
            }
        }
    }
    false
}

fn rewire_value(
    ctx: &mut Context,
    old_val: lift_core::values::ValueKey,
    new_val: lift_core::values::ValueKey,
    _skip_op: lift_core::operations::OpKey,
) {
    let op_keys: Vec<_> = ctx.ops.keys().collect();
    for ok in op_keys {
        if let Some(op) = ctx.ops.get_mut(ok) {
            for input in &mut op.inputs {
                if *input == old_val {
                    *input = new_val;
                }
            }
        }
    }
}