[Mlir-commits] [mlir] 30d4f6a - Make createReadOrMaskedRead and isValidMaskedInputVector vector utilities (#89119)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Apr 22 17:18:50 PDT 2024
Author: Lubomir Litchev
Date: 2024-04-22T17:18:45-07:00
New Revision: 30d4f6afc9d594fff309121116d8373b66fb1a24
URL: https://github.com/llvm/llvm-project/commit/30d4f6afc9d594fff309121116d8373b66fb1a24
DIFF: https://github.com/llvm/llvm-project/commit/30d4f6afc9d594fff309121116d8373b66fb1a24.diff
LOG: Make createReadOrMaskedRead and isValidMaskedInputVector vector utilities (#89119)
Made the createReadOrMaskedRead and isValidMaskedInputVector utility
functions - to be accessible outside of the CU. Needed by the IREE new
TopK implementation.
Added:
Modified:
mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index f88fbdf9e62765..8a57c6094c41c0 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -180,6 +180,30 @@ struct MaskableOpRewritePattern : OpRewritePattern<SourceOp> {
/// are not linearizable.
bool isLinearizableVector(VectorType type);
+/// Create a TransferReadOp from `source` with static shape `readShape`. If the
+/// vector type for the read is not the same as the type of `source`, then a
+/// mask is created on the read, if use of mask is specified or the bounds on a
+/// dimension are
diff erent.
+///
+/// `useInBoundsInsteadOfMasking` if false, the inBoundsVal values are set
+/// properly, based on
+/// the rank dimensions of the source and destination tensors. And that is
+/// what determines if masking is done.
+///
+/// Note that the internal `vector::TransferReadOp` always read at indices zero
+/// for each dimension of the passed in tensor.
+Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
+ ArrayRef<int64_t> readShape, Value padValue,
+ bool useInBoundsInsteadOfMasking = true);
+
+/// Returns success if `inputVectorSizes` is a valid masking configuraion for
+/// given `shape`, i.e., it meets:
+/// 1. The numbers of elements in both array are equal.
+/// 2. `inputVectorSizes` does not have dynamic dimensions.
+/// 3. All the values in `inputVectorSizes` are greater than or equal to
+/// static sizes in `shape`.
+LogicalResult isValidMaskedInputVector(ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> inputVectorSizes);
} // namespace vector
/// Constructs a permutation map of invariant memref indices to vector
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index fbff2088637f44..e836f0dc63b4f9 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1410,46 +1410,6 @@ static SmallVector<int64_t> getTiledPackShape(tensor::PackOp packOp,
return applyPermutation(destShape, tensor::getPackInverseDestPerm(packOp));
}
-/// Create a TransferReadOp from `source` with static shape `readShape`. If the
-/// vector type for the read is not the same as the type of `source`, then a
-/// mask is created on the read. If `doMasking` parameter is set to false we
-/// update the `inBounds` attribute instead of masking.
-static Value createReadOrMaskedRead(OpBuilder &builder, Location loc,
- Value source, ArrayRef<int64_t> readShape,
- Value padValue, bool doMasking = true) {
- assert(llvm::none_of(readShape,
- [](int64_t s) { return s == ShapedType::kDynamic; }));
- auto sourceShape = dyn_cast<ShapedType>(source.getType()).getShape();
- assert(sourceShape.size() == readShape.size());
- auto maskType = VectorType::get(readShape, builder.getI1Type());
- auto vectorType = VectorType::get(readShape, padValue.getType());
- int64_t readRank = readShape.size();
- auto zero = builder.create<arith::ConstantIndexOp>(loc, 0);
- SmallVector<bool> inBoundsVal(readRank, true);
- if (!doMasking) {
- // Update the inBounds attribute.
- for (unsigned i = 0; i < readRank; i++)
- inBoundsVal[i] = sourceShape[i] == readShape[i];
- }
- auto transferReadOp = builder.create<vector::TransferReadOp>(
- loc,
- /*vectorType=*/vectorType,
- /*source=*/source,
- /*indices=*/SmallVector<Value>(readRank, zero),
- /*padding=*/padValue,
- /*inBounds=*/inBoundsVal);
-
- if (llvm::equal(readShape, sourceShape) || !doMasking) {
- return transferReadOp;
- }
- SmallVector<OpFoldResult> mixedSourceDims =
- tensor::getMixedSizes(builder, loc, source);
- Value mask =
- builder.create<vector::CreateMaskOp>(loc, maskType, mixedSourceDims);
- return mlir::vector::maskOperation(builder, transferReadOp, mask)
- ->getResult(0);
-}
-
/// Given an input, the mixed destSizes, and the vector sizes for vectorization,
/// create an empty destination tensor and create a TransferWriteOp from the
/// input to the empty tensor. If the destination shape is not the same as the
@@ -1539,11 +1499,11 @@ vectorizeAsTensorPackOp(RewriterBase &rewriter, tensor::PackOp packOp,
// If the input vector sizes are not provided, then the vector sizes are
// determined by the result tensor shape. In case the vector sizes aren't
// provided, we update the inBounds attribute instead of masking.
- bool doMasking = true;
+ bool useInBoundsInsteadOfMasking = true;
if (inputVectorSizes.empty()) {
ArrayRef<int64_t> resultTensorShape = packOp.getDestType().getShape();
inputVectorSizes = resultTensorShape.take_front(packOp.getSourceRank());
- doMasking = false;
+ useInBoundsInsteadOfMasking = false;
}
// Create masked TransferReadOp.
@@ -1556,8 +1516,9 @@ vectorizeAsTensorPackOp(RewriterBase &rewriter, tensor::PackOp packOp,
invertPermutationVector(outerDimsPerm));
for (auto [idx, size] : enumerate(innerTiles))
inputShape[innerDimsPos[idx]] *= size;
- auto maskedRead = createReadOrMaskedRead(rewriter, loc, packOp.getSource(),
- inputShape, padValue, doMasking);
+ auto maskedRead = vector::createReadOrMaskedRead(
+ rewriter, loc, packOp.getSource(), inputShape, padValue,
+ useInBoundsInsteadOfMasking);
// Create ShapeCastOp.
SmallVector<int64_t> destShape(inputVectorSizes);
@@ -1649,7 +1610,7 @@ vectorizeAsTensorUnpackOp(RewriterBase &rewriter, tensor::UnPackOp unpackOp,
// Read result, mask if necessary. If transferReadOp shape is not equal
// to shape of source, then a mask is necessary.
- Value readResult = createReadOrMaskedRead(
+ Value readResult = vector::createReadOrMaskedRead(
rewriter, loc, unpackOp.getSource(),
ArrayRef<int64_t>(readMaskShape.begin(), readMaskShape.end()), padValue);
@@ -1707,8 +1668,8 @@ vectorizeAsTensorPadOp(RewriterBase &rewriter, tensor::PadOp padOp,
.reifyResultShapes(rewriter, reifiedReturnShapes);
(void)status; // prevent unused variable warning on non-assert builds
assert(succeeded(status) && "failed to reify result shapes");
- auto maskedRead = createReadOrMaskedRead(rewriter, loc, padOp.getSource(),
- inputVectorSizes, padValue);
+ auto maskedRead = vector::createReadOrMaskedRead(
+ rewriter, loc, padOp.getSource(), inputVectorSizes, padValue);
Operation *write = createWriteOrMaskedWrite(
rewriter, loc, maskedRead, reifiedReturnShapes[0], inputVectorSizes);
newResults.push_back(write->getResult(0));
@@ -1781,41 +1742,6 @@ vectorizeDynamicLinalgOpPrecondition(linalg::LinalgOp op,
return success();
}
-/// Returns success if `inputVectorSizes` is a valid masking configuraion for
-/// given `shape`, i.e., it meets:
-/// 1. The numbers of elements in both array are equal.
-/// 2. `inputVectorSizes` does not have dynamic dimensions.
-/// 3. All the values in `inputVectorSizes` are greater than or equal to
-/// static sizes in `shape`.
-static LogicalResult
-isValidMaskedInputVector(ArrayRef<int64_t> shape,
- ArrayRef<int64_t> inputVectorSizes) {
- LDBG("Iteration space static sizes:");
- LLVM_DEBUG(llvm::interleaveComma(shape, llvm::dbgs()));
- LLVM_DEBUG(llvm::dbgs() << "\n");
-
- if (inputVectorSizes.size() != shape.size()) {
- LDBG("Input vector sizes don't match the number of loops");
- return failure();
- }
- if (ShapedType::isDynamicShape(inputVectorSizes)) {
- LDBG("Input vector sizes can't have dynamic dimensions");
- return failure();
- }
- if (!llvm::all_of(llvm::zip(shape, inputVectorSizes),
- [](std::tuple<int64_t, int64_t> sizePair) {
- int64_t staticSize = std::get<0>(sizePair);
- int64_t inputSize = std::get<1>(sizePair);
- return ShapedType::isDynamic(staticSize) ||
- staticSize <= inputSize;
- })) {
- LDBG("Input vector sizes must be greater than or equal to iteration space "
- "static sizes");
- return failure();
- }
- return success();
-}
-
/// Need to check if the inner-tiles are static/constant.
static LogicalResult
vectorizeUnPackOpPrecondition(tensor::UnPackOp unpackOp,
@@ -1829,7 +1755,7 @@ vectorizeUnPackOpPrecondition(tensor::UnPackOp unpackOp,
}
llvm::ArrayRef<int64_t> resultShape = unpackOp.getDestType().getShape();
if (!inputVectorSizes.empty() &&
- failed(isValidMaskedInputVector(resultShape, inputVectorSizes)))
+ failed(vector::isValidMaskedInputVector(resultShape, inputVectorSizes)))
return failure();
return success();
@@ -1843,8 +1769,8 @@ static LogicalResult vectorizeLinalgOpPrecondition(
return failure();
// Check API contract for input vector sizes.
if (!inputVectorSizes.empty() &&
- failed(isValidMaskedInputVector(linalgOp.getStaticLoopRanges(),
- inputVectorSizes)))
+ failed(vector::isValidMaskedInputVector(linalgOp.getStaticLoopRanges(),
+ inputVectorSizes)))
return failure();
if (linalgOp.hasDynamicShape() && failed(vectorizeDynamicLinalgOpPrecondition(
@@ -1920,7 +1846,7 @@ vectorizePackOpPrecondition(tensor::PackOp packOp,
}
if (!satisfyEmptyCond &&
- failed(isValidMaskedInputVector(
+ failed(vector::isValidMaskedInputVector(
resultTensorShape.take_front(packOp.getSourceRank()),
inputVectorSizes)))
return failure();
@@ -1945,7 +1871,8 @@ vectorizePadOpPrecondition(tensor::PadOp padOp,
}
ArrayRef<int64_t> resultTensorShape = padOp.getResultType().getShape();
- if (failed(isValidMaskedInputVector(resultTensorShape, inputVectorSizes)))
+ if (failed(vector::isValidMaskedInputVector(resultTensorShape,
+ inputVectorSizes)))
return failure();
if (llvm::any_of(padOp.getLow(), [](Value v) {
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index ebc6f5cbcaa9ed..fcaf1ec944b479 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -30,6 +30,11 @@
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
+#define DEBUG_TYPE "vector-utils"
+
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
+
using namespace mlir;
/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
@@ -322,3 +327,73 @@ bool vector::isLinearizableVector(VectorType type) {
auto numScalableDims = llvm::count(type.getScalableDims(), true);
return (type.getRank() > 1) && (numScalableDims <= 1);
}
+
+Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
+ Value source, ArrayRef<int64_t> readShape,
+ Value padValue,
+ bool useInBoundsInsteadOfMasking) {
+ assert(llvm::none_of(readShape,
+ [](int64_t s) { return s == ShapedType::kDynamic; }) &&
+ "expected static shape");
+ auto sourceShapedType = cast<ShapedType>(source.getType());
+ auto sourceShape = sourceShapedType.getShape();
+ assert(sourceShape.size() == readShape.size() && "expected same ranks.");
+ auto maskType = VectorType::get(readShape, builder.getI1Type());
+ auto vectorType = VectorType::get(readShape, padValue.getType());
+ assert(padValue.getType() == sourceShapedType.getElementType() &&
+ "expected same pad element type to match source element type");
+ int64_t readRank = readShape.size();
+ auto zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+ SmallVector<bool> inBoundsVal(readRank, true);
+ if (!useInBoundsInsteadOfMasking) {
+ // Update the inBounds attribute.
+ for (unsigned i = 0; i < readRank; i++)
+ inBoundsVal[i] = (sourceShape[i] == readShape[i]) &&
+ !ShapedType::isDynamic(sourceShape[i]);
+ }
+ auto transferReadOp = builder.create<vector::TransferReadOp>(
+ loc,
+ /*vectorType=*/vectorType,
+ /*source=*/source,
+ /*indices=*/SmallVector<Value>(readRank, zero),
+ /*padding=*/padValue,
+ /*inBounds=*/inBoundsVal);
+
+ if (llvm::equal(readShape, sourceShape) || !useInBoundsInsteadOfMasking)
+ return transferReadOp;
+ SmallVector<OpFoldResult> mixedSourceDims =
+ tensor::getMixedSizes(builder, loc, source);
+ Value mask =
+ builder.create<vector::CreateMaskOp>(loc, maskType, mixedSourceDims);
+ return mlir::vector::maskOperation(builder, transferReadOp, mask)
+ ->getResult(0);
+}
+
+LogicalResult
+vector::isValidMaskedInputVector(ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> inputVectorSizes) {
+ LDBG("Iteration space static sizes:");
+ LLVM_DEBUG(llvm::interleaveComma(shape, llvm::dbgs()));
+ LLVM_DEBUG(llvm::dbgs() << "\n");
+
+ if (inputVectorSizes.size() != shape.size()) {
+ LDBG("Input vector sizes don't match the number of loops");
+ return failure();
+ }
+ if (ShapedType::isDynamicShape(inputVectorSizes)) {
+ LDBG("Input vector sizes can't have dynamic dimensions");
+ return failure();
+ }
+ if (!llvm::all_of(llvm::zip(shape, inputVectorSizes),
+ [](std::tuple<int64_t, int64_t> sizePair) {
+ int64_t staticSize = std::get<0>(sizePair);
+ int64_t inputSize = std::get<1>(sizePair);
+ return ShapedType::isDynamic(staticSize) ||
+ staticSize <= inputSize;
+ })) {
+ LDBG("Input vector sizes must be greater than or equal to iteration space "
+ "static sizes");
+ return failure();
+ }
+ return success();
+}
More information about the Mlir-commits
mailing list