Skip to content

Commit b1cb86d

Browse files
Prakhar-DixitGeorgeARM
authored andcommitted
[mlir][affine] Modify assertion into a user visible diagnostic (llvm#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]).
1 parent a5cb4fa commit b1cb86d

File tree

2 files changed

+63
-0
lines changed

2 files changed

+63
-0
lines changed

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1185,6 +1185,32 @@ static Value vectorizeOperand(Value operand, VectorizationState &state) {
11851185
return nullptr;
11861186
}
11871187

1188+
/// Returns true if any vectorized loop IV drives more than one index.
1189+
static bool isIVMappedToMultipleIndices(
1190+
ArrayRef<Value> indices,
1191+
const DenseMap<Operation *, unsigned> &loopToVectorDim) {
1192+
for (auto &kvp : loopToVectorDim) {
1193+
AffineForOp forOp = cast<AffineForOp>(kvp.first);
1194+
// Find which indices are invariant w.r.t. this loop IV.
1195+
llvm::DenseSet<Value> invariants =
1196+
affine::getInvariantAccesses(forOp.getInductionVar(), indices);
1197+
// Count how many vary (i.e. are not invariant).
1198+
unsigned nonInvariant = 0;
1199+
for (Value idx : indices) {
1200+
if (invariants.count(idx))
1201+
continue;
1202+
1203+
if (++nonInvariant > 1) {
1204+
LLVM_DEBUG(dbgs() << "[early‑vect] Bail out: IV "
1205+
<< forOp.getInductionVar() << " drives "
1206+
<< nonInvariant << " indices\n");
1207+
return true;
1208+
}
1209+
}
1210+
}
1211+
return false;
1212+
}
1213+
11881214
/// Vectorizes an affine load with the vectorization strategy in 'state' by
11891215
/// generating a 'vector.transfer_read' op with the proper permutation map
11901216
/// inferred from the indices of the load. The new 'vector.transfer_read' is
@@ -1217,6 +1243,9 @@ static Operation *vectorizeAffineLoad(AffineLoadOp loadOp,
12171243
indices.append(mapOperands.begin(), mapOperands.end());
12181244
}
12191245

1246+
if (isIVMappedToMultipleIndices(indices, state.vecLoopToVecDim))
1247+
return nullptr;
1248+
12201249
// Compute permutation map using the information of new vector loops.
12211250
auto permutationMap = makePermutationMap(state.builder.getInsertionBlock(),
12221251
indices, state.vecLoopToVecDim);
@@ -1262,6 +1291,9 @@ static Operation *vectorizeAffineStore(AffineStoreOp storeOp,
12621291
else
12631292
indices.append(mapOperands.begin(), mapOperands.end());
12641293

1294+
if (isIVMappedToMultipleIndices(indices, state.vecLoopToVecDim))
1295+
return nullptr;
1296+
12651297
// Compute permutation map using the information of new vector loops.
12661298
auto permutationMap = makePermutationMap(state.builder.getInsertionBlock(),
12671299
indices, state.vecLoopToVecDim);

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,34 @@ func.func @unparallel_loop_reduction_unsupported(%in: memref<256x512xf32>, %out:
99
}
1010
return
1111
}
12+
13+
// -----
14+
15+
#map = affine_map<(d0)[s0] -> (d0 mod s0)>
16+
#map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)>
17+
18+
func.func @iv_mapped_to_multiple_indices_unsupported(%arg0: index) -> memref<2x2xf32> {
19+
%c2 = arith.constant 2 : index
20+
%cst = arith.constant 1.0 : f32
21+
%alloc = memref.alloc() : memref<2x2xf32>
22+
23+
affine.for %i = 0 to 4 {
24+
%row = affine.apply #map1(%i)[%c2]
25+
%col = affine.apply #map(%i)[%c2]
26+
affine.store %cst, %alloc[%row, %col] : memref<2x2xf32>
27+
}
28+
29+
return %alloc : memref<2x2xf32>
30+
}
31+
32+
// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
33+
// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0)[s0] -> (d0 mod s0)>
34+
35+
// CHECK-LABEL: func.func @iv_mapped_to_multiple_indices_unsupported(
36+
// CHECK-SAME: %[[VAL_0:.*]]: index) -> memref<2x2xf32> {
37+
// CHECK: %[[VAL_1:.*]] = arith.constant 2 : index
38+
// CHECK: affine.for %[[VAL_4:.*]] = 0 to 4 {
39+
// CHECK: %[[VAL_5:.*]] = affine.apply #[[$ATTR_0]](%[[VAL_4]]){{\[}}%[[VAL_1]]]
40+
// CHECK: %[[VAL_6:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_4]]){{\[}}%[[VAL_1]]]
41+
// CHECK: }
42+
// CHECK: }

0 commit comments

Comments
 (0)