StableHLO Quantization

Quantization Types in StableHLO

Quantization is a technique to optimize machine learning models by converting floating-point numbers (like those used in original models) into lower-precision integers. This reduces memory usage and speeds up computations, making models more efficient for deployment on devices with limited resources.

StableHLO quantization follows the LiteRT quantization specification, using a uniform quantization scheme with support for both per-tensor and per-axis quantization. It inherits its type expression from MLIR's Quant dialect, providing a standardized way to represent quantized data types.

Uniform quantization maps floating-point values to integers using a uniform step size, resulting in evenly spaced quantized values. This is achieved through an affine ralationship using two key quantization parameters.

Uniform quantization simplifies the representation of floating-point numbers by mapping them to integers that are evenly spaced. This mapping is achieved through an affine transformation that uses two key parameters: scale and zero point. The scale determines determines the step size between consecutive quantized values. A smaller scale means the quantized values are closer together. The zero point defines the integer value that represents zero in the original floating-point space.

The relationship between the original floating-point value (real_value) and the quantized integer value (quantized_value) in uniform quantization is:

real_value = scale * (quantized_value - zero_point)

Per-tensor Quantization

In per-tensor quantization, a single scale and zero point are used for all the values within the tensor. A per-tensor quantized type is expressed in StableHLO as:

quant.uniform scale:zero_point>

Example: !quant.uniform<i8:f32, 0.01:50>

This represents an 8-bit integer (i8) used to store a 32-bit floating-point number (f32) using a scale of 0.01 and a zero point of 50.

Per-axis Quantization

Per-axis quantization offers a more fine-grained approach compared to per-tensor quantization. Instead of using a single scale and zero point for the entire tensor, per-axis quantization assigns separate scales and zero points to slices along a specific dimension quantized_dimension of the tensor. This is particularly useful when values vary significantly across different dimensions, allowing for better preservation of information and accuracy.

Consider a tensor t with dimensions sizes [4, 3, 2]. We choose to quantize this tensor along the second dimension (quantized_dimension = 1). This means we'll have three slices (since the second dimension has a size of 3), each with its own scale and zero point:

t[:, 0, :]: This slice gets scale[0] and zero_point[0].
t[:, 1, :]: This slice gets scale[1] and zero_point[1].
t[:, 2, :]: This slice gets scale[2] and zero_point[2].

In StableHLO, per-axis quantized type is expressed as:

quant.uniform {scale0:zero_point0, scale1:zero_point1, ...}>

where the length of the scale:zero_point matches the number of slices along the quantized_dimension of the containing tensor.

Example: tensor<4x3x2x!quant.uniform<i8:f32:1, {0.2:20, 0.1:10, 0.3:30}>>

Quantization Passes in StableHLO

StableHLO provides several compiler passes which allow for different transformations and optimizations related to quantization, giving you flexibility in how you handle quantized models. These passes are:

stablehlo-legalize-qdq-to-quantized-op

This pass fuses a common pattern in quantized models, a dequantize operation followed by a floating-point operation, and finally a quantize operation, into a single quantized operation. details

stablehlo-legalize-quantized-op-to-qdq

This pass does the opposite of the previous pass. It decomposes a quantized operation into its equivalent sequence of dequantize, floating-point operation, and quantize operations. details

stablehlo-legalize-quant-to-math

This pass converts StableHLO operations on quantized types into equivalent operations on integer types. It essentially implements the quantization arithmetic using standard mathematical operations. This decompsition is useful for systems that do not support quantization natively, but can still use the quantization arithmetic to express the semantics of quantized models. details

stablehlo-quant-legalize-to-tosa-rescale

StableHLO offers the capability to legalize quantized operations to their corresponding representations in the TOSA dialect. This legalization facilitates compatibility and interoperability between StableHLO and TOSA. This pass strategically converts StableHLO quantized operations into a combination of StableHLO and TOSA operations, with the TOSA dialect primarily employed for the rescale operation. The tosa.rescale op plays a crucial role in adjusting the scale and zero point of quantized values, enabling accurate representation of quantized data within the TOSA framework. details

tosa-rescale-legalize-to-stablehlo

This pass rewrites TOSA rescale operations to StableHLO primitive math operations. One of the main use cases for this pass is to allow the StableHLO interpreter to evaluate programs containing TOSA rescale operations. details

Evaluating Quantized Programs

The StableHLO reference interpreter can efficiently execute programs containing quantized operations. To achieve this, it first lowers the program to an equivalent representation using only integer operations. This lowering process involves a series of compiler passes that transform the program before interpretation.

Essentially, the interpreter leverages the stablehlo-legalize-quant-to-math pass to convert quantized operations into their corresponding integer arithmetic implementations. This pass introduces CHLO broadcast operations for handling scale multiplication/division and zero-point addition. To ensure compatibility with the StableHLO interpreter, these CHLO operations are then legalized to StableHLO operations. This introduces shape-related operations that are subsequently canonicalized and optimized using a series of canonicalization passes.

The complete sequence of passes involved in this lowering process is as follows:

stablehlo-legalize-quant-to-math
chlo-legalize-to-stablehlo
canonicalize
shape-legalize-to-stablehlo
stablehlo-canonicalize-dynamism

Quantized Test Cases

StableHLO provides a comprehensive suite of quantized test cases to validate the correctness and behavior of quantized operations. These test cases serve as unit tests, covering various StableHLO operations in quantized scenarios.

A typical example of a quantized test case looks like

func.func @main() -> tensor<11xf32> {
    %operand_0 = stablehlo.constant dense<...> : tensor<11xf32>
    %operand_1 = stablehlo.constant dense<...> : tensor<11xf32>
    %golden = stablehlo.constant dense<...> : tensor<11xf32>

    %0 = stablehlo.uniform_quantize %operand_0 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>
    %1 = stablehlo.uniform_quantize %operand_1 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>

    %2 = stablehlo.add %1, %0 : tensor<11x!quant.uniform<i8:f32, 0.3>>

    %result = stablehlo.uniform_dequantize %2 : (tensor<11x!quant.uniform<i8:f32, 0.3>>) -> tensor<11xf32>

    %4 = stablehlo.custom_call @check.eq(%golden, %result) : (tensor<11xf32>, tensor<11xf32>) -> tensor<i1>

    return %3 : tensor<11xf32>
  }

and includes:

These test cases are valuable for:

  • Validating StableHLO quantization: Ensuring that the quantization behavior of StableHLO operations aligns with the expected results.
  • Cross-validation: Comparing the behavior of StableHLO quantization with other implementations or frameworks.
  • Debugging and development: Aiding in the development and debugging of new quantization features or optimizations.