[MLIR][Affine] Fix addInductionVarOrTerminalSymbol to return status (#129476)

Fixes: #64287

This method is failable on valid IR and we should've been returning
failure instead of asserting, and checking status at its users.

Author: bondhugula

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

@@ -106,11 +106,14 @@ class FlatAffineValueConstraints : public FlatLinearValueConstraints {
 
   /// Add the specified values as a dim or symbol var depending on its nature,
   /// if it already doesn't exist in the system. `val` has to be either a
-  /// terminal symbol or a loop IV, i.e., it cannot be the result affine.apply
-  /// of any symbols or loop IVs. The variable is added to the end of the
-  /// existing dims or symbols. Additional information on the variable is
-  /// extracted from the IR and added to the constraint system.
-  void addInductionVarOrTerminalSymbol(Value val);
+  /// terminal symbol or a loop IV, i.e., it cannot be the result of an
+  /// affine.apply of any symbols or loop IVs. Return failure if the addition
+  /// wasn't possible due to the above conditions not being met. This method can
+  /// also fail if the addition of the domain of an affine IV fails. The
+  /// variable is added to the end of the existing dims or symbols. Additional
+  /// information on the variable is extracted from the IR and added to the
+  /// constraint system.
+  LogicalResult addInductionVarOrTerminalSymbol(Value val);
 
   /// Adds slice lower bounds represented by lower bounds in `lbMaps` and upper
   /// bounds in `ubMaps` to each variable in the constraint system which has

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

@@ -34,35 +34,45 @@ using namespace mlir;
 using namespace affine;
 using namespace presburger;
 
-
-void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
+LogicalResult
+FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
   if (containsVar(val))
-    return;
+    return success();
 
   // Caller is expected to fully compose map/operands if necessary.
-  assert((isTopLevelValue(val) || isAffineInductionVar(val)) &&
-         "non-terminal symbol / loop IV expected");
+  if (val.getDefiningOp<affine::AffineApplyOp>() ||
+      (!isValidSymbol(val) && !isAffineInductionVar(val))) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "only valid terminal symbols and affine IVs supported\n");
+    return failure();
+  }
   // Outer loop IVs could be used in forOp's bounds.
   if (auto loop = getForInductionVarOwner(val)) {
     appendDimVar(val);
-    if (failed(this->addAffineForOpDomain(loop)))
+    if (failed(this->addAffineForOpDomain(loop))) {
       LLVM_DEBUG(
           loop.emitWarning("failed to add domain info to constraint system"));
-    return;
+      return failure();
+    }
+    return success();
   }
+
   if (auto parallel = getAffineParallelInductionVarOwner(val)) {
     appendDimVar(parallel.getIVs());
-    if (failed(this->addAffineParallelOpDomain(parallel)))
+    if (failed(this->addAffineParallelOpDomain(parallel))) {
       LLVM_DEBUG(parallel.emitWarning(
           "failed to add domain info to constraint system"));
-    return;
+      return failure();
+    }
+    return success();
   }
 
   // Add top level symbol.
   appendSymbolVar(val);
   // Check if the symbol is a constant.
   if (std::optional<int64_t> constOp = getConstantIntValue(val))
     addBound(BoundType::EQ, val, constOp.value());
+  return success();
 }
 
 LogicalResult
@@ -222,8 +232,10 @@ LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,
   fullyComposeAffineMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
   canonicalizeMapAndOperands(&map, &operands);
-  for (auto operand : operands)
-    addInductionVarOrTerminalSymbol(operand);
+  for (Value operand : operands) {
+    if (failed(addInductionVarOrTerminalSymbol(operand)))
+      return failure();
+  }
   return addBound(type, pos, computeAlignedMap(map, operands));
 }
 

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

@@ -1255,7 +1255,8 @@ LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   if (sliceState != nullptr) {
     // Add dim and symbol slice operands.
     for (auto operand : sliceState->lbOperands[0]) {
-      cst.addInductionVarOrTerminalSymbol(operand);
+      if (failed(cst.addInductionVarOrTerminalSymbol(operand)))
+        return failure();
     }
     // Add upper/lower bounds from 'sliceState' to 'cst'.
     LogicalResult ret =

mlir/test/Dialect/Affine/parallelize.mlir

@@ -323,3 +323,21 @@ func.func @iter_arg_memrefs(%in: memref<10xf32>) {
   }
   return
 }
+
+// Test affine analysis machinery to ensure it generates valid IR and doesn't
+// crash on this combination of ops.
+
+// CHECK-LABEL: @test_add_inv_or_terminal_symbol
+func.func @test_add_inv_or_terminal_symbol(%arg0: memref<9x9xi32>, %arg1: i1) {
+  %idx0 = index.constant 1
+  %29 = tensor.empty() : tensor<10xf16>
+  memref.alloca_scope {
+    %dim_30 = tensor.dim %29, %idx0 : tensor<10xf16>
+    %alloc_31 = memref.alloc(%idx0, %idx0) {alignment = 64 : i64} : memref<?x?xf16>
+    affine.for %arg3 = 0 to %dim_30 {
+      %207 = affine.load %alloc_31[%idx0, %idx0] : memref<?x?xf16>
+      affine.store %207, %alloc_31[%idx0, %idx0] : memref<?x?xf16>
+    }
+  }
+  return
+}