shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[mlir][affine] Bail out when store permutation map has broadcast dimensions (#184618)

Browse Source 
When the vectorized loop is an outer loop and the store index uses an
inner loop's IV (which is invariant w.r.t. the outer loop),
makePermutationMap
produces a broadcast map (e.g. (d0) -> (0)). A vector.transfer_write
with
broadcast dimensions is invalid, causing a verifier error.

Fix: check for AffineConstantExpr results in the store permutation map
and
bail out of vectorization, preserving the scalar loop nest.

Fixes #131135

Assisted-by: Claude Code
This commit is contained in:
Mehdi Amini 2026-03-11 16:59:38 +01:00 committed by GitHub
parent b1bc05d5b9
commit 5842a0f76c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 34 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
View File

@ -1312,6 +1312,15 @@ static Operation *vectorizeAffineStore(AffineStoreOp storeOp,
LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
LLVM_DEBUG(permutationMap.print(dbgs()));
// A transfer_write with a broadcast dimension (constant expr in the
// permutation map) is invalid. Bail out to avoid producing invalid IR.
if (llvm::any_of(permutationMap.getResults(),
llvm::IsaPred<AffineConstantExpr>)) {
LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ store permutation map has "
"broadcast dims, bailing out\n");
return nullptr;
}
auto transfer = vector::TransferWriteOp::create(
state.builder, storeOp.getLoc(), vectorValue, storeOp.getMemRef(),
indices, permutationMap);

25
mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
View File

@ -43,6 +43,31 @@ func.func @iv_mapped_to_multiple_indices_unsupported(%arg0: index) -> memref<2x2
// -----
// Regression test: when the store's permutation map has a broadcast dimension
// (because the index is invariant w.r.t. the vectorized loop), vectorization
// must bail out gracefully instead of emitting an invalid transfer_write.
#map_id = affine_map<(d0) -> (d0)>
#map_id_p1 = affine_map<(d0) -> (d0 + 1)>
// CHECK-LABEL: func.func @store_broadcast_perm_map_unsupported
// CHECK: affine.for
// CHECK: affine.for
// CHECK: affine.load
// CHECK: affine.store
// CHECK-NOT: vector.transfer_write
func.func @store_broadcast_perm_map_unsupported(%arg0: memref<4x4xf32>, %arg1: memref<4xf32>) {
affine.for %i = 0 to 4 {
affine.for %j = #map_id(%i) to #map_id_p1(%i) {
%0 = affine.load %arg0[%j, %j] : memref<4x4xf32>
affine.store %0, %arg1[%j] : memref<4xf32>
}
}
return
}
// -----
// Regression test for https://github.com/llvm/llvm-project/issues/128334
// The vectorizer test utility used to crash when a reduction loop with a
// dynamic upper bound was vectorized via 'vectorizeAffineLoopNest', because