[MLIR] Allow recovering from semi-affine sets in FlatLinearValueConst… (#212346)

…raints

The IntegerSet constructor of FlatLinearValueConstraints asserts when
flattening fails, which happens for semi-affine integer sets. Since the
failure is signalled from inside a constructor, callers have no way to
detect the unsupported case ahead of time and no way to work around it;
they simply crash.

Add an optional `bool *error` parameter to the constructor. When it is
non-null, hitting the unimplemented case sets `*error` to true and
returns instead of asserting, letting callers bail out gracefully.
`*error` is set to false on success. When the pointer is null the
previous assert behavior is retained, so existing callers are
unaffected.

mlir/include/mlir/Analysis/FlatLinearValueConstraints.h

@@ -334,8 +334,10 @@ public:
         setValue(i, *valArgs[i]);
   }
 
-  /// Creates an affine constraint system from an IntegerSet.
-  explicit FlatLinearValueConstraints(IntegerSet set, ValueRange operands = {});
+  /// Creates an affine constraint system from an IntegerSet. Returns failure
+  /// if `set` is semi-affine, as flattening is not implemented for those.
+  static FailureOr<FlatLinearValueConstraints>
+  create(IntegerSet set, ValueRange operands = {});
 
   /// Return the kind of this object.
   Kind getKind() const override { return Kind::FlatLinearValueConstraints; }
@@ -517,6 +519,12 @@ public:
   ///    output = {0 <= d0 <= 6, 1 <= d1 <= 15}
   LogicalResult unionBoundingBox(const FlatLinearValueConstraints &other);
   using IntegerPolyhedron::unionBoundingBox;
+
+protected:
+  /// Creates an affine constraint system from an IntegerSet. `error` is set to
+  /// true if `set` could not be flattened, in which case the constraint system
+  /// is left without the constraints of `set`. Use `create` instead.
+  FlatLinearValueConstraints(IntegerSet set, ValueRange operands, bool *error);
 };
 
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the

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

@@ -45,6 +45,11 @@ class FlatAffineValueConstraints : public FlatLinearValueConstraints {
 public:
   using FlatLinearValueConstraints::FlatLinearValueConstraints;
 
+  /// Creates an affine constraint system from an IntegerSet. Returns failure
+  /// if `set` is semi-affine, as flattening is not implemented for those.
+  static FailureOr<FlatAffineValueConstraints>
+  create(IntegerSet set, ValueRange operands = {});
+
   /// Return the kind of this object.
   Kind getKind() const override { return Kind::FlatAffineValueConstraints; }
 

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -1145,13 +1145,16 @@ IntegerSet FlatLinearConstraints::getAsIntegerSet(MLIRContext *context) const {
 
 // Construct from an IntegerSet.
 FlatLinearValueConstraints::FlatLinearValueConstraints(IntegerSet set,
-                                                       ValueRange operands)
+                                                       ValueRange operands,
+                                                       bool *error)
     : FlatLinearConstraints(set.getNumInequalities(), set.getNumEqualities(),
                             set.getNumDims() + set.getNumSymbols() + 1,
                             set.getNumDims(), set.getNumSymbols(),
                             /*numLocals=*/0) {
+  assert(error && "expected a valid error flag");
   assert((operands.empty() || set.getNumInputs() == operands.size()) &&
          "operand count mismatch");
+  *error = false;
   // Set the values for the non-local variables.
   for (unsigned i = 0, e = operands.size(); i < e; ++i)
     setValue(i, operands[i]);
@@ -1160,7 +1163,8 @@ FlatLinearValueConstraints::FlatLinearValueConstraints(IntegerSet set,
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   FlatLinearConstraints localVarCst;
   if (failed(getFlattenedAffineExprs(set, &flatExprs, &localVarCst))) {
-    assert(false && "flattening unimplemented for semi-affine integer sets");
+    // Flattening is unimplemented for semi-affine integer sets.
+    *error = true;
     return;
   }
   assert(flatExprs.size() == set.getNumConstraints());
@@ -1180,6 +1184,15 @@ FlatLinearValueConstraints::FlatLinearValueConstraints(IntegerSet set,
   append(localVarCst);
 }
 
+FailureOr<FlatLinearValueConstraints>
+FlatLinearValueConstraints::create(IntegerSet set, ValueRange operands) {
+  bool error = false;
+  FlatLinearValueConstraints cst(set, operands, &error);
+  if (error)
+    return failure();
+  return cst;
+}
+
 unsigned FlatLinearValueConstraints::appendDimVar(ValueRange vals) {
   unsigned pos = getNumDimVars();
   return insertVar(VarKind::SetDim, pos, vals);

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

@@ -29,6 +29,15 @@ using namespace mlir;
 using namespace affine;
 using namespace presburger;
 
+FailureOr<FlatAffineValueConstraints>
+FlatAffineValueConstraints::create(IntegerSet set, ValueRange operands) {
+  bool error = false;
+  FlatAffineValueConstraints cst(set, operands, &error);
+  if (error)
+    return failure();
+  return cst;
+}
+
 LogicalResult
 FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
   if (containsVar(val))
@@ -208,13 +217,18 @@ void FlatAffineValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
   canonicalizeSetAndOperands(&set, &operands);
 
   // Create the base constraints from the integer set attached to ifOp.
-  FlatAffineValueConstraints cst(set, operands);
+  FailureOr<FlatAffineValueConstraints> cst =
+      FlatAffineValueConstraints::create(set, operands);
+  if (failed(cst)) {
+    assert(false && "semi-affine integer sets are unsupported here");
+    return;
+  }
 
   // Merge the constraints from ifOp to the current domain. We need first merge
   // and align the IDs from both constraints, and then append the constraints
   // from the ifOp into the current one.
-  mergeAndAlignVarsWithOther(0, &cst);
-  append(cst);
+  mergeAndAlignVarsWithOther(0, &*cst);
+  append(*cst);
 }
 
 LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,

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

@@ -2195,13 +2195,17 @@ void mlir::affine::getSequentialLoops(
 }
 
 IntegerSet mlir::affine::simplifyIntegerSet(IntegerSet set) {
-  FlatAffineValueConstraints fac(set);
-  if (fac.isEmpty())
+  FailureOr<FlatAffineValueConstraints> fac =
+      FlatAffineValueConstraints::create(set);
+  // Semi-affine sets can't be flattened; return them as is.
+  if (failed(fac))
+    return set;
+  if (fac->isEmpty())
     return IntegerSet::getEmptySet(set.getNumDims(), set.getNumSymbols(),
                                    set.getContext());
-  fac.removeTrivialRedundancy();
+  fac->removeTrivialRedundancy();
 
-  auto simplifiedSet = fac.getAsIntegerSet(set.getContext());
+  auto simplifiedSet = fac->getAsIntegerSet(set.getContext());
   assert(simplifiedSet && "guaranteed to succeed while roundtripping");
   return simplifiedSet;
 }

mlir/unittests/Analysis/Presburger/Parser.h

@@ -30,7 +30,10 @@ namespace presburger {
 /// represents a valid IntegerSet.
 inline IntegerPolyhedron parseIntegerPolyhedron(StringRef str) {
   MLIRContext context(MLIRContext::Threading::DISABLED);
-  return affine::FlatAffineValueConstraints(parseIntegerSet(str, &context));
+  FailureOr<affine::FlatAffineValueConstraints> cst =
+      affine::FlatAffineValueConstraints::create(parseIntegerSet(str, &context));
+  assert(succeeded(cst) && "expected a valid IntegerSet");
+  return *cst;
 }
 
 /// Parse a list of StringRefs to IntegerRelation and combine them into a