[mlir][vector] Prevent folding non memref-type gather into maskedload #135371

sagarkulkarni19 · 2025-04-11T14:05:52Z

This patch fixes an issue in the FoldContiguousGather pattern which was incorrectly folding vector.gather operations with contiguous indices into vector.maskedload operations regardless of the base operand type.

While vector.gather operations can work on both tensor and memref types, vector.maskedload operations are only valid for memref types. The pattern was incorrectly lowering a tensor-based gather into a masked-load, which is invalid.

This fix adds a type check to ensure the pattern only applies to memref-based gather operations.

github-actions · 2025-04-11T14:06:18Z

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 · 2025-04-11T14:06:51Z

@llvm/pr-subscribers-mlir-vector

Author: Sagar Kulkarni (sagarkulkarni19)

Changes

This patch fixes an issue in the FoldContiguousGather pattern which was incorrectly folding vector.gather operations with contiguous indices into vector.maskedload operations regardless of the base operand type.

While vector.gather operations can work on both tensor and memref types, vector.maskedload operations are only valid for memref types. The pattern was incorrectly lowering a tensor-based gather into a masked-load, which is invalid.

This fix adds a type check to ensure the pattern only applies to memref-based gather operations.

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

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+3)
(modified) mlir/test/Dialect/Vector/canonicalize.mlir (+12)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 98d98f067de14..8955438b57343 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5340,6 +5340,9 @@ class FoldContiguousGather final : public OpRewritePattern<GatherOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(GatherOp op,
                                 PatternRewriter &rewriter) const override {
+    if (!op.getBase().getType().isa<MemRefType>())
+      return failure();
+
     if (failed(isZeroBasedContiguousSeq(op.getIndexVec())))
       return failure();
 
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index b7db8ec834be7..7d9223696712d 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3149,6 +3149,18 @@ func.func @contiguous_gather_step(%base: memref<?xf32>,
 
 // -----
 
+// CHECK-LABEL: @dont_fold_tensor_type_contiguous_gather
+func.func @dont_fold_tensor_type_contiguous_gather(%base: tensor<8xf32>, %mask: vector<4xi1>, %pass_thru: vector<4xf32>) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+  %indices = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
+  // CHECK: vector.gather
+  // CHECK-NOT: vector.maskedload
+  %0 = vector.gather %base[%c0][%indices], %mask, %pass_thru : tensor<8xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @gather_broadcast(
 // TODO: Broadcast is not supported yet
 //       CHECK:   %[[R:.*]] = vector.gather

llvmbot · 2025-04-11T14:06:52Z

@llvm/pr-subscribers-mlir

Author: Sagar Kulkarni (sagarkulkarni19)

Changes

This patch fixes an issue in the FoldContiguousGather pattern which was incorrectly folding vector.gather operations with contiguous indices into vector.maskedload operations regardless of the base operand type.

While vector.gather operations can work on both tensor and memref types, vector.maskedload operations are only valid for memref types. The pattern was incorrectly lowering a tensor-based gather into a masked-load, which is invalid.

This fix adds a type check to ensure the pattern only applies to memref-based gather operations.

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

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+3)
(modified) mlir/test/Dialect/Vector/canonicalize.mlir (+12)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 98d98f067de14..8955438b57343 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5340,6 +5340,9 @@ class FoldContiguousGather final : public OpRewritePattern<GatherOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(GatherOp op,
                                 PatternRewriter &rewriter) const override {
+    if (!op.getBase().getType().isa<MemRefType>())
+      return failure();
+
     if (failed(isZeroBasedContiguousSeq(op.getIndexVec())))
       return failure();
 
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index b7db8ec834be7..7d9223696712d 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3149,6 +3149,18 @@ func.func @contiguous_gather_step(%base: memref<?xf32>,
 
 // -----
 
+// CHECK-LABEL: @dont_fold_tensor_type_contiguous_gather
+func.func @dont_fold_tensor_type_contiguous_gather(%base: tensor<8xf32>, %mask: vector<4xi1>, %pass_thru: vector<4xf32>) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+  %indices = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
+  // CHECK: vector.gather
+  // CHECK-NOT: vector.maskedload
+  %0 = vector.gather %base[%c0][%indices], %mask, %pass_thru : tensor<8xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @gather_broadcast(
 // TODO: Broadcast is not supported yet
 //       CHECK:   %[[R:.*]] = vector.gather

sagarkulkarni19 · 2025-04-11T14:16:08Z

@Hardcode84 @ashay

Hardcode84

This is a weird discrepancy between vector ops I wasn't aware of. Fix itself is LGTM.

Please also wait for @banach-space or @dcaballe review.

banach-space

The fix makes sense to me — thanks for taking care of it! I've left a few nits, but otherwise LGTM.

@dcaballe and I were actually just talking about the “duality” of read/write ops in the Vector dialect — we agree it's something we should clean up and this PR confirms that.

banach-space · 2025-04-11T16:33:36Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -5340,6 +5340,9 @@ class FoldContiguousGather final : public OpRewritePattern<GatherOp> {
  using OpRewritePattern::OpRewritePattern;
  LogicalResult matchAndRewrite(GatherOp op,
                                PatternRewriter &rewriter) const override {
+    if (!op.getBase().getType().isa<MemRefType>())
+      return failure();


[nit] Could you use notifyMatchFailure? Thanks!

Good point, done.

banach-space · 2025-04-11T16:35:02Z

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -3149,6 +3149,18 @@ func.func @contiguous_gather_step(%base: memref<?xf32>,

 // -----

+// CHECK-LABEL: @dont_fold_tensor_type_contiguous_gather
+func.func @dont_fold_tensor_type_contiguous_gather(%base: tensor<8xf32>, %mask: vector<4xi1>, %pass_thru: vector<4xf32>) -> vector<4xf32> {


no_fold is shorter than dont_fold :)

The naming format seems to be @contiguous_gather_{other_stuff}, so @dont_fold_tensor_type_contiguous_gather -> @no_fold_contiguous_gather_tensor (or something similar)

Changed it to @no_fold_contiguous_gather_tensor.

banach-space · 2025-04-11T16:35:32Z

mlir/test/Dialect/Vector/canonicalize.mlir

+  %indices = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
+  // CHECK: vector.gather
+  // CHECK-NOT: vector.maskedload
+  %0 = vector.gather %base[%c0][%indices], %mask, %pass_thru : tensor<8xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>


[nit] Would you mind splitting into two lines like we did above?

Hardcode84 · 2025-04-11T16:48:25Z

Side comment: gather takes AnyShaped so gather "loading" from vector will be accepted as well, which was probably never intended.

This patch fixes an issue in the FoldContiguousGather pattern which was incorrectly folding vector.gather operations with contiguous indices into vector.maskedload operations regardless of the base operand type. While vector.gather operations can work on both tensor and memref types, vector.maskedload operations are only valid for memref types. The pattern was incorrectly lowering a tensor-based gather into a masked-load, which is invalid. This fix adds a type check to ensure the pattern only applies to memref-based gather operations.

sagarkulkarni19 · 2025-04-11T20:55:08Z

@banach-space I've addressed the comments. Lmk if it looks good.

@Hardcode84 @banach-space Can one of you help me merge it in since I do not have merge permission. Thanks!

github-actions · 2025-04-12T01:16:10Z

@sagarkulkarni19 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!

* origin/main: [mlir][vector] Prevent folding non memref-type gather into maskedload (llvm#135371) [mlir][SMT] remove custom forall/exists builder because of asan memory leak [bazel] Fix a typo (llvm#135460) [bazel] Add support for SMT Dialect (llvm#135454) [clang] ASTImporter: fix SubstNonTypeTemplateParmExpr source location (llvm#135450) [RISCV] Don't fold offsets into auipc if offset is larger than the reference global variable. (llvm#135297) [gn] port d1fd977 [NFC][LLVM] Apply std::move to object being pushed back in findSymbolCommon (llvm#135290) [AMDGPU] Teach iterative schedulers about IGLP (llvm#134953)

This was introduced in: llvm#135371

This was introduced in: #135371

…135512) This was introduced in: llvm/llvm-project#135371

banach-space · 2025-04-15T05:56:36Z

Side comment: gather takes AnyShaped so gather "loading" from vector will be accepted as well, which was probably never intended.

Thanks for pointing this out: #135749

…llvm#135371) This patch fixes an issue in the FoldContiguousGather pattern which was incorrectly folding vector.gather operations with contiguous indices into vector.maskedload operations regardless of the base operand type. While vector.gather operations can work on both tensor and memref types, vector.maskedload operations are only valid for memref types. The pattern was incorrectly lowering a tensor-based gather into a masked-load, which is invalid. This fix adds a type check to ensure the pattern only applies to memref-based gather operations. Co-authored-by: Sagar Kulkarni <[email protected]>

This was introduced in: llvm#135371

sagarkulkarni19 requested review from banach-space, dcaballe, nicolasvasilache and Groverkss as code owners April 11, 2025 14:05

llvmbot added mlir:vectorops mlir mlir:vector labels Apr 11, 2025

Hardcode84 reviewed Apr 11, 2025

View reviewed changes

banach-space approved these changes Apr 11, 2025

View reviewed changes

sagarkulkarni19 force-pushed the fix-fold-contiguous-gather-pattern branch from 1e50115 to 34b56da Compare April 11, 2025 20:51

Merge branch 'main' into fix-fold-contiguous-gather-pattern

7d58b14

Hardcode84 merged commit 357e380 into llvm:main Apr 12, 2025
11 checks passed

kuhar added a commit to kuhar/llvm-project that referenced this pull request Apr 13, 2025

[mlir][vector] Fix deprecation warning for .isa. NFC.

1fb559f

This was introduced in: llvm#135371

kuhar mentioned this pull request Apr 13, 2025

[mlir][vector] Fix deprecation warning for .isa. NFC. #135512

Merged

kuhar added a commit that referenced this pull request Apr 13, 2025

[mlir][vector] Fix deprecation warning for .isa. NFC. (#135512)

b0acbbe

This was introduced in: #135371

llvm-sync bot pushed a commit to arm/arm-toolchain that referenced this pull request Apr 13, 2025

Automerge: [mlir][vector] Fix deprecation warning for .isa. NFC. (#…

41df357

…135512) This was introduced in: llvm/llvm-project#135371

sagarkulkarni19 deleted the fix-fold-contiguous-gather-pattern branch April 14, 2025 14:21

var-const pushed a commit to ldionne/llvm-project that referenced this pull request Apr 17, 2025

[mlir][vector] Fix deprecation warning for .isa. NFC. (llvm#135512)

cc13297

This was introduced in: llvm#135371

[mlir][vector] Prevent folding non memref-type gather into maskedload #135371

[mlir][vector] Prevent folding non memref-type gather into maskedload #135371

Uh oh!

Conversation

sagarkulkarni19 commented Apr 11, 2025

Uh oh!

github-actions bot commented Apr 11, 2025

Uh oh!

llvmbot commented Apr 11, 2025

Uh oh!

llvmbot commented Apr 11, 2025

Uh oh!

sagarkulkarni19 commented Apr 11, 2025

Uh oh!

Hardcode84 left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

sagarkulkarni19 Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

banach-space Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

sagarkulkarni19 Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

banach-space Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

sagarkulkarni19 Apr 11, 2025

Choose a reason for hiding this comment

Uh oh!

Hardcode84 commented Apr 11, 2025

Uh oh!

sagarkulkarni19 commented Apr 11, 2025

Uh oh!

Uh oh!

github-actions bot commented Apr 12, 2025

Uh oh!

banach-space commented Apr 15, 2025

Uh oh!

Uh oh!