Lightweight Neural Network Quantization Pipeline using MLIR

The project aimed to build a compiler pipeline in MLIR that automatically quantizes floating-point neural network models, eg. FP32 to INT8/INT4, and includes a testing framework to benchmark accuracy vs. performance tradeoffs.

Research Intern: Aaron Gabayoyo 

Project supervisor: Ajitha Rajan, Jackson Woodruff  

PhD Tutor: Eymen Unay 

Quantisation is the process of reducing the precision of the numbers inside a neural network — for example, converting 32-bit floating point numbers to 8-bit integers. This makes models smaller and faster to run, at the cost of a small drop in accuracy. The project aimed to build a compiler pipeline in MLIR that automatically quantizes floating-point neural network models, eg. FP32 to INT8/INT4, and includes a testing framework to benchmark accuracy vs. performance tradeoffs. 

Steps: 

  1. MLIR Dialect & Passes
    1. Define a simple quant dialect or extend the existing mlir::quant dialect 
    2. Write analysis passes to identify quantization-sensitive layers (e.g. skip first/last layers) 
    3. Implement a lowering pass: FP32 ops → quantized integer ops 
  1. Quantization Strategies 
    1. Implement post-training quantization (PTQ) — quantize a pre-trained model without retraining 
    2. Add per-channel vs per-tensor quantization modes as a comparison (time permitting) 
    3. Calibrate scale factors using a small representative dataset 
  1. Testing & Evaluation Framework 
    1. Write unit tests for each MLIR pass using mlir-opt + FileCheck 
    2. Build an end-to-end benchmark: run a small model (e.g. MobileNet) through your pipeline 
    3. Measure and report: model size, inference latency, and accuracy degradation 
  1. Analysis 
    1. Compare INT8 vs INT4 results 
    2. Discuss which layers are most sensitive to quantization 
    3. Reflect on MLIR's strengths/limitations for this task 

This article was published on