Add MLIR integration with sub-byte data types for mixed-precision quantization

[Copilot]

• Explore the repository to understand existing code generation approach
• Create MLIR integration proposal document (doc/MLIR_INTEGRATION.md)
• Implement MiCoMLIRGen.py - New MLIR code generator module
• Create custom MiCo MLIR dialect definitions for sub-byte quantized operations
• Add example script demonstrating MLIR generation (examples/mlir_example.py)
• Add unit tests for MLIR generator (tests/test_mlir_codegen.py)
• Run tests to validate the implementation
• Update README.md with MLIR documentation
• Run code review and fix issues (fixed MLIR type signatures)
• Run CodeQL security scanner (no alerts)
• Add torch-mlir backend support (MiCoTorchMLIRGen.py)
• Add torch-mlir example and tests
• Update documentation for torch-mlir integration

Summary

This PR adds MLIR (Multi-Level Intermediate Representation) integration to the MiCo framework with two backends:

Backend 1: Standalone MiCoMLIRGen

• Custom MLIR generator with MiCo dialect
• Sub-byte data types (!mico.int<N> for N ∈ {1,2,4,8})
• Quantized operations (mico.bitlinear, mico.bitconv2d, etc.)

Backend 2: MiCoTorchMLIRGen (NEW)

• Uses torch-mlir as first pass for PyTorch → MLIR conversion
• Adds MiCo dialect overlay for sub-byte quantization metadata
• Falls back to standalone mode if torch-mlir is not installed
• Supports output types: torch, linalg, stablehlo

Files Added/Modified

• MiCoMLIRGen.py: Standalone MLIR code generator
• MiCoTorchMLIRGen.py: Torch-mlir based MLIR generator (NEW)
• doc/MLIR_INTEGRATION.md: Comprehensive proposal and dialect specification
• examples/mlir_example.py: Standalone MLIR examples
• examples/torch_mlir_example.py: Torch-mlir backend examples (NEW)
• tests/test_mlir_codegen.py: 19 unit tests for standalone generator
• tests/test_torch_mlir_codegen.py: 10 unit tests for torch-mlir generator (NEW)
• requirements.txt: Added torch-mlir as optional dependency
• readme.md: Updated with both backends

Original prompt

---

This section details on the original issue you should resolve

<issue_title>MLIR Integration</issue_title>
<issue_description>It would be good to use MLIR as a mid-end representation, which is used in many AI compiler systems.

To achieve that, I believe we need to add a new dialect to MLIR and also enable the sub-byte data types in MLIR, then we shall create another codegen or add the feature to the existing codegen.

We will need a proposal for the MLIR Integration, and hopefully some examples.</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes MLIR Integration #28

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

[Jzjerry]

@copilot Current MLIR generator seems to be working well, but less formatted. Can we utilize https://github.com/llvm/torch-mlir directly? Maybe we don't have to write all the codegen ourselves.

[Copilot]

@copilot Current MLIR generator seems to be working well, but less formatted. Can we utilize https://github.com/llvm/torch-mlir directly? Maybe we don't have to write all the codegen ourselves.

Good point! I looked into torch-mlir and here's the trade-off:

torch-mlir advantages:

• Production-ready PyTorch → MLIR lowering via FX/ONNX
• Integrates with downstream compilers (IREE, Blade)
• Active community and maintained infrastructure

Why current custom approach:

• MiCo's sub-byte quantization (1-8 bit types like !mico.int<4>) isn't directly supported in torch-mlir's standard dialects
• Custom ops (mico.bitlinear, mico.bitconv2d) carry mixed-precision metadata (weight_bits, act_bits per layer)
• The generated MLIR is meant to represent the quantization scheme for downstream hardware-aware optimizations

Possible hybrid approach:
We could use torch-mlir as a dependency for the initial PyTorch → Torch dialect conversion, then add a custom lowering pass that transforms to our MiCo dialect with sub-byte types. This would reduce codegen maintenance while keeping MiCo-specific quantization semantics.

Should I explore adding torch-mlir integration as an optional backend? We could keep the current standalone generator as default and add torch-mlir-based generation as an alternative path.

[Jzjerry]

@copilot How about checking the https://github.com/HKUSTGZ-MICS-LYU/MiCo-Lib to see how far we can go?

[Copilot]

@copilot How about checking the https://github.com/HKUSTGZ-MICS-LYU/MiCo-Lib to see how far we can go?

I checked MiCo-Lib and here's the integration path:

Current flow (MiCoCodeGen):

PyTorch → MiCoCodeGen → C header (model.h) → MiCo-Lib (mico_qnn.h kernels)

MLIR integration opportunity:

PyTorch → MiCoMLIRGen → MLIR → MLIR Lowering → C/LLVM → MiCo-Lib kernels

The MLIR operations in MiCoMLIRGen map directly to MiCo-Lib's quantized kernels:

MLIR Op	MiCo-Lib Kernel
mico.bitlinear {weight_bits=4}	MiCo_Q4_MatMul()
mico.bitlinear {weight_bits=2}	MiCo_Q2_MatMul()
mico.bitlinear {weight_bits=1}	MiCo_Q1_MatMul()

MiCo-Lib's packed types (int4x2_t, int2x4_t, int1x8_t) align with our MLIR sub-byte types (!mico.int<4>, etc.).

Next steps to consider:

1. Add MLIR lowering pass: mico.bitlinear → func.call @MiCo_Q{N}_MatMul
2. Generate C code from lowered MLIR (via mlir-translate or custom emitter)
3. Link with MiCo-Lib for execution

This would give us MLIR as the optimization layer while reusing MiCo-Lib's hand-tuned kernels. Want me to prototype this lowering path?

[Jzjerry]

@copilot Great! I still want to share the torch-mlir as one of the first passes to do. In our dialect, we can create those low-bitwidth integer datatypes, or tensor datatypes.

[Jzjerry]

@copilot Now, let's start going? Create a mlir codegen with torch-mlir as a pass.

[Copilot]

@copilot Now, let's start going? Create a mlir codegen with torch-mlir as a pass.

Created MiCoTorchMLIRGen in commit c14eafb.

Usage:

from MiCoTorchMLIRGen import MiCoTorchMLIRGen

mlir_gen = MiCoTorchMLIRGen(model, output_type="torch")
mlir_gen.forward(torch.randn(1, 1, 28, 28))
mlir_gen.convert("output", "lenet_torch_mlir")

Features:

• Uses torch-mlir as first pass (if installed)
• Supports output types: torch, linalg, stablehlo
• Falls back to standalone mode if torch-mlir unavailable
• Adds MiCo dialect overlay for sub-byte quantization metadata

Install torch-mlir:

pip install torch-mlir -f https://github.com/llvm/torch-mlir-release/releases

See examples/torch_mlir_example.py for complete examples.