[mlir][affine] Modify assertion into a user visible diagnostic #136474

Prakhar-Dixit · 2025-04-20T06:45:35Z

Fixes #122227
The loop’s induction variable (%i) is used to compute two different indices via affine.apply.
And the Vectorization Assumption is Violated i.e, Each vectorized loop should contribute at most one non-invariant index.

Minimal example crashing :

#map = affine_map<(d0)[s0] -> (d0 mod s0)>
#map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)>

func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf32> {
  %c2 = arith.constant 2 : index
  %cst = arith.constant 1.0 : f32
  %alloc = memref.alloc() : memref<2x2xf32>

    affine.for %i = 0 to 4 {
      %row = affine.apply #map1(%i)[%c2]  
      %col = affine.apply #map(%i)[%c2]  
      affine.store %cst, %alloc[%row, %col] : memref<2x2xf32>
    }

    return %alloc : memref<2x2xf32>
  }

The single loop %i contributes two indices (%row and %col) to the 2D memref access.

The permutation map expects one index per vectorized loop dimension, but here one loop (%i) maps to two indices (dim=0 and dim=1).

The code detects this when trying to assign the second index (dim=1) to the same vector dimension (perm[0]).

llvmbot · 2025-04-20T06:46:08Z

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Prakhar Dixit (Prakhar-Dixit)

Changes

Fixes #122227
The loop’s induction variable (%i) is used to compute two different indices via affine.apply.
And the Vectorization Assumption is Violated i.e, Each vectorized loop should contribute at most one non-invariant index.

Minimal example crashing :

#map = affine_map&lt;(d0)[s0] -&gt; (d0 mod s0)&gt;
#map1 = affine_map&lt;(d0)[s0] -&gt; (d0 floordiv s0)&gt;

func.func @<!-- -->single_loop_unrolling_2D_access_pattern(%arg0: index) -&gt; memref&lt;2x2xf32&gt; {
  %c2 = arith.constant 2 : index
  %cst = arith.constant 1.0 : f32
  %alloc = memref.alloc() : memref&lt;2x2xf32&gt;

    affine.for %i = 0 to 4 {
      %row = affine.apply #map1(%i)[%c2]  
      %col = affine.apply #map(%i)[%c2]  
      affine.store %cst, %alloc[%row, %col] : memref&lt;2x2xf32&gt;
    }

    return %alloc : memref&lt;2x2xf32&gt;
  }

The single loop %i contributes two indices (%row and %col) to the 2D memref access.

The permutation map expects one index per vectorized loop dimension, but here one loop (%i) maps to two indices (dim=0 and dim=1).

The code detects this when trying to assign the second index (dim=1) to the same vector dimension (perm[0]).

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

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp (+8-2)
(modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir (+35)

diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index 024e8ccb901de..f15ed4b0d5f84 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -141,8 +141,14 @@ static AffineMap makePermutationMap(
     unsigned countInvariantIndices = 0;
     for (unsigned dim = 0; dim < numIndices; ++dim) {
       if (!invariants.count(indices[dim])) {
-        assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
-               "permutationMap already has an entry along dim");
+        if (perm[kvp.second] != getAffineConstantExpr(0, context)) {
+          auto loopOp = cast<affine::AffineForOp>(kvp.first);
+          loopOp->emitError(
+              "loop induction variable is used in multiple indices, which is "
+              "unsupported for vectorization. Consider using nested loops "
+              "instead of a single loop with affine.apply.");
+          return AffineMap();
+        }
         perm[kvp.second] = getAffineDimExpr(dim, context);
       } else {
         ++countInvariantIndices;
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
index 6c1a7c48c4cb1..c8c009f02212b 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
@@ -9,3 +9,38 @@ func.func @unparallel_loop_reduction_unsupported(%in: memref<256x512xf32>, %out:
  }
  return
 }
+
+// -----
+
+#map = affine_map<(d0)[s0] -> (d0 mod s0)>
+#map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+
+func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf32> {
+  %c2 = arith.constant 2 : index
+  %cst = arith.constant 1.0 : f32
+  %alloc = memref.alloc() : memref<2x2xf32>
+    
+    affine.for %i = 0 to 4 {
+      %row = affine.apply #map1(%i)[%c2]  
+      %col = affine.apply #map(%i)[%c2]  
+      affine.store %cst, %alloc[%row, %col] : memref<2x2xf32>
+    }
+    
+    return %alloc : memref<2x2xf32>
+  }
+
+// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0)[s0] -> (d0 mod s0)>
+
+// CHECK-LABEL:   func.func @single_loop_unrolling_2D_access_pattern(
+// CHECK-SAME:                            %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) -> memref<2x2xf32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant 2 : index
+// CHECK:           %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK:           %[[VAL_3:.*]] = memref.alloc() : memref<2x2xf32>
+// CHECK:           affine.for %[[VAL_4:.*]] = 0 to 4 {
+// CHECK:             %[[VAL_5:.*]] = affine.apply #[[$ATTR_0]](%[[VAL_4]]){{\[}}%[[VAL_1]]]
+// CHECK:             %[[VAL_6:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_4]]){{\[}}%[[VAL_1]]]
+// CHECK:             affine.store %[[VAL_2]], %[[VAL_3]]{{\[}}%[[VAL_5]], %[[VAL_6]]] : memref<2x2xf32>
+// CHECK:           }
+// CHECK:           return %[[VAL_3]] : memref<2x2xf32>
+// CHECK:         }
\ No newline at end of file

llvmbot · 2025-04-20T06:46:09Z

@llvm/pr-subscribers-mlir-affine

Author: Prakhar Dixit (Prakhar-Dixit)

Changes

Fixes #122227
The loop’s induction variable (%i) is used to compute two different indices via affine.apply.
And the Vectorization Assumption is Violated i.e, Each vectorized loop should contribute at most one non-invariant index.

Minimal example crashing :

#map = affine_map&lt;(d0)[s0] -&gt; (d0 mod s0)&gt;
#map1 = affine_map&lt;(d0)[s0] -&gt; (d0 floordiv s0)&gt;

func.func @<!-- -->single_loop_unrolling_2D_access_pattern(%arg0: index) -&gt; memref&lt;2x2xf32&gt; {
  %c2 = arith.constant 2 : index
  %cst = arith.constant 1.0 : f32
  %alloc = memref.alloc() : memref&lt;2x2xf32&gt;

    affine.for %i = 0 to 4 {
      %row = affine.apply #map1(%i)[%c2]  
      %col = affine.apply #map(%i)[%c2]  
      affine.store %cst, %alloc[%row, %col] : memref&lt;2x2xf32&gt;
    }

    return %alloc : memref&lt;2x2xf32&gt;
  }

The single loop %i contributes two indices (%row and %col) to the 2D memref access.

The permutation map expects one index per vectorized loop dimension, but here one loop (%i) maps to two indices (dim=0 and dim=1).

The code detects this when trying to assign the second index (dim=1) to the same vector dimension (perm[0]).

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

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp (+8-2)
(modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir (+35)

diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index 024e8ccb901de..f15ed4b0d5f84 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -141,8 +141,14 @@ static AffineMap makePermutationMap(
     unsigned countInvariantIndices = 0;
     for (unsigned dim = 0; dim < numIndices; ++dim) {
       if (!invariants.count(indices[dim])) {
-        assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
-               "permutationMap already has an entry along dim");
+        if (perm[kvp.second] != getAffineConstantExpr(0, context)) {
+          auto loopOp = cast<affine::AffineForOp>(kvp.first);
+          loopOp->emitError(
+              "loop induction variable is used in multiple indices, which is "
+              "unsupported for vectorization. Consider using nested loops "
+              "instead of a single loop with affine.apply.");
+          return AffineMap();
+        }
         perm[kvp.second] = getAffineDimExpr(dim, context);
       } else {
         ++countInvariantIndices;
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
index 6c1a7c48c4cb1..c8c009f02212b 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
@@ -9,3 +9,38 @@ func.func @unparallel_loop_reduction_unsupported(%in: memref<256x512xf32>, %out:
  }
  return
 }
+
+// -----
+
+#map = affine_map<(d0)[s0] -> (d0 mod s0)>
+#map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+
+func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf32> {
+  %c2 = arith.constant 2 : index
+  %cst = arith.constant 1.0 : f32
+  %alloc = memref.alloc() : memref<2x2xf32>
+    
+    affine.for %i = 0 to 4 {
+      %row = affine.apply #map1(%i)[%c2]  
+      %col = affine.apply #map(%i)[%c2]  
+      affine.store %cst, %alloc[%row, %col] : memref<2x2xf32>
+    }
+    
+    return %alloc : memref<2x2xf32>
+  }
+
+// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0)[s0] -> (d0 mod s0)>
+
+// CHECK-LABEL:   func.func @single_loop_unrolling_2D_access_pattern(
+// CHECK-SAME:                            %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) -> memref<2x2xf32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant 2 : index
+// CHECK:           %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK:           %[[VAL_3:.*]] = memref.alloc() : memref<2x2xf32>
+// CHECK:           affine.for %[[VAL_4:.*]] = 0 to 4 {
+// CHECK:             %[[VAL_5:.*]] = affine.apply #[[$ATTR_0]](%[[VAL_4]]){{\[}}%[[VAL_1]]]
+// CHECK:             %[[VAL_6:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_4]]){{\[}}%[[VAL_1]]]
+// CHECK:             affine.store %[[VAL_2]], %[[VAL_3]]{{\[}}%[[VAL_5]], %[[VAL_6]]] : memref<2x2xf32>
+// CHECK:           }
+// CHECK:           return %[[VAL_3]] : memref<2x2xf32>
+// CHECK:         }
\ No newline at end of file

Prakhar-Dixit · 2025-04-20T06:53:47Z

@ftynse

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

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir

Prakhar-Dixit · 2025-04-24T02:50:33Z

Just a gentle reminder!
@joker-eph

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir

…136474) Fixes llvm#122227 The loop’s induction variable (%i) is used to compute two different indices via affine.apply. And the Vectorization Assumption is Violated i.e, Each vectorized loop should contribute at most one non-invariant index. **Minimal example crashing :** ``` #map = affine_map<(d0)[s0] -> (d0 mod s0)> #map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)> func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf32> { %c2 = arith.constant 2 : index %cst = arith.constant 1.0 : f32 %alloc = memref.alloc() : memref<2x2xf32> affine.for %i = 0 to 4 { %row = affine.apply #map1(%i)[%c2] %col = affine.apply #map(%i)[%c2] affine.store %cst, %alloc[%row, %col] : memref<2x2xf32> } return %alloc : memref<2x2xf32> } ``` The single loop %i contributes two indices (%row and %col) to the 2D memref access. The permutation map expects one index per vectorized loop dimension, but here one loop (%i) maps to two indices (dim=0 and dim=1). The code detects this when trying to assign the second index (dim=1) to the same vector dimension (perm[0]).

[mlir][affine] Modify assertion into a user visible diagnostic

84157c0

Prakhar-Dixit requested review from banach-space, dcaballe, nicolasvasilache and Groverkss as code owners April 20, 2025 06:45

llvmbot added mlir:affine mlir:vectorops mlir mlir:vector labels Apr 20, 2025

joker-eph reviewed Apr 20, 2025

View reviewed changes

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp Outdated Show resolved Hide resolved

Prakhar-Dixit added 2 commits April 22, 2025 13:09

Add check to ensure pass fails gracefully

8112d1a

typo fix

24acbf9

Prakhar-Dixit requested a review from joker-eph April 22, 2025 08:19

joker-eph reviewed Apr 22, 2025

View reviewed changes

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp Outdated Show resolved Hide resolved

joker-eph reviewed Apr 22, 2025

View reviewed changes

mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir Show resolved Hide resolved

mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir Outdated Show resolved Hide resolved

Create a helper function to avoid code duplication

99f6885

Prakhar-Dixit requested a review from joker-eph April 22, 2025 15:32

ftynse approved these changes Apr 29, 2025

View reviewed changes

Prakhar-Dixit added 2 commits April 30, 2025 11:10

modify code with best practices

2052b9d

modify

827b352

Prakhar-Dixit requested a review from ftynse April 30, 2025 06:15

ftynse approved these changes Apr 30, 2025

View reviewed changes

ftynse merged commit c26db58 into llvm:main Apr 30, 2025
11 checks passed

Prakhar-Dixit deleted the pd/affine-122227 branch May 1, 2025 02:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir][affine] Modify assertion into a user visible diagnostic #136474

[mlir][affine] Modify assertion into a user visible diagnostic #136474

Prakhar-Dixit commented Apr 20, 2025

llvmbot commented Apr 20, 2025 •

edited

Loading

llvmbot commented Apr 20, 2025

Prakhar-Dixit commented Apr 20, 2025

Prakhar-Dixit commented Apr 24, 2025 •

edited

Loading

[mlir][affine] Modify assertion into a user visible diagnostic #136474

[mlir][affine] Modify assertion into a user visible diagnostic #136474

Conversation

Prakhar-Dixit commented Apr 20, 2025

llvmbot commented Apr 20, 2025 • edited Loading

llvmbot commented Apr 20, 2025

Prakhar-Dixit commented Apr 20, 2025

Prakhar-Dixit commented Apr 24, 2025 • edited Loading

llvmbot commented Apr 20, 2025 •

edited

Loading

Prakhar-Dixit commented Apr 24, 2025 •

edited

Loading