[MLIR] Add support for defining and using Op specific analysis

[jurahul]

Summary

• Add variants of getAnalysis() and friends that operate on a specific derived operation types.
• Add OpPassManager::getAnalysis() to always call the base getAnalysis() with OpT.
• With this, an OperationPass can call getAnalysis<> using an analysis type that is generic (works on Operation *) or specific to the OpT for the pass. Anything else will fail to compile.
• Extend AnalysisManager unit test to test this, and add a new PassManager unit test to test this functionality in the context of an OperationPass.

[jurahul]

Rebase

[jurahul]

Address clang-tidy warnings

[rriddle]

Thanks Rahul.

mlir/include/mlir/Pass/AnalysisManager.h

132

Why is this guarded by a template?

136–142

typo: assuming

its -> it's

137

nit: Use triple tick comments here.

172

Why is the argument list variadic?

254

This method is unnecessary, please remove it. *Op methods already implicitly convert to Operation*, there is nothing gained here by providing the derived operation type.

254

Please fix all of the comments to be triple tick.

276

.getOperation() shouldn't be necessary.

291

This method is unnecessary, please remove it.

mlir/include/mlir/Pass/Pass.h

206

Same here, this method should be removed.

221

Same here, this method should be removed.

[jurahul]

Address review comments

[jurahul]

mlir/include/mlir/Pass/AnalysisManager.h

132

Removed.

[rriddle]

mlir/include/mlir/Pass/AnalysisManager.h

194

nit: Remove the double private:

mlir/unittests/Pass/PassManagerTest.cpp

66

I'm not sure that any of this is necessary to test, if the op specific analysis doesn't provide a proper constructor none of this should compile in the first place.

[jurahul]

mlir/unittests/Pass/PassManagerTest.cpp

66

Beyond just being able to compile, I think it would be good to test basic functionality. For instance, the correct op is passed into the constructor etc.

[jurahul]

Remove double private

mlir/include/mlir/Pass/AnalysisManager.h

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) {
-    TypeID id = TypeID::get<AnalysisT>();
-
-    typename ConceptMap::iterator it;
-    bool wasInserted;
-    std::tie(it, wasInserted) = analyses.try_emplace(id);
-
-    // If we don't have a cached analysis for this function, compute it directly
-    // and add it to the cache.
-    if (wasInserted) {
-      if (pi)
-        pi->runBeforeAnalysis(getAnalysisName<AnalysisT>(), id, ir);
-
-      it->second = std::make_unique<AnalysisModel<AnalysisT>>(ir);
-
-      if (pi)
-        pi->runAfterAnalysis(getAnalysisName<AnalysisT>(), id, ir);
-    }
-    return static_cast<AnalysisModel<AnalysisT> &>(*it->second).analysis;
+  template <typename AnalysisT>
+  AnalysisT &getAnalysis(PassInstrumentor *pi) {
+    return getAnalysisImpl<AnalysisT, Operation *>(pi, ir);
+
+  /// 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));
   }
 
   /// 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;
void invalidate(const PreservedAnalyses &pa) {
     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) {
+    TypeID id = TypeID::get<AnalysisT>();
+
+    typename ConceptMap::iterator it;
+    bool wasInserted;
+    std::tie(it, wasInserted) = analyses.try_emplace(id);
+
+    // If we don't have a cached analysis for this function, compute it directly
+    // and add it to the cache.
+    if (wasInserted) {
+      if (pi)
+        pi->runBeforeAnalysis(getAnalysisName<AnalysisT>(), id, ir);
+
+      it->second = std::make_unique<AnalysisModel<AnalysisT>>(op);
+
+      if (pi)
+        pi->runAfterAnalysis(getAnalysisName<AnalysisT>(), id, ir);
+    }
+    return static_cast<AnalysisModel<AnalysisT> &>(*it->second).analysis;
+  }
+
+private:
   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) : analyses(op) {}

 /// and cannot be constructed directly.
 class AnalysisManager {
   using ParentPointerT =
       PointerUnion<const ModuleAnalysisManager *, const AnalysisManager *>;
 
 public:
   using PreservedAnalyses = detail::PreservedAnalyses;
 
-  // Query for a cached analysis on the given parent operation. The analysis may
-  // not exist and if it does it may be out-of-date.
+  /// Query for a cached analysis on the given parent operation. The analysis
+  /// may not exist and if it does it may be out-of-date.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>>
   getCachedParentAnalysis(Operation *parentOp) const {
     ParentPointerT curParent = parent;
     while (auto *parentAM = curParent.dyn_cast<const AnalysisManager *>()) {
       if (parentAM->impl->getOperation() == parentOp)
         return parentAM->getCachedAnalysis<AnalysisT>();
       curParent = parentAM->parent;
     }
     return None;
   }
 
-  // Query for the given analysis for the current operation.
+  /// Query for the given analysis for the current operation.
   template <typename AnalysisT> AnalysisT &getAnalysis() {
     return impl->analyses.getAnalysis<AnalysisT>(getPassInstrumentor());
   }
 
-  // Query for a cached entry of the given analysis on the current operation.
+  /// 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());
+  }
+
+  /// 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>();
   }
 
   /// Query for an analysis of a child operation, constructing it if necessary.
   template <typename AnalysisT> AnalysisT &getChildAnalysis(Operation *op) {
     return slice(op).template getAnalysis<AnalysisT>();
   }
 
+  /// Query for an analysis of a child operation of a specifc derived operation
+  /// type, constructing it if necessary.
+  template <typename AnalysisT, typename OpT>
+  AnalysisT &getChildAnalysis(OpT child) {
+    return slice(child).template getAnalysis<AnalysisT, OpT>();
+  }
+
   /// Query for a cached analysis of a child operation, or return null.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>>
   getCachedChildAnalysis(Operation *op) const {
     assert(op->getParentOp() == impl->getOperation());
     auto it = impl->childAnalyses.find(op);
     if (it == impl->childAnalyses.end())
       return llvm::None;
     return it->second->analyses.getCachedAnalysis<AnalysisT>();
   }
 
   /// Get an analysis manager for the given child operation.
   AnalysisManager slice(Operation *op);
 
   /// Invalidate any non preserved analyses,
   void invalidate(const PreservedAnalyses &pa) { impl->invalidate(pa); }
 
   /// Clear any held analyses.
   void clear() {
     impl->analyses.clear();
     impl->childAnalyses.clear();

mlir/include/mlir/Pass/Pass.h

protected:
   /// be in an invalid state.
   void signalPassFailure() { getPassState().irAndPassFailed.setInt(true); }
 
   /// Query an analysis for the current ir unit.
   template <typename AnalysisT> AnalysisT &getAnalysis() {
     return getAnalysisManager().getAnalysis<AnalysisT>();
   }
 
+  /// Query an analysis for the current ir unit of a specific derived operation
+  /// type.
+  template <typename AnalysisT, typename OpT>
+  AnalysisT &getAnalysis() {
+    return getAnalysisManager().getAnalysis<AnalysisT, OpT>();
+  }
+
   /// Query a cached instance of an analysis for the current ir unit if one
   /// exists.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>> getCachedAnalysis() {
     return getAnalysisManager().getCachedAnalysis<AnalysisT>();
   }
 
   /// Mark all analyses as preserved.
   void markAllAnalysesPreserved() {
     getPassState().preservedAnalyses.preserveAll();
   }
 
   /// Mark the provided analyses as preserved.
   template <typename... AnalysesT> void markAnalysesPreserved() {
     getPassState().preservedAnalyses.preserve<AnalysesT...>();
   }
   void markAnalysesPreserved(TypeID id) {
     getPassState().preservedAnalyses.preserve(id);
   }
 
-  /// Returns the analysis for the parent operation if it exists.
+  /// Returns the analysis for the given parent operation if it exists.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>>
   getCachedParentAnalysis(Operation *parent) {
     return getAnalysisManager().getCachedParentAnalysis<AnalysisT>(parent);
   }
+
+  /// Returns the analysis for the parent operation if it exists.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>> getCachedParentAnalysis() {
     return getAnalysisManager().getCachedParentAnalysis<AnalysisT>(
         getOperation()->getParentOp());
   }
 
   /// Returns the analysis for the given child operation if it exists.
   template <typename AnalysisT>
   Optional<std::reference_wrapper<AnalysisT>>
   getCachedChildAnalysis(Operation *child) {
     return getAnalysisManager().getCachedChildAnalysis<AnalysisT>(child);
   }
 
   /// Returns the analysis for the given child operation, or creates it if it
   /// doesn't exist.
   template <typename AnalysisT> AnalysisT &getChildAnalysis(Operation *child) {
     return getAnalysisManager().getChildAnalysis<AnalysisT>(child);
   }
 
+  /// Returns the analysis for the given child operation of specific derived
+  /// operation type, or creates it if it doesn't exist.
+  template <typename AnalysisT, typename OpTy>
+  AnalysisT &getChildAnalysis(OpTy child) {
+    return getAnalysisManager().getChildAnalysis<AnalysisT>(child);
+  }
+
   /// Returns the current analysis manager.
   AnalysisManager getAnalysisManager() {
     return getPassState().analysisManager;
   }
 
   /// Create a copy of this pass, ignoring statistics and options.
   virtual std::unique_ptr<Pass> clonePass() const = 0;
 
protected:
 
   /// Support isa/dyn_cast functionality.
   static bool classof(const Pass *pass) {
     return pass->getOpName() == OpT::getOperationName();
   }
 
   /// Return the current operation being transformed.
   OpT getOperation() { return cast<OpT>(Pass::getOperation()); }
+
+  /// Query an analysis for the current operation of the specific derived
+  /// operation type.
+  template <typename AnalysisT>
+  AnalysisT &getAnalysis() {
+    return Pass::getAnalysis<AnalysisT, OpT>();
+  }
 };
 
 /// Pass to transform an operation.
 ///
 /// Operation passes must not:
 ///   - modify any other operations within the parent region, as other threads
 ///     may be manipulating them concurrently.
 ///   - modify any state within the parent operation, this includes adding

mlir/unittests/Pass/AnalysisManagerTest.cpp

 //===- AnalysisManagerTest.cpp - AnalysisManager unit tests ---------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Pass/AnalysisManager.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 namespace {
 /// Minimal class definitions for two analyses.
 struct MyAnalysis {
   MyAnalysis(Operation *) {}
 };
 struct OtherAnalysis {
   OtherAnalysis(Operation *) {}
 };
+struct OpSpecificAnalysis {
+  OpSpecificAnalysis(ModuleOp) {}
+};
 
 TEST(AnalysisManagerTest, FineGrainModuleAnalysisPreservation) {
   MLIRContext context;
 
   // Test fine grain invalidation of the module analysis manager.
   OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
   ModuleAnalysisManager mam(*module, /*passInstrumentor=*/nullptr);
   AnalysisManager am = mam;
TEST(AnalysisManagerTest, CustomInvalidation) {
   EXPECT_FALSE(am.getCachedAnalysis<CustomInvalidatingAnalysis>().hasValue());
 
   // Check that the analysis is preserved properly.
   am.getAnalysis<CustomInvalidatingAnalysis>();
   pa.preserve<TestAnalysisSet>();
   am.invalidate(pa);
   EXPECT_TRUE(am.getCachedAnalysis<CustomInvalidatingAnalysis>().hasValue());
 }
+
+TEST(AnalysisManagerTest, OpSpecificAnalysis) {
+  MLIRContext context;
+
+  // Create a module.
+  OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
+  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());
+}
+
 } // end namespace

mlir/unittests/Pass/CMakeLists.txt

 add_mlir_unittest(MLIRPassTests
   AnalysisManagerTest.cpp
+  PassManagerTest.cpp
 )
 target_link_libraries(MLIRPassTests
   PRIVATE
   MLIRPass)

mlir/unittests/Pass/PassManagerTest.cpp

+//===- PassManagerTest.cpp - PassManager unit tests -----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Pass/PassManager.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Function.h"
+#include "mlir/Pass/Pass.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::detail;
+
+namespace {
+/// Analysis that operates on any operation.
+struct GenericAnalysis {
+  GenericAnalysis(Operation *op) : isFunc(isa<FuncOp>(op)) {}
+  const bool isFunc;
+};
+
+/// Analysis that operates on a specific operation.
+struct OpSpecificAnalysis {
+  OpSpecificAnalysis(FuncOp op) : isSecret(op.getName() == "secret") {}
+  const bool isSecret;
+};
+
+/// Simple pass to annotate a FuncOp with the results of analysis.
+/// Note: not using FunctionPass as it skip external functions.
+struct AnnotateFunctionPass
+    : public PassWrapper<AnnotateFunctionPass, OperationPass<FuncOp>> {
+  void runOnOperation() override {
+    FuncOp op = getOperation();
+    Builder builder(op.getParentOfType<ModuleOp>());
+
+    auto &ga = getAnalysis<GenericAnalysis>();
+    auto &sa = getAnalysis<OpSpecificAnalysis>();
+
+    op.setAttr("isFunc", builder.getBoolAttr(ga.isFunc));
+    op.setAttr("isSecret", builder.getBoolAttr(sa.isSecret));
+  }
+};
+
+TEST(PassManagerTest, OpSpecificAnalysis) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  // Create a module with 2 functions.
+  OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
+  for (StringRef name : {"secret", "not_secret"}) {
+    FuncOp func =
+        FuncOp::create(builder.getUnknownLoc(), name,
+                       builder.getFunctionType(llvm::None, llvm::None));
+    module->push_back(func);
+  }
+
+  // Instantiate and run our pass.
+  PassManager pm(&context);
+  pm.addNestedPass<FuncOp>(std::make_unique<AnnotateFunctionPass>());
+  LogicalResult result = pm.run(module.get());
+  EXPECT_TRUE(succeeded(result));
+
+  // Verify that each function got annotated with expected attributes.
+  for (FuncOp func : module->getOps<FuncOp>()) {
+    ASSERT_TRUE(func.getAttr("isFunc").isa<BoolAttr>());
+    EXPECT_TRUE(func.getAttr("isFunc").cast<BoolAttr>().getValue());
+
+    bool isSecret = func.getName() == "secret";
+    ASSERT_TRUE(func.getAttr("isSecret").isa<BoolAttr>());
+    EXPECT_EQ(func.getAttr("isSecret").cast<BoolAttr>().getValue(), isSecret);
+  }
+}
+
+} // end namespace