[mlir][vector] Add support for linearizing Extract, ExtractStridedSlice, Shuffle VectorOps in VectorLinearize #88204
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -13,9 +13,14 @@ | |||||
| #include "mlir/Dialect/Arith/IR/Arith.h" | ||||||
| #include "mlir/Dialect/Vector/IR/VectorOps.h" | ||||||
| #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h" | ||||||
| #include "mlir/IR/Attributes.h" | ||||||
| #include "mlir/IR/Operation.h" | ||||||
| #include "mlir/IR/PatternMatch.h" | ||||||
| #include "mlir/IR/TypeUtilities.h" | ||||||
| #include "mlir/Support/LogicalResult.h" | ||||||
| #include "mlir/Transforms/DialectConversion.h" | ||||||
| #include <cstdint> | ||||||
| #include <numeric> | ||||||
|
|
||||||
| using namespace mlir; | ||||||
|
|
||||||
|
|
@@ -103,6 +108,252 @@ struct LinearizeVectorizable final | |||||
| return success(); | ||||||
| } | ||||||
|
|
||||||
| private: | ||||||
| unsigned targetVectorBitWidth; | ||||||
| }; | ||||||
|
|
||||||
| /// This pattern converts the ExtractStridedSliceOp into a ShuffleOp that works | ||||||
| /// on a linearized vector. | ||||||
| /// Following, | ||||||
| /// vector.extract_strided_slice %source | ||||||
| /// { offsets = [..], strides = [..], sizes = [..] } | ||||||
| /// is converted to : | ||||||
| /// %source_1d = vector.shape_cast %source | ||||||
| /// %out_1d = vector.shuffle %source_1d, %source_1d [ shuffle_indices_1d ] | ||||||
| /// %out_nd = vector.shape_cast %out_1d | ||||||
| /// `shuffle_indices_1d` is computed using the offsets and sizes of the | ||||||
| /// extraction. | ||||||
| struct LinearizeVectorExtractStridedSlice final | ||||||
banach-space marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||
| : public mlir::OpConversionPattern<mlir::vector::ExtractStridedSliceOp> { | ||||||
| using OpConversionPattern::OpConversionPattern; | ||||||
| LinearizeVectorExtractStridedSlice( | ||||||
| const TypeConverter &typeConverter, MLIRContext *context, | ||||||
| unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(), | ||||||
| PatternBenefit benefit = 1) | ||||||
| : OpConversionPattern(typeConverter, context, benefit), | ||||||
| targetVectorBitWidth(targetVectBitWidth) {} | ||||||
|
There was a problem hiding this comment. Thanks for adding support for the target vector bitwidth! |
||||||
|
|
||||||
| LogicalResult | ||||||
| matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor, | ||||||
| ConversionPatternRewriter &rewriter) const override { | ||||||
| auto dstType = getTypeConverter()->convertType(extractOp.getType()); | ||||||
|
There was a problem hiding this comment. yes. this is not needed. I removed the |
||||||
| assert(!(extractOp.getVector().getType().isScalable() || | ||||||
| dstType.cast<VectorType>().isScalable()) && | ||||||
| "scalable vectors are not supported."); | ||||||
| if (!isLessThanTargetBitWidth(extractOp, targetVectorBitWidth)) | ||||||
| return rewriter.notifyMatchFailure( | ||||||
| extractOp, "Can't flatten since targetBitWidth <= OpSize"); | ||||||
banach-space marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||
|
|
||||||
| auto offsets = extractOp.getOffsets(); | ||||||
| auto sizes = extractOp.getSizes(); | ||||||
| auto strides = extractOp.getStrides(); | ||||||
| if (!isConstantIntValue(strides[0], 1)) | ||||||
| return rewriter.notifyMatchFailure( | ||||||
| extractOp, "Strided slice with stride != 1 is not supported."); | ||||||
| Value srcVector = adaptor.getVector(); | ||||||
| // If kD offsets are specified for nd source vector (n > k), the granularity | ||||||
|
There was a problem hiding this comment.
Suggested change
|
||||||
| // of the extraction is greater than 1. In this case last (n-k) dimensions | ||||||
| // form the extraction granularity. | ||||||
| // Example : | ||||||
| // vector.extract_strided_slice %src { | ||||||
| // offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : | ||||||
| // vector<4x8x8xf32> to vector<2x2x8xf32> | ||||||
| // Here, extraction granularity is 8. | ||||||
| int64_t extractGranularitySize = 1; | ||||||
| auto n = extractOp.getSourceVectorType().getRank(); | ||||||
|
There was a problem hiding this comment. nit: please, spell out the type when it's not literally redundant in the statement. |
||||||
| int64_t k = (int64_t)offsets.size(); | ||||||
| if (n > k) { | ||||||
|
There was a problem hiding this comment. Is this There was a problem hiding this comment. addressed above. n > k is true for the following example. In this case, only first 2 dims are extracted using the offsets, sizes etc. last dim is extracted in full. There was a problem hiding this comment. Right, so there are 2 cases:
right? I think that you can get rid of int64_t outRank = (int64_t)offsets.size();
int64_t k = outRank;
while (k < n) {
extractGranularitySize *=
extractOp.getSourceVectorType().getShape()[outRank + (n-k)];
++k;
}This way there's only one case. Also, I'd avoid variable names like |
||||||
| for (unsigned i = 0; i < n - k; i++) { | ||||||
|
There was a problem hiding this comment. nit: use pre-increment per coding guidelines |
||||||
| extractGranularitySize *= | ||||||
| extractOp.getSourceVectorType().getShape()[i + k]; | ||||||
| } | ||||||
| } | ||||||
| // Get total number of extracted slices. | ||||||
| int64_t nExtractedSlices = 1; | ||||||
| llvm::for_each(sizes, [&](Attribute size) { | ||||||
|
There was a problem hiding this comment.
There was a problem hiding this comment. @banach-space recommended using |
||||||
| nExtractedSlices *= size.cast<IntegerAttr>().getInt(); | ||||||
| }); | ||||||
| // Compute the strides of the source vector considering first k dimensions. | ||||||
| llvm::SmallVector<int64_t, 4> sourceStrides(k, extractGranularitySize); | ||||||
| for (int i = k - 2; i >= 0; --i) { | ||||||
| sourceStrides[i] = sourceStrides[i + 1] * | ||||||
| extractOp.getSourceVectorType().getShape()[i + 1]; | ||||||
| } | ||||||
| // Final shuffle indices has nExtractedSlices * extractGranularitySize | ||||||
| // elements. | ||||||
| llvm::SmallVector<int64_t, 4> indices(nExtractedSlices * | ||||||
| extractGranularitySize); | ||||||
| // Compute the strides of the extracted kD vector. | ||||||
| llvm::SmallVector<int64_t, 4> extractedStrides(k, 1); | ||||||
| // Compute extractedStrides. | ||||||
| for (int i = k - 2; i >= 0; --i) { | ||||||
| extractedStrides[i] = | ||||||
| extractedStrides[i + 1] * sizes[i + 1].cast<IntegerAttr>().getInt(); | ||||||
| } | ||||||
| // Iterate over all extracted slices from 0 to nExtractedSlices - 1 | ||||||
| // and compute the multi-dimensional index and the corresponding linearized | ||||||
| // index within the source vector. | ||||||
| for (int64_t i = 0; i < nExtractedSlices; ++i) { | ||||||
| int64_t index = i; | ||||||
| // Compute the corresponding multi-dimensional index. | ||||||
| llvm::SmallVector<int64_t, 4> multiDimIndex(k, 0); | ||||||
| for (int64_t j = 0; j < k; ++j) { | ||||||
| multiDimIndex[j] = (index / extractedStrides[j]); | ||||||
| index -= multiDimIndex[j] * extractedStrides[j]; | ||||||
| } | ||||||
| // Compute the corresponding linearized index in the source vector | ||||||
| // i.e. shift the multiDimIndex by the offsets. | ||||||
| int64_t linearizedIndex = 0; | ||||||
| for (int64_t j = 0; j < k; ++j) { | ||||||
| linearizedIndex += | ||||||
| (offsets[j].cast<IntegerAttr>().getInt() + multiDimIndex[j]) * | ||||||
| sourceStrides[j]; | ||||||
| } | ||||||
| // Fill the indices array form linearizedIndex to linearizedIndex + | ||||||
| // extractGranularitySize. | ||||||
| for (int64_t j = 0; j < extractGranularitySize; ++j) { | ||||||
| indices[i * extractGranularitySize + j] = linearizedIndex + j; | ||||||
| } | ||||||
| } | ||||||
| // Perform a shuffle to extract the kD vector. | ||||||
| rewriter.replaceOpWithNewOp<vector::ShuffleOp>( | ||||||
| extractOp, dstType, srcVector, srcVector, | ||||||
| rewriter.getI64ArrayAttr(indices)); | ||||||
| return success(); | ||||||
| } | ||||||
|
|
||||||
| private: | ||||||
| unsigned targetVectorBitWidth; | ||||||
| }; | ||||||
|
|
||||||
| /// This pattern converts the ShuffleOp that works on nD (n > 1) | ||||||
| /// vectors to a ShuffleOp that works on linearized vectors. | ||||||
| /// Following, | ||||||
| /// vector.shuffle %v1, %v2 [ shuffle_indices ] | ||||||
| /// is converted to : | ||||||
| /// %v1_1d = vector.shape_cast %v1 | ||||||
| /// %v2_1d = vector.shape_cast %v2 | ||||||
| /// %out_1d = vector.shuffle %v1_1d, %v2_1d [ shuffle_indices_1d ] | ||||||
| /// %out_nd = vector.shape_cast %out_1d | ||||||
| // `shuffle_indices_1d` is computed using the sizes and `shuffle_indices` | ||||||
| /// of the original shuffle operation. | ||||||
| struct LinearizeVectorShuffle final | ||||||
| : public OpConversionPattern<vector::ShuffleOp> { | ||||||
| using OpConversionPattern::OpConversionPattern; | ||||||
| LinearizeVectorShuffle( | ||||||
| const TypeConverter &typeConverter, MLIRContext *context, | ||||||
| unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(), | ||||||
| PatternBenefit benefit = 1) | ||||||
| : OpConversionPattern(typeConverter, context, benefit), | ||||||
| targetVectorBitWidth(targetVectBitWidth) {} | ||||||
|
|
||||||
| LogicalResult | ||||||
| matchAndRewrite(vector::ShuffleOp shuffleOp, OpAdaptor adaptor, | ||||||
| ConversionPatternRewriter &rewriter) const override { | ||||||
| auto dstType = getTypeConverter()->convertType(shuffleOp.getType()); | ||||||
| assert(!(shuffleOp.getV1VectorType().isScalable() || | ||||||
| shuffleOp.getV2VectorType().isScalable() || | ||||||
| dstType.cast<VectorType>().isScalable()) && | ||||||
| "scalable vectors are not supported."); | ||||||
| if (!isLessThanTargetBitWidth(shuffleOp, targetVectorBitWidth)) | ||||||
| return rewriter.notifyMatchFailure( | ||||||
| shuffleOp, "Can't flatten since targetBitWidth <= OpSize"); | ||||||
|
|
||||||
| auto vec1 = adaptor.getV1(); | ||||||
| auto vec2 = adaptor.getV2(); | ||||||
|
|
||||||
| int shuffleSliceLen = 1; | ||||||
| int rank = shuffleOp.getV1().getType().getRank(); | ||||||
|
|
||||||
| // If rank > 1, we need to do the shuffle in the granularity of slices | ||||||
| // instead of scalars. Size of the slice is equal to the rank-1 innermost | ||||||
| // dims. Mask of the shuffle op specifies which slice to take from the | ||||||
| // outermost dim. | ||||||
| if (rank > 1) { | ||||||
| auto shape = shuffleOp.getV1().getType().getShape(); | ||||||
| for (unsigned i = 1; i < shape.size(); i++) { | ||||||
|
|
||||||
| shuffleSliceLen *= shape[i]; | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| // For each value in the mask, we generate the indices of the source vectors | ||||||
| // that needs to be shuffled to the destination vector. If shuffleSliceLen > | ||||||
| // 1 we need to shuffle the slices (consecutive shuffleSliceLen number of | ||||||
| // elements) instead of scalars. | ||||||
| auto mask = shuffleOp.getMask(); | ||||||
| auto totalSizeOfShuffledElmnts = mask.size() * shuffleSliceLen; | ||||||
| llvm::SmallVector<int64_t, 2> indices(totalSizeOfShuffledElmnts); | ||||||
| for (auto [i, value] : | ||||||
| llvm::enumerate(mask.getAsValueRange<IntegerAttr>())) { | ||||||
|
|
||||||
| int64_t v = value.getZExtValue(); | ||||||
| std::iota(indices.begin() + shuffleSliceLen * i, | ||||||
| indices.begin() + shuffleSliceLen * (i + 1), | ||||||
| shuffleSliceLen * v); | ||||||
| } | ||||||
|
|
||||||
| rewriter.replaceOpWithNewOp<vector::ShuffleOp>( | ||||||
| shuffleOp, dstType, vec1, vec2, rewriter.getI64ArrayAttr(indices)); | ||||||
| return success(); | ||||||
| } | ||||||
|
|
||||||
| private: | ||||||
| unsigned targetVectorBitWidth; | ||||||
| }; | ||||||
|
|
||||||
| /// This pattern converts the ExtractOp to a ShuffleOp that works on a | ||||||
| /// linearized vector. | ||||||
| /// Following, | ||||||
| /// vector.extract %source [ position ] | ||||||
| /// is converted to : | ||||||
| /// %source_1d = vector.shape_cast %source | ||||||
| /// %out_1d = vector.shuffle %source_1d, %source_1d [ shuffle_indices_1d ] | ||||||
| /// %out_nd = vector.shape_cast %out_1d | ||||||
| /// `shuffle_indices_1d` is computed using the position of the original extract. | ||||||
| struct LinearizeVectorExtract final | ||||||
| : public OpConversionPattern<vector::ExtractOp> { | ||||||
| using OpConversionPattern::OpConversionPattern; | ||||||
| LinearizeVectorExtract( | ||||||
| const TypeConverter &typeConverter, MLIRContext *context, | ||||||
| unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(), | ||||||
| PatternBenefit benefit = 1) | ||||||
| : OpConversionPattern(typeConverter, context, benefit), | ||||||
| targetVectorBitWidth(targetVectBitWidth) {} | ||||||
| LogicalResult | ||||||
| matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor, | ||||||
| ConversionPatternRewriter &rewriter) const override { | ||||||
| auto dstTy = getTypeConverter()->convertType(extractOp.getType()); | ||||||
| assert(!(extractOp.getVector().getType().isScalable() || | ||||||
| dstTy.cast<VectorType>().isScalable()) && | ||||||
| "scalable vectors are not supported."); | ||||||
| if (!isLessThanTargetBitWidth(extractOp, targetVectorBitWidth)) | ||||||
| return rewriter.notifyMatchFailure( | ||||||
| extractOp, "Can't flatten since targetBitWidth <= OpSize"); | ||||||
|
|
||||||
| // Dynamic position is not supported. | ||||||
| if (extractOp.hasDynamicPosition()) | ||||||
| return rewriter.notifyMatchFailure(extractOp, | ||||||
| "dynamic position is not supported."); | ||||||
|
|
||||||
| auto shape = extractOp.getVector().getType().getShape(); | ||||||
| auto size = extractOp.getVector().getType().getNumElements(); | ||||||
|
|
||||||
| // Compute linearized offset. | ||||||
| int64_t linearizedOffset = 0; | ||||||
| auto offsets = extractOp.getStaticPosition(); | ||||||
| for (auto [i, off] : llvm::enumerate(offsets)) { | ||||||
| size /= shape[i]; | ||||||
| linearizedOffset += offsets[i] * size; | ||||||
| } | ||||||
|
|
||||||
| llvm::SmallVector<int64_t, 2> indices(size); | ||||||
| std::iota(indices.begin(), indices.end(), linearizedOffset); | ||||||
| rewriter.replaceOpWithNewOp<vector::ShuffleOp>( | ||||||
| extractOp, dstTy, adaptor.getVector(), adaptor.getVector(), | ||||||
| rewriter.getI64ArrayAttr(indices)); | ||||||
|
|
||||||
| return success(); | ||||||
| } | ||||||
|
|
||||||
| private: | ||||||
| unsigned targetVectorBitWidth; | ||||||
| }; | ||||||
|
|
@@ -145,3 +396,21 @@ void mlir::vector::populateVectorLinearizeTypeConversionsAndLegality( | |||||
| patterns.add<LinearizeConstant, LinearizeVectorizable>( | ||||||
| typeConverter, patterns.getContext(), targetBitWidth); | ||||||
| } | ||||||
|
|
||||||
| void mlir::vector::populateVectorLinearizeToShuffleRewritePatterns( | ||||||
| TypeConverter &typeConverter, RewritePatternSet &patterns, | ||||||
| ConversionTarget &target, unsigned int targetBitWidth) { | ||||||
| target.addDynamicallyLegalOp<vector::ShuffleOp>( | ||||||
| [=](vector::ShuffleOp shuffleOp) -> bool { | ||||||
| return isLessThanTargetBitWidth(shuffleOp, targetBitWidth) | ||||||
| ? (typeConverter.isLegal(shuffleOp) && | ||||||
| shuffleOp.getResult() | ||||||
| .getType() | ||||||
| .cast<mlir::VectorType>() | ||||||
| .getRank() == 1) | ||||||
| : true; | ||||||
| }); | ||||||
| patterns.add<LinearizeVectorShuffle, LinearizeVectorExtract, | ||||||
| LinearizeVectorExtractStridedSlice>( | ||||||
| typeConverter, patterns.getContext(), targetBitWidth); | ||||||
| } | ||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Extract ops are "shuffle-like" ops so perhaps we can call this
populateVectorLinearizeShuffleLikeOpsPatternsor similarThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
renamed to
populateVectorLinearizeShuffleLikeOpsPatterns