| // 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 |