[mlir] Swap integer range inference to the new framework

[Mogball]

Summary

Integer range inference has been swapped to the new framework. The integer value range lattices automatically updates the corresponding constant value on update.

Depends on D127173

[krzysz00]

So, overall comments - especially since it looks like all the details are basically fine ... wouldn't it still be a good idea to have a wrapper around IntRangeAnalysis that hides the fact that it's a dataflow analysis? Or are we abandoning that wrapper here?

Also, how's this meant to interact with SCCP? (Is that future work?)

[Mogball]

The details can be hidden, but then the analysis would not be composable because... it would be hidden.

It can interact with SCCP by just loading both analyses at the same time. This is better than running them separately because the partial results of each analysis can help refine each other. Although there are some imperfections in that interaction with the current setup but you can check https://reviews.llvm.org/D128868 for a fully-functioning composable setup.

But this doesn't address River's complaint in your patch of the analysis making two passes over the IR. A monolithic analysis will inherently be more efficient, but that's the price paid for extensibility. At least in this case, they can work together instead of just run separately.

[krzysz00]

LGTM, at least as far as the integer analysis parts go.

I'd want to make sure @rriddle's fine with the bigger API surface.

[rriddle]

It's a bit sad that unnecessary API is being exposed here. I think we'll want to have cleaner wrappers around the lattice state though, feels like a pattern is emerging of grabbing the state -> checking null + uninitialized -> calling getValue.

mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp

173–175

If you want to include unregistered as potential terminators, this should use mightHaveTrait which returns true if the op is unregistered.

[Mogball]

mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp

173–175

That would give a false positive for unregistered operations that aren't the last operation in the block. Visiting these operations would be a no-op in the analysis, but that's extra computation that can be avoided.

[Mogball]

change isTerminator to be more precise

[zero9178]

Just as a heads up: I have bisected miscompiles in my downstream compiler to this patch. I have figured out a small reproducer that can be run
via mlir-opt %s --pass-pipeline="func.func(sccp)":

func.func private @foo() -> index {
    %0 = arith.constant 10 : index
    return %0 : index
}

func.func private @bar(%arg0: index) -> index {
  %c0 = arith.constant 0 : index
  %1 = arith.constant 420 : index
  %7 = arith.cmpi eq, %arg0, %c0 : index
  cf.cond_br %7, ^bb1(%1 : index), ^bb2

^bb1(%8: index):  // 2 preds: ^bb0, ^bb4
  return %8 : index

^bb2:
  %13 = call @foo() : () -> index
  cf.br ^bb1(%13 : index)
}

this will incorrectly fold the return %8 : index to return %c420 : index, despite ^bb2 branching to ^bb1 with a different value as block argument.

Tested this on both the current revision: 9ad082eb5a948ffd5d8fc959967d184c3433aca0, the one before this commit and the revision of this commit.

[zero9178]

I have done some further investigation and the issue isn't 100% with this patch itself but more with the Dataflow analysis itself I believe.
Specifically, the reason that this patch causes this miscompile seems to be because of the lines removed at DataFlowFramework.cpp:90 which served as a sort of mitigataion I believe.

The real underlying issue is that:
https://github.com/llvm/llvm-project/blob/ab701975e7f3b63bb474afbdeb8c474950d41074/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp#L109
queries for a predecessor state of a call that is never created/written to by DCE. Since PredecessorsState optimistically assumes all predecessors are know, and also has no known predecessors at the same time, it does nothing but return at line 117.
I can only guess, but it seems to me the intention was that this predecessor state should be set in either:
https://github.com/llvm/llvm-project/blob/ab701975e7f3b63bb474afbdeb8c474950d41074/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp#L410
or
https://github.com/llvm/llvm-project/blob/ab701975e7f3b63bb474afbdeb8c474950d41074/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp#L300
or maybe during priming of the analysis.
The former does not work as the analysis is scheduled on func.func operations, hence it never seeing the func.return of @foo.
The latter only sets the call as predecessor to the callable.
I think what is needed is to mark the predecessors of the call as unknown if the callable is not part of the analysis.

[Mogball]

Thanks for finding this! I have a fix in https://reviews.llvm.org/D130829 (please take a look!)

If you could copy-paste that into a git issue, it'd be nice too

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

+//===-IntegerRangeAnalysis.h - Integer range analysis -----------*- 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 the dataflow analysis class for integer range inference
+// so that it can be used in transformations over the `arith` dialect such as
+// branch elimination or signed->unsigned rewriting
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_DATAFLOW_INTEGERANGEANALYSIS_H
+#define MLIR_ANALYSIS_DATAFLOW_INTEGERANGEANALYSIS_H
+
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+
+namespace mlir {
+namespace dataflow {
+
+/// This lattice value represents the integer range of an SSA value.
+class IntegerValueRange {
+public:
+  /// Create a maximal range ([0, uint_max(t)] / [int_min(t), int_max(t)])
+  /// range that is used to mark the value as unable to be analyzed further,
+  /// where `t` is the type of `value`.
+  static IntegerValueRange getPessimisticValueState(Value value);
+
+  /// Create an integer value range lattice value.
+  IntegerValueRange(ConstantIntRanges value) : value(std::move(value)) {}
+
+  /// Get the known integer value range.
+  const ConstantIntRanges &getValue() const { return value; }
+
+  /// Compare two ranges.
+  bool operator==(const IntegerValueRange &rhs) const {
+    return value == rhs.value;
+  }
+
+  /// Take the union of two ranges.
+  static IntegerValueRange join(const IntegerValueRange &lhs,
+                                const IntegerValueRange &rhs) {
+    return lhs.value.rangeUnion(rhs.value);
+  }
+
+  /// Print the integer value range.
+  void print(raw_ostream &os) const { os << value; }
+
+private:
+  /// The known integer value range.
+  ConstantIntRanges value;
+};
+
+/// This lattice element represents the integer value range of an SSA value.
+/// When this lattice is updated, it automatically updates the constant value
+/// of the SSA value (if the range can be narrowed to one).
+class IntegerValueRangeLattice : public Lattice<IntegerValueRange> {
+public:
+  using Lattice::Lattice;
+
+  /// If the range can be narrowed to an integer constant, update the constant
+  /// value of the SSA value.
+  void onUpdate(DataFlowSolver *solver) const override;
+};
+
+/// Integer range analysis determines the integer value range of SSA values
+/// using operations that define `InferIntRangeInterface` and also sets the
+/// range of iteration indices of loops with known bounds.
+class IntegerRangeAnalysis
+    : public SparseDataFlowAnalysis<IntegerValueRangeLattice> {
+public:
+  using SparseDataFlowAnalysis::SparseDataFlowAnalysis;
+
+  /// Visit an operation. Invoke the transfer function on each operation that
+  /// implements `InferIntRangeInterface`.
+  void visitOperation(Operation *op,
+                      ArrayRef<const IntegerValueRangeLattice *> operands,
+                      ArrayRef<IntegerValueRangeLattice *> results) override;
+
+  /// Visit block arguments or operation results of an operation with region
+  /// control-flow for which values are not defined by region control-flow. This
+  /// function calls `InferIntRangeInterface` to provide values for block
+  /// arguments or tries to reduce the range on loop induction variables with
+  /// known bounds.
+  void
+  visitNonControlFlowArguments(Operation *op, const RegionSuccessor &successor,
+                               ArrayRef<IntegerValueRangeLattice *> argLattices,
+                               unsigned firstIndex) override;
+};
+
+} // end namespace dataflow
+} // end namespace mlir
+
+#endif // MLIR_ANALYSIS_DATAFLOW_INTEGERANGEANALYSIS_H

mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h

 // dead code analysis to prune dead code during the analysis.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_DATAFLOW_SPARSEANALYSIS_H
 #define MLIR_ANALYSIS_DATAFLOW_SPARSEANALYSIS_H
 
 #include "mlir/Analysis/DataFlowFramework.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 namespace mlir {
-
-class RegionBranchOpInterface;
-
 namespace dataflow {
 
 //===----------------------------------------------------------------------===//
 // AbstractSparseLattice
 //===----------------------------------------------------------------------===//
 
 /// This class represents an abstract lattice. A lattice contains information
 /// about an SSA value and is what's propagated across the IR by sparse
protected:
 
   /// The operation transfer function. Given the operand lattices, this
   /// function is expected to set the result lattices.
   virtual void
   visitOperationImpl(Operation *op,
                      ArrayRef<const AbstractSparseLattice *> operandLattices,
                      ArrayRef<AbstractSparseLattice *> resultLattices) = 0;
 
+  /// Given an operation with region control-flow, the lattices of the operands,
+  /// and a region successor, compute the lattice values for block arguments
+  /// that are not accounted for by the branching control flow (ex. the bounds
+  /// of loops).
+  virtual void visitNonControlFlowArgumentsImpl(
+      Operation *op, const RegionSuccessor &successor,
+      ArrayRef<AbstractSparseLattice *> argLattices, unsigned firstIndex) = 0;
+
   /// Get the lattice element of a value.
   virtual AbstractSparseLattice *getLatticeElement(Value value) = 0;
 
   /// Get a read-only lattice element for a value and add it as a dependency to
   /// a program point.
   const AbstractSparseLattice *getLatticeElementFor(ProgramPoint point,
                                                     Value value);
 
public:
   explicit SparseDataFlowAnalysis(DataFlowSolver &solver)
       : AbstractSparseDataFlowAnalysis(solver) {}
 
   /// Visit an operation with the lattices of its operands. This function is
   /// expected to set the lattices of the operation's results.
   virtual void visitOperation(Operation *op, ArrayRef<const StateT *> operands,
                               ArrayRef<StateT *> results) = 0;
 
+  /// Given an operation with possible region control-flow, the lattices of the
+  /// operands, and a region successor, compute the lattice values for block
+  /// arguments that are not accounted for by the branching control flow (ex.
+  /// the bounds of loops). By default, this method marks all such lattice
+  /// elements as having reached a pessimistic fixpoint. `firstIndex` is the
+  /// index of the first element of `argLattices` that is set by control-flow.
+  virtual void visitNonControlFlowArguments(Operation *op,
+                                            const RegionSuccessor &successor,
+                                            ArrayRef<StateT *> argLattices,
+                                            unsigned firstIndex) {
+    markAllPessimisticFixpoint(argLattices.take_front(firstIndex));
+    markAllPessimisticFixpoint(argLattices.drop_front(
+        firstIndex + successor.getSuccessorInputs().size()));
+  }
+
 protected:
   /// Get the lattice element for a value.
   StateT *getLatticeElement(Value value) override {
     return getOrCreate<StateT>(value);
   }
 
   /// Get the lattice element for a value and create a dependency on the
   /// provided program point.
void visitOperationImpl(
       ArrayRef<AbstractSparseLattice *> resultLattices) override {
     visitOperation(
         op,
         {reinterpret_cast<const StateT *const *>(operandLattices.begin()),
          operandLattices.size()},
         {reinterpret_cast<StateT *const *>(resultLattices.begin()),
          resultLattices.size()});
   }
+  void visitNonControlFlowArgumentsImpl(
+      Operation *op, const RegionSuccessor &successor,
+      ArrayRef<AbstractSparseLattice *> argLattices,
+      unsigned firstIndex) override {
+    visitNonControlFlowArguments(
+        op, successor,
+        {reinterpret_cast<StateT *const *>(argLattices.begin()),
+         argLattices.size()},
+        firstIndex);
+  }
 };
 
 } // end namespace dataflow
 } // end namespace mlir
 
 #endif // MLIR_ANALYSIS_DATAFLOW_SPARSEANALYSIS_H

mlir/include/mlir/Analysis/DataFlowFramework.h

public:
 
   /// A work item on the solver queue is a program point, child analysis pair.
   /// Each item is processed by invoking the child analysis at the program
   /// point.
   using WorkItem = std::pair<ProgramPoint, DataFlowAnalysis *>;
   /// Push a work item onto the worklist.
   void enqueue(WorkItem item) { worklist.push(std::move(item)); }
 
-protected:
   /// Get the state associated with the given program point. If it does not
   /// exist, create an uninitialized state.
   template <typename StateT, typename PointT>
   StateT *getOrCreateState(PointT point);
 
   /// Propagate an update to an analysis state if it changed by pushing
   /// dependent work items to the back of the queue.
   void propagateIfChanged(AnalysisState *state, ChangeResult changed);

mlir/include/mlir/Analysis/IntRangeAnalysis.h

-//===- IntRangeAnalysis.h - Infer Ranges Interfaces --*- 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 the dataflow analysis class for integer range inference
-// so that it can be used in transformations over the `arith` dialect such as
-// branch elimination or signed->unsigned rewriting
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_ANALYSIS_INTRANGEANALYSIS_H
-#define MLIR_ANALYSIS_INTRANGEANALYSIS_H
-
-#include "mlir/Interfaces/InferIntRangeInterface.h"
-
-namespace mlir {
-namespace detail {
-class IntRangeAnalysisImpl;
-} // end namespace detail
-
-class IntRangeAnalysis {
-public:
-  /// Analyze all operations rooted under (but not including)
-  /// `topLevelOperation`.
-  IntRangeAnalysis(Operation *topLevelOperation);
-  IntRangeAnalysis(IntRangeAnalysis &&other);
-  ~IntRangeAnalysis();
-
-  /// Get inferred range for value `v` if one exists.
-  Optional<ConstantIntRanges> getResult(Value v);
-
-private:
-  std::unique_ptr<detail::IntRangeAnalysisImpl> impl;
-};
-} // end namespace mlir
-
-#endif

mlir/include/mlir/Interfaces/InferIntRangeInterface.td

def InferIntRangeInterface : OpInterface<"InferIntRangeInterface"> {
 
   let methods = [
     InterfaceMethod<[{
       Infer the bounds on the results of this op given the bounds on its arguments.
       For each result value or block argument (that isn't a branch argument,
       since the dataflow analysis handles those case), the method should call
       `setValueRange` with that `Value` as an argument. When `setValueRange`
       is not called for some value, it will recieve a default value of the mimimum
-      and maximum values forits type (the unbounded range).
+      and maximum values for its type (the unbounded range).
 
       When called on an op that also implements the RegionBranchOpInterface
       or BranchOpInterface, this method should not attempt to infer the values
       of the branch results, as this will be handled by the analyses that use
       this interface.
 
       This function will only be called when at least one result of the op is a
       scalar integer value or the op has a region.

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
-  IntRangeAnalysis.cpp
   Liveness.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   DataFlow/ConstantPropagationAnalysis.cpp
   DataFlow/DeadCodeAnalysis.cpp
   DataFlow/DenseAnalysis.cpp
+  DataFlow/IntegerRangeAnalysis.cpp
   DataFlow/SparseAnalysis.cpp
   )
 
 add_mlir_library(MLIRAnalysis
   AliasAnalysis.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataFlowFramework.cpp
   DataLayoutAnalysis.cpp
-  IntRangeAnalysis.cpp
   Liveness.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   DataFlow/ConstantPropagationAnalysis.cpp
   DataFlow/DeadCodeAnalysis.cpp
   DataFlow/DenseAnalysis.cpp
+  DataFlow/IntegerRangeAnalysis.cpp
   DataFlow/SparseAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 
   DEPENDS
   mlir-headers
 

mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp

for (const SymbolTable::SymbolUse &use : *uses) {
       auto *state = getOrCreate<PredecessorState>(symbol);
       propagateIfChanged(state, state->setHasUnknownPredecessors());
     }
   };
   SymbolTable::walkSymbolTables(top, /*allSymUsesVisible=*/!top->getBlock(),
                                 walkFn);
 }
 
+/// Returns true if the operation is a returning terminator in region
+/// control-flow or the terminator of a callable region.
+static bool isRegionOrCallableReturn(Operation *op) {
+  return !op->getNumSuccessors() &&
+         isa<RegionBranchOpInterface, CallableOpInterface>(op->getParentOp()) &&
+         op->getBlock()->getTerminator() == op;
+}
+
 LogicalResult DeadCodeAnalysis::initializeRecursively(Operation *op) {
   // Initialize the analysis by visiting every op with control-flow semantics.
   if (op->getNumRegions() || op->getNumSuccessors() ||
-      op->hasTrait<OpTrait::IsTerminator>() || isa<CallOpInterface>(op)) {
+      isRegionOrCallableReturn(op) || isa<CallOpInterface>(op)) {
     // When the liveness of the parent block changes, make sure to re-invoke the
     // analysis on the op.
     if (op->getBlock())
       getOrCreate<Executable>(op->getBlock())->blockContentSubscribe(this);
     // Visit the op.
     if (failed(visit(op)))
       return failure();
   }
if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
         markEntryBlocksLive(callable);
 
       // Otherwise, conservatively mark all entry blocks as executable.
     } else {
       markEntryBlocksLive(op);
     }
   }
 
-  if (op->hasTrait<OpTrait::IsTerminator>() && !op->getNumSuccessors()) {
+  if (isRegionOrCallableReturn(op)) {
     if (auto branch = dyn_cast<RegionBranchOpInterface>(op->getParentOp())) {
       // Visit the exiting terminator of a region.
       visitRegionTerminator(op, branch);
     } else if (auto callable =
                    dyn_cast<CallableOpInterface>(op->getParentOp())) {
       // Visit the exiting terminator of a callable.
       visitCallableTerminator(op, callable);
     }

mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp

+//===- IntegerRangeAnalysis.cpp - Integer range analysis --------*- 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 defines the dataflow analysis class for integer range inference
+// which is used in transformations over the `arith` dialect such as
+// branch elimination or signed->unsigned rewriting
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "int-range-analysis"
+
+using namespace mlir;
+using namespace mlir::dataflow;
+
+IntegerValueRange IntegerValueRange::getPessimisticValueState(Value value) {
+  unsigned width = ConstantIntRanges::getStorageBitwidth(value.getType());
+  APInt umin = APInt::getMinValue(width);
+  APInt umax = APInt::getMaxValue(width);
+  APInt smin = width != 0 ? APInt::getSignedMinValue(width) : umin;
+  APInt smax = width != 0 ? APInt::getSignedMaxValue(width) : umax;
+  return {{umin, umax, smin, smax}};
+}
+
+void IntegerValueRangeLattice::onUpdate(DataFlowSolver *solver) const {
+  Lattice::onUpdate(solver);
+
+  // If the integer range can be narrowed to a constant, update the constant
+  // value of the SSA value.
+  Optional<APInt> constant = getValue().getValue().getConstantValue();
+  auto value = point.get<Value>();
+  auto *cv = solver->getOrCreateState<Lattice<ConstantValue>>(value);
+  if (!constant)
+    return solver->propagateIfChanged(cv, cv->markPessimisticFixpoint());
+
+  Dialect *dialect;
+  if (auto *parent = value.getDefiningOp())
+    dialect = parent->getDialect();
+  else
+    dialect = value.getParentBlock()->getParentOp()->getDialect();
+  solver->propagateIfChanged(
+      cv, cv->join(ConstantValue(IntegerAttr::get(value.getType(), *constant),
+                                 dialect)));
+}
+
+void IntegerRangeAnalysis::visitOperation(
+    Operation *op, ArrayRef<const IntegerValueRangeLattice *> operands,
+    ArrayRef<IntegerValueRangeLattice *> results) {
+  // Ignore non-integer outputs - return early if the op has no scalar
+  // integer results
+  bool hasIntegerResult = false;
+  for (auto it : llvm::zip(results, op->getResults())) {
+    if (std::get<1>(it).getType().isIntOrIndex()) {
+      hasIntegerResult = true;
+    } else {
+      propagateIfChanged(std::get<0>(it),
+                         std::get<0>(it)->markPessimisticFixpoint());
+    }
+  }
+  if (!hasIntegerResult)
+    return;
+
+  auto inferrable = dyn_cast<InferIntRangeInterface>(op);
+  if (!inferrable)
+    return markAllPessimisticFixpoint(results);
+
+  LLVM_DEBUG(llvm::dbgs() << "Inferring ranges for " << *op << "\n");
+  SmallVector<ConstantIntRanges> argRanges(
+      llvm::map_range(operands, [](const IntegerValueRangeLattice *val) {
+        return val->getValue().getValue();
+      }));
+
+  auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
+    auto result = v.dyn_cast<OpResult>();
+    if (!result)
+      return;
+    assert(llvm::find(op->getResults(), result) != op->result_end());
+
+    LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
+    IntegerValueRangeLattice *lattice = results[result.getResultNumber()];
+    Optional<IntegerValueRange> oldRange;
+    if (!lattice->isUninitialized())
+      oldRange = lattice->getValue();
+
+    ChangeResult changed = lattice->join(attrs);
+
+    // Catch loop results with loop variant bounds and conservatively make
+    // them [-inf, inf] so we don't circle around infinitely often (because
+    // the dataflow analysis in MLIR doesn't attempt to work out trip counts
+    // and often can't).
+    bool isYieldedResult = llvm::any_of(v.getUsers(), [](Operation *op) {
+      return op->hasTrait<OpTrait::IsTerminator>();
+    });
+    if (isYieldedResult && oldRange.hasValue() &&
+        !(lattice->getValue() == *oldRange)) {
+      LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
+      changed |= lattice->markPessimisticFixpoint();
+    }
+    propagateIfChanged(lattice, changed);
+  };
+
+  inferrable.inferResultRanges(argRanges, joinCallback);
+}
+
+void IntegerRangeAnalysis::visitNonControlFlowArguments(
+    Operation *op, const RegionSuccessor &successor,
+    ArrayRef<IntegerValueRangeLattice *> argLattices, unsigned firstIndex) {
+  if (auto inferrable = dyn_cast<InferIntRangeInterface>(op)) {
+    LLVM_DEBUG(llvm::dbgs() << "Inferring ranges for " << *op << "\n");
+    SmallVector<ConstantIntRanges> argRanges(
+        llvm::map_range(op->getOperands(), [&](Value value) {
+          return getLatticeElementFor(op, value)->getValue().getValue();
+        }));
+
+    auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
+      auto arg = v.dyn_cast<BlockArgument>();
+      if (!arg)
+        return;
+      if (llvm::find(successor.getSuccessor()->getArguments(), arg) ==
+          successor.getSuccessor()->args_end())
+        return;
+
+      LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
+      IntegerValueRangeLattice *lattice = argLattices[arg.getArgNumber()];
+      Optional<IntegerValueRange> oldRange;
+      if (!lattice->isUninitialized())
+        oldRange = lattice->getValue();
+
+      ChangeResult changed = lattice->join(attrs);
+
+      // Catch loop results with loop variant bounds and conservatively make
+      // them [-inf, inf] so we don't circle around infinitely often (because
+      // the dataflow analysis in MLIR doesn't attempt to work out trip counts
+      // and often can't).
+      bool isYieldedValue = llvm::any_of(v.getUsers(), [](Operation *op) {
+        return op->hasTrait<OpTrait::IsTerminator>();
+      });
+      if (isYieldedValue && oldRange && !(lattice->getValue() == *oldRange)) {
+        LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
+        changed |= lattice->markPessimisticFixpoint();
+      }
+      propagateIfChanged(lattice, changed);
+    };
+
+    inferrable.inferResultRanges(argRanges, joinCallback);
+    return;
+  }
+
+  /// Given the results of getConstant{Lower,Upper}Bound() or getConstantStep()
+  /// on a LoopLikeInterface return the lower/upper bound for that result if
+  /// possible.
+  auto getLoopBoundFromFold = [&](Optional<OpFoldResult> loopBound,
+                                  Type boundType, bool getUpper) {
+    unsigned int width = ConstantIntRanges::getStorageBitwidth(boundType);
+    if (loopBound.hasValue()) {
+      if (loopBound->is<Attribute>()) {
+        if (auto bound =
+                loopBound->get<Attribute>().dyn_cast_or_null<IntegerAttr>())
+          return bound.getValue();
+      } else if (auto value = loopBound->dyn_cast<Value>()) {
+        const IntegerValueRangeLattice *lattice =
+            getLatticeElementFor(op, value);
+        if (lattice != nullptr)
+          return getUpper ? lattice->getValue().getValue().smax()
+                          : lattice->getValue().getValue().smin();
+      }
+    }
+    // Given the results of getConstant{Lower,Upper}Bound()
+    // or getConstantStep() on a LoopLikeInterface return the lower/upper
+    // bound
+    return getUpper ? APInt::getSignedMaxValue(width)
+                    : APInt::getSignedMinValue(width);
+  };
+
+  // Infer bounds for loop arguments that have static bounds
+  if (auto loop = dyn_cast<LoopLikeOpInterface>(op)) {
+    Optional<Value> iv = loop.getSingleInductionVar();
+    if (!iv) {
+      return SparseDataFlowAnalysis ::visitNonControlFlowArguments(
+          op, successor, argLattices, firstIndex);
+    }
+    Optional<OpFoldResult> lowerBound = loop.getSingleLowerBound();
+    Optional<OpFoldResult> upperBound = loop.getSingleUpperBound();
+    Optional<OpFoldResult> step = loop.getSingleStep();
+    APInt min = getLoopBoundFromFold(lowerBound, iv->getType(),
+                                     /*getUpper=*/false);
+    APInt max = getLoopBoundFromFold(upperBound, iv->getType(),
+                                     /*getUpper=*/true);
+    // Assume positivity for uniscoverable steps by way of getUpper = true.
+    APInt stepVal =
+        getLoopBoundFromFold(step, iv->getType(), /*getUpper=*/true);
+
+    if (stepVal.isNegative()) {
+      std::swap(min, max);
+    } else {
+      // Correct the upper bound by subtracting 1 so that it becomes a <=
+      // bound, because loops do not generally include their upper bound.
+      max -= 1;
+    }
+
+    IntegerValueRangeLattice *ivEntry = getLatticeElement(*iv);
+    auto ivRange = ConstantIntRanges::fromSigned(min, max);
+    propagateIfChanged(ivEntry, ivEntry->join(ivRange));
+    return;
+  }
+
+  return SparseDataFlowAnalysis::visitNonControlFlowArguments(
+      op, successor, argLattices, firstIndex);
+}

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

 //===- SparseAnalysis.cpp - Sparse data-flow analysis ---------------------===//
 //
 // 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/Analysis/DataFlow/SparseAnalysis.h"
 #include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
 #include "mlir/Interfaces/CallInterfaces.h"
-#include "mlir/Interfaces/ControlFlowInterfaces.h"
 
 using namespace mlir;
 using namespace mlir::dataflow;
 
 //===----------------------------------------------------------------------===//
 // AbstractSparseLattice
 //===----------------------------------------------------------------------===//
 
if (block->isEntryBlock()) {
 
     // Check if the lattices can be determined from region control flow.
     if (auto branch = dyn_cast<RegionBranchOpInterface>(block->getParentOp())) {
       return visitRegionSuccessors(
           block, branch, block->getParent()->getRegionNumber(), argLattices);
     }
 
     // Otherwise, we can't reason about the data-flow.
-    return markAllPessimisticFixpoint(argLattices);
+    return visitNonControlFlowArgumentsImpl(block->getParentOp(),
+                                            RegionSuccessor(block->getParent()),
+                                            argLattices, /*firstIndex=*/0);
   }
 
   // Iterate over the predecessors of the non-entry block.
   for (Block::pred_iterator it = block->pred_begin(), e = block->pred_end();
        it != e; ++it) {
     Block *predecessor = *it;
 
     // If the edge from the predecessor block to the current block is not live,
for (Operation *op : predecessors->getKnownPredecessors()) {
     // Get the incoming successor operands.
     Optional<OperandRange> operands;
 
     // Check if the predecessor is the parent op.
     if (op == branch) {
       operands = branch.getSuccessorEntryOperands(successorIndex);
       // Otherwise, try to deduce the operands from a region return-like op.
     } else {
-      assert(op->hasTrait<OpTrait::IsTerminator>() && "expected a terminator");
       if (isRegionReturnLike(op))
         operands = getRegionBranchSuccessorOperands(op, successorIndex);
     }
 
     if (!operands) {
       // We can't reason about the data-flow.
       return markAllPessimisticFixpoint(lattices);
     }
 
     ValueRange inputs = predecessors->getSuccessorInputs(op);
     assert(inputs.size() == operands->size() &&
            "expected the same number of successor inputs as operands");
 
-    // TODO: This was updated to be exposed upstream.
     unsigned firstIndex = 0;
     if (inputs.size() != lattices.size()) {
-      if (inputs.empty()) {
-        markAllPessimisticFixpoint(lattices);
-        return;
-      }
+      if (auto *op = point.dyn_cast<Operation *>()) {
+        if (!inputs.empty())
+          firstIndex = inputs.front().cast<OpResult>().getResultNumber();
+        visitNonControlFlowArgumentsImpl(
+            branch,
+            RegionSuccessor(
+                branch->getResults().slice(firstIndex, inputs.size())),
+            lattices, firstIndex);
+      } else {
+        if (!inputs.empty())
           firstIndex = inputs.front().cast<BlockArgument>().getArgNumber();
-      markAllPessimisticFixpoint(lattices.take_front(firstIndex));
-      markAllPessimisticFixpoint(
-          lattices.drop_front(firstIndex + inputs.size()));
+        Region *region = point.get<Block *>()->getParent();
+        visitNonControlFlowArgumentsImpl(
+            branch,
+            RegionSuccessor(region, region->getArguments().slice(
+                                        firstIndex, inputs.size())),
+            lattices, firstIndex);
+      }
     }
 
     for (auto it : llvm::zip(*operands, lattices.drop_front(firstIndex)))
       join(std::get<1>(it), *getLatticeElementFor(point, std::get<0>(it)));
   }
 }
 
 const AbstractSparseLattice *

mlir/lib/Analysis/DataFlowFramework.cpp

while (!worklist.empty()) {
       worklist.pop();
 
       DATAFLOW_DEBUG(llvm::dbgs() << "Invoking '" << analysis->debugName
                                   << "' on: " << point << "\n");
       if (failed(analysis->visit(point)))
         return failure();
     }
 
-    // "Nudge" the state of the analysis by forcefully initializing states that
-    // are still uninitialized. All uninitialized states in the graph can be
-    // initialized in any order because the analysis reached fixpoint, meaning
-    // that there are no work items that would have further nudged the analysis.
-    for (AnalysisState &state :
-         llvm::make_pointee_range(llvm::make_second_range(analysisStates))) {
-      if (!state.isUninitialized())
-        continue;
-      DATAFLOW_DEBUG(llvm::dbgs() << "Default initializing " << state.debugName
-                                  << " of " << state.point << "\n");
-      propagateIfChanged(&state, state.defaultInitialize());
-    }
-
     // Iterate until all states are in some initialized state and the worklist
     // is exhausted.
   } while (!worklist.empty());
 
   return success();
 }
 
 void DataFlowSolver::propagateIfChanged(AnalysisState *state,

mlir/lib/Analysis/IntRangeAnalysis.cpp

-//===- IntRangeAnalysis.cpp - Infer Ranges Interfaces --*- 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 defines the dataflow analysis class for integer range inference
-// which is used in transformations over the `arith` dialect such as
-// branch elimination or signed->unsigned rewriting
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Analysis/IntRangeAnalysis.h"
-#include "mlir/Analysis/DataFlowAnalysis.h"
-#include "mlir/Interfaces/InferIntRangeInterface.h"
-#include "mlir/Interfaces/LoopLikeInterface.h"
-#include "llvm/Support/Debug.h"
-
-#define DEBUG_TYPE "int-range-analysis"
-
-using namespace mlir;
-
-namespace {
-/// A wrapper around ConstantIntRanges that provides the lattice functions
-/// expected by dataflow analysis.
-struct IntRangeLattice {
-  IntRangeLattice(const ConstantIntRanges &value) : value(value){};
-  IntRangeLattice(ConstantIntRanges &&value) : value(value){};
-
-  bool operator==(const IntRangeLattice &other) const {
-    return value == other.value;
-  }
-
-  /// wrapper around rangeUnion()
-  static IntRangeLattice join(const IntRangeLattice &a,
-                              const IntRangeLattice &b) {
-    return a.value.rangeUnion(b.value);
-  }
-
-  /// Creates a range with bitwidth 0 to represent that we don't know if the
-  /// value being marked overdefined is even an integer.
-  static IntRangeLattice getPessimisticValueState(MLIRContext *context) {
-    APInt noIntValue = APInt::getZeroWidth();
-    return ConstantIntRanges(noIntValue, noIntValue, noIntValue, noIntValue);
-  }
-
-  /// Create a maximal range ([0, uint_max(t)] / [int_min(t), int_max(t)])
-  /// range that is used to mark the value v as unable to be analyzed further,
-  /// where t is the type of v.
-  static IntRangeLattice getPessimisticValueState(Value v) {
-    unsigned int width = ConstantIntRanges::getStorageBitwidth(v.getType());
-    APInt umin = APInt::getMinValue(width);
-    APInt umax = APInt::getMaxValue(width);
-    APInt smin = width != 0 ? APInt::getSignedMinValue(width) : umin;
-    APInt smax = width != 0 ? APInt::getSignedMaxValue(width) : umax;
-    return ConstantIntRanges{umin, umax, smin, smax};
-  }
-
-  ConstantIntRanges value;
-};
-} // end anonymous namespace
-
-namespace mlir {
-namespace detail {
-class IntRangeAnalysisImpl : public ForwardDataFlowAnalysis<IntRangeLattice> {
-  using ForwardDataFlowAnalysis<IntRangeLattice>::ForwardDataFlowAnalysis;
-
-public:
-  /// Define bounds on the results or block arguments of the operation
-  /// based on the bounds on the arguments given in `operands`
-  ChangeResult
-  visitOperation(Operation *op,
-                 ArrayRef<LatticeElement<IntRangeLattice> *> operands) final;
-
-  /// Skip regions of branch ops when we can statically infer constant
-  /// values for operands to the branch op and said op tells us it's safe to do
-  /// so.
-  LogicalResult
-  getSuccessorsForOperands(BranchOpInterface branch,
-                           ArrayRef<LatticeElement<IntRangeLattice> *> operands,
-                           SmallVectorImpl<Block *> &successors) final;
-
-  /// Skip regions of branch or loop ops when we can statically infer constant
-  /// values for operands to the branch op and said op tells us it's safe to do
-  /// so.
-  void
-  getSuccessorsForOperands(RegionBranchOpInterface branch,
-                           Optional<unsigned> sourceIndex,
-                           ArrayRef<LatticeElement<IntRangeLattice> *> operands,
-                           SmallVectorImpl<RegionSuccessor> &successors) final;
-
-  /// Call the InferIntRangeInterface implementation for region-using ops
-  /// that implement it, and infer the bounds of loop induction variables
-  /// for ops that implement LoopLikeOPInterface.
-  ChangeResult visitNonControlFlowArguments(
-      Operation *op, const RegionSuccessor &region,
-      ArrayRef<LatticeElement<IntRangeLattice> *> operands) final;
-};
-} // end namespace detail
-} // end namespace mlir
-
-/// Given the results of getConstant{Lower,Upper}Bound()
-/// or getConstantStep() on a LoopLikeInterface return the lower/upper bound for
-/// that result if possible.
-static APInt getLoopBoundFromFold(Optional<OpFoldResult> loopBound,
-                                  Type boundType,
-                                  detail::IntRangeAnalysisImpl &analysis,
-                                  bool getUpper) {
-  unsigned int width = ConstantIntRanges::getStorageBitwidth(boundType);
-  if (loopBound) {
-    if (loopBound->is<Attribute>()) {
-      if (auto bound =
-              loopBound->get<Attribute>().dyn_cast_or_null<IntegerAttr>())
-        return bound.getValue();
-    } else if (loopBound->is<Value>()) {
-      LatticeElement<IntRangeLattice> *lattice =
-          analysis.lookupLatticeElement(loopBound->get<Value>());
-      if (lattice != nullptr)
-        return getUpper ? lattice->getValue().value.smax()
-                        : lattice->getValue().value.smin();
-    }
-  }
-  return getUpper ? APInt::getSignedMaxValue(width)
-                  : APInt::getSignedMinValue(width);
-}
-
-ChangeResult detail::IntRangeAnalysisImpl::visitOperation(
-    Operation *op, ArrayRef<LatticeElement<IntRangeLattice> *> operands) {
-  ChangeResult result = ChangeResult::NoChange;
-  // Ignore non-integer outputs - return early if the op has no scalar
-  // integer results
-  bool hasIntegerResult = false;
-  for (Value v : op->getResults()) {
-    if (v.getType().isIntOrIndex())
-      hasIntegerResult = true;
-    else
-      result |= markAllPessimisticFixpoint(v);
-  }
-  if (!hasIntegerResult)
-    return result;
-
-  if (auto inferrable = dyn_cast<InferIntRangeInterface>(op)) {
-    LLVM_DEBUG(llvm::dbgs() << "Inferring ranges for ");
-    LLVM_DEBUG(inferrable->print(llvm::dbgs()));
-    LLVM_DEBUG(llvm::dbgs() << "\n");
-    SmallVector<ConstantIntRanges> argRanges(
-        llvm::map_range(operands, [](LatticeElement<IntRangeLattice> *val) {
-          return val->getValue().value;
-        }));
-
-    auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
-      LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
-      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(v);
-      Optional<IntRangeLattice> oldRange;
-      if (!lattice.isUninitialized())
-        oldRange = lattice.getValue();
-      result |= lattice.join(IntRangeLattice(attrs));
-
-      // Catch loop results with loop variant bounds and conservatively make
-      // them [-inf, inf] so we don't circle around infinitely often (because
-      // the dataflow analysis in MLIR doesn't attempt to work out trip counts
-      // and often can't).
-      bool isYieldedResult = llvm::any_of(v.getUsers(), [](Operation *op) {
-        return op->hasTrait<OpTrait::IsTerminator>();
-      });
-      if (isYieldedResult && oldRange && !(lattice.getValue() == *oldRange)) {
-        LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
-        result |= lattice.markPessimisticFixpoint();
-      }
-    };
-
-    inferrable.inferResultRanges(argRanges, joinCallback);
-    for (Value opResult : op->getResults()) {
-      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(opResult);
-      // setResultRange() not called, make pessimistic.
-      if (lattice.isUninitialized())
-        result |= lattice.markPessimisticFixpoint();
-    }
-  } else if (op->getNumRegions() == 0) {
-    // No regions + no result inference method -> unbounded results (ex. memory
-    // ops)
-    result |= markAllPessimisticFixpoint(op->getResults());
-  }
-  return result;
-}
-
-LogicalResult detail::IntRangeAnalysisImpl::getSuccessorsForOperands(
-    BranchOpInterface branch,
-    ArrayRef<LatticeElement<IntRangeLattice> *> operands,
-    SmallVectorImpl<Block *> &successors) {
-  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
-    Optional<APInt> maybeConstValue =
-        enumPair.value()->getValue().value.getConstantValue();
-
-    if (maybeConstValue) {
-      return IntegerAttr::get(branch->getOperand(enumPair.index()).getType(),
-                              *maybeConstValue);
-    }
-    return {};
-  };
-  SmallVector<Attribute> inferredConsts(
-      llvm::map_range(llvm::enumerate(operands), toConstantAttr));
-  if (Block *singleSucc = branch.getSuccessorForOperands(inferredConsts)) {
-    successors.push_back(singleSucc);
-    return success();
-  }
-  return failure();
-}
-
-void detail::IntRangeAnalysisImpl::getSuccessorsForOperands(
-    RegionBranchOpInterface branch, Optional<unsigned> sourceIndex,
-    ArrayRef<LatticeElement<IntRangeLattice> *> operands,
-    SmallVectorImpl<RegionSuccessor> &successors) {
-  // Get a type with which to construct a constant.
-  auto getOperandType = [branch, sourceIndex](unsigned index) {
-    // The types of all return-like operations are the same.
-    if (!sourceIndex)
-      return branch->getOperand(index).getType();
-
-    for (Block &block : branch->getRegion(*sourceIndex)) {
-      Operation *terminator = block.getTerminator();
-      if (getRegionBranchSuccessorOperands(terminator, *sourceIndex))
-        return terminator->getOperand(index).getType();
-    }
-    return Type();
-  };
-
-  auto toConstantAttr = [&getOperandType](auto enumPair) -> Attribute {
-    if (Optional<APInt> maybeConstValue =
-            enumPair.value()->getValue().value.getConstantValue()) {
-      return IntegerAttr::get(getOperandType(enumPair.index()),
-                              *maybeConstValue);
-    }
-    return {};
-  };
-  SmallVector<Attribute> inferredConsts(
-      llvm::map_range(llvm::enumerate(operands), toConstantAttr));
-  branch.getSuccessorRegions(sourceIndex, inferredConsts, successors);
-}
-
-ChangeResult detail::IntRangeAnalysisImpl::visitNonControlFlowArguments(
-    Operation *op, const RegionSuccessor &region,
-    ArrayRef<LatticeElement<IntRangeLattice> *> operands) {
-  if (auto inferrable = dyn_cast<InferIntRangeInterface>(op)) {
-    LLVM_DEBUG(llvm::dbgs() << "Inferring ranges for ");
-    LLVM_DEBUG(inferrable->print(llvm::dbgs()));
-    LLVM_DEBUG(llvm::dbgs() << "\n");
-    SmallVector<ConstantIntRanges> argRanges(
-        llvm::map_range(operands, [](LatticeElement<IntRangeLattice> *val) {
-          return val->getValue().value;
-        }));
-
-    ChangeResult result = ChangeResult::NoChange;
-    auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
-      LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
-      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(v);
-      Optional<IntRangeLattice> oldRange;
-      if (!lattice.isUninitialized())
-        oldRange = lattice.getValue();
-      result |= lattice.join(IntRangeLattice(attrs));
-
-      // Catch loop results with loop variant bounds and conservatively make
-      // them [-inf, inf] so we don't circle around infinitely often (because
-      // the dataflow analysis in MLIR doesn't attempt to work out trip counts
-      // and often can't).
-      bool isYieldedValue = llvm::any_of(v.getUsers(), [](Operation *op) {
-        return op->hasTrait<OpTrait::IsTerminator>();
-      });
-      if (isYieldedValue && oldRange && !(lattice.getValue() == *oldRange)) {
-        LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
-        result |= lattice.markPessimisticFixpoint();
-      }
-    };
-
-    inferrable.inferResultRanges(argRanges, joinCallback);
-    for (Value regionArg : region.getSuccessor()->getArguments()) {
-      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(regionArg);
-      // setResultRange() not called, make pessimistic.
-      if (lattice.isUninitialized())
-        result |= lattice.markPessimisticFixpoint();
-    }
-
-    return result;
-  }
-
-  // Infer bounds for loop arguments that have static bounds
-  if (auto loop = dyn_cast<LoopLikeOpInterface>(op)) {
-    Optional<Value> iv = loop.getSingleInductionVar();
-    if (!iv) {
-      return ForwardDataFlowAnalysis<
-          IntRangeLattice>::visitNonControlFlowArguments(op, region, operands);
-    }
-    Optional<OpFoldResult> lowerBound = loop.getSingleLowerBound();
-    Optional<OpFoldResult> upperBound = loop.getSingleUpperBound();
-    Optional<OpFoldResult> step = loop.getSingleStep();
-    APInt min = getLoopBoundFromFold(lowerBound, iv->getType(), *this,
-                                     /*getUpper=*/false);
-    APInt max = getLoopBoundFromFold(upperBound, iv->getType(), *this,
-                                     /*getUpper=*/true);
-    // Assume positivity for uniscoverable steps by way of getUpper = true.
-    APInt stepVal =
-        getLoopBoundFromFold(step, iv->getType(), *this, /*getUpper=*/true);
-
-    if (stepVal.isNegative()) {
-      std::swap(min, max);
-    } else {
-      // Correct the upper bound by subtracting 1 so that it becomes a <= bound,
-      // because loops do not generally include their upper bound.
-      max -= 1;
-    }
-
-    LatticeElement<IntRangeLattice> &ivEntry = getLatticeElement(*iv);
-    return ivEntry.join(ConstantIntRanges::fromSigned(min, max));
-  }
-  return ForwardDataFlowAnalysis<IntRangeLattice>::visitNonControlFlowArguments(
-      op, region, operands);
-}
-
-IntRangeAnalysis::IntRangeAnalysis(Operation *topLevelOperation) {
-  impl = std::make_unique<mlir::detail::IntRangeAnalysisImpl>(
-      topLevelOperation->getContext());
-  impl->run(topLevelOperation);
-}
-
-IntRangeAnalysis::~IntRangeAnalysis() = default;
-IntRangeAnalysis::IntRangeAnalysis(IntRangeAnalysis &&other) = default;
-
-Optional<ConstantIntRanges> IntRangeAnalysis::getResult(Value v) {
-  LatticeElement<IntRangeLattice> *result = impl->lookupLatticeElement(v);
-  if (result == nullptr || result->isUninitialized())
-    return llvm::None;
-  return result->getValue().value;
-}

mlir/lib/Dialect/Arithmetic/Transforms/UnsignedWhenEquivalent.cpp

 //===- UnsignedWhenEquivalent.cpp - Pass to replace signed operations with
 // unsigned
 // ones when all their arguments and results are statically non-negative --===//
 //
 // 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 "PassDetail.h"
-#include "mlir/Analysis/IntRangeAnalysis.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::arith;
+using namespace mlir::dataflow;
 
 /// Succeeds when a value is statically non-negative in that it has a lower
 /// bound on its value (if it is treated as signed) and that bound is
 /// non-negative.
-static LogicalResult staticallyNonNegative(IntRangeAnalysis &analysis,
-                                           Value v) {
-  Optional<ConstantIntRanges> result = analysis.getResult(v);
-  if (!result.hasValue())
+static LogicalResult staticallyNonNegative(DataFlowSolver &solver, Value v) {
+  auto *result = solver.lookupState<IntegerValueRangeLattice>(v);
+  if (!result)
     return failure();
-  const ConstantIntRanges &range = result.getValue();
+  const ConstantIntRanges &range = result->getValue().getValue();
   return success(range.smin().isNonNegative());
 }
 
 /// Succeeds if an op can be converted to its unsigned equivalent without
 /// changing its semantics. This is the case when none of its openands or
 /// results can be below 0 when analyzed from a signed perspective.
-static LogicalResult staticallyNonNegative(IntRangeAnalysis &analysis,
+static LogicalResult staticallyNonNegative(DataFlowSolver &solver,
                                            Operation *op) {
-  auto nonNegativePred = [&analysis](Value v) -> bool {
-    return succeeded(staticallyNonNegative(analysis, v));
+  auto nonNegativePred = [&solver](Value v) -> bool {
+    return succeeded(staticallyNonNegative(solver, v));
   };
   return success(llvm::all_of(op->getOperands(), nonNegativePred) &&
                  llvm::all_of(op->getResults(), nonNegativePred));
 }
 
 /// Succeeds when the comparison predicate is a signed operation and all the
 /// operands are non-negative, indicating that the cmpi operation `op` can have
 /// its predicate changed to an unsigned equivalent.
-static LogicalResult isCmpIConvertable(IntRangeAnalysis &analysis, CmpIOp op) {
+static LogicalResult isCmpIConvertable(DataFlowSolver &solver, CmpIOp op) {
   CmpIPredicate pred = op.getPredicate();
   switch (pred) {
   case CmpIPredicate::sle:
   case CmpIPredicate::slt:
   case CmpIPredicate::sge:
   case CmpIPredicate::sgt:
-    return success(llvm::all_of(op.getOperands(), [&analysis](Value v) -> bool {
-      return succeeded(staticallyNonNegative(analysis, v));
+    return success(llvm::all_of(op.getOperands(), [&solver](Value v) -> bool {
+      return succeeded(staticallyNonNegative(solver, v));
     }));
   default:
     return failure();
   }
 }
 
 /// Return the unsigned equivalent of a signed comparison predicate,
 /// or the predicate itself if there is none.
struct ArithmeticUnsignedWhenEquivalentPass
     : public ArithmeticUnsignedWhenEquivalentBase<
           ArithmeticUnsignedWhenEquivalentPass> {
   /// Implementation structure: first find all equivalent ops and collect them,
   /// then perform all the rewrites in a second pass over the target op. This
   /// ensures that analysis results are not invalidated during rewriting.
   void runOnOperation() override {
     Operation *op = getOperation();
     MLIRContext *ctx = op->getContext();
-    IntRangeAnalysis analysis(op);
+    DataFlowSolver solver;
+    solver.load<DeadCodeAnalysis>();
+    solver.load<IntegerRangeAnalysis>();
+    if (failed(solver.initializeAndRun(op)))
+      return signalPassFailure();
 
     ConversionTarget target(*ctx);
     target.addLegalDialect<ArithmeticDialect>();
     target
         .addDynamicallyLegalOp<DivSIOp, CeilDivSIOp, CeilDivUIOp, FloorDivSIOp,
                                RemSIOp, MinSIOp, MaxSIOp, ExtSIOp>(
-            [&analysis](Operation *op) -> Optional<bool> {
-              return failed(staticallyNonNegative(analysis, op));
+            [&solver](Operation *op) -> Optional<bool> {
+              return failed(staticallyNonNegative(solver, op));
             });
     target.addDynamicallyLegalOp<CmpIOp>(
-        [&analysis](CmpIOp op) -> Optional<bool> {
-          return failed(isCmpIConvertable(analysis, op));
+        [&solver](CmpIOp op) -> Optional<bool> {
+          return failed(isCmpIConvertable(solver, op));
         });
 
     RewritePatternSet patterns(ctx);
     patterns.add<ConvertOpToUnsigned<DivSIOp, DivUIOp>,
                  ConvertOpToUnsigned<CeilDivSIOp, CeilDivUIOp>,
                  ConvertOpToUnsigned<FloorDivSIOp, DivUIOp>,
                  ConvertOpToUnsigned<RemSIOp, RemUIOp>,
                  ConvertOpToUnsigned<MinSIOp, MinUIOp>,

mlir/lib/Transforms/SCCP.cpp

 
 /// Replace the given value with a constant if the corresponding lattice
 /// represents a constant. Returns success if the value was replaced, failure
 /// otherwise.
 static LogicalResult replaceWithConstant(DataFlowSolver &solver,
                                          OpBuilder &builder,
                                          OperationFolder &folder, Value value) {
   auto *lattice = solver.lookupState<Lattice<ConstantValue>>(value);
-  if (!lattice)
+  if (!lattice || lattice->isUninitialized())
     return failure();
   const ConstantValue &latticeValue = lattice->getValue();
   if (!latticeValue.getConstantValue())
     return failure();
 
   // Attempt to materialize a constant for the given value.
   Dialect *dialect = latticeValue.getConstantDialect();
   Value constant = folder.getOrCreateConstant(builder, dialect,

mlir/test/lib/Analysis/DataFlow/TestDeadCodeAnalysis.cpp

 namespace {
 /// This is a simple analysis that implements a transfer function for constant
 /// operations.
 struct ConstantAnalysis : public DataFlowAnalysis {
   using DataFlowAnalysis::DataFlowAnalysis;
 
   LogicalResult initialize(Operation *top) override {
     WalkResult result = top->walk([&](Operation *op) {
-      if (op->hasTrait<OpTrait::ConstantLike>())
-        if (failed(visit(op)))
-          return WalkResult::interrupt();
       return WalkResult::advance();
     });
     return success(!result.wasInterrupted());
   }
 
   LogicalResult visit(ProgramPoint point) override {
     Operation *op = point.get<Operation *>();
     Attribute value;
     if (matchPattern(op, m_Constant(&value))) {
       auto *constant = getOrCreate<Lattice<ConstantValue>>(op->getResult(0));
       propagateIfChanged(
           constant, constant->join(ConstantValue(value, op->getDialect())));
+      return success();
     }
+    markAllPessimisticFixpoint(op->getResults());
+    for (Region &region : op->getRegions())
+      markAllPessimisticFixpoint(region.getArguments());
     return success();
   }
+
+  /// Mark the constant values of all given values as having reached a
+  /// pessimistic fixpoint.
+  void markAllPessimisticFixpoint(ValueRange values) {
+    for (Value value : values) {
+      auto *constantValue = getOrCreate<Lattice<ConstantValue>>(value);
+      propagateIfChanged(constantValue,
+                         constantValue->markPessimisticFixpoint());
+    }
+  }
 };
 
-/// This is a simple pass that runs dead code analysis with no constant value
-/// provider. It marks everything as live.
+/// This is a simple pass that runs dead code analysis with a constant value
+/// provider that only understands constant operations.
 struct TestDeadCodeAnalysisPass
     : public PassWrapper<TestDeadCodeAnalysisPass, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDeadCodeAnalysisPass)
 
   StringRef getArgument() const override { return "test-dead-code-analysis"; }
 
   void runOnOperation() override {
     Operation *op = getOperation();

mlir/test/lib/Transforms/TestIntRangeInference.cpp

 //===- TestIntRangeInference.cpp - Create consts from range inference ---===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 // TODO: This pass is needed to test integer range inference until that
 // functionality has been integrated into SCCP.
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Analysis/IntRangeAnalysis.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/TypeID.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 using namespace mlir;
+using namespace mlir::dataflow;
 
 /// Patterned after SCCP
-static LogicalResult replaceWithConstant(IntRangeAnalysis &analysis,
-                                         OpBuilder &b, OperationFolder &folder,
-                                         Value value) {
-  Optional<ConstantIntRanges> maybeInferredRange = analysis.getResult(value);
-  if (!maybeInferredRange)
+static LogicalResult replaceWithConstant(DataFlowSolver &solver, OpBuilder &b,
+                                         OperationFolder &folder, Value value) {
+  auto *maybeInferredRange =
+      solver.lookupState<IntegerValueRangeLattice>(value);
+  if (!maybeInferredRange || maybeInferredRange->isUninitialized())
     return failure();
-  const ConstantIntRanges &inferredRange = maybeInferredRange.getValue();
+  const ConstantIntRanges &inferredRange =
+      maybeInferredRange->getValue().getValue();
   Optional<APInt> maybeConstValue = inferredRange.getConstantValue();
   if (!maybeConstValue.hasValue())
     return failure();
 
   Operation *maybeDefiningOp = value.getDefiningOp();
   Dialect *valueDialect =
       maybeDefiningOp ? maybeDefiningOp->getDialect()
                       : value.getParentRegion()->getParentOp()->getDialect();
   Attribute constAttr = b.getIntegerAttr(value.getType(), *maybeConstValue);
   Value constant = folder.getOrCreateConstant(b, valueDialect, constAttr,
                                               value.getType(), value.getLoc());
   if (!constant)
     return failure();
 
   value.replaceAllUsesWith(constant);
   return success();
 }
 
-static void rewrite(IntRangeAnalysis &analysis, MLIRContext *context,
+static void rewrite(DataFlowSolver &solver, MLIRContext *context,
                     MutableArrayRef<Region> initialRegions) {
   SmallVector<Block *> worklist;
   auto addToWorklist = [&](MutableArrayRef<Region> regions) {
     for (Region &region : regions)
       for (Block &block : llvm::reverse(region))
         worklist.push_back(&block);
   };
 
   OpBuilder builder(context);
   OperationFolder folder(context);
 
   addToWorklist(initialRegions);
   while (!worklist.empty()) {
     Block *block = worklist.pop_back_val();
 
     for (Operation &op : llvm::make_early_inc_range(*block)) {
       builder.setInsertionPoint(&op);
 
       // Replace any result with constants.
       bool replacedAll = op.getNumResults() != 0;
       for (Value res : op.getResults())
         replacedAll &=
-            succeeded(replaceWithConstant(analysis, builder, folder, res));
+            succeeded(replaceWithConstant(solver, builder, folder, res));
 
       // If all of the results of the operation were replaced, try to erase
       // the operation completely.
       if (replacedAll && wouldOpBeTriviallyDead(&op)) {
         assert(op.use_empty() && "expected all uses to be replaced");
         op.erase();
         continue;
       }
 
       // Add any the regions of this operation to the worklist.
       addToWorklist(op.getRegions());
     }
 
     // Replace any block arguments with constants.
     builder.setInsertionPointToStart(block);
     for (BlockArgument arg : block->getArguments())
-      (void)replaceWithConstant(analysis, builder, folder, arg);
+      (void)replaceWithConstant(solver, builder, folder, arg);
   }
 }
 
 namespace {
 struct TestIntRangeInference
     : PassWrapper<TestIntRangeInference, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestIntRangeInference)
 
   StringRef getArgument() const final { return "test-int-range-inference"; }
   StringRef getDescription() const final {
     return "Test integer range inference analysis";
   }
 
   void runOnOperation() override {
     Operation *op = getOperation();
-    IntRangeAnalysis analysis(op);
-    rewrite(analysis, op->getContext(), op->getRegions());
+    DataFlowSolver solver;
+    solver.load<DeadCodeAnalysis>();
+    solver.load<IntegerRangeAnalysis>();
+    if (failed(solver.initializeAndRun(op)))
+      return signalPassFailure();
+    rewrite(solver, op->getContext(), op->getRegions());
   }
 };
 } // end anonymous namespace
 
 namespace mlir {
 namespace test {
 void registerTestIntRangeInference() {
   PassRegistration<TestIntRangeInference>();
 }
 } // end namespace test
 } // end namespace mlir