[mlir][NFC] Refactor eraseState to take constant time

[IanWood1]

Refactors analysisStates to use two nested maps . This prevents eraseState from having to scan through every analysis state which can be costly when there are many analysis states and/or eraseState is called frequently.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ian Wood (IanWood1)

Changes

Refactors analysisStates to use two nested maps . This prevents eraseState from having to scan through every analysis state which can be costly when there are many analysis states and/or eraseState is called frequently.

---

Full diff: https://github.com/llvm/llvm-project/pull/121670.diff

1 Files Affected:

• (modified) mlir/include/mlir/Analysis/DataFlowFramework.h (+9-14)

diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index dfd358e7017a4e..b6d10ba0bea2d8 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -332,9 +332,11 @@ class DataFlowSolver {
   /// does not exist.
   template <typename StateT, typename AnchorT>
   const StateT *lookupState(AnchorT anchor) const {
-    auto it =
-        analysisStates.find({LatticeAnchor(anchor), TypeID::get<StateT>()});
-    if (it == analysisStates.end())
+    const auto &mapIt = analysisStates.find(LatticeAnchor(anchor));
+    if (mapIt == analysisStates.end())
+      return nullptr;
+    auto it = mapIt->second.find(TypeID::get<StateT>());
+    if (it == mapIt->second.end())
       return nullptr;
     return static_cast<const StateT *>(it->second.get());
   }
@@ -343,11 +345,7 @@ class DataFlowSolver {
   template <typename AnchorT>
   void eraseState(AnchorT anchor) {
     LatticeAnchor la(anchor);
-
-    for (auto it = analysisStates.begin(); it != analysisStates.end(); ++it) {
-      if (it->first.first == la)
-        analysisStates.erase(it);
-    }
+    analysisStates.erase(LatticeAnchor(anchor));
   }
 
   // Erase all analysis states
@@ -426,7 +424,8 @@ class DataFlowSolver {
 
   /// A type-erased map of lattice anchors to associated analysis states for
   /// first-class lattice anchors.
-  DenseMap<std::pair<LatticeAnchor, TypeID>, std::unique_ptr<AnalysisState>>
+  DenseMap<LatticeAnchor, DenseMap<TypeID, std::unique_ptr<AnalysisState>>,
+           DenseMapInfo<LatticeAnchor::ParentTy>>
       analysisStates;
 
   /// Allow the base child analysis class to access the internals of the solver.
@@ -643,7 +642,7 @@ AnalysisT *DataFlowSolver::load(Args &&...args) {
 template <typename StateT, typename AnchorT>
 StateT *DataFlowSolver::getOrCreateState(AnchorT anchor) {
   std::unique_ptr<AnalysisState> &state =
-      analysisStates[{LatticeAnchor(anchor), TypeID::get<StateT>()}];
+      analysisStates[LatticeAnchor(anchor)][TypeID::get<StateT>()];
   if (!state) {
     state = std::unique_ptr<StateT>(new StateT(anchor));
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
@@ -689,10 +688,6 @@ struct DenseMapInfo<mlir::ProgramPoint> {
   }
 };
 
-template <>
-struct DenseMapInfo<mlir::LatticeAnchor>
-    : public DenseMapInfo<mlir::LatticeAnchor::ParentTy> {};
-
 // Allow llvm::cast style functions.
 template <typename To>
 struct CastInfo<To, mlir::LatticeAnchor>

[stellaraccident]

(edit - this comment was written in the wrong tab)

[stellaraccident]

Thanks - this helps some cases I was hitting when trying to use this.