lift_opt/
quantisation_pass.rs

1use lift_core::context::Context;
2use lift_core::pass::{AnalysisCache, Pass, PassResult};
3
4/// Quantisation pass: inserts quantize/dequantize pairs around compute-heavy ops
5/// to reduce memory footprint and accelerate inference.
6///
7/// Supported modes:
8/// - Dynamic: insert Q/DQ around MatMul and Linear
9/// - Static: insert Q/DQ with pre-computed scales from calibration
10#[derive(Debug)]
11pub struct QuantisationPass {
12    pub target_dtype: QuantTarget,
13    pub mode: QuantMode,
14}
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum QuantTarget {
18    Int8,
19    Int4,
20    Fp8E4M3,
21    Fp8E5M2,
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum QuantMode {
26    Dynamic,
27    Static,
28}
29
30impl Default for QuantisationPass {
31    fn default() -> Self {
32        Self {
33            target_dtype: QuantTarget::Int8,
34            mode: QuantMode::Dynamic,
35        }
36    }
37}
38
39impl Pass for QuantisationPass {
40    fn name(&self) -> &str {
41        "quantisation"
42    }
43
44    fn run(&self, ctx: &mut Context, _cache: &mut AnalysisCache) -> PassResult {
45        let mut quantised = 0usize;
46
47        // Find ops that benefit from quantisation
48        let op_keys: Vec<_> = ctx.ops.keys().collect();
49        let target_ops: Vec<_> = op_keys
50            .into_iter()
51            .filter(|&ok| {
52                if let Some(op) = ctx.ops.get(ok) {
53                    let name = ctx.strings.resolve(op.name);
54                    // Target compute-heavy ops
55                    matches!(
56                        name,
57                        "tensor.matmul"
58                            | "tensor.linear"
59                            | "tensor.conv2d"
60                            | "tensor.conv1d"
61                            | "tensor.multi_head_attention"
62                            | "tensor.attention"
63                    )
64                } else {
65                    false
66                }
67            })
68            .collect();
69
70        let _quant_name = match self.target_dtype {
71            QuantTarget::Int8 => "tensor.quantize",
72            QuantTarget::Int4 => "tensor.quantize_int4",
73            QuantTarget::Fp8E4M3 | QuantTarget::Fp8E5M2 => "tensor.quantize_fp8",
74        };
75        let _dequant_name = match self.target_dtype {
76            QuantTarget::Int8 => "tensor.dequantize",
77            QuantTarget::Int4 => "tensor.dequantize_int4",
78            QuantTarget::Fp8E4M3 | QuantTarget::Fp8E5M2 => "tensor.dequantize_fp8",
79        };
80
81        // Mark ops for quantisation via attributes
82        for op_key in target_ops {
83            if let Some(op) = ctx.ops.get_mut(op_key) {
84                // Skip already quantised ops
85                if op.attrs.get_bool("quantised").unwrap_or(false) {
86                    continue;
87                }
88
89                op.attrs
90                    .set("quantised", lift_core::attributes::Attribute::Bool(true));
91                op.attrs.set(
92                    "quant_method",
93                    lift_core::attributes::Attribute::Integer(match self.target_dtype {
94                        QuantTarget::Int8 => 8,
95                        QuantTarget::Int4 => 4,
96                        QuantTarget::Fp8E4M3 => 83,
97                        QuantTarget::Fp8E5M2 => 82,
98                    }),
99                );
100                op.attrs.set(
101                    "quant_bits",
102                    lift_core::attributes::Attribute::Integer(match self.target_dtype {
103                        QuantTarget::Int8 | QuantTarget::Fp8E4M3 | QuantTarget::Fp8E5M2 => 8,
104                        QuantTarget::Int4 => 4,
105                    }),
106                );
107
108                quantised += 1;
109            }
110        }
111
112        if quantised > 0 {
113            tracing::info!(
114                pass = "quantisation",
115                quantised = quantised,
116                target = ?self.target_dtype,
117                mode = ?self.mode,
118                "Quantisation annotations applied"
119            );
120            PassResult::Changed
121        } else {
122            PassResult::Unchanged
123        }
124    }
125
126    fn invalidates(&self) -> Vec<&str> {
127        vec!["analysis"]
128    }
129}
lift_opt/quantisation_pass.rs

lift_opt/
quantisation_pass.rs
