[mlir] [dataflow] : Improve the time and space footprint of data flow.

mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h

@@ -73,6 +73,14 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
   /// may modify the program state; that is, every operation and block.
   LogicalResult initialize(Operation *top) override;
 
+  /// Initialize lattice anchor equivalence class from the provided top-level
+  /// operation.
+  ///
+  /// This function will union lattice anchor to same equivalent class if the
+  /// analysis can determine the lattice content of lattice anchor is
+  /// necessarily identical under the corrensponding lattice type.
+  virtual void initializeEquivalentLatticeAnchor(Operation *top) override;
+
   /// Visit a program point that modifies the state of the program. If the
   /// program point is at the beginning of a block, then the state is propagated
   /// from control-flow predecessors or callsites.  If the operation before
@@ -114,6 +122,11 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
   /// operation transfer function.
   virtual LogicalResult processOperation(Operation *op);
 
+  /// Visit an operation. If this analysis can confirm that lattice content
+  /// of lattice anchors around operation are necessarily identical, join
+  /// them into the same equivalent class.
+  virtual void buildOperationEquivalentLatticeAnchor(Operation *op) { return; }
+
   /// Propagate the dense lattice forward along the control flow edge from
   /// `regionFrom` to `regionTo` regions of the `branch` operation. `nullopt`
   /// values correspond to control flow branches originating at or targeting the
@@ -310,6 +323,14 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// may modify the program state; that is, every operation and block.
   LogicalResult initialize(Operation *top) override;
 
+  /// Initialize lattice anchor equivalence class from the provided top-level
+  /// operation.
+  ///
+  /// This function will union lattice anchor to same equivalent class if the
+  /// analysis can determine the lattice content of lattice anchor is
+  /// necessarily identical under the corrensponding lattice type.
+  virtual void initializeEquivalentLatticeAnchor(Operation *top) override;
+
   /// Visit a program point that modifies the state of the program. The state is
   /// propagated along control flow directions for branch-, region- and
   /// call-based control flow using the respective interfaces. For other
@@ -353,6 +374,11 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// transfer function.
   virtual LogicalResult processOperation(Operation *op);
 
+  /// Visit an operation. If this analysis can confirm that lattice content
+  /// of lattice anchors around operation are necessarily identical, join
+  /// them into the same equivalent class.
+  virtual void buildOperationEquivalentLatticeAnchor(Operation *op) { return; }
+
   /// Propagate the dense lattice backwards along the control flow edge from
   /// `regionFrom` to `regionTo` regions of the `branch` operation. `nullopt`
   /// values correspond to control flow branches originating at or targeting the

mlir/include/mlir/Analysis/DataFlowFramework.h

@@ -18,6 +18,7 @@
 
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/StorageUniquer.h"
+#include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/Compiler.h"
@@ -265,6 +266,14 @@ struct LatticeAnchor
 /// Forward declaration of the data-flow analysis class.
 class DataFlowAnalysis;
 
+} // namespace mlir
+
+template <>
+struct llvm::DenseMapInfo<mlir::LatticeAnchor>
+    : public llvm::DenseMapInfo<mlir::LatticeAnchor::ParentTy> {};
+
+namespace mlir {
+
 //===----------------------------------------------------------------------===//
 // DataFlowConfig
 //===----------------------------------------------------------------------===//
@@ -332,7 +341,9 @@ class DataFlowSolver {
   /// does not exist.
   template <typename StateT, typename AnchorT>
   const StateT *lookupState(AnchorT anchor) const {
-    const auto &mapIt = analysisStates.find(LatticeAnchor(anchor));
+    LatticeAnchor latticeAnchor =
+        getLeaderAnchorOrSelf<StateT>(LatticeAnchor(anchor));
+    const auto &mapIt = analysisStates.find(latticeAnchor);
     if (mapIt == analysisStates.end())
       return nullptr;
     auto it = mapIt->second.find(TypeID::get<StateT>());
@@ -344,12 +355,34 @@ class DataFlowSolver {
   /// Erase any analysis state associated with the given lattice anchor.
   template <typename AnchorT>
   void eraseState(AnchorT anchor) {
-    LatticeAnchor la(anchor);
-    analysisStates.erase(LatticeAnchor(anchor));
+    LatticeAnchor latticeAnchor(anchor);
+
+    // Update equivalentAnchorMap.
+    for (auto &&[TypeId, eqClass] : equivalentAnchorMap) {
+      if (!eqClass.contains(latticeAnchor)) {
+        continue;
+      }
+      llvm::EquivalenceClasses<LatticeAnchor>::member_iterator leaderIt =
+          eqClass.findLeader(latticeAnchor);
+
+      // Update analysis states with new leader if needed.
+      if (*leaderIt == latticeAnchor && ++leaderIt != eqClass.member_end()) {
+        analysisStates[*leaderIt][TypeId] =
+            std::move(analysisStates[latticeAnchor][TypeId]);
+      }
+
+      eqClass.erase(latticeAnchor);
+    }
+
+    // Update analysis states.
+    analysisStates.erase(latticeAnchor);
   }
 
-  // Erase all analysis states
-  void eraseAllStates() { analysisStates.clear(); }
+  // Erase all analysis states.
+  void eraseAllStates() {
+    analysisStates.clear();
+    equivalentAnchorMap.clear();
+  }
 
   /// Get a uniqued lattice anchor instance. If one is not present, it is
   /// created with the provided arguments.
@@ -399,6 +432,20 @@ class DataFlowSolver {
   template <typename StateT, typename AnchorT>
   StateT *getOrCreateState(AnchorT anchor);
 
+  /// Get leader lattice anchor in equivalence lattice anchor group, return
+  /// input lattice anchor if input not found in equivalece lattice anchor
+  /// group.
+  template <typename StateT>
+  LatticeAnchor getLeaderAnchorOrSelf(LatticeAnchor latticeAnchor) const;
+
+  /// Union input anchors under the given state.
+  template <typename StateT, typename AnchorT>
+  void unionLatticeAnchors(AnchorT anchor, AnchorT other);
+
+  /// Return given lattice is equivalent on given state.
+  template <typename StateT>
+  bool isEquivalent(LatticeAnchor lhs, LatticeAnchor rhs) const;
+
   /// Propagate an update to an analysis state if it changed by pushing
   /// dependent work items to the back of the queue.
   /// This should only be used when DataFlowSolver is running.
@@ -429,10 +476,15 @@ class DataFlowSolver {
 
   /// A type-erased map of lattice anchors to associated analysis states for
   /// first-class lattice anchors.
-  DenseMap<LatticeAnchor, DenseMap<TypeID, std::unique_ptr<AnalysisState>>,
-           DenseMapInfo<LatticeAnchor::ParentTy>>
+  DenseMap<LatticeAnchor, DenseMap<TypeID, std::unique_ptr<AnalysisState>>>
       analysisStates;
 
+  /// A type-erased map of lattice type to the equivalet lattice anchors.
+  /// Lattice anchors are considered equivalent under a certain lattice type if
+  /// and only if, under this lattice type, the lattices pointed to by these
+  /// lattice anchors necessarily contain identical value.
+  DenseMap<TypeID, llvm::EquivalenceClasses<LatticeAnchor>> equivalentAnchorMap;
+
   /// Allow the base child analysis class to access the internals of the solver.
   friend class DataFlowAnalysis;
 };
@@ -564,6 +616,14 @@ class DataFlowAnalysis {
   /// will provide a value for then.
   virtual LogicalResult visit(ProgramPoint *point) = 0;
 
+  /// Initialize lattice anchor equivalence class from the provided top-level
+  /// operation.
+  ///
+  /// This function will union lattice anchor to same equivalent class if the
+  /// analysis can determine the lattice content of lattice anchor is
+  /// necessarily identical under the corrensponding lattice type.
+  virtual void initializeEquivalentLatticeAnchor(Operation *top) { return; }
+
 protected:
   /// Create a dependency between the given analysis state and lattice anchor
   /// on this analysis.
@@ -584,6 +644,12 @@ class DataFlowAnalysis {
     return solver.getLatticeAnchor<AnchorT>(std::forward<Args>(args)...);
   }
 
+  /// Union input anchors under the given state.
+  template <typename StateT, typename AnchorT>
+  void unionLatticeAnchors(AnchorT anchor, AnchorT other) {
+    return solver.unionLatticeAnchors<StateT>(anchor, other);
+  }
+
   /// Get the analysis state associated with the lattice anchor. The returned
   /// state is expected to be "write-only", and any updates need to be
   /// propagated by `propagateIfChanged`.
@@ -598,7 +664,9 @@ class DataFlowAnalysis {
   template <typename StateT, typename AnchorT>
   const StateT *getOrCreateFor(ProgramPoint *dependent, AnchorT anchor) {
     StateT *state = getOrCreate<StateT>(anchor);
-    addDependency(state, dependent);
+    if (!solver.isEquivalent<StateT>(LatticeAnchor(anchor),
+                                     LatticeAnchor(dependent)))
+      addDependency(state, dependent);
     return state;
   }
 
@@ -644,10 +712,26 @@ AnalysisT *DataFlowSolver::load(Args &&...args) {
   return static_cast<AnalysisT *>(childAnalyses.back().get());
 }
 
+template <typename StateT>
+LatticeAnchor
+DataFlowSolver::getLeaderAnchorOrSelf(LatticeAnchor latticeAnchor) const {
+  const llvm::EquivalenceClasses<LatticeAnchor> eqClass =
+      equivalentAnchorMap.lookup(TypeID::get<StateT>());
+  llvm::EquivalenceClasses<LatticeAnchor>::member_iterator leaderIt =
+      eqClass.findLeader(latticeAnchor);
+  if (leaderIt != eqClass.member_end()) {
+    return *leaderIt;
+  }
+  return latticeAnchor;
+}
+
 template <typename StateT, typename AnchorT>
 StateT *DataFlowSolver::getOrCreateState(AnchorT anchor) {
+  // Replace to leader anchor if found.
+  LatticeAnchor latticeAnchor(anchor);
+  latticeAnchor = getLeaderAnchorOrSelf<StateT>(latticeAnchor);
   std::unique_ptr<AnalysisState> &state =
-      analysisStates[LatticeAnchor(anchor)][TypeID::get<StateT>()];
+      analysisStates[latticeAnchor][TypeID::get<StateT>()];
   if (!state) {
     state = std::unique_ptr<StateT>(new StateT(anchor));
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
@@ -657,6 +741,22 @@ StateT *DataFlowSolver::getOrCreateState(AnchorT anchor) {
   return static_cast<StateT *>(state.get());
 }
 
+template <typename StateT>
+bool DataFlowSolver::isEquivalent(LatticeAnchor lhs, LatticeAnchor rhs) const {
+  const llvm::EquivalenceClasses<LatticeAnchor> eqClass =
+      equivalentAnchorMap.lookup(TypeID::get<StateT>());
+  if (!eqClass.contains(lhs) || !eqClass.contains(rhs))
+    return false;
+  return eqClass.isEquivalent(lhs, rhs);
+}
+
+template <typename StateT, typename AnchorT>
+void DataFlowSolver::unionLatticeAnchors(AnchorT anchor, AnchorT other) {
+  llvm::EquivalenceClasses<LatticeAnchor> &eqClass =
+      equivalentAnchorMap[TypeID::get<StateT>()];
+  eqClass.unionSets(LatticeAnchor(anchor), LatticeAnchor(other));
+}
+
 inline raw_ostream &operator<<(raw_ostream &os, const AnalysisState &state) {
   state.print(os);
   return os;

mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp

@@ -28,6 +28,15 @@ using namespace mlir::dataflow;
 // AbstractDenseForwardDataFlowAnalysis
 //===----------------------------------------------------------------------===//
 
+void AbstractDenseForwardDataFlowAnalysis::initializeEquivalentLatticeAnchor(
+    Operation *top) {
+  top->walk([&](Operation *op) {
+    if (isa<RegionBranchOpInterface, CallOpInterface>(op))
+      return;
+    buildOperationEquivalentLatticeAnchor(op);
+  });
+}
+
 LogicalResult AbstractDenseForwardDataFlowAnalysis::initialize(Operation *top) {
   // Visit every operation and block.
   if (failed(processOperation(top)))
@@ -252,6 +261,15 @@ AbstractDenseForwardDataFlowAnalysis::getLatticeFor(ProgramPoint *dependent,
 // AbstractDenseBackwardDataFlowAnalysis
 //===----------------------------------------------------------------------===//
 
+void AbstractDenseBackwardDataFlowAnalysis::initializeEquivalentLatticeAnchor(
+    Operation *top) {
+  top->walk([&](Operation *op) {
+    if (isa<RegionBranchOpInterface, CallOpInterface>(op))
+      return;
+    buildOperationEquivalentLatticeAnchor(op);
+  });
+}
+
 LogicalResult
 AbstractDenseBackwardDataFlowAnalysis::initialize(Operation *top) {
   // Visit every operation and block.

mlir/lib/Analysis/DataFlowFramework.cpp

@@ -109,6 +109,11 @@ LogicalResult DataFlowSolver::initializeAndRun(Operation *top) {
   isRunning = true;
   auto guard = llvm::make_scope_exit([&]() { isRunning = false; });
 
+  // Initialize equivalent lattice anchors.
+  for (DataFlowAnalysis &analysis : llvm::make_pointee_range(childAnalyses)) {
+    analysis.initializeEquivalentLatticeAnchor(top);
+  }
+
   // Initialize the analyses.
   for (DataFlowAnalysis &analysis : llvm::make_pointee_range(childAnalyses)) {
     DATAFLOW_DEBUG(llvm::dbgs()

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

@@ -76,6 +76,11 @@ class NextAccessAnalysis : public DenseBackwardDataFlowAnalysis<NextAccess> {
     propagateIfChanged(lattice, lattice->setKnownToUnknown());
   }
 
+  /// Visit an operation. If this analysis can confirm that lattice content
+  /// of lattice anchors around operation are necessarily identical, join
+  /// them into the same equivalent class.
+  void buildOperationEquivalentLatticeAnchor(Operation *op) override;
+
   const bool assumeFuncReads;
 };
 } // namespace
@@ -141,6 +146,13 @@ LogicalResult NextAccessAnalysis::visitOperation(Operation *op,
   return success();
 }
 
+void NextAccessAnalysis::buildOperationEquivalentLatticeAnchor(Operation *op) {
+  if (isMemoryEffectFree(op)) {
+    unionLatticeAnchors<NextAccess>(getProgramPointBefore(op),
+                                    getProgramPointAfter(op));
+  }
+}
+
 void NextAccessAnalysis::visitCallControlFlowTransfer(
     CallOpInterface call, CallControlFlowAction action, const NextAccess &after,
     NextAccess *before) {

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

@@ -72,6 +72,11 @@ class LastModifiedAnalysis
                                             const LastModification &before,
                                             LastModification *after) override;
 
+  /// Visit an operation. If this analysis can confirm that lattice content
+  /// of lattice anchors around operation are necessarily identical, join
+  /// them into the same equivalent class.
+  void buildOperationEquivalentLatticeAnchor(Operation *op) override;
+
   /// At an entry point, the last modifications of all memory resources are
   /// unknown.
   void setToEntryState(LastModification *lattice) override {
@@ -147,6 +152,14 @@ LogicalResult LastModifiedAnalysis::visitOperation(
   return success();
 }
 
+void LastModifiedAnalysis::buildOperationEquivalentLatticeAnchor(
+    Operation *op) {
+  if (isMemoryEffectFree(op)) {
+    unionLatticeAnchors<LastModification>(getProgramPointBefore(op),
+                                          getProgramPointAfter(op));
+  }
+}
+
 void LastModifiedAnalysis::visitCallControlFlowTransfer(
     CallOpInterface call, CallControlFlowAction action,
     const LastModification &before, LastModification *after) {