[mlir][Index] Implement InferIntRangeInterface

[krzysz00]

Summary

Implement InferIntRangeInterface for all operations in the Index dialect. The
inference implementation, unlike the one for Arith, accounts for the
fact that Index can be either 64 or 32 bits long by evaluating both
cases. Bounds are stored as if index were i64, but when inferring new
bounds, we compute both f(...) and f(trunc(...)). We then compare
trunc(f(...)) to f(trunc(...)). If they are equal in the relevant
range components, we use the 64-bit range computation, otherwise we
give the range ext(f(trunc(...))) union f(...).

Note that this can cause surprising behavior as seen in the tests,
where, for example, the order of min and max operations impacts the
behavior of the inference. The inference could perhaps be made more
precise in the future (ex. by tracking 32 and 64-bit results
separately and having them influence each other somehow) butt, since
my project targets an index=i32 platform and doesn't see index-valued
values > uint32_max, I'm not too concerned about it.

Depends on https://reviews.llvm.org/D141299

Depends on https://reviews.llvm.org/D141296

[Mogball]

At a glance, this seems correct. Please split up the patch and I'll take a deeper look

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

309

please split out the addition of these ops from this patch

mlir/test/Dialect/Arith/int-range-interface.mlir

720 ↗

(On Diff #486010)

can this be split out of the current patch too?

[krzysz00]

(Comment seen, I'll probably split the patch on Monday)

[krzysz00]

Split patch

[krzysz00]

Add depends lines to commit message

[Mogball]

mlir/include/mlir/Interfaces/Utils/InferIntRangeCommon.h

29–30

32

I don't think it's important to explicitly specify the size here

34

106–115

can you drop the explicit values?

mlir/lib/Dialect/Index/IR/InferIntRangeInterfaceImpls.cpp

2

can you name this IndexOpsInferIntRange.cpp?

34

please drop const on these

165

can you move this comment down to the function?

185

please drop trivial braces

215

[krzysz00]

Rebase, adress most review comments

mlir/include/mlir/Interfaces/Utils/InferIntRangeCommon.h

32

I'd say it's important so anyone needing to forward-declare this enum knows what to do.

mlir/lib/Dialect/Index/IR/InferIntRangeInterfaceImpls.cpp

2

I was trying to match the convention over in Arith and other dialects where you have BlahOpInterfaceImpls.cpp, but I'm not strongly attached to the name and can change it if you'd like

[Mogball]

LGTM, thanks!

[krzysz00]

Rebase for std::optional migration

[krzysz00]

Fix accidental build error

mlir/include/mlir/Dialect/Index/IR/IndexOps.h

 //===- IndexOps.h - Index operation declarations ------------------*- C++-*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_INDEX_IR_INDEXOPS_H
 #define MLIR_DIALECT_INDEX_IR_INDEXOPS_H
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CastInterfaces.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Forward Declarations
 //===----------------------------------------------------------------------===//
 
 namespace mlir::index {

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

 //===- IndexOps.td - Index operation definitions -----------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef INDEX_OPS
 #define INDEX_OPS
 
 include "mlir/Dialect/Index/IR/IndexDialect.td"
 include "mlir/Dialect/Index/IR/IndexEnums.td"
 include "mlir/Interfaces/CastInterfaces.td"
+include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/OpBase.td"
 
 //===----------------------------------------------------------------------===//
 // IndexOp
 //===----------------------------------------------------------------------===//
 
 /// Base class for Index dialect operations.
 class IndexOp<string mnemonic, list<Trait> traits = []>
-    : Op<IndexDialect, mnemonic, [Pure] # traits>;
+    : Op<IndexDialect, mnemonic,
+      [Pure, DeclareOpInterfaceMethods<InferIntRangeInterface>] # traits>;
 
 //===----------------------------------------------------------------------===//
 // IndexBinaryOp
 //===----------------------------------------------------------------------===//
 
 /// Base class for binary Index dialect operations.
 class IndexBinaryOp<string mnemonic, list<Trait> traits = []>
     : IndexOp<mnemonic, traits> {
let description = [{
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // MinUOp
 //===----------------------------------------------------------------------===//
 
 def Index_MinUOp : IndexBinaryOp<"minu"> {
   let summary = "index unsigned minimum";
   let description = [{
     The `index.minu` operation takes two index values and computes their
     unsigned minimum value. Treats the leading bit as the most significant, i.e.
     `min(15, 6) = 6` or `min(-2, 6) = 6`.
 
     Example:
 

mlir/include/mlir/Interfaces/Utils/InferIntRangeCommon.h

+//===- InferIntRangeCommon.cpp - Inference for common ops --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares implementations of range inference for operations that are
+// common to both the `arith` and `index` dialects to facilitate reuse.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_UTILS_INFERINTRANGECOMMON_H
+#define MLIR_INTERFACES_UTILS_INFERINTRANGECOMMON_H
+
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+#include "llvm/ADT/ArrayRef.h"
+
+namespace mlir {
+namespace intrange {
+/// Function that performs inference on an array of `ConstantIntRanges`,
+/// abstracted away here to permit writing the function that handles both
+/// 64- and 32-bit index types.
+using InferRangeFn =
+    function_ref<ConstantIntRanges(ArrayRef<ConstantIntRanges>)>;
+
+static constexpr unsigned indexMinWidth = 32;
+static constexpr unsigned indexMaxWidth = 64;
+
+enum class CmpMode : uint32_t { Both, Signed, Unsigned };
+
+/// Compute `inferFn` on `ranges`, whose size should be the index storage
+/// bitwidth. Then, compute the function on `argRanges` again after truncating
+/// the ranges to 32 bits. Finally, if the truncation of the 64-bit result is
+/// equal to the 32-bit result, use it (to preserve compatibility with folders
+/// and inference precision), and take the union of the results otherwise.
+///
+/// The `mode` argument specifies if the unsigned, signed, or both results of
+/// the inference computation should be used when comparing the results.
+ConstantIntRanges inferIndexOp(InferRangeFn inferFn,
+                               ArrayRef<ConstantIntRanges> argRanges,
+                               CmpMode mode);
+
+/// Independently zero-extend the unsigned values and sign-extend the signed
+/// values in `range` to `destWidth` bits, returning the resulting range.
+ConstantIntRanges extRange(const ConstantIntRanges &range, unsigned destWidth);
+
+/// Use the unsigned values in `range` to zero-extend it to `destWidth`.
+ConstantIntRanges extUIRange(const ConstantIntRanges &range,
+                             unsigned destWidth);
+
+/// Use the signed values in `range` to sign-extend it to `destWidth`.
+ConstantIntRanges extSIRange(const ConstantIntRanges &range,
+                             unsigned destWidth);
+
+/// Truncate `range` to `destWidth` bits, taking care to handle cases such as
+/// the truncation of [255, 256] to i8 not being a uniform range.
+ConstantIntRanges truncRange(const ConstantIntRanges &range,
+                             unsigned destWidth);
+
+ConstantIntRanges inferAdd(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferSub(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferMul(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferDivS(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferDivU(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferCeilDivS(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferCeilDivU(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferFloorDivS(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferRemS(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferRemU(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferMaxS(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferMaxU(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferMinS(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferMinU(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferAnd(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferOr(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferXor(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferShl(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferShrS(ArrayRef<ConstantIntRanges> argRanges);
+
+ConstantIntRanges inferShrU(ArrayRef<ConstantIntRanges> argRanges);
+
+/// Copy of the enum from `arith` and `index` to allow the common integer range
+/// infrastructure to not depend on either dialect.
+enum class CmpPredicate : uint64_t {
+  eq,
+  ne,
+  slt,
+  sle,
+  sgt,
+  sge,
+  ult,
+  ule,
+  ugt,
+  uge,
+};
+
+/// Returns a boolean value if `pred` is statically true or false for
+/// anypossible inputs falling within `lhs` and `rhs`, and std::nullopt if the
+/// value of the predicate cannot be determined.
+Optional<bool> evaluatePred(CmpPredicate pred, const ConstantIntRanges &lhs,
+                            const ConstantIntRanges &rhs);
+
+} // namespace intrange
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_UTILS_INFERINTRANGECOMMON_H

mlir/lib/Dialect/Arith/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRArithDialect
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith
 
   DEPENDS
   MLIRArithOpsIncGen
   MLIRArithOpsInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
+  MLIRInferIntRangeCommon
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRIR
   )

mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp

 //===- InferIntRangeInterfaceImpls.cpp - Integer range impls for arith -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
+#include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
 
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 #define DEBUG_TYPE "int-range-analysis"
 
 using namespace mlir;
 using namespace mlir::arith;
-
-/// Function that evaluates the result of doing something on arithmetic
-/// constants and returns std::nullopt on overflow.
-using ConstArithFn =
-    function_ref<std::optional<APInt>(const APInt &, const APInt &)>;
-
-/// Return the maxmially wide signed or unsigned range for a given bitwidth.
-
-/// Compute op(minLeft, minRight) and op(maxLeft, maxRight) if possible,
-/// If either computation overflows, make the result unbounded.
-static ConstantIntRanges computeBoundsBy(ConstArithFn op, const APInt &minLeft,
-                                         const APInt &minRight,
-                                         const APInt &maxLeft,
-                                         const APInt &maxRight, bool isSigned) {
-  std::optional<APInt> maybeMin = op(minLeft, minRight);
-  std::optional<APInt> maybeMax = op(maxLeft, maxRight);
-  if (maybeMin && maybeMax)
-    return ConstantIntRanges::range(*maybeMin, *maybeMax, isSigned);
-  return ConstantIntRanges::maxRange(minLeft.getBitWidth());
-}
-
-/// Compute the minimum and maximum of `(op(l, r) for l in lhs for r in rhs)`,
-/// ignoring unbounded values. Returns the maximal range if `op` overflows.
-static ConstantIntRanges minMaxBy(ConstArithFn op, ArrayRef<APInt> lhs,
-                                  ArrayRef<APInt> rhs, bool isSigned) {
-  unsigned width = lhs[0].getBitWidth();
-  APInt min =
-      isSigned ? APInt::getSignedMaxValue(width) : APInt::getMaxValue(width);
-  APInt max =
-      isSigned ? APInt::getSignedMinValue(width) : APInt::getZero(width);
-  for (const APInt &left : lhs) {
-    for (const APInt &right : rhs) {
-      std::optional<APInt> maybeThisResult = op(left, right);
-      if (!maybeThisResult)
-        return ConstantIntRanges::maxRange(width);
-      APInt result = std::move(*maybeThisResult);
-      min = (isSigned ? result.slt(min) : result.ult(min)) ? result : min;
-      max = (isSigned ? result.sgt(max) : result.ugt(max)) ? result : max;
-    }
-  }
-  return ConstantIntRanges::range(min, max, isSigned);
-}
+using namespace mlir::intrange;
 
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
 void arith::ConstantOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                           SetIntRangeFn setResultRange) {
   auto constAttr = getValue().dyn_cast_or_null<IntegerAttr>();
   if (constAttr) {
     const APInt &value = constAttr.getValue();
     setResultRange(getResult(), ConstantIntRanges::constant(value));
   }
 }
 
 //===----------------------------------------------------------------------===//
 // AddIOp
 //===----------------------------------------------------------------------===//
 
 void arith::AddIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                       SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-  ConstArithFn uadd = [](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-    bool overflowed = false;
-    APInt result = a.uadd_ov(b, overflowed);
-    return overflowed ? std::optional<APInt>() : result;
-  };
-  ConstArithFn sadd = [](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-    bool overflowed = false;
-    APInt result = a.sadd_ov(b, overflowed);
-    return overflowed ? std::optional<APInt>() : result;
-  };
-
-  ConstantIntRanges urange = computeBoundsBy(
-      uadd, lhs.umin(), rhs.umin(), lhs.umax(), rhs.umax(), /*isSigned=*/false);
-  ConstantIntRanges srange = computeBoundsBy(
-      sadd, lhs.smin(), rhs.smin(), lhs.smax(), rhs.smax(), /*isSigned=*/true);
-  setResultRange(getResult(), urange.intersection(srange));
+  setResultRange(getResult(), inferAdd(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // SubIOp
 //===----------------------------------------------------------------------===//
 
 void arith::SubIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                       SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  ConstArithFn usub = [](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-    bool overflowed = false;
-    APInt result = a.usub_ov(b, overflowed);
-    return overflowed ? std::optional<APInt>() : result;
-  };
-  ConstArithFn ssub = [](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-    bool overflowed = false;
-    APInt result = a.ssub_ov(b, overflowed);
-    return overflowed ? std::optional<APInt>() : result;
-  };
-  ConstantIntRanges urange = computeBoundsBy(
-      usub, lhs.umin(), rhs.umax(), lhs.umax(), rhs.umin(), /*isSigned=*/false);
-  ConstantIntRanges srange = computeBoundsBy(
-      ssub, lhs.smin(), rhs.smax(), lhs.smax(), rhs.smin(), /*isSigned=*/true);
-  setResultRange(getResult(), urange.intersection(srange));
+  setResultRange(getResult(), inferSub(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // MulIOp
 //===----------------------------------------------------------------------===//
 
 void arith::MulIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                       SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  ConstArithFn umul = [](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-    bool overflowed = false;
-    APInt result = a.umul_ov(b, overflowed);
-    return overflowed ? std::optional<APInt>() : result;
-  };
-  ConstArithFn smul = [](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-    bool overflowed = false;
-    APInt result = a.smul_ov(b, overflowed);
-    return overflowed ? std::optional<APInt>() : result;
-  };
-
-  ConstantIntRanges urange =
-      minMaxBy(umul, {lhs.umin(), lhs.umax()}, {rhs.umin(), rhs.umax()},
-               /*isSigned=*/false);
-  ConstantIntRanges srange =
-      minMaxBy(smul, {lhs.smin(), lhs.smax()}, {rhs.smin(), rhs.smax()},
-               /*isSigned=*/true);
-
-  setResultRange(getResult(), urange.intersection(srange));
+  setResultRange(getResult(), inferMul(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // DivUIOp
 //===----------------------------------------------------------------------===//
 
-/// Fix up division results (ex. for ceiling and floor), returning an APInt
-/// if there has been no overflow
-using DivisionFixupFn = function_ref<std::optional<APInt>(
-    const APInt &lhs, const APInt &rhs, const APInt &result)>;
-
-static ConstantIntRanges inferDivUIRange(const ConstantIntRanges &lhs,
-                                         const ConstantIntRanges &rhs,
-                                         DivisionFixupFn fixup) {
-  const APInt &lhsMin = lhs.umin(), &lhsMax = lhs.umax(), &rhsMin = rhs.umin(),
-              &rhsMax = rhs.umax();
-
-  if (!rhsMin.isZero()) {
-    auto udiv = [&fixup](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-      return fixup(a, b, a.udiv(b));
-    };
-    return minMaxBy(udiv, {lhsMin, lhsMax}, {rhsMin, rhsMax},
-                    /*isSigned=*/false);
-  }
-  // Otherwise, it's possible we might divide by 0.
-  return ConstantIntRanges::maxRange(rhsMin.getBitWidth());
-}
-
 void arith::DivUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  setResultRange(getResult(),
-                 inferDivUIRange(argRanges[0], argRanges[1],
-                                 [](const APInt &lhs, const APInt &rhs,
-                                    const APInt &result) { return result; }));
+  setResultRange(getResult(), inferDivU(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // DivSIOp
 //===----------------------------------------------------------------------===//
 
-static ConstantIntRanges inferDivSIRange(const ConstantIntRanges &lhs,
-                                         const ConstantIntRanges &rhs,
-                                         DivisionFixupFn fixup) {
-  const APInt &lhsMin = lhs.smin(), &lhsMax = lhs.smax(), &rhsMin = rhs.smin(),
-              &rhsMax = rhs.smax();
-  bool canDivide = rhsMin.isStrictlyPositive() || rhsMax.isNegative();
-
-  if (canDivide) {
-    auto sdiv = [&fixup](const APInt &a,
-                         const APInt &b) -> std::optional<APInt> {
-      bool overflowed = false;
-      APInt result = a.sdiv_ov(b, overflowed);
-      return overflowed ? std::optional<APInt>() : fixup(a, b, result);
-    };
-    return minMaxBy(sdiv, {lhsMin, lhsMax}, {rhsMin, rhsMax},
-                    /*isSigned=*/true);
-  }
-  return ConstantIntRanges::maxRange(rhsMin.getBitWidth());
-}
-
 void arith::DivSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  setResultRange(getResult(),
-                 inferDivSIRange(argRanges[0], argRanges[1],
-                                 [](const APInt &lhs, const APInt &rhs,
-                                    const APInt &result) { return result; }));
+  setResultRange(getResult(), inferDivS(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // CeilDivUIOp
 //===----------------------------------------------------------------------===//
 
 void arith::CeilDivUIOp::inferResultRanges(
     ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  DivisionFixupFn ceilDivUIFix =
-      [](const APInt &lhs, const APInt &rhs,
-         const APInt &result) -> std::optional<APInt> {
-    if (!lhs.urem(rhs).isZero()) {
-      bool overflowed = false;
-      APInt corrected =
-          result.uadd_ov(APInt(result.getBitWidth(), 1), overflowed);
-      return overflowed ? std::optional<APInt>() : corrected;
-    }
-    return result;
-  };
-  setResultRange(getResult(), inferDivUIRange(lhs, rhs, ceilDivUIFix));
+  setResultRange(getResult(), inferCeilDivU(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // CeilDivSIOp
 //===----------------------------------------------------------------------===//
 
 void arith::CeilDivSIOp::inferResultRanges(
     ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  DivisionFixupFn ceilDivSIFix =
-      [](const APInt &lhs, const APInt &rhs,
-         const APInt &result) -> std::optional<APInt> {
-    if (!lhs.srem(rhs).isZero() && lhs.isNonNegative() == rhs.isNonNegative()) {
-      bool overflowed = false;
-      APInt corrected =
-          result.sadd_ov(APInt(result.getBitWidth(), 1), overflowed);
-      return overflowed ? std::optional<APInt>() : corrected;
-    }
-    return result;
-  };
-  setResultRange(getResult(), inferDivSIRange(lhs, rhs, ceilDivSIFix));
+  setResultRange(getResult(), inferCeilDivS(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // FloorDivSIOp
 //===----------------------------------------------------------------------===//
 
 void arith::FloorDivSIOp::inferResultRanges(
     ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  DivisionFixupFn floorDivSIFix =
-      [](const APInt &lhs, const APInt &rhs,
-         const APInt &result) -> std::optional<APInt> {
-    if (!lhs.srem(rhs).isZero() && lhs.isNonNegative() != rhs.isNonNegative()) {
-      bool overflowed = false;
-      APInt corrected =
-          result.ssub_ov(APInt(result.getBitWidth(), 1), overflowed);
-      return overflowed ? std::optional<APInt>() : corrected;
-    }
-    return result;
-  };
-  setResultRange(getResult(), inferDivSIRange(lhs, rhs, floorDivSIFix));
+  return setResultRange(getResult(), inferFloorDivS(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // RemUIOp
 //===----------------------------------------------------------------------===//
 
 void arith::RemUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-  const APInt &rhsMin = rhs.umin(), &rhsMax = rhs.umax();
-
-  unsigned width = rhsMin.getBitWidth();
-  APInt umin = APInt::getZero(width);
-  APInt umax = APInt::getMaxValue(width);
-
-  if (!rhsMin.isZero()) {
-    umax = rhsMax - 1;
-    // Special case: sweeping out a contiguous range in N/[modulus]
-    if (rhsMin == rhsMax) {
-      const APInt &lhsMin = lhs.umin(), &lhsMax = lhs.umax();
-      if ((lhsMax - lhsMin).ult(rhsMax)) {
-        APInt minRem = lhsMin.urem(rhsMax);
-        APInt maxRem = lhsMax.urem(rhsMax);
-        if (minRem.ule(maxRem)) {
-          umin = minRem;
-          umax = maxRem;
-        }
-      }
-    }
-  }
-  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(umin, umax));
+  setResultRange(getResult(), inferRemU(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // RemSIOp
 //===----------------------------------------------------------------------===//
 
 void arith::RemSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-  const APInt &lhsMin = lhs.smin(), &lhsMax = lhs.smax(), &rhsMin = rhs.smin(),
-              &rhsMax = rhs.smax();
-
-  unsigned width = rhsMax.getBitWidth();
-  APInt smin = APInt::getSignedMinValue(width);
-  APInt smax = APInt::getSignedMaxValue(width);
-  // No bounds if zero could be a divisor.
-  bool canBound = (rhsMin.isStrictlyPositive() || rhsMax.isNegative());
-  if (canBound) {
-    APInt maxDivisor = rhsMin.isStrictlyPositive() ? rhsMax : rhsMin.abs();
-    bool canNegativeDividend = lhsMin.isNegative();
-    bool canPositiveDividend = lhsMax.isStrictlyPositive();
-    APInt zero = APInt::getZero(maxDivisor.getBitWidth());
-    APInt maxPositiveResult = maxDivisor - 1;
-    APInt minNegativeResult = -maxPositiveResult;
-    smin = canNegativeDividend ? minNegativeResult : zero;
-    smax = canPositiveDividend ? maxPositiveResult : zero;
-    // Special case: sweeping out a contiguous range in N/[modulus].
-    if (rhsMin == rhsMax) {
-      if ((lhsMax - lhsMin).ult(maxDivisor)) {
-        APInt minRem = lhsMin.srem(maxDivisor);
-        APInt maxRem = lhsMax.srem(maxDivisor);
-        if (minRem.sle(maxRem)) {
-          smin = minRem;
-          smax = maxRem;
-        }
-      }
-    }
-  }
-  setResultRange(getResult(), ConstantIntRanges::fromSigned(smin, smax));
+  setResultRange(getResult(), inferRemS(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // AndIOp
 //===----------------------------------------------------------------------===//
 
-/// "Widen" bounds - if 0bvvvvv??? <= a <= 0bvvvvv???,
-/// relax the bounds to 0bvvvvv000 <= a <= 0bvvvvv111, where vvvvv are the bits
-/// that both bonuds have in common. This gives us a consertive approximation
-/// for what values can be passed to bitwise operations.
-static std::tuple<APInt, APInt>
-widenBitwiseBounds(const ConstantIntRanges &bound) {
-  APInt leftVal = bound.umin(), rightVal = bound.umax();
-  unsigned bitwidth = leftVal.getBitWidth();
-  unsigned differingBits = bitwidth - (leftVal ^ rightVal).countLeadingZeros();
-  leftVal.clearLowBits(differingBits);
-  rightVal.setLowBits(differingBits);
-  return std::make_tuple(std::move(leftVal), std::move(rightVal));
-}
-
 void arith::AndIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                       SetIntRangeFn setResultRange) {
-  auto [lhsZeros, lhsOnes] = widenBitwiseBounds(argRanges[0]);
-  auto [rhsZeros, rhsOnes] = widenBitwiseBounds(argRanges[1]);
-  auto andi = [](const APInt &a, const APInt &b) -> std::optional<APInt> {
-    return a & b;
-  };
-  setResultRange(getResult(),
-                 minMaxBy(andi, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
-                          /*isSigned=*/false));
+  setResultRange(getResult(), inferAnd(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // OrIOp
 //===----------------------------------------------------------------------===//
 
 void arith::OrIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                      SetIntRangeFn setResultRange) {
-  auto [lhsZeros, lhsOnes] = widenBitwiseBounds(argRanges[0]);
-  auto [rhsZeros, rhsOnes] = widenBitwiseBounds(argRanges[1]);
-  auto ori = [](const APInt &a, const APInt &b) -> std::optional<APInt> {
-    return a | b;
-  };
-  setResultRange(getResult(),
-                 minMaxBy(ori, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
-                          /*isSigned=*/false));
+  setResultRange(getResult(), inferOr(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // XOrIOp
 //===----------------------------------------------------------------------===//
 
 void arith::XOrIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                       SetIntRangeFn setResultRange) {
-  auto [lhsZeros, lhsOnes] = widenBitwiseBounds(argRanges[0]);
-  auto [rhsZeros, rhsOnes] = widenBitwiseBounds(argRanges[1]);
-  auto xori = [](const APInt &a, const APInt &b) -> std::optional<APInt> {
-    return a ^ b;
-  };
-  setResultRange(getResult(),
-                 minMaxBy(xori, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
-                          /*isSigned=*/false));
+  setResultRange(getResult(), inferXor(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // MaxSIOp
 //===----------------------------------------------------------------------===//
 
 void arith::MaxSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  const APInt &smin = lhs.smin().sgt(rhs.smin()) ? lhs.smin() : rhs.smin();
-  const APInt &smax = lhs.smax().sgt(rhs.smax()) ? lhs.smax() : rhs.smax();
-  setResultRange(getResult(), ConstantIntRanges::fromSigned(smin, smax));
+  setResultRange(getResult(), inferMaxS(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // MaxUIOp
 //===----------------------------------------------------------------------===//
 
 void arith::MaxUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  const APInt &umin = lhs.umin().ugt(rhs.umin()) ? lhs.umin() : rhs.umin();
-  const APInt &umax = lhs.umax().ugt(rhs.umax()) ? lhs.umax() : rhs.umax();
-  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(umin, umax));
+  setResultRange(getResult(), inferMaxU(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // MinSIOp
 //===----------------------------------------------------------------------===//
 
 void arith::MinSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  const APInt &smin = lhs.smin().slt(rhs.smin()) ? lhs.smin() : rhs.smin();
-  const APInt &smax = lhs.smax().slt(rhs.smax()) ? lhs.smax() : rhs.smax();
-  setResultRange(getResult(), ConstantIntRanges::fromSigned(smin, smax));
+  setResultRange(getResult(), inferMinS(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // MinUIOp
 //===----------------------------------------------------------------------===//
 
 void arith::MinUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  const APInt &umin = lhs.umin().ult(rhs.umin()) ? lhs.umin() : rhs.umin();
-  const APInt &umax = lhs.umax().ult(rhs.umax()) ? lhs.umax() : rhs.umax();
-  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(umin, umax));
+  setResultRange(getResult(), inferMinU(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // ExtUIOp
 //===----------------------------------------------------------------------===//
 
-static ConstantIntRanges extUIRange(const ConstantIntRanges &range,
-                                    Type destType) {
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-  APInt umin = range.umin().zext(destWidth);
-  APInt umax = range.umax().zext(destWidth);
-  return ConstantIntRanges::fromUnsigned(umin, umax);
-}
-
 void arith::ExtUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  Type destType = getResult().getType();
-  setResultRange(getResult(), extUIRange(argRanges[0], destType));
+  unsigned destWidth =
+      ConstantIntRanges::getStorageBitwidth(getResult().getType());
+  setResultRange(getResult(), extUIRange(argRanges[0], destWidth));
 }
 
 //===----------------------------------------------------------------------===//
 // ExtSIOp
 //===----------------------------------------------------------------------===//
 
-static ConstantIntRanges extSIRange(const ConstantIntRanges &range,
-                                    Type destType) {
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-  APInt smin = range.smin().sext(destWidth);
-  APInt smax = range.smax().sext(destWidth);
-  return ConstantIntRanges::fromSigned(smin, smax);
-}
-
 void arith::ExtSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  Type destType = getResult().getType();
-  setResultRange(getResult(), extSIRange(argRanges[0], destType));
+  unsigned destWidth =
+      ConstantIntRanges::getStorageBitwidth(getResult().getType());
+  setResultRange(getResult(), extSIRange(argRanges[0], destWidth));
 }
 
 //===----------------------------------------------------------------------===//
 // TruncIOp
 //===----------------------------------------------------------------------===//
 
-static ConstantIntRanges truncIRange(const ConstantIntRanges &range,
-                                     Type destType) {
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-  // If you truncate the first four bytes in [0xaaaabbbb, 0xccccbbbb],
-  // the range of the resulting value is not contiguous ind includes 0.
-  // Ex. If you truncate [256, 258] from i16 to i8, you validly get [0, 2],
-  // but you can't truncate [255, 257] similarly.
-  bool hasUnsignedRollover =
-      range.umin().lshr(destWidth) != range.umax().lshr(destWidth);
-  APInt umin = hasUnsignedRollover ? APInt::getZero(destWidth)
-                                   : range.umin().trunc(destWidth);
-  APInt umax = hasUnsignedRollover ? APInt::getMaxValue(destWidth)
-                                   : range.umax().trunc(destWidth);
-
-  // Signed post-truncation rollover will not occur when either:
-  // - The high parts of the min and max, plus the sign bit, are the same
-  // - The high halves + sign bit of the min and max are either all 1s or all 0s
-  //  and you won't create a [positive, negative] range by truncating.
-  // For example, you can truncate the ranges [256, 258]_i16 to [0, 2]_i8
-  // but not [255, 257]_i16 to a range of i8s. You can also truncate
-  // [-256, -256]_i16 to [-2, 0]_i8, but not [-257, -255]_i16.
-  // You can also truncate [-130, 0]_i16 to i8 because -130_i16 (0xff7e)
-  // will truncate to 0x7e, which is greater than 0
-  APInt sminHighPart = range.smin().ashr(destWidth - 1);
-  APInt smaxHighPart = range.smax().ashr(destWidth - 1);
-  bool hasSignedOverflow =
-      (sminHighPart != smaxHighPart) &&
-      !(sminHighPart.isAllOnes() &&
-        (smaxHighPart.isAllOnes() || smaxHighPart.isZero())) &&
-      !(sminHighPart.isZero() && smaxHighPart.isZero());
-  APInt smin = hasSignedOverflow ? APInt::getSignedMinValue(destWidth)
-                                 : range.smin().trunc(destWidth);
-  APInt smax = hasSignedOverflow ? APInt::getSignedMaxValue(destWidth)
-                                 : range.smax().trunc(destWidth);
-  return {umin, umax, smin, smax};
-}
-
 void arith::TruncIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                         SetIntRangeFn setResultRange) {
-  Type destType = getResult().getType();
-  setResultRange(getResult(), truncIRange(argRanges[0], destType));
+  unsigned destWidth =
+      ConstantIntRanges::getStorageBitwidth(getResult().getType());
+  setResultRange(getResult(), truncRange(argRanges[0], destWidth));
 }
 
 //===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
 void arith::IndexCastOp::inferResultRanges(
     ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
   Type sourceType = getOperand().getType();
   Type destType = getResult().getType();
   unsigned srcWidth = ConstantIntRanges::getStorageBitwidth(sourceType);
   unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
 
   if (srcWidth < destWidth)
-    setResultRange(getResult(), extSIRange(argRanges[0], destType));
+    setResultRange(getResult(), extSIRange(argRanges[0], destWidth));
   else if (srcWidth > destWidth)
-    setResultRange(getResult(), truncIRange(argRanges[0], destType));
+    setResultRange(getResult(), truncRange(argRanges[0], destWidth));
   else
     setResultRange(getResult(), argRanges[0]);
 }
 
 //===----------------------------------------------------------------------===//
 // IndexCastUIOp
 //===----------------------------------------------------------------------===//
 
 void arith::IndexCastUIOp::inferResultRanges(
     ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
   Type sourceType = getOperand().getType();
   Type destType = getResult().getType();
   unsigned srcWidth = ConstantIntRanges::getStorageBitwidth(sourceType);
   unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
 
   if (srcWidth < destWidth)
-    setResultRange(getResult(), extUIRange(argRanges[0], destType));
+    setResultRange(getResult(), extUIRange(argRanges[0], destWidth));
   else if (srcWidth > destWidth)
-    setResultRange(getResult(), truncIRange(argRanges[0], destType));
+    setResultRange(getResult(), truncRange(argRanges[0], destWidth));
   else
     setResultRange(getResult(), argRanges[0]);
 }
 
 //===----------------------------------------------------------------------===//
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
-bool isStaticallyTrue(arith::CmpIPredicate pred, const ConstantIntRanges &lhs,
-                      const ConstantIntRanges &rhs) {
-  switch (pred) {
-  case arith::CmpIPredicate::sle:
-  case arith::CmpIPredicate::slt:
-    return (applyCmpPredicate(pred, lhs.smax(), rhs.smin()));
-  case arith::CmpIPredicate::ule:
-  case arith::CmpIPredicate::ult:
-    return applyCmpPredicate(pred, lhs.umax(), rhs.umin());
-  case arith::CmpIPredicate::sge:
-  case arith::CmpIPredicate::sgt:
-    return applyCmpPredicate(pred, lhs.smin(), rhs.smax());
-  case arith::CmpIPredicate::uge:
-  case arith::CmpIPredicate::ugt:
-    return applyCmpPredicate(pred, lhs.umin(), rhs.umax());
-  case arith::CmpIPredicate::eq: {
-    std::optional<APInt> lhsConst = lhs.getConstantValue();
-    std::optional<APInt> rhsConst = rhs.getConstantValue();
-    return lhsConst && rhsConst && lhsConst == rhsConst;
-  }
-  case arith::CmpIPredicate::ne: {
-    // While equality requires that there is an interpration of the preceeding
-    // computations that produces equal constants, whether that be signed or
-    // unsigned, statically determining inequality requires that neither
-    // interpretation produce potentially overlapping ranges.
-    bool sne = isStaticallyTrue(CmpIPredicate::slt, lhs, rhs) ||
-               isStaticallyTrue(CmpIPredicate::sgt, lhs, rhs);
-    bool une = isStaticallyTrue(CmpIPredicate::ult, lhs, rhs) ||
-               isStaticallyTrue(CmpIPredicate::ugt, lhs, rhs);
-    return sne && une;
-  }
-  }
-  return false;
-}
-
 void arith::CmpIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                       SetIntRangeFn setResultRange) {
-  arith::CmpIPredicate pred = getPredicate();
+  arith::CmpIPredicate arithPred = getPredicate();
+  intrange::CmpPredicate pred = static_cast<intrange::CmpPredicate>(arithPred);
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
 
   APInt min = APInt::getZero(1);
   APInt max = APInt::getAllOnesValue(1);
-  if (isStaticallyTrue(pred, lhs, rhs))
+
+  Optional<bool> truthValue = intrange::evaluatePred(pred, lhs, rhs);
+  if (truthValue.has_value() && *truthValue)
     min = max;
-  else if (isStaticallyTrue(invertPredicate(pred), lhs, rhs))
+  else if (truthValue.has_value() && !(*truthValue))
     max = min;
 
   setResultRange(getResult(), ConstantIntRanges::fromUnsigned(min, max));
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//

 }
 
 //===----------------------------------------------------------------------===//
 // ShLIOp
 //===----------------------------------------------------------------------===//
 
 void arith::ShLIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                       SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-  ConstArithFn shl = [](const APInt &l,
-                        const APInt &r) -> std::optional<APInt> {
-    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.shl(r);
-  };
-  ConstantIntRanges urange =
-      minMaxBy(shl, {lhs.umin(), lhs.umax()}, {rhs.umin(), rhs.umax()},
-               /*isSigned=*/false);
-  ConstantIntRanges srange =
-      minMaxBy(shl, {lhs.smin(), lhs.smax()}, {rhs.umin(), rhs.umax()},
-               /*isSigned=*/true);
-  setResultRange(getResult(), urange.intersection(srange));
+  setResultRange(getResult(), inferShl(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // ShRUIOp
 //===----------------------------------------------------------------------===//
 
 void arith::ShRUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  ConstArithFn lshr = [](const APInt &l,
-                         const APInt &r) -> std::optional<APInt> {
-    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.lshr(r);
-  };
-  setResultRange(getResult(), minMaxBy(lshr, {lhs.umin(), lhs.umax()},
-                                       {rhs.umin(), rhs.umax()},
-                                       /*isSigned=*/false));
+  setResultRange(getResult(), inferShrU(argRanges));
 }
 
 //===----------------------------------------------------------------------===//
 // ShRSIOp
 //===----------------------------------------------------------------------===//
 
 void arith::ShRSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  ConstArithFn ashr = [](const APInt &l,
-                         const APInt &r) -> std::optional<APInt> {
-    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.ashr(r);
-  };
-
-  setResultRange(getResult(),
-                 minMaxBy(ashr, {lhs.smin(), lhs.smax()},
-                          {rhs.umin(), rhs.umax()}, /*isSigned=*/true));
+  setResultRange(getResult(), inferShrS(argRanges));
 }

mlir/lib/Dialect/Index/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRIndexDialect
   IndexAttrs.cpp
   IndexDialect.cpp
   IndexOps.cpp
+  InferIntRangeInterfaceImpls.cpp
 
   DEPENDS
   MLIRIndexOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
   MLIRIR
   MLIRCastInterfaces
+  MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   )

mlir/lib/Dialect/Index/IR/InferIntRangeInterfaceImpls.cpp

+//===- InferIntRangeInterfaceImpls.cpp - Integer range impls for arith -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Index/IR/IndexOps.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+#include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
+
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "int-range-analysis"
+
+using namespace mlir;
+using namespace mlir::index;
+using namespace mlir::intrange;
+
+//===----------------------------------------------------------------------===//
+// Constants
+//===----------------------------------------------------------------------===//
+
+void ConstantOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                   SetIntRangeFn setResultRange) {
+  const APInt &value = getValue();
+  setResultRange(getResult(), ConstantIntRanges::constant(value));
+}
+
+void BoolConstantOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                       SetIntRangeFn setResultRange) {
+  bool value = getValue();
+  APInt asInt(/*numBits=*/1, value);
+  setResultRange(getResult(), ConstantIntRanges::constant(asInt));
+}
+
+//===----------------------------------------------------------------------===//
+// Arithmec operations. All of these operations will have their results inferred
+// using both the 64-bit values and truncated 32-bit values of their inputs,
+// with the results being the union of those inferences, except where the
+// truncation of the 64-bit result is equal to the 32-bit result (at which time
+// we take the 64-bit result).
+//===----------------------------------------------------------------------===//
+
+void AddOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                              SetIntRangeFn setResultRange) {
+  setResultRange(getResult(), inferIndexOp(inferAdd, argRanges, CmpMode::Both));
+}
+
+void SubOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                              SetIntRangeFn setResultRange) {
+  setResultRange(getResult(), inferIndexOp(inferSub, argRanges, CmpMode::Both));
+}
+
+void MulOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                              SetIntRangeFn setResultRange) {
+  setResultRange(getResult(), inferIndexOp(inferMul, argRanges, CmpMode::Both));
+}
+
+void DivUOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferDivU, argRanges, CmpMode::Unsigned));
+}
+
+void DivSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferDivS, argRanges, CmpMode::Signed));
+}
+
+void CeilDivUOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                   SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferCeilDivU, argRanges, CmpMode::Unsigned));
+}
+
+void CeilDivSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                   SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferCeilDivS, argRanges, CmpMode::Signed));
+}
+
+void FloorDivSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                    SetIntRangeFn setResultRange) {
+  return setResultRange(
+      getResult(), inferIndexOp(inferFloorDivS, argRanges, CmpMode::Signed));
+}
+
+void RemSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferRemS, argRanges, CmpMode::Signed));
+}
+
+void RemUOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferRemU, argRanges, CmpMode::Unsigned));
+}
+
+void MaxSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferMaxS, argRanges, CmpMode::Signed));
+}
+
+void MaxUOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferMaxU, argRanges, CmpMode::Unsigned));
+}
+
+void MinSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferMinS, argRanges, CmpMode::Signed));
+}
+
+void MinUOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferMinU, argRanges, CmpMode::Unsigned));
+}
+
+void ShlOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                              SetIntRangeFn setResultRange) {
+  setResultRange(getResult(), inferIndexOp(inferShl, argRanges, CmpMode::Both));
+}
+
+void ShrSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferShrS, argRanges, CmpMode::Signed));
+}
+
+void ShrUOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                               SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferShrU, argRanges, CmpMode::Unsigned));
+}
+
+void AndOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                              SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferAnd, argRanges, CmpMode::Unsigned));
+}
+
+void OrOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                             SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferOr, argRanges, CmpMode::Unsigned));
+}
+
+void XOrOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                              SetIntRangeFn setResultRange) {
+  setResultRange(getResult(),
+                 inferIndexOp(inferXor, argRanges, CmpMode::Unsigned));
+}
+
+//===----------------------------------------------------------------------===//
+// Casts
+//===----------------------------------------------------------------------===//
+
+static ConstantIntRanges makeLikeDest(const ConstantIntRanges &range,
+                                      unsigned srcWidth, unsigned destWidth,
+                                      bool isSigned) {
+  if (srcWidth < destWidth)
+    return isSigned ? extSIRange(range, destWidth)
+                    : extUIRange(range, destWidth);
+  if (srcWidth > destWidth)
+    return truncRange(range, destWidth);
+  return range;
+}
+
+// When casting to `index`, we will take the union of the possible fixed-width
+// casts.
+static ConstantIntRanges inferIndexCast(const ConstantIntRanges &range,
+                                        Type sourceType, Type destType,
+                                        bool isSigned) {
+  unsigned srcWidth = ConstantIntRanges::getStorageBitwidth(sourceType);
+  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
+  if (sourceType.isIndex())
+    return makeLikeDest(range, srcWidth, destWidth, isSigned);
+  // We are casting to indexs, so use the union of the 32-bit and 64-bit casts
+  ConstantIntRanges storageRange =
+      makeLikeDest(range, srcWidth, destWidth, isSigned);
+  ConstantIntRanges minWidthRange =
+      makeLikeDest(range, srcWidth, indexMinWidth, isSigned);
+  ConstantIntRanges minWidthExt = extRange(minWidthRange, destWidth);
+  ConstantIntRanges ret = storageRange.rangeUnion(minWidthExt);
+  return ret;
+}
+
+void CastSOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                SetIntRangeFn setResultRange) {
+  Type sourceType = getOperand().getType();
+  Type destType = getResult().getType();
+  setResultRange(getResult(), inferIndexCast(argRanges[0], sourceType, destType,
+                                             /*isSigned=*/true));
+}
+
+void CastUOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                SetIntRangeFn setResultRange) {
+  Type sourceType = getOperand().getType();
+  Type destType = getResult().getType();
+  setResultRange(getResult(), inferIndexCast(argRanges[0], sourceType, destType,
+                                             /*isSigned=*/false));
+}
+
+//===----------------------------------------------------------------------===//
+// CmpOp
+//===----------------------------------------------------------------------===//
+
+void CmpOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                              SetIntRangeFn setResultRange) {
+  index::IndexCmpPredicate indexPred = getPred();
+  intrange::CmpPredicate pred = static_cast<intrange::CmpPredicate>(indexPred);
+  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
+
+  APInt min = APInt::getZero(1);
+  APInt max = APInt::getAllOnesValue(1);
+
+  Optional<bool> truthValue64 = intrange::evaluatePred(pred, lhs, rhs);
+
+  ConstantIntRanges lhsTrunc = truncRange(lhs, indexMinWidth),
+                    rhsTrunc = truncRange(rhs, indexMinWidth);
+  Optional<bool> truthValue32 =
+      intrange::evaluatePred(pred, lhsTrunc, rhsTrunc);
+
+  if (truthValue64 == truthValue32) {
+    if (truthValue64.has_value() && *truthValue64)
+      min = max;
+    else if (truthValue64.has_value() && !(*truthValue64))
+      max = min;
+  }
+  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(min, max));
+}
+
+//===----------------------------------------------------------------------===//
+// SizeOf, which is bounded between the two supported bitwidth (32 and 64).
+//===----------------------------------------------------------------------===//
+
+void SizeOfOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                 SetIntRangeFn setResultRange) {
+  unsigned storageWidth =
+      ConstantIntRanges::getStorageBitwidth(getResult().getType());
+  APInt min(/*numBits=*/storageWidth, indexMinWidth);
+  APInt max(/*numBits=*/storageWidth, indexMaxWidth);
+  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(min, max));
+}

mlir/lib/Interfaces/CMakeLists.txt

 add_mlir_interface_library(LoopLikeInterface)
 add_mlir_interface_library(ParallelCombiningOpInterface)
 add_mlir_interface_library(RuntimeVerifiableOpInterface)
 add_mlir_interface_library(ShapedOpInterfaces)
 add_mlir_interface_library(SideEffectInterfaces)
 add_mlir_interface_library(TilingInterface)
 add_mlir_interface_library(VectorInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
+
+add_subdirectory(Utils)

mlir/lib/Interfaces/Utils/CMakeLists.txt

+add_mlir_library(MLIRInferIntRangeCommon
+    InferIntRangeCommon.cpp
+
+    ADDITIONAL_HEADER_DIRS
+    ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces/Utils
+
+    DEPENDS
+    MLIRInferIntRangeInterfaceIncGen
+
+    LINK_LIBS PUBLIC
+    MLIRInferIntRangeInterface
+    MLIRIR
+)

mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp

-//===- InferIntRangeInterfaceImpls.cpp - Integer range impls for arith -===//
+//===- InferIntRangeCommon.cpp - Inference for common ops ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+//
+// This file contains implementations of range inference for operations that are
+// common to both the `arith` and `index` dialects to facilitate reuse.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
 
-#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+
 #include "llvm/Support/Debug.h"
+
+#include <iterator>
 #include <optional>
 
+using namespace mlir;
+
 #define DEBUG_TYPE "int-range-analysis"
 
-using namespace mlir;
-using namespace mlir::arith;
+//===----------------------------------------------------------------------===//
+// General utilities
+//===----------------------------------------------------------------------===//
 
 /// Function that evaluates the result of doing something on arithmetic
 /// constants and returns std::nullopt on overflow.
 using ConstArithFn =
     function_ref<std::optional<APInt>(const APInt &, const APInt &)>;
 
-/// Return the maxmially wide signed or unsigned range for a given bitwidth.
-
 /// Compute op(minLeft, minRight) and op(maxLeft, maxRight) if possible,
 /// If either computation overflows, make the result unbounded.
 static ConstantIntRanges computeBoundsBy(ConstArithFn op, const APInt &minLeft,
                                          const APInt &minRight,
                                          const APInt &maxLeft,
                                          const APInt &maxRight, bool isSigned) {
   std::optional<APInt> maybeMin = op(minLeft, minRight);
   std::optional<APInt> maybeMax = op(maxLeft, maxRight);
for (const APInt &right : rhs) {
       min = (isSigned ? result.slt(min) : result.ult(min)) ? result : min;
       max = (isSigned ? result.sgt(max) : result.ugt(max)) ? result : max;
     }
   }
   return ConstantIntRanges::range(min, max, isSigned);
 }
 
 //===----------------------------------------------------------------------===//
-// ConstantOp
+// Ext, trunc, index op handling
 //===----------------------------------------------------------------------===//
 
-void arith::ConstantOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                          SetIntRangeFn setResultRange) {
-  auto constAttr = getValue().dyn_cast_or_null<IntegerAttr>();
-  if (constAttr) {
-    const APInt &value = constAttr.getValue();
-    setResultRange(getResult(), ConstantIntRanges::constant(value));
-  }
+ConstantIntRanges
+mlir::intrange::inferIndexOp(InferRangeFn inferFn,
+                             ArrayRef<ConstantIntRanges> argRanges,
+                             intrange::CmpMode mode) {
+  ConstantIntRanges sixtyFour = inferFn(argRanges);
+  SmallVector<ConstantIntRanges, 2> truncated;
+  llvm::transform(argRanges, std::back_inserter(truncated),
+                  [](const ConstantIntRanges &range) {
+                    return truncRange(range, /*destWidth=*/indexMinWidth);
+                  });
+  ConstantIntRanges thirtyTwo = inferFn(truncated);
+  ConstantIntRanges thirtyTwoAsSixtyFour =
+      extRange(thirtyTwo, /*destWidth=*/indexMaxWidth);
+  ConstantIntRanges sixtyFourAsThirtyTwo =
+      truncRange(sixtyFour, /*destWidth=*/indexMinWidth);
+
+  LLVM_DEBUG(llvm::dbgs() << "Index handling: 64-bit result = " << sixtyFour
+                          << " 32-bit = " << thirtyTwo << "\n");
+  bool truncEqual = false;
+  switch (mode) {
+  case intrange::CmpMode::Both:
+    truncEqual = (thirtyTwo == sixtyFourAsThirtyTwo);
+    break;
+  case intrange::CmpMode::Signed:
+    truncEqual = (thirtyTwo.smin() == sixtyFourAsThirtyTwo.smin() &&
+                  thirtyTwo.smax() == sixtyFourAsThirtyTwo.smax());
+    break;
+  case intrange::CmpMode::Unsigned:
+    truncEqual = (thirtyTwo.umin() == sixtyFourAsThirtyTwo.umin() &&
+                  thirtyTwo.umax() == sixtyFourAsThirtyTwo.umax());
+    break;
+  }
+  if (truncEqual)
+    // Returing the 64-bit result preserves more information.
+    return sixtyFour;
+  ConstantIntRanges merged = sixtyFour.rangeUnion(thirtyTwoAsSixtyFour);
+  return merged;
+}
+
+ConstantIntRanges mlir::intrange::extRange(const ConstantIntRanges &range,
+                                           unsigned int destWidth) {
+  APInt umin = range.umin().zext(destWidth);
+  APInt umax = range.umax().zext(destWidth);
+  APInt smin = range.smin().sext(destWidth);
+  APInt smax = range.smax().sext(destWidth);
+  return {umin, umax, smin, smax};
+}
+
+ConstantIntRanges mlir::intrange::extUIRange(const ConstantIntRanges &range,
+                                             unsigned destWidth) {
+  APInt umin = range.umin().zext(destWidth);
+  APInt umax = range.umax().zext(destWidth);
+  return ConstantIntRanges::fromUnsigned(umin, umax);
+}
+
+ConstantIntRanges mlir::intrange::extSIRange(const ConstantIntRanges &range,
+                                             unsigned destWidth) {
+  APInt smin = range.smin().sext(destWidth);
+  APInt smax = range.smax().sext(destWidth);
+  return ConstantIntRanges::fromSigned(smin, smax);
+
+ConstantIntRanges mlir::intrange::truncRange(const ConstantIntRanges &range,
+                                             unsigned int destWidth) {
+  // If you truncate the first four bytes in [0xaaaabbbb, 0xccccbbbb],
+  // the range of the resulting value is not contiguous ind includes 0.
+  // Ex. If you truncate [256, 258] from i16 to i8, you validly get [0, 2],
+  // but you can't truncate [255, 257] similarly.
+  bool hasUnsignedRollover =
+      range.umin().lshr(destWidth) != range.umax().lshr(destWidth);
+  APInt umin = hasUnsignedRollover ? APInt::getZero(destWidth)
+                                   : range.umin().trunc(destWidth);
+  APInt umax = hasUnsignedRollover ? APInt::getMaxValue(destWidth)
+                                   : range.umax().trunc(destWidth);
+
+  // Signed post-truncation rollover will not occur when either:
+  // - The high parts of the min and max, plus the sign bit, are the same
+  // - The high halves + sign bit of the min and max are either all 1s or all 0s
+  //  and you won't create a [positive, negative] range by truncating.
+  // For example, you can truncate the ranges [256, 258]_i16 to [0, 2]_i8
+  // but not [255, 257]_i16 to a range of i8s. You can also truncate
+  // [-256, -256]_i16 to [-2, 0]_i8, but not [-257, -255]_i16.
+  // You can also truncate [-130, 0]_i16 to i8 because -130_i16 (0xff7e)
+  // will truncate to 0x7e, which is greater than 0
+  APInt sminHighPart = range.smin().ashr(destWidth - 1);
+  APInt smaxHighPart = range.smax().ashr(destWidth - 1);
+  bool hasSignedOverflow =
+      (sminHighPart != smaxHighPart) &&
+      !(sminHighPart.isAllOnes() &&
+        (smaxHighPart.isAllOnes() || smaxHighPart.isZero())) &&
+      !(sminHighPart.isZero() && smaxHighPart.isZero());
+  APInt smin = hasSignedOverflow ? APInt::getSignedMinValue(destWidth)
+                                 : range.smin().trunc(destWidth);
+  APInt smax = hasSignedOverflow ? APInt::getSignedMaxValue(destWidth)
+                                 : range.smax().trunc(destWidth);
+  return {umin, umax, smin, smax};
 }
 
 //===----------------------------------------------------------------------===//
-// AddIOp
+// Addition
 //===----------------------------------------------------------------------===//
 
-void arith::AddIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                      SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferAdd(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
   ConstArithFn uadd = [](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
     bool overflowed = false;
     APInt result = a.uadd_ov(b, overflowed);
     return overflowed ? std::optional<APInt>() : result;
   };
   ConstArithFn sadd = [](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
     bool overflowed = false;
     APInt result = a.sadd_ov(b, overflowed);
     return overflowed ? std::optional<APInt>() : result;
   };
 
   ConstantIntRanges urange = computeBoundsBy(
       uadd, lhs.umin(), rhs.umin(), lhs.umax(), rhs.umax(), /*isSigned=*/false);
   ConstantIntRanges srange = computeBoundsBy(
       sadd, lhs.smin(), rhs.smin(), lhs.smax(), rhs.smax(), /*isSigned=*/true);
-  setResultRange(getResult(), urange.intersection(srange));
+  return urange.intersection(srange);
 }
 
 //===----------------------------------------------------------------------===//
-// SubIOp
+// Subtraction
 //===----------------------------------------------------------------------===//
 
-void arith::SubIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                      SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferSub(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
 
   ConstArithFn usub = [](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
     bool overflowed = false;
     APInt result = a.usub_ov(b, overflowed);
     return overflowed ? std::optional<APInt>() : result;
   };
   ConstArithFn ssub = [](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
     bool overflowed = false;
     APInt result = a.ssub_ov(b, overflowed);
     return overflowed ? std::optional<APInt>() : result;
   };
   ConstantIntRanges urange = computeBoundsBy(
       usub, lhs.umin(), rhs.umax(), lhs.umax(), rhs.umin(), /*isSigned=*/false);
   ConstantIntRanges srange = computeBoundsBy(
       ssub, lhs.smin(), rhs.smax(), lhs.smax(), rhs.smin(), /*isSigned=*/true);
-  setResultRange(getResult(), urange.intersection(srange));
+  return urange.intersection(srange);
 }
 
 //===----------------------------------------------------------------------===//
-// MulIOp
+// Multiplication
 //===----------------------------------------------------------------------===//
 
-void arith::MulIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                      SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferMul(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
 
   ConstArithFn umul = [](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
     bool overflowed = false;
     APInt result = a.umul_ov(b, overflowed);
     return overflowed ? std::optional<APInt>() : result;
   };
   ConstArithFn smul = [](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
     bool overflowed = false;
     APInt result = a.smul_ov(b, overflowed);
     return overflowed ? std::optional<APInt>() : result;
   };
 
   ConstantIntRanges urange =
       minMaxBy(umul, {lhs.umin(), lhs.umax()}, {rhs.umin(), rhs.umax()},
                /*isSigned=*/false);
   ConstantIntRanges srange =
       minMaxBy(smul, {lhs.smin(), lhs.smax()}, {rhs.smin(), rhs.smax()},
                /*isSigned=*/true);
-
-  setResultRange(getResult(), urange.intersection(srange));
+  return urange.intersection(srange);
 }
 
 //===----------------------------------------------------------------------===//
-// DivUIOp
+// DivU, CeilDivU (Unsigned division)
 //===----------------------------------------------------------------------===//
 
 /// Fix up division results (ex. for ceiling and floor), returning an APInt
 /// if there has been no overflow
 using DivisionFixupFn = function_ref<std::optional<APInt>(
     const APInt &lhs, const APInt &rhs, const APInt &result)>;
 
-static ConstantIntRanges inferDivUIRange(const ConstantIntRanges &lhs,
-                                         const ConstantIntRanges &rhs,
-                                         DivisionFixupFn fixup) {
+static ConstantIntRanges inferDivURange(const ConstantIntRanges &lhs,
   const APInt &lhsMin = lhs.umin(), &lhsMax = lhs.umax(), &rhsMin = rhs.umin(),
               &rhsMax = rhs.umax();
 
   if (!rhsMin.isZero()) {
     auto udiv = [&fixup](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
       return fixup(a, b, a.udiv(b));
     };
     return minMaxBy(udiv, {lhsMin, lhsMax}, {rhsMin, rhsMax},
                     /*isSigned=*/false);
   }
   // Otherwise, it's possible we might divide by 0.
   return ConstantIntRanges::maxRange(rhsMin.getBitWidth());
 }
 
-void arith::DivUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
-  setResultRange(getResult(),
-                 inferDivUIRange(argRanges[0], argRanges[1],
-                                 [](const APInt &lhs, const APInt &rhs,
-                                    const APInt &result) { return result; }));
+ConstantIntRanges
+mlir::intrange::inferDivU(ArrayRef<ConstantIntRanges> argRanges) {
+  return inferDivURange(argRanges[0], argRanges[1],
+                        [](const APInt &lhs, const APInt &rhs,
+                           const APInt &result) { return result; });
+}
+
+ConstantIntRanges
+mlir::intrange::inferCeilDivU(ArrayRef<ConstantIntRanges> argRanges) {
+  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
+
+  DivisionFixupFn ceilDivUIFix =
+         const APInt &result) -> std::optional<APInt> {
+    if (!lhs.urem(rhs).isZero()) {
+      bool overflowed = false;
+      APInt corrected =
+          result.uadd_ov(APInt(result.getBitWidth(), 1), overflowed);
+      return overflowed ? std::optional<APInt>() : corrected;
+    }
+    return result;
+  };
+  return inferDivURange(lhs, rhs, ceilDivUIFix);
 }
 
 //===----------------------------------------------------------------------===//
-// DivSIOp
+// DivS, CeilDivS, FloorDivS (Signed division)
 //===----------------------------------------------------------------------===//
 
-static ConstantIntRanges inferDivSIRange(const ConstantIntRanges &lhs,
-                                         const ConstantIntRanges &rhs,
-                                         DivisionFixupFn fixup) {
+static ConstantIntRanges inferDivSRange(const ConstantIntRanges &lhs,
   const APInt &lhsMin = lhs.smin(), &lhsMax = lhs.smax(), &rhsMin = rhs.smin(),
               &rhsMax = rhs.smax();
   bool canDivide = rhsMin.isStrictlyPositive() || rhsMax.isNegative();
 
   if (canDivide) {
     auto sdiv = [&fixup](const APInt &a,
                          const APInt &b) -> std::optional<APInt> {
       bool overflowed = false;
       APInt result = a.sdiv_ov(b, overflowed);
       return overflowed ? std::optional<APInt>() : fixup(a, b, result);
     };
     return minMaxBy(sdiv, {lhsMin, lhsMax}, {rhsMin, rhsMax},
                     /*isSigned=*/true);
   }
   return ConstantIntRanges::maxRange(rhsMin.getBitWidth());
 }
 
-void arith::DivSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
-  setResultRange(getResult(),
-                 inferDivSIRange(argRanges[0], argRanges[1],
-                                 [](const APInt &lhs, const APInt &rhs,
-                                    const APInt &result) { return result; }));
+ConstantIntRanges
+mlir::intrange::inferDivS(ArrayRef<ConstantIntRanges> argRanges) {
+  return inferDivSRange(argRanges[0], argRanges[1],
+                           const APInt &result) { return result; });
 }
 
-//===----------------------------------------------------------------------===//
-// CeilDivUIOp
-//===----------------------------------------------------------------------===//
-
-void arith::CeilDivUIOp::inferResultRanges(
-    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferCeilDivS(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
 
-  DivisionFixupFn ceilDivUIFix =
+  DivisionFixupFn ceilDivSIFix =
       [](const APInt &lhs, const APInt &rhs,
          const APInt &result) -> std::optional<APInt> {
-    if (!lhs.urem(rhs).isZero()) {
+    if (!lhs.srem(rhs).isZero() && lhs.isNonNegative() == rhs.isNonNegative()) {
       bool overflowed = false;
       APInt corrected =
-          result.uadd_ov(APInt(result.getBitWidth(), 1), overflowed);
+          result.sadd_ov(APInt(result.getBitWidth(), 1), overflowed);
       return overflowed ? std::optional<APInt>() : corrected;
     }
     return result;
   };
-  setResultRange(getResult(), inferDivUIRange(lhs, rhs, ceilDivUIFix));
+  return inferDivSRange(lhs, rhs, ceilDivSIFix);
 }
 
-//===----------------------------------------------------------------------===//
-// CeilDivSIOp
-//===----------------------------------------------------------------------===//
-
-void arith::CeilDivSIOp::inferResultRanges(
-    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferFloorDivS(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
 
-  DivisionFixupFn ceilDivSIFix =
+  DivisionFixupFn floorDivSIFix =
       [](const APInt &lhs, const APInt &rhs,
          const APInt &result) -> std::optional<APInt> {
-    if (!lhs.srem(rhs).isZero() && lhs.isNonNegative() == rhs.isNonNegative()) {
+    if (!lhs.srem(rhs).isZero() && lhs.isNonNegative() != rhs.isNonNegative()) {
       bool overflowed = false;
       APInt corrected =
-          result.sadd_ov(APInt(result.getBitWidth(), 1), overflowed);
+          result.ssub_ov(APInt(result.getBitWidth(), 1), overflowed);
       return overflowed ? std::optional<APInt>() : corrected;
     }
     return result;
   };
-  setResultRange(getResult(), inferDivSIRange(lhs, rhs, ceilDivSIFix));
+  return inferDivSRange(lhs, rhs, floorDivSIFix);
 }
 
 //===----------------------------------------------------------------------===//
-// FloorDivSIOp
+// Signed remainder (RemS)
 //===----------------------------------------------------------------------===//
 
-void arith::FloorDivSIOp::inferResultRanges(
-    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferRemS(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
+  const APInt &lhsMin = lhs.smin(), &lhsMax = lhs.smax(), &rhsMin = rhs.smin(),
+              &rhsMax = rhs.smax();
 
-  DivisionFixupFn floorDivSIFix =
-      [](const APInt &lhs, const APInt &rhs,
-         const APInt &result) -> std::optional<APInt> {
-    if (!lhs.srem(rhs).isZero() && lhs.isNonNegative() != rhs.isNonNegative()) {
-      bool overflowed = false;
-      APInt corrected =
-          result.ssub_ov(APInt(result.getBitWidth(), 1), overflowed);
-      return overflowed ? std::optional<APInt>() : corrected;
+  unsigned width = rhsMax.getBitWidth();
+  APInt smin = APInt::getSignedMinValue(width);
+  APInt smax = APInt::getSignedMaxValue(width);
+  // No bounds if zero could be a divisor.
+  bool canBound = (rhsMin.isStrictlyPositive() || rhsMax.isNegative());
+  if (canBound) {
+    APInt maxDivisor = rhsMin.isStrictlyPositive() ? rhsMax : rhsMin.abs();
+    bool canNegativeDividend = lhsMin.isNegative();
+    bool canPositiveDividend = lhsMax.isStrictlyPositive();
+    APInt zero = APInt::getZero(maxDivisor.getBitWidth());
+    APInt maxPositiveResult = maxDivisor - 1;
+    APInt minNegativeResult = -maxPositiveResult;
+    smin = canNegativeDividend ? minNegativeResult : zero;
+    smax = canPositiveDividend ? maxPositiveResult : zero;
+    // Special case: sweeping out a contiguous range in N/[modulus].
+    if (rhsMin == rhsMax) {
+      if ((lhsMax - lhsMin).ult(maxDivisor)) {
+        APInt minRem = lhsMin.srem(maxDivisor);
+        APInt maxRem = lhsMax.srem(maxDivisor);
+        if (minRem.sle(maxRem)) {
+          smin = minRem;
+          smax = maxRem;
+        }
+      }
     }
-    return result;
-  };
-  setResultRange(getResult(), inferDivSIRange(lhs, rhs, floorDivSIFix));
+  }
+  return ConstantIntRanges::fromSigned(smin, smax);
 }
 
 //===----------------------------------------------------------------------===//
-// RemUIOp
+// Unsigned remainder (RemU)
 //===----------------------------------------------------------------------===//
 
-void arith::RemUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferRemU(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
   const APInt &rhsMin = rhs.umin(), &rhsMax = rhs.umax();
 
   unsigned width = rhsMin.getBitWidth();
   APInt umin = APInt::getZero(width);
   APInt umax = APInt::getMaxValue(width);
 
   if (!rhsMin.isZero()) {
     umax = rhsMax - 1;
     // Special case: sweeping out a contiguous range in N/[modulus]
     if (rhsMin == rhsMax) {
       const APInt &lhsMin = lhs.umin(), &lhsMax = lhs.umax();
       if ((lhsMax - lhsMin).ult(rhsMax)) {
         APInt minRem = lhsMin.urem(rhsMax);
         APInt maxRem = lhsMax.urem(rhsMax);
         if (minRem.ule(maxRem)) {
           umin = minRem;
           umax = maxRem;
         }
       }
     }
   }
-  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(umin, umax));
+  return ConstantIntRanges::fromUnsigned(umin, umax);
 }
 
 //===----------------------------------------------------------------------===//
-// RemSIOp
+// Max and min (MaxS, MaxU, MinS, MinU)
 //===----------------------------------------------------------------------===//
 
-void arith::RemSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferMaxS(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-  const APInt &lhsMin = lhs.smin(), &lhsMax = lhs.smax(), &rhsMin = rhs.smin(),
-              &rhsMax = rhs.smax();
 
-  unsigned width = rhsMax.getBitWidth();
-  APInt smin = APInt::getSignedMinValue(width);
-  APInt smax = APInt::getSignedMaxValue(width);
-  // No bounds if zero could be a divisor.
-  bool canBound = (rhsMin.isStrictlyPositive() || rhsMax.isNegative());
-  if (canBound) {
-    APInt maxDivisor = rhsMin.isStrictlyPositive() ? rhsMax : rhsMin.abs();
-    bool canNegativeDividend = lhsMin.isNegative();
-    bool canPositiveDividend = lhsMax.isStrictlyPositive();
-    APInt zero = APInt::getZero(maxDivisor.getBitWidth());
-    APInt maxPositiveResult = maxDivisor - 1;
-    APInt minNegativeResult = -maxPositiveResult;
-    smin = canNegativeDividend ? minNegativeResult : zero;
-    smax = canPositiveDividend ? maxPositiveResult : zero;
-    // Special case: sweeping out a contiguous range in N/[modulus].
-    if (rhsMin == rhsMax) {
-      if ((lhsMax - lhsMin).ult(maxDivisor)) {
-        APInt minRem = lhsMin.srem(maxDivisor);
-        APInt maxRem = lhsMax.srem(maxDivisor);
-        if (minRem.sle(maxRem)) {
-          smin = minRem;
-          smax = maxRem;
-        }
-      }
-    }
-  }
-  setResultRange(getResult(), ConstantIntRanges::fromSigned(smin, smax));
+  const APInt &smin = lhs.smin().sgt(rhs.smin()) ? lhs.smin() : rhs.smin();
+  const APInt &smax = lhs.smax().sgt(rhs.smax()) ? lhs.smax() : rhs.smax();
+  return ConstantIntRanges::fromSigned(smin, smax);
+
+ConstantIntRanges
+mlir::intrange::inferMaxU(ArrayRef<ConstantIntRanges> argRanges) {
+  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
+
+  const APInt &umin = lhs.umin().ugt(rhs.umin()) ? lhs.umin() : rhs.umin();
+  const APInt &umax = lhs.umax().ugt(rhs.umax()) ? lhs.umax() : rhs.umax();
+  return ConstantIntRanges::fromUnsigned(umin, umax);
+
+ConstantIntRanges
+mlir::intrange::inferMinS(ArrayRef<ConstantIntRanges> argRanges) {
+  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
+
+  const APInt &smin = lhs.smin().slt(rhs.smin()) ? lhs.smin() : rhs.smin();
+  const APInt &smax = lhs.smax().slt(rhs.smax()) ? lhs.smax() : rhs.smax();
+  return ConstantIntRanges::fromSigned(smin, smax);
+
+ConstantIntRanges
+mlir::intrange::inferMinU(ArrayRef<ConstantIntRanges> argRanges) {
+  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
+
+  const APInt &umin = lhs.umin().ult(rhs.umin()) ? lhs.umin() : rhs.umin();
+  const APInt &umax = lhs.umax().ult(rhs.umax()) ? lhs.umax() : rhs.umax();
+  return ConstantIntRanges::fromUnsigned(umin, umax);
 }
 
 //===----------------------------------------------------------------------===//
-// AndIOp
+// Bitwise operators (And, Or, Xor)
 //===----------------------------------------------------------------------===//
 
 /// "Widen" bounds - if 0bvvvvv??? <= a <= 0bvvvvv???,
 /// relax the bounds to 0bvvvvv000 <= a <= 0bvvvvv111, where vvvvv are the bits
 /// that both bonuds have in common. This gives us a consertive approximation
 /// for what values can be passed to bitwise operations.
 static std::tuple<APInt, APInt>
 widenBitwiseBounds(const ConstantIntRanges &bound) {
   APInt leftVal = bound.umin(), rightVal = bound.umax();
   unsigned bitwidth = leftVal.getBitWidth();
   unsigned differingBits = bitwidth - (leftVal ^ rightVal).countLeadingZeros();
   leftVal.clearLowBits(differingBits);
   rightVal.setLowBits(differingBits);
   return std::make_tuple(std::move(leftVal), std::move(rightVal));
 }
 
-void arith::AndIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                      SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferAnd(ArrayRef<ConstantIntRanges> argRanges) {
   auto [lhsZeros, lhsOnes] = widenBitwiseBounds(argRanges[0]);
   auto [rhsZeros, rhsOnes] = widenBitwiseBounds(argRanges[1]);
   auto andi = [](const APInt &a, const APInt &b) -> std::optional<APInt> {
     return a & b;
   };
-  setResultRange(getResult(),
-                 minMaxBy(andi, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
-                          /*isSigned=*/false));
+  return minMaxBy(andi, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
+                  /*isSigned=*/false);
 }
 
-//===----------------------------------------------------------------------===//
-// OrIOp
-//===----------------------------------------------------------------------===//
-
-void arith::OrIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                     SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferOr(ArrayRef<ConstantIntRanges> argRanges) {
   auto [lhsZeros, lhsOnes] = widenBitwiseBounds(argRanges[0]);
   auto [rhsZeros, rhsOnes] = widenBitwiseBounds(argRanges[1]);
   auto ori = [](const APInt &a, const APInt &b) -> std::optional<APInt> {
     return a | b;
   };
-  setResultRange(getResult(),
-                 minMaxBy(ori, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
-                          /*isSigned=*/false));
+  return minMaxBy(ori, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
+                  /*isSigned=*/false);
 }
 
-//===----------------------------------------------------------------------===//
-// XOrIOp
-//===----------------------------------------------------------------------===//
-
-void arith::XOrIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                      SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferXor(ArrayRef<ConstantIntRanges> argRanges) {
   auto [lhsZeros, lhsOnes] = widenBitwiseBounds(argRanges[0]);
   auto [rhsZeros, rhsOnes] = widenBitwiseBounds(argRanges[1]);
   auto xori = [](const APInt &a, const APInt &b) -> std::optional<APInt> {
     return a ^ b;
   };
-  setResultRange(getResult(),
-                 minMaxBy(xori, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
-                          /*isSigned=*/false));
+  return minMaxBy(xori, {lhsZeros, lhsOnes}, {rhsZeros, rhsOnes},
+                  /*isSigned=*/false);
 }
 
 //===----------------------------------------------------------------------===//
-// MaxSIOp
+// Shifts (Shl, ShrS, ShrU)
 //===----------------------------------------------------------------------===//
 
-void arith::MaxSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferShl(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  const APInt &smin = lhs.smin().sgt(rhs.smin()) ? lhs.smin() : rhs.smin();
-  const APInt &smax = lhs.smax().sgt(rhs.smax()) ? lhs.smax() : rhs.smax();
-  setResultRange(getResult(), ConstantIntRanges::fromSigned(smin, smax));
+  ConstArithFn shl = [](const APInt &l,
+                        const APInt &r) -> std::optional<APInt> {
+    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.shl(r);
+  };
+  ConstantIntRanges urange =
+      minMaxBy(shl, {lhs.umin(), lhs.umax()}, {rhs.umin(), rhs.umax()},
+               /*isSigned=*/false);
+  ConstantIntRanges srange =
+      minMaxBy(shl, {lhs.smin(), lhs.smax()}, {rhs.umin(), rhs.umax()},
+               /*isSigned=*/true);
+  return urange.intersection(srange);
 }
 
-//===----------------------------------------------------------------------===//
-// MaxUIOp
-//===----------------------------------------------------------------------===//
-
-void arith::MaxUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferShrS(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
 
-  const APInt &umin = lhs.umin().ugt(rhs.umin()) ? lhs.umin() : rhs.umin();
-  const APInt &umax = lhs.umax().ugt(rhs.umax()) ? lhs.umax() : rhs.umax();
-  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(umin, umax));
-}
-
-//===----------------------------------------------------------------------===//
-// MinSIOp
-//===----------------------------------------------------------------------===//
-
-void arith::MinSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
+  ConstArithFn ashr = [](const APInt &l,
+                         const APInt &r) -> std::optional<APInt> {
+    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.ashr(r);
+  };
 
-  const APInt &smin = lhs.smin().slt(rhs.smin()) ? lhs.smin() : rhs.smin();
-  const APInt &smax = lhs.smax().slt(rhs.smax()) ? lhs.smax() : rhs.smax();
-  setResultRange(getResult(), ConstantIntRanges::fromSigned(smin, smax));
+  return minMaxBy(ashr, {lhs.smin(), lhs.smax()}, {rhs.umin(), rhs.umax()},
+                  /*isSigned=*/true);
 }
 
-//===----------------------------------------------------------------------===//
-// MinUIOp
-//===----------------------------------------------------------------------===//
-
-void arith::MinUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
+ConstantIntRanges
+mlir::intrange::inferShrU(ArrayRef<ConstantIntRanges> argRanges) {
   const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
 
-  const APInt &umin = lhs.umin().ult(rhs.umin()) ? lhs.umin() : rhs.umin();
-  const APInt &umax = lhs.umax().ult(rhs.umax()) ? lhs.umax() : rhs.umax();
-  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(umin, umax));
-}
-
-//===----------------------------------------------------------------------===//
-// ExtUIOp
-//===----------------------------------------------------------------------===//
-
-static ConstantIntRanges extUIRange(const ConstantIntRanges &range,
-                                    Type destType) {
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-  APInt umin = range.umin().zext(destWidth);
-  APInt umax = range.umax().zext(destWidth);
-  return ConstantIntRanges::fromUnsigned(umin, umax);
-}
-
-void arith::ExtUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
-  Type destType = getResult().getType();
-  setResultRange(getResult(), extUIRange(argRanges[0], destType));
-}
-
-//===----------------------------------------------------------------------===//
-// ExtSIOp
-//===----------------------------------------------------------------------===//
-
-static ConstantIntRanges extSIRange(const ConstantIntRanges &range,
-                                    Type destType) {
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-  APInt smin = range.smin().sext(destWidth);
-  APInt smax = range.smax().sext(destWidth);
-  return ConstantIntRanges::fromSigned(smin, smax);
-}
-
-void arith::ExtSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
-  Type destType = getResult().getType();
-  setResultRange(getResult(), extSIRange(argRanges[0], destType));
-}
-
-//===----------------------------------------------------------------------===//
-// TruncIOp
-//===----------------------------------------------------------------------===//
-
-static ConstantIntRanges truncIRange(const ConstantIntRanges &range,
-                                     Type destType) {
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-  // If you truncate the first four bytes in [0xaaaabbbb, 0xccccbbbb],
-  // the range of the resulting value is not contiguous ind includes 0.
-  // Ex. If you truncate [256, 258] from i16 to i8, you validly get [0, 2],
-  // but you can't truncate [255, 257] similarly.
-  bool hasUnsignedRollover =
-      range.umin().lshr(destWidth) != range.umax().lshr(destWidth);
-  APInt umin = hasUnsignedRollover ? APInt::getZero(destWidth)
-                                   : range.umin().trunc(destWidth);
-  APInt umax = hasUnsignedRollover ? APInt::getMaxValue(destWidth)
-                                   : range.umax().trunc(destWidth);
-
-  // Signed post-truncation rollover will not occur when either:
-  // - The high parts of the min and max, plus the sign bit, are the same
-  // - The high halves + sign bit of the min and max are either all 1s or all 0s
-  //  and you won't create a [positive, negative] range by truncating.
-  // For example, you can truncate the ranges [256, 258]_i16 to [0, 2]_i8
-  // but not [255, 257]_i16 to a range of i8s. You can also truncate
-  // [-256, -256]_i16 to [-2, 0]_i8, but not [-257, -255]_i16.
-  // You can also truncate [-130, 0]_i16 to i8 because -130_i16 (0xff7e)
-  // will truncate to 0x7e, which is greater than 0
-  APInt sminHighPart = range.smin().ashr(destWidth - 1);
-  APInt smaxHighPart = range.smax().ashr(destWidth - 1);
-  bool hasSignedOverflow =
-      (sminHighPart != smaxHighPart) &&
-      !(sminHighPart.isAllOnes() &&
-        (smaxHighPart.isAllOnes() || smaxHighPart.isZero())) &&
-      !(sminHighPart.isZero() && smaxHighPart.isZero());
-  APInt smin = hasSignedOverflow ? APInt::getSignedMinValue(destWidth)
-                                 : range.smin().trunc(destWidth);
-  APInt smax = hasSignedOverflow ? APInt::getSignedMaxValue(destWidth)
-                                 : range.smax().trunc(destWidth);
-  return {umin, umax, smin, smax};
-}
-
-void arith::TruncIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                        SetIntRangeFn setResultRange) {
-  Type destType = getResult().getType();
-  setResultRange(getResult(), truncIRange(argRanges[0], destType));
+  ConstArithFn lshr = [](const APInt &l,
+                         const APInt &r) -> std::optional<APInt> {
+    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.lshr(r);
+  };
+  return minMaxBy(lshr, {lhs.umin(), lhs.umax()}, {rhs.umin(), rhs.umax()},
+                  /*isSigned=*/false);
 }
 
 //===----------------------------------------------------------------------===//
-// IndexCastOp
+// Comparisons (Cmp)
 //===----------------------------------------------------------------------===//
 
-void arith::IndexCastOp::inferResultRanges(
-    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
-  Type sourceType = getOperand().getType();
-  Type destType = getResult().getType();
-  unsigned srcWidth = ConstantIntRanges::getStorageBitwidth(sourceType);
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-
-  if (srcWidth < destWidth)
-    setResultRange(getResult(), extSIRange(argRanges[0], destType));
-  else if (srcWidth > destWidth)
-    setResultRange(getResult(), truncIRange(argRanges[0], destType));
-  else
-    setResultRange(getResult(), argRanges[0]);
+static intrange::CmpPredicate invertPredicate(intrange::CmpPredicate pred) {
+  switch (pred) {
+  case intrange::CmpPredicate::eq:
+    return intrange::CmpPredicate::ne;
+  case intrange::CmpPredicate::ne:
+    return intrange::CmpPredicate::eq;
+  case intrange::CmpPredicate::slt:
+    return intrange::CmpPredicate::sge;
+  case intrange::CmpPredicate::sle:
+    return intrange::CmpPredicate::sgt;
+  case intrange::CmpPredicate::sgt:
+    return intrange::CmpPredicate::sle;
+  case intrange::CmpPredicate::sge:
+    return intrange::CmpPredicate::slt;
+  case intrange::CmpPredicate::ult:
+    return intrange::CmpPredicate::uge;
+  case intrange::CmpPredicate::ule:
+    return intrange::CmpPredicate::ugt;
+  case intrange::CmpPredicate::ugt:
+    return intrange::CmpPredicate::ule;
+  case intrange::CmpPredicate::uge:
+    return intrange::CmpPredicate::ult;
   }
-
-//===----------------------------------------------------------------------===//
-// IndexCastUIOp
-//===----------------------------------------------------------------------===//
-
-void arith::IndexCastUIOp::inferResultRanges(
-    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRange) {
-  Type sourceType = getOperand().getType();
-  Type destType = getResult().getType();
-  unsigned srcWidth = ConstantIntRanges::getStorageBitwidth(sourceType);
-  unsigned destWidth = ConstantIntRanges::getStorageBitwidth(destType);
-
-  if (srcWidth < destWidth)
-    setResultRange(getResult(), extUIRange(argRanges[0], destType));
-  else if (srcWidth > destWidth)
-    setResultRange(getResult(), truncIRange(argRanges[0], destType));
-  else
-    setResultRange(getResult(), argRanges[0]);
+  llvm_unreachable("unknown cmp predicate value");
 }
 
-//===----------------------------------------------------------------------===//
-// CmpIOp
-//===----------------------------------------------------------------------===//
-
-bool isStaticallyTrue(arith::CmpIPredicate pred, const ConstantIntRanges &lhs,
+static bool isStaticallyTrue(intrange::CmpPredicate pred,
+                             const ConstantIntRanges &lhs,
                              const ConstantIntRanges &rhs) {
   switch (pred) {
-  case arith::CmpIPredicate::sle:
-  case arith::CmpIPredicate::slt:
-    return (applyCmpPredicate(pred, lhs.smax(), rhs.smin()));
-  case arith::CmpIPredicate::ule:
-  case arith::CmpIPredicate::ult:
-    return applyCmpPredicate(pred, lhs.umax(), rhs.umin());
-  case arith::CmpIPredicate::sge:
-  case arith::CmpIPredicate::sgt:
-    return applyCmpPredicate(pred, lhs.smin(), rhs.smax());
-  case arith::CmpIPredicate::uge:
-  case arith::CmpIPredicate::ugt:
-    return applyCmpPredicate(pred, lhs.umin(), rhs.umax());
-  case arith::CmpIPredicate::eq: {
+  case intrange::CmpPredicate::sle:
+    return lhs.smax().sle(rhs.smin());
+  case intrange::CmpPredicate::slt:
+    return lhs.smax().slt(rhs.smin());
+  case intrange::CmpPredicate::ule:
+    return lhs.umax().ule(rhs.umin());
+  case intrange::CmpPredicate::ult:
+    return lhs.umax().ult(rhs.umin());
+  case intrange::CmpPredicate::sge:
+    return lhs.smin().sge(rhs.smax());
+  case intrange::CmpPredicate::sgt:
+    return lhs.smin().sgt(rhs.smax());
+  case intrange::CmpPredicate::uge:
+    return lhs.umin().uge(rhs.umax());
+  case intrange::CmpPredicate::ugt:
+    return lhs.umin().ugt(rhs.umax());
+  case intrange::CmpPredicate::eq: {
     std::optional<APInt> lhsConst = lhs.getConstantValue();
     std::optional<APInt> rhsConst = rhs.getConstantValue();
     return lhsConst && rhsConst && lhsConst == rhsConst;
   }
-  case arith::CmpIPredicate::ne: {
+  case intrange::CmpPredicate::ne: {
     // While equality requires that there is an interpration of the preceeding
     // computations that produces equal constants, whether that be signed or
     // unsigned, statically determining inequality requires that neither
     // interpretation produce potentially overlapping ranges.
-    bool sne = isStaticallyTrue(CmpIPredicate::slt, lhs, rhs) ||
-               isStaticallyTrue(CmpIPredicate::sgt, lhs, rhs);
-    bool une = isStaticallyTrue(CmpIPredicate::ult, lhs, rhs) ||
-               isStaticallyTrue(CmpIPredicate::ugt, lhs, rhs);
+    bool sne = isStaticallyTrue(intrange::CmpPredicate::slt, lhs, rhs) ||
+               isStaticallyTrue(intrange::CmpPredicate::sgt, lhs, rhs);
+    bool une = isStaticallyTrue(intrange::CmpPredicate::ult, lhs, rhs) ||
+               isStaticallyTrue(intrange::CmpPredicate::ugt, lhs, rhs);
     return sne && une;
   }
   }
   return false;
 }
 
-void arith::CmpIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                      SetIntRangeFn setResultRange) {
-  arith::CmpIPredicate pred = getPredicate();
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  APInt min = APInt::getZero(1);
-  APInt max = APInt::getAllOnesValue(1);
+std::optional<bool> mlir::intrange::evaluatePred(CmpPredicate pred,
+                                                 const ConstantIntRanges &lhs,
+                                                 const ConstantIntRanges &rhs) {
   if (isStaticallyTrue(pred, lhs, rhs))
-    min = max;
-  else if (isStaticallyTrue(invertPredicate(pred), lhs, rhs))
-    max = min;
-
-  setResultRange(getResult(), ConstantIntRanges::fromUnsigned(min, max));
-}
-
-//===----------------------------------------------------------------------===//
-// SelectOp
-//===----------------------------------------------------------------------===//
-
-void arith::SelectOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                        SetIntRangeFn setResultRange) {
-  std::optional<APInt> mbCondVal = argRanges[0].getConstantValue();
-
-  if (mbCondVal) {
-    if (mbCondVal->isZero())
-      setResultRange(getResult(), argRanges[2]);
-    else
-      setResultRange(getResult(), argRanges[1]);
-    return;
-  }
-  setResultRange(getResult(), argRanges[1].rangeUnion(argRanges[2]));
-}
-
-//===----------------------------------------------------------------------===//
-// ShLIOp
-//===----------------------------------------------------------------------===//
-
-void arith::ShLIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                      SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-  ConstArithFn shl = [](const APInt &l,
-                        const APInt &r) -> std::optional<APInt> {
-    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.shl(r);
-  };
-  ConstantIntRanges urange =
-      minMaxBy(shl, {lhs.umin(), lhs.umax()}, {rhs.umin(), rhs.umax()},
-               /*isSigned=*/false);
-  ConstantIntRanges srange =
-      minMaxBy(shl, {lhs.smin(), lhs.smax()}, {rhs.umin(), rhs.umax()},
-               /*isSigned=*/true);
-  setResultRange(getResult(), urange.intersection(srange));
-}
-
-//===----------------------------------------------------------------------===//
-// ShRUIOp
-//===----------------------------------------------------------------------===//
-
-void arith::ShRUIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  ConstArithFn lshr = [](const APInt &l,
-                         const APInt &r) -> std::optional<APInt> {
-    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.lshr(r);
-  };
-  setResultRange(getResult(), minMaxBy(lshr, {lhs.umin(), lhs.umax()},
-                                       {rhs.umin(), rhs.umax()},
-                                       /*isSigned=*/false));
-}
-
-//===----------------------------------------------------------------------===//
-// ShRSIOp
-//===----------------------------------------------------------------------===//
-
-void arith::ShRSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
-                                       SetIntRangeFn setResultRange) {
-  const ConstantIntRanges &lhs = argRanges[0], &rhs = argRanges[1];
-
-  ConstArithFn ashr = [](const APInt &l,
-                         const APInt &r) -> std::optional<APInt> {
-    return r.uge(r.getBitWidth()) ? std::optional<APInt>() : l.ashr(r);
-  };
-
-  setResultRange(getResult(),
-                 minMaxBy(ashr, {lhs.smin(), lhs.smax()},
-                          {rhs.umin(), rhs.umax()}, /*isSigned=*/true));
+    return true;
+  if (isStaticallyTrue(invertPredicate(pred), lhs, rhs))
+    return false;
+  return std::nullopt;
 }

mlir/test/Dialect/Index/int-range-inference.mlir

+// RUN: mlir-opt -test-int-range-inference -canonicalize %s | FileCheck %s
+
+// Most operations are covered by the `arith` tests, which use the same code
+// Here, we add a few tests to ensure the "index can be 32- or 64-bit" handling
+// code is operating as expected.
+
+// CHECK-LABEL: func @add_same_for_both
+// CHECK: %[[true:.*]] = index.bool.constant true
+// CHECK: return %[[true]]
+func.func @add_same_for_both(%arg0 : index) -> i1 {
+  %c1 = index.constant 1
+  %calmostBig = index.constant 0xfffffffe
+  %0 = index.minu %arg0, %calmostBig
+  %1 = index.add %0, %c1
+  %2 = index.cmp uge(%1, %c1)
+  func.return %2 : i1
+}
+
+// CHECK-LABEL: func @add_unsigned_ov
+// CHECK: %[[uge:.*]] = index.cmp uge
+// CHECK: return %[[uge]]
+func.func @add_unsigned_ov(%arg0 : index) -> i1 {
+  %c1 = index.constant 1
+  %cu32_max = index.constant 0xffffffff
+  %0 = index.minu %arg0, %cu32_max
+  %1 = index.add %0, %c1
+  // On 32-bit, the add could wrap, so the result doesn't have to be >= 1
+  %2 = index.cmp uge(%1, %c1)
+  func.return %2 : i1
+}
+
+// CHECK-LABEL: func @add_signed_ov
+// CHECK: %[[sge:.*]] = index.cmp sge
+// CHECK: return %[[sge]]
+func.func @add_signed_ov(%arg0 : index) -> i1 {
+  %c0 = index.constant 0
+  %c1 = index.constant 1
+  %ci32_max = index.constant 0x7fffffff
+  %0 = index.minu %arg0, %ci32_max
+  %1 = index.add %0, %c1
+  // On 32-bit, the add could wrap, so the result doesn't have to be positive
+  %2 = index.cmp sge(%1, %c0)
+  func.return %2 : i1
+}
+
+// CHECK-LABEL: func @add_big
+// CHECK: %[[true:.*]] = index.bool.constant true
+// CHECK: return %[[true]]
+func.func @add_big(%arg0 : index) -> i1 {
+  %c1 = index.constant 1
+  %cmin = index.constant 0x300000000
+  %cmax = index.constant 0x30000ffff
+  // Note: the order of the clamps matters.
+  // If you go max, then min, you infer the ranges [0x300...0, 0xff..ff]
+  // and then [0x30...0000, 0x30...ffff]
+  // If you switch the order of the below operations, you instead first infer
+  // the range [0,0x3...ffff]. Then, the min inference can't constraint
+  // this intermediate, since in the 32-bit case we could have, for example
+  // trunc(%arg0 = 0x2ffffffff) = 0xffffffff > trunc(0x30000ffff) = 0x0000ffff
+  // which means we can't do any inference.
+  %0 = index.maxu %arg0, %cmin
+  %1 = index.minu %0, %cmax
+  %2 = index.add %1, %c1
+  %3 = index.cmp uge(%1, %cmin)
+  func.return %3 : i1
+}