Add support for preserving specific analyses in the analysis manager. Passes can now preserve specific analyses via 'markAnalysesPreserved'.

Example:

markAnalysesPreserved<DominanceInfo>();
markAnalysesPreserved<DominanceInfo, PostDominanceInfo>();

PiperOrigin-RevId: 237081454

mlir/include/mlir/Pass/AnalysisManager.h

@@ -44,9 +44,33 @@ public:
   /// Mark all analyses as preserved.
   void preserveAll() { preservedIDs.insert(&allAnalysesID); }
 
-  /// Returns if all analyses were marked preserved.
+  /// Returns true if all analyses were marked preserved.
   bool isAll() const { return preservedIDs.count(&allAnalysesID); }
 
+  /// Returns true if no analyses were marked preserved.
+  bool isNone() const { return preservedIDs.empty(); }
+
+  /// Preserve the given analyses.
+  template <typename AnalysisT> void preserve() {
+    preserve(AnalysisID::getID<AnalysisT>());
+  }
+  template <typename AnalysisT, typename AnalysisT2, typename... OtherAnalysesT>
+  void preserve() {
+    preserve<AnalysisT>();
+    preserve<AnalysisT2, OtherAnalysesT...>();
+  }
+  void preserve(const AnalysisID *id) { preservedIDs.insert(id); }
+
+  /// Returns 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(AnalysisID::getID<AnalysisT>());
+  }
+  bool isPreserved(const AnalysisID *id) const {
+    return preservedIDs.count(id);
+  }
+
 private:
   /// An identifier used to represent all potential analyses.
   constexpr static AnalysisID allAnalysesID = {};
@@ -112,11 +136,12 @@ public:
   /// Invalidate any cached analyses based upon the given set of preserved
   /// analyses.
   void invalidate(const detail::PreservedAnalyses &pa) {
-    // If all analyses were preserved, then there is nothing to do here.
-    if (pa.isAll())
-      return;
-    // TODO: Fine grain invalidation of analyses.
-    clear();
+    // Remove any analyses not marked as preserved.
+    for (auto it = results.begin(), e = results.end(); it != e;) {
+      auto curIt = it++;
+      if (!pa.isPreserved(curIt->first))
+        results.erase(curIt);
+    }
   }
 
 private:
@@ -154,7 +179,12 @@ public:
   }
 
   /// Invalidate any non preserved analyses,
-  void invalidate(const detail::PreservedAnalyses &pa) { impl->invalidate(pa); }
+  void invalidate(const detail::PreservedAnalyses &pa) {
+    // If all analyses were preserved, then there is nothing to do here.
+    if (pa.isAll())
+      return;
+    impl->invalidate(pa);
+  }
 
   /// Clear any held analyses.
   void clear() { impl->clear(); }
@@ -189,13 +219,23 @@ public:
     return slice(function).getResult<AnalysisT>();
   }
 
+  /// Query for a cached analysis of a function, or return null.
+  template <typename AnalysisT>
+  llvm::Optional<std::reference_wrapper<AnalysisT>>
+  getCachedFunctionResult(Function *function) const {
+    auto it = functionAnalyses.find(function);
+    if (it == functionAnalyses.end())
+      return llvm::None;
+    return it->second.getCachedResult<AnalysisT>();
+  }
+
   /// Query for the analysis of a module. The analysis is computed if it does
   /// not exist.
   template <typename AnalysisT> AnalysisT &getResult() {
     return moduleAnalyses.getResult<AnalysisT>();
   }
 
-  /// Query for a cached analysis for the module, or return nullptr.
+  /// Query for a cached analysis for the module, or return null.
   template <typename AnalysisT>
   llvm::Optional<std::reference_wrapper<AnalysisT>> getCachedResult() const {
     return moduleAnalyses.getCachedResult<AnalysisT>();

mlir/include/mlir/Pass/Pass.h

@@ -207,6 +207,14 @@ protected:
   void markAllAnalysesPreserved() {
     this->getPassState().preservedAnalyses.preserveAll();
   }
+
+  /// Mark the provided analyses as preserved.
+  template <typename... AnalysesT> void markAnalysesPreserved() {
+    this->getPassState().preservedAnalyses.template preserve<AnalysesT...>();
+  }
+  void markAnalysesPreserved(const AnalysisID *id) {
+    this->getPassState().preservedAnalyses.preserve(id);
+  }
 };
 } // end namespace detail
 
@@ -242,6 +250,14 @@ struct ModulePass : public detail::PassModel<Module, T, ModulePassBase> {
   AnalysisT &getFunctionAnalysisResult(Function *f) {
     return this->getAnalysisManager().template getFunctionResult<AnalysisT>(f);
   }
+
+  /// Returns an existing analysis result for a child function if it exists.
+  template <typename AnalysisT>
+  llvm::Optional<std::reference_wrapper<AnalysisT>>
+  getCachedFunctionAnalysisResult(Function *f) {
+    return this->getAnalysisManager()
+        .template getCachedFunctionResult<AnalysisT>(f);
+  }
 };
 } // end namespace mlir
 

mlir/lib/Pass/Pass.cpp

@@ -309,10 +309,21 @@ FunctionAnalysisManager ModuleAnalysisManager::slice(Function *function) {
 
 /// Invalidate any non preserved analyses.
 void ModuleAnalysisManager::invalidate(const detail::PreservedAnalyses &pa) {
+  // If all analyses were preserved, then there is nothing to do here.
   if (pa.isAll())
     return;
 
-  // TODO: Fine grain invalidation of analyses.
-  moduleAnalyses.clear();
-  functionAnalyses.clear();
+  // Invalidate the module analyses directly.
+  moduleAnalyses.invalidate(pa);
+
+  // If no analyses were preserved, then just simply clear out the function
+  // analysis results.
+  if (pa.isNone()) {
+    functionAnalyses.clear();
+    return;
+  }
+
+  // Otherwise, invalidate each function analyses.
+  for (auto &analysisPair : functionAnalyses)
+    analysisPair.second.invalidate(pa);
 }

mlir/lib/Transforms/CSE.cpp

@@ -124,6 +124,12 @@ private:
 
 /// Attempt to eliminate a redundant operation.
 bool CSE::simplifyOperation(Instruction *op) {
+  // Don't simplify operations with nested blocks. We don't currently model
+  // equality comparisons correctly among other things. It is also unclear
+  // whether we would want to CSE such operations.
+  if (op->getNumBlockLists() != 0)
+    return false;
+
   // TODO(riverriddle) We currently only eliminate non side-effecting
   // operations.
   if (!op->hasNoSideEffect())
@@ -230,6 +236,10 @@ void CSE::runOnFunction() {
   for (auto *op : opsToErase)
     op->erase();
   opsToErase.clear();
+
+  // We currently don't remove region operations, so mark dominance as
+  // preserved.
+  markAnalysesPreserved<DominanceInfo, PostDominanceInfo>();
 }
 
 FunctionPassBase *mlir::createCSEPass() { return new CSE(); }

mlir/unittests/Pass/AnalysisManagerTest.cpp

@@ -0,0 +1,112 @@
+//===- AnalysisManagerTest.cpp - AnalysisManager unit tests ---------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+#include "mlir/Pass/AnalysisManager.h"
+#include "mlir/IR/Builders.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::detail;
+
+namespace {
+/// Minimal class definitions for two analyses.
+struct MyAnalysis {
+  MyAnalysis(Function *) {}
+  MyAnalysis(Module *) {}
+};
+struct OtherAnalysis {
+  OtherAnalysis(Function *) {}
+  OtherAnalysis(Module *) {}
+};
+
+TEST(AnalysisManagerTest, FineGrainModuleAnalysisPreservation) {
+  MLIRContext context;
+
+  // Test fine grain invalidation of the module analysis manager.
+  std::unique_ptr<Module> module(new Module(&context));
+  ModuleAnalysisManager mam(&*module);
+
+  // Query two different analyses, but only preserve one before invalidating.
+  mam.getResult<MyAnalysis>();
+  mam.getResult<OtherAnalysis>();
+
+  detail::PreservedAnalyses pa;
+  pa.preserve<MyAnalysis>();
+  mam.invalidate(pa);
+
+  // Check that only MyAnalysis is preserved.
+  EXPECT_TRUE(mam.getCachedResult<MyAnalysis>().hasValue());
+  EXPECT_FALSE(mam.getCachedResult<OtherAnalysis>().hasValue());
+}
+
+TEST(AnalysisManagerTest, FineGrainFunctionAnalysisPreservation) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  // Create a function and a module.
+  std::unique_ptr<Module> module(new Module(&context));
+  Function *func1 =
+      new Function(builder.getUnknownLoc(), "foo",
+                   builder.getFunctionType(llvm::None, llvm::None));
+  module->getFunctions().push_back(func1);
+
+  // Test fine grain invalidation of the function analysis manager.
+  ModuleAnalysisManager mam(&*module);
+  FunctionAnalysisManager fam = mam.slice(func1);
+
+  // Query two different analyses, but only preserve one before invalidating.
+  fam.getResult<MyAnalysis>();
+  fam.getResult<OtherAnalysis>();
+
+  detail::PreservedAnalyses pa;
+  pa.preserve<MyAnalysis>();
+  fam.invalidate(pa);
+
+  // Check that only MyAnalysis is preserved.
+  EXPECT_TRUE(fam.getCachedResult<MyAnalysis>().hasValue());
+  EXPECT_FALSE(fam.getCachedResult<OtherAnalysis>().hasValue());
+}
+
+TEST(AnalysisManagerTest, FineGrainChildFunctionAnalysisPreservation) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  // Create a function and a module.
+  std::unique_ptr<Module> module(new Module(&context));
+  Function *func1 =
+      new Function(builder.getUnknownLoc(), "foo",
+                   builder.getFunctionType(llvm::None, llvm::None));
+  module->getFunctions().push_back(func1);
+
+  // Test fine grain invalidation of a function analysis from within a module
+  // analysis manager.
+  ModuleAnalysisManager mam(&*module);
+
+  // Query two different analyses, but only preserve one before invalidating.
+  mam.getFunctionResult<MyAnalysis>(func1);
+  mam.getFunctionResult<OtherAnalysis>(func1);
+
+  detail::PreservedAnalyses pa;
+  pa.preserve<MyAnalysis>();
+  mam.invalidate(pa);
+
+  // Check that only MyAnalysis is preserved.
+  EXPECT_TRUE(mam.getCachedFunctionResult<MyAnalysis>(func1).hasValue());
+  EXPECT_FALSE(mam.getCachedFunctionResult<OtherAnalysis>(func1).hasValue());
+}
+
+} // end namespace