Request description
How quantized models arrive today
torch-mlir represents quantized models in QDQ form. PT2E models are already exported this way, while quantized ops from other frontends, such as ONNX, are lowered to the same representation.
QDQ describes the model semantics, but it is not an efficient execution form. Executing the graph directly means dequantizing the operands to floating point and performing the computation in floating point.
Today this happens in two stages. First, torch-mlir's FuseQuantizedOps finds operations such as matmul and convolution whose inputs are dequantized values. It removes the dequantization from the compute path and rewrites the operation to consume the quantized operands directly.
TorchToLinalg then lowers these quantized Torch operations to linalg.quantized_matmul, or the corresponding quantized convolution. The zero points are recovered from the operations that produced the quantized operands.
IREE later lowers these quantized Linalg operations again during global optimization. QuantizedMatmulToMatmul and QuantizedConvToConv replace them with regular linalg.matmul or convolution operations together with the required zero-point correction terms.
So the linalg.quantized_* operations are only an intermediate representation. They are introduced during Torch-to-Linalg conversion and removed again before execution.
|
struct QuantizedMatmulToMatmul : OpInterfaceRewritePattern<linalg::LinalgOp> { |
|
struct QuantizedConvToConv : OpRewritePattern<linalg::Conv2DNhwcHwcfQOp> { |
Limitations of the current approach
The current implementation only handles a fairly narrow set of cases:
https://github.com/llvm/torch-mlir/blob/6d0186baf569e771ea8f1a8e7914aa0f65aadc75/lib/Conversion/TorchToLinalg/Linear.cpp#L270
These limitations mostly come from the representation itself. The linalg.quantized_* operations use scalar zero points and have fixed ranks, fixed operand counts, and a single element type per operand.
As a result, quantization schemes such as per-channel or blockwise quantization cannot be represented directly. Cases that do not match the expected form fall back to floating-point execution.
Why the rewrite should happen below torch-mlir
The limitations above are just one reason to move away from the current torch-mlir rewrite.
There is also a separate architectural reason. Even if those representation limitations were fixed, torch-mlir is still too early in the compilation pipeline to decide how a QDQ graph should be executed.
The discussion around first-class QDQ operations in torch-mlir already raised this issue (llvm/torch-mlir#4600): different backends may want to lower the same quantized model differently.
For example:
- quantize both operands, accumulate in i32, and requantize using integer arithmetic;
- quantize both operands, accumulate in i32, and apply the scales in floating point;
- use weight-only quantization: keep the activations in floating point, dequantize only the weights, and perform the contraction in floating point;
- execute the QDQ graph directly.
The preferred strategy depends on the target. GPUs may benefit more from reducing weight memory traffic than from integer arithmetic itself, making weight-only quantization attractive. CPUs, on the other hand, often have to use some from of integer dot-product instructions for performance, making an integer contraction preferable.
Accuracy requirements can also influence the choice, since these execution strategies are not necessarily numerically equivalent.
Performing the rewrite in torch-mlir therefore makes a backend-specific decision before the backend and its capabilities are known. Keeping QDQ intact until Linalg/IREE allows that decision to be made at the level where target information is available.
The rewrite can be expressed generically
The algebra for converting QDQ contractions to integer arithmetic is described here: https://hackmd.io/@Afn5isQ9RRKCQuXiomYR4A/H1OLIAxvGg
The derivation is not specific to matmul or convolution. It applies to contractions in general. The main difference between individual cases is which dimensions the quantization parameters vary over.
This makes a per-op pattern set unnecessary. Instead, the rewrite can operate on the contraction structure itself.
Linalg is a suitable level for this because contractions are represented using indexing maps and iterator types. The same rewrite can therefore be applied to named and generic contraction-like operations.
Required pieces
The implementation can be split into 4 mostly independent parts.
1. QDQ operations
Add:
iree_linalg_ext.quantize_affineiree_linalg_ext.dequantize_affine
Each operand has an indexing map over a shared iteration space corresponding to the value being quantized.
per-tensor affine_map<(d0, d1) -> ()>
per-channel affine_map<(d0, d1) -> (d0)>
per-block affine_map<(d0, d1, d2) -> (d0, d1)>
This allows the same pair of operations to represent different quantization granularities.
2. Frontend conversion
Add a pass converting the QDQ operations in torch-mlir
quantized_decomposed.{,de}quantize_per_{tensor,channel}...
to the new linalg_ext operations.
3. Propagation
For the contraction rewrite to match, the dequantize_affine should be immediate producers of the contraction operands. Exported models frequently contain shape or indexing operations between the dequantization and the contraction.
Examples include:
aten.linear, where the weights are dequantized and then transposed- padded convolution, where the dequantized activation is padded
- classifiers that reshape a pooled activation before a matmul.
Two kinds of propagation are needed.
Unit extent dim folding, reshape propagation and transpose propagation are already generic over ops that state an indexing map per operand, so they only have to be applied to the new ops.
Pad needs a new pattern, since padding the real valued tensor with zero is the same as padding the quantized tensor with the zero point, and only the quantization op knows its zero point.
4. Contraction rewrite
Linalg based implementation of: https://hackmd.io/@Afn5isQ9RRKCQuXiomYR4A/H1OLIAxvGg
contraction(dequantized operands)
into an integer contraction followed by the required zero-point correction epilogue.
Sign shifting
Storage signedness should remain a lowering decision because different targets have different efficient integer instruction forms.
For example, the 8-bit VNNI dot-product instruction VPDPBUSD / _mm512_dpbusd_epi32 operates on an unsigned LHS and signed RHS. There is no equivalent s8 × s8 form.
IREE's s8 × s8 VNNI micro-kernel therefore does not use this instruction directly. Instead,iree_uk_mmt4d_tile_s8s8s32_*_x86_64_avx512_vnni
widens the inputs to s16 and uses _mm512_dpwssd_epi32. _mm512_dpwssd_epi32 performs two products per 32-bit lane, whereas the 8-bit instruction performs four. The s8 × s8 path therefore has roughly half the multiply-accumulate rate, in addition to the widening overhead.
Why iree_linalg_ext and not linalg
Could also be moved directly to linalg but this way we can get some milage on this approach first.
Assisted by Claude/Codex
What component(s) does this issue relate to?
Frontends, Compiler
Additional context
No response
Request description
How quantized models arrive today
torch-mlir represents quantized models in QDQ form. PT2E models are already exported this way, while quantized ops from other frontends, such as ONNX, are lowered to the same representation.
QDQ describes the model semantics, but it is not an efficient execution form. Executing the graph directly means dequantizing the operands to floating point and performing the computation in floating point.
Today this happens in two stages. First, torch-mlir's
FuseQuantizedOpsfinds operations such as matmul and convolution whose inputs are dequantized values. It removes the dequantization from the compute path and rewrites the operation to consume the quantized operands directly.TorchToLinalgthen lowers these quantized Torch operations tolinalg.quantized_matmul, or the corresponding quantized convolution. The zero points are recovered from the operations that produced the quantized operands.IREE later lowers these quantized Linalg operations again during global optimization.
QuantizedMatmulToMatmulandQuantizedConvToConvreplace them with regularlinalg.matmulor convolution operations together with the required zero-point correction terms.So the
linalg.quantized_*operations are only an intermediate representation. They are introduced during Torch-to-Linalg conversion and removed again before execution.iree/compiler/src/iree/compiler/GlobalOptimization/QuantizedMatmulToMatmul.cpp
Line 37 in a82330a
iree/compiler/src/iree/compiler/GlobalOptimization/QuantizedConvToConv.cpp
Line 127 in a82330a
Limitations of the current approach
The current implementation only handles a fairly narrow set of cases:
https://github.com/llvm/torch-mlir/blob/6d0186baf569e771ea8f1a8e7914aa0f65aadc75/lib/Conversion/TorchToLinalg/Linear.cpp#L270
These limitations mostly come from the representation itself. The
linalg.quantized_*operations use scalar zero points and have fixed ranks, fixed operand counts, and a single element type per operand.As a result, quantization schemes such as per-channel or blockwise quantization cannot be represented directly. Cases that do not match the expected form fall back to floating-point execution.
Why the rewrite should happen below torch-mlir
The limitations above are just one reason to move away from the current torch-mlir rewrite.
There is also a separate architectural reason. Even if those representation limitations were fixed, torch-mlir is still too early in the compilation pipeline to decide how a QDQ graph should be executed.
The discussion around first-class QDQ operations in torch-mlir already raised this issue (llvm/torch-mlir#4600): different backends may want to lower the same quantized model differently.
For example:
The preferred strategy depends on the target. GPUs may benefit more from reducing weight memory traffic than from integer arithmetic itself, making weight-only quantization attractive. CPUs, on the other hand, often have to use some from of integer dot-product instructions for performance, making an integer contraction preferable.
Accuracy requirements can also influence the choice, since these execution strategies are not necessarily numerically equivalent.
Performing the rewrite in torch-mlir therefore makes a backend-specific decision before the backend and its capabilities are known. Keeping QDQ intact until Linalg/IREE allows that decision to be made at the level where target information is available.
The rewrite can be expressed generically
The algebra for converting QDQ contractions to integer arithmetic is described here: https://hackmd.io/@Afn5isQ9RRKCQuXiomYR4A/H1OLIAxvGg
The derivation is not specific to matmul or convolution. It applies to contractions in general. The main difference between individual cases is which dimensions the quantization parameters vary over.
This makes a per-op pattern set unnecessary. Instead, the rewrite can operate on the contraction structure itself.
Linalg is a suitable level for this because contractions are represented using indexing maps and iterator types. The same rewrite can therefore be applied to named and generic contraction-like operations.
Required pieces
The implementation can be split into 4 mostly independent parts.
1. QDQ operations
Add:
iree_linalg_ext.quantize_affineiree_linalg_ext.dequantize_affineEach operand has an indexing map over a shared iteration space corresponding to the value being quantized.
This allows the same pair of operations to represent different quantization granularities.
2. Frontend conversion
Add a pass converting the QDQ operations in torch-mlir
to the new linalg_ext operations.
3. Propagation
For the contraction rewrite to match, the
dequantize_affineshould be immediate producers of the contraction operands. Exported models frequently contain shape or indexing operations between the dequantization and the contraction.Examples include:
aten.linear, where the weights are dequantized and then transposedTwo kinds of propagation are needed.
Unit extent dim folding, reshape propagation and transpose propagation are already generic over ops that state an indexing map per operand, so they only have to be applied to the new ops.
Pad needs a new pattern, since padding the real valued tensor with zero is the same as padding the quantized tensor with the zero point, and only the quantization op knows its zero point.
4. Contraction rewrite
Linalg based implementation of: https://hackmd.io/@Afn5isQ9RRKCQuXiomYR4A/H1OLIAxvGg
into an integer contraction followed by the required zero-point correction epilogue.
Sign shifting
Storage signedness should remain a lowering decision because different targets have different efficient integer instruction forms.
For example, the 8-bit VNNI dot-product instruction
VPDPBUSD/_mm512_dpbusd_epi32operates on an unsigned LHS and signed RHS. There is no equivalent s8 × s8 form.IREE's s8 × s8 VNNI micro-kernel therefore does not use this instruction directly. Instead,
iree_uk_mmt4d_tile_s8s8s32_*_x86_64_avx512_vnniwidens the inputs to s16 and uses
_mm512_dpwssd_epi32._mm512_dpwssd_epi32performs two products per 32-bit lane, whereas the 8-bit instruction performs four. The s8 × s8 path therefore has roughly half the multiply-accumulate rate, in addition to the widening overhead.Why iree_linalg_ext and not linalg
Could also be moved directly to linalg but this way we can get some milage on this approach first.
Assisted by Claude/Codex
What component(s) does this issue relate to?
Frontends, Compiler
Additional context
No response