[MLIR][Affine] Fix addInductionVarOrTerminalSymbol

Update affine analysis method `addInductionVarOrTerminalSymbol` for
affine.parallel IV.

Fixes https://github.com/llvm/llvm-project/issues/61371

Reviewed By: dcaballe

Differential Revision: https://reviews.llvm.org/D146493
This commit is contained in:
Uday Bondhugula
2023-03-21 10:22:19 +05:30
parent d9661d79f4
commit 71d97df6af
2 changed files with 33 additions and 1 deletions

View File

@@ -537,7 +537,7 @@ void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
return;
// Caller is expected to fully compose map/operands if necessary.
assert((isTopLevelValue(val) || isAffineForInductionVar(val)) &&
assert((isTopLevelValue(val) || isAffineInductionVar(val)) &&
"non-terminal symbol / loop IV expected");
// Outer loop IVs could be used in forOp's bounds.
if (auto loop = getForInductionVarOwner(val)) {
@@ -547,6 +547,14 @@ void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
loop.emitWarning("failed to add domain info to constraint system"));
return;
}
if (auto parallel = getAffineParallelInductionVarOwner(val)) {
appendDimVar(parallel.getIVs());
if (failed(this->addAffineParallelOpDomain(parallel)))
LLVM_DEBUG(parallel.emitWarning(
"failed to add domain info to constraint system"));
return;
}
// Add top level symbol.
appendSymbolVar(val);
// Check if the symbol is a constant.

View File

@@ -286,3 +286,27 @@ func.func @empty_loops(%arg0: memref<1024x1024xf64>) {
// CHECK-NOT: memref.alloc
// CHECK: return
}
#map16 = affine_map<(d0, d1, d2) -> (d0 * 40 + d1 * 8 + d2 * 2)>
#map17 = affine_map<(d0, d1, d2) -> (d0 * 40 + d1 * 8 + d2 * 2 + 2)>
// CHECK-LABEL: func @affine_parallel
func.func @affine_parallel(%85:memref<2x5x4x2xi64>) {
affine.for %arg0 = 0 to 2 {
affine.parallel (%arg1) = (0) to (5) {
affine.parallel (%arg2) = (0) to (4) {
affine.for %arg3 = #map16(%arg0, %arg1, %arg2) to #map17(%arg0, %arg1, %arg2) {
%105 = affine.load %85[((%arg3 floordiv 2) floordiv 4) floordiv 5, ((%arg3 floordiv 2) floordiv 4) mod 5, (%arg3 floordiv 2) mod 4, %arg3 mod 2] : memref<2x5x4x2xi64>
}
}
}
}
// CHECK: affine.for
// CHECK-NEXT: affine.for %{{.*}} = 0 to 5
// CHECK-NEXT: affine.for %{{.*}} = 0 to 4
// CHECK-NEXT: affine.for %{{.*}} = 0 to 2
// CHECK: affine.for
// CHECK-NEXT: affine.parallel
// CHECK-NEXT: affine.parallel
return
}