[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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,7 +15,9 @@ | |
| #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h" | ||
| #include "mlir/IR/PatternMatch.h" | ||
| #include "mlir/IR/TypeUtilities.h" | ||
| #include "mlir/Support/LogicalResult.h" | ||
| #include "mlir/Transforms/DialectConversion.h" | ||
| #include <numeric> | ||
|
|
||
| using namespace mlir; | ||
|
|
||
|
|
@@ -103,6 +105,234 @@ struct LinearizeVectorizable final | |
| return success(); | ||
| } | ||
|
|
||
| private: | ||
| unsigned targetVectorBitWidth; | ||
| }; | ||
|
|
||
| struct LinearizeVectorExtractStridedSlice final | ||
| : 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 |
||
| auto loc = extractOp.getLoc(); | ||
| if (!dstType) | ||
| return rewriter.notifyMatchFailure(loc, "cannot convert type."); | ||
| if (extractOp.getVector().getType().isScalable() || | ||
| dstType.cast<VectorType>().isScalable()) | ||
| return rewriter.notifyMatchFailure(loc, | ||
| "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().getValue(); | ||
| auto sizes = extractOp.getSizes().getValue(); | ||
| auto strides = extractOp.getStrides().getValue(); | ||
banach-space marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| 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 | ||
| // of the extraction is greater than 1. In this case last (n-k) dimensions | ||
| // form the extraction granularity. example : %0 = | ||
| // 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 extractSliceLen = 1; | ||
|
There was a problem hiding this comment. This represents "flattened slice len", right? Why not There was a problem hiding this comment. addressed above. basically the extraction can happen only at first k dimensions of an n-D vectors. the remaining dimensions are extracted in full. so the last n-k dimensions form the extraction granularity. this is the meaning I wanted to capture in the naming. |
||
| 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. |
||
| auto k = (int64_t)offsets.size(); | ||
|
There was a problem hiding this comment. nit: we should spell out |
||
| 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 |
||
| extractSliceLen *= extractOp.getSourceVectorType().getShape()[i + k]; | ||
| } | ||
| } | ||
|
|
||
| // get total number of extracted slices | ||
|
There was a problem hiding this comment. nit: we should use capital letters and periods in comments per coding guidelines. There was a problem hiding this comment. fixed all comments. |
||
| int64_t nExtractedSlices = 1; | ||
| for (auto size : sizes) { | ||
| nExtractedSlices *= size.cast<IntegerAttr>().getInt(); | ||
| } | ||
|
There was a problem hiding this comment. Could you try using There was a problem hiding this comment. tried this. but this require constructing an |
||
|
|
||
| // compute the strides of the source vector considering first k dimensions | ||
| llvm::SmallVector<int64_t, 4> sourceStrides(k, extractSliceLen); | ||
| for (int i = k - 2; i >= 0; --i) { | ||
| sourceStrides[i] = sourceStrides[i + 1] * | ||
| extractOp.getSourceVectorType().getShape()[i + 1]; | ||
| } | ||
| // final shuffle indices has nExtractedElems * extractSliceLen elements | ||
| llvm::SmallVector<int64_t, 4> indices(nExtractedSlices * extractSliceLen); | ||
| // 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 nExtractedElems-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 + | ||
| // sliceLen | ||
| for (int64_t j = 0; j < extractSliceLen; ++j) { | ||
| indices[i * extractSliceLen + 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; | ||
| }; | ||
|
|
||
| struct LinearizeVectorShffle final | ||
|
|
||
| : public OpConversionPattern<vector::ShuffleOp> { | ||
| using OpConversionPattern::OpConversionPattern; | ||
| LinearizeVectorShffle( | ||
| 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()); | ||
| auto loc = shuffleOp.getLoc(); | ||
| if (!dstType) | ||
| return rewriter.notifyMatchFailure(loc, "cannot convert type."); | ||
|
|
||
| if (shuffleOp.getV1VectorType().isScalable() || | ||
| shuffleOp.getV2VectorType().isScalable() || | ||
| dstType.cast<VectorType>().isScalable()) | ||
| return rewriter.notifyMatchFailure(loc, | ||
| "scalable vectors are not supported."); | ||
banach-space marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| 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]; | ||
| } | ||
| } | ||
|
|
||
| auto mask = shuffleOp.getMask(); | ||
| auto totalSize = mask.size() * shuffleSliceLen; | ||
banach-space marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| llvm::SmallVector<int64_t, 2> indices(totalSize); | ||
| 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; | ||
| }; | ||
|
|
||
| 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()); | ||
| if (!dstTy) | ||
| return rewriter.notifyMatchFailure(extractOp, "cannot convert type."); | ||
|
|
||
| if (extractOp.getVector().getType().isScalable() || | ||
| dstTy.cast<VectorType>().isScalable()) | ||
| return rewriter.notifyMatchFailure(extractOp, | ||
| "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; | ||
| }; | ||
|
|
@@ -139,9 +369,19 @@ void mlir::vector::populateVectorLinearizeTypeConversionsAndLegality( | |
| ? typeConverter.isLegal(op) | ||
| : true); | ||
| } | ||
| if (isa<vector::ShuffleOp>(op)) { | ||
| return (isLessThanTargetBitWidth(op, targetBitWidth) | ||
| ? (typeConverter.isLegal(op) && | ||
| op->getResult(0) | ||
| .getType() | ||
| .cast<mlir::VectorType>() | ||
| .getRank() == 1) | ||
| : true); | ||
| } | ||
| return std::nullopt; | ||
| }); | ||
|
|
||
| patterns.add<LinearizeConstant, LinearizeVectorizable, LinearizeVectorShffle, | ||
| LinearizeVectorExtract, LinearizeVectorExtractStridedSlice>( | ||
|
There was a problem hiding this comment. Could we move this to a different There was a problem hiding this comment. I added a new populate method called |
||
| 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
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.
Uh oh!
There was an error while loading. Please reload this page.