diff --git a/MiCoMLIRGen.py b/MiCoMLIRGen.py new file mode 100644 index 0000000..e7b477c --- /dev/null +++ b/MiCoMLIRGen.py @@ -0,0 +1,696 @@ +""" +MiCo MLIR Code Generator + +This module provides MLIR code generation for MiCo mixed-precision quantized models. +It generates MLIR code using a custom MiCo dialect that supports sub-byte data types +and quantized neural network operations. + +Example Usage: + from models import LeNet + from MiCoMLIRGen import MiCoMLIRGen + from MiCoUtils import fuse_model + import torch + + model = LeNet(1) + model.set_qscheme([[8, 6, 6, 4, 4], [8, 8, 8, 8, 8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_gen.convert("output", "lenet_mnist") +""" + +import os +import logging +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.fx + +from MiCoCodeGen import MiCoCodeGen, MiCoTrace +from MiCoQLayers import BitLinear, BitConv1d, BitConv2d, BitQLayer, weight_quant +from MiCoRegistry import MiCoOpRegistry + + +class MiCoMLIRGen(MiCoCodeGen): + """ + MLIR code generator for MiCo models. + + Generates MLIR code using the MiCo dialect for mixed-precision quantized models. + The generated MLIR uses custom operations that represent sub-byte quantized + neural network layers. + + Attributes: + model: The PyTorch model to convert + dialect: The MLIR dialect to use (default: "mico") + mlir_ops: List of MLIR operations generated + mlir_weights: Dictionary of weight tensors with metadata + ssa_counter: Counter for SSA value naming + """ + + MLIR_TEMPLATE = """// MiCo MLIR Module - {model_name} +// Generated by MiCoMLIRGen +// This file uses the MiCo dialect for mixed-precision quantized neural networks + +module @{model_name} {{ + +{weight_declarations} + +{forward_function} + +}} +""" + + def __init__(self, model: nn.Module, dialect: str = "mico", + log_level: int = logging.INFO): + """ + Initialize the MLIR code generator. + + Args: + model: PyTorch model to convert + dialect: MLIR dialect to use (default: "mico") + log_level: Logging level + """ + # Initialize parent class (handles graph extraction) + super().__init__(model, align_to=1, log_level=log_level, gemmini_mode=False) + + self.dialect = dialect + self.mlir_ops: List[str] = [] + self.mlir_weights: Dict[str, Dict[str, Any]] = {} + self.mlir_weight_declarations: List[str] = [] + self.ssa_counter: int = 0 + self.ssa_mapping: Dict[str, str] = {} # Maps node names to SSA values + + self.logger = logging.getLogger("MiCoMLIRGen") + self.logger.setLevel(log_level) + + def reset(self): + """Reset the generator state.""" + super().reset() + self.mlir_ops = [] + self.mlir_weights = {} + self.mlir_weight_declarations = [] + self.ssa_counter = 0 + self.ssa_mapping = {} + self.example_inputs = None + + def _get_ssa_name(self, prefix: str = "v") -> str: + """Generate a new SSA value name.""" + name = f"%{prefix}{self.ssa_counter}" + self.ssa_counter += 1 + return name + + def _register_ssa(self, node_name: str, ssa_name: str): + """Register a mapping from node name to SSA value.""" + self.ssa_mapping[node_name] = ssa_name + + def _get_ssa(self, node_name: str) -> str: + """Get the SSA value for a node name.""" + return self.ssa_mapping.get(node_name, f"%{node_name}") + + def _format_tensor_type(self, tensor: torch.Tensor, qbit: int = 0) -> str: + """ + Format a tensor type for MLIR. + + Args: + tensor: The tensor to get type for + qbit: Quantization bits (0 for full precision) + + Returns: + MLIR type string (e.g., "tensor<1x64x28x28xf32>") + """ + shape = "x".join(str(d) for d in tensor.shape) + + if qbit > 0: + elem_type = f"!mico.int<{qbit}>" + elif tensor.dtype == torch.float32: + elem_type = "f32" + elif tensor.dtype == torch.float16: + elem_type = "f16" + elif tensor.dtype == torch.int8: + elem_type = "i8" + elif tensor.dtype == torch.int32: + elem_type = "i32" + else: + elem_type = "f32" + + return f"tensor<{shape}x{elem_type}>" + + def _format_shape(self, tensor: torch.Tensor) -> str: + """Format tensor shape as MLIR dimension list.""" + return "x".join(str(d) for d in tensor.shape) + + def _format_list(self, values: List, prefix: str = "") -> str: + """Format a list of values for MLIR attributes.""" + formatted = ", ".join(str(v) for v in values) + return f"[{formatted}]" + + def _add_weight_constant(self, name: str, tensor: torch.Tensor, + qbit: int = 0, scale: float = 0.0): + """ + Add a weight constant declaration. + + Args: + name: Name of the weight constant + tensor: Weight tensor + qbit: Quantization bits (0 for full precision) + scale: Quantization scale factor + """ + tensor_type = self._format_tensor_type(tensor, qbit) + + # Store weight metadata + self.mlir_weights[name] = { + "tensor": tensor, + "type": tensor_type, + "qbit": qbit, + "scale": scale + } + + # Generate constant declaration + # For now, use opaque constant with placeholder + if qbit > 0: + decl = f' mico.constant @{name} : {tensor_type} {{scale = {scale:.6f} : f32}}' + else: + decl = f' mico.constant @{name} : {tensor_type}' + + self.mlir_weight_declarations.append(decl) + + def _add_mlir_op(self, op: str): + """Add an MLIR operation to the list.""" + self.mlir_ops.append(f" {op}") + + def handle_placeholder(self, n: torch.fx.node.Node, out: torch.Tensor): + """Handle input placeholder nodes.""" + super().handle_placeholder(n, out) + + # Register input as function argument + ssa_name = f"%{n.name}" + self._register_ssa(n.name, ssa_name) + + self.logger.debug(f"MLIR placeholder: {n.name} -> {ssa_name}") + + def handle_call_module(self, n: torch.fx.node.Node, out: torch.Tensor): + """Handle module calls (Conv2d, Linear, etc.).""" + # Call parent for tensor tracking + super().handle_call_module(n, out) + + module = self.get_module(n.target) + input_node = self.node_info[n.name][0][0] + input_ssa = self._get_ssa(input_node.name) + result_ssa = self._get_ssa_name() + result_type = self._format_tensor_type(out) + + self._register_ssa(n.name, result_ssa) + + if isinstance(module, BitLinear): + self._handle_bitlinear(n, module, input_ssa, result_ssa, result_type, out) + elif isinstance(module, BitConv2d): + self._handle_bitconv2d(n, module, input_ssa, result_ssa, result_type, out) + elif isinstance(module, BitConv1d): + self._handle_bitconv1d(n, module, input_ssa, result_ssa, result_type, out) + elif isinstance(module, nn.Linear): + self._handle_linear(n, module, input_ssa, result_ssa, result_type, out) + elif isinstance(module, nn.Conv2d): + self._handle_conv2d(n, module, input_ssa, result_ssa, result_type, out) + elif isinstance(module, nn.Conv1d): + self._handle_conv1d(n, module, input_ssa, result_ssa, result_type, out) + elif isinstance(module, nn.ReLU): + self._handle_relu(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, nn.ReLU6): + self._handle_relu6(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, nn.MaxPool2d): + self._handle_maxpool2d(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, nn.AvgPool2d): + self._handle_avgpool2d(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, nn.AdaptiveAvgPool2d): + self._handle_adaptive_avgpool2d(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, nn.Flatten): + self._handle_flatten(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, (nn.Identity, nn.Dropout)): + # Pass-through operations + self._register_ssa(n.name, input_ssa) + elif isinstance(module, nn.BatchNorm2d): + self._handle_batchnorm2d(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, nn.Tanh): + self._handle_tanh(n, module, input_ssa, result_ssa, result_type) + elif isinstance(module, nn.ELU): + self._handle_elu(n, module, input_ssa, result_ssa, result_type) + else: + self.logger.warning(f"Unsupported module type: {type(module).__name__}") + # Create a generic operation comment + self._add_mlir_op(f"// Unsupported: {type(module).__name__}") + self._register_ssa(n.name, input_ssa) + + def _handle_bitlinear(self, n, module: BitLinear, input_ssa: str, + result_ssa: str, result_type: str, out: torch.Tensor): + """Handle BitLinear quantized linear layer.""" + weight_name = f"{n.name}_weight" + bias_name = f"{n.name}_bias" + + # Add weight and bias constants + self._add_weight_constant(weight_name, module.weight, + qbit=module.qtype, scale=float(module.qw_scale) if module.qw_scale else 1.0) + if module.bias is not None: + self._add_weight_constant(bias_name, module.bias) + + # Generate operation + attrs = f"weight_bits = {module.qtype} : i32, act_bits = {module.act_q} : i32" + input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type + op = f"{result_ssa} = mico.bitlinear({input_ssa}, @{weight_name}, @{bias_name}) {{{attrs}}} : {input_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_bitconv2d(self, n, module: BitConv2d, input_ssa: str, + result_ssa: str, result_type: str, out: torch.Tensor): + """Handle BitConv2d quantized 2D convolution.""" + weight_name = f"{n.name}_weight" + bias_name = f"{n.name}_bias" + + # Add weight and bias constants + self._add_weight_constant(weight_name, module.weight, + qbit=module.qtype, scale=float(module.qw_scale) if module.qw_scale else 1.0) + if module.bias is not None: + self._add_weight_constant(bias_name, module.bias) + + # Extract convolution parameters + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride, module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding, module.padding] + dilation = module.dilation if isinstance(module.dilation, (list, tuple)) else [module.dilation, module.dilation] + + # Generate operation + attrs = (f"weight_bits = {module.qtype} : i32, act_bits = {module.act_q} : i32, " + f"stride = {self._format_list(stride)}, padding = {self._format_list(padding)}, " + f"dilation = {self._format_list(dilation)}, groups = {module.groups} : i32") + input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type + op = f"{result_ssa} = mico.bitconv2d({input_ssa}, @{weight_name}, @{bias_name}) {{{attrs}}} : {input_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_bitconv1d(self, n, module: BitConv1d, input_ssa: str, + result_ssa: str, result_type: str, out: torch.Tensor): + """Handle BitConv1d quantized 1D convolution.""" + weight_name = f"{n.name}_weight" + bias_name = f"{n.name}_bias" + + # Add weight and bias constants + self._add_weight_constant(weight_name, module.weight, + qbit=module.qtype, scale=float(module.qw_scale) if module.qw_scale else 1.0) + if module.bias is not None: + self._add_weight_constant(bias_name, module.bias) + + # Extract convolution parameters + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding] + dilation = module.dilation if isinstance(module.dilation, (list, tuple)) else [module.dilation] + + # Generate operation + attrs = (f"weight_bits = {module.qtype} : i32, act_bits = {module.act_q} : i32, " + f"stride = {self._format_list(stride)}, padding = {self._format_list(padding)}, " + f"dilation = {self._format_list(dilation)}, groups = {module.groups} : i32") + input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type + op = f"{result_ssa} = mico.bitconv1d({input_ssa}, @{weight_name}, @{bias_name}) {{{attrs}}} : {input_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_linear(self, n, module: nn.Linear, input_ssa: str, + result_ssa: str, result_type: str, out: torch.Tensor): + """Handle standard Linear layer.""" + weight_name = f"{n.name}_weight" + bias_name = f"{n.name}_bias" + + self._add_weight_constant(weight_name, module.weight) + if module.bias is not None: + self._add_weight_constant(bias_name, module.bias) + + input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type + op = f"{result_ssa} = mico.linear({input_ssa}, @{weight_name}, @{bias_name}) : {input_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_conv2d(self, n, module: nn.Conv2d, input_ssa: str, + result_ssa: str, result_type: str, out: torch.Tensor): + """Handle standard Conv2d layer.""" + weight_name = f"{n.name}_weight" + bias_name = f"{n.name}_bias" + + self._add_weight_constant(weight_name, module.weight) + if module.bias is not None: + self._add_weight_constant(bias_name, module.bias) + + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride, module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding, module.padding] + dilation = module.dilation if isinstance(module.dilation, (list, tuple)) else [module.dilation, module.dilation] + + attrs = (f"stride = {self._format_list(stride)}, padding = {self._format_list(padding)}, " + f"dilation = {self._format_list(dilation)}, groups = {module.groups} : i32") + input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type + op = f"{result_ssa} = mico.conv2d({input_ssa}, @{weight_name}, @{bias_name}) {{{attrs}}} : {input_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_conv1d(self, n, module: nn.Conv1d, input_ssa: str, + result_ssa: str, result_type: str, out: torch.Tensor): + """Handle standard Conv1d layer.""" + weight_name = f"{n.name}_weight" + bias_name = f"{n.name}_bias" + + self._add_weight_constant(weight_name, module.weight) + if module.bias is not None: + self._add_weight_constant(bias_name, module.bias) + + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding] + dilation = module.dilation if isinstance(module.dilation, (list, tuple)) else [module.dilation] + + attrs = (f"stride = {self._format_list(stride)}, padding = {self._format_list(padding)}, " + f"dilation = {self._format_list(dilation)}, groups = {module.groups} : i32") + input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type + op = f"{result_ssa} = mico.conv1d({input_ssa}, @{weight_name}, @{bias_name}) {{{attrs}}} : {input_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_relu(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle ReLU activation.""" + op = f"{result_ssa} = mico.relu({input_ssa}) : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_relu6(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle ReLU6 activation.""" + op = f"{result_ssa} = mico.relu6({input_ssa}) : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_tanh(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle Tanh activation.""" + op = f"{result_ssa} = mico.tanh({input_ssa}) : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_elu(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle ELU activation.""" + op = f"{result_ssa} = mico.elu({input_ssa}) {{alpha = {module.alpha} : f32}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_maxpool2d(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle MaxPool2d.""" + kernel_size = module.kernel_size if isinstance(module.kernel_size, (list, tuple)) else [module.kernel_size, module.kernel_size] + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride, module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding, module.padding] + + attrs = f"kernel_size = {self._format_list(kernel_size)}, stride = {self._format_list(stride)}, padding = {self._format_list(padding)}" + op = f"{result_ssa} = mico.maxpool2d({input_ssa}) {{{attrs}}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_avgpool2d(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle AvgPool2d.""" + kernel_size = module.kernel_size if isinstance(module.kernel_size, (list, tuple)) else [module.kernel_size, module.kernel_size] + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride, module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding, module.padding] + + attrs = f"kernel_size = {self._format_list(kernel_size)}, stride = {self._format_list(stride)}, padding = {self._format_list(padding)}" + op = f"{result_ssa} = mico.avgpool2d({input_ssa}) {{{attrs}}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_adaptive_avgpool2d(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle AdaptiveAvgPool2d.""" + output_size = module.output_size if isinstance(module.output_size, (list, tuple)) else [module.output_size, module.output_size] + + attrs = f"output_size = {self._format_list(output_size)}" + op = f"{result_ssa} = mico.adaptive_avgpool2d({input_ssa}) {{{attrs}}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_flatten(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle Flatten operation.""" + start_dim = module.start_dim if hasattr(module, 'start_dim') else 1 + op = f"{result_ssa} = mico.flatten({input_ssa}) {{start_dim = {start_dim} : i32}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_batchnorm2d(self, n, module, input_ssa: str, result_ssa: str, result_type: str): + """Handle BatchNorm2d.""" + weight_name = f"{n.name}_weight" + bias_name = f"{n.name}_bias" + mean_name = f"{n.name}_running_mean" + var_name = f"{n.name}_running_var" + + self._add_weight_constant(weight_name, module.weight) + self._add_weight_constant(bias_name, module.bias) + self._add_weight_constant(mean_name, module.running_mean) + self._add_weight_constant(var_name, module.running_var) + + attrs = f"eps = {module.eps} : f32" + op = f"{result_ssa} = mico.batchnorm2d({input_ssa}, @{weight_name}, @{bias_name}, @{mean_name}, @{var_name}) {{{attrs}}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def handle_call_function(self, n: torch.fx.node.Node, out: torch.Tensor): + """Handle function calls (torch.add, F.relu, etc.).""" + super().handle_call_function(n, out) + + function = n.target + input_names = self._extract_input_names(n) + input_args = n.args + + result_ssa = self._get_ssa_name() + result_type = self._format_tensor_type(out) + self._register_ssa(n.name, result_ssa) + + if function in (torch.add, __import__('operator').__add__): + self._handle_add_function(n, input_names, result_ssa, result_type) + elif function in (torch.nn.functional.relu,): + input_ssa = self._get_ssa(input_names[0]) + op = f"{result_ssa} = mico.relu({input_ssa}) : {result_type} -> {result_type}" + self._add_mlir_op(op) + elif function in (torch.nn.functional.relu6,): + input_ssa = self._get_ssa(input_names[0]) + op = f"{result_ssa} = mico.relu6({input_ssa}) : {result_type} -> {result_type}" + self._add_mlir_op(op) + elif function in (torch.flatten,): + input_ssa = self._get_ssa(input_names[0]) + start_dim = input_args[1] if len(input_args) > 1 else 1 + op = f"{result_ssa} = mico.flatten({input_ssa}) {{start_dim = {start_dim} : i32}} : {result_type}" + self._add_mlir_op(op) + elif function in (torch.cat,): + input_ssas = [self._get_ssa(name) for name in input_names] + dim = input_args[1] if len(input_args) > 1 else 0 + inputs_str = ", ".join(input_ssas) + op = f"{result_ssa} = mico.concat({inputs_str}) {{dim = {dim} : i32}} : {result_type}" + self._add_mlir_op(op) + elif function in (torch.nn.functional.max_pool2d,): + self._handle_maxpool2d_function(n, input_names, input_args, result_ssa, result_type) + elif function in (torch.nn.functional.avg_pool2d,): + self._handle_avgpool2d_function(n, input_names, input_args, result_ssa, result_type) + elif function in (torch.nn.functional.adaptive_avg_pool2d,): + self._handle_adaptive_avgpool2d_function(n, input_names, input_args, result_ssa, result_type) + else: + # Fallback: log warning and pass through + func_name = getattr(function, '__name__', str(function)) + self.logger.warning(f"Unsupported function: {func_name}") + if input_names: + self._register_ssa(n.name, self._get_ssa(input_names[0])) + self._add_mlir_op(f"// Unsupported function: {func_name}") + + def _handle_add_function(self, n, input_names, result_ssa, result_type): + """Handle element-wise addition.""" + lhs_ssa = self._get_ssa(input_names[0]) + rhs_ssa = self._get_ssa(input_names[1]) if len(input_names) > 1 else lhs_ssa + op = f"{result_ssa} = mico.add({lhs_ssa}, {rhs_ssa}) : {result_type}, {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_maxpool2d_function(self, n, input_names, input_args, result_ssa, result_type): + """Handle max_pool2d function.""" + input_ssa = self._get_ssa(input_names[0]) + kernel_size = input_args[1] if len(input_args) > 1 else [2, 2] + stride = input_args[2] if len(input_args) > 2 else kernel_size + padding = input_args[3] if len(input_args) > 3 else [0, 0] + + if isinstance(kernel_size, int): + kernel_size = [kernel_size, kernel_size] + if isinstance(stride, int): + stride = [stride, stride] + if isinstance(padding, int): + padding = [padding, padding] + + attrs = f"kernel_size = {self._format_list(kernel_size)}, stride = {self._format_list(stride)}, padding = {self._format_list(padding)}" + op = f"{result_ssa} = mico.maxpool2d({input_ssa}) {{{attrs}}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_avgpool2d_function(self, n, input_names, input_args, result_ssa, result_type): + """Handle avg_pool2d function.""" + input_ssa = self._get_ssa(input_names[0]) + kernel_size = input_args[1] if len(input_args) > 1 else [2, 2] + stride = input_args[2] if len(input_args) > 2 else kernel_size + padding = input_args[3] if len(input_args) > 3 else [0, 0] + + if isinstance(kernel_size, int): + kernel_size = [kernel_size, kernel_size] + if isinstance(stride, int): + stride = [stride, stride] + if isinstance(padding, int): + padding = [padding, padding] + + attrs = f"kernel_size = {self._format_list(kernel_size)}, stride = {self._format_list(stride)}, padding = {self._format_list(padding)}" + op = f"{result_ssa} = mico.avgpool2d({input_ssa}) {{{attrs}}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def _handle_adaptive_avgpool2d_function(self, n, input_names, input_args, result_ssa, result_type): + """Handle adaptive_avg_pool2d function.""" + input_ssa = self._get_ssa(input_names[0]) + output_size = input_args[1] if len(input_args) > 1 else [1, 1] + + if isinstance(output_size, int): + output_size = [output_size, output_size] + + attrs = f"output_size = {self._format_list(output_size)}" + op = f"{result_ssa} = mico.adaptive_avgpool2d({input_ssa}) {{{attrs}}} : {result_type} -> {result_type}" + self._add_mlir_op(op) + + def handle_output(self, n: torch.fx.node.Node, out: torch.Tensor): + """Handle output node.""" + super().handle_output(n, out) + + # Get the input to the output node + if n.args and hasattr(n.args[0], 'name'): + output_ssa = self._get_ssa(n.args[0].name) + result_type = self._format_tensor_type(out) + self._add_mlir_op(f"return {output_ssa} : {result_type}") + + def forward(self, *args): + """ + Run forward pass to trace the model and collect MLIR operations. + + Args: + *args: Input tensors + + Returns: + Output tensor from the model + """ + self.reset() + self.example_inputs = args + return self.run(*args) + + def convert(self, output_directory: str = "output", + model_name: str = "model", + verbose: bool = False) -> str: + """ + Convert the traced model to MLIR code. + + Args: + output_directory: Directory to write output files + model_name: Name for the generated model + verbose: Enable verbose output + + Returns: + Path to the generated MLIR file + """ + if self.example_inputs is None: + raise ValueError("No example inputs provided. Please call forward() first.") + + # Create output directory + os.makedirs(output_directory, exist_ok=True) + + # Determine input type for function signature + if len(self.example_inputs) > 0: + input_tensor = self.example_inputs[0] + input_type = self._format_tensor_type(input_tensor) + input_name = "input" + else: + input_type = "tensor<1x1xf32>" + input_name = "input" + + # Get output type from last operation + output_type = "tensor<1x10xf32>" # Default + for op in reversed(self.mlir_ops): + if "return" in op: + # Extract return type + parts = op.split(":") + if len(parts) > 1: + output_type = parts[-1].strip() + break + + # Build forward function + func_signature = f" func.func @forward(%{input_name}: {input_type}) -> {output_type} {{" + operations = "\n".join(self.mlir_ops) + func_end = " }" + forward_function = f"{func_signature}\n{operations}\n{func_end}" + + # Build weight declarations + weight_declarations = "\n".join(self.mlir_weight_declarations) + + # Generate complete MLIR code + mlir_code = MiCoMLIRGen.MLIR_TEMPLATE.format( + model_name=model_name, + weight_declarations=weight_declarations, + forward_function=forward_function + ) + + # Write to file + mlir_path = os.path.join(output_directory, f"{model_name}.mlir") + with open(mlir_path, "w") as f: + f.write(mlir_code) + + if verbose: + print(f"Generated MLIR code:\n{mlir_code}") + + print(f"MLIR code written to: {mlir_path}") + print(f"Number of operations: {len(self.mlir_ops)}") + print(f"Number of weight constants: {len(self.mlir_weights)}") + + return mlir_path + + def get_mlir_code(self) -> str: + """ + Get the generated MLIR code as a string without writing to file. + + Returns: + MLIR code as a string + """ + if len(self.mlir_ops) == 0: + raise ValueError("No operations generated. Please call forward() first.") + + # Simplified version for inspection + weight_declarations = "\n".join(self.mlir_weight_declarations) + operations = "\n".join(self.mlir_ops) + + return f"// Weight declarations:\n{weight_declarations}\n\n// Operations:\n{operations}" + + +if __name__ == "__main__": + # Example usage + import sys + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + + from models import MLP, LeNet + from MiCoUtils import fuse_model + + print("=" * 60) + print("MiCo MLIR Code Generator Example") + print("=" * 60) + + # Example 1: MLP + print("\n--- Example 1: MLP ---") + mlp = MLP(in_features=64, config={"Layers": [32, 16, 10]}) + weight_q = [8, 6, 4] # Mixed precision + activation_q = [8, 8, 8] + mlp.set_qscheme([weight_q, activation_q]) + mlp = fuse_model(mlp) + mlp.eval() + + mlir_gen = MiCoMLIRGen(mlp) + mlir_gen.forward(torch.randn(1, 64)) + mlir_path = mlir_gen.convert("output", "mlp_mnist_mlir", verbose=False) + + print(f"\nGenerated MLIR at: {mlir_path}") + print("\nMLIR Code Preview:") + print(mlir_gen.get_mlir_code()) + + # Example 2: LeNet + print("\n--- Example 2: LeNet ---") + lenet = LeNet(1) + weight_q = [8, 6, 6, 4, 4] # Mixed precision for conv1, conv2, fc1, fc2, fc3 + activation_q = [8, 8, 8, 8, 8] + lenet.set_qscheme([weight_q, activation_q]) + lenet = fuse_model(lenet) + lenet.eval() + + mlir_gen = MiCoMLIRGen(lenet) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_path = mlir_gen.convert("output", "lenet_mnist_mlir", verbose=False) + + print(f"\nGenerated MLIR at: {mlir_path}") + + print("\n" + "=" * 60) + print("MLIR Code Generation Complete!") + print("=" * 60) diff --git a/MiCoTorchMLIRGen.py b/MiCoTorchMLIRGen.py new file mode 100644 index 0000000..ecb9117 --- /dev/null +++ b/MiCoTorchMLIRGen.py @@ -0,0 +1,446 @@ +""" +MiCo Torch-MLIR Code Generator + +This module provides MLIR code generation for MiCo mixed-precision quantized models +using torch-mlir as the first pass for PyTorch to MLIR conversion, followed by +custom MiCo dialect lowering for sub-byte quantization support. + +The pipeline is: + PyTorch Model → torch-mlir (Torch dialect) → MiCo dialect (sub-byte types) + +Example Usage: + from models import LeNet + from MiCoTorchMLIRGen import MiCoTorchMLIRGen + from MiCoUtils import fuse_model + import torch + + model = LeNet(1) + model.set_qscheme([[8, 6, 6, 4, 4], [8, 8, 8, 8, 8]]) + model = fuse_model(model) + model.eval() + + # Use torch-mlir backend + mlir_gen = MiCoTorchMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_gen.convert("output", "lenet_mnist") +""" + +import os +import logging +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.fx + +# Check if torch-mlir is available +TORCH_MLIR_AVAILABLE = False +try: + import torch_mlir + from torch_mlir import torchscript + TORCH_MLIR_AVAILABLE = True +except ImportError: + pass + +from MiCoCodeGen import MiCoCodeGen, MiCoTrace +from MiCoMLIRGen import MiCoMLIRGen +from MiCoQLayers import BitLinear, BitConv1d, BitConv2d, BitQLayer, weight_quant + + +class MiCoTorchMLIRGen(MiCoMLIRGen): + """ + MLIR code generator using torch-mlir as the first pass. + + This generator uses torch-mlir to convert PyTorch models to MLIR's Torch dialect, + then applies custom lowering to the MiCo dialect for sub-byte quantization. + + If torch-mlir is not available, falls back to the standalone MiCoMLIRGen. + + Attributes: + model: The PyTorch model to convert + use_torch_mlir: Whether to use torch-mlir (True if available) + output_type: The torch-mlir output type ("torch", "linalg", "stablehlo") + """ + + TORCH_MLIR_TEMPLATE = """// MiCo MLIR Module - {model_name} +// Generated by MiCoTorchMLIRGen with torch-mlir backend +// Pipeline: PyTorch → torch-mlir (Torch dialect) → MiCo dialect + +// ============================================================================ +// Torch-MLIR Generated IR (Torch Dialect) +// ============================================================================ +{torch_mlir_ir} + +// ============================================================================ +// MiCo Dialect Overlay (Sub-byte Quantization) +// ============================================================================ +// The following module contains MiCo-specific quantization metadata that +// should be used in conjunction with the Torch dialect IR above. + +module @{model_name}_mico {{ + +{weight_declarations} + +{mico_ops} + +}} +""" + + MICO_STANDALONE_TEMPLATE = """// MiCo MLIR Module - {model_name} +// Generated by MiCoTorchMLIRGen (standalone mode - torch-mlir not available) +// This file uses the MiCo dialect for mixed-precision quantized neural networks + +module @{model_name} {{ + +{weight_declarations} + +{forward_function} + +}} +""" + + def __init__(self, model: nn.Module, dialect: str = "mico", + output_type: str = "torch", log_level: int = logging.INFO): + """ + Initialize the Torch-MLIR code generator. + + Args: + model: PyTorch model to convert + dialect: MLIR dialect to use for MiCo ops (default: "mico") + output_type: torch-mlir output type: "torch", "linalg", or "stablehlo" + log_level: Logging level + """ + super().__init__(model, dialect=dialect, log_level=log_level) + + self.use_torch_mlir = TORCH_MLIR_AVAILABLE + self.output_type = output_type + self.torch_mlir_ir: Optional[str] = None + self.mico_quantization_ops: List[str] = [] + + self.logger = logging.getLogger("MiCoTorchMLIRGen") + self.logger.setLevel(log_level) + + if not TORCH_MLIR_AVAILABLE: + self.logger.warning( + "torch-mlir not available. Install with: " + "pip install torch-mlir -f https://github.com/llvm/torch-mlir-release/releases" + ) + self.logger.info("Falling back to standalone MiCoMLIRGen mode") + + def reset(self): + """Reset the generator state.""" + super().reset() + self.torch_mlir_ir = None + self.mico_quantization_ops = [] + + def _convert_with_torch_mlir(self, example_inputs) -> str: + """ + Convert the model using torch-mlir. + + Args: + example_inputs: Example input tensors for tracing (tuple or single tensor) + + Returns: + MLIR string in Torch dialect + """ + if not TORCH_MLIR_AVAILABLE: + raise RuntimeError("torch-mlir is not available") + + # Get the output type enum + if self.output_type == "torch": + output_type = torchscript.OutputType.TORCH + elif self.output_type == "linalg": + output_type = torchscript.OutputType.LINALG_ON_TENSORS + elif self.output_type == "stablehlo": + output_type = torchscript.OutputType.STABLEHLO + else: + output_type = torchscript.OutputType.TORCH + + # Create a wrapper model that doesn't have MiCo-specific quantization + # torch-mlir expects standard PyTorch operations + class ModelWrapper(nn.Module): + def __init__(self, original_model): + super().__init__() + self.model = original_model + + def forward(self, x): + return self.model(x) + + wrapped_model = ModelWrapper(self.model) + wrapped_model.eval() + + try: + # Compile with torch-mlir + module = torchscript.compile( + wrapped_model, + example_inputs, + output_type=output_type, + use_tracing=True + ) + + # Get the MLIR string representation + mlir_str = str(module) + return mlir_str + + except Exception as e: + self.logger.warning(f"torch-mlir compilation failed: {e}") + self.logger.info("Falling back to standalone mode") + self.use_torch_mlir = False + return "" + + def _generate_mico_quantization_overlay(self) -> str: + """ + Generate MiCo dialect operations for quantization metadata. + + This creates MiCo-specific operations that annotate the torch-mlir + generated IR with sub-byte quantization information. + + Returns: + MLIR string with MiCo quantization operations + """ + ops = [] + + # Add comments explaining the quantization scheme + ops.append(" // Quantization scheme for mixed-precision inference") + ops.append(" // Each layer has specific weight_bits and activation_bits") + ops.append("") + + # Generate MiCo quantization ops based on the model's QLayers + layer_idx = 0 + for name, module in self.model.named_modules(): + if isinstance(module, BitLinear): + safe_name = name.replace(".", "_") + ops.append(f" // Layer: {name}") + ops.append(f" // mico.quantize_linear @{safe_name} {{") + ops.append(f" // weight_bits = {module.qtype} : i32,") + ops.append(f" // act_bits = {module.act_q} : i32,") + if module.qw_scale is not None: + ops.append(f" // scale = {float(module.qw_scale):.6f} : f32") + ops.append(f" // }}") + ops.append("") + layer_idx += 1 + + elif isinstance(module, BitConv2d): + safe_name = name.replace(".", "_") + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride, module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding, module.padding] + + ops.append(f" // Layer: {name}") + ops.append(f" // mico.quantize_conv2d @{safe_name} {{") + ops.append(f" // weight_bits = {module.qtype} : i32,") + ops.append(f" // act_bits = {module.act_q} : i32,") + ops.append(f" // stride = [{stride[0]}, {stride[1]}],") + ops.append(f" // padding = [{padding[0]}, {padding[1]}],") + if module.qw_scale is not None: + ops.append(f" // scale = {float(module.qw_scale):.6f} : f32") + ops.append(f" // }}") + ops.append("") + layer_idx += 1 + + elif isinstance(module, BitConv1d): + safe_name = name.replace(".", "_") + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride] + padding = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding] + + ops.append(f" // Layer: {name}") + ops.append(f" // mico.quantize_conv1d @{safe_name} {{") + ops.append(f" // weight_bits = {module.qtype} : i32,") + ops.append(f" // act_bits = {module.act_q} : i32,") + ops.append(f" // stride = [{stride[0]}],") + ops.append(f" // padding = [{padding[0]}],") + if module.qw_scale is not None: + ops.append(f" // scale = {float(module.qw_scale):.6f} : f32") + ops.append(f" // }}") + ops.append("") + layer_idx += 1 + + return "\n".join(ops) + + def _generate_mico_type_definitions(self) -> str: + """ + Generate MiCo sub-byte type definitions. + + Returns: + MLIR string with MiCo type definitions + """ + types = [] + types.append(" // MiCo Sub-byte Integer Types") + types.append(" // !mico.int where N ∈ {1, 2, 4, 8}") + types.append("") + + # Collect all unique bit widths used in the model + bit_widths = set() + for name, module in self.model.named_modules(): + if isinstance(module, (BitLinear, BitConv1d, BitConv2d)): + bit_widths.add(module.qtype) + bit_widths.add(module.act_q) + + for bw in sorted(bit_widths): + types.append(f" // !mico.int<{bw}> - {bw}-bit quantized integer") + + return "\n".join(types) + + def convert(self, output_directory: str = "output", + model_name: str = "model", + verbose: bool = False) -> str: + """ + Convert the traced model to MLIR code. + + Args: + output_directory: Directory to write output files + model_name: Name for the generated model + verbose: Enable verbose output + + Returns: + Path to the generated MLIR file + """ + if self.example_inputs is None: + raise ValueError("No example inputs provided. Please call forward() first.") + + # Create output directory + os.makedirs(output_directory, exist_ok=True) + + if self.use_torch_mlir: + # Use torch-mlir for conversion + self.logger.info("Using torch-mlir backend for MLIR generation") + + try: + # Convert with torch-mlir + self.torch_mlir_ir = self._convert_with_torch_mlir(self.example_inputs) + + if self.torch_mlir_ir: + # Generate MiCo quantization overlay + mico_type_defs = self._generate_mico_type_definitions() + mico_quant_ops = self._generate_mico_quantization_overlay() + + # Combine weight declarations + weight_declarations = "\n".join(self.mlir_weight_declarations) if self.mlir_weight_declarations else " // No weight constants (weights embedded in torch-mlir IR)" + + # Generate combined output + mlir_code = self.TORCH_MLIR_TEMPLATE.format( + model_name=model_name, + torch_mlir_ir=self.torch_mlir_ir, + weight_declarations=mico_type_defs + "\n\n" + weight_declarations, + mico_ops=mico_quant_ops + ) + else: + # Fallback to standalone mode + self.use_torch_mlir = False + + except Exception as e: + self.logger.warning(f"torch-mlir conversion failed: {e}") + self.use_torch_mlir = False + + if not self.use_torch_mlir: + # Use standalone MiCoMLIRGen + self.logger.info("Using standalone MiCoMLIRGen mode") + return super().convert(output_directory, model_name, verbose) + + # Write to file + mlir_path = os.path.join(output_directory, f"{model_name}_torch_mlir.mlir") + with open(mlir_path, "w") as f: + f.write(mlir_code) + + if verbose: + print(f"Generated MLIR code:\n{mlir_code}") + + print(f"MLIR code written to: {mlir_path}") + print(f"Backend: torch-mlir ({self.output_type})") + + return mlir_path + + def get_torch_mlir_ir(self) -> Optional[str]: + """ + Get the torch-mlir generated IR. + + Returns: + The torch-mlir MLIR string, or None if not available + """ + return self.torch_mlir_ir + + def is_torch_mlir_available(self) -> bool: + """ + Check if torch-mlir is available. + + Returns: + True if torch-mlir is available and can be used + """ + return TORCH_MLIR_AVAILABLE + + +def check_torch_mlir_installation() -> Dict[str, Any]: + """ + Check torch-mlir installation status. + + Returns: + Dictionary with installation status and version info + """ + result = { + "available": TORCH_MLIR_AVAILABLE, + "version": None, + "output_types": [], + "install_command": "pip install torch-mlir -f https://github.com/llvm/torch-mlir-release/releases" + } + + if TORCH_MLIR_AVAILABLE: + try: + result["version"] = torch_mlir.__version__ + except AttributeError: + result["version"] = "unknown" + + result["output_types"] = ["torch", "linalg", "stablehlo"] + + return result + + +if __name__ == "__main__": + # Example usage and installation check + import sys + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + + print("=" * 60) + print("MiCo Torch-MLIR Code Generator") + print("=" * 60) + + # Check installation + status = check_torch_mlir_installation() + print(f"\ntorch-mlir available: {status['available']}") + if status['available']: + print(f"torch-mlir version: {status['version']}") + print(f"Supported output types: {status['output_types']}") + else: + print(f"Install with: {status['install_command']}") + + # Try to run an example + print("\n" + "=" * 60) + print("Running Example") + print("=" * 60) + + try: + from models import MLP, LeNet + from MiCoUtils import fuse_model + + # Create a simple model + model = MLP(in_features=64, config={"Layers": [32, 16, 10]}) + weight_q = [8, 6, 4] # Mixed precision + activation_q = [8, 8, 8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + # Generate MLIR + mlir_gen = MiCoTorchMLIRGen(model, output_type="torch") + mlir_gen.forward(torch.randn(1, 64)) + mlir_path = mlir_gen.convert("output", "mlp_torch_mlir", verbose=False) + + print(f"\nGenerated MLIR at: {mlir_path}") + print(f"Backend used: {'torch-mlir' if mlir_gen.use_torch_mlir else 'standalone'}") + + except Exception as e: + print(f"Example failed: {e}") + import traceback + traceback.print_exc() + + print("\n" + "=" * 60) + print("Done!") + print("=" * 60) diff --git a/doc/MLIR_INTEGRATION.md b/doc/MLIR_INTEGRATION.md new file mode 100644 index 0000000..6246f8b --- /dev/null +++ b/doc/MLIR_INTEGRATION.md @@ -0,0 +1,441 @@ +# MLIR Integration Proposal for MiCo Framework + +## Overview + +This document proposes the integration of MLIR (Multi-Level Intermediate Representation) as a mid-end representation for the MiCo framework. MLIR is widely used in modern AI compiler systems and provides a flexible, extensible infrastructure for building domain-specific compilers. + +## Goals + +1. **Unified IR**: Use MLIR as a common intermediate representation between PyTorch models and various backend targets +2. **Sub-byte Data Types**: Enable mixed-precision quantization with sub-byte (1-8 bit) weight and activation representations +3. **Hardware Portability**: Leverage MLIR's infrastructure for targeting multiple hardware backends +4. **Optimization Pipeline**: Enable MLIR-based optimization passes for quantized neural networks + +## Architecture + +MiCo provides two MLIR generation paths: + +### Path 1: Standalone MiCoMLIRGen + +``` +┌─────────────────────┐ +│ PyTorch Model │ +│ (MiCo Framework) │ +└─────────┬───────────┘ + │ MiCoMLIRGen + ▼ +┌─────────────────────┐ +│ MiCo Dialect │ +│ (Sub-byte types, │ +│ Quant operations) │ +└─────────┬───────────┘ + │ Lowering + ▼ +┌─────────────────────┐ +│ Arith + Tensor │ +│ + MemRef │ +└─────────┬───────────┘ + │ Lowering + ▼ +┌─────────────────────┐ +│ LLVM IR │ +│ or Hardware IR │ +└─────────┴───────────┘ +``` + +### Path 2: Torch-MLIR Integration (MiCoTorchMLIRGen) + +``` +┌─────────────────────┐ +│ PyTorch Model │ +│ (MiCo Framework) │ +└─────────┬───────────┘ + │ torch-mlir (first pass) + ▼ +┌─────────────────────┐ +│ Torch Dialect │ +│ (Standard PyTorch │ +│ operations) │ +└─────────┬───────────┘ + │ MiCo quantization overlay + ▼ +┌─────────────────────┐ +│ MiCo Dialect │ +│ (Sub-byte types, │ +│ Quant metadata) │ +└─────────┬───────────┘ + │ Lowering + ▼ +┌─────────────────────┐ +│ Linalg / StableHLO │ +│ + MiCo-Lib │ +└─────────┴───────────┘ +``` + +## Torch-MLIR Integration + +### Installation + +```bash +pip install torch-mlir -f https://github.com/llvm/torch-mlir-release/releases +``` + +### Usage with Torch-MLIR Backend + +```python +from models import LeNet +from MiCoTorchMLIRGen import MiCoTorchMLIRGen +from MiCoUtils import fuse_model +import torch + +model = LeNet(1) +model.set_qscheme([[8, 6, 6, 4, 4], [8, 8, 8, 8, 8]]) +model = fuse_model(model) +model.eval() + +# Use torch-mlir backend with Torch dialect output +mlir_gen = MiCoTorchMLIRGen(model, output_type="torch") +mlir_gen.forward(torch.randn(1, 1, 28, 28)) +mlir_path = mlir_gen.convert("output", "lenet_mnist") + +# Check if torch-mlir was used +print(f"Backend: {'torch-mlir' if mlir_gen.use_torch_mlir else 'standalone'}") +``` + +### Supported Output Types + +| Output Type | Description | +|-------------|-------------| +| `torch` | Torch dialect (default) - closest to PyTorch semantics | +| `linalg` | Linalg on Tensors - good for custom lowering | +| `stablehlo` | StableHLO - compatible with XLA ecosystem | + +### Benefits of Torch-MLIR Integration + +1. **Production-ready lowering**: Leverages torch-mlir's mature PyTorch → MLIR conversion +2. **Ecosystem compatibility**: Integrates with IREE, Blade, and other MLIR-based compilers +3. **Reduced maintenance**: Core operation lowering handled by torch-mlir +4. **MiCo-specific extensions**: Sub-byte quantization metadata preserved via MiCo dialect overlay + +## MiCo MLIR Dialect + +### Custom Types + +The MiCo dialect introduces sub-byte integer types for mixed-precision quantization: + +```mlir +// Sub-byte integer types +!mico.int<1> // 1-bit (binary) +!mico.int<2> // 2-bit (ternary) +!mico.int<4> // 4-bit +!mico.int<8> // 8-bit (standard) + +// Quantized tensor type with scale factor +!mico.qtensor +``` + +### Operations + +#### Quantized Linear Operations + +```mlir +// Quantized matrix multiplication +// Inputs: activation tensor, weight tensor, bias tensor +// Attributes: weight_bits, activation_bits +%result = mico.bitlinear(%activation, %weight, %bias) { + weight_bits = 4 : i32, + act_bits = 8 : i32 +} : (tensor<1x256xf32>, tensor<64x256x!mico.int<4>>, tensor<64xf32>) -> tensor<1x64xf32> + +// Quantized 2D convolution +%result = mico.bitconv2d(%input, %weight, %bias) { + weight_bits = 4 : i32, + act_bits = 8 : i32, + stride = [1, 1], + padding = [1, 1], + dilation = [1, 1], + groups = 1 +} : (tensor<1x32x28x28xf32>, tensor<64x32x3x3x!mico.int<4>>, tensor<64xf32>) -> tensor<1x64x28x28xf32> + +// Quantized 1D convolution +%result = mico.bitconv1d(%input, %weight, %bias) { + weight_bits = 8 : i32, + act_bits = 8 : i32, + stride = [1], + padding = [0], + dilation = [1], + groups = 1 +} : (tensor<1x32x128xf32>, tensor<64x32x3x!mico.int<8>>, tensor<64xf32>) -> tensor<1x64x126xf32> +``` + +#### Quantization Operations + +```mlir +// Weight quantization (per-tensor symmetric) +%qweight, %scale = mico.weight_quant(%weight) { + bits = 4 : i32, + mode = "max" // or "mean" +} : (tensor<64x256xf32>) -> (tensor<64x256x!mico.int<4>>, f32) + +// Activation quantization (per-tensor symmetric) +%qact = mico.activation_quant(%activation) { + bits = 8 : i32 +} : (tensor<1x256xf32>) -> tensor<1x256xf32> +``` + +#### Standard Operations + +```mlir +// ReLU activation +%result = mico.relu(%input) : (tensor<1x64x28x28xf32>) -> tensor<1x64x28x28xf32> + +// ReLU6 activation +%result = mico.relu6(%input) : (tensor<1x64x28x28xf32>) -> tensor<1x64x28x28xf32> + +// Max pooling +%result = mico.maxpool2d(%input) { + kernel_size = [2, 2], + stride = [2, 2] +} : (tensor<1x64x28x28xf32>) -> tensor<1x64x14x14xf32> + +// Average pooling +%result = mico.avgpool2d(%input) { + kernel_size = [2, 2], + stride = [2, 2] +} : (tensor<1x64x28x28xf32>) -> tensor<1x64x14x14xf32> + +// Flatten +%result = mico.flatten(%input) { + start_dim = 1 : i32 +} : (tensor<1x64x7x7xf32>) -> tensor<1x3136xf32> + +// Element-wise add +%result = mico.add(%lhs, %rhs) : (tensor<1x64x28x28xf32>, tensor<1x64x28x28xf32>) -> tensor<1x64x28x28xf32> + +// Concatenation +%result = mico.concat(%t1, %t2) { + dim = 1 : i32 +} : (tensor<1x32x28x28xf32>, tensor<1x32x28x28xf32>) -> tensor<1x64x28x28xf32> +``` + +## Implementation + +### MiCoMLIRGen Class + +The `MiCoMLIRGen` class extends the existing code generation infrastructure: + +```python +from MiCoCodeGen import MiCoCodeGen + +class MiCoMLIRGen(MiCoCodeGen): + """ + MLIR code generator for MiCo models. + Generates MLIR code using the MiCo dialect for mixed-precision quantized models. + """ + + def __init__(self, model, dialect="mico"): + super().__init__(model) + self.dialect = dialect + self.mlir_operations = [] + self.mlir_types = {} + + def convert(self, output_directory, model_name): + """Generate MLIR code from the traced model.""" + # Generate MLIR module + mlir_code = self.generate_mlir_module(model_name) + + # Write to file + mlir_path = os.path.join(output_directory, f"{model_name}.mlir") + with open(mlir_path, "w") as f: + f.write(mlir_code) + + return mlir_path +``` + +### File Structure + +``` +MiCo-python/ +├── MiCoMLIRGen.py # Main MLIR code generator +├── mlir/ # MLIR dialect definitions (optional, for MLIR-based compilation) +│ └── MiCoDialect.td # TableGen dialect definition +├── doc/ +│ └── MLIR_INTEGRATION.md # This document +├── examples/ +│ └── mlir_example.py # Example usage +└── tests/ + └── test_mlir_codegen.py # Unit tests +``` + +## Usage Example + +```python +import torch +from models import LeNet +from MiCoMLIRGen import MiCoMLIRGen +from MiCoUtils import fuse_model + +# Load and prepare model +model = LeNet(1) +model.set_qscheme([[8, 6, 6, 4, 4], [8, 8, 8, 8, 8]]) # Mixed precision config +model.load_state_dict(torch.load('output/ckpt/lenet_mnist.pth')) +model = fuse_model(model) +model.eval() + +# Generate MLIR +mlir_gen = MiCoMLIRGen(model) +mlir_gen.forward(torch.randn(1, 1, 28, 28)) +mlir_path = mlir_gen.convert("output", "lenet_mnist") + +print(f"MLIR code generated at: {mlir_path}") +``` + +### Generated MLIR Output Example + +```mlir +// MiCo MLIR Module for LeNet-5 with Mixed Precision +module @lenet_mnist { + // Weight constants + mico.constant @conv1_weight : tensor<6x1x5x5x!mico.int<8>> = dense<...> + mico.constant @conv1_bias : tensor<6xf32> = dense<...> + mico.constant @conv2_weight : tensor<16x6x5x5x!mico.int<6>> = dense<...> + mico.constant @conv2_bias : tensor<16xf32> = dense<...> + mico.constant @fc1_weight : tensor<120x256x!mico.int<6>> = dense<...> + mico.constant @fc1_bias : tensor<120xf32> = dense<...> + mico.constant @fc2_weight : tensor<84x120x!mico.int<4>> = dense<...> + mico.constant @fc2_bias : tensor<84xf32> = dense<...> + mico.constant @fc3_weight : tensor<10x84x!mico.int<4>> = dense<...> + mico.constant @fc3_bias : tensor<10xf32> = dense<...> + + func.func @forward(%input: tensor<1x1x28x28xf32>) -> tensor<1x10xf32> { + // Conv1: W8A8 + %conv1 = mico.bitconv2d(%input, @conv1_weight, @conv1_bias) { + weight_bits = 8, act_bits = 8, stride = [1, 1], padding = [0, 0] + } : (tensor<1x1x28x28xf32>) -> tensor<1x6x24x24xf32> + %relu1 = mico.relu(%conv1) : tensor<1x6x24x24xf32> -> tensor<1x6x24x24xf32> + %pool1 = mico.maxpool2d(%relu1) {kernel_size = [2, 2], stride = [2, 2]} + : tensor<1x6x24x24xf32> -> tensor<1x6x12x12xf32> + + // Conv2: W6A8 + %conv2 = mico.bitconv2d(%pool1, @conv2_weight, @conv2_bias) { + weight_bits = 6, act_bits = 8, stride = [1, 1], padding = [0, 0] + } : (tensor<1x6x12x12xf32>) -> tensor<1x16x8x8xf32> + %relu2 = mico.relu(%conv2) : tensor<1x16x8x8xf32> -> tensor<1x16x8x8xf32> + %pool2 = mico.maxpool2d(%relu2) {kernel_size = [2, 2], stride = [2, 2]} + : tensor<1x16x8x8xf32> -> tensor<1x16x4x4xf32> + + // Flatten + %flat = mico.flatten(%pool2) {start_dim = 1} + : tensor<1x16x4x4xf32> -> tensor<1x256xf32> + + // FC1: W6A8 + %fc1 = mico.bitlinear(%flat, @fc1_weight, @fc1_bias) { + weight_bits = 6, act_bits = 8 + } : (tensor<1x256xf32>) -> tensor<1x120xf32> + %relu3 = mico.relu(%fc1) : tensor<1x120xf32> -> tensor<1x120xf32> + + // FC2: W4A8 + %fc2 = mico.bitlinear(%relu3, @fc2_weight, @fc2_bias) { + weight_bits = 4, act_bits = 8 + } : (tensor<1x120xf32>) -> tensor<1x84xf32> + %relu4 = mico.relu(%fc2) : tensor<1x84xf32> -> tensor<1x84xf32> + + // FC3: W4A8 + %fc3 = mico.bitlinear(%relu4, @fc3_weight, @fc3_bias) { + weight_bits = 4, act_bits = 8 + } : (tensor<1x10xf32>) -> tensor<1x10xf32> + + return %fc3 : tensor<1x10xf32> + } +} +``` + +## Lowering Passes + +### Stage 1: MiCo Dialect → Linalg + Arith + +Lower MiCo high-level operations to Linalg generic operations with quantization-aware transformations: + +```mlir +// Before (MiCo dialect) +%result = mico.bitlinear(%act, %weight, %bias) {weight_bits = 4, act_bits = 8} + +// After (Linalg + Arith) +%qact = linalg.generic ... // Quantize activation +%dequant_weight = arith.mulf %weight, %scale : f32 +%matmul = linalg.matmul ins(%qact, %dequant_weight) outs(%result) +%biased = linalg.generic ... // Add bias +``` + +### Stage 2: Linalg → Affine/SCF + +Lower Linalg operations to loop-based representations for further optimization: + +```mlir +// Affine loops with tiling for cache optimization +affine.for %i = 0 to 64 step 16 { + affine.for %j = 0 to 256 step 32 { + // Tiled matrix multiplication + } +} +``` + +### Stage 3: Affine → LLVM/Hardware IR + +Final lowering to LLVM IR or target-specific IR (e.g., for VexiiRiscv, Gemmini): + +```mlir +// LLVM dialect +llvm.func @forward(%arg0: !llvm.ptr) -> !llvm.ptr { + // Low-level implementation +} +``` + +## Hardware Backend Targets + +### 1. CPU (LLVM) +Standard LLVM-based compilation for x86, ARM, RISC-V CPUs. + +### 2. VexiiRiscv (Custom SIMD) +Target MiCo's custom SIMD instructions through LLVM RISC-V backend with custom intrinsics. + +### 3. Gemmini Accelerator +Generate code compatible with Gemmini's systolic array architecture. + +### 4. BitFusion Accelerator +Lower to BitFusion-compatible operations using existing `MiCoGraphGen` as reference. + +## Benefits + +1. **Portability**: MLIR provides a standard infrastructure for targeting multiple hardware backends +2. **Optimization**: Leverage MLIR's optimization passes (loop fusion, tiling, vectorization) +3. **Interoperability**: Compatible with other MLIR-based tools (IREE, TVM-MLIR, torch-mlir) +4. **Extensibility**: Easy to add new operations and lowering passes +5. **Debugging**: MLIR's textual format enables easier debugging and analysis + +## Future Work + +1. **Group Quantization**: Add support for group-wise quantization in the dialect +2. **Transformer Support**: Add operations for attention mechanisms (for LLaMa, ViT) +3. **Auto-tuning**: Integration with auto-tuning frameworks for optimal lowering decisions +4. **JIT Compilation**: Enable just-in-time compilation using MLIR's execution engine +5. **Hardware Synthesis**: Explore MLIR-to-RTL lowering for FPGA/ASIC targets + +## Dependencies + +- **mlir-python-bindings**: Python bindings for MLIR (optional, for advanced compilation) +- **torch-mlir**: For potential integration with torch-mlir ecosystem (optional) + +## References + +1. [MLIR: Multi-Level Intermediate Representation](https://mlir.llvm.org/) +2. [MLIR Dialects Documentation](https://mlir.llvm.org/docs/Dialects/) +3. [torch-mlir Project](https://github.com/llvm/torch-mlir) +4. [IREE: Intermediate Representation Execution Environment](https://github.com/iree-org/iree) +5. [MiCo Framework Paper](https://github.com/HKUSTGZ-MICS-LYU/MiCo-python) + +--- + +**Author**: MiCo Development Team +**Version**: 1.0 +**Date**: January 2026 diff --git a/examples/mlir_example.py b/examples/mlir_example.py new file mode 100644 index 0000000..880bdf5 --- /dev/null +++ b/examples/mlir_example.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +""" +MLIR Code Generation Example + +This script demonstrates how to use MiCoMLIRGen to generate MLIR code +from PyTorch models with mixed-precision quantization. + +Example usage: + python examples/mlir_example.py + +The script will: +1. Load a LeNet model with mixed-precision quantization +2. Generate MLIR code with the MiCo dialect +3. Save the output to output/lenet_mnist.mlir +""" + +import os +import sys +import torch + +# Add parent directory to path for imports +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from models import MLP, LeNet, VGG +from MiCoMLIRGen import MiCoMLIRGen +from MiCoUtils import fuse_model + + +def example_mlp(): + """Generate MLIR for a simple MLP model.""" + print("=" * 60) + print("Example 1: MLP with Mixed Precision") + print("=" * 60) + + # Create MLP model + model = MLP(in_features=256, config={"Layers": [128, 64, 10]}) + + # Set mixed precision: different bit widths per layer + # Weights: [8-bit, 6-bit, 4-bit] + # Activations: [8-bit, 8-bit, 8-bit] + weight_bits = [8, 6, 4] + activation_bits = [8, 8, 8] + model.set_qscheme([weight_bits, activation_bits]) + + # Fuse batch normalization if present + model = fuse_model(model) + model.eval() + + # Create MLIR generator + mlir_gen = MiCoMLIRGen(model) + + # Trace the model with example input + example_input = torch.randn(1, 256) + mlir_gen.forward(example_input) + + # Generate MLIR code + output_dir = "output" + os.makedirs(output_dir, exist_ok=True) + mlir_path = mlir_gen.convert(output_dir, "mlp_example") + + print(f"\nGenerated MLIR file: {mlir_path}") + + # Print first few lines of the generated code + with open(mlir_path, 'r') as f: + lines = f.readlines() + print("\nPreview of generated MLIR:") + print("-" * 40) + for line in lines[:30]: + print(line, end='') + if len(lines) > 30: + print("... (truncated)") + + print() + return mlir_path + + +def example_lenet(): + """Generate MLIR for LeNet model.""" + print("=" * 60) + print("Example 2: LeNet with Mixed Precision") + print("=" * 60) + + # Create LeNet model (1 input channel for MNIST) + model = LeNet(in_channels=1) + + # Set mixed precision configuration + # LeNet has 5 quantizable layers: conv1, conv2, fc1, fc2, fc3 + # Using progressively lower precision for later layers + weight_bits = [8, 6, 6, 4, 4] + activation_bits = [8, 8, 8, 8, 8] + model.set_qscheme([weight_bits, activation_bits]) + + # Fuse batch normalization + model = fuse_model(model) + model.eval() + + # Create MLIR generator + mlir_gen = MiCoMLIRGen(model) + + # Trace with 28x28 input (MNIST dimensions) + example_input = torch.randn(1, 1, 28, 28) + mlir_gen.forward(example_input) + + # Generate MLIR code + output_dir = "output" + os.makedirs(output_dir, exist_ok=True) + mlir_path = mlir_gen.convert(output_dir, "lenet_mnist_mlir") + + print(f"\nGenerated MLIR file: {mlir_path}") + + # Print the generated code + with open(mlir_path, 'r') as f: + content = f.read() + print("\nGenerated MLIR code:") + print("-" * 40) + print(content) + + print() + return mlir_path + + +def example_vgg_partial(): + """Generate MLIR for VGG (partial - first few layers).""" + print("=" * 60) + print("Example 3: VGG with Mixed Precision") + print("=" * 60) + + # Create VGG model for CIFAR-10 + model = VGG(in_channels=3, num_class=10) + + # Set uniform 8-bit precision (can be modified for MPQ) + n_layers = model.n_layers + weight_bits = [8] * n_layers + activation_bits = [8] * n_layers + model.set_qscheme([weight_bits, activation_bits]) + + # Fuse batch normalization + model = fuse_model(model) + model.eval() + + # Create MLIR generator + mlir_gen = MiCoMLIRGen(model) + + # Trace with CIFAR-10 dimensions (32x32) + example_input = torch.randn(1, 3, 32, 32) + mlir_gen.forward(example_input) + + # Generate MLIR code + output_dir = "output" + os.makedirs(output_dir, exist_ok=True) + mlir_path = mlir_gen.convert(output_dir, "vgg_cifar10_mlir") + + print(f"\nGenerated MLIR file: {mlir_path}") + print(f"Number of operations: {len(mlir_gen.mlir_ops)}") + print(f"Number of weights: {len(mlir_gen.mlir_weights)}") + + print() + return mlir_path + + +def main(): + """Run all examples.""" + print("\n" + "=" * 60) + print("MiCo MLIR Code Generation Examples") + print("=" * 60 + "\n") + + # Run examples + mlp_path = example_mlp() + lenet_path = example_lenet() + + # VGG example (optional, can be slow) + try: + vgg_path = example_vgg_partial() + except Exception as e: + print(f"VGG example skipped due to: {e}") + vgg_path = None + + # Summary + print("=" * 60) + print("Summary") + print("=" * 60) + print("\nGenerated MLIR files:") + print(f" - MLP: {mlp_path}") + print(f" - LeNet: {lenet_path}") + if vgg_path: + print(f" - VGG: {vgg_path}") + + print("\nThe generated MLIR files use the MiCo dialect which supports:") + print(" - Sub-byte integer types (!mico.int)") + print(" - Quantized operations (mico.bitlinear, mico.bitconv2d)") + print(" - Standard neural network operations (mico.relu, mico.maxpool2d)") + print("\nSee doc/MLIR_INTEGRATION.md for dialect specification and usage.") + + print("\n" + "=" * 60) + print("Examples Complete!") + print("=" * 60 + "\n") + + +if __name__ == "__main__": + main() diff --git a/examples/torch_mlir_example.py b/examples/torch_mlir_example.py new file mode 100644 index 0000000..1cef86f --- /dev/null +++ b/examples/torch_mlir_example.py @@ -0,0 +1,219 @@ +#!/usr/bin/env python3 +""" +Torch-MLIR Code Generation Example + +This script demonstrates how to use MiCoTorchMLIRGen to generate MLIR code +from PyTorch models using torch-mlir as the first pass. + +Example usage: + python examples/torch_mlir_example.py + +The script will: +1. Check if torch-mlir is installed +2. Load a LeNet model with mixed-precision quantization +3. Generate MLIR code using torch-mlir backend (or fallback to standalone) +4. Save the output to output/lenet_torch_mlir.mlir +""" + +import os +import sys + +# Add parent directory to path for imports +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import torch +from models import MLP, LeNet, VGG +from MiCoTorchMLIRGen import MiCoTorchMLIRGen, check_torch_mlir_installation +from MiCoUtils import fuse_model + + +def check_installation(): + """Check and report torch-mlir installation status.""" + print("=" * 60) + print("Torch-MLIR Installation Status") + print("=" * 60) + + status = check_torch_mlir_installation() + + print(f"Available: {status['available']}") + if status['available']: + print(f"Version: {status['version']}") + print(f"Supported output types: {', '.join(status['output_types'])}") + else: + print(f"\nTo install torch-mlir:") + print(f" {status['install_command']}") + + print() + return status['available'] + + +def example_mlp(): + """Generate MLIR for a simple MLP model.""" + print("=" * 60) + print("Example 1: MLP with Torch-MLIR Backend") + print("=" * 60) + + # Create MLP model + model = MLP(in_features=256, config={"Layers": [128, 64, 10]}) + + # Set mixed precision + weight_bits = [8, 6, 4] + activation_bits = [8, 8, 8] + model.set_qscheme([weight_bits, activation_bits]) + + model = fuse_model(model) + model.eval() + + # Create torch-mlir generator + mlir_gen = MiCoTorchMLIRGen(model, output_type="torch") + + # Trace the model + example_input = torch.randn(1, 256) + mlir_gen.forward(example_input) + + # Generate MLIR code + output_dir = "output" + os.makedirs(output_dir, exist_ok=True) + mlir_path = mlir_gen.convert(output_dir, "mlp_torch_mlir") + + print(f"\nGenerated MLIR file: {mlir_path}") + print(f"Backend used: {'torch-mlir' if mlir_gen.use_torch_mlir else 'standalone (fallback)'}") + + # Print first few lines + with open(mlir_path, 'r') as f: + lines = f.readlines() + print("\nPreview of generated MLIR:") + print("-" * 40) + for line in lines[:25]: + print(line, end='') + if len(lines) > 25: + print("... (truncated)") + + print() + return mlir_path + + +def example_lenet(): + """Generate MLIR for LeNet model.""" + print("=" * 60) + print("Example 2: LeNet with Torch-MLIR Backend") + print("=" * 60) + + # Create LeNet model + model = LeNet(in_channels=1) + + # Set mixed precision + weight_bits = [8, 6, 6, 4, 4] + activation_bits = [8, 8, 8, 8, 8] + model.set_qscheme([weight_bits, activation_bits]) + + model = fuse_model(model) + model.eval() + + # Create torch-mlir generator with linalg output + mlir_gen = MiCoTorchMLIRGen(model, output_type="torch") + + # Trace with MNIST dimensions + example_input = torch.randn(1, 1, 28, 28) + mlir_gen.forward(example_input) + + # Generate MLIR code + output_dir = "output" + os.makedirs(output_dir, exist_ok=True) + mlir_path = mlir_gen.convert(output_dir, "lenet_torch_mlir") + + print(f"\nGenerated MLIR file: {mlir_path}") + print(f"Backend used: {'torch-mlir' if mlir_gen.use_torch_mlir else 'standalone (fallback)'}") + + print() + return mlir_path + + +def example_compare_backends(): + """Compare torch-mlir and standalone backend outputs.""" + print("=" * 60) + print("Example 3: Compare Backends") + print("=" * 60) + + # Create model + model = MLP(in_features=64, config={"Layers": [32, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + example_input = torch.randn(1, 64) + output_dir = "output" + os.makedirs(output_dir, exist_ok=True) + + # Generate with torch-mlir backend (reusing the imported class from line 26) + torch_mlir_gen = MiCoTorchMLIRGen(model, output_type="torch") + torch_mlir_gen.forward(example_input) + torch_mlir_path = torch_mlir_gen.convert(output_dir, "compare_torch_mlir") + torch_mlir_used = torch_mlir_gen.use_torch_mlir + + # Generate with standalone backend + from MiCoMLIRGen import MiCoMLIRGen + standalone_gen = MiCoMLIRGen(model) + standalone_gen.forward(example_input) + standalone_path = standalone_gen.convert(output_dir, "compare_standalone") + + print(f"\nTorch-MLIR output: {torch_mlir_path}") + print(f" Backend actually used: {'torch-mlir' if torch_mlir_used else 'standalone (fallback)'}") + print(f"\nStandalone output: {standalone_path}") + + # Get file sizes + torch_mlir_size = os.path.getsize(torch_mlir_path) + standalone_size = os.path.getsize(standalone_path) + + print(f"\nFile sizes:") + print(f" Torch-MLIR: {torch_mlir_size} bytes") + print(f" Standalone: {standalone_size} bytes") + + print() + return torch_mlir_path, standalone_path + + +def main(): + """Run all examples.""" + print("\n" + "=" * 60) + print("MiCo Torch-MLIR Code Generation Examples") + print("=" * 60 + "\n") + + # Check installation + is_available = check_installation() + + if not is_available: + print("Note: torch-mlir not installed. Examples will use standalone fallback.") + print() + + # Run examples + mlp_path = example_mlp() + lenet_path = example_lenet() + torch_mlir_path, standalone_path = example_compare_backends() + + # Summary + print("=" * 60) + print("Summary") + print("=" * 60) + print("\nGenerated MLIR files:") + print(f" - MLP: {mlp_path}") + print(f" - LeNet: {lenet_path}") + print(f" - Comparison: {torch_mlir_path}") + print(f" {standalone_path}") + + print("\nThe torch-mlir backend provides:") + print(" - Production-ready PyTorch → MLIR lowering") + print(" - Torch dialect output compatible with IREE, Blade, etc.") + print(" - MiCo dialect overlay for sub-byte quantization metadata") + + if not is_available: + print("\nTo enable torch-mlir backend:") + print(" pip install torch-mlir -f https://github.com/llvm/torch-mlir-release/releases") + + print("\n" + "=" * 60) + print("Examples Complete!") + print("=" * 60 + "\n") + + +if __name__ == "__main__": + main() diff --git a/readme.md b/readme.md index b73c1ad..5777f47 100644 --- a/readme.md +++ b/readme.md @@ -102,6 +102,8 @@ Here are the main components/modules of MiCo. **Codegen** + `MiCoCodeGen`： C code generator for MiCo models. ++ `MiCoMLIRGen`: MLIR code generator for MiCo models (see [doc/MLIR_INTEGRATION.md](doc/MLIR_INTEGRATION.md)). ++ `MiCoTorchMLIRGen`: MLIR code generator using torch-mlir as first pass (see [doc/MLIR_INTEGRATION.md](doc/MLIR_INTEGRATION.md)). + `MiCoGraphGen`: DNN Weaver Op graph generator for MiCo models. + `MiCoLLaMaGen`: C code generator for MiCo TinyLLaMa models. @@ -125,6 +127,39 @@ Here are the main components/modules of MiCo. [doc/CHIPYARD_INTEGRATION.md](doc/CHIPYARD_INTEGRATION.md) +## MLIR Integration + +MiCo supports generating MLIR (Multi-Level Intermediate Representation) code for mixed-precision quantized models. The MiCo MLIR dialect provides sub-byte integer types and quantized neural network operations. + +**Standalone MLIR generation:** +```python +from models import LeNet +from MiCoMLIRGen import MiCoMLIRGen +from MiCoUtils import fuse_model +import torch + +model = LeNet(1) +model.set_qscheme([[8, 6, 6, 4, 4], [8, 8, 8, 8, 8]]) +model = fuse_model(model) +model.eval() + +mlir_gen = MiCoMLIRGen(model) +mlir_gen.forward(torch.randn(1, 1, 28, 28)) +mlir_gen.convert("output", "lenet_mnist") +``` + +**With torch-mlir backend:** +```python +from MiCoTorchMLIRGen import MiCoTorchMLIRGen + +# Use torch-mlir as first pass (requires: pip install torch-mlir) +mlir_gen = MiCoTorchMLIRGen(model, output_type="torch") +mlir_gen.forward(torch.randn(1, 1, 28, 28)) +mlir_gen.convert("output", "lenet_torch_mlir") +``` + +For detailed specification, see [doc/MLIR_INTEGRATION.md](doc/MLIR_INTEGRATION.md). + ## Publication Please refer to the paper for details. diff --git a/requirements.txt b/requirements.txt index 1ebbdd4..d5eda78 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,3 +19,7 @@ torchao transformers datasets accelerate + +# Optional: for MLIR generation with torch-mlir backend +# Install separately with: pip install torch-mlir -f https://github.com/llvm/torch-mlir-release/releases +# torch-mlir diff --git a/tests/test_mlir_codegen.py b/tests/test_mlir_codegen.py new file mode 100644 index 0000000..8cd3095 --- /dev/null +++ b/tests/test_mlir_codegen.py @@ -0,0 +1,478 @@ +#!/usr/bin/env python3 +""" +Test suite for MiCoMLIRGen class. + +Tests the MLIR code generation functionality including: +- Basic initialization and tracing +- Module handler tests (BitLinear, BitConv2d, etc.) +- Function handler tests (relu, add, etc.) +- Full model conversion tests +""" + +import os +import sys +import tempfile +import shutil +import unittest + +import torch +import torch.nn as nn + +# Add parent directory to path +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from models import MLP, LeNet +from MiCoMLIRGen import MiCoMLIRGen +from MiCoUtils import fuse_model + + +class TestMiCoMLIRGenBasics(unittest.TestCase): + """Test basic MiCoMLIRGen functionality.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_initialization(self): + """Test MiCoMLIRGen initialization.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + + self.assertIsNotNone(mlir_gen) + self.assertEqual(mlir_gen.dialect, "mico") + self.assertEqual(len(mlir_gen.mlir_ops), 0) + self.assertEqual(len(mlir_gen.mlir_weights), 0) + + def test_reset(self): + """Test the reset method.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + + # Add some data + mlir_gen.mlir_ops.append("test_op") + mlir_gen.mlir_weights["test"] = {} + + # Reset + mlir_gen.reset() + + # Verify cleared + self.assertEqual(len(mlir_gen.mlir_ops), 0) + self.assertEqual(len(mlir_gen.mlir_weights), 0) + self.assertEqual(mlir_gen.ssa_counter, 0) + + def test_format_tensor_type_float32(self): + """Test tensor type formatting for float32.""" + model = MLP(in_features=16, config={"Layers": [8]}) + model.set_qscheme([[8], [8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + + tensor = torch.randn(1, 64, 28, 28) + type_str = mlir_gen._format_tensor_type(tensor) + + self.assertEqual(type_str, "tensor<1x64x28x28xf32>") + + def test_format_tensor_type_quantized(self): + """Test tensor type formatting for quantized types.""" + model = MLP(in_features=16, config={"Layers": [8]}) + model.set_qscheme([[8], [8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + + tensor = torch.randn(64, 32) + type_str = mlir_gen._format_tensor_type(tensor, qbit=4) + + self.assertEqual(type_str, "tensor<64x32x!mico.int<4>>") + + def test_ssa_naming(self): + """Test SSA value naming.""" + model = MLP(in_features=16, config={"Layers": [8]}) + model.set_qscheme([[8], [8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + + name1 = mlir_gen._get_ssa_name() + name2 = mlir_gen._get_ssa_name() + name3 = mlir_gen._get_ssa_name("conv") + + self.assertEqual(name1, "%v0") + self.assertEqual(name2, "%v1") + self.assertEqual(name3, "%conv2") + + +class TestMiCoMLIRGenForward(unittest.TestCase): + """Test forward pass functionality.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_forward_mlp(self): + """Test forward pass with MLP model.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8, 6] # Mixed precision + activation_q = [8, 8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + example_input = torch.randn(1, 32) + + # Forward pass + output = mlir_gen.forward(example_input) + + self.assertIsNotNone(output) + self.assertEqual(output.shape, (1, 10)) + + # Check that operations were generated + self.assertGreater(len(mlir_gen.mlir_ops), 0) + + def test_forward_lenet(self): + """Test forward pass with LeNet model.""" + model = LeNet(1) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + example_input = torch.randn(1, 1, 28, 28) + + # Forward pass + output = mlir_gen.forward(example_input) + + self.assertIsNotNone(output) + self.assertEqual(output.shape[0], 1) + + # Check operations generated + self.assertGreater(len(mlir_gen.mlir_ops), 0) + + # Check weight constants generated + self.assertGreater(len(mlir_gen.mlir_weights), 0) + + +class TestMiCoMLIRGenConvert(unittest.TestCase): + """Test MLIR code conversion functionality.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_convert_generates_file(self): + """Test that convert generates MLIR file.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + + # Convert + mlir_path = mlir_gen.convert(self.temp_dir, "test_model") + + # Check file exists + self.assertTrue(os.path.exists(mlir_path)) + self.assertTrue(mlir_path.endswith(".mlir")) + + def test_convert_without_forward_raises_error(self): + """Test that convert raises error if forward not called.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + + with self.assertRaises(ValueError): + mlir_gen.convert(self.temp_dir, "test_model") + + def test_convert_content_has_module(self): + """Test that generated MLIR has module structure.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + mlir_path = mlir_gen.convert(self.temp_dir, "test_model") + + with open(mlir_path, 'r') as f: + content = f.read() + + # Check for essential MLIR structures + self.assertIn("module @test_model", content) + self.assertIn("func.func @forward", content) + self.assertIn("return", content) + + def test_convert_content_has_mico_ops(self): + """Test that generated MLIR contains mico dialect operations.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + mlir_path = mlir_gen.convert(self.temp_dir, "test_model") + + with open(mlir_path, 'r') as f: + content = f.read() + + # Check for MiCo dialect operations + self.assertIn("mico.bitlinear", content) + self.assertIn("mico.relu", content) + + def test_convert_mixed_precision(self): + """Test MLIR generation with mixed precision.""" + model = MLP(in_features=64, config={"Layers": [32, 16, 10]}) + weight_q = [8, 6, 4] # Different precisions + activation_q = [8, 8, 8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 64)) + mlir_path = mlir_gen.convert(self.temp_dir, "mixed_precision_model") + + with open(mlir_path, 'r') as f: + content = f.read() + + # Check that different bit widths are present + self.assertIn("weight_bits = 8", content) + self.assertIn("weight_bits = 6", content) + self.assertIn("weight_bits = 4", content) + + +class TestMiCoMLIRGenLeNet(unittest.TestCase): + """Test MLIR generation for LeNet model.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_lenet_conv_ops(self): + """Test that LeNet generates conv2d operations.""" + model = LeNet(1) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_path = mlir_gen.convert(self.temp_dir, "lenet_test") + + with open(mlir_path, 'r') as f: + content = f.read() + + # Check for conv2d operations + self.assertIn("mico.bitconv2d", content) + + def test_lenet_pooling_ops(self): + """Test that LeNet generates pooling operations.""" + model = LeNet(1) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_path = mlir_gen.convert(self.temp_dir, "lenet_test") + + with open(mlir_path, 'r') as f: + content = f.read() + + # Check for pooling operations (LeNet uses avgpool) + self.assertIn("mico.avgpool2d", content) + + def test_lenet_weight_constants(self): + """Test that LeNet generates weight constants.""" + model = LeNet(1) + weight_q = [8, 6, 6, 4, 4] # Mixed precision + activation_q = [8, 8, 8, 8, 8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_path = mlir_gen.convert(self.temp_dir, "lenet_test") + + with open(mlir_path, 'r') as f: + content = f.read() + + # Check for weight constants + self.assertIn("mico.constant", content) + self.assertIn("_weight", content) + self.assertIn("_bias", content) + + +class TestMiCoMLIRGenGetMLIRCode(unittest.TestCase): + """Test get_mlir_code method.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + + def test_get_mlir_code_before_forward_raises(self): + """Test that get_mlir_code raises error before forward.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + + with self.assertRaises(ValueError): + mlir_gen.get_mlir_code() + + def test_get_mlir_code_returns_string(self): + """Test that get_mlir_code returns a string.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + + code = mlir_gen.get_mlir_code() + + self.assertIsInstance(code, str) + self.assertGreater(len(code), 0) + + +class TestMiCoMLIRGenCustomOperations(unittest.TestCase): + """Test handling of various PyTorch operations.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_relu6_operation(self): + """Test ReLU6 activation handling.""" + class ModelWithReLU6(nn.Module): + def __init__(self): + super().__init__() + self.fc = nn.Linear(32, 16) + self.relu6 = nn.ReLU6() + self.n_layers = 1 + + def forward(self, x): + x = self.fc(x) + x = self.relu6(x) + return x + + from MiCoUtils import replace_quantize_layers + + model = ModelWithReLU6() + replace_quantize_layers(model, [8], [8]) # Modifies model in-place + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + mlir_path = mlir_gen.convert(self.temp_dir, "relu6_test") + + with open(mlir_path, 'r') as f: + content = f.read() + + self.assertIn("mico.relu6", content) + + def test_flatten_operation(self): + """Test Flatten operation handling.""" + class ModelWithFlatten(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(1, 8, 3) + self.flatten = nn.Flatten() + self.fc = nn.Linear(8 * 26 * 26, 10) + self.n_layers = 2 + + def forward(self, x): + x = self.conv(x) + x = self.flatten(x) + x = self.fc(x) + return x + + from MiCoUtils import replace_quantize_layers + + model = ModelWithFlatten() + replace_quantize_layers(model, [8, 8], [8, 8]) # Modifies model in-place + model.eval() + + mlir_gen = MiCoMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_path = mlir_gen.convert(self.temp_dir, "flatten_test") + + with open(mlir_path, 'r') as f: + content = f.read() + + self.assertIn("mico.flatten", content) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_torch_mlir_codegen.py b/tests/test_torch_mlir_codegen.py new file mode 100644 index 0000000..54e5140 --- /dev/null +++ b/tests/test_torch_mlir_codegen.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +""" +Test suite for MiCoTorchMLIRGen class. + +Tests the torch-mlir based MLIR code generation functionality including: +- Fallback behavior when torch-mlir is not available +- Basic initialization and configuration +- Output file generation +""" + +import os +import sys +import tempfile +import shutil +import unittest + +import torch +import torch.nn as nn + +# Add parent directory to path +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from models import MLP, LeNet +from MiCoTorchMLIRGen import MiCoTorchMLIRGen, check_torch_mlir_installation, TORCH_MLIR_AVAILABLE +from MiCoUtils import fuse_model + + +class TestMiCoTorchMLIRGenBasics(unittest.TestCase): + """Test basic MiCoTorchMLIRGen functionality.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_initialization(self): + """Test MiCoTorchMLIRGen initialization.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + + self.assertIsNotNone(mlir_gen) + self.assertEqual(mlir_gen.dialect, "mico") + self.assertEqual(mlir_gen.output_type, "torch") + + def test_initialization_with_output_type(self): + """Test initialization with different output types.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + for output_type in ["torch", "linalg", "stablehlo"]: + mlir_gen = MiCoTorchMLIRGen(model, output_type=output_type) + self.assertEqual(mlir_gen.output_type, output_type) + + def test_check_torch_mlir_installation(self): + """Test the installation check function.""" + status = check_torch_mlir_installation() + + self.assertIn("available", status) + self.assertIn("version", status) + self.assertIn("output_types", status) + self.assertIn("install_command", status) + + self.assertIsInstance(status["available"], bool) + self.assertIsInstance(status["output_types"], list) + + def test_is_torch_mlir_available(self): + """Test the availability check method.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + + # Should match the global availability + self.assertEqual(mlir_gen.is_torch_mlir_available(), TORCH_MLIR_AVAILABLE) + + +class TestMiCoTorchMLIRGenConvert(unittest.TestCase): + """Test MLIR code conversion functionality.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_convert_generates_file(self): + """Test that convert generates an MLIR file.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + + # Convert + mlir_path = mlir_gen.convert(self.temp_dir, "test_model") + + # Check file exists (either torch_mlir or standalone) + self.assertTrue(os.path.exists(mlir_path)) + self.assertTrue(mlir_path.endswith(".mlir")) + + def test_convert_without_forward_raises_error(self): + """Test that convert raises error if forward not called.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + + with self.assertRaises(ValueError): + mlir_gen.convert(self.temp_dir, "test_model") + + def test_convert_fallback_mode(self): + """Test that fallback to standalone mode works.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + + # Force fallback mode + original_use_torch_mlir = mlir_gen.use_torch_mlir + mlir_gen.use_torch_mlir = False + + mlir_path = mlir_gen.convert(self.temp_dir, "fallback_test") + + # Restore + mlir_gen.use_torch_mlir = original_use_torch_mlir + + self.assertTrue(os.path.exists(mlir_path)) + + with open(mlir_path, 'r') as f: + content = f.read() + + # Should have MiCo dialect content + self.assertIn("module", content) + + +class TestMiCoTorchMLIRGenLeNet(unittest.TestCase): + """Test torch-mlir generation for LeNet model.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_lenet_conversion(self): + """Test LeNet model conversion.""" + model = LeNet(1) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_path = mlir_gen.convert(self.temp_dir, "lenet_test") + + self.assertTrue(os.path.exists(mlir_path)) + + def test_lenet_mixed_precision(self): + """Test LeNet with mixed precision conversion.""" + model = LeNet(1) + weight_q = [8, 6, 6, 4, 4] # Mixed precision + activation_q = [8, 8, 8, 8, 8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + mlir_gen.forward(torch.randn(1, 1, 28, 28)) + mlir_path = mlir_gen.convert(self.temp_dir, "lenet_mixed") + + self.assertTrue(os.path.exists(mlir_path)) + + +class TestMiCoTorchMLIRGenReset(unittest.TestCase): + """Test reset functionality.""" + + def setUp(self): + """Set up test fixtures.""" + torch.manual_seed(42) + + def test_reset_clears_state(self): + """Test that reset clears the generator state.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + model.set_qscheme([[8, 8], [8, 8]]) + model = fuse_model(model) + model.eval() + + mlir_gen = MiCoTorchMLIRGen(model) + mlir_gen.forward(torch.randn(1, 32)) + + # Check state is populated + self.assertIsNotNone(mlir_gen.example_inputs) + + # Reset + mlir_gen.reset() + + # Check state is cleared + self.assertIsNone(mlir_gen.example_inputs) + self.assertIsNone(mlir_gen.torch_mlir_ir) + self.assertEqual(len(mlir_gen.mico_quantization_ops), 0) + + +if __name__ == "__main__": + unittest.main()