[MLIR][Affine] Fix addInductionVarOrTerminalSymbol

[bondhugula]

Summary

Update affine analysis method addInductionVarOrTerminalSymbol for
affine.parallel IV.

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

[bondhugula]

Add CHECK-LABEL line.

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

void FlatAffineValueConstraints::convertLoopIVSymbolsToDims() {
   }
 }
 
 void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
   if (containsVar(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)) {
     appendDimVar(val);
     if (failed(this->addAffineForOpDomain(loop)))
       LLVM_DEBUG(
           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.
   if (auto constOp = val.getDefiningOp<arith::ConstantIndexOp>())
     addBound(BoundType::EQ, val, constOp.value());
 }
 
 LogicalResult

mlir/test/Dialect/Affine/affine-data-copy.mlir

func.func @empty_loops(%arg0: memref<1024x1024xf64>) {
   }
   affine.for %i = 0 to -16 {
     affine.load %arg0[0, %i] : memref<1024x1024xf64>
   }
   return
   // 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
+}