| //===- LowerQuantOps.cpp - Lower 'quant' dialect ops ----------------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Transforms `quant.dcast` and `quant.qcast` into lower-level ops. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Dialect/Func/IR/FuncOps.h" |
| #include "mlir/Dialect/Linalg/IR/Linalg.h" |
| #include "mlir/Dialect/Quant/IR/Quant.h" |
| #include "mlir/Dialect/Quant/IR/QuantTypes.h" |
| #include "mlir/Dialect/Quant/Transforms/Passes.h" |
| #include "mlir/Dialect/Shape/IR/Shape.h" |
| #include "mlir/Dialect/Tensor/IR/Tensor.h" |
| #include "mlir/IR/Matchers.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| |
| namespace mlir { |
| namespace quant { |
| |
| #define GEN_PASS_DEF_LOWERQUANTOPS |
| #include "mlir/Dialect/Quant/Transforms/Passes.h.inc" |
| |
| namespace { |
| |
| // If 'inputType' is a tensor, return its element type. If it is a scalar, |
| // return it as is. |
| Type getScalarType(Type inputType) { |
| if (auto tensorType = dyn_cast<TensorType>(inputType)) |
| return tensorType.getElementType(); |
| return inputType; |
| } |
| |
| // Return the shape of an input value as a list of attributes (static |
| // dimensions) and values (dynamic dimensions). If 'input' is a scalar, an empty |
| // list is returned. If 'input' is a tensor, its shape is returned. |
| SmallVector<OpFoldResult> getScalarOrTensorShape(OpBuilder &builder, |
| Location loc, Value input) { |
| if (isa<TensorType>(input.getType())) |
| return tensor::getMixedSizes(builder, loc, input); |
| return {}; |
| } |
| |
| // If 'referenceType' is a scalar, return 'elementType' as is. If |
| // 'referenceType' is a tensor, return another tensor with the same shape and |
| // elements of type 'elementType'. |
| Type getScalarOrTensorType(Type elementType, Type referenceType) { |
| if (auto tensorType = dyn_cast<TensorType>(referenceType)) |
| return tensorType.clone(elementType); |
| return elementType; |
| } |
| |
| // Return a constant with the given value. If 'referenceType' is a tensor, a |
| // tensor splat of shape 'referenceShape' is returned. If 'referenceType' is a |
| // scalar, 'referenceShape' is ignored and a scalar constant is returned. |
| Value getScalarOrTensorConstant(OpBuilder &builder, Location loc, Value scalar, |
| Type referenceType, |
| ArrayRef<OpFoldResult> referenceShape) { |
| // If the result type is a scalar, return the unmodified scalar constant. |
| auto tensorType = dyn_cast<TensorType>(referenceType); |
| if (!tensorType) { |
| assert(referenceShape.empty()); |
| return scalar; |
| } |
| |
| // Create tensor splat |
| auto tensorConstant = |
| tensor::SplatOp::create(builder, loc, scalar, referenceShape); |
| return tensorConstant; |
| } |
| |
| // Reshape an unranked tensor into a 1D ranked tensor. |
| // |
| // - input |
| // Unranked tensor. |
| // |
| // Return values: |
| // |
| // - flatInput |
| // 1D ranked, dynamically shaped tensor. |
| // |
| // - inputShape |
| // 1D extent tensor containing the shape of the original unranked input. |
| // |
| std::pair<Value, Value> flattenUnrankedTensor(OpBuilder &builder, Location loc, |
| Value input) { |
| // Get unranked input shape and total size |
| auto *context = builder.getContext(); |
| auto shapeType = shape::getExtentTensorType(context); |
| auto inputShape = shape::ShapeOfOp::create(builder, loc, shapeType, input); |
| Value inputSize = shape::NumElementsOp::create( |
| builder, loc, builder.getIndexType(), inputShape); |
| |
| // Turn input size into 1D tensor |
| auto flatShapeType = shape::getExtentTensorType(context, 1); |
| auto flatInputShape = |
| tensor::FromElementsOp::create(builder, loc, flatShapeType, inputSize); |
| |
| // Reshape input tensor into 1D |
| auto inputType = cast<UnrankedTensorType>(input.getType()); |
| auto elementType = inputType.getElementType(); |
| auto flatInputType = |
| RankedTensorType::get({ShapedType::kDynamic}, elementType); |
| auto flatInput = tensor::ReshapeOp::create(builder, loc, flatInputType, input, |
| flatInputShape); |
| return std::make_pair(flatInput, inputShape); |
| } |
| |
| // Reshape an unranked tensor into a 3D ranked tensor where the central |
| // dimension of the result tensor corresponds to dimension 'axis' of the input |
| // tensor. |
| // |
| // - input |
| // Unranked tensor. |
| // |
| // - axis |
| // Index of the input dimension around which other input dimiensions will be |
| // collapsed. |
| // |
| // - axisSize |
| // Size of input dimension 'axis'. |
| // |
| // Return values: |
| // |
| // - flatInput |
| // 3D ranked tensor of shape [?, axisSize, ?]. |
| // |
| // - inputShape |
| // 1D extent tensor containing the shape of the original unranked input. |
| // |
| std::pair<Value, Value> |
| flattenUnrankedTensorAroundAxis(OpBuilder &builder, Location loc, Value input, |
| int64_t axis, int64_t axisSize) { |
| // Get full tensor shape |
| auto *context = builder.getContext(); |
| auto indexType = builder.getIndexType(); |
| auto shapeType = shape::getExtentTensorType(context); |
| auto inputShape = shape::ShapeOfOp::create(builder, loc, shapeType, input); |
| |
| // Get shape and sizes on left and right of axis |
| auto axisValue = arith::ConstantIndexOp::create(builder, loc, axis); |
| auto axisNextValue = arith::ConstantIndexOp::create(builder, loc, axis + 1); |
| auto shapeLeft = |
| builder |
| .create<shape::SplitAtOp>(loc, TypeRange{shapeType, shapeType}, |
| inputShape, axisValue) |
| .getResult(0); |
| auto sizeLeft = |
| shape::NumElementsOp::create(builder, loc, indexType, shapeLeft); |
| auto shapeRight = |
| builder |
| .create<shape::SplitAtOp>(loc, TypeRange{shapeType, shapeType}, |
| inputShape, axisNextValue) |
| .getResult(1); |
| auto sizeRight = |
| shape::NumElementsOp::create(builder, loc, indexType, shapeRight); |
| |
| // Compute flat input shape as a 3-element 1D tensor |
| auto axisSizeValue = arith::ConstantIndexOp::create(builder, loc, axisSize); |
| auto flatShapeType = shape::getExtentTensorType(context, 3); |
| auto flatInputShape = tensor::FromElementsOp::create( |
| builder, loc, flatShapeType, |
| ValueRange{sizeLeft, axisSizeValue, sizeRight}); |
| |
| // Reshape input to 3D tensor |
| auto inputType = cast<UnrankedTensorType>(input.getType()); |
| auto elementType = inputType.getElementType(); |
| auto flatInputType = RankedTensorType::get( |
| {ShapedType::kDynamic, axisSize, ShapedType::kDynamic}, elementType); |
| auto flatInput = tensor::ReshapeOp::create(builder, loc, flatInputType, input, |
| flatInputShape); |
| |
| return std::make_pair(flatInput, inputShape); |
| } |
| |
| // Reshape an input tensor into its original unranked shape. |
| // |
| // - input |
| // Ranked tensor. |
| // |
| // - inputShape |
| // 1D extent tensor. |
| // |
| Value restoreUnrankedTensorShape(OpBuilder &builder, Location loc, Value input, |
| Value inputShape) { |
| auto inputType = cast<RankedTensorType>(input.getType()); |
| auto elementType = inputType.getElementType(); |
| auto unrankedType = UnrankedTensorType::get(elementType); |
| return tensor::ReshapeOp::create(builder, loc, unrankedType, input, |
| inputShape); |
| } |
| |
| // Create a tensor constant containing all scales in a per-channel quantized |
| // type. Example: |
| // |
| // !quant.uniform<i8:f32:1, {2.0:10, 3.0:20}> |
| // |
| // produces |
| // |
| // %cst = arith.constant dense<[2.0, 3.0]> : tensor<2xf32> |
| // |
| Value materializePerChannelScales(OpBuilder &builder, Location loc, |
| UniformQuantizedPerAxisType quantizedType) { |
| auto scales = quantizedType.getScales(); |
| auto expressedType = quantizedType.getExpressedType(); |
| auto scaleAttrs = llvm::map_to_vector(scales, [&](double scale) -> Attribute { |
| return builder.getFloatAttr(expressedType, scale); |
| }); |
| auto tensorType = |
| RankedTensorType::get({(int64_t)scales.size()}, expressedType); |
| auto scalesAttr = DenseElementsAttr::get(tensorType, scaleAttrs); |
| return arith::ConstantOp::create(builder, loc, tensorType, scalesAttr); |
| } |
| |
| // Create a tensor constant containing all zero points in a per-channel |
| // quantized type. Example: |
| // |
| // !quant.uniform<i8:f32:1, {2.0:10, 3.0:20}> |
| // |
| // produces |
| // |
| // %cst = arith.constant dense<[10, 20]> : tensor<2xi8> |
| // |
| Value materializePerChannelZeroPoints( |
| OpBuilder &builder, Location loc, |
| UniformQuantizedPerAxisType quantizedType) { |
| auto zeroPoints = quantizedType.getZeroPoints(); |
| auto storageType = quantizedType.getStorageType(); |
| auto zeroPointAttrs = |
| llvm::map_to_vector(zeroPoints, [&](int64_t zeroPoint) -> Attribute { |
| return builder.getIntegerAttr(storageType, zeroPoint); |
| }); |
| auto tensorType = |
| RankedTensorType::get({(int64_t)zeroPoints.size()}, storageType); |
| auto zeroPointsAttr = DenseElementsAttr::get(tensorType, zeroPointAttrs); |
| return arith::ConstantOp::create(builder, loc, tensorType, zeroPointsAttr); |
| } |
| |
| // Create a tensor constant containing all scales in a sub-channel quantized |
| // type. Example: |
| // |
| // !quant.uniform<i8:f32:{0:1,1:2}, {{2.0:10, 3.0:20}, {4.0:30, 5.0:40}}> |
| // |
| // produces |
| // |
| // %cst = arith.constant dense<[[2.0, 3.0], [4.0, 5.0]]> : tensor<2x2xf32> |
| // |
| Value materializeSubChannelScales( |
| OpBuilder &builder, Location loc, |
| UniformQuantizedSubChannelType quantizedType) { |
| auto scales = quantizedType.getScales(); |
| auto expressedType = quantizedType.getExpressedType(); |
| auto scaleAttrs = llvm::map_to_vector( |
| scales.getValues<APFloat>(), [&](APFloat scale) -> Attribute { |
| return builder.getFloatAttr(expressedType, scale); |
| }); |
| auto tensorType = |
| RankedTensorType::get(scales.getType().getShape(), expressedType); |
| auto scalesAttr = DenseElementsAttr::get(tensorType, scaleAttrs); |
| return arith::ConstantOp::create(builder, loc, tensorType, scalesAttr); |
| } |
| |
| // Create a tensor constant containing all zero points in a sub-channel |
| // quantized type. Example: |
| // |
| // !quant.uniform<i8:f32:{0:1,1:2}, {{2.0:10, 3.0:20}, {4.0:30, 5.0:40}}> |
| // |
| // produces |
| // |
| // %cst = arith.constant dense<[[10, 20], [30, 40]]> : tensor<2x2xi8> |
| // |
| Value materializeSubChannelZeroPoints( |
| OpBuilder &builder, Location loc, |
| UniformQuantizedSubChannelType quantizedType) { |
| auto zeroPoints = quantizedType.getZeroPoints(); |
| auto storageType = quantizedType.getStorageType(); |
| auto zeroPointAttrs = llvm::map_to_vector( |
| zeroPoints.getValues<APInt>(), [&](APInt zeroPoint) -> Attribute { |
| return builder.getIntegerAttr(storageType, zeroPoint); |
| }); |
| auto tensorType = |
| RankedTensorType::get(zeroPoints.getType().getShape(), storageType); |
| auto zeroPointsAttr = DenseElementsAttr::get(tensorType, zeroPointAttrs); |
| return arith::ConstantOp::create(builder, loc, tensorType, zeroPointsAttr); |
| } |
| |
| // Clamp the given scalar or tensor input using the storage bounds encoded in |
| // the given quantized type, if present. |
| // |
| // - input |
| // Scalar or ranked tensor input. The element type must match the storage type |
| // of 'quantizedType'. |
| // |
| // - inputShape |
| // If 'input' is a tensor, combination of attributes/values representing its |
| // static/dynamic dimensions. If 'input' is a scalar, empty list. |
| // |
| // - quantizedType |
| // Per-axis or per-channel quantized type. |
| Value clampScalarOrTensor(OpBuilder &builder, Location loc, Value input, |
| ArrayRef<OpFoldResult> inputShape, |
| QuantizedType quantizedType) { |
| // If quantized type does not narrow down the storage type range, there is |
| // nothing to do. |
| if (!quantizedType.hasStorageTypeBounds()) |
| return input; |
| |
| // Materialize bounds |
| auto inputType = input.getType(); |
| auto storageType = quantizedType.getStorageType(); |
| auto storageMinScalar = arith::ConstantIntOp::create( |
| builder, loc, storageType, quantizedType.getStorageTypeMin()); |
| auto storageMaxScalar = arith::ConstantIntOp::create( |
| builder, loc, storageType, quantizedType.getStorageTypeMax()); |
| auto storageMin = getScalarOrTensorConstant(builder, loc, storageMinScalar, |
| inputType, inputShape); |
| auto storageMax = getScalarOrTensorConstant(builder, loc, storageMaxScalar, |
| inputType, inputShape); |
| |
| // Clamp |
| if (quantizedType.isSigned()) { |
| input = arith::MaxSIOp::create(builder, loc, input, storageMin); |
| input = arith::MinSIOp::create(builder, loc, input, storageMax); |
| } else { |
| input = arith::MaxUIOp::create(builder, loc, input, storageMin); |
| input = arith::MinUIOp::create(builder, loc, input, storageMax); |
| } |
| return input; |
| } |
| |
| // Emit op 'arith.fptosi' or 'arith.fptoui'. |
| Value convertFloatToInteger(OpBuilder &builder, Location loc, Value input, |
| Type resultType, bool isSigned) { |
| if (isSigned) |
| return arith::FPToSIOp::create(builder, loc, resultType, input); |
| return arith::FPToUIOp::create(builder, loc, resultType, input); |
| } |
| |
| // Emit op 'arith.sitofp' or 'arith.uitofp'. |
| Value convertIntegerToFloat(OpBuilder &builder, Location loc, Value input, |
| Type resultType, bool isSigned) { |
| if (isSigned) |
| return arith::SIToFPOp::create(builder, loc, resultType, input); |
| return arith::UIToFPOp::create(builder, loc, resultType, input); |
| } |
| |
| // Quantize a scalar or ranked tensor value. The stored value is clamped using |
| // the storage bounds encoded in the given quantized type. |
| // |
| // See function 'convertRanked()' below for a description of the arguments. |
| Value quantizeValue(OpBuilder &builder, Location loc, Value input, |
| ArrayRef<OpFoldResult> inputShape, Value scale, |
| Value zeroPoint, QuantizedType quantizedType) { |
| // Convert scale to tensor if necessary |
| auto inputType = input.getType(); |
| scale = getScalarOrTensorConstant(builder, loc, scale, inputType, inputShape); |
| |
| // Scale input |
| auto scaledValue = arith::DivFOp::create(builder, loc, input, scale); |
| |
| // Skip unnecessary computations if no zero point is given |
| Value storedValueFloat = scaledValue; |
| if (!matchPattern(zeroPoint, m_Zero())) { |
| // Convert zero point to tensor if necessary |
| zeroPoint = getScalarOrTensorConstant(builder, loc, zeroPoint, inputType, |
| inputShape); |
| |
| // Convert zero point from storage to expressed type |
| zeroPoint = convertIntegerToFloat(builder, loc, zeroPoint, scale.getType(), |
| quantizedType.isSigned()); |
| |
| // Add zero point to stored value |
| storedValueFloat = |
| arith::AddFOp::create(builder, loc, scaledValue, zeroPoint); |
| } |
| |
| // Convert stored value to storage type |
| auto storageScalarOrTensorType = |
| getScalarOrTensorType(quantizedType.getStorageType(), inputType); |
| auto storedValueInt = convertFloatToInteger(builder, loc, storedValueFloat, |
| storageScalarOrTensorType, |
| quantizedType.isSigned()); |
| |
| // Clamp stored value it if the storage type is bound |
| auto storedValueClamped = clampScalarOrTensor(builder, loc, storedValueInt, |
| inputShape, quantizedType); |
| return storedValueClamped; |
| } |
| |
| // Dequantize a scalar or ranked tensor input. |
| // |
| // See function 'convertRanked()' below for a description of the arguments. |
| Value dequantizeValue(OpBuilder &builder, Location loc, Value input, |
| ArrayRef<OpFoldResult> inputShape, Value scale, |
| Value zeroPoint, QuantizedType quantizedType) { |
| // Convert scale to tensor if necessary |
| auto inputType = input.getType(); |
| scale = getScalarOrTensorConstant(builder, loc, scale, inputType, inputShape); |
| |
| // Convert stored value to float |
| auto result = convertIntegerToFloat(builder, loc, input, scale.getType(), |
| quantizedType.isSigned()); |
| |
| // Skip unnecessary computations if no zero point is given |
| if (!matchPattern(zeroPoint, m_Zero())) { |
| // Convert zero point to tensor if necessary |
| zeroPoint = getScalarOrTensorConstant(builder, loc, zeroPoint, inputType, |
| inputShape); |
| |
| // Convert zero point from storage to expressed type |
| zeroPoint = convertIntegerToFloat(builder, loc, zeroPoint, scale.getType(), |
| quantizedType.isSigned()); |
| |
| // Subtract zero point to stored value |
| result = arith::SubFOp::create(builder, loc, result, zeroPoint); |
| } |
| |
| // Multiply by scale |
| result = arith::MulFOp::create(builder, loc, result, scale); |
| return result; |
| } |
| |
| // Convert a scalar or ranked tensor input with the given scale and zero point |
| // values. |
| // |
| // - input |
| // Scalar or ranked tensor value. |
| // |
| // - inputShape |
| // If 'input' is a tensor, combination or attributes/values representing its |
| // static/dynamic dimensions. If 'input' is a scalar, empty list. |
| // |
| // - scale |
| // Scale as a floating-point scalar value. |
| // |
| // - zeroPoint |
| // Zero point as an integer scalar value. |
| // |
| // - quantizedType |
| // Scalar quantized type of the result ('quant.qcast') or of the input |
| // ('quant.dcast'). |
| // |
| Value convertRanked(OpBuilder &builder, Location loc, Operation *op, |
| Value input, ArrayRef<OpFoldResult> inputShape, Value scale, |
| Value zeroPoint, QuantizedType quantizedType) { |
| if (isa<QuantizeCastOp>(op)) |
| return quantizeValue(builder, loc, input, inputShape, scale, zeroPoint, |
| quantizedType); |
| if (isa<DequantizeCastOp>(op)) |
| return dequantizeValue(builder, loc, input, inputShape, scale, zeroPoint, |
| quantizedType); |
| llvm_unreachable("unexpected quant op"); |
| } |
| |
| // Convert an operation using per-layer quantization with a scalar or ranked |
| // tensor input. |
| // |
| // - op |
| // 'quant.dcast' or 'quant.qcast' op. |
| // |
| // - input |
| // Scalar or ranked tensor. |
| // |
| // - quantizedType |
| // Per-layer quantized type. |
| // |
| Value convertPerLayerRanked(OpBuilder &builder, Location loc, Operation *op, |
| Value input, UniformQuantizedType quantizedType) { |
| // Create scale and zero point constants |
| auto expressedType = quantizedType.getExpressedType(); |
| auto storageType = quantizedType.getStorageType(); |
| auto scaleAttr = |
| builder.getFloatAttr(expressedType, quantizedType.getScale()); |
| auto scale = |
| arith::ConstantOp::create(builder, loc, expressedType, scaleAttr); |
| auto zeroPointAttr = |
| builder.getIntegerAttr(storageType, quantizedType.getZeroPoint()); |
| auto zeroPoint = |
| arith::ConstantOp::create(builder, loc, storageType, zeroPointAttr); |
| |
| auto inputShape = getScalarOrTensorShape(builder, loc, input); |
| return convertRanked(builder, loc, op, input, inputShape, scale, zeroPoint, |
| quantizedType); |
| } |
| |
| // Convert an operation using per-layer quantization. |
| // |
| // - op |
| // 'quant.dcast' or 'quant.qcast' op. |
| // |
| // - input |
| // Scalar, ranked tensor, or unranked tensor. |
| // |
| // - quantizedType |
| // Per-layer quantized type. |
| // |
| Value convertPerLayer(OpBuilder &builder, Location loc, Operation *op, |
| Value input, UniformQuantizedType quantizedType) { |
| // Flatten input if unranked |
| bool isUnranked = isa<UnrankedTensorType>(input.getType()); |
| Value inputShape; |
| if (isUnranked) |
| std::tie(input, inputShape) = flattenUnrankedTensor(builder, loc, input); |
| |
| // Process ranked tensor |
| auto result = convertPerLayerRanked(builder, loc, op, input, quantizedType); |
| |
| // Restore original shape if unranked |
| if (isUnranked) |
| result = restoreUnrankedTensorShape(builder, loc, result, inputShape); |
| |
| return result; |
| } |
| |
| // Convert an operation using per-channel quantization and a scalar or ranked |
| // tensor as an input. |
| // |
| // - op |
| // 'quant.dcast' or 'quant.qcast' op. |
| // |
| // - input |
| // Scalar or ranked tensor. |
| // |
| // - quantizedType |
| // Per-channel quantized type. |
| // |
| Value convertPerChannelRanked(OpBuilder &builder, Location loc, Operation *op, |
| Value input, |
| UniformQuantizedPerAxisType quantizedType, |
| int64_t channelAxis) { |
| auto *context = builder.getContext(); |
| |
| auto inputType = cast<RankedTensorType>(input.getType()); |
| auto inputRank = inputType.getRank(); |
| |
| auto scales = materializePerChannelScales(builder, loc, quantizedType); |
| auto zeroPoints = |
| materializePerChannelZeroPoints(builder, loc, quantizedType); |
| |
| auto elementType = isa<FloatType>(inputType.getElementType()) |
| ? quantizedType.getStorageType() |
| : quantizedType.getExpressedType(); |
| auto initShape = tensor::getMixedSizes(builder, loc, input); |
| Value init = tensor::EmptyOp::create(builder, loc, initShape, elementType); |
| |
| SmallVector<utils::IteratorType> iteratorTypes(inputRank, |
| utils::IteratorType::parallel); |
| auto channelAxisAffineMap = AffineMap::get( |
| inputRank, 0, builder.getAffineDimExpr(channelAxis), context); |
| SmallVector<AffineMap> indexingMaps{ |
| builder.getMultiDimIdentityMap(inputRank), channelAxisAffineMap, |
| channelAxisAffineMap, builder.getMultiDimIdentityMap(inputRank)}; |
| auto result = builder |
| .create<linalg::GenericOp>( |
| loc, |
| init.getType(), // resultType |
| ValueRange{input, scales, zeroPoints}, // inputs |
| ValueRange{init}, // outputs |
| indexingMaps, iteratorTypes, |
| [&](OpBuilder &builder, Location loc, ValueRange args) { |
| assert(args.size() == 4); |
| auto input = args[0]; |
| auto scale = args[1]; |
| auto zeroPoint = args[2]; |
| |
| auto result = |
| convertRanked(builder, loc, op, input, {}, scale, |
| zeroPoint, quantizedType); |
| |
| linalg::YieldOp::create(builder, loc, result); |
| }) |
| .getResult(0); |
| |
| return result; |
| } |
| |
| // Convert an operation using per-channel quantization. |
| // |
| // - op |
| // 'quant.dcast' or 'quant.qcast' op. |
| // |
| // - input |
| // Scalar, ranked tensor, or unranked tensor. |
| // |
| // - quantizedType |
| // Per-channel quantized type. |
| // |
| Value convertPerChannel(OpBuilder &builder, Location loc, Operation *op, |
| Value input, |
| UniformQuantizedPerAxisType quantizedType) { |
| // Flatten unranked tensor into a 3D ranked tensor if necessary |
| bool isUnranked = isa<UnrankedTensorType>(input.getType()); |
| int64_t channelAxis = quantizedType.getQuantizedDimension(); |
| int64_t channelAxisSize = (int64_t)quantizedType.getScales().size(); |
| Value inputShape; |
| if (isUnranked) { |
| std::tie(input, inputShape) = flattenUnrankedTensorAroundAxis( |
| builder, loc, input, channelAxis, channelAxisSize); |
| channelAxis = 1; |
| } |
| |
| // Work on a ranked tensor |
| auto result = convertPerChannelRanked(builder, loc, op, input, quantizedType, |
| channelAxis); |
| |
| // Restore original tensor shape if unranked |
| if (isUnranked) |
| result = restoreUnrankedTensorShape(builder, loc, result, inputShape); |
| |
| return result; |
| } |
| |
| // Convert an operation using sub-channel quantization. |
| // |
| // - op |
| // 'quant.dcast' or 'quant.qcast' op. |
| // |
| // - input |
| // Scalar, ranked tensor. |
| // |
| // - quantizedType |
| // Sub-channel quantized type. |
| // |
| Value convertSubChannel(OpBuilder &builder, Location loc, Operation *op, |
| Value input, |
| UniformQuantizedSubChannelType quantizedType) { |
| auto *context = builder.getContext(); |
| |
| auto inputType = cast<RankedTensorType>(input.getType()); |
| auto inputRank = inputType.getRank(); |
| |
| auto scales = materializeSubChannelScales(builder, loc, quantizedType); |
| auto zeroPoints = |
| materializeSubChannelZeroPoints(builder, loc, quantizedType); |
| |
| auto elementType = isa<FloatType>(inputType.getElementType()) |
| ? quantizedType.getStorageType() |
| : quantizedType.getExpressedType(); |
| auto initShape = tensor::getMixedSizes(builder, loc, input); |
| Value init = tensor::EmptyOp::create(builder, loc, initShape, elementType); |
| |
| SmallVector<utils::IteratorType> iteratorTypes(inputRank, |
| utils::IteratorType::parallel); |
| const SmallVector<std::pair<int32_t, int64_t>> &blockSizeInfo = |
| quantizedType.getBlockSizeInfo(); |
| SmallVector<AffineExpr> affineExprs(inputRank, |
| builder.getAffineConstantExpr(0)); |
| for (auto [quantizedDimension, blockSize] : blockSizeInfo) { |
| affineExprs[quantizedDimension] = |
| builder.getAffineDimExpr(quantizedDimension).floorDiv(blockSize); |
| } |
| auto affineMap = AffineMap::get(inputRank, 0, affineExprs, context); |
| SmallVector<AffineMap> indexingMaps{ |
| builder.getMultiDimIdentityMap(inputRank), affineMap, affineMap, |
| builder.getMultiDimIdentityMap(inputRank)}; |
| auto result = builder |
| .create<linalg::GenericOp>( |
| loc, |
| init.getType(), // resultType |
| ValueRange{input, scales, zeroPoints}, // inputs |
| ValueRange{init}, // outputs |
| indexingMaps, iteratorTypes, |
| [&](OpBuilder &builder, Location loc, ValueRange args) { |
| assert(args.size() == 4); |
| auto input = args[0]; |
| auto scale = args[1]; |
| auto zeroPoint = args[2]; |
| |
| auto result = |
| convertRanked(builder, loc, op, input, {}, scale, |
| zeroPoint, quantizedType); |
| |
| linalg::YieldOp::create(builder, loc, result); |
| }) |
| .getResult(0); |
| |
| return result; |
| } |
| |
| // Convert a quantization operation. |
| // |
| // - op |
| // 'quant.dcast' or 'quant.qcast' op. |
| // |
| // - input |
| // Scalar, ranked tensor, or unranked tensor. The element type matches |
| // the storage type (quant.dcast) or expressed type (quant.qcast) of |
| // 'quantizedType'. |
| // |
| // - quantizedType |
| // Per-layer or per-channel quantized type. |
| // |
| Value convertQuantized(OpBuilder &builder, Location loc, Operation *op, |
| Value input, Type quantizedType) { |
| if (auto uniformQuantizedType = dyn_cast<UniformQuantizedType>(quantizedType)) |
| return convertPerLayer(builder, loc, op, input, uniformQuantizedType); |
| |
| if (auto uniformQuantizedPerAxisType = |
| dyn_cast<UniformQuantizedPerAxisType>(quantizedType)) |
| return convertPerChannel(builder, loc, op, input, |
| uniformQuantizedPerAxisType); |
| |
| if (auto uniformQuantizedSubChannelType = |
| dyn_cast<UniformQuantizedSubChannelType>(quantizedType)) |
| return convertSubChannel(builder, loc, op, input, |
| uniformQuantizedSubChannelType); |
| |
| llvm_unreachable("unexpected quantized type"); |
| } |
| |
| // Lowering pattern for 'quant.dcast' |
| struct DequantizeCastOpConversion |
| : public OpConversionPattern<quant::DequantizeCastOp> { |
| using OpConversionPattern<quant::DequantizeCastOp>::OpConversionPattern; |
| |
| LogicalResult |
| matchAndRewrite(quant::DequantizeCastOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto loc = op.getLoc(); |
| auto input = op.getInput(); |
| auto quantizedType = |
| cast<QuantizedType>(getScalarType(op.getInput().getType())); |
| |
| // Convert quantized input to storage type |
| auto storageScalarOrTensorType = |
| getScalarOrTensorType(quantizedType.getStorageType(), input.getType()); |
| input = quant::StorageCastOp::create(rewriter, loc, |
| storageScalarOrTensorType, input); |
| |
| auto result = convertQuantized(rewriter, loc, op, input, quantizedType); |
| |
| rewriter.replaceOp(op, result); |
| return success(); |
| } |
| }; |
| |
| // Lowering pattern for 'quant.qcast' |
| struct QuantizeCastOpConversion |
| : public OpConversionPattern<quant::QuantizeCastOp> { |
| using OpConversionPattern<quant::QuantizeCastOp>::OpConversionPattern; |
| |
| LogicalResult |
| matchAndRewrite(quant::QuantizeCastOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto loc = op.getLoc(); |
| auto input = op.getInput(); |
| auto quantizedType = getScalarType(op.getResult().getType()); |
| |
| // Flatten unranked tensor input |
| auto result = convertQuantized(rewriter, loc, op, input, quantizedType); |
| |
| // Cast stored value to result quantized value |
| rewriter.replaceOpWithNewOp<quant::StorageCastOp>( |
| op, op.getResult().getType(), result); |
| return success(); |
| } |
| }; |
| |
| struct LowerQuantOps : public impl::LowerQuantOpsBase<LowerQuantOps> { |
| void runOnOperation() override { |
| RewritePatternSet patterns(&getContext()); |
| populateLowerQuantOpsPatterns(patterns); |
| |
| ConversionTarget target(getContext()); |
| target.addLegalOp<quant::StorageCastOp>(); |
| target.addIllegalDialect<quant::QuantDialect>(); |
| target.addLegalDialect<arith::ArithDialect, linalg::LinalgDialect, |
| shape::ShapeDialect, tensor::TensorDialect>(); |
| |
| if (failed(applyPartialConversion(getOperation(), target, |
| std::move(patterns)))) |
| signalPassFailure(); |
| } |
| }; |
| |
| } // namespace |
| |
| void populateLowerQuantOpsPatterns(RewritePatternSet &patterns) { |
| patterns.add<DequantizeCastOpConversion, QuantizeCastOpConversion>( |
| patterns.getContext()); |
| } |
| |
| } // namespace quant |
| } // namespace mlir |
