[Mlir-commits] [mlir] [mlir][vector] Add support for linearizing Extract, ExtractStridedSlice, Shuffle VectorOps in VectorLinearize (PR #88204)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Apr 9 15:40:14 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector
Author: Charitha Saumya (charithaintc)
<details>
<summary>Changes</summary>
---
Full diff: https://github.com/llvm/llvm-project/pull/88204.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp (+241-1)
- (modified) mlir/test/Dialect/Vector/linearize.mlir (+92)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index b59e9062e5a08e..e5157abd245b5d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -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) {}
+
+ LogicalResult
+ matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto dstType = getTypeConverter()->convertType(extractOp.getType());
+ 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");
+
+ auto offsets = extractOp.getOffsets().getValue();
+ auto sizes = extractOp.getSizes().getValue();
+ auto strides = extractOp.getStrides().getValue();
+
+ 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;
+ auto n = extractOp.getSourceVectorType().getRank();
+ auto k = (int64_t)offsets.size();
+ if (n > k) {
+ for (unsigned i = 0; i < n - k; i++) {
+ extractSliceLen *= extractOp.getSourceVectorType().getShape()[i + k];
+ }
+ }
+
+ // get total number of extracted slices
+ int64_t nExtractedSlices = 1;
+ for (auto size : sizes) {
+ nExtractedSlices *= size.cast<IntegerAttr>().getInt();
+ }
+
+ // 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.");
+ 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;
+
+ 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>(
+ patterns.add<LinearizeConstant, LinearizeVectorizable, LinearizeVectorShffle,
+ LinearizeVectorExtract, LinearizeVectorExtractStridedSlice>(
typeConverter, patterns.getContext(), targetBitWidth);
}
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 22be78cd682057..67f0f667a6b205 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -153,3 +153,95 @@ func.func @test_0d_vector() -> vector<f32> {
// ALL: return %[[CST]]
return %0 : vector<f32>
}
+
+// -----
+// ALL-LABEL: test_extract_strided_slice_1
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<4x8xf32>) -> vector<2x2xf32> {
+func.func @test_extract_strided_slice_1(%arg0 : vector<4x8xf32>) -> vector<2x2xf32> {
+ // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<4x8xf32> to vector<32xf32>
+ // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+ // DEFAULT: [4, 5, 12, 13] : vector<32xf32>, vector<32xf32>
+ // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<4xf32> to vector<2x2xf32>
+ // DEFAULT: return %[[RES]] : vector<2x2xf32
+
+ // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<4x8xf32> to vector<32xf32>
+ // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+ // BW-128: [4, 5, 12, 13] : vector<32xf32>, vector<32xf32>
+ // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<4xf32> to vector<2x2xf32>
+ // BW-128: return %[[RES]] : vector<2x2xf32>
+
+ // BW-0: %[[RES:.*]] = vector.extract_strided_slice %[[ARG:.*]] {offsets = [0, 4], sizes = [2, 2], strides = [1, 1]} : vector<4x8xf32> to vector<2x2xf32>
+ // BW-0: return %[[RES]] : vector<2x2xf32>
+ %0 = vector.extract_strided_slice %arg0 { sizes = [2, 2], strides = [1, 1], offsets = [0, 4]}
+ : vector<4x8xf32> to vector<2x2xf32>
+ return %0 : vector<2x2xf32>
+}
+
+// -----
+// ALL-LABEL: test_extract_strided_slice_2
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<2x8x2xf32>) -> vector<1x4x2xf32> {
+func.func @test_extract_strided_slice_2(%arg0 : vector<2x8x2xf32>) -> vector<1x4x2xf32> {
+ // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+ // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+ // DEFAULT: [20, 21, 22, 23, 24, 25, 26, 27] : vector<32xf32>, vector<32xf32>
+ // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<8xf32> to vector<1x4x2xf32>
+ // DEFAULT: return %[[RES]] : vector<1x4x2xf32>
+
+ // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+ // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+ // BW-128: [20, 21, 22, 23, 24, 25, 26, 27] : vector<32xf32>, vector<32xf32>
+ // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<8xf32> to vector<1x4x2xf32>
+ // BW-128: return %[[RES]] : vector<1x4x2xf32>
+
+ // BW-0: %[[RES:.*]] = vector.extract_strided_slice %[[ORIG_ARG]] {offsets = [1, 2], sizes = [1, 4], strides = [1, 1]} : vector<2x8x2xf32> to vector<1x4x2xf32>
+ // BW-0: return %[[RES]] : vector<1x4x2xf32>
+ %0 = vector.extract_strided_slice %arg0 { offsets = [1, 2], strides = [1, 1], sizes = [1, 4] }
+ : vector<2x8x2xf32> to vector<1x4x2xf32>
+ return %0 : vector<1x4x2xf32>
+}
+
+// -----
+// ALL-LABEL: test_vector_shuffle
+// ALL-SAME: (%[[ORIG_ARG0:.*]]: vector<4x2xf32>, %[[ORIG_ARG1:.*]]: vector<4x2xf32>) -> vector<8x2xf32> {
+func.func @test_vector_shuffle(%arg0: vector<4x2xf32>, %arg1: vector<4x2xf32>) -> vector<8x2xf32> {
+ // DEFAULT: %[[ARG0:.*]] = vector.shape_cast %[[ORIG_ARG0]] : vector<4x2xf32> to vector<8xf32>
+ // DEFAULT: %[[ARG1:.*]] = vector.shape_cast %[[ORIG_ARG1]] : vector<4x2xf32> to vector<8xf32>
+ // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG0]], %[[ARG1]]
+ // DEFAULT: [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15] : vector<8xf32>, vector<8xf32>
+ // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+ // DEFAULT: return %[[RES]] : vector<8x2xf32>
+
+ // BW-128: %[[ARG0:.*]] = vector.shape_cast %[[ORIG_ARG0]] : vector<4x2xf32> to vector<8xf32>
+ // BW-128: %[[ARG1:.*]] = vector.shape_cast %[[ORIG_ARG1]] : vector<4x2xf32> to vector<8xf32>
+ // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG0]], %[[ARG1]]
+ // BW-128: [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15] : vector<8xf32>, vector<8xf32>
+ // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+ // BW-128: return %[[RES]] : vector<8x2xf32>
+
+ // BW-0: %[[RES:.*]] = vector.shuffle %[[ORIG_ARG0]], %[[ORIG_ARG1]] [0, 4, 1, 5, 2, 6, 3, 7] : vector<4x2xf32>, vector<4x2xf32>
+ // BW-0: return %[[RES]] : vector<8x2xf32>
+ %0 = vector.shuffle %arg0, %arg1 [0, 4, 1, 5, 2, 6, 3, 7] : vector<4x2xf32>, vector<4x2xf32>
+ return %0 : vector<8x2xf32>
+}
+
+// -----
+// ALL-LABEL: test_vector_extract
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<2x8x2xf32>) -> vector<8x2xf32> {
+func.func @test_vector_extract(%arg0: vector<2x8x2xf32>) -> vector<8x2xf32> {
+ // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+ // DEFAULT: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+ // DEFAULT: [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31] : vector<32xf32>, vector<32xf32>
+ // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+ // DEFAULT: return %[[RES]] : vector<8x2xf32>
+
+ // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>
+ // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]
+ // BW-128: [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31] : vector<32xf32>, vector<32xf32>
+ // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>
+ // BW-128: return %[[RES]] : vector<8x2xf32>
+
+ // BW-0: %[[RES:.*]] = vector.extract %[[ORIG_ARG]][1] : vector<8x2xf32> from vector<2x8x2xf32>
+ // BW-0: return %[[RES]] : vector<8x2xf32>
+ %0 = vector.extract %arg0[1]: vector<8x2xf32> from vector<2x8x2xf32>
+ return %0 : vector<8x2xf32>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/88204
More information about the Mlir-commits
mailing list