Improved range analysis fixes

Author: sabauma

mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h

@@ -33,7 +33,7 @@ class IntegerValueRange {
   static IntegerValueRange getMaxRange(Value value);
 
   /// Create an integer value range lattice value.
-  IntegerValueRange(std::optional<ConstantIntRanges> value = std::nullopt)
+  IntegerValueRange(OptionalIntRanges value = std::nullopt)
       : value(std::move(value)) {}
 
   /// Whether the range is uninitialized. This happens when the state hasn't

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

@@ -105,10 +105,11 @@ class ConstantIntRanges {
 
 raw_ostream &operator<<(raw_ostream &, const ConstantIntRanges &);
 
+using OptionalIntRanges = std::optional<ConstantIntRanges>;
 /// The type of the `setResultRanges` callback provided to ops implementing
 /// InferIntRangeInterface. It should be called once for each integer result
 /// value and be passed the ConstantIntRanges corresponding to that value.
-using SetIntRangeFn = function_ref<void(Value, const ConstantIntRanges &)>;
+using SetIntRangeFn = function_ref<void(Value, const OptionalIntRanges &)>;
 } // end namespace mlir
 
 #include "mlir/Interfaces/InferIntRangeInterface.h.inc"

mlir/include/mlir/Interfaces/InferIntRangeInterface.td

@@ -45,7 +45,7 @@ def InferIntRangeInterface : OpInterface<"InferIntRangeInterface"> {
        APInts in their `argRanges` element.
     }],
     "void", "inferResultRanges", (ins
-      "::llvm::ArrayRef<::mlir::ConstantIntRanges>":$argRanges,
+      "::llvm::ArrayRef<::std::optional<::mlir::ConstantIntRanges>>":$argRanges,
       "::mlir::SetIntRangeFn":$setResultRanges)
   >];
 }

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

@@ -27,6 +27,9 @@ namespace intrange {
 using InferRangeFn =
     function_ref<ConstantIntRanges(ArrayRef<ConstantIntRanges>)>;
 
+using OptionalRangeFn =
+    std::function<OptionalIntRanges(ArrayRef<OptionalIntRanges>)>;
+
 static constexpr unsigned indexMinWidth = 32;
 static constexpr unsigned indexMaxWidth = 64;
 
@@ -44,6 +47,8 @@ enum class OverflowFlags : uint32_t {
 using InferRangeWithOvfFlagsFn =
     function_ref<ConstantIntRanges(ArrayRef<ConstantIntRanges>, OverflowFlags)>;
 
+OptionalRangeFn inferFromOptionals(intrange::InferRangeFn inferFn);
+
 /// 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

mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp

@@ -36,8 +36,25 @@
 using namespace mlir;
 using namespace mlir::dataflow;
 
+namespace {
+
+OptionalIntRanges getOptionalRange(const IntegerValueRange &range) {
+  if (range.isUninitialized())
+    return std::nullopt;
+  return range.getValue();
+}
+
+OptionalIntRanges getOptionalRangeFromLattice(const IntegerValueRangeLattice* lattice) {
+  return getOptionalRange(lattice->getValue());
+}
+
+} // end namespace
+
 IntegerValueRange IntegerValueRange::getMaxRange(Value value) {
   unsigned width = ConstantIntRanges::getStorageBitwidth(value.getType());
+  if (width == 0)
+    return {};
+
   APInt umin = APInt::getMinValue(width);
   APInt umax = APInt::getMaxValue(width);
   APInt smin = width != 0 ? APInt::getSignedMinValue(width) : umin;
@@ -71,23 +88,15 @@ void IntegerRangeAnalysis::visitOperation(
     Operation *op, ArrayRef<const IntegerValueRangeLattice *> operands,
     ArrayRef<IntegerValueRangeLattice *> results) {
   // If the lattice on any operand is unitialized, bail out.
-  if (llvm::any_of(operands, [](const IntegerValueRangeLattice *lattice) {
-        return lattice->getValue().isUninitialized();
-      })) {
-    return;
-  }
-
   auto inferrable = dyn_cast<InferIntRangeInterface>(op);
   if (!inferrable)
     return setAllToEntryStates(results);
 
   LLVM_DEBUG(llvm::dbgs() << "Inferring ranges for " << *op << "\n");
-  SmallVector<ConstantIntRanges> argRanges(
-      llvm::map_range(operands, [](const IntegerValueRangeLattice *val) {
-        return val->getValue().getValue();
-      }));
+  SmallVector<OptionalIntRanges> argRanges(llvm::map_range(
+      operands, getOptionalRangeFromLattice));
 
-  auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
+  auto joinCallback = [&](Value v, const OptionalIntRanges &attrs) {
     auto result = dyn_cast<OpResult>(v);
     if (!result)
       return;
@@ -97,7 +106,9 @@ void IntegerRangeAnalysis::visitOperation(
     IntegerValueRangeLattice *lattice = results[result.getResultNumber()];
     IntegerValueRange oldRange = lattice->getValue();
 
-    ChangeResult changed = lattice->join(IntegerValueRange{attrs});
+    ChangeResult changed =
+        attrs ? lattice->join(IntegerValueRange{attrs})
+              : lattice->join(IntegerValueRange::getMaxRange(v));
 
     // Catch loop results with loop variant bounds and conservatively make
     // them [-inf, inf] so we don't circle around infinitely often (because
@@ -127,12 +138,12 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
           return getLatticeElementFor(op, value)->getValue().isUninitialized();
         }))
       return;
-    SmallVector<ConstantIntRanges> argRanges(
+    SmallVector<OptionalIntRanges> argRanges(
         llvm::map_range(op->getOperands(), [&](Value value) {
-          return getLatticeElementFor(op, value)->getValue().getValue();
+          return getOptionalRangeFromLattice(getLatticeElementFor(op, value));
         }));
 
-    auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
+    auto joinCallback = [&](Value v, const OptionalIntRanges &attrs) {
       auto arg = dyn_cast<BlockArgument>(v);
       if (!arg)
         return;
@@ -143,7 +154,9 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
       IntegerValueRangeLattice *lattice = argLattices[arg.getArgNumber()];
       IntegerValueRange oldRange = lattice->getValue();
 
-      ChangeResult changed = lattice->join(IntegerValueRange{attrs});
+      ChangeResult changed =
+          attrs ? lattice->join(IntegerValueRange{attrs})
+                : lattice->join(IntegerValueRange::getMaxRange(v));
 
       // Catch loop results with loop variant bounds and conservatively make
       // them [-inf, inf] so we don't circle around infinitely often (because