[mlir] Pass AnalysisManager as optional parameter to analysis ctor, so it can request any other analysis as dependency

[Hardcode84]

fix getAnalysis

[Hardcode84]

Templates fix

[rriddle]

You'll also need to make sure that when an analysis is invalidated that all of the things that depend on it check to see if they also need to be invalidated. I would take a look at the Analysis Manager in LLVM as an inspiration for inter-analysis dependencies.

[rriddle]

Thanks for taking a look at this, it has been pushed off for a while.

[Hardcode84]

In D100274#2684577, @rriddle wrote:

You'll also need to make sure that when an analysis is invalidated that all of the things that depend on it check to see if they also need to be invalidated. I would take a look at the Analysis Manager in LLVM as an inspiration for inter-analysis dependencies.

This should be already covered by isInvalidated (see AnalysisWithDependency::isInvalidated in tests). It is probably a good idea to have a template helper which can generate isInvalidated automatically, but its better make as separate patch.

[rriddle]

In D100274#2685158, @Hardcode84 wrote:

In D100274#2684577, @rriddle wrote:

You'll also need to make sure that when an analysis is invalidated that all of the things that depend on it check to see if they also need to be invalidated. I would take a look at the Analysis Manager in LLVM as an inspiration for inter-analysis dependencies.

This should be already covered by isInvalidated (see AnalysisWithDependency::isInvalidated in tests). It is probably a good idea to have a template helper which can generate isInvalidated automatically, but its better make as separate patch.

That only handles very simple cases, not the more complex. For example, say you have an analysis A that depends on an analysis B, which depends on analysis C. A -> B -> C. What happens if a pass marks analysis B as preserved but not C? It seems that B would be invalidated properly (because C is not preserved), but A would still see that B has been preserved so it wouldn't invalidate itself.

[Hardcode84]

Properly handle analysis dependency chains

[rriddle]

Added more comments, looking good.

mlir/include/mlir/Pass/AnalysisManager.h

64

Can you document this?

117–121

The method is already marked final, why add an override?

118

nit: Spell out to result

183

If we changed analyses to use a MapVector, can we rely on dependencies being ordered first to avoid the need for this? (MapVector guarantees deterministic iteration based on insertion time)

213

Please document these methods.

[Hardcode84]

MapVector, comments

[Hardcode84]

mlir/include/mlir/Pass/AnalysisManager.h

64

done. also moved to private for now

117–121

you are rigth

183

good idea, done

[rriddle]

Great!

mlir/include/mlir/Pass/AnalysisManager.h

65

nit: Comments should have proper punctuation.

72

nit: Add a comment on why this is a friend.

116

Can you add it that we also unpreserve any analyses from pa that were invalidated? Maybe we should also rename this method to invalidate, given that it mutates as well.

182

[Hardcode84]

https://mlir.llvm.org/docs/PassManagement/#analysis-management
Docs probably also need to be updated and I will need some help with that. Should it go to separate review?

[rriddle]

In D100274#2695307, @Hardcode84 wrote:

https://mlir.llvm.org/docs/PassManagement/#analysis-management
Docs probably also need to be updated and I will need some help with that. Should it go to separate review?

Let's roll that into this one. (Will mark as changes requested for now)

Can you also update the description?

[Hardcode84]

review comments, docs

[Hardcode84]

fix code in docs

[rriddle]

LG thanks.

mlir/docs/PassManagement.md

159–165 ↗

(On Diff #338207)

174–176 ↗

(On Diff #338207)

208 ↗

(On Diff #338207)

Add comments here for the interesting bits.

917–920 ↗

(On Diff #338207)

[Hardcode84]

update docs

mlir/include/mlir/Pass/AnalysisManager.h

 //===- AnalysisManager.h - Analysis Management Infrastructure ---*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_PASS_ANALYSISMANAGER_H
 #define MLIR_PASS_ANALYSISMANAGER_H
 
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/PassInstrumentation.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/TypeName.h"
 
 namespace mlir {
+class AnalysisManager;
+
 //===----------------------------------------------------------------------===//
 // Analysis Preservation and Concept Modeling
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// A utility class to represent the analyses that are known to be preserved.
 class PreservedAnalyses {
   /// A type used to represent all potential analyses.
public:
   /// Returns true if the given analysis has been marked as preserved. Note that
   /// this simply checks for the presence of a given analysis ID and should not
   /// be used as a general preservation checker.
   template <typename AnalysisT> bool isPreserved() const {
     return isPreserved(TypeID::get<AnalysisT>());
   }
   bool isPreserved(TypeID id) const { return preservedIDs.count(id); }
 
 private:
+  /// Remove the analysis from preserved set
+  template <typename AnalysisT>
+  void unpreserve() {
+    preservedIDs.erase(TypeID::get<AnalysisT>());
+  }
+
+  template <typename>
+  friend struct AnalysisModel;
+
   /// The set of analyses that are known to be preserved.
   SmallPtrSet<TypeID, 2> preservedIDs;
 };
 
 namespace analysis_impl {
 /// Trait to check if T provides a static 'isInvalidated' method.
 template <typename T, typename... Args>
 using has_is_invalidated = decltype(std::declval<T &>().isInvalidated(

 struct AnalysisConcept {
   virtual ~AnalysisConcept() = default;
 
   /// A hook used to query analyses for invalidation. Given a preserved analysis
   /// set, returns true if it should truly be invalidated. This allows for more
   /// fine-tuned invalidation in cases where an analysis wasn't explicitly
   /// marked preserved, but may be preserved(or invalidated) based upon other
   /// properties such as analyses sets.
-  virtual bool isInvalidated(const PreservedAnalyses &pa) = 0;
+  virtual bool isInvalidated(PreservedAnalyses &pa) = 0;
 };
 
 /// A derived analysis model used to hold a specific analysis object.
 template <typename AnalysisT> struct AnalysisModel : public AnalysisConcept {
   template <typename... Args>
   explicit AnalysisModel(Args &&...args)
       : analysis(std::forward<Args>(args)...) {}
 
   /// A hook used to query analyses for invalidation.
-  bool isInvalidated(const PreservedAnalyses &pa) final {
-    return analysis_impl::isInvalidated(analysis, pa);
+  bool isInvalidated(PreservedAnalyses &pa) final {
+    bool result = analysis_impl::isInvalidated(analysis, pa);
+    if (result)
+      pa.unpreserve<AnalysisT>();
+    return result;
   }
 
   /// The actual analysis object.
   AnalysisT analysis;
 };
 
 /// This class represents a cache of analyses for a single operation. All
 /// computation, caching, and invalidation of analyses takes place here.
 class AnalysisMap {
   /// A mapping between an analysis id and an existing analysis instance.
-  using ConceptMap = DenseMap<TypeID, std::unique_ptr<AnalysisConcept>>;
+  using ConceptMap = llvm::MapVector<TypeID, std::unique_ptr<AnalysisConcept>>;
 
   /// Utility to return the name of the given analysis class.
   template <typename AnalysisT> static StringRef getAnalysisName() {
     StringRef name = llvm::getTypeName<AnalysisT>();
     if (!name.consume_front("mlir::"))
       name.consume_front("(anonymous namespace)::");
     return name;
   }
 
 public:
   explicit AnalysisMap(Operation *ir) : ir(ir) {}
 
   /// Get an analysis for the current IR unit, computing it if necessary.
   template <typename AnalysisT>
-  AnalysisT &getAnalysis(PassInstrumentor *pi) {
-    return getAnalysisImpl<AnalysisT, Operation *>(pi, ir);
+  AnalysisT &getAnalysis(PassInstrumentor *pi, AnalysisManager &am) {
+    return getAnalysisImpl<AnalysisT, Operation *>(pi, ir, am);
   }
 
   /// Get an analysis for the current IR unit assuming it's of specific derived
   /// operation type.
   template <typename AnalysisT, typename OpT>
-  typename std::enable_if<std::is_constructible<AnalysisT, OpT>::value,
-                          AnalysisT &>::type
-  getAnalysis(PassInstrumentor *pi) {
-    return getAnalysisImpl<AnalysisT, OpT>(pi, cast<OpT>(ir));
+  std::enable_if_t<
+      std::is_constructible<AnalysisT, OpT>::value ||
+          std::is_constructible<AnalysisT, OpT, AnalysisManager &>::value,
+      AnalysisT &>
+  getAnalysis(PassInstrumentor *pi, AnalysisManager &am) {
+    return getAnalysisImpl<AnalysisT, OpT>(pi, cast<OpT>(ir), am);
   }
 
   /// Get a cached analysis instance if one exists, otherwise return null.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>> getCachedAnalysis() const {
     auto res = analyses.find(TypeID::get<AnalysisT>());
     if (res == analyses.end())
       return llvm::None;
     return {static_cast<AnalysisModel<AnalysisT> &>(*res->second).analysis};
   }
 
   /// Returns the operation that this analysis map represents.
   Operation *getOperation() const { return ir; }
 
   /// Clear any held analyses.
   void clear() { analyses.clear(); }
 
   /// Invalidate any cached analyses based upon the given set of preserved
   /// analyses.
   void invalidate(const PreservedAnalyses &pa) {
+    PreservedAnalyses paCopy(pa);
     // Remove any analyses that were invalidated.
-    for (auto it = analyses.begin(), e = analyses.end(); it != e;) {
-      auto curIt = it++;
-      if (curIt->second->isInvalidated(pa))
-        analyses.erase(curIt);
-    }
+    // As we using MapVector, order of insertion is preserved and
+    // dependencies always go before users, so we need only one iteration
+    analyses.remove_if(
+        [&](auto &val) { return val.second->isInvalidated(paCopy); });
   }
 
 private:
   template <typename AnalysisT, typename OpT>
-  AnalysisT &getAnalysisImpl(PassInstrumentor *pi, OpT op) {
+  AnalysisT &getAnalysisImpl(PassInstrumentor *pi, OpT op,
+                             AnalysisManager &am) {
     TypeID id = TypeID::get<AnalysisT>();
 
-    typename ConceptMap::iterator it;
-    bool wasInserted;
-    std::tie(it, wasInserted) = analyses.try_emplace(id);
-
+    auto it = analyses.find(id);
     // If we don't have a cached analysis for this operation, compute it
     // directly and add it to the cache.
-    if (wasInserted) {
+    if (analyses.end() == it) {
       if (pi)
         pi->runBeforeAnalysis(getAnalysisName<AnalysisT>(), id, ir);
 
-      it->second = std::make_unique<AnalysisModel<AnalysisT>>(op);
+      bool wasInserted;
+      std::tie(it, wasInserted) =
+          analyses.insert({id, constructAnalysis<AnalysisT>(am, op)});
+      assert(wasInserted);
 
       if (pi)
         pi->runAfterAnalysis(getAnalysisName<AnalysisT>(), id, ir);
     }
     return static_cast<AnalysisModel<AnalysisT> &>(*it->second).analysis;
   }
 
+  /// Construct analysis using two arguments contructor (OpT, AnalysisManager)
+  template <typename AnalysisT, typename OpT,
+            std::enable_if_t<std::is_constructible<
+                AnalysisT, OpT, AnalysisManager &>::value> * = nullptr>
+  static auto constructAnalysis(AnalysisManager &am, OpT op) {
+    return std::make_unique<AnalysisModel<AnalysisT>>(op, am);
+  }
+
+  /// Construct analysis using single argument contructor (OpT)
+  template <typename AnalysisT, typename OpT,
+            std::enable_if_t<!std::is_constructible<
+                AnalysisT, OpT, AnalysisManager &>::value> * = nullptr>
+  static auto constructAnalysis(AnalysisManager &, OpT op) {
+    return std::make_unique<AnalysisModel<AnalysisT>>(op);
+  }
+
   Operation *ir;
   ConceptMap analyses;
 };
 
 /// An analysis map that contains a map for the current operation, and a set of
 /// maps for any child operations.
 struct NestedAnalysisMap {
   NestedAnalysisMap(Operation *op, PassInstrumentor *instrumentor)
while (auto *parentAM = curParent->getParent()) {
         return parentAM->analyses.getCachedAnalysis<AnalysisT>();
       curParent = parentAM;
     }
     return None;
   }
 
   /// Query for the given analysis for the current operation.
   template <typename AnalysisT> AnalysisT &getAnalysis() {
-    return impl->analyses.getAnalysis<AnalysisT>(getPassInstrumentor());
+    return impl->analyses.getAnalysis<AnalysisT>(getPassInstrumentor(), *this);
   }
 
   /// Query for the given analysis for the current operation of a specific
   /// derived operation type.
   template <typename AnalysisT, typename OpT>
   AnalysisT &getAnalysis() {
-    return impl->analyses.getAnalysis<AnalysisT, OpT>(getPassInstrumentor());
+    return impl->analyses.getAnalysis<AnalysisT, OpT>(getPassInstrumentor(),
+                                                      *this);
   }
 
   /// Query for a cached entry of the given analysis on the current operation.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>> getCachedAnalysis() const {
     return impl->analyses.getCachedAnalysis<AnalysisT>();
   }
 

mlir/unittests/Pass/AnalysisManagerTest.cpp

TEST(AnalysisManagerTest, OpSpecificAnalysis) {
   ModuleAnalysisManager mam(*module, /*passInstrumentor=*/nullptr);
   AnalysisManager am = mam;
 
   // Query the op specific analysis for the module and verify that its cached.
   am.getAnalysis<OpSpecificAnalysis, ModuleOp>();
   EXPECT_TRUE(am.getCachedAnalysis<OpSpecificAnalysis>().hasValue());
 }
 
+struct AnalysisWithDependency {
+  AnalysisWithDependency(Operation *, AnalysisManager &am) {
+    am.getAnalysis<MyAnalysis>();
+  }
+
+  bool isInvalidated(const AnalysisManager::PreservedAnalyses &pa) {
+    return !pa.isPreserved<AnalysisWithDependency>() ||
+           !pa.isPreserved<MyAnalysis>();
+  }
+};
+
+TEST(AnalysisManagerTest, DependentAnalysis) {
+  MLIRContext context;
+
+  // Create a module.
+  OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
+  ModuleAnalysisManager mam(*module, /*passInstrumentor=*/nullptr);
+  AnalysisManager am = mam;
+
+  am.getAnalysis<AnalysisWithDependency>();
+  EXPECT_TRUE(am.getCachedAnalysis<AnalysisWithDependency>().hasValue());
+  EXPECT_TRUE(am.getCachedAnalysis<MyAnalysis>().hasValue());
+
+  detail::PreservedAnalyses pa;
+  pa.preserve<AnalysisWithDependency>();
+  am.invalidate(pa);
+
+  EXPECT_FALSE(am.getCachedAnalysis<AnalysisWithDependency>().hasValue());
+  EXPECT_FALSE(am.getCachedAnalysis<MyAnalysis>().hasValue());
+}
+
+struct AnalysisWithNestedDependency {
+  AnalysisWithNestedDependency(Operation *, AnalysisManager &am) {
+    am.getAnalysis<AnalysisWithDependency>();
+  }
+
+  bool isInvalidated(const AnalysisManager::PreservedAnalyses &pa) {
+    return !pa.isPreserved<AnalysisWithNestedDependency>() ||
+           !pa.isPreserved<AnalysisWithDependency>();
+  }
+};
+
+TEST(AnalysisManagerTest, NestedDependentAnalysis) {
+  MLIRContext context;
+
+  // Create a module.
+  OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
+  ModuleAnalysisManager mam(*module, /*passInstrumentor=*/nullptr);
+  AnalysisManager am = mam;
+
+  am.getAnalysis<AnalysisWithNestedDependency>();
+  EXPECT_TRUE(am.getCachedAnalysis<AnalysisWithNestedDependency>().hasValue());
+  EXPECT_TRUE(am.getCachedAnalysis<AnalysisWithDependency>().hasValue());
+  EXPECT_TRUE(am.getCachedAnalysis<MyAnalysis>().hasValue());
+
+  detail::PreservedAnalyses pa;
+  pa.preserve<AnalysisWithDependency>();
+  pa.preserve<AnalysisWithNestedDependency>();
+  am.invalidate(pa);
+
+  EXPECT_FALSE(am.getCachedAnalysis<AnalysisWithNestedDependency>().hasValue());
+  EXPECT_FALSE(am.getCachedAnalysis<AnalysisWithDependency>().hasValue());
+  EXPECT_FALSE(am.getCachedAnalysis<MyAnalysis>().hasValue());
+}
+
+struct AnalysisWith2Ctors {
+  AnalysisWith2Ctors(Operation *) { ctor1called = true; }
+
+  AnalysisWith2Ctors(Operation *, AnalysisManager &) { ctor2called = true; }
+
+  bool ctor1called = false;
+  bool ctor2called = false;
+};
+
+TEST(AnalysisManagerTest, DependentAnalysis2Ctors) {
+  MLIRContext context;
+
+  // Create a module.
+  OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
+  ModuleAnalysisManager mam(*module, /*passInstrumentor=*/nullptr);
+  AnalysisManager am = mam;
+
+  auto &an = am.getAnalysis<AnalysisWith2Ctors>();
+  EXPECT_FALSE(an.ctor1called);
+  EXPECT_TRUE(an.ctor2called);
+}
+
 } // end namespace