From 6c4d25e5529f2766de171faac64c002634986bcb Mon Sep 17 00:00:00 2001 From: Prashant Kumar Date: Thu, 25 Apr 2024 22:02:09 +0530 Subject: [PATCH] [mlir][linalg] Fix the semantic use of a flag `useInBoundsInsteadOfMasking` was doing the opposite i.e., when set to true; was updating the mask instead of updating the inBounds. --- mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h | 2 +- mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp | 10 ++++++---- mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 4 ++-- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h index 8a57c6094c41c..030be328e97fd 100644 --- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h +++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h @@ -194,7 +194,7 @@ bool isLinearizableVector(VectorType type); /// for each dimension of the passed in tensor. Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source, ArrayRef readShape, Value padValue, - bool useInBoundsInsteadOfMasking = true); + bool useInBoundsInsteadOfMasking); /// Returns success if `inputVectorSizes` is a valid masking configuraion for /// given `shape`, i.e., it meets: diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp index e836f0dc63b4f..ef9a30be9a015 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -1499,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 useInBoundsInsteadOfMasking = true; + bool useInBoundsInsteadOfMasking = false; if (inputVectorSizes.empty()) { ArrayRef resultTensorShape = packOp.getDestType().getShape(); inputVectorSizes = resultTensorShape.take_front(packOp.getSourceRank()); - useInBoundsInsteadOfMasking = false; + useInBoundsInsteadOfMasking = true; } // Create masked TransferReadOp. @@ -1612,7 +1612,8 @@ vectorizeAsTensorUnpackOp(RewriterBase &rewriter, tensor::UnPackOp unpackOp, // to shape of source, then a mask is necessary. Value readResult = vector::createReadOrMaskedRead( rewriter, loc, unpackOp.getSource(), - ArrayRef(readMaskShape.begin(), readMaskShape.end()), padValue); + ArrayRef(readMaskShape.begin(), readMaskShape.end()), padValue, + /*useInBoundsInsteadOfMasking=*/false); PackingMetadata packMetadata; SmallVector lastDimToInsertPosPerm = @@ -1669,7 +1670,8 @@ vectorizeAsTensorPadOp(RewriterBase &rewriter, tensor::PadOp padOp, (void)status; // prevent unused variable warning on non-assert builds assert(succeeded(status) && "failed to reify result shapes"); auto maskedRead = vector::createReadOrMaskedRead( - rewriter, loc, padOp.getSource(), inputVectorSizes, padValue); + rewriter, loc, padOp.getSource(), inputVectorSizes, padValue, + /*useInBoundsInsteadOfMasking=*/false); Operation *write = createWriteOrMaskedWrite( rewriter, loc, maskedRead, reifiedReturnShapes[0], inputVectorSizes); newResults.push_back(write->getResult(0)); diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp index fcaf1ec944b47..6727f3f461722 100644 --- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp +++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp @@ -345,7 +345,7 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc, int64_t readRank = readShape.size(); auto zero = builder.create(loc, 0); SmallVector inBoundsVal(readRank, true); - if (!useInBoundsInsteadOfMasking) { + if (useInBoundsInsteadOfMasking) { // Update the inBounds attribute. for (unsigned i = 0; i < readRank; i++) inBoundsVal[i] = (sourceShape[i] == readShape[i]) && @@ -359,7 +359,7 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc, /*padding=*/padValue, /*inBounds=*/inBoundsVal); - if (llvm::equal(readShape, sourceShape) || !useInBoundsInsteadOfMasking) + if (llvm::equal(readShape, sourceShape) || useInBoundsInsteadOfMasking) return transferReadOp; SmallVector mixedSourceDims = tensor::getMixedSizes(builder, loc, source);