Keep changes within FlatLinearValueConstraints

MacDue · MacDue · commit 179b27bcef7f · 2024-05-29T17:08:52.000Z
diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -454,20 +454,6 @@ class IntegerRelation {
     addLocalFloorDiv(getMPIntVec(dividend), MPInt(divisor));
   }
 
-  /// Adds a new local variable as the mod of an affine function of other
-  /// variables. The coefficients of the operands of the mod are provided in
-  /// `lhs` and `rhs` respectively. Three constraints are added to provide a
-  /// conservative bound for the mod:
-  ///  1. rhs > 0 (assumption/precondition)
-  ///  2. lhs % rhs < rhs
-  ///  3. lhs % rhs >= 0
-  /// We ensure the rhs is positive so we can assume the result is positive.
-  void addLocalModConservativeBounds(ArrayRef<MPInt> lhs, ArrayRef<MPInt> rhs);
-  void addLocalModConservativeBounds(ArrayRef<int64_t> lhs,
-                                     ArrayRef<int64_t> rhs) {
-    addLocalModConservativeBounds(getMPIntVec(lhs), getMPIntVec(rhs));
-  }
-
   /// Projects out (aka eliminates) `num` variables starting at position
   /// `pos`. The resulting constraint system is the shadow along the dimensions
   /// that still exist. This method may not always be integer exact.
diff --git a/mlir/include/mlir/IR/AffineExprVisitor.h b/mlir/include/mlir/IR/AffineExprVisitor.h
@@ -413,8 +413,9 @@ class SimpleAffineExprFlattener
   /// lhs of the mod, floordiv, ceildiv or mul expression and with respect to a
   /// symbolic rhs expression. `localExpr` is the simplified tree expression
   /// (AffineExpr) corresponding to the quantifier.
-  virtual void addLocalIdSemiAffine(AffineExpr localExpr, ArrayRef<int64_t> lhs,
-                                    ArrayRef<int64_t> rhs);
+  virtual LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+                                             ArrayRef<int64_t> rhs,
+                                             AffineExpr localExpr);
 
 private:
   /// Adds `expr`, which may be mod, ceildiv, floordiv or mod expression
@@ -423,10 +424,11 @@ class SimpleAffineExprFlattener
   /// quantifier is already present, we put the coefficient in the proper index
   /// of `result`, otherwise we add a new local variable and put the coefficient
   /// there.
-  void addLocalVariableSemiAffine(AffineExpr expr, ArrayRef<int64_t> lhs,
-                                  ArrayRef<int64_t> rhs,
-                                  SmallVectorImpl<int64_t> &result,
-                                  unsigned long resultSize);
+  LogicalResult addLocalVariableSemiAffine(AffineExpr expr,
+                                           ArrayRef<int64_t> lhs,
+                                           ArrayRef<int64_t> rhs,
+                                           SmallVectorImpl<int64_t> &result,
+                                           unsigned long resultSize);
 
   // t = expr floordiv c   <=> t = q, c * q <= expr <= c * q + c - 1
   // A floordiv is thus flattened by introducing a new local variable q, and
diff --git a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -36,25 +36,20 @@ using namespace presburger;
 namespace {
 
 // See comments for SimpleAffineExprFlattener.
-// An AffineExprFlattener extends a SimpleAffineExprFlattener by recording
-// constraint information associated with mod's, floordiv's, and ceildiv's
-// in FlatLinearConstraints 'localVarCst'.
-struct AffineExprFlattener : public SimpleAffineExprFlattener {
-public:
+// An AffineExprFlattenerWithLocalVars extends a SimpleAffineExprFlattener by
+// recording constraint information associated with mod's, floordiv's, and
+// ceildiv's in FlatLinearConstraints 'localVarCst'.
+struct AffineExprFlattenerWithLocalVars : public SimpleAffineExprFlattener {
+  using SimpleAffineExprFlattener::SimpleAffineExprFlattener;
+
   // Constraints connecting newly introduced local variables (for mod's and
   // div's) to existing (dimensional and symbolic) ones. These are always
   // inequalities.
   IntegerPolyhedron localVarCst;
 
-  AffineExprFlattener(unsigned nDims, unsigned nSymbols,
-                      bool addConservativeSemiAffineBounds = false)
+  AffineExprFlattenerWithLocalVars(unsigned nDims, unsigned nSymbols)
       : SimpleAffineExprFlattener(nDims, nSymbols),
-        localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)),
-        addConservativeSemiAffineBounds(addConservativeSemiAffineBounds) {}
-
-  bool hasUnhandledSemiAffineExpressions() const {
-    return unhandledSemiAffineExpressions;
-  }
+        localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {};
 
 private:
   // Add a local variable (needed to flatten a mod, floordiv, ceildiv expr).
@@ -70,30 +65,71 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
     localVarCst.addLocalFloorDiv(dividend, divisor);
   }
 
-  // Add a local identifier (needed to flatten a mod, floordiv, ceildiv, mul
-  // expr) when the rhs is a symbolic expression. The local identifier added
-  // may be a floordiv, ceildiv, mul or mod of a pure affine/semi-affine
-  // function of other identifiers, coefficients of which are specified in the
-  // lhs of the mod, floordiv, ceildiv or mul expression and with respect to a
-  // symbolic rhs expression. `localExpr` is the simplified tree expression
-  // (AffineExpr) corresponding to the quantifier.
-  void addLocalIdSemiAffine(AffineExpr localExpr, ArrayRef<int64_t> lhs,
-                            ArrayRef<int64_t> rhs) override {
-    SimpleAffineExprFlattener::addLocalIdSemiAffine(localExpr, lhs, rhs);
-    if (!addConservativeSemiAffineBounds) {
-      unhandledSemiAffineExpressions = true;
-      return;
-    }
+  // Semi-affine expressions are not supported by all flatteners.
+  LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+                                     ArrayRef<int64_t> rhs,
+                                     AffineExpr localExpr) override = 0;
+};
+
+// An AffineExprFlattener is an AffineExprFlattenerWithLocalVars that explicitly
+// disallows semi-affine expressions. Flattening will fail if a semi-affine
+// expression is encountered.
+struct AffineExprFlattener : public AffineExprFlattenerWithLocalVars {
+  using AffineExprFlattenerWithLocalVars::AffineExprFlattenerWithLocalVars;
+
+  LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+                                     ArrayRef<int64_t> rhs,
+                                     AffineExpr localExpr) override {
+    // AffineExprFlattener does not support semi-affine expressions.
+    return failure();
+  }
+};
+
+// A SemiAffineExprFlattener is an AffineExprFlattenerWithLocalVars that adds
+// conservative bounds for semi-affine expressions (given assumptions hold). If
+// the assumptions required to add the semi-affine bounds are found not to hold
+// the final constraints set will be empty/inconsistent. If the assumptions are
+// never contradicted the final bounds still only will be correct if the
+// assumptions hold.
+struct SemiAffineExprFlattener : public AffineExprFlattenerWithLocalVars {
+  using AffineExprFlattenerWithLocalVars::AffineExprFlattenerWithLocalVars;
+
+  LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+                                     ArrayRef<int64_t> rhs,
+                                     AffineExpr localExpr) override {
+    auto result =
+        SimpleAffineExprFlattener::addLocalIdSemiAffine(lhs, rhs, localExpr);
+    assert(succeeded(result) &&
+           "unexpected failure in SimpleAffineExprFlattener");
+    (void)result;
+
     if (localExpr.getKind() == AffineExprKind::Mod) {
-      localVarCst.addLocalModConservativeBounds(lhs, rhs);
-      return;
+      localVarCst.appendVar(VarKind::Local);
+      // Add a conservative bound for `mod` assuming the rhs is > 0.
+
+      // Note: If the rhs is later found to be < 0 the following two constraints
+      // will contradict each other (and lead to the final constraints set
+      // becoming empty). If the sign of the rhs is never specified the bound
+      // will assume it is positive.
+
+      // Upper bound: rhs - (lhs % rhs) - 1 >= 0 i.e. lhs % rhs < rhs
+      // This only holds if the rhs is > 0.
+      SmallVector<int64_t, 8> resultUpperBound(rhs);
+      resultUpperBound.insert(resultUpperBound.end() - 1, -1);
+      --resultUpperBound.back();
+      localVarCst.addInequality(resultUpperBound);
+
+      // Lower bound: lhs % rhs >= 0 (always holds)
+      SmallVector<int64_t, 8> resultLowerBound(rhs.size());
+      resultLowerBound.insert(resultLowerBound.end() - 1, 1);
+      localVarCst.addInequality(resultLowerBound);
+
+      return success();
     }
+
     // TODO: Support other semi-affine expressions.
-    unhandledSemiAffineExpressions = true;
+    return failure();
   }
-
-  bool addConservativeSemiAffineBounds = false;
-  bool unhandledSemiAffineExpressions = false;
 };
 
 } // namespace
@@ -114,27 +150,34 @@ getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
     return success();
   }
 
-  AffineExprFlattener flattener(numDims, numSymbols,
-                                addConservativeSemiAffineBounds);
-  // Use the same flattener to simplify each expression successively. This way
-  // local variables / expressions are shared.
-  for (auto expr : exprs) {
-    auto flattenResult = flattener.walkPostOrder(expr);
-    if (failed(flattenResult))
-      return failure();
-    if (flattener.hasUnhandledSemiAffineExpressions())
-      return failure();
-  }
+  auto flattenExprs =
+      [&](AffineExprFlattenerWithLocalVars &flattener) -> LogicalResult {
+    // Use the same flattener to simplify each expression successively. This way
+    // local variables / expressions are shared.
+    for (auto expr : exprs) {
+      auto flattenResult = flattener.walkPostOrder(expr);
+      if (failed(flattenResult))
+        return failure();
+    }
+
+    assert(flattener.operandExprStack.size() == exprs.size());
+    flattenedExprs->clear();
+    flattenedExprs->assign(flattener.operandExprStack.begin(),
+                           flattener.operandExprStack.end());
 
-  assert(flattener.operandExprStack.size() == exprs.size());
-  flattenedExprs->clear();
-  flattenedExprs->assign(flattener.operandExprStack.begin(),
-                         flattener.operandExprStack.end());
+    if (localVarCst)
+      localVarCst->clearAndCopyFrom(flattener.localVarCst);
+
+    return success();
+  };
 
-  if (localVarCst)
-    localVarCst->clearAndCopyFrom(flattener.localVarCst);
+  if (addConservativeSemiAffineBounds) {
+    SemiAffineExprFlattener flattener(numDims, numSymbols);
+    return flattenExprs(flattener);
+  }
 
-  return success();
+  AffineExprFlattener flattener(numDims, numSymbols);
+  return flattenExprs(flattener);
 }
 
 // Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -1521,37 +1521,6 @@ void IntegerRelation::addLocalFloorDiv(ArrayRef<MPInt> dividend,
       getDivUpperBound(dividendCopy, divisor, dividendCopy.size() - 2));
 }
 
-/// Adds a new local variable as the mod of an affine function of other
-/// variables. The coefficients of the operands of the mod are provided in `lhs`
-/// and `rhs` respectively. Three constraints are added to provide a
-/// conservative bound for the mod:
-///  1. rhs > 0 (assumption/precondition)
-///  2. lhs % rhs < rhs
-///  3. lhs % rhs >= 0
-/// We ensure the rhs is positive so we can assume the result is positive.
-void IntegerRelation::addLocalModConservativeBounds(ArrayRef<MPInt> lhs,
-                                                    ArrayRef<MPInt> rhs) {
-  appendVar(VarKind::Local);
-
-  // Ensure the rhs is > 0 (most likely case).
-  // If this constraint does not hold the following bounds are incorrect.
-  SmallVector<MPInt, 8> rhsCopy(rhs);
-  rhsCopy.insert(rhsCopy.end() - 1, MPInt(0));
-  rhsCopy.back() -= MPInt(1);
-  addInequality(rhsCopy);
-
-  // rhs - (lhs % rhs) - 1 >= 0 i.e. lhs % rhs < rhs
-  SmallVector<MPInt, 8> resultUpperBound(rhs);
-  resultUpperBound.insert(resultUpperBound.end() - 1, MPInt(-1));
-  resultUpperBound.back() -= MPInt(1);
-  addInequality(resultUpperBound);
-
-  // lhs % rhs >= 0
-  SmallVector<MPInt, 8> resultLowerBound(rhs.size());
-  resultLowerBound.insert(resultLowerBound.end() - 1, MPInt(1));
-  addInequality(resultLowerBound);
-}
-
 /// Finds an equality that equates the specified variable to a constant.
 /// Returns the position of the equality row. If 'symbolic' is set to true,
 /// symbols are also treated like a constant, i.e., an affine function of the
diff --git a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
@@ -7,10 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
-
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "llvm/Support/Debug.h"
-
 namespace mlir::vector {
 
 FailureOr<ConstantOrScalableBound::BoundSize>
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
@@ -1248,8 +1248,7 @@ LogicalResult SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
                                              localExprs, context);
     AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
                                              localExprs, context);
-    addLocalVariableSemiAffine(a * b, mulLhs, rhs, lhs, lhs.size());
-    return success();
+    return addLocalVariableSemiAffine(a * b, mulLhs, rhs, lhs, lhs.size());
   }
 
   // Get the RHS constant.
@@ -1302,8 +1301,7 @@ LogicalResult SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
     AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
                                                        localExprs, context);
     AffineExpr modExpr = dividendExpr % divisorExpr;
-    addLocalVariableSemiAffine(modExpr, modLhs, rhs, lhs, lhs.size());
-    return success();
+    return addLocalVariableSemiAffine(modExpr, modLhs, rhs, lhs, lhs.size());
   }
 
   int64_t rhsConst = rhs[getConstantIndex()];
@@ -1387,19 +1385,22 @@ SimpleAffineExprFlattener::visitConstantExpr(AffineConstantExpr expr) {
   return success();
 }
 
-void SimpleAffineExprFlattener::addLocalVariableSemiAffine(
+LogicalResult SimpleAffineExprFlattener::addLocalVariableSemiAffine(
     AffineExpr expr, ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs,
     SmallVectorImpl<int64_t> &result, unsigned long resultSize) {
   assert(result.size() == resultSize &&
          "`result` vector passed is not of correct size");
   int loc;
-  if ((loc = findLocalId(expr)) == -1)
-    addLocalIdSemiAffine(expr, lhs, rhs);
+  if ((loc = findLocalId(expr)) == -1) {
+    if (failed(addLocalIdSemiAffine(lhs, rhs, expr)))
+      return failure();
+  }
   std::fill(result.begin(), result.end(), 0);
   if (loc == -1)
     result[getLocalVarStartIndex() + numLocals - 1] = 1;
   else
     result[getLocalVarStartIndex() + loc] = 1;
+  return success();
 }
 
 // t = expr floordiv c   <=> t = q, c * q <= expr <= c * q + c - 1
@@ -1434,8 +1435,7 @@ LogicalResult SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
     AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
                                              localExprs, context);
     AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
-    addLocalVariableSemiAffine(divExpr, divLhs, rhs, lhs, lhs.size());
-    return success();
+    return addLocalVariableSemiAffine(divExpr, divLhs, rhs, lhs, lhs.size());
   }
 
   // This is a pure affine expr; the RHS is a positive constant.
@@ -1506,14 +1506,14 @@ void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend,
   // dividend and divisor are not used here; an override of this method uses it.
 }
 
-void SimpleAffineExprFlattener::addLocalIdSemiAffine(AffineExpr localExpr,
-                                                     ArrayRef<int64_t> lhs,
-                                                     ArrayRef<int64_t> rhs) {
+LogicalResult SimpleAffineExprFlattener::addLocalIdSemiAffine(
+    ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr) {
   for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
     subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
   localExprs.push_back(localExpr);
   ++numLocals;
   // lhs and rhs are not used here; an override of this method uses them.
+  return success();
 }
 
 int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) {
