[Torch] Preserve PT2E per-channel-group quantization as LinalgExt ops (#25008) Extends `ConvertTorchQuantizationToLinalgExt` to the PT2E `quantized_decomposed` per-channel-group quantize/dequantize ops. All granularities, including per-tensor and per-channel, derive their layout from a torchao-style block size and share one operand check. Bool zero points are treated as unsigned. Part of #24862. --------- Signed-off-by: Roberto Laudani <laudani@roofline.ai>
diff --git a/compiler/plugins/input/Torch/InputConversion/BUILD.bazel b/compiler/plugins/input/Torch/InputConversion/BUILD.bazel index 2ed3b1c..1ea16d3 100644 --- a/compiler/plugins/input/Torch/InputConversion/BUILD.bazel +++ b/compiler/plugins/input/Torch/InputConversion/BUILD.bazel
@@ -85,6 +85,7 @@ "@llvm-project//mlir:SCFDialect", "@llvm-project//mlir:Support", "@llvm-project//mlir:TensorDialect", + "@llvm-project//mlir:TensorUtils", "@llvm-project//mlir:TransformUtils", "@llvm-project//mlir:Transforms", "@torch-mlir//:TorchMLIRConversionPasses",
diff --git a/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt b/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt index 4c325b1..604d945 100644 --- a/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt +++ b/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt
@@ -65,6 +65,7 @@ MLIRSCFDialect MLIRSupport MLIRTensorDialect + MLIRTensorUtils MLIRTransformUtils MLIRTransforms iree::compiler::Dialect::Flow::IR
diff --git a/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp b/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp index f4cd70a..790322a 100644 --- a/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp +++ b/compiler/plugins/input/Torch/InputConversion/ConvertTorchQuantizationToLinalgExt.cpp
@@ -10,6 +10,7 @@ #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Tensor/Utils/Utils.h" #include "mlir/IR/TypeUtilities.h" #include "mlir/Transforms/DialectConversion.h" #include "torch-mlir/Conversion/TorchToLinalg/Utils.h" @@ -30,23 +31,24 @@ namespace { // The affine quantization described by a PT2E quantize or dequantize op. It -// does not depend on the direction of the op, and per-tensor and per-channel -// quantization differ only in whether `axis` is set. +// does not depend on the direction of the op, and `blockSize` describes its +// granularity. struct PT2EQuantization { // Builtin tensor that is quantized or dequantized. Value input; - // Builtin scale: an f64 scalar for per-tensor quantization or a rank-one - // float tensor for per-channel quantization. + // Builtin scale: an f64 scalar for per-tensor quantization or a float tensor + // with one value per block. Value scale; - // Builtin zero point: an i64 scalar or a rank-one integer tensor; null for - // symmetric quantization. + // Builtin zero point: an i64 scalar or an integer tensor shaped like the + // scale; null for symmetric quantization. Value zeroPoint; // Whether the zero-point elements are unsigned. Read from the Torch type // because builtin integer types are signless. bool zeroPointUnsigned = false; - // Non-negative input dimension indexed by per-channel parameters; nullopt - // for per-tensor quantization. - std::optional<int64_t> axis; + // Number of consecutive input elements along each dimension that share + // parameters, as in torchao's `block_size`. A block spanning a dynamic + // dimension is `ShapedType::kDynamic`. + SmallVector<int64_t> blockSize; // Inclusive lower bound of the quantized value range. int64_t quantMin = 0; // Inclusive upper bound of the quantized value range. @@ -71,7 +73,8 @@ } // LinalgExt indexes per-channel parameters with an affine map, so the axis -// must be a Torch constant. Requires `quantization.input` to be set. +// must be a Torch constant. Sets `blockSize` to one along the axis and to the +// dimension size elsewhere. Requires `quantization.input` to be set. static LogicalResult matchConstantAxis(Operation *op, Value axis, PT2EQuantization &quantization, ConversionPatternRewriter &rewriter) { @@ -79,15 +82,51 @@ if (!matchPattern(axis, Torch::m_TorchConstantInt(&dimension))) { return rewriter.notifyMatchFailure(op, "expected a constant axis"); } - int64_t rank = cast<RankedTensorType>(quantization.input.getType()).getRank(); + auto inputType = cast<RankedTensorType>(quantization.input.getType()); + int64_t rank = inputType.getRank(); int64_t positiveDimension = Torch::toPositiveDim(dimension, rank); if (!Torch::isValidDim(positiveDimension, rank)) { return rewriter.notifyMatchFailure(op, "expected an axis within the rank"); } - quantization.axis = positiveDimension; + quantization.blockSize = llvm::to_vector(inputType.getShape()); + quantization.blockSize[positiveDimension] = 1; return success(); } +// LinalgExt indexes per-channel-group parameters on the input with its last +// dimension split into groups, so the group size must be a Torch constant that +// divides the static last dimension of a rank-2 input. Sets `blockSize` to one +// along the first dimension and to the group size along the last. Requires +// `quantization.input` to be set. +static LogicalResult +matchConstantGroupSize(Operation *op, Value groupSize, + PT2EQuantization &quantization, + ConversionPatternRewriter &rewriter) { + int64_t size; + if (!matchPattern(groupSize, Torch::m_TorchConstantInt(&size)) || size <= 1) { + return rewriter.notifyMatchFailure( + op, "expected a constant group size greater than one"); + } + auto inputType = cast<RankedTensorType>(quantization.input.getType()); + if (inputType.getRank() != 2 || inputType.isDynamicDim(1)) { + return rewriter.notifyMatchFailure( + op, "expected a rank-2 input with a static last dimension"); + } + if (inputType.getDimSize(1) % size != 0) { + return rewriter.notifyMatchFailure( + op, "expected groups that evenly split the last dimension"); + } + quantization.blockSize = {1, size}; + return success(); +} + +// 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) { + return torch::torch_to_linalg::isUnsignedTorchType(zeroPointType) || + cast<Torch::ValueTensorType>(zeroPointType).getDtype().isInteger(1); +} + // How each PT2E op class spells its operands is the only thing that differs // between the ops of one direction, so each class has its own reader. @@ -99,6 +138,8 @@ quantization.input = adaptor.getInput(); quantization.scale = adaptor.getScale(); quantization.zeroPoint = adaptor.getZeroPoint(); + quantization.blockSize = llvm::to_vector( + cast<RankedTensorType>(quantization.input.getType()).getShape()); if (failed(matchConstantBounds(op, op.getQuantMin(), op.getQuantMax(), quantization, rewriter))) { return failure(); @@ -114,6 +155,8 @@ quantization.input = adaptor.getInput(); quantization.scale = adaptor.getScale(); quantization.zeroPoint = adaptor.getZeroPoint(); + quantization.blockSize = llvm::to_vector( + cast<RankedTensorType>(quantization.input.getType()).getShape()); if (failed(matchConstantBounds(op, op.getQuantMin(), op.getQuantMax(), quantization, rewriter))) { return failure(); @@ -130,7 +173,7 @@ quantization.scale = adaptor.getScales(); quantization.zeroPoint = adaptor.getZeroPoints(); quantization.zeroPointUnsigned = - torch::torch_to_linalg::isUnsignedTorchType(op.getZeroPoints().getType()); + isUnsignedZeroPoint(op.getZeroPoints().getType()); if (failed(matchConstantAxis(op, op.getAxis(), quantization, rewriter)) || failed(matchConstantBounds(op, op.getQuantMin(), op.getQuantMax(), quantization, rewriter))) { @@ -150,8 +193,7 @@ if (!isa<Torch::NoneType>(op.getZeroPoints().getType())) { quantization.zeroPoint = adaptor.getZeroPoints(); quantization.zeroPointUnsigned = - torch::torch_to_linalg::isUnsignedTorchType( - op.getZeroPoints().getType()); + isUnsignedZeroPoint(op.getZeroPoints().getType()); } if (failed(matchConstantAxis(op, op.getAxis(), quantization, rewriter)) || failed(matchConstantBounds(op, op.getQuantMin(), op.getQuantMax(), @@ -161,6 +203,104 @@ return quantization; } +static FailureOr<PT2EQuantization> readPT2EQuantization( + Torch::QuantizedDecomposedQuantizePerChannelGroupOp op, + Torch::QuantizedDecomposedQuantizePerChannelGroupOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) { + PT2EQuantization quantization; + quantization.input = adaptor.getInput(); + quantization.scale = adaptor.getScales(); + quantization.zeroPoint = adaptor.getZeroPoints(); + quantization.zeroPointUnsigned = + isUnsignedZeroPoint(op.getZeroPoints().getType()); + if (failed(matchConstantGroupSize(op, op.getGroupSize(), quantization, + rewriter)) || + failed(matchConstantBounds(op, op.getQuantMin(), op.getQuantMax(), + quantization, rewriter))) { + return failure(); + } + return quantization; +} + +static FailureOr<PT2EQuantization> readPT2EQuantization( + Torch::QuantizedDecomposedDequantizePerChannelGroupOp op, + Torch::QuantizedDecomposedDequantizePerChannelGroupOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) { + PT2EQuantization quantization; + quantization.input = adaptor.getInput(); + quantization.scale = adaptor.getScales(); + // A None zero point denotes symmetric quantization. + if (!isa<Torch::NoneType>(op.getZeroPoints().getType())) { + quantization.zeroPoint = adaptor.getZeroPoints(); + quantization.zeroPointUnsigned = + isUnsignedZeroPoint(op.getZeroPoints().getType()); + } + if (failed(matchConstantGroupSize(op, op.getGroupSize(), quantization, + rewriter)) || + 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. +static SmallVector<ReassociationIndices> +getBlockedReassociation(const PT2EQuantization &quantization) { + ArrayRef<int64_t> inputShape = + cast<RankedTensorType>(quantization.input.getType()).getShape(); + SmallVector<ReassociationIndices> reassociation; + int64_t dim = 0; + for (auto [size, block] : + llvm::zip_equal(inputShape, quantization.blockSize)) { + if (block == 1 || block == size) { + reassociation.push_back({dim++}); + } else { + reassociation.push_back({dim, dim + 1}); + dim += 2; + } + } + return reassociation; +} + +// Returns the type the LinalgExt op works on: the input type, split as +// `getBlockedReassociation` describes. +static RankedTensorType getBlockedType(const PT2EQuantization &quantization) { + auto inputType = cast<RankedTensorType>(quantization.input.getType()); + SmallVector<int64_t> shape; + for (auto [dims, size, block] : + llvm::zip_equal(getBlockedReassociation(quantization), + inputType.getShape(), quantization.blockSize)) { + if (dims.size() == 1) { + shape.push_back(size); + } else { + shape.append({size / block, block}); + } + } + return inputType.clone(shape); +} + +// 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. +static SmallVector<unsigned> +getParameterDims(const PT2EQuantization &quantization) { + ArrayRef<int64_t> inputShape = + cast<RankedTensorType>(quantization.input.getType()).getShape(); + SmallVector<unsigned> parameterDims; + for (auto [dims, size, block] : + llvm::zip_equal(getBlockedReassociation(quantization), inputShape, + quantization.blockSize)) { + if (block != size) { + parameterDims.push_back(dims.front()); + } + } + return parameterDims; +} + // Checks whether LinalgExt can use the scale and zero point, reporting why // unsupported parameters prevent conversion. Per-tensor scalars need no // additional checks. `realType` is the element type of the floating-point @@ -169,7 +309,10 @@ checkScaleAndZeroPointSupport(Operation *op, Type realType, const PT2EQuantization &quantization, ConversionPatternRewriter &rewriter) { - if (!quantization.axis) { + // Only per-tensor scales are scalars. The block size alone cannot identify + // per-tensor quantization, since per-channel quantization along a dimension + // of size one has the same block size. + if (isa<FloatType>(quantization.scale.getType())) { return success(); } auto scaleType = dyn_cast<RankedTensorType>(quantization.scale.getType()); @@ -211,14 +354,9 @@ // point when present, and output. static ArrayAttr getIndexingMaps(Builder &builder, const PT2EQuantization &quantization) { - int64_t rank = cast<RankedTensorType>(quantization.input.getType()).getRank(); - // The scale and zero point are scalars for per-tensor quantization and are - // indexed by the channel dimension for per-channel quantization. - AffineMap parameterMap = - quantization.axis - ? AffineMap::get(rank, 0, - builder.getAffineDimExpr(*quantization.axis)) - : AffineMap::get(rank, 0, builder.getContext()); + int64_t rank = getBlockedType(quantization).getRank(); + AffineMap parameterMap = AffineMap::getMultiDimMapWithTargets( + rank, getParameterDims(quantization), builder.getContext()); AffineMap identity = builder.getMultiDimIdentityMap(rank); SmallVector<AffineMap> maps = {identity, parameterMap}; if (quantization.zeroPoint) { @@ -228,56 +366,109 @@ return builder.getAffineMapArrayAttr(maps); } -static Value createInit(OpBuilder &builder, Location location, - const PT2EQuantization &quantization, +// Returns an empty tensor shaped like `input`, with the element type of +// `resultType`. +static Value createInit(OpBuilder &builder, Location location, Value input, RankedTensorType resultType) { return tensor::EmptyOp::create( - builder, location, - tensor::getMixedSizes(builder, location, quantization.input), + builder, location, tensor::getMixedSizes(builder, location, input), resultType.getElementType()); } +// Returns `quantization.input`, expanded into the blocked type. +static Value createBlockedInput(OpBuilder &builder, Location location, + const PT2EQuantization &quantization) { + RankedTensorType blockedType = getBlockedType(quantization); + if (blockedType == quantization.input.getType()) { + return quantization.input; + } + return tensor::ExpandShapeOp::create(builder, location, blockedType, + quantization.input, + getBlockedReassociation(quantization)); +} + +// Collapses a `result` with the blocked shape back into the input shape. +static Value collapseBlockedResult(OpBuilder &builder, Location location, + Value result, + const PT2EQuantization &quantization) { + if (getBlockedType(quantization) == quantization.input.getType()) { + return result; + } + return tensor::CollapseShapeOp::create(builder, location, result, + getBlockedReassociation(quantization)); +} + +// Drops the static unit dimensions of a tensor `parameter`, since LinalgExt +// does not index them, and returns scalars and an absent zero point unchanged. +static Value collapseParameter(OpBuilder &builder, Location location, + Value parameter) { + auto parameterType = + parameter ? dyn_cast<RankedTensorType>(parameter.getType()) : nullptr; + if (!parameterType) { + return parameter; + } + llvm::SmallBitVector unitDims(parameterType.getRank()); + for (auto [dim, size] : llvm::enumerate(parameterType.getShape())) { + unitDims[dim] = size == 1; + } + if (unitDims.none()) { + return parameter; + } + return tensor::dropGivenUnitDims(builder, location, parameter, unitDims); +} + static Value createQuantizeAffine(OpBuilder &builder, Location location, RankedTensorType resultType, const PT2EQuantization &quantization, bool storageUnsigned) { // 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 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. + Value input = createBlockedInput(builder, location, quantization); Value scale = truncateScalarScaleToF32(builder, location, quantization.scale); + scale = collapseParameter(builder, location, scale); + Value zeroPoint = + collapseParameter(builder, location, quantization.zeroPoint); + Value init = createInit(builder, location, input, resultType); auto quantize = IREE::LinalgExt::QuantizeAffineOp::create( - builder, location, resultType, quantization.input, scale, - quantization.zeroPoint, - createInit(builder, location, quantization, resultType), + builder, location, init.getType(), input, scale, zeroPoint, init, getIndexingMaps(builder, quantization), builder.getI64IntegerAttr(quantization.quantMin), builder.getI64IntegerAttr(quantization.quantMax), storageUnsigned ? builder.getUnitAttr() : nullptr, quantization.zeroPointUnsigned ? builder.getUnitAttr() : nullptr); - return quantize->getResult(0); + return collapseBlockedResult(builder, location, quantize->getResult(0), + quantization); } static Value createDequantizeAffine(OpBuilder &builder, Location location, RankedTensorType resultType, const PT2EQuantization &quantization, bool inputUnsigned) { + Value input = createBlockedInput(builder, location, quantization); Value scale = truncateScalarScaleToF32(builder, location, quantization.scale); + scale = collapseParameter(builder, location, scale); + Value zeroPoint = + collapseParameter(builder, location, quantization.zeroPoint); + Value init = createInit(builder, location, input, resultType); auto dequantize = IREE::LinalgExt::DequantizeAffineOp::create( - builder, location, resultType, quantization.input, scale, - quantization.zeroPoint, - createInit(builder, location, quantization, resultType), + builder, location, init.getType(), input, scale, zeroPoint, init, getIndexingMaps(builder, quantization), builder.getI64IntegerAttr(quantization.quantMin), builder.getI64IntegerAttr(quantization.quantMax), inputUnsigned ? builder.getUnitAttr() : nullptr, quantization.zeroPointUnsigned ? builder.getUnitAttr() : nullptr); - return dequantize->getResult(0); + return collapseBlockedResult(builder, location, dequantize->getResult(0), + quantization); } namespace { -// Converts `quantize_per_tensor` or `quantize_per_channel`, whose input holds -// the real values and whose result holds the quantized storage. +// 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. template <typename OpTy> struct ConvertQuantize final : OpConversionPattern<OpTy> { using OpConversionPattern<OpTy>::OpConversionPattern; @@ -320,8 +511,9 @@ } }; -// Converts `dequantize_per_tensor` or `dequantize_per_channel`, whose input -// holds the quantized storage and whose result holds the real values. +// 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. template <typename OpTy> struct ConvertDequantize final : OpConversionPattern<OpTy> { using OpConversionPattern<OpTy>::OpConversionPattern; @@ -389,8 +581,11 @@ patterns.add< ConvertQuantize<Torch::QuantizedDecomposedQuantizePerTensorOp>, ConvertQuantize<Torch::QuantizedDecomposedQuantizePerChannelOp>, + ConvertQuantize<Torch::QuantizedDecomposedQuantizePerChannelGroupOp>, ConvertDequantize<Torch::QuantizedDecomposedDequantizePerTensorOp>, - ConvertDequantize<Torch::QuantizedDecomposedDequantizePerChannelOp>>( + ConvertDequantize<Torch::QuantizedDecomposedDequantizePerChannelOp>, + ConvertDequantize< + Torch::QuantizedDecomposedDequantizePerChannelGroupOp>>( 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 828c5d9..6d98e53 100644 --- a/compiler/plugins/input/Torch/InputConversion/Passes.td +++ b/compiler/plugins/input/Torch/InputConversion/Passes.td
@@ -33,10 +33,10 @@ InterfacePass<"torch-iree-torch-quantization-to-linalg-ext", "mlir::FunctionOpInterface"> { let summary = "Convert PT2E quantization operations to LinalgExt"; let description = [{ - Preserves PT2E per-tensor and per-channel 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 + 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. }]; }
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 2ee9add..eb4afb0 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
@@ -101,6 +101,111 @@ // ----- +// The parameter maps of a single channel index no dimension, so the unit scale +// is collapsed to a rank-0 tensor. +// CHECK: #map = affine_map<(d0, d1) -> (d0, d1)> +// CHECK: #map1 = affine_map<(d0, d1) -> ()> +// CHECK-LABEL: func.func @single_channel( +// CHECK: %[[SCALE:.+]] = tensor.collapse_shape %{{.+}} [] +// CHECK-SAME: tensor<1xf32> into tensor<f32> +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: indexing_maps = [#map, #map1, #map] +// CHECK-SAME: ins(%{{.+}}, %[[SCALE]] : tensor<1x8xi8>, tensor<f32>) +func.func @single_channel( + %input: !torch.vtensor<[1,8],si8>, + %scales: !torch.vtensor<[1],f32>) -> !torch.vtensor<[1,8],f32> { + %axis = torch.constant.int 0 + %minimum = torch.constant.int -128 + %maximum = torch.constant.int 127 + %dtype = torch.constant.int 1 + %none = torch.constant.none + %out_dtype = torch.derefine %none : !torch.none to !torch.optional<int> + %result = torch.quantized_decomposed.dequantize_per_channel + %input, %scales, %none, %axis, %minimum, %maximum, %dtype, %out_dtype + : !torch.vtensor<[1,8],si8>, !torch.vtensor<[1],f32>, !torch.none, + !torch.int, !torch.int, !torch.int, !torch.int, !torch.optional<int> + -> !torch.vtensor<[1,8],f32> + return %result : !torch.vtensor<[1,8],f32> +} + +// ----- + +// CHECK: #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK: #map1 = affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-LABEL: func.func @per_channel_group( +// CHECK: %[[BLOCKED_INPUT:.+]] = tensor.expand_shape %{{.+}} {{\[}}[0], [1, 2]] +// CHECK-SAME: output_shape [4, 2, 4] : tensor<4x8xf32> into tensor<4x2x4xf32> +// CHECK: %[[QUANTIZED:.+]] = iree_linalg_ext.quantize_affine +// CHECK-SAME: indexing_maps = [#map, #map1, #map1, #map] +// CHECK-SAME: ins(%[[BLOCKED_INPUT]], %{{.+}}, %{{.+}} : tensor<4x2x4xf32>, +// CHECK-SAME: tensor<4x2xf32>, tensor<4x2xi8>) +// CHECK: %[[STORAGE:.+]] = tensor.collapse_shape %[[QUANTIZED]] +// CHECK-SAME: {{\[}}[0], [1, 2]] : tensor<4x2x4xi8> into tensor<4x8xi8> +// CHECK: %[[BLOCKED_STORAGE:.+]] = tensor.expand_shape %[[STORAGE]] +// CHECK-SAME: {{\[}}[0], [1, 2]] output_shape [4, 2, 4] +// CHECK-SAME: tensor<4x8xi8> into tensor<4x2x4xi8> +// CHECK: %[[DEQUANTIZED:.+]] = iree_linalg_ext.dequantize_affine +// CHECK-SAME: indexing_maps = [#map, #map1, #map1, #map] +// CHECK-SAME: ins(%[[BLOCKED_STORAGE]], %{{.+}}, %{{.+}} : tensor<4x2x4xi8>, +// CHECK-SAME: tensor<4x2xf32>, tensor<4x2xi8>) +// CHECK: tensor.collapse_shape %[[DEQUANTIZED]] +// CHECK-SAME: {{\[}}[0], [1, 2]] : tensor<4x2x4xf32> into tensor<4x8xf32> +func.func @per_channel_group( + %input: !torch.vtensor<[4,8],f32>, + %scales: !torch.vtensor<[4,2],f32>, + %zero_points: !torch.vtensor<[4,2],si8>) -> !torch.vtensor<[4,8],f32> { + %minimum = torch.constant.int -128 + %maximum = torch.constant.int 127 + %dtype = torch.constant.int 1 + %group_size = torch.constant.int 4 + %out_dtype = torch.constant.int 6 + %quantized = torch.quantized_decomposed.quantize_per_channel_group + %input, %scales, %zero_points, %minimum, %maximum, %dtype, %group_size + : !torch.vtensor<[4,8],f32>, !torch.vtensor<[4,2],f32>, + !torch.vtensor<[4,2],si8>, !torch.int, !torch.int, !torch.int, + !torch.int -> !torch.vtensor<[4,8],si8> + %result = torch.quantized_decomposed.dequantize_per_channel_group + %quantized, %scales, %zero_points, %minimum, %maximum, %dtype, + %group_size, %out_dtype + : !torch.vtensor<[4,8],si8>, !torch.vtensor<[4,2],f32>, + !torch.vtensor<[4,2],si8>, !torch.int, !torch.int, !torch.int, + !torch.int, !torch.int -> !torch.vtensor<[4,8],f32> + return %result : !torch.vtensor<[4,8],f32> +} + +// ----- + +// A group spanning the whole row has one scale per row, so the scales are +// collapsed and indexed by the row. +// CHECK: #map = affine_map<(d0, d1) -> (d0, d1)> +// CHECK: #map1 = affine_map<(d0, d1) -> (d0)> +// CHECK-LABEL: func.func @group_spanning_row( +// CHECK-NOT: tensor.expand_shape +// CHECK: %[[SCALES:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0, 1]] +// CHECK-SAME: tensor<4x1xf32> into tensor<4xf32> +// CHECK: iree_linalg_ext.dequantize_affine +// CHECK-SAME: indexing_maps = [#map, #map1, #map] +// CHECK-SAME: ins(%{{.+}}, %[[SCALES]] : tensor<4x8xi8>, tensor<4xf32>) +func.func @group_spanning_row( + %input: !torch.vtensor<[4,8],si8>, + %scales: !torch.vtensor<[4,1],f32>) -> !torch.vtensor<[4,8],f32> { + %minimum = torch.constant.int -128 + %maximum = torch.constant.int 127 + %dtype = torch.constant.int 1 + %group_size = torch.constant.int 8 + %none = torch.constant.none + %out_dtype = torch.constant.int 6 + %result = torch.quantized_decomposed.dequantize_per_channel_group + %input, %scales, %none, %minimum, %maximum, %dtype, %group_size, + %out_dtype + : !torch.vtensor<[4,8],si8>, !torch.vtensor<[4,1],f32>, !torch.none, + !torch.int, !torch.int, !torch.int, !torch.int, !torch.int + -> !torch.vtensor<[4,8],f32> + return %result : !torch.vtensor<[4,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( @@ -146,6 +251,31 @@ // ----- +// A group size that does not divide the last dimension cannot split the input +// into groups, so the op is left for the Torch-to-Linalg lowering. +// CHECK-LABEL: func.func @indivisible_group_size( +// CHECK: torch.quantized_decomposed.dequantize_per_channel_group +// CHECK-NOT: iree_linalg_ext.dequantize_affine +func.func @indivisible_group_size( + %input: !torch.vtensor<[4,8],si8>, + %scales: !torch.vtensor<[4,2],f32>) -> !torch.vtensor<[4,8],f32> { + %minimum = torch.constant.int -128 + %maximum = torch.constant.int 127 + %dtype = torch.constant.int 1 + %group_size = torch.constant.int 3 + %none = torch.constant.none + %out_dtype = torch.constant.int 6 + %result = torch.quantized_decomposed.dequantize_per_channel_group + %input, %scales, %none, %minimum, %maximum, %dtype, %group_size, + %out_dtype + : !torch.vtensor<[4,8],si8>, !torch.vtensor<[4,2],f32>, !torch.none, + !torch.int, !torch.int, !torch.int, !torch.int, !torch.int + -> !torch.vtensor<[4,8],f32> + return %result : !torch.vtensor<[4,8],f32> +} + +// ----- + // LinalgExt storage is a signless integer, so fp8 storage is left for the // Torch-to-Linalg lowering. // CHECK-LABEL: func.func @float_storage(