[mlir][vector] Add a check to ensure input vector rank equals target shape rank #127706

Prakhar-Dixit · 2025-02-18T22:26:22Z

Fixes issue #126197

The crash is caused because, during IR transformation, the vector-unrolling pass (using ExtractStridedSliceOp) attempts to slice an input vector of higher rank using a target vector of lower rank, which is not supported.

Specific example :

module {
  func.func @func1() {
    %cst_25 = arith.constant dense<3.718400e+04> : vector<4x2x2xf16>
    %cst_26 = arith.constant dense<1.000000e+00> : vector<24x2x2xf32>
    %47 = vector.fma %cst_26, %cst_26, %cst_26 : vector<24x2x2xf32>
    %818 = scf.execute_region -> vector<24x2x2xf32> {
        scf.yield %47 : vector<24x2x2xf32>
      }
    %823 = vector.extract_strided_slice %cst_25 {offsets = [2], sizes = [1], strides = [1]} : vector<4x2x2xf16> to vector<1x2x2xf16>
    return
  }
}

github-actions · 2025-02-18T22:26:40Z

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-02-18T22:27:14Z

@llvm/pr-subscribers-mlir

Author: Prakhar Dixit (Prakhar-Dixit)

Changes

Fixes issue #126197

The crash is caused because, during IR transformation, the vector-unrolling pass (using ExtractStridedSliceOp) attempts to slice an input vector of higher rank using a target vector of lower rank, which is not supported.

Specific example :

module {
  func.func @<!-- -->func1() {
    %cst_25 = arith.constant dense&lt;3.718400e+04&gt; : vector&lt;4x2x2xf16&gt;
    %cst_26 = arith.constant dense&lt;1.000000e+00&gt; : vector&lt;24x2x2xf32&gt;
    %47 = vector.fma %cst_26, %cst_26, %cst_26 : vector&lt;24x2x2xf32&gt;
    %818 = scf.execute_region -&gt; vector&lt;24x2x2xf32&gt; {
        scf.yield %47 : vector&lt;24x2x2xf32&gt;
      }
    %823 = vector.extract_strided_slice %cst_25 {offsets = [2], sizes = [1], strides = [1]} : vector&lt;4x2x2xf16&gt; to vector&lt;1x2x2xf16&gt;
    return
  }
}

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

1 Files Affected:

(modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+2)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index c1e3850f05c5e..82e473ef7e3b0 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -437,6 +437,8 @@ struct UnrollElementwisePattern : public RewritePattern {
     auto dstVecType = cast<VectorType>(op->getResult(0).getType());
     SmallVector<int64_t> originalSize =
         *cast<VectorUnrollOpInterface>(op).getShapeForUnroll();
+    if (originalSize.size() != targetShape->size())
+      return failure();
     Location loc = op->getLoc();
     // Prepare the result vector.
     Value result = rewriter.create<arith::ConstantOp>(

llvmbot · 2025-02-18T22:27:15Z

@llvm/pr-subscribers-mlir-vector

Author: Prakhar Dixit (Prakhar-Dixit)

Changes

Fixes issue #126197

The crash is caused because, during IR transformation, the vector-unrolling pass (using ExtractStridedSliceOp) attempts to slice an input vector of higher rank using a target vector of lower rank, which is not supported.

Specific example :

module {
  func.func @<!-- -->func1() {
    %cst_25 = arith.constant dense&lt;3.718400e+04&gt; : vector&lt;4x2x2xf16&gt;
    %cst_26 = arith.constant dense&lt;1.000000e+00&gt; : vector&lt;24x2x2xf32&gt;
    %47 = vector.fma %cst_26, %cst_26, %cst_26 : vector&lt;24x2x2xf32&gt;
    %818 = scf.execute_region -&gt; vector&lt;24x2x2xf32&gt; {
        scf.yield %47 : vector&lt;24x2x2xf32&gt;
      }
    %823 = vector.extract_strided_slice %cst_25 {offsets = [2], sizes = [1], strides = [1]} : vector&lt;4x2x2xf16&gt; to vector&lt;1x2x2xf16&gt;
    return
  }
}

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

1 Files Affected:

(modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+2)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index c1e3850f05c5e..82e473ef7e3b0 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -437,6 +437,8 @@ struct UnrollElementwisePattern : public RewritePattern {
     auto dstVecType = cast<VectorType>(op->getResult(0).getType());
     SmallVector<int64_t> originalSize =
         *cast<VectorUnrollOpInterface>(op).getShapeForUnroll();
+    if (originalSize.size() != targetShape->size())
+      return failure();
     Location loc = op->getLoc();
     // Prepare the result vector.
     Value result = rewriter.create<arith::ConstantOp>(

mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp

banach-space

Thank you! Please add a test 🙏🏻

Lewuathe · 2025-02-25T04:27:04Z

mlir/test/Dialect/Vector/invalid.mlir

@@ -768,6 +768,15 @@ func.func @extract_strided_slice(%arg0: vector<4x8x16xf32>) {

 // -----

+func.func @extract_strided_slice(%arg0: vector<3x2x2xf32>) {
+  // expected-error@+1 {{expected input vector rank to match target shape rank}}
+  %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]}:


Lewuathe

The change itself looks good to me. But the build failed. Let me trigger the CI again with the latest main branch.

Prakhar-Dixit · 2025-02-25T05:11:35Z

Thanks for the ping!

I've left one small request for an additional comment (please address). Otherwise LGTM, thanks!

Do you have commit access?

I've added the comment for better clarity!
I don't have commit access can you please merge it for me ?

Lewuathe · 2025-02-25T06:22:10Z

I'm going to merge this PR once the CI in buildkite passes.

Lewuathe · 2025-02-25T06:42:43Z

Hmm, the following tests we have added are failing in Windows platform build.

https://buildkite.com/llvm-project/github-pull-requests/builds/150531#01953bc4-51bf-4b66-a7e0-020a1f8f74bb

  MLIR :: Dialect/Vector/invalid.mlir
  MLIR :: Dialect/Vector/vector-unroll-options.mlir

Lewuathe · 2025-02-25T06:46:10Z

For the invalid.mlir case, the error message is probably incorrect.

# .---command stderr------------
# | within split at C:\ws\src\mlir\test\Dialect\Vector\invalid.mlir:769 offset :5:8: error: unexpected error: 'vector.extract_strided_slice' op expected result type to be 'vector<2x2x2xf32>'
# |   %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]}:
# |        ^
# | within split at C:\ws\src\mlir\test\Dialect\Vector\invalid.mlir:769 offset :4:6: error: expected error "expected input vector rank to match target shape rank" was not produced
# |   // expected-error@+1 {{expected input vector rank to match target shape rank}}
# |      ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# `-----------------------------

Lewuathe · 2025-02-25T06:50:25Z

In vector-unrole-option case, the signature of the variable can be changed. We need to use the wildcard (e.g. or assignment %[[ARG0:.*]]) instead of hard-coding %a.

# .---command stderr------------
# | C:\ws\src\mlir\test\Dialect\Vector\vector-unroll-options.mlir:198:11: error: CHECK: expected string not found in input
# | // CHECK: %0 = vector.fma %a, %a, %a : vector<3x2x2xf32>
# |           ^
# | <stdin>:104:35: note: scanning from here
# |  func.func @negative_vector_fma_3d(%arg0: vector<3x2x2xf32>) {
# |                                   ^
# |
# | Input file: <stdin>
# | Check file: C:\ws\src\mlir\test\Dialect\Vector\vector-unroll-options.mlir
# |
# | -dump-input=help explains the following input dump.
# |
# | Input was:
# | <<<<<<
# |            .
# |            .
# |            .
# |           99:  %17 = vector.extract_strided_slice %arg2 {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
# |          100:  %18 = vector.fma %15, %16, %17 : vector<2x2xf32>
# |          101:  %19 = vector.insert_strided_slice %18, %14 {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
# |          102:  return %19 : vector<4x4xf32>
# |          103:  }
# |          104:  func.func @negative_vector_fma_3d(%arg0: vector<3x2x2xf32>) {
# | check:198                                       X~~~~~~~~~~~~~~~~~~~~~~~~~~~~ error: no match found
# |          105:  return
# | check:198     ~~~~~~~~
# |          106:  }
# | check:198     ~~~
# |          107:  func.func @vector_multi_reduction(%arg0: vector<4x6xf32>, %arg1: vector<4xf32>) -> vector<4xf32> {
# | check:198     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# |          108:  %cst = arith.constant dense<0.000000e+00> : vector<4xf32>
# |          109:  %0 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x6xf32> to vector<2x2xf32>
# |            .
# |            .
# |            .
# | >>>>>>
# `-----------------------------

Lewuathe · 2025-02-25T07:02:07Z

@Prakhar-Dixit Could you try to fix the error of tests?

Prakhar-Dixit · 2025-02-25T07:06:50Z

@Prakhar-Dixit Could you try to fix the error of tests?

I was wondering—since we are not emitting any error messages when the conditions fail—whether a diagnostic test like the one in invalid.mlir is required ?

if (originalSize.size() != targetShape->size())
      return failure();

Prakhar-Dixit · 2025-02-25T07:19:06Z

For the vector-unroll-option case i guess this might work ?
// CHECK: %[[R0:.*]] = vector.fma %{{.+}}, %{{.+}}, %{{.+}} : vector<3x2x2xf32>

banach-space · 2025-02-25T10:56:21Z

I was wondering—since we are not emitting any error messages when the conditions fail—whether a diagnostic test like the one in invalid.mlir is required ?

From what I can tell, your test in invalid.mlir simply covers a case that we are not testing ATM. So it's something that's missing and it's good that you are adding it. Strictly speaking, it's not required for this PR (you have another, more targeted test), but the noise level is low and it's a positive addition. A dedicated note in the summary would be helpful.

Now, since you mentioned a diagnostic, could you replace return failure() with return rewriter.notifyMatchFailure("SOME MSG HERE");. Note, you won't be able to test it :)

For the vector-unroll-option case i guess this might work ?
// CHECK: %[[R0:.*]] = vector.fma %{{.+}}, %{{.+}}, %{{.+}} : vector<3x2x2xf32>

You can simplify further (since the actual arguments don't matter here):

// CHECK: %[[R0:.*]] = vector.fma %{{.+}} : vector<3x2x2xf32>

Thanks!

banach-space · 2025-02-25T11:07:38Z

mlir/test/Dialect/Vector/invalid.mlir

@@ -768,6 +768,15 @@ func.func @extract_strided_slice(%arg0: vector<4x8x16xf32>) {

 // -----

+func.func @extract_strided_slice(%arg0: vector<3x2x2xf32>) {
+  // expected-error@+1 {{expected input vector rank to match target shape rank}}


I looked at the error in CI:

https://buildkite.com/llvm-project/github-pull-requests/builds/150611#annotation-windows-x64-test-results

# .---command stderr------------ # | within split at C:\ws\src\mlir\test\Dialect\Vector\invalid.mlir:769 offset :5:8: error: unexpected error: 'vector.extract_strided_slice' op expected result type to be 'vector<2x2x2xf32>' # | %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]}: # | ^ # | within split at C:\ws\src\mlir\test\Dialect\Vector\invalid.mlir:769 offset :4:6: error: expected error "expected input vector rank to match target shape rank" was not produced # | // expected-error@+1 {{expected input vector rank to match target shape rank}} # | ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # `-----------------------------

I can reproduce this locally on non-Windows machine. @Prakhar-Dixit , does this test pass for you?

It failed on a non-Windows machine. Trying to figure it out.
😢

I looked at the error in CI:

https://buildkite.com/llvm-project/github-pull-requests/builds/150611#annotation-windows-x64-test-results

# .---command stderr------------ # | within split at C:\ws\src\mlir\test\Dialect\Vector\invalid.mlir:769 offset :5:8: error: unexpected error: 'vector.extract_strided_slice' op expected result type to be 'vector<2x2x2xf32>' # | %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]}: # | ^ # | within split at C:\ws\src\mlir\test\Dialect\Vector\invalid.mlir:769 offset :4:6: error: expected error "expected input vector rank to match target shape rank" was not produced # | // expected-error@+1 {{expected input vector rank to match target shape rank}} # | ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # `-----------------------------

I can reproduce this locally on non-Windows machine. @Prakhar-Dixit , does this test pass for you?

I have no clue how to handle this. I would be very grateful if you could help 🙏

Just to double-check - did it use to work? I am wondering whether it's some recent change to MLIR.

No, I guess not, because extracting a 2D slice from a 3D vector is not supported. As a result, an error is thrown indicating that the resulting vector type must be a 3D vector. This behavior is already verified by the following test:

func.func @extract_strided_slice(%arg0: vector<4x8x16xf32>) { // expected-error@+1 {{op expected result type to be 'vector<2x8x16xf32>'}} %1 = vector.extract_strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4x8x16xf32> to vector<3x1xf32> }

And i guess this test is not required since it's been taken care of.

OK, so you've made this error up:

// expected-error@+1 {{expected input vector rank to match target shape rank}}

? :) As in, it wasn't an actual error that you saw? That's what I was going by when reviewing 😅

And i guess this test is not required since it's been taken care of.

Indeed

Newbie struglles 🤧

…roll-options.mlir

Lewuathe

Now, the tests running in CI passed and the code seems cleaner. I'm going to merge this.

@Prakhar-Dixit Thank you for the fix!

@banach-space Thanks for your help!

github-actions · 2025-02-26T01:39:45Z

@Prakhar-Dixit 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!

Prakhar-Dixit requested review from banach-space, dcaballe, hanhanW and nicolasvasilache as code owners February 18, 2025 22:26

llvmbot added mlir:vectorops mlir mlir:vector labels Feb 18, 2025

Prakhar-Dixit changed the title ~~Add check to ensure input vector rank equals target shape rank~~ [mlir][vector] Add a check to ensure input vector rank equals target shape rank Feb 18, 2025

Prakhar-Dixit mentioned this pull request Feb 18, 2025

[mlir] Crash when using --test-vector-unrolling-patterns #126197

Closed

banach-space reviewed Feb 19, 2025

View reviewed changes

mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp Show resolved Hide resolved

banach-space reviewed Feb 19, 2025

View reviewed changes

Prakhar-Dixit requested review from ftynse, makslevental, stellaraccident, rengolin, matthias-springer, JDevlieghere, nikic and a team as code owners February 19, 2025 18:44

Prakhar-Dixit marked this pull request as draft February 19, 2025 18:54

ldionne removed the request for review from a team February 19, 2025 19:47

Prakhar-Dixit force-pushed the pd/bug-fix branch from a332d31 to bd5e0b1 Compare February 20, 2025 12:17

Prakhar-Dixit marked this pull request as ready for review February 20, 2025 12:18

Prakhar-Dixit requested a review from banach-space February 20, 2025 12:19

banach-space removed request for nicolasvasilache, ftynse and makslevental February 20, 2025 12:53

Lewuathe approved these changes Feb 25, 2025

View reviewed changes

Lewuathe reviewed Feb 25, 2025

View reviewed changes

Add a comment for better clarity in the negative test case

e0860bd

Prakhar-Dixit force-pushed the pd/bug-fix branch from d705753 to e0860bd Compare February 25, 2025 05:06

modify test

62961c2

Prakhar-Dixit force-pushed the pd/bug-fix branch from 865b584 to 62961c2 Compare February 25, 2025 10:10

modify test

77b4b65

banach-space reviewed Feb 25, 2025

View reviewed changes

Prakhar-Dixit and others added 3 commits February 25, 2025 19:11

Update return method in VectorUnroll.cpp and modify test in vector-un…

2535f83

…roll-options.mlir

modify test

8bfb3eb

Merge branch 'main' into pd/bug-fix

76e358d

Lewuathe approved these changes Feb 26, 2025

View reviewed changes

Lewuathe merged commit da37c76 into llvm:main Feb 26, 2025
11 checks passed

Prakhar-Dixit deleted the pd/bug-fix branch February 26, 2025 07:19

[mlir][vector] Add a check to ensure input vector rank equals target shape rank #127706

[mlir][vector] Add a check to ensure input vector rank equals target shape rank #127706

Uh oh!

Conversation

Prakhar-Dixit commented Feb 18, 2025

Uh oh!

github-actions bot commented Feb 18, 2025

Uh oh!

llvmbot commented Feb 18, 2025

Uh oh!

llvmbot commented Feb 18, 2025

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Lewuathe Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

Lewuathe left a comment

Choose a reason for hiding this comment

Uh oh!

Prakhar-Dixit commented Feb 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Lewuathe commented Feb 25, 2025

Uh oh!

Lewuathe commented Feb 25, 2025

Uh oh!

Lewuathe commented Feb 25, 2025

Uh oh!

Lewuathe commented Feb 25, 2025

Uh oh!

Lewuathe commented Feb 25, 2025

Uh oh!

Prakhar-Dixit commented Feb 25, 2025

Uh oh!

Prakhar-Dixit commented Feb 25, 2025

Uh oh!

banach-space commented Feb 25, 2025

Uh oh!

banach-space Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

Prakhar-Dixit Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

Prakhar-Dixit Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

banach-space Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

Prakhar-Dixit Feb 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

banach-space Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

Prakhar-Dixit Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

Lewuathe left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Feb 26, 2025

Uh oh!

Uh oh!

Prakhar-Dixit commented Feb 25, 2025 •

edited

Loading

Prakhar-Dixit Feb 25, 2025 •

edited

Loading