[mlir][affine] Fix an error when symbol is not defined at top level

[sotoshigoto]

Summary

In the execution of the AffineParallelize pass, when Symbols of
affine.loop are defined in a top-level region without the AffineScope
trait, an assertion error occurs. By changing the Symbol check during
addBound to isValidSymbol, we address this issue by extending the Symbol
condition beyond just top-level region conditions.

This issue was reported at
https://github.com/llvm/llvm-project/issues/64287

[sotoshigoto]

cleanup

[sotoshigoto]

cleanup

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

 using namespace presburger;
 
 
 void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
   if (containsVar(val))
     return;
 
   // Caller is expected to fully compose map/operands if necessary.
-  assert((isTopLevelValue(val) || isAffineInductionVar(val)) &&
-         "non-terminal symbol / loop IV expected");
+  assert((isValidSymbol(val) || isAffineInductionVar(val)) &&
+         "valid 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.
+  // Add valid symbol.
   appendSymbolVar(val);
   // Check if the symbol is a constant.
   if (std::optional<int64_t> constOp = getConstantIntValue(val))
     addBound(BoundType::EQ, val, constOp.value());
 }
 
 LogicalResult
 FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) {

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

if (isTopLevelValue(value))
     return true;
 
   if (auto *defOp = value.getDefiningOp())
     return isValidSymbol(value, getAffineScope(defOp));
 
   return false;
 }
 
-/// A value can be used as a symbol for `region` iff it meets one of the
+/// A value can be used as a symbol for `region` if it meets one of the
 /// following conditions:
 /// *) It is a constant.
 /// *) It is the result of an affine apply operation with symbol arguments.
 /// *) It is a result of the dim op on a memref whose corresponding size is
 ///    a valid symbol.
 /// *) It is defined at the top level of 'region' or is its argument.
 /// *) It dominates `region`'s parent op.
 /// If `region` is null, conservatively assume the symbol definition scope does

mlir/test/Dialect/Affine/parallelize.mlir

func.func @iter_arg_memrefs(%in: memref<10xf32>) {
   %mi = memref.alloc() : memref<f32>
   // Loop-carried memrefs are treated as serializing the loop.
   // CHECK: affine.for
   %mo = affine.for %i = 0 to 10 iter_args(%m_arg = %mi) -> (memref<f32>) {
     affine.yield %m_arg : memref<f32>
   }
   return
 }
+
+// CHECK-LABEL: @no_toplevel_block_parallelize
+func.func @no_toplevel_block_parallelize(%arg0: i1) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %alloc_0 = memref.alloc() : memref<2x3x6xi1>
+
+  scf.if %arg0 {
+    %alloc_1 = memref.alloc() : memref<2x3x6x9xi1>
+    %dim_0 = memref.dim %alloc_0, %c0 : memref<2x3x6xi1>
+    %dim_1 = memref.dim %alloc_0, %c1 : memref<2x3x6xi1>
+    affine.for %arg2 = 0 to %dim_0 {
+      affine.for %arg3 = 0 to %dim_1 {
+        affine.for %arg4 = 0 to 6 {
+          affine.for %arg5 = 0 to 9 {
+            %load_0 = affine.load %alloc_0[%arg2, %arg3, %arg4] : memref<2x3x6xi1>
+            affine.store %load_0, %alloc_1[%arg2, %arg3, %arg4, %arg5] : memref<2x3x6x9xi1>
+          }
+        }
+      }
+    }
+  }
+  // CHECK: affine.parallel (%{{.*}}) = (0) to (symbol(%{{.*}})) {
+  // CHECK:   affine.parallel (%{{.*}}) = (0) to (symbol(%{{.*}})) {
+  // CHECK:    affine.parallel (%{{.*}}) = (0) to (6) {
+  // CHECK:      affine.parallel (%{{.*}}) = (0) to (9) {
+  return
+}