[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?

118–124

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

119

nit: Spell out to result

186

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)

216

Please document these methods.

[Hardcode84]

MapVector, comments

[Hardcode84]

mlir/include/mlir/Pass/AnalysisManager.h

64

done. also moved to private for now

118–124

you are rigth

186

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.

117–118

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.

185–188

[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

160–166

175–178

209

Add comments here for the interesting bits.

920–925

[Hardcode84]

update docs

mlir/docs/PassManagement.md

 
 ## Analysis Management
 
 An important concept, along with transformation passes, are analyses. These are
 conceptually similar to transformation passes, except that they compute
 information on a specific operation without modifying it. In MLIR, analyses are
 not passes but free-standing classes that are computed lazily on-demand and
 cached to avoid unnecessary recomputation. An analysis in MLIR must adhere to
 the following:
 
-*   Provide a valid constructor taking an `Operation*`.
+*   Provide a valid constructor taking either an `Operation*` or `Operation*`
+    and `AnalysisManager &`.
+    *   The provided `AnalysisManager &` should be used to query any necessary
+        analysis dependencies.
 *   Must not modify the given operation.
 
 An analysis may provide additional hooks to control various behavior:
 
 *   `bool isInvalidated(const AnalysisManager::PreservedAnalyses &)`
 
 Given a preserved analysis set, the analysis 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.
+invalidated) based upon other properties such as analyses sets. If the analysis
+uses any other analysis as a dependency, it must also check if the dependency
+was invalidated.
 
 ### Querying Analyses
 
 The base `OperationPass` class provides utilities for querying and preserving
 analyses for the current operation being processed.
 
 *   OperationPass automatically provides the following utilities for querying
     analyses:
     *   `getAnalysis<>`

 
 ```c++
 /// An interesting analysis.
 struct MyOperationAnalysis {
   // Compute this analysis with the provided operation.
   MyOperationAnalysis(Operation *op);
 };
 
+struct MyOperationAnalysisWithDependency {
+  MyOperationAnalysisWithDependency(Operation *op, AnalysisManager &am) {
+    // Request other analysis as dependency
+    MyOperationAnalysis &otherAnalysis = am.getAnalysis<MyOperationAnalysis>();
+    ...
+  }
+
+  bool isInvalidated(const AnalysisManager::PreservedAnalyses &pa) {
+    // Check if analysis or its dependency were invalidated
+    return !pa.isPreserved<MyOperationAnalysisWithDependency>() ||
+           !pa.isPreserved<MyOperationAnalysis>();
+  }
+};
+
 void MyOperationPass::runOnOperation() {
   // Query MyOperationAnalysis for the current operation.
   MyOperationAnalysis &myAnalysis = getAnalysis<MyOperationAnalysis>();
 
   // Query a cached instance of MyOperationAnalysis for the current operation.
   // It will not be computed if it doesn't exist.
   auto optionalAnalysis = getCachedAnalysis<MyOperationAnalysis>();
   if (optionalAnalysis)
*   `runBeforePass`
     *   This callback is run just before a pass is executed.
 *   `runAfterPass`
     *   This callback is run right after a pass has been successfully executed.
         If this hook is executed, `runAfterPassFailed` will *not* be.
 *   `runAfterPassFailed`
     *   This callback is run right after a pass execution fails. If this hook is
         executed, `runAfterPass` will *not* be.
 *   `runBeforeAnalysis`
     *   This callback is run just before an analysis is computed.
+    *   If the analysis requested another analysis as a dependency, the
+        `runBeforeAnalysis`/`runAfterAnalysis` pair for the dependency can be
+        called from inside of the current `runBeforeAnalysis`/`runAfterAnalysis`
+        pair.
 *   `runAfterAnalysis`
     *   This callback is run right after an analysis is computed.
 
 PassInstrumentation instances may be registered directly with a
 [PassManager](#pass-manager) instance via the `addInstrumentation` method.
 Instrumentations added to the PassManager are run in a stack like fashion, i.e.
 the last instrumentation to execute a `runBefore*` hook will be the first to
 execute the respective `runAfter*` hook. The hooks of a `PassInstrumentation`
 class are guaranteed to be executed in a thread safe fashion, so additional

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>());
+  }
+
+  /// AnalysisModel need access to unpreserve().
+  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(

 /// The abstract polymorphic base class representing an analysis.
 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;
+  /// properties such as analyses sets. Invalidated analyses must also be
+  /// removed from pa.
+  virtual bool invalidate(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);
+  /// A hook used to query analyses for invalidation. Removes invalidated
+  /// analyses from pa.
+  bool invalidate(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 are 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->invalidate(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