[mlir][dataflow] Unify dependency management in AnalysisState.

[phisiart]

Summary

In the MLIR dataflow analysis framework, when an AnalysisState is updated, it's dependents are enqueued to be visited.

Currently, there are two ways dependents are managed:

• AnalysisState::dependents stores a list of dependents. DataFlowSolver::propagateIfChanged() reads this list and enqueues them to the worklist.

• AnalysisState::onUpdate() allows custom logic to enqueue more to the worklist. This is called by DataFlowSolver::propagateIfChanged().

This cleanup diff consolidates the two into AnalysisState::onUpdate(). This way, DataFlowSolver does not need to know the detail about AnalysisState::dependents, and the logic of dependency management is entirely handled by AnalysisState.

[Mogball]

Much better, thanks!

mlir/include/mlir/Analysis/DataFlowFramework.h

302

Please spell out this auto. The MLIR guidelines for using auto are that:

1. The type is visible verbatim on the RHS, like auto thing = dyn_cast<Thing>(foo) or
2. It's "too complicated", like std::vector<FooBar>::iterator.

[phisiart]

Explicitly spell out WorkItem instead of auto.

[phisiart]

@Mogball Hi, I don't have commit access. Could you commit this patch for me? Thanks!

[Mogball]

@phisiart Yes. Also you should email chris and ask for commit access

mlir/include/mlir/Analysis/DataFlowFramework.h

public:
   /// 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);
 
-  /// Add a dependency to an analysis state on a child analysis and program
-  /// point. If the state is updated, the child analysis must be invoked on the
-  /// given program point again.
-  void addDependency(AnalysisState *state, DataFlowAnalysis *analysis,
-                     ProgramPoint point);
-
 private:
   /// The solver's work queue. Work items can be inserted to the front of the
   /// queue to be processed greedily, speeding up computations that otherwise
   /// quickly degenerate to quadratic due to propagation of state updates.
   std::queue<WorkItem> worklist;
 
   /// Type-erased instances of the children analyses.
   SmallVector<std::unique_ptr<DataFlowAnalysis>> childAnalyses;
public:
   AnalysisState(ProgramPoint point) : point(point) {}
 
   /// Returns the program point this static is located at.
   ProgramPoint getPoint() const { return point; }
 
   /// Print the contents of the analysis state.
   virtual void print(raw_ostream &os) const = 0;
 
+  /// Add a dependency to this analysis state on a program point and an
+  /// analysis. If this state is updated, the analysis will be invoked on the
+  /// given program point again (in onUpdate()).
+  void addDependency(ProgramPoint dependent, DataFlowAnalysis *analysis);
+
 protected:
   /// This function is called by the solver when the analysis state is updated
-  /// to optionally enqueue more work items. For example, if a state tracks
-  /// dependents through the IR (e.g. use-def chains), this function can be
-  /// implemented to push those dependents on the worklist.
-  virtual void onUpdate(DataFlowSolver *solver) const {}
+  /// to enqueue more work items. For example, if a state tracks dependents
+  /// through the IR (e.g. use-def chains), this function can be implemented to
+  /// push those dependents on the worklist.
+  virtual void onUpdate(DataFlowSolver *solver) const {
+    for (const DataFlowSolver::WorkItem &item : dependents)
+      solver->enqueue(item);
+  }
+
+  /// The program point to which the state belongs.
+  ProgramPoint point;
+
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+  /// When compiling with debugging, keep a name for the analysis state.
+  StringRef debugName;
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
+private:
   /// The dependency relations originating from this analysis state. An entry
   /// `state -> (analysis, point)` is created when `analysis` queries `state`
   /// when updating `point`.
   ///
   /// When this state is updated, all dependent child analysis invocations are
   /// pushed to the back of the queue. Use a `SetVector` to keep the analysis
   /// deterministic.
   ///
   /// Store the dependents on the analysis state for efficiency.
   SetVector<DataFlowSolver::WorkItem> dependents;
 
-  /// The program point to which the state belongs.
-  ProgramPoint point;
-
-#if LLVM_ENABLE_ABI_BREAKING_CHECKS
-  /// When compiling with debugging, keep a name for the analysis state.
-  StringRef debugName;
-#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
-
   /// Allow the framework to access the dependents.
   friend class DataFlowSolver;
 };
 
 //===----------------------------------------------------------------------===//
 // DataFlowAnalysis
 //===----------------------------------------------------------------------===//
 

mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp

 //===- DeadCodeAnalysis.cpp - Dead code 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/DeadCodeAnalysis.h"
 #include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
+#include "mlir/Analysis/DataFlowFramework.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::dataflow;
 
 //===----------------------------------------------------------------------===//
 // Executable
 //===----------------------------------------------------------------------===//
 
 ChangeResult Executable::setToLive() {
   if (live)
     return ChangeResult::NoChange;
   live = true;
   return ChangeResult::Change;
 }
 
 void Executable::print(raw_ostream &os) const {
   os << (live ? "live" : "dead");
 }
 
 void Executable::onUpdate(DataFlowSolver *solver) const {
+  AnalysisState::onUpdate(solver);
+
   if (auto *block = llvm::dyn_cast_if_present<Block *>(point)) {
     // Re-invoke the analyses on the block itself.
     for (DataFlowAnalysis *analysis : subscribers)
       solver->enqueue({block, analysis});
     // Re-invoke the analyses on all operations in the block.
     for (DataFlowAnalysis *analysis : subscribers)
       for (Operation &op : *block)
         solver->enqueue({&op, analysis});

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/Analysis/DataFlowFramework.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 
 using namespace mlir;
 using namespace mlir::dataflow;
 
 //===----------------------------------------------------------------------===//
 // AbstractSparseLattice
 //===----------------------------------------------------------------------===//
 
 void AbstractSparseLattice::onUpdate(DataFlowSolver *solver) const {
+  AnalysisState::onUpdate(solver);
+
   // Push all users of the value to the queue.
   for (Operation *user : point.get<Value>().getUsers())
     for (DataFlowAnalysis *analysis : useDefSubscribers)
       solver->enqueue({user, analysis});
 }
 
 //===----------------------------------------------------------------------===//
 // AbstractSparseDataFlowAnalysis

mlir/lib/Analysis/DataFlowFramework.cpp

 GenericProgramPoint::~GenericProgramPoint() = default;
 
 //===----------------------------------------------------------------------===//
 // AnalysisState
 //===----------------------------------------------------------------------===//
 
 AnalysisState::~AnalysisState() = default;
 
+void AnalysisState::addDependency(ProgramPoint dependent,
+                                  DataFlowAnalysis *analysis) {
+  auto inserted = dependents.insert({dependent, analysis});
+  (void)inserted;
+  DATAFLOW_DEBUG({
+    if (inserted) {
+      llvm::dbgs() << "Creating dependency between " << debugName << " of "
+                   << point << "\nand " << debugName << " on " << dependent
+                   << "\n";
+    }
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // ProgramPoint
 //===----------------------------------------------------------------------===//
 
 void ProgramPoint::print(raw_ostream &os) const {
   if (isNull()) {
     os << "<NULL POINT>";
     return;

 }
 
 void DataFlowSolver::propagateIfChanged(AnalysisState *state,
                                         ChangeResult changed) {
   if (changed == ChangeResult::Change) {
     DATAFLOW_DEBUG(llvm::dbgs() << "Propagating update to " << state->debugName
                                 << " of " << state->point << "\n"
                                 << "Value: " << *state << "\n");
-    for (const WorkItem &item : state->dependents)
-      enqueue(item);
     state->onUpdate(this);
   }
 }
 
-void DataFlowSolver::addDependency(AnalysisState *state,
-                                   DataFlowAnalysis *analysis,
-                                   ProgramPoint point) {
-  auto inserted = state->dependents.insert({point, analysis});
-  (void)inserted;
-  DATAFLOW_DEBUG({
-    if (inserted) {
-      llvm::dbgs() << "Creating dependency between " << state->debugName
-                   << " of " << state->point << "\nand " << analysis->debugName
-                   << " on " << point << "\n";
-    }
-  });
-}
-
 //===----------------------------------------------------------------------===//
 // DataFlowAnalysis
 //===----------------------------------------------------------------------===//
 
 DataFlowAnalysis::~DataFlowAnalysis() = default;
 
 DataFlowAnalysis::DataFlowAnalysis(DataFlowSolver &solver) : solver(solver) {}
 
 void DataFlowAnalysis::addDependency(AnalysisState *state, ProgramPoint point) {
-  solver.addDependency(state, this, point);
+  state->addDependency(point, this);
 }
 
 void DataFlowAnalysis::propagateIfChanged(AnalysisState *state,
                                           ChangeResult changed) {
   solver.propagateIfChanged(state, changed);
 }