[Mlir-commits] [mlir] [mlir][affine][Analysis] Add conservative bounds for semi-affine mods (PR #93576)

[Benjamin Maxwell]

================
@@ -63,76 +64,150 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
     // Update localVarCst.
     localVarCst.addLocalFloorDiv(dividend, divisor);
   }
+
+  // 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.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.
----------------
MacDue wrote:

Tried to tidy this up :+1: 

https://github.com/llvm/llvm-project/pull/93576