[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

diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
index 9c9d089..03b8b1d 100644 (file)
--- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
@@ -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.

diff --git a/mlir/test/Dialect/Affine/affine-data-copy.mlir b/mlir/test/Dialect/Affine/affine-data-copy.mlir
index a5fa22a..22fbd73 100644 (file)
--- a/mlir/test/Dialect/Affine/affine-data-copy.mlir
+++ b/mlir/test/Dialect/Affine/affine-data-copy.mlir
@@ -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
+}