Representing Quantized Operations in IR

Effectively optimizing low-precision models requires the compiler's intermediate representation (IR) to explicitly capture the nuances of quantization. Simply using standard integer types like i8 is insufficient because it omits critical information about how these integers map back to the real-number domain they approximate. The IR must encode the quantization scheme, parameters, and the operations that convert between floating-point and quantized domains.

A representation acts as the contract between high-level model descriptions (potentially annotated with quantization information) and the low-level optimization and code generation passes. Without this explicit representation, the compiler cannot reason about the numerical properties of quantized operations or target specialized low-precision hardware instructions effectively.

Encoding Quantization Parameters in the IR

The core challenge is representing the affine mapping between the real numbers ($r$) and the quantized integers ($q$):

$$r = s \times (q - Z)$$

Here, $s$ is the scale (a positive float) and $Z$ is the zero-point (an integer matching the range of $q$). Both $s$ and $Z$ must be somehow associated with the quantized tensor data within the IR. Several strategies exist:

1. Dedicated Quantized Types: This is arguably the cleanest approach, often seen in multi-level IRs like MLIR. New types are defined that directly bundle the storage type (e.g., i8, u8) with the quantization parameters and potentially the "expressed" floating-point type it represents.

   • Example (MLIR-like syntax): A tensor type might look like tensor<1x256x256x3x !quant.uniform<i8:f32, 0.015:128>>. This clearly defines a tensor with dimensions 1x256x256x3, where each element is stored as an i8. This i8 represents an f32 value using a uniform affine quantization scheme (!quant.uniform) with a scale of 0.015 and a zero-point of 128.
   • Advantages: Type safety ensures that quantization information is always present and consistently propagated. Operations can be defined to operate directly on these specific quantized types, simplifying analysis and transformation.
   • Disadvantages: Requires extending the IR's type system, which can be a significant undertaking if starting from a more traditional IR.

2. Type Attributes or Metadata: An alternative is to use standard integer types (e.g., tensor<1x256x256x3 x i8>) but attach the quantization parameters (scale, zero_point, axis for per-channel) as attributes or metadata to the tensor value or the operations producing/consuming it.

• Example (IR): %input_quant = MyDialect.Quantize(%input_fp32) {scale=0.015, zero_point=128, storage_type=i8} : (tensor<1x...xf32>) -> tensor<1x...xi8> %weight_quant = GetQuantizedWeight() {scale=0.008, zero_point=0, storage_type=i8, axis=0} : () -> tensor<64x...xi8> %conv_output_quant = MyDialect.Conv2D(%input_quant, %weight_quant) {output_scale=0.1, output_zero_point=110, storage_type=i8} : (tensor<...xi8>, tensor<...xi8>) -> tensor<...xi8>
  • Advantages: Avoids complex type system changes. Can be retrofitted onto existing IRs more easily.
  • Disadvantages: Information is less tightly coupled to the data. Requires careful handling by optimization passes to ensure parameters are correctly propagated and not accidentally dropped. Less compile-time safety.

Representing Quantization Granularity

Models often use different quantization parameters for different parts of a tensor, particularly weights (per-channel or per-axis quantization). The IR must support this:

• Per-Tensor: A single scale and zero-point apply to the entire tensor. This is simpler and is often represented as scalar attributes or directly within the quantized type.
• Per-Axis (Per-Channel): Multiple scales and zero-points apply, typically one for each channel along a specific axis (e.g., the output channel axis for convolution weights). The IR needs to store an array of scales/zero-points and the corresponding axis index.
  • In dedicated type systems, the type itself might encode the axis and reference a constant array of parameters.
  • In attribute systems, attributes would hold the array values and the axis index.

Representing Quantization and Dequantization Operations

Explicit operations are needed in the IR to represent the transitions between floating-point and quantized domains:

• Quantize (quant): Takes a floating-point tensor and quantization parameters ($s$, $Z$) as input, and produces a quantized integer tensor. $$q = \text{round}(r / s) + Z$$ The IR operation node would reference the input tensor and the scale/zero-point values (or the target quantized type which implies them).
• Dequantize (dequant): Takes a quantized integer tensor and its associated parameters ($s$, $Z$) as input, and produces a floating-point tensor. $$r = s \times (q - Z)$$ Similarly, the IR node references the quantized input and its parameters.
• Requantize (requant): Takes a quantized tensor (often an intermediate result with higher precision, e.g., i32 accumulator) and parameters for both the input and desired output types. It performs the scale adjustment and potential down-casting in the quantized domain, avoiding a costly round-trip through floating-point. $$q_{\text{out}} = \text{round}((s_{\text{in}} / s_{\text{out}}) \times (q_{\text{in}} - Z_{\text{in}})) + Z_{\text{out}}$$ This operation is fundamental for optimizing sequences of quantized computations, especially convolutions and matrix multiplications where intermediate accumulations occur.

Below is a diagram showing how these operations might appear in a graph IR fragment:

A flow showing quantization of input, a quantized convolution producing a higher-precision accumulator, requantization back to INT8, and final dequantization to FP32.

Implications for Optimization

Having this explicit representation allows the compiler to:

• Fuse Operations: Combine quantize -> op -> dequantize sequences into dedicated quantized kernel calls. Fuse requantize operations with preceding computations.
• Propagate Parameters: Deduce the correct quantization parameters for intermediate tensors resulting from operations like element-wise addition or concatenation.
• Select Optimal Kernels: Map high-level quantized operations (like MyDialect.Conv2D operating on quantized types) to specific low-precision hardware instructions (e.g., INT8 dot products) or optimized library calls (e.g., cuDNN, MIOpen, oneDNN).
• Verify Correctness: Perform static checks to ensure quantization parameters are used consistently and that operations are valid for the given quantized types.

In summary, representing quantized operations and their associated parameters directly and explicitly within the compiler's IR is fundamental. Whether through dedicated types or attribute systems, this representation provides the necessary information for sophisticated optimization passes to analyze, transform, and generate highly efficient low-precision code tailored for modern hardware. It bridges the gap between the high-level intent of using quantization and the low-level execution details required for performance.