[Mlir-commits] [mlir] [mlir][affine] Add an integer range interface to `affine.apply` (PR #174277)

[Krzysztof Drewniak]

================
@@ -768,3 +769,130 @@ mlir::intrange::inferShapedDimOpInterface(ShapedDimOpInterface op,
   }
   return result.value_or(ConstantIntRanges::fromSigned(zero, typeMax));
 }
+
+//===----------------------------------------------------------------------===//
+// Affine expression inference
+//===----------------------------------------------------------------------===//
+
+ConstantIntRanges
+mlir::intrange::inferAffineExpr(AffineExpr expr,
+                                ArrayRef<ConstantIntRanges> dimRanges,
+                                ArrayRef<ConstantIntRanges> symbolRanges) {
+  switch (expr.getKind()) {
+  case AffineExprKind::Constant: {
+    auto constExpr = cast<AffineConstantExpr>(expr);
+    APInt value(indexMaxWidth, constExpr.getValue(), /*isSigned=*/true);
+    return ConstantIntRanges::constant(value);
+  }
+  case AffineExprKind::DimId: {
+    auto dimExpr = cast<AffineDimExpr>(expr);
+    unsigned pos = dimExpr.getPosition();
+    assert(pos < dimRanges.size() && "Dimension index out of bounds");
+    return dimRanges[pos];
+  }
+  case AffineExprKind::SymbolId: {
+    auto symbolExpr = cast<AffineSymbolExpr>(expr);
+    unsigned pos = symbolExpr.getPosition();
+    assert(pos < symbolRanges.size() && "Symbol index out of bounds");
+    return symbolRanges[pos];
+  }
+  case AffineExprKind::Add: {
+    auto binExpr = cast<AffineBinaryOpExpr>(expr);
+    ConstantIntRanges lhs =
+        inferAffineExpr(binExpr.getLHS(), dimRanges, symbolRanges);
+    ConstantIntRanges rhs =
+        inferAffineExpr(binExpr.getRHS(), dimRanges, symbolRanges);
+    return inferAdd({lhs, rhs}, OverflowFlags::Nsw);
+  }
+  case AffineExprKind::Mul: {
+    auto binExpr = cast<AffineBinaryOpExpr>(expr);
+    ConstantIntRanges lhs =
+        inferAffineExpr(binExpr.getLHS(), dimRanges, symbolRanges);
+    ConstantIntRanges rhs =
+        inferAffineExpr(binExpr.getRHS(), dimRanges, symbolRanges);
+    return inferMul({lhs, rhs}, OverflowFlags::Nsw);
+  }
+  case AffineExprKind::Mod: {
+    auto binExpr = cast<AffineBinaryOpExpr>(expr);
+    ConstantIntRanges lhs =
+        inferAffineExpr(binExpr.getLHS(), dimRanges, symbolRanges);
+    ConstantIntRanges rhs =
+        inferAffineExpr(binExpr.getRHS(), dimRanges, symbolRanges);
+    // Affine mod is Euclidean modulo: result is always in [0, rhs-1].
+    // This assumes RHS is positive (enforced by affine expr semantics).
+    const APInt &lhsMin = lhs.smin(), &lhsMax = lhs.smax();
+    const APInt &rhsMin = rhs.smin(), &rhsMax = rhs.smax();
+    unsigned width = rhsMin.getBitWidth();
+    APInt zero = APInt::getZero(width);
+
+    // Guard against division by zero.
+    if (rhsMax.isZero())
+      return ConstantIntRanges::maxRange(width);
+
+    // For Euclidean mod, result is in [0, max(rhs)-1].
+    APInt umin = zero;
+    APInt umax = rhsMax - 1;
+
+    // Special case: if dividend is already in [0, min(rhs)), result equals
+    // dividend. We use rhsMin to ensure this is safe for all possible divisor
+    // values.
+    if (rhsMin.isStrictlyPositive() && lhsMin.isNonNegative() &&
+        lhsMax.ult(rhsMin)) {
+      umin = lhsMin;
+      umax = lhsMax;
+    }
+    // Special case: sweeping out a contiguous range with constant divisor.
+    // Only applies when dividend is non-negative to ensure result range is
+    // contiguous.
+    else if (rhsMin == rhsMax && lhsMin.isNonNegative() &&
+             (lhsMax - lhsMin).ult(rhsMax)) {
+      // For non-negative dividends, Euclidean mod is same as unsigned
+      // remainder.
+      umin = lhsMin.urem(rhsMax);
+      umax = lhsMax.urem(rhsMax);
----------------
krzysz00 wrote:

See comment in this bit of code - if we're here, we know everything's non-negative

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