[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

63

Can you document this?

108

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

109

nit: Spell out to result

173

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)

200

Please document these methods.

[Hardcode84]

MapVector, comments

[Hardcode84]

mlir/include/mlir/Pass/AnalysisManager.h

63

done. also moved to private for now

108

you are rigth

173

good idea, done

[rriddle]

Great!

mlir/include/mlir/Pass/AnalysisManager.h

64

nit: Comments should have proper punctuation.

71

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

106–107

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.

172

[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

 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/PassInstrumentation.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMap.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:
   /// 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(
     std::declval<const PreservedAnalyses &>()));
 
 /// Implementation of 'isInvalidated' if the analysis provides a definition.
 template <typename AnalysisT>
 std::enable_if_t<llvm::is_detected<has_is_invalidated, AnalysisT>::value, bool>
 isInvalidated(AnalysisT &analysis, const PreservedAnalyses &pa) {
   return analysis.isInvalidated(pa);
 }

 };
 
 /// 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);
   }
 
   /// 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>>;
 
   /// 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;
public:
   /// Invalidate any cached analyses based upon the given set of preserved
   /// analyses.
   void invalidate(const PreservedAnalyses &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);
     }
   }
 
 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.try_emplace(id, constructAnalysis<AnalysisT>(am, op));
+      assert(wasInserted);
 
       if (pi)
         pi->runAfterAnalysis(getAnalysisName<AnalysisT>(), id, ir);
     }
     return static_cast<AnalysisModel<AnalysisT> &>(*it->second).analysis;
   }
 
+  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);
+  }
+
+  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 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