| // Copyright 2019 The IREE Authors |
| // |
| // Licensed 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 |
| |
| // Implements logic for lowering StableHLO convolution ops to Linalg dialect. |
| |
| #include "compiler/plugins/input/StableHLO/Conversion/LegalizeToLinalgUtils.h" |
| #include "compiler/plugins/input/StableHLO/Conversion/Rewriters.h" |
| #include "mlir/Dialect/Linalg/IR/Linalg.h" |
| #include "mlir/IR/BuiltinTypes.h" |
| #include "mlir/Support/LogicalResult.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| #include "stablehlo/dialect/StablehloOps.h" |
| |
| namespace mlir::iree_compiler::stablehlo { |
| namespace { |
| /// Apply dilation and padding to the input of a convolution. |
| Value applyConvolutionPadding(Location loc, Value input, |
| DenseIntElementsAttr padding, |
| std::optional<ArrayRef<int64_t>> lhsDilation, |
| llvm::ArrayRef<int64_t> dimMappings, |
| OpBuilder &rewriter) { |
| SmallVector<int64_t> lhsDilationValues; |
| if (lhsDilation.has_value()) |
| lhsDilationValues = llvm::to_vector(lhsDilation.value()); |
| bool noPadding = !padding || isSplatValue(padding, 0); |
| bool noDilation = !lhsDilation || hlo::isSplatArray(lhsDilationValues, 1); |
| if (noPadding && noDilation) |
| return input; |
| |
| auto inputType = cast<ShapedType>(input.getType()); |
| int64_t rank = inputType.getRank(); |
| |
| // Translate window padding into low/high padding. |
| SmallVector<int64_t, 8> padLow(rank, 0); |
| SmallVector<int64_t, 8> padHigh(rank, 0); |
| if (padding) { |
| // The padding attribute contains two values per dimension, but excludes the |
| // batch and feature dimensions. |
| assert(rank * 2 == padding.size() + 4 && |
| "There should be 2 padding values per dimension, i.e low and high."); |
| for (int64_t i : llvm::seq<int64_t>(0, padding.size() / 2)) { |
| int64_t dim = dimMappings[i]; |
| padLow[dim] = padding.getValues<int64_t>()[i * 2]; |
| padHigh[dim] = padding.getValues<int64_t>()[i * 2 + 1]; |
| } |
| } |
| |
| // Translate input dilation into interior padding. |
| SmallVector<int64_t, 8> padInterior(rank, 0); |
| if (lhsDilation) { |
| assert(rank == static_cast<int64_t>(lhsDilationValues.size()) + 2); |
| for (int64_t i : llvm::seq<int64_t>(0, lhsDilationValues.size())) { |
| int64_t dim = dimMappings[i]; |
| padInterior[dim] = lhsDilationValues[i] - 1; |
| } |
| } |
| |
| Value zero; |
| if (auto complexType = dyn_cast<ComplexType>(inputType.getElementType())) { |
| auto zeroElement = rewriter.getZeroAttr(complexType.getElementType()); |
| auto zeroAttr = rewriter.getArrayAttr({zeroElement, zeroElement}); |
| zero = rewriter.create<complex::ConstantOp>(loc, complexType, zeroAttr); |
| zero = rewriter.create<tensor::FromElementsOp>( |
| loc, RankedTensorType::get({}, complexType), zero); |
| } else { |
| zero = rewriter.create<arith::ConstantOp>( |
| loc, rewriter.getZeroAttr( |
| RankedTensorType::get({}, inputType.getElementType()))); |
| } |
| |
| return rewriter.create<mlir::stablehlo::PadOp>(loc, input, zero, padLow, |
| padHigh, padInterior); |
| } |
| |
| /// If the ConvolutionOp has a window reversal, applies it to the filter. |
| Value applyConvolutionReversal(Location loc, OpBuilder &b, |
| mlir::stablehlo::ConvolutionOp op, |
| Value filter) { |
| std::optional reversals = op.getWindowReversal(); |
| if (!reversals.has_value()) { |
| return filter; |
| } |
| llvm::SmallVector<int64_t> reversedDims; |
| for (auto [idx, reversed] : llvm::enumerate(reversals.value())) { |
| if (reversed) { |
| reversedDims.push_back( |
| op.getDimensionNumbers().getKernelSpatialDimensions()[idx]); |
| } |
| } |
| |
| return b.create<mlir::stablehlo::ReverseOp>( |
| loc, filter, b.getDenseI64ArrayAttr(reversedDims)); |
| } |
| |
| /// Returns true if the given `dimensionNumbers` from a stablehlo.convolution op |
| /// follows a canonical form: |
| /// |
| /// * Input dimensions have order: (batch_count, spatial_dims, |
| /// input_channel_count). |
| /// * Filter dimensions have order: (spatial_dims, input_channel_count, |
| /// output_channel_count). |
| /// * Output dimensions have order: (batch_count, spatial_dims, |
| /// output_channel_count). |
| bool hasCanonicalDimensionNumbers( |
| mlir::stablehlo::ConvDimensionNumbersAttr dimensionNumbers) { |
| const int64_t inputSpatialRank = |
| dimensionNumbers.getInputSpatialDimensions().size(); |
| // The dimensions for input should follow the order of |
| // batch_count, spatial_dims..., input_feature_count. |
| if (dimensionNumbers.getInputBatchDimension() != 0 || |
| dimensionNumbers.getInputFeatureDimension() != (inputSpatialRank + 1)) { |
| return false; |
| } |
| |
| const int64_t kernelSpatialRank = |
| dimensionNumbers.getKernelSpatialDimensions().size(); |
| // The dimensions for filter should follow the order of |
| // spatial_dims..., input_feature_count, num_output_feature_count. |
| if (dimensionNumbers.getKernelInputFeatureDimension() != kernelSpatialRank || |
| dimensionNumbers.getKernelOutputFeatureDimension() != |
| (kernelSpatialRank + 1)) { |
| return false; |
| } |
| |
| const int64_t outputSpatialRank = |
| dimensionNumbers.getOutputSpatialDimensions().size(); |
| // The dimensions for output should follow the order of |
| // batch_count, spatial_dims.., output_feature_count. |
| if (dimensionNumbers.getOutputBatchDimension() != 0 || |
| dimensionNumbers.getOutputFeatureDimension() != (outputSpatialRank + 1)) { |
| return false; |
| } |
| |
| if (inputSpatialRank != outputSpatialRank || |
| inputSpatialRank != kernelSpatialRank) { |
| return false; |
| } |
| |
| const int64_t *inputSpatialDim = |
| dimensionNumbers.getInputSpatialDimensions().data(); |
| const int64_t *kernelSpatialDim = |
| dimensionNumbers.getKernelSpatialDimensions().data(); |
| const int64_t *outputSpatialDim = |
| dimensionNumbers.getOutputSpatialDimensions().data(); |
| // Check spatial dims are ordered correctly. |
| for (int64_t i = 0; i < inputSpatialRank; ++i) { |
| const int64_t dim = i + 1; |
| if ((*inputSpatialDim++) != dim || (*outputSpatialDim++) != dim || |
| (*kernelSpatialDim++) != i) { |
| return false; |
| } |
| } |
| |
| return true; |
| } |
| |
| /// Converts stablehlo.conv operation to linalg named op. This only covers |
| /// normal convolution cases. The op must have canonical dimension numbers. |
| /// Depthwise convolution and pointwise convolution are not handled in the |
| /// conversion. |
| struct NormalConvolutionOpConversion final |
| : OpConversionPattern<mlir::stablehlo::ConvolutionOp> { |
| using OpConversionPattern::OpConversionPattern; |
| |
| LogicalResult |
| matchAndRewrite(mlir::stablehlo::ConvolutionOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| if (!hasCanonicalDimensionNumbers(op.getDimensionNumbers())) { |
| return failure(); |
| } |
| if (op.getFeatureGroupCount() != 1u) |
| return failure(); |
| if (op.getBatchGroupCount() != 1u) |
| return failure(); |
| |
| Location loc = op.getLoc(); |
| Value input = adaptor.getLhs(); |
| Value filter = adaptor.getRhs(); |
| filter = applyConvolutionReversal(loc, rewriter, op, filter); |
| auto resultType = dyn_cast_or_null<ShapedType>( |
| getTypeConverter()->convertType(op.getResult().getType())); |
| if (!resultType) { |
| return rewriter.notifyMatchFailure(op, "type conversion failed"); |
| } |
| |
| int64_t rank = resultType.getRank(); |
| |
| // Immediately emit an EmptyOp for output tensors with zero dimension. |
| if (llvm::is_contained(resultType.getShape(), 0)) { |
| rewriter.replaceOpWithNewOp<tensor::EmptyOp>(op, resultType.getShape(), |
| resultType.getElementType()); |
| return success(); |
| } |
| |
| // The output shape is N spatial_dims F. |
| SmallVector<Value, 8> dynSizes; |
| if (resultType.isDynamicDim(0)) { |
| dynSizes.push_back(rewriter.create<tensor::DimOp>(loc, input, 0)); |
| } |
| for (int64_t i = 1, e = rank - 1; i < e; ++i) { |
| if (resultType.isDynamicDim(i)) { |
| return rewriter.notifyMatchFailure( |
| op, "expected output spatial dims to be static shapes"); |
| } |
| } |
| if (resultType.isDynamicDim(rank - 1)) { |
| dynSizes.push_back(rewriter.create<tensor::DimOp>(loc, filter, rank - 1)); |
| } |
| Value emptyTensor = rewriter.create<tensor::EmptyOp>( |
| loc, resultType.getShape(), resultType.getElementType(), dynSizes); |
| Value zeroTensor = fillTensorWithZeros(rewriter, loc, emptyTensor); |
| linalg::LinalgOp res; |
| Attribute strides; |
| if (auto s = op.getWindowStrides()) |
| strides = rewriter.getI64TensorAttr(*s); |
| Attribute dilations; |
| if (auto d = op.getRhsDilation()) |
| dilations = rewriter.getI64TensorAttr(*d); |
| |
| // Apply padding and input dilation. |
| llvm::SmallVector<int64_t> spatialDimMapping(rank - 2); |
| std::iota(spatialDimMapping.begin(), spatialDimMapping.end(), 1); |
| input = applyConvolutionPadding(loc, input, op.getPaddingAttr(), |
| op.getLhsDilation(), spatialDimMapping, |
| rewriter); |
| |
| switch (rank) { |
| case 2: { |
| res = rewriter.create<linalg::MatmulOp>( |
| loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor}, |
| linalg::getPrunedAttributeList(op)); |
| break; |
| } |
| case 3: { |
| res = rewriter.create<linalg::Conv1DNwcWcfOp>( |
| loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor}, |
| strides, dilations, linalg::getPrunedAttributeList(op)); |
| break; |
| } |
| case 4: { |
| res = rewriter.create<linalg::Conv2DNhwcHwcfOp>( |
| loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor}, |
| strides, dilations, linalg::getPrunedAttributeList(op)); |
| break; |
| } |
| case 5: { |
| res = rewriter.create<linalg::Conv3DNdhwcDhwcfOp>( |
| loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor}, |
| strides, dilations, linalg::getPrunedAttributeList(op)); |
| break; |
| } |
| default: { |
| return rewriter.notifyMatchFailure(op, "expected 1/2/3D conv op"); |
| } |
| } |
| rewriter.replaceOp(op, res.getOperation()->getResults()); |
| return success(); |
| } |
| }; |
| |
| /// Handles all possible inputs for the mlir::stablehlo::ConvolutionOp |
| struct ConvolutionOpGeneralConversion final |
| : OpConversionPattern<mlir::stablehlo::ConvolutionOp> { |
| using OpConversionPattern::OpConversionPattern; |
| |
| /// This lowering proceeds with the following steps: |
| /// 1. Handle padding and dilation of the input |
| /// 2. Handle padding and dilation of the window |
| /// 3. Handle reversal of the window |
| /// 4. If feature_group_count != 1: |
| /// - Reshape the input feature dimension, kernel output feature dimension, |
| /// and output feature dimension. |
| /// - Create the AffineExpr for the new dimension |
| /// - Conceptually, this splits the input feature and both output feature |
| /// dimensions and computes sets of convolutions with these partial views |
| /// of the values as if they were multiple convolutions combined in a |
| /// batch. |
| /// 5: If batch_group_count != 1: |
| /// - Reshape the input batch dimension, kernel output feature dimension, |
| /// and output feature dimension. |
| /// - Create the AffineExpr for the new dimension |
| /// - Conceptually, this splits the input batch and both output feature |
| /// dimensions and computes sets of convolutions with these partial views |
| /// of the values as if they were multiple convolutions combined in a |
| /// batch. |
| /// 6. For all dimensions not newly created by a reshape, create the |
| /// appropriate parallel and reduction dimensions to create a convolution. |
| /// 7. Create the linalg.generic that computes the multiply-add |
| /// 8. Reshape the output to the original shape if it was reshaped by the |
| /// feature or group count attributes. |
| LogicalResult |
| matchAndRewrite(mlir::stablehlo::ConvolutionOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| Location loc = op.getLoc(); |
| MLIRContext *ctx = op.getContext(); |
| |
| auto resultType = dyn_cast_or_null<ShapedType>( |
| getTypeConverter()->convertType(op.getResult().getType())); |
| if (!resultType) { |
| return rewriter.notifyMatchFailure(op, "type conversion failed"); |
| } |
| |
| auto reshapedResultShape = resultType.getShape().vec(); |
| if (!resultType.hasStaticShape()) |
| return failure(); |
| |
| // Immediately emit an EmptyOp for output tensors with zero dimension. |
| if (llvm::is_contained(reshapedResultShape, 0)) { |
| rewriter.replaceOpWithNewOp<tensor::EmptyOp>(op, reshapedResultShape, |
| resultType.getElementType()); |
| return success(); |
| } |
| |
| mlir::stablehlo::ConvDimensionNumbersAttr dimensionNumbers = |
| op.getDimensionNumbers(); |
| int64_t inputBatchDimension = dimensionNumbers.getInputBatchDimension(); |
| int64_t inputFeatureDimension = dimensionNumbers.getInputFeatureDimension(); |
| ArrayRef<int64_t> inputSpatialDimensions = |
| dimensionNumbers.getInputSpatialDimensions(); |
| |
| int64_t kernelInputFeatureDimension = |
| dimensionNumbers.getKernelInputFeatureDimension(); |
| int64_t kernelOutputFeatureDimension = |
| dimensionNumbers.getKernelOutputFeatureDimension(); |
| ArrayRef<int64_t> kernelSpatialDimensions = |
| dimensionNumbers.getKernelSpatialDimensions(); |
| |
| int64_t outputBatchDimension = dimensionNumbers.getOutputBatchDimension(); |
| int64_t outputFeatureDimension = |
| dimensionNumbers.getOutputFeatureDimension(); |
| ArrayRef<int64_t> outputSpatialDimensions = |
| dimensionNumbers.getOutputSpatialDimensions(); |
| |
| size_t featureGroupCount = op.getFeatureGroupCount(); |
| size_t batchGroupCount = op.getBatchGroupCount(); |
| |
| if (op.getFeatureGroupCount() != 1 && op.getBatchGroupCount() != 1) { |
| return rewriter.notifyMatchFailure( |
| op, "only one of feature and batch group counts can be non-one"); |
| } |
| |
| // Decompose the convolution into an initial padding |
| Value modifiedLhs = applyConvolutionPadding( |
| op.getLoc(), adaptor.getLhs(), adaptor.getPaddingAttr(), |
| adaptor.getLhsDilation(), |
| op.getDimensionNumbers().getInputSpatialDimensions(), rewriter); |
| Value modifiedRhs = applyConvolutionPadding( |
| op.getLoc(), adaptor.getRhs(), nullptr, adaptor.getRhsDilation(), |
| op.getDimensionNumbers().getKernelSpatialDimensions(), rewriter); |
| modifiedRhs = applyConvolutionReversal(loc, rewriter, op, modifiedRhs); |
| |
| // Non-one values for feature or batch group counts will result in reshaped |
| // inputs and outputs. These mappings are used to keep track of the the new |
| // index after reshaping has possibly inserted new dimensions. |
| auto paddedLhsType = cast<ShapedType>(modifiedLhs.getType()); |
| auto paddedRhsType = cast<ShapedType>(modifiedRhs.getType()); |
| SmallVector<int64_t> lhsIndexMapping(paddedLhsType.getRank()); |
| std::iota(lhsIndexMapping.begin(), lhsIndexMapping.end(), 0); |
| SmallVector<int64_t> rhsIndexMapping(paddedRhsType.getRank()); |
| std::iota(rhsIndexMapping.begin(), rhsIndexMapping.end(), 0); |
| SmallVector<int64_t> resultIndexMapping(resultType.getRank()); |
| std::iota(resultIndexMapping.begin(), resultIndexMapping.end(), 0); |
| auto updateDimMappingFromOffset = |
| [](llvm::SmallVectorImpl<int64_t> &mapping, int64_t offset) { |
| for (auto &mappingElt : llvm::drop_begin(mapping, offset)) { |
| mappingElt += 1; |
| } |
| }; |
| |
| // The rest of this code prepares the inputs and a single linalg::GenericOp |
| // to execute the convolution. The final linalg::GenericOp will be iterated |
| // through based on the following eventual maps. |
| SmallVector<AffineExpr, 2> srcExprs(paddedLhsType.getRank()); |
| SmallVector<AffineExpr, 2> windowExprs(paddedRhsType.getRank()); |
| SmallVector<AffineExpr, 2> dstExprs(reshapedResultShape.size()); |
| int64_t nextDim = 0; |
| int64_t rank = resultType.getRank(); |
| |
| auto reshapeShapeVector = [](llvm::ArrayRef<int64_t> oldShape, |
| llvm::SmallVectorImpl<int64_t> &newShape, |
| int64_t reshapedDim, int64_t factor) { |
| newShape.reserve(oldShape.size() + 1); |
| for (int64_t i : llvm::seq<int64_t>(0, oldShape.size())) { |
| if (i == reshapedDim) { |
| newShape.push_back(factor); |
| newShape.push_back(oldShape[reshapedDim] / factor); |
| } else { |
| newShape.push_back(oldShape[i]); |
| } |
| } |
| }; |
| |
| // If batch or feature count groupings exist, represent this through |
| // reshaping the input to have an additional dimension that these groupings |
| // exist along, and reduce in that dimension |
| SmallVector<utils::IteratorType, 3> iterationLoops; |
| if (featureGroupCount != 1) { |
| AffineExpr parallelDim = mlir::getAffineDimExpr(nextDim++, ctx); |
| iterationLoops.push_back(utils::IteratorType::parallel); |
| // Reshape LHS |
| { |
| srcExprs.insert(srcExprs.begin() + inputFeatureDimension, parallelDim); |
| auto prevDimsRef = paddedLhsType.getShape(); |
| llvm::SmallVector<int64_t> newShape; |
| reshapeShapeVector(prevDimsRef, newShape, inputFeatureDimension, |
| featureGroupCount); |
| updateDimMappingFromOffset(lhsIndexMapping, inputFeatureDimension); |
| modifiedLhs = rewriter.create<mlir::stablehlo::ReshapeOp>( |
| loc, |
| RankedTensorType::get(newShape, paddedLhsType.getElementType()), |
| modifiedLhs); |
| } |
| |
| // Reshape RHS |
| { |
| windowExprs.insert(windowExprs.begin() + kernelOutputFeatureDimension, |
| parallelDim); |
| auto prevDimsRef = paddedRhsType.getShape(); |
| llvm::SmallVector<int64_t> newShape; |
| reshapeShapeVector(prevDimsRef, newShape, kernelOutputFeatureDimension, |
| featureGroupCount); |
| updateDimMappingFromOffset(rhsIndexMapping, |
| kernelOutputFeatureDimension); |
| modifiedRhs = rewriter.create<mlir::stablehlo::ReshapeOp>( |
| loc, |
| RankedTensorType::get(newShape, paddedRhsType.getElementType()), |
| modifiedRhs); |
| } |
| // Prepare reshaped output shape |
| { |
| dstExprs.insert(dstExprs.begin() + outputFeatureDimension, parallelDim); |
| updateDimMappingFromOffset(resultIndexMapping, outputFeatureDimension); |
| reshapedResultShape.insert(reshapedResultShape.begin() + |
| outputFeatureDimension, |
| featureGroupCount); |
| reshapedResultShape[outputFeatureDimension + 1] /= featureGroupCount; |
| } |
| } |
| |
| if (batchGroupCount != 1) { |
| iterationLoops.push_back(utils::IteratorType::parallel); |
| AffineExpr parallelDim = mlir::getAffineDimExpr(nextDim++, ctx); |
| // Reshape LHS |
| { |
| srcExprs.insert(srcExprs.begin() + inputBatchDimension, parallelDim); |
| ArrayRef<int64_t> prevDimsRef = paddedLhsType.getShape(); |
| llvm::SmallVector<int64_t> newShape; |
| reshapeShapeVector(prevDimsRef, newShape, inputBatchDimension, |
| batchGroupCount); |
| updateDimMappingFromOffset(lhsIndexMapping, inputBatchDimension); |
| modifiedLhs = rewriter.create<mlir::stablehlo::ReshapeOp>( |
| op.getLoc(), |
| RankedTensorType::get(newShape, paddedLhsType.getElementType()), |
| modifiedLhs); |
| } |
| |
| // Reshape RHS |
| { |
| windowExprs.insert(windowExprs.begin() + kernelOutputFeatureDimension, |
| parallelDim); |
| ArrayRef<int64_t> prevDimsRef = paddedRhsType.getShape(); |
| llvm::SmallVector<int64_t> newShape; |
| reshapeShapeVector(prevDimsRef, newShape, kernelOutputFeatureDimension, |
| batchGroupCount); |
| updateDimMappingFromOffset(rhsIndexMapping, |
| kernelOutputFeatureDimension); |
| modifiedRhs = rewriter.create<mlir::stablehlo::ReshapeOp>( |
| op.getLoc(), |
| RankedTensorType::get(newShape, paddedRhsType.getElementType()), |
| modifiedRhs); |
| } |
| // Prepare reshaped output shape |
| { |
| int64_t outputFeatureDim = resultIndexMapping[outputFeatureDimension]; |
| dstExprs.insert(dstExprs.begin() + outputFeatureDim, parallelDim); |
| updateDimMappingFromOffset(resultIndexMapping, outputFeatureDimension); |
| reshapedResultShape.insert( |
| reshapedResultShape.begin() + outputFeatureDim, batchGroupCount); |
| reshapedResultShape[outputFeatureDim + 1] /= batchGroupCount; |
| } |
| } |
| |
| // Handle input feature dimension |
| { |
| iterationLoops.push_back(utils::IteratorType::reduction); |
| AffineExpr inputFeatureDim = mlir::getAffineDimExpr(nextDim++, ctx); |
| srcExprs[lhsIndexMapping[inputFeatureDimension]] = inputFeatureDim; |
| windowExprs[rhsIndexMapping[kernelInputFeatureDimension]] = |
| inputFeatureDim; |
| } |
| |
| // Handle output feature dimension |
| { |
| iterationLoops.push_back(utils::IteratorType::parallel); |
| AffineExpr outputFeatureDim = mlir::getAffineDimExpr(nextDim++, ctx); |
| dstExprs[resultIndexMapping[outputFeatureDimension]] = outputFeatureDim; |
| windowExprs[rhsIndexMapping[kernelOutputFeatureDimension]] = |
| outputFeatureDim; |
| } |
| |
| // Handle spatial Dimensions |
| int64_t numSpatialDims = rank - 2; |
| for (int64_t i = 0; i < numSpatialDims; ++i) { |
| iterationLoops.push_back(utils::IteratorType::parallel); |
| iterationLoops.push_back(utils::IteratorType::reduction); |
| AffineExpr dim0 = mlir::getAffineDimExpr(nextDim++, ctx); |
| AffineExpr dim1 = mlir::getAffineDimExpr(nextDim++, ctx); |
| |
| AffineExpr stride = dim0; |
| if (op.getWindowStrides().has_value()) |
| stride = stride * op.getWindowStrides().value()[i]; |
| AffineExpr srcExpr = stride + dim1; |
| |
| srcExprs[lhsIndexMapping[inputSpatialDimensions[i]]] = srcExpr; |
| dstExprs[resultIndexMapping[outputSpatialDimensions[i]]] = dim0; |
| windowExprs[rhsIndexMapping[kernelSpatialDimensions[i]]] = dim1; |
| } |
| |
| // Handle batch dimension |
| { |
| iterationLoops.push_back(utils::IteratorType::parallel); |
| AffineExpr batchDim = mlir::getAffineDimExpr(nextDim++, ctx); |
| |
| srcExprs[lhsIndexMapping[inputBatchDimension]] = batchDim; |
| dstExprs[resultIndexMapping[outputBatchDimension]] = batchDim; |
| } |
| |
| // Finally, create the computation |
| auto inferredMaps = |
| AffineMap::inferFromExprList({srcExprs, windowExprs, dstExprs}, ctx); |
| |
| Value emptyTensor = rewriter.create<tensor::EmptyOp>( |
| loc, reshapedResultShape, resultType.getElementType()); |
| Value zeroTensor = fillTensorWithZeros(rewriter, loc, emptyTensor); |
| |
| Value convolved = |
| rewriter |
| .create<linalg::GenericOp>( |
| loc, |
| /*resultTensors=*/ |
| llvm::ArrayRef<Type>(zeroTensor.getType()), |
| /*inputs=*/ |
| llvm::ArrayRef<Value>({modifiedLhs, modifiedRhs}), |
| /*outputs=*/llvm::ArrayRef<Value>(zeroTensor), inferredMaps, |
| iterationLoops, |
| /*bodyBuild=*/ |
| [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange) { |
| ImplicitLocOpBuilder builder(nestedLoc, nestedBuilder); |
| linalg::Conv2DOp::regionBuilder( |
| builder, *builder.getInsertionBlock(), {}); |
| }, |
| linalg::getPrunedAttributeList(op)) |
| .getResult(0); |
| rewriter.replaceOpWithNewOp<mlir::stablehlo::ReshapeOp>(op, resultType, |
| convolved); |
| |
| return success(); |
| } |
| }; |
| |
| /// Converts stablehlo.convolution operation to |
| /// linalg.depthwise_conv_2d_input_nhwc_filter_hwcf op or |
| /// depthwise_conv_2d_input_nhwc_filter_hwc op. |
| struct DepthwiseConvolutionOpConversion final |
| : OpConversionPattern<mlir::stablehlo::ConvolutionOp> { |
| using OpConversionPattern::OpConversionPattern; |
| |
| LogicalResult |
| matchAndRewrite(mlir::stablehlo::ConvolutionOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| if (op.getBatchGroupCount() != 1) |
| return failure(); |
| // Fall into the normal convolution cases. |
| if (op.getFeatureGroupCount() == 1) |
| return failure(); |
| |
| const mlir::stablehlo::ConvDimensionNumbersAttr &dimensionNumbers = |
| op.getDimensionNumbers(); |
| const int64_t spatialRank = |
| dimensionNumbers.getInputSpatialDimensions().size(); |
| if (spatialRank == 0 || spatialRank > 3) { |
| return rewriter.notifyMatchFailure(op, "only support up to 3D for now"); |
| } |
| |
| // Make sure that this is depthwise convolution. |
| int64_t inputFeatureDim = dimensionNumbers.getInputFeatureDimension(); |
| int64_t inputFeatureCount = |
| cast<ShapedType>(op.getLhs().getType()).getDimSize(inputFeatureDim); |
| if (static_cast<int64_t>(op.getFeatureGroupCount()) != inputFeatureCount) { |
| return rewriter.notifyMatchFailure(op, "not depth-wise convolution"); |
| } |
| |
| // Make sure that this convolution has a canonical form. |
| if (!hasCanonicalDimensionNumbers(dimensionNumbers)) { |
| return rewriter.notifyMatchFailure(op, "does not have canonical form"); |
| } |
| |
| Attribute windowStrides; |
| if (op.getWindowStrides()) { |
| windowStrides = rewriter.getI64TensorAttr(op.getWindowStrides().value()); |
| } else { |
| windowStrides = SplatElementsAttr::get( |
| VectorType::get({spatialRank}, rewriter.getI64Type()), |
| rewriter.getI64IntegerAttr(1)); |
| } |
| |
| Attribute rhsDilation; |
| if (op.getRhsDilation()) { |
| rhsDilation = rewriter.getI64TensorAttr(op.getRhsDilation().value()); |
| } else { |
| rhsDilation = SplatElementsAttr::get( |
| VectorType::get({spatialRank}, rewriter.getI64Type()), |
| rewriter.getI64IntegerAttr(1)); |
| } |
| |
| Location loc = op.getLoc(); |
| Value input = adaptor.getLhs(); |
| Value filter = adaptor.getRhs(); |
| auto resultType = dyn_cast_or_null<RankedTensorType>( |
| getTypeConverter()->convertType(op.getResult().getType())); |
| if (!resultType) { |
| return rewriter.notifyMatchFailure(op, "type conversion failed"); |
| } |
| if (!resultType.hasStaticShape()) { |
| return rewriter.notifyMatchFailure(op, |
| "expected output has static shapes"); |
| } |
| |
| // Immediately emit an EmptyOp for output tensors with zero dimension. |
| if (llvm::is_contained(resultType.getShape(), 0)) { |
| rewriter.replaceOpWithNewOp<tensor::EmptyOp>(op, resultType.getShape(), |
| resultType.getElementType()); |
| return success(); |
| } |
| |
| // Apply padding and input dilation. |
| llvm::SmallVector<int64_t> spatialDimMapping(spatialRank); |
| std::iota(spatialDimMapping.begin(), spatialDimMapping.end(), 1); |
| input = applyConvolutionPadding(loc, input, op.getPaddingAttr(), |
| op.getLhsDilation(), spatialDimMapping, |
| rewriter); |
| |
| auto filterDims = |
| llvm::to_vector(cast<ShapedType>(op.getRhs().getType()).getShape()); |
| |
| auto getReassociationIndicesToCollapseLastTwoDims = [](Value v) { |
| SmallVector<ReassociationIndices> reassociations; |
| int64_t rank = cast<ShapedType>(v.getType()).getRank(); |
| for (int64_t i = 0; i < rank - 1; ++i) |
| reassociations.emplace_back(1, i); |
| reassociations.back().push_back(rank - 1); |
| return reassociations; |
| }; |
| |
| int64_t kernelInputFeatureDimension = |
| dimensionNumbers.getKernelInputFeatureDimension(); |
| int64_t kernelOutputFeatureDimension = |
| dimensionNumbers.getKernelOutputFeatureDimension(); |
| if (filterDims[kernelInputFeatureDimension] * |
| filterDims[kernelOutputFeatureDimension] != |
| static_cast<int64_t>(op.getFeatureGroupCount())) { |
| // For cases where channel multiplier != 1 |
| |
| // Reshaping filter shape |
| // [filter_height, filter_width, 1, kernel-output-feature]. |
| // to |
| // [filter_height, filter_width, feature_group_count, |
| // kernel-output-feature/feature_group_count ] |
| SmallVector<int64_t> reshapedFilterDims; |
| reshapedFilterDims.assign(filterDims.begin(), filterDims.end()); |
| Value reshapedFilter = filter; |
| if (filterDims[kernelInputFeatureDimension] == 1) { |
| reshapedFilterDims[kernelInputFeatureDimension] = |
| op.getFeatureGroupCount(); |
| reshapedFilterDims[kernelOutputFeatureDimension] /= |
| op.getFeatureGroupCount(); |
| auto reshapedFilterType = RankedTensorType::get( |
| reshapedFilterDims, |
| cast<ShapedType>(op.getRhs().getType()).getElementType()); |
| |
| reshapedFilter = rewriter.create<mlir::stablehlo::ReshapeOp>( |
| loc, reshapedFilterType, filter); |
| } |
| |
| ArrayRef<int64_t> outputDims = resultType.getShape(); |
| int64_t channelMultiplier = reshapedFilterDims.back(); |
| SmallVector<int64_t> reshapedOutputDims; |
| reshapedOutputDims.assign(outputDims.begin(), outputDims.end()); |
| reshapedOutputDims.push_back(channelMultiplier); |
| reshapedOutputDims[reshapedOutputDims.size() - 2] /= channelMultiplier; |
| |
| Value emptyTensor = rewriter.create<tensor::EmptyOp>( |
| loc, reshapedOutputDims, resultType.getElementType()); |
| Value zeroTensor = fillTensorWithZeros(rewriter, loc, emptyTensor); |
| |
| auto reshapedOutputType = RankedTensorType::get( |
| reshapedOutputDims, resultType.getElementType()); |
| Value conv; |
| switch (spatialRank) { |
| case 1: { |
| conv = |
| rewriter |
| .create<linalg::DepthwiseConv1DNwcWcmOp>( |
| loc, reshapedOutputType, ValueRange{input, reshapedFilter}, |
| ValueRange{zeroTensor}, windowStrides, rhsDilation, |
| linalg::getPrunedAttributeList(op)) |
| .getResult(0); |
| break; |
| } |
| case 2: { |
| conv = |
| rewriter |
| .create<linalg::DepthwiseConv2DNhwcHwcmOp>( |
| loc, reshapedOutputType, ValueRange{input, reshapedFilter}, |
| ValueRange{zeroTensor}, windowStrides, rhsDilation, |
| linalg::getPrunedAttributeList(op)) |
| .getResult(0); |
| break; |
| } |
| case 3: { |
| conv = |
| rewriter |
| .create<linalg::DepthwiseConv3DNdhwcDhwcmOp>( |
| loc, reshapedOutputType, ValueRange{input, reshapedFilter}, |
| ValueRange{zeroTensor}, windowStrides, rhsDilation, |
| linalg::getPrunedAttributeList(op)) |
| .getResult(0); |
| break; |
| } |
| default: |
| llvm_unreachable("Unhandled case"); |
| } |
| |
| // Create a Linalg reshape op that converts the output from 5 dimensions |
| // into 4 dimensions (by collapsing the last two dimensions). This is |
| // needed because linalg.depthwise_conv_2d_input_nhwc_filter_hwcf returns |
| // 5 dimensions for the output. |
| rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>( |
| op, resultType, conv, |
| getReassociationIndicesToCollapseLastTwoDims(conv)); |
| } else { |
| // For cases where channel multiplier == 1 |
| Value emptyTensor = rewriter.create<tensor::EmptyOp>( |
| loc, resultType.getShape(), resultType.getElementType()); |
| Value zeroTensor = fillTensorWithZeros(rewriter, loc, emptyTensor); |
| |
| // Create a Linalg reshape op that converts the filter from 4 dimensions |
| // into 3 dimensions (by droping the unit dimension). This is needed |
| // because linalg.depthwise_conv_2d_input_nhwc_filter_hwc expects 3 |
| // dimensions for the filter. |
| |
| filterDims[filterDims.size() - 2] = |
| static_cast<int64_t>(op.getFeatureGroupCount()); |
| filterDims.pop_back(); |
| |
| RankedTensorType filterShape = |
| RankedTensorType::get(filterDims, op.getType().getElementType()); |
| |
| Value reshapedFilter = rewriter.create<tensor::CollapseShapeOp>( |
| loc, filterShape, filter, |
| getReassociationIndicesToCollapseLastTwoDims(filter)); |
| |
| switch (spatialRank) { |
| case 1: |
| rewriter.replaceOpWithNewOp<linalg::DepthwiseConv1DNwcWcOp>( |
| op, resultType, ValueRange{input, reshapedFilter}, |
| ValueRange{zeroTensor}, windowStrides, rhsDilation, |
| linalg::getPrunedAttributeList(op)); |
| break; |
| case 2: |
| rewriter.replaceOpWithNewOp<linalg::DepthwiseConv2DNhwcHwcOp>( |
| op, resultType, ValueRange{input, reshapedFilter}, |
| ValueRange{zeroTensor}, windowStrides, rhsDilation, |
| linalg::getPrunedAttributeList(op)); |
| break; |
| case 3: |
| rewriter.replaceOpWithNewOp<linalg::DepthwiseConv3DNdhwcDhwcOp>( |
| op, resultType, ValueRange{input, reshapedFilter}, |
| ValueRange{zeroTensor}, windowStrides, rhsDilation, |
| linalg::getPrunedAttributeList(op)); |
| break; |
| } |
| } |
| |
| return success(); |
| } |
| }; |
| |
| } // namespace |
| |
| namespace detail { |
| void populateStableHloConvolutionToLinalgConversionPatterns( |
| MLIRContext *context, TypeConverter &typeConverter, |
| RewritePatternSet *patterns) { |
| // Ensure specialized patterns are higher priority than their generic |
| // versions. |
| patterns |
| ->add<NormalConvolutionOpConversion, DepthwiseConvolutionOpConversion>( |
| typeConverter, context, PatternBenefit(2)); |
| |
| patterns->add<ConvolutionOpGeneralConversion>(typeConverter, context); |
| } |
| } // namespace detail |
| } // namespace mlir::iree_compiler::stablehlo |
