[Torch] Convert PT2E per-token quantization to LinalgExt (#25024) Extends `ConvertTorchQuantizationToLinalgExt` (#24918) to `quantized_decomposed.quantize_per_token` and `dequantize_per_token`. Depends on #25008 and #25046. Part of #24862. Signed-off-by: Roberto Laudani <laudani@roofline.ai>
diff --git a/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp b/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp index 790322a..853c18e 100644 --- a/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp +++ b/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp
@@ -120,6 +120,15 @@ return success(); } +// Each vector along the last input dimension is a token with its own +// parameters. Sets `blockSize` to the dimension size along the last dimension +// and to one elsewhere. Requires `quantization.input` to be set. +static void setTokenBlockSize(PT2EQuantization &quantization) { + auto inputType = cast<RankedTensorType>(quantization.input.getType()); + quantization.blockSize = llvm::to_vector(inputType.getShape()); + std::fill_n(quantization.blockSize.begin(), inputType.getRank() - 1, 1); +} + // Returns whether the elements of a Torch zero point are unsigned. Bools count // as unsigned, since LinalgExt would otherwise sign-extend `true` to -1. static bool isUnsignedZeroPoint(Type zeroPointType) { @@ -244,6 +253,42 @@ return quantization; } +static FailureOr<PT2EQuantization> readPT2EQuantization( + Torch::QuantizedDecomposedQuantizePerTokenOp op, + Torch::QuantizedDecomposedQuantizePerTokenOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) { + PT2EQuantization quantization; + quantization.input = adaptor.getInput(); + quantization.scale = adaptor.getScales(); + quantization.zeroPoint = adaptor.getZeroPoints(); + quantization.zeroPointUnsigned = + isUnsignedZeroPoint(op.getZeroPoints().getType()); + setTokenBlockSize(quantization); + if (failed(matchConstantBounds(op, op.getQuantMin(), op.getQuantMax(), + quantization, rewriter))) { + return failure(); + } + return quantization; +} + +static FailureOr<PT2EQuantization> readPT2EQuantization( + Torch::QuantizedDecomposedDequantizePerTokenOp op, + Torch::QuantizedDecomposedDequantizePerTokenOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) { + PT2EQuantization quantization; + quantization.input = adaptor.getInput(); + quantization.scale = adaptor.getScales(); + quantization.zeroPoint = adaptor.getZeroPoints(); + quantization.zeroPointUnsigned = + isUnsignedZeroPoint(op.getZeroPoints().getType()); + setTokenBlockSize(quantization); + if (failed(matchConstantBounds(op, op.getQuantMin(), op.getQuantMax(), + quantization, rewriter))) { + return failure(); + } + return quantization; +} + // Returns, for each input dimension, its dimensions in the blocked type. A // dimension with several blocks of several elements is split into the block // and the position within it, so that LinalgExt can index the block. @@ -284,8 +329,9 @@ // Returns the dimensions of the blocked type that select a block and so index // the scale and zero point: none for per-tensor quantization, the channel -// dimension for per-channel quantization, and the row and group for -// per-channel-group quantization. +// dimension for per-channel quantization, the row and group for +// per-channel-group quantization, and every dimension but the last for +// per-token quantization. static SmallVector<unsigned> getParameterDims(const PT2EQuantization &quantization) { ArrayRef<int64_t> inputShape = @@ -424,8 +470,9 @@ // The PT2E reference upcasts f16 and bf16 inputs and computes in f32. It // multiplies by `1.0 / scale` rounded to f32, whereas LinalgExt divides by // the scale, so results can differ by one ulp before rounding. The - // per-channel-group reference also adds the zero point before rounding, so - // ties with odd zero points differ, as in Torch-to-Linalg. + // per-channel-group and per-token references also add the zero point before + // rounding, so ties with odd zero points differ. Torch-to-Linalg follows only + // the per-token one. Value input = createBlockedInput(builder, location, quantization); Value scale = truncateScalarScaleToF32(builder, location, quantization.scale); scale = collapseParameter(builder, location, scale); @@ -466,9 +513,9 @@ namespace { -// Converts `quantize_per_tensor`, `quantize_per_channel` or -// `quantize_per_channel_group`, whose input holds the real values and whose -// result holds the quantized storage. +// Converts `quantize_per_tensor`, `quantize_per_channel`, +// `quantize_per_channel_group` or `quantize_per_token`, whose input holds the +// real values and whose result holds the quantized storage. template <typename OpTy> struct ConvertQuantize final : OpConversionPattern<OpTy> { using OpConversionPattern<OpTy>::OpConversionPattern; @@ -511,9 +558,9 @@ } }; -// Converts `dequantize_per_tensor`, `dequantize_per_channel` or -// `dequantize_per_channel_group`, whose input holds the quantized storage and -// whose result holds the real values. +// Converts `dequantize_per_tensor`, `dequantize_per_channel`, +// `dequantize_per_channel_group` or `dequantize_per_token`, whose input holds +// the quantized storage and whose result holds the real values. template <typename OpTy> struct ConvertDequantize final : OpConversionPattern<OpTy> { using OpConversionPattern<OpTy>::OpConversionPattern; @@ -582,10 +629,12 @@ ConvertQuantize<Torch::QuantizedDecomposedQuantizePerTensorOp>, ConvertQuantize<Torch::QuantizedDecomposedQuantizePerChannelOp>, ConvertQuantize<Torch::QuantizedDecomposedQuantizePerChannelGroupOp>, + ConvertQuantize<Torch::QuantizedDecomposedQuantizePerTokenOp>, ConvertDequantize<Torch::QuantizedDecomposedDequantizePerTensorOp>, ConvertDequantize<Torch::QuantizedDecomposedDequantizePerChannelOp>, ConvertDequantize< - Torch::QuantizedDecomposedDequantizePerChannelGroupOp>>( + Torch::QuantizedDecomposedDequantizePerChannelGroupOp>, + ConvertDequantize<Torch::QuantizedDecomposedDequantizePerTokenOp>>( typeConverter, context); // The PT2E ops are deliberately not marked illegal: ops whose metadata // LinalgExt cannot represent remain for the Torch-to-Linalg lowering.
diff --git a/compiler/plugins/input/Torch/InputConversion/Passes.td b/compiler/plugins/input/Torch/InputConversion/Passes.td index 6d98e53..55ba824 100644 --- a/compiler/plugins/input/Torch/InputConversion/Passes.td +++ b/compiler/plugins/input/Torch/InputConversion/Passes.td
@@ -33,11 +33,11 @@ InterfacePass<"torch-iree-torch-quantization-to-linalg-ext", "mlir::FunctionOpInterface"> { let summary = "Convert PT2E quantization operations to LinalgExt"; let description = [{ - Preserves PT2E per-tensor, per-channel, and per-channel-group quantization - as `iree_linalg_ext.quantize_affine` and `dequantize_affine` operations. - This pass runs before Torch-to-Linalg expands the PT2E operations into - generic arithmetic so later IREE optimization can still identify quantized - computations. + Preserves PT2E per-tensor, per-channel, per-channel-group, and per-token + quantization as `iree_linalg_ext.quantize_affine` and `dequantize_affine` + operations. This pass runs before Torch-to-Linalg expands the PT2E + operations into generic arithmetic so later IREE optimization can still + identify quantized computations. }]; }
diff --git a/compiler/plugins/input/Torch/InputConversion/test/quantization_to_linalg_ext.mlir b/compiler/plugins/input/Torch/InputConversion/test/quantization_to_linalg_ext.mlir index eb4afb0..6ca034f 100644 --- a/compiler/plugins/input/Torch/InputConversion/test/quantization_to_linalg_ext.mlir +++ b/compiler/plugins/input/Torch/InputConversion/test/quantization_to_linalg_ext.mlir
@@ -206,6 +206,49 @@ // ----- +// Each vector along the last dimension is a token with its own scale and zero +// point. PT2E keeps a unit last dimension in their shape, which is dropped so +// that LinalgExt indexes them by the leading dimensions. +// CHECK: #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK: #map1 = affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-LABEL: func.func @per_token( +// CHECK-NOT: tensor.expand_shape +// CHECK: %[[SCALES:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0], [1, 2]] +// CHECK-SAME: tensor<2x?x1xf64> into tensor<2x?xf64> +// CHECK: %[[ZERO_POINTS:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0], [1, 2]] +// CHECK-SAME: tensor<2x?x1xi64> into tensor<2x?xi64> +// CHECK: %[[QUANTIZED:.+]] = iree_linalg_ext.quantize_affine +// CHECK-SAME: indexing_maps = [#map, #map1, #map1, #map] +// CHECK-SAME: ins(%{{.+}}, %[[SCALES]], %[[ZERO_POINTS]] : tensor<2x?x8xf32>, +// CHECK-SAME: tensor<2x?xf64>, tensor<2x?xi64>) +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: indexing_maps = [#map, #map1, #map1, #map] +// CHECK-SAME: ins(%[[QUANTIZED]], %{{.+}}, %{{.+}} : tensor<2x?x8xi8>, +// CHECK-SAME: tensor<2x?xf64>, tensor<2x?xi64>) +func.func @per_token( + %input: !torch.vtensor<[2,?,8],f32>, + %scales: !torch.vtensor<[2,?,1],f64>, + %zero_points: !torch.vtensor<[2,?,1],si64>) -> !torch.vtensor<[2,?,8],f32> { + %minimum = torch.constant.int -128 + %maximum = torch.constant.int 127 + %dtype = torch.constant.int 1 + %out_dtype = torch.constant.int 6 + %quantized = torch.quantized_decomposed.quantize_per_token + %input, %scales, %zero_points, %minimum, %maximum, %dtype + : !torch.vtensor<[2,?,8],f32>, !torch.vtensor<[2,?,1],f64>, + !torch.vtensor<[2,?,1],si64>, !torch.int, !torch.int, !torch.int + -> !torch.vtensor<[2,?,8],si8> + %result = torch.quantized_decomposed.dequantize_per_token + %quantized, %scales, %zero_points, %minimum, %maximum, %dtype, + %out_dtype + : !torch.vtensor<[2,?,8],si8>, !torch.vtensor<[2,?,1],f64>, + !torch.vtensor<[2,?,1],si64>, !torch.int, !torch.int, !torch.int, + !torch.int -> !torch.vtensor<[2,?,8],f32> + return %result : !torch.vtensor<[2,?,8],f32> +} + +// ----- + // Dynamic bounds cannot be represented as LinalgExt attributes. The pass // leaves this op for the general Torch-to-Linalg lowering that follows it. // CHECK-LABEL: func.func @dynamic_bounds(