[Mlir-commits] [mlir] [mlir][affine] Modify assertion into a user visible diagnostic (PR #136474)
Prakhar Dixit
llvmlistbot at llvm.org
Tue Apr 22 00:40:14 PDT 2025
https://github.com/Prakhar-Dixit updated https://github.com/llvm/llvm-project/pull/136474
>From 84157c03453f2c23faa086541a0cc3421ae7711a Mon Sep 17 00:00:00 2001
From: Prakhar Dixit <dixitprakhar11 at gmail.com>
Date: Sun, 20 Apr 2025 12:01:42 +0530
Subject: [PATCH 1/2] [mlir][affine] Modify assertion into a user visible
diagnostic
---
mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 10 ++++--
.../SuperVectorize/vectorize_unsupported.mlir | 35 +++++++++++++++++++
2 files changed, 43 insertions(+), 2 deletions(-)
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
>From 8112d1a04b248b4bd6098ad3d9e70ef06232868d Mon Sep 17 00:00:00 2001
From: Prakhar Dixit <dixitprakhar11 at gmail.com>
Date: Tue, 22 Apr 2025 13:09:52 +0530
Subject: [PATCH 2/2] Add check to ensure pass fails gracefully
---
.../Affine/Transforms/SuperVectorize.cpp | 30 +++++++++++++++
mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 10 +----
.../SuperVectorize/vectorize_unsupported.mlir | 37 ++++++++++++++++++-
3 files changed, 67 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index eaaafaf68767e..cf1341084021b 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1217,6 +1217,21 @@ static Operation *vectorizeAffineLoad(AffineLoadOp loadOp,
indices.append(mapOperands.begin(), mapOperands.end());
}
+ for (auto &kvp : state.vecLoopToVecDim) {
+ AffineForOp forOp = cast<AffineForOp>(kvp.first);
+ auto invariants =
+ affine::getInvariantAccesses(forOp.getInductionVar(), indices);
+ unsigned nonInvariant = 0;
+ for (Value idx : indices)
+ if (!invariants.count(idx))
+ ++nonInvariant;
+ if (nonInvariant > 1) {
+ LLVM_DEBUG(dbgs() << "\n[early-vect] Bail out: loop IV "
+ << forOp.getInductionVar() << " drives " << nonInvariant
+ << " indices (must be ≤1)\n");
+ return nullptr;
+ }
+ }
// Compute permutation map using the information of new vector loops.
auto permutationMap = makePermutationMap(state.builder.getInsertionBlock(),
indices, state.vecLoopToVecDim);
@@ -1262,6 +1277,21 @@ static Operation *vectorizeAffineStore(AffineStoreOp storeOp,
else
indices.append(mapOperands.begin(), mapOperands.end());
+ for (auto &kvp : state.vecLoopToVecDim) {
+ AffineForOp forOp = cast<AffineForOp>(kvp.first);
+ auto invariants =
+ affine::getInvariantAccesses(forOp.getInductionVar(), indices);
+ unsigned nonInvariant = 0;
+ for (Value idx : indices)
+ if (!invariants.count(idx))
+ ++nonInvariant;
+ if (nonInvariant > 1) {
+ LLVM_DEBUG(dbgs() << "\n[early-vect] Bail out: loop IV "
+ << forOp.getInductionVar() << " drives " << nonInvariant
+ << " indices (must be ≤1)\n");
+ return nullptr;
+ }
+ }
// Compute permutation map using the information of new vector loops.
auto permutationMap = makePermutationMap(state.builder.getInsertionBlock(),
indices, state.vecLoopToVecDim);
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index f15ed4b0d5f84..024e8ccb901de 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -141,14 +141,8 @@ static AffineMap makePermutationMap(
unsigned countInvariantIndices = 0;
for (unsigned dim = 0; dim < numIndices; ++dim) {
if (!invariants.count(indices[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();
- }
+ assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
+ "permutationMap already has an entry along dim");
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 c8c009f02212b..51b7020f30943 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
@@ -15,7 +15,7 @@ func.func @unparallel_loop_reduction_unsupported(%in: memref<256x512xf32>, %out:
#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> {
+func.func @single_loop_unrolling_2D_access_pattern_storeOp(%arg0: index) -> memref<2x2xf32> {
%c2 = arith.constant 2 : index
%cst = arith.constant 1.0 : f32
%alloc = memref.alloc() : memref<2x2xf32>
@@ -33,7 +33,7 @@ func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf
// 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-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>
@@ -43,4 +43,37 @@ func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf
// CHECK: affine.store %[[VAL_2]], %[[VAL_3]]{{\[}}%[[VAL_5]], %[[VAL_6]]] : memref<2x2xf32>
// CHECK: }
// CHECK: return %[[VAL_3]] : memref<2x2xf32>
+// CHECK: }
+
+// -----
+
+#map = affine_map<(d0)[s0] -> (d0 mod s0)>
+#map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+
+func.func @single_loop_unrolling_2D_access_pattern_loadOp(%arg0: index) -> memref<2x2xf32> {
+ %c2 = arith.constant 2 : index
+ %alloc = memref.alloc() : memref<2x2xf32>
+
+ affine.for %i = 0 to 4 {
+ %row = affine.apply #map1(%i)[%c2]
+ %col = affine.apply #map(%i)[%c2]
+ %val = affine.load %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:.*]] = memref.alloc() : memref<2x2xf32>
+// CHECK: affine.for %[[VAL_3:.*]] = 0 to 4 {
+// CHECK: %[[VAL_4:.*]] = affine.apply #[[$ATTR_0]](%[[VAL_3]]){{\[}}%[[VAL_1]]]
+// CHECK: %[[VAL_5:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_3]]){{\[}}%[[VAL_1]]]
+// CHECK: %[[VAL_6:.*]] = affine.load %[[VAL_2]]{{\[}}%[[VAL_4]], %[[VAL_5]]] : memref<2x2xf32>
+// CHECK: }
+// CHECK: return %[[VAL_2]] : memref<2x2xf32>
// CHECK: }
\ No newline at end of file
More information about the Mlir-commits
mailing list