[MLIR] Minor fixes to FoldTransposeBroadcast rewrite (NFC) #140083

momchil-velikov · 2025-05-15T15:32:25Z

This patch contains two minor changes, which I believe were the original author's intent.

when folding transpose(broadcast(x)) emit broadcast(x) instead of broadcast(broadcast(x)). The latter causes intermittent verifier failures, e.g.

mlir-asm-printer: 'func.func' failed to verify and will be printed in generic form
"func.func"() <{function_type = (vector<4x1x1x7xi8>) -> vector<3x2x4x5x6x7xi8>, sym_name = "broadcast_transpose_mixed_example"}> ({
^bb0(%arg0: vector<4x1x1x7xi8>):
  %0 = "vector.broadcast"(%arg0) : (vector<4x1x1x7xi8>) -> vector<2x3x4x5x6x7xi8>
  %1 = "vector.broadcast"(%0) : (vector<2x3x4x5x6x7xi8>) -> vector<3x2x4x5x6x7xi8>
  "func.return"(%1) : (vector<3x2x4x5x6x7xi8>) -> ()
}) : () -> ()

when checking permutation groups the variable low was set just once to zero, thus checking was quadratic. It looks the intent was for low to track the beginning of each dimension groups. (Nevertheless the check was correct).

This patch contains to minor changes, which I believe were the original author's intent. * when folding `transpose(broadcast(x))` emit `broadcast(x)` instead of `broadcast(broadca(x))`. The later causes intermittent verifier failures, e.g. ``` mlir-asm-printer: 'func.func' failed to verify and will be printed in generic form "func.func"() <{function_type = (vector<4x1x1x7xi8>) -> vector<3x2x4x5x6x7xi8>, sym_name = "broadcast_transpose_mixed_example"}> ({ ^bb0(%arg0: vector<4x1x1x7xi8>): %0 = "vector.broadcast"(%arg0) : (vector<4x1x1x7xi8>) -> vector<2x3x4x5x6x7xi8> %1 = "vector.broadcast"(%0) : (vector<2x3x4x5x6x7xi8>) -> vector<3x2x4x5x6x7xi8> "func.return"(%1) : (vector<3x2x4x5x6x7xi8>) -> () }) : () -> () ``` * when checking permutation groups the variable `low` was set just once to zero, thus checking was quadratic. It looks the intent was for `low` to track the beginning of each dimension groups. (Nevertheless the check was correct).

llvmbot · 2025-05-15T15:32:58Z

@llvm/pr-subscribers-mlir-vector

Author: Momchil Velikov (momchil-velikov)

Changes

This patch contains two minor changes, which I believe were the original author's intent.

when folding transpose(broadcast(x)) emit broadcast(x) instead of broadcast(broadcast(x)). The latter causes intermittent verifier failures, e.g.

mlir-asm-printer: 'func.func' failed to verify and will be printed in generic form
"func.func"() &lt;{function_type = (vector&lt;4x1x1x7xi8&gt;) -&gt; vector&lt;3x2x4x5x6x7xi8&gt;, sym_name = "broadcast_transpose_mixed_example"}&gt; ({
^bb0(%arg0: vector&lt;4x1x1x7xi8&gt;):
  %0 = "vector.broadcast"(%arg0) : (vector&lt;4x1x1x7xi8&gt;) -&gt; vector&lt;2x3x4x5x6x7xi8&gt;
  %1 = "vector.broadcast"(%0) : (vector&lt;2x3x4x5x6x7xi8&gt;) -&gt; vector&lt;3x2x4x5x6x7xi8&gt;
  "func.return"(%1) : (vector&lt;3x2x4x5x6x7xi8&gt;) -&gt; ()
}) : () -&gt; ()

when checking permutation groups the variable low was set just once to zero, thus checking was quadratic. It looks the intent was for low to track the beginning of each dimension groups. (Nevertheless the check was correct).

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

1 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+3-2)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 79bf87ccd34af..7ae43b64a5deb 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6201,7 +6201,7 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
     bool inputIsScalar = !inputType;
     if (inputIsScalar) {
       rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
-                                                       transpose.getVector());
+                                                       broadcast.getSource());
       return success();
     }
 
@@ -6227,6 +6227,7 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
                 transpose, "permutation not local to group");
           }
         }
+        low = high;
       }
     }
 
@@ -6241,7 +6242,7 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
            "not broadcastable directly to transpose output");
 
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
-                                                     transpose.getVector());
+                                                     broadcast.getSource());
 
     return success();
   }

llvmbot · 2025-05-15T15:32:59Z

@llvm/pr-subscribers-mlir

Author: Momchil Velikov (momchil-velikov)

Changes

This patch contains two minor changes, which I believe were the original author's intent.

when folding transpose(broadcast(x)) emit broadcast(x) instead of broadcast(broadcast(x)). The latter causes intermittent verifier failures, e.g.

mlir-asm-printer: 'func.func' failed to verify and will be printed in generic form
"func.func"() &lt;{function_type = (vector&lt;4x1x1x7xi8&gt;) -&gt; vector&lt;3x2x4x5x6x7xi8&gt;, sym_name = "broadcast_transpose_mixed_example"}&gt; ({
^bb0(%arg0: vector&lt;4x1x1x7xi8&gt;):
  %0 = "vector.broadcast"(%arg0) : (vector&lt;4x1x1x7xi8&gt;) -&gt; vector&lt;2x3x4x5x6x7xi8&gt;
  %1 = "vector.broadcast"(%0) : (vector&lt;2x3x4x5x6x7xi8&gt;) -&gt; vector&lt;3x2x4x5x6x7xi8&gt;
  "func.return"(%1) : (vector&lt;3x2x4x5x6x7xi8&gt;) -&gt; ()
}) : () -&gt; ()

when checking permutation groups the variable low was set just once to zero, thus checking was quadratic. It looks the intent was for low to track the beginning of each dimension groups. (Nevertheless the check was correct).

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

1 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+3-2)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 79bf87ccd34af..7ae43b64a5deb 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6201,7 +6201,7 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
     bool inputIsScalar = !inputType;
     if (inputIsScalar) {
       rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
-                                                       transpose.getVector());
+                                                       broadcast.getSource());
       return success();
     }
 
@@ -6227,6 +6227,7 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
                 transpose, "permutation not local to group");
           }
         }
+        low = high;
       }
     }
 
@@ -6241,7 +6242,7 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
            "not broadcastable directly to transpose output");
 
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
-                                                     transpose.getVector());
+                                                     broadcast.getSource());
 
     return success();
   }

joker-eph · 2025-05-15T15:37:08Z

Can you provide a test that triggered the verification error?

momchil-velikov · 2025-05-15T15:48:03Z

Can you provide a test that triggered the verification error?

Yes, it from the test from the original PR #135096
https://github.com/llvm/llvm-project/blob/540cf25a6df56fa1810a7411477dca9896aeed20/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir

newling · 2025-05-15T16:05:39Z

Thanks @momchil-velikov I think you are correct. What I don't understand is why there was no verification error, every time. For example in

llvm-project/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir

Line 25 in 540cf25

func.func @transpose_scalar_broadcast2(%value: f32) -> vector<1x8xf32> {

It was converting

  %bcast = vector.broadcast %value : f32 to vector<8x1xf32>
  %t = vector.transpose %bcast, [1, 0] : vector<8x1xf32> to vector<1x8xf32>

to

 %bcast = vector.broadcast %value : f32 to vector<8x1xf32>
  %t = vector.broadcast %bcast  : vector<8x1xf32> to vector<1x8xf32>

but that %t is invalid. Unless the folder runs before verification? But yeah, it should have going directly to

  %t = vector.broadcast %value : f32 to vector<1x8xf32>

newling · 2025-05-15T16:18:37Z

when checking permutation groups the variable low was set just once to zero, thus checking was quadratic. It looks the intent was for low to track the beginning of each dimension groups. (Nevertheless the check was correct).

Yes, good catch thanks.

newling

Thanks for the fix. I find it worrying that this didn't create a verification error every time (as per comment above). I'll spend some time trying to figure out why not (any suggestions @joker-eph ?) maybe if you can too @momchil-velikov that'll be useful (I'm not sure I'll manage to reproduce any verification error with my setup).

nit: technically not (NFC) ?.

newling · 2025-05-15T16:38:23Z

mlir-asm-printer: 'func.func' failed to verify and will be printed in generic form
"func.func"() <{function_type = (vector<4x1x1x7xi8>) -> vector<3x2x4x5x6x7xi8>, sym_name = "broadcast_transpose_mixed_example"}> ({
^bb0(%arg0: vector<4x1x1x7xi8>):
%0 = "vector.broadcast"(%arg0) : (vector<4x1x1x7xi8>) -> vector<2x3x4x5x6x7xi8>
%1 = "vector.broadcast"(%0) : (vector<2x3x4x5x6x7xi8>) -> vector<3x2x4x5x6x7xi8>
"func.return"(%1) : (vector<3x2x4x5x6x7xi8>) -> ()
}) : () -> ()

I can see this everytime in when running mlir-opt with the --debug flag.

newling · 2025-05-15T16:58:12Z

The latter causes intermittent verifier failures, e.g.

I guess, what do you mean by intermittent? This fixes a bug (thanks again!) I'm just not sure it would cause anything observable like intermittent compilation failure. My rough understanding is that the verifier is only called with --debug because the printer verifies, but as the 2 broadcasts fold, by the time the pattern rewrite is done the IR is valid again.

momchil-velikov · 2025-05-15T17:05:25Z

The latter causes intermittent verifier failures, e.g.

I guess, what do you mean by intermittent?

I guess bad wording. I saw it with --debug happening at one point and disappearing before the end of the compilation.

newling · 2025-05-15T17:32:46Z

The latter causes intermittent verifier failures, e.g.

I guess, what do you mean by intermittent?

I guess bad wording. I saw it with --debug happening at one point and disappearing before the end of the compilation.

Ok, good. Thanks!

banach-space · 2025-05-15T19:31:49Z

Can you try building mlir-opt with -DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=On? With expensive checks on, mlir-opt should fail.

newling · 2025-05-16T00:57:07Z

Can you try building mlir-opt with -DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=On? With expensive checks on, mlir-opt should fail.

I've just tried this and yes, it fails. That's a useful flag, thanks for point this out @banach-space! Should there be a build in CI with this on?

banach-space · 2025-05-16T16:58:17Z

Can you try building mlir-opt with -DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=On? With expensive checks on, mlir-opt should fail.

I've just tried this and yes, it fails. That's a useful flag, thanks for point this out @banach-space! Should there be a build in CI with this on?

All the builders are defined here:

https://github.com/llvm/llvm-zorg/tree/main

Unfortunately, searching for MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS gives no hits. I'm hesitant to enable this in one of the existing Arm builders, since I worry it might be too fragile and break frequently.

I'll look into whether we can add a dedicated builder for this - but please don’t hold your breath; these things can take time.

dcaballe · 2025-05-16T21:19:20Z

Let's see if I follow... we initially generate:

 %bcast = vector.broadcast %value : f32 to vector<8x1xf32>
  %t = vector.broadcast %bcast  : vector<8x1xf32> to vector<1x8xf32>

but the CHECK rules don't fail:

//  CHECK-SAME: (%[[ARG:.+]]: f32)
//       CHECK:   %[[V:.+]] = vector.broadcast %[[ARG]] : f32 to vector<1x8xf32>
//       CHECK:   return %[[V]] : vector<1x8xf32>

because there is another folding that is turning the two invalid vector.broadcast ops into a single but valid one, and the intermediate state is not verified without the expensive checks?

Yeah, having a buildbot covering this would be awesome.

newling · 2025-05-19T14:51:06Z

Can you try building mlir-opt with -DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=On? With expensive checks on, mlir-opt should fail.

I've just tried this and yes, it fails. That's a useful flag, thanks for point this out @banach-space! Should there be a build in CI with this on?

All the builders are defined here:

https://github.com/llvm/llvm-zorg/tree/main

Unfortunately, searching for MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS gives no hits. I'm hesitant to enable this in one of the existing Arm builders, since I worry it might be too fragile and break frequently.

I'll look into whether we can add a dedicated builder for this - but please don’t hold your breath; these things can take time.

Thanks @banach-space that might be useful, do let me know if I can help

momchil-velikov requested a review from newling May 15, 2025 15:32

momchil-velikov requested review from banach-space, dcaballe, nicolasvasilache and Groverkss as code owners May 15, 2025 15:32

llvmbot added mlir:vectorops mlir mlir:vector labels May 15, 2025

newling approved these changes May 15, 2025

View reviewed changes

dcaballe approved these changes May 16, 2025

View reviewed changes

momchil-velikov merged commit 38d2306 into llvm:main May 19, 2025
15 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR] Minor fixes to FoldTransposeBroadcast rewrite (NFC) #140083

[MLIR] Minor fixes to FoldTransposeBroadcast rewrite (NFC) #140083

momchil-velikov commented May 15, 2025 •

edited

Loading

llvmbot commented May 15, 2025

llvmbot commented May 15, 2025

joker-eph commented May 15, 2025

momchil-velikov commented May 15, 2025

newling commented May 15, 2025 •

edited

Loading

newling commented May 15, 2025

newling left a comment •

edited

Loading

newling commented May 15, 2025

newling commented May 15, 2025

momchil-velikov commented May 15, 2025

newling commented May 15, 2025

banach-space commented May 15, 2025

newling commented May 16, 2025

banach-space commented May 16, 2025

dcaballe commented May 16, 2025

newling commented May 19, 2025

[MLIR] Minor fixes to FoldTransposeBroadcast rewrite (NFC) #140083

[MLIR] Minor fixes to FoldTransposeBroadcast rewrite (NFC) #140083

Conversation

momchil-velikov commented May 15, 2025 • edited Loading

llvmbot commented May 15, 2025

llvmbot commented May 15, 2025

joker-eph commented May 15, 2025

momchil-velikov commented May 15, 2025

newling commented May 15, 2025 • edited Loading

newling commented May 15, 2025

newling left a comment • edited Loading

Choose a reason for hiding this comment

newling commented May 15, 2025

newling commented May 15, 2025

momchil-velikov commented May 15, 2025

newling commented May 15, 2025

banach-space commented May 15, 2025

newling commented May 16, 2025

banach-space commented May 16, 2025

dcaballe commented May 16, 2025

newling commented May 19, 2025

momchil-velikov commented May 15, 2025 •

edited

Loading

newling commented May 15, 2025 •

edited

Loading

newling left a comment •

edited

Loading