Make createReadOrMaskedRead and isValidMaskedInputVector vector utilities #89119

LLITCHEV · 2024-04-17T18:43:46Z

Made the createReadOrMaskedRead and isValidMaskedInputVector utility functions - to be accessible outside of the CU. Needed by the IREE new TopK implementation.

Made the createReadOrMaskedRead a utility function - to be accessible outside of the CU. Needed by the IREE new TopK implementation.

github-actions · 2024-04-17T18:44:04Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2024-04-17T18:44:34Z

@llvm/pr-subscribers-mlir-vector
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Lubomir Litchev (LLITCHEV)

Changes

Made the createReadOrMaskedRead a utility function - to be accessible outside of the CU. Needed by the IREE new TopK implementation.

Full diff: https://github.com/llvm/llvm-project/pull/89119.diff

3 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h (+6)
(modified) mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp (+29)
(modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (-40)

diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index feb3b3f03cf538..f4c56b671e9d7e 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1616,6 +1616,12 @@ void populateSplitReductionPattern(
     const ControlSplitReductionFn &controlSplitReductionFn,
     bool useAlloc = false);
 
+/// 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.
+Value createReadOrMaskedRead(OpBuilder &builder, Location loc,
+                                    Value source, ArrayRef<int64_t> readShape,
+                                    Value padValue);
 } // namespace linalg
 } // namespace mlir
 
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index a17bc8e4cd318f..b32ebfc380fcfb 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -1593,3 +1593,32 @@ void linalg::populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
       DownscaleSizeOneWindowed2DConvolution<PoolingNchwMaxOp, PoolingNcwMaxOp>>(
       patterns.getContext(), benefit);
 }
+
+Value mlir::linalg::createReadOrMaskedRead(OpBuilder &builder, Location loc,
+                                    Value source, ArrayRef<int64_t> readShape,
+                                    Value padValue) {
+  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);
+  auto transferReadOp = builder.create<vector::TransferReadOp>(
+      loc,
+      /*vectorType=*/vectorType,
+      /*source=*/source,
+      /*indices=*/SmallVector<Value>(readRank, zero),
+      /*padding=*/padValue,
+      /*inBounds=*/SmallVector<bool>(readRank, true));
+  if (llvm::equal(readShape, sourceShape)) {
+    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);
+}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index df61381432921b..e2ca5e14377286 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

github-actions · 2024-04-17T18:46:45Z

✅ With the latest revision this PR passed the C/C++ code formatter.

matthias-springer

This is just moving code, but I think it's worth cleaning it up a bit.

matthias-springer · 2024-04-18T08:13:10Z

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

+                                           ArrayRef<int64_t> readShape,
+                                           Value padValue, bool doMasking) {
+  assert(llvm::none_of(readShape,
+                       [](int64_t s) { return s == ShapedType::kDynamic; }));


nit: add && "expected static shape" for better error messages

matthias-springer · 2024-04-18T08:13:38Z

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

+                                           Value padValue, bool doMasking) {
+  assert(llvm::none_of(readShape,
+                       [](int64_t s) { return s == ShapedType::kDynamic; }));
+  auto sourceShape = dyn_cast<ShapedType>(source.getType()).getShape();


Turn into cast<...> because there is no check that the dyn_cast succeeded.

matthias-springer · 2024-04-18T08:16:11Z

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

+  if (!doMasking) {
+    // Update the inBounds attribute.
+    for (unsigned i = 0; i < readRank; i++)
+      inBoundsVal[i] = sourceShape[i] == readShape[i];


Can we change this so that the respective in_bounds is set to "false" if readShape > sourceShape or sourceShape is dynamic?

matthias-springer · 2024-04-18T08:21:43Z

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

@@ -1616,6 +1616,12 @@ void populateSplitReductionPattern(
    const ControlSplitReductionFn &controlSplitReductionFn,
    bool useAlloc = false);

+/// Create a TransferReadOp from `source` with static shape `readShape`. If the


Can you update the comment and describe doMasking. (Maybe I'd also rename it to just masking.)

Could you also include a comment on indices? These are hard-coded to be 0 - that's worth documenting.

(Maybe I'd also rename it to just masking.)

This is a bit of bikesheding and I don't mind that much ... But I do like when a bool variable includes a verb so that it's effectively a question with a binary answer (i.e. "yes"/"no") :) Personally I'd use enableMasking instead. Again, don't really mind 😅

Changed to enableMasking.

matthias-springer · 2024-04-18T08:22:22Z

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

+  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());


Add && "expected same rank"

matthias-springer · 2024-04-18T08:23:11Z

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

+  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());


I'd also put an assert(padValue.getType() == shapedType.getElementType() && "expected same pad element type to match source element type").

matthias-springer · 2024-04-18T08:24:16Z

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

+      /*padding=*/padValue,
+      /*inBounds=*/inBoundsVal);
+
+  if (llvm::equal(readShape, sourceShape) || !doMasking) {


nit: Trivial braces not needed

banach-space

Shouldn't createReadOrMaskedRead be a Vector dialect Util?

hanhanW

I think you also need to refactor isValidMaskedInputVector for your topk implementation?

matthias-springer · 2024-04-19T07:09:20Z

Shouldn't createReadOrMaskedRead be a Vector dialect Util?

Oh right, this shouldn't be in Linalg.

Moved code around and other CR.

LLITCHEV · 2024-04-19T14:53:47Z

This is just moving code, but I think it's worth cleaning it up a bit.

Fixed.

LLITCHEV · 2024-04-19T14:54:26Z

Shouldn't createReadOrMaskedRead be a Vector dialect Util?

Oh right, this shouldn't be in Linalg.

Fixed.

LLITCHEV · 2024-04-19T14:56:59Z

I think you also need to refactor isValidMaskedInputVector for your topk implementation?

I thought to do two separate PRs, but there is no reason why not to do it in one. Thanks!

LLITCHEV · 2024-04-19T14:57:29Z

Shouldn't createReadOrMaskedRead be a Vector dialect Util?

Oh right, this shouldn't be in Linalg.

Moved to VectorUtils.

banach-space · 2024-04-19T15:22:13Z

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

+      inBoundsVal[i] = (sourceShape[i] == readShape[i]) &&
+                       !ShapedType::isDynamic(sourceShape[] i);


Why not:

Suggested change

inBoundsVal[i] = (sourceShape[i] == readShape[i]) &&

!ShapedType::isDynamic(sourceShape[] i);

inBoundsVal[i] = (sourceShape[i] == readShape[i]) &&

!ShapedType::isDynamic(sourceShape[i]);

?

inBoundsVal[i] = (sourceShape[i] == readShape[i]) &&
!ShapedType::isDynamic(sourceShape[i]);

Yes... Fixed... That was a typo... :) Not sure how it got there. I formatted after building maybe some accidental undo. The CI cought the issue. :)

banach-space · 2024-04-19T15:23:32Z

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

+  if (llvm::equal(readShape, sourceShape) || !doMasking)
+    return transferReadOp;


Suggested change

if (llvm::equal(readShape, sourceShape) || !doMasking)

return transferReadOp;

if (llvm::equal(readShape, sourceShape) || !enableMasking)

return transferReadOp;

Same as above. :) Fixed.

banach-space · 2024-04-19T15:40:23Z

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

+/// 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.
+/// enableMasking if false, the inBoundsVal values are set properly, based on
+///     the rank dimensions of the source and destination tensors.


Please update this comment.

Fix formatting

"then a mask is created on the read." - the mask is never created if enableMasking is false.

Missing note on indices being hard-coded as 0.

Point 2. is rather counter-intuitive to me. @pashu123, you introduced doMasking in #88249. Looking at the implementation and you PR, it feels like you meant something like bool useInBoundsInsteadOfMasking rather than doMasking/enableMasking? As in:

BEFORE #88249:

if (llvm::equal(readShape, sourceShape)) // use masking

AFTER #88249:

if (llvm::equal(readShape, sourceShape) && !useInBoundsInsteadOfMasking) // use masking

Put differently, the extra logic added to createReadOrMaskedRead in #88249 reads to me as:

"In some cases I want to disable masking and use in_bounds attr instead", rather than

"I want a dedicated toggle to enable/disable masking".

I'm asking to better understand the current logic 😅

"I feel the same way too. Using useInBoundsInsteadOfMasking provides a better understanding of the entire context." I can modify the variables.

@banach-space To me this is just a potential optimization (maybe I'm wrong). It looks like the code is setting up for a later optimization to turn read_transfer into a simple load. But if we do masking or the bounds of a dimension is different, the optimization can't be done. It alsocan't be done if masking is done. So, it looked to me that if no masking and bounds the same, set it up for later pass to do load instead of a transfer (with skipping the masking op). Or just maybe just removing an unnecessary op - masking. Maybe I misunderstood the purpose of it too. That's what I thought... :)
The rest is fixed, I hope... :)
Thanks!

useInBoundsInsteadOfMasking

I renamed them, since I'm moving the code as utility methods. Thanks!

pashu123 · 2024-04-19T18:00:17Z

@LLITCHEV If you could rename the variables here:

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Line 1502 in 31a0b27

bool doMasking = true;

.

pashu123 · 2024-04-19T18:10:55Z

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

+/// mask is created on the read, if use of mask is specified or the bounds on a
+/// dimension are different.
+///
+/// `enableMasking` if false, the inBoundsVal values are set properly, based on


We need to update the doc comments w.r.t useInBoundsInsteadOfMasking.

Ops... Missed that one. Good catch. Fixed. Thanks!

LLITCHEV · 2024-04-22T17:27:18Z

@LLITCHEV If you could rename the variables here:

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Line 1502 in 31a0b27

bool doMasking = true;

.

Renamed.

banach-space

Thanks for addressing my comments 🙏🏻

github-actions · 2024-04-23T00:19:12Z

@LLITCHEV Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested
by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as
the builds can include changes from many authors. It is not uncommon for your
change to be included in a build that fails due to someone else's changes, or
infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself.
This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

The PR llvm#89119 renamed a flag, which inverted it's previous meaning, but the logic was not updated to reflact that. Fixed the logic to reflect the inversion of the meaning of the name.

Make createReadOrMaskedRead a utility

ba6d1ec

Made the createReadOrMaskedRead a utility function - to be accessible outside of the CU. Needed by the IREE new TopK implementation.

LLITCHEV requested review from hanhanW and nicolasvasilache as code owners April 17, 2024 18:43

llvmbot added mlir:linalg mlir labels Apr 17, 2024

LLITCHEV added 2 commits April 17, 2024 18:48

Formatting fixes.

10823ad

Merge of the latest.

0898983

LLITCHEV mentioned this pull request Apr 17, 2024

[Codegen][TopK] Faster TopK implementation iree-org/iree#17045

Open

matthias-springer approved these changes Apr 18, 2024

View reviewed changes

matthias-springer reviewed Apr 18, 2024

View reviewed changes

banach-space reviewed Apr 18, 2024

View reviewed changes

LLITCHEV requested a review from jpienaar April 18, 2024 18:02

hanhanW reviewed Apr 18, 2024

View reviewed changes

Addressed CR requests.

d31a4cf

Moved code around and other CR.

llvmbot added mlir:vectorops mlir:vector labels Apr 19, 2024

LLITCHEV changed the title ~~Make createReadOrMaskedRead a utility~~ Make createReadOrMaskedRead and isValidMaskedInputVector vector utilities Apr 19, 2024

dcaballe approved these changes Apr 19, 2024

View reviewed changes

banach-space reviewed Apr 19, 2024

View reviewed changes

banach-space mentioned this pull request Apr 19, 2024

[mlir][linalg] Vectorize unpack op without masking #89067

Merged

LLITCHEV marked this pull request as draft April 19, 2024 16:46

Fixed some syntax errors, comments and other CRF.

1897a65

More CR addressed.

31a0b27

pashu123 reviewed Apr 19, 2024

View reviewed changes

LLITCHEV added 2 commits April 19, 2024 21:01

Fixed a missed comment update.

01d1e46

Renamed a var per CR feedback.

0a064df

LLITCHEV marked this pull request as ready for review April 22, 2024 17:28

hanhanW approved these changes Apr 22, 2024

View reviewed changes

banach-space approved these changes Apr 22, 2024

View reviewed changes

LLITCHEV merged commit 30d4f6a into llvm:main Apr 23, 2024
6 checks passed

LLITCHEV mentioned this pull request Apr 25, 2024

Update logic to reflect name of flag #90084

Closed

		inBoundsVal[i] = (sourceShape[i] == readShape[i]) &&
		!ShapedType::isDynamic(sourceShape[] i);

		if (llvm::equal(readShape, sourceShape) \|\| !doMasking)
		return transferReadOp;

Make createReadOrMaskedRead and isValidMaskedInputVector vector utilities #89119

Make createReadOrMaskedRead and isValidMaskedInputVector vector utilities #89119

Conversation

LLITCHEV commented Apr 17, 2024 • edited Loading

github-actions bot commented Apr 17, 2024

llvmbot commented Apr 17, 2024 • edited Loading

github-actions bot commented Apr 17, 2024 • edited Loading

matthias-springer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

banach-space left a comment

Choose a reason for hiding this comment

hanhanW left a comment

Choose a reason for hiding this comment

matthias-springer commented Apr 19, 2024

LLITCHEV commented Apr 19, 2024

LLITCHEV commented Apr 19, 2024

LLITCHEV commented Apr 19, 2024

LLITCHEV commented Apr 19, 2024

Choose a reason for hiding this comment

LLITCHEV Apr 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pashu123 commented Apr 19, 2024

Choose a reason for hiding this comment

LLITCHEV Apr 19, 2024 • edited Loading

Choose a reason for hiding this comment

LLITCHEV commented Apr 22, 2024

banach-space left a comment

Choose a reason for hiding this comment

github-actions bot commented Apr 23, 2024

LLITCHEV commented Apr 17, 2024 •

edited

Loading

llvmbot commented Apr 17, 2024 •

edited

Loading

github-actions bot commented Apr 17, 2024 •

edited

Loading

LLITCHEV Apr 19, 2024 •

edited

Loading

LLITCHEV Apr 19, 2024 •

edited

Loading