WIP: thread AnalysisManager through verifier

[ftynse]

Summary

https://discourse.llvm.org/t/connecting-op-verifiers-to-analyses/3553

[ftynse]

Rebase

mlir/include/mlir/IR/OpDefinition.h

public:
   static void getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {}
 
 protected:
   /// If the concrete type didn't implement a custom verifier hook, just fall
   /// back to this one which accepts everything.
   LogicalResult verify() { return success(); }
 
+  LogicalResult verifyAnalyzed(AnalysisManager *analysisManager) {
+    return verify();
+  }
+
   /// Unless overridden, the custom assembly form of an op is always rejected.
   /// Op implementations should implement this to return failure.
   /// On success, they should fill in result with the fields to use.
   static ParseResult parse(OpAsmParser &parser, OperationState &result);
 
   // The fallback for the printer is to print it the generic assembly form.
   static void print(Operation *op, OpAsmPrinter &p);
 
private:
 
   static constexpr bool hasNoDataMembers() {
     // Checking that the derived class does not define any member by comparing
     // its size to an ad-hoc EmptyOp.
     class EmptyOp : public Op<EmptyOp, Traits...> {};
     return sizeof(ConcreteType) == sizeof(EmptyOp);
   }
 
-  static LogicalResult verifyInvariants(Operation *op) {
+  static LogicalResult verifyInvariants(Operation *op,
+                                        AnalysisManager *analysisManager) {
     static_assert(hasNoDataMembers(),
                   "Op class shouldn't define new data members");
     return failure(
         failed(op_definition_impl::verifyTraits<VerifiableTraitsTupleT>(op)) ||
-        failed(cast<ConcreteType>(op).verify()));
+        failed(cast<ConcreteType>(op).verifyAnalyzed(analysisManager)));
   }
 
   /// Allow access to internal implementation methods.
   friend AbstractOperation;
 };
 
 /// This class represents the base of an operation interface. See the definition
 /// of `detail::Interface` for requirements on the `Traits` type.

mlir/include/mlir/IR/OperationSupport.h

 #include "llvm/Support/TrailingObjects.h"
 #include <memory>
 
 namespace llvm {
 class BitVector;
 } // end namespace llvm
 
 namespace mlir {
+class AnalysisManager;
 class Dialect;
 class DictionaryAttr;
 class ElementsAttr;
 class MutableOperandRangeRange;
 class Operation;
 struct OperationState;
 class OpAsmParser;
 class OpAsmParserResult;
using GetCanonicalizationPatternsFn =
       llvm::unique_function<void(RewritePatternSet &, MLIRContext *) const>;
   using FoldHookFn = llvm::unique_function<LogicalResult(
       Operation *, ArrayRef<Attribute>, SmallVectorImpl<OpFoldResult> &) const>;
   using HasTraitFn = llvm::unique_function<bool(TypeID) const>;
   using ParseAssemblyFn =
       llvm::unique_function<ParseResult(OpAsmParser &, OperationState &) const>;
   using PrintAssemblyFn =
       llvm::unique_function<void(Operation *, OpAsmPrinter &) const>;
-  using VerifyInvariantsFn =
-      llvm::unique_function<LogicalResult(Operation *) const>;
+  using VerifyInvariantsFn = llvm::unique_function<LogicalResult(
+      Operation *, AnalysisManager *) const>;
 
   /// This is the name of the operation.
   const Identifier name;
 
   /// This is the dialect that this operation belongs to.
   Dialect &dialect;
 
   /// The unique identifier of the derived Op class.
   TypeID typeID;
 
   /// Use the specified object to parse this ops custom assembly format.
   ParseResult parseAssembly(OpAsmParser &parser, OperationState &result) const;
 
   /// Return the static hook for parsing this operation assembly.
   const ParseAssemblyFn &getParseAssemblyFn() const { return parseAssemblyFn; }
 
   /// This hook implements the AsmPrinter for this operation.
   void printAssembly(Operation *op, OpAsmPrinter &p) const {
     return printAssemblyFn(op, p);
   }
 
   /// This hook implements the verifier for this operation.  It should emits an
   /// error message and returns failure if a problem is detected, or returns
   /// success if everything is ok.
-  LogicalResult verifyInvariants(Operation *op) const {
-    return verifyInvariantsFn(op);
+  LogicalResult verifyInvariants(Operation *op,
+                                 AnalysisManager *analysisManager) const {
+    return verifyInvariantsFn(op, analysisManager);
   }
 
   /// This hook implements a generalized folder for this operation.  Operations
   /// can implement this to provide simplifications rules that are applied by
   /// the Builder::createOrFold API and the canonicalization pass.
   ///
   /// This is an intentionally limited interface - implementations of this hook
   /// can only perform the following changes to the operation:

mlir/include/mlir/IR/Verifier.h

 //===- Verifier.h - Verifier analysis for MLIR structures -------*- 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_IR_VERIFIER_H
 #define MLIR_IR_VERIFIER_H
 
 namespace mlir {
+class AnalysisManager;
 struct LogicalResult;
 class Operation;
 
 /// Perform (potentially expensive) checks of invariants, used to detect
 /// compiler bugs, on this operation and any nested operations. On error, this
 /// reports the error through the MLIRContext and returns failure.
-LogicalResult verify(Operation *op);
+LogicalResult verify(Operation *op, AnalysisManager *analysisManager = nullptr);
 } //  end namespace mlir
 
 #endif

mlir/lib/IR/Verifier.cpp

 #include <atomic>
 
 using namespace mlir;
 
 namespace {
 /// This class encapsulates all the state used to verify an operation region.
 class OperationVerifier {
 public:
+  explicit OperationVerifier(AnalysisManager *analysisManager)
+      : analysisManager(analysisManager) {}
+
   /// Verify the given operation.
   LogicalResult verifyOpAndDominance(Operation &op);
 
 private:
   LogicalResult
   verifyBlock(Block &block,
               SmallVectorImpl<Operation *> &opsWithIsolatedRegions);
   /// Verify the properties and dominance relationships of this operation,
   /// stopping region recursion at any "isolated from above operations".  Any
   /// such ops are returned in the opsWithIsolatedRegions vector.
   LogicalResult
   verifyOperation(Operation &op,
                   SmallVectorImpl<Operation *> &opsWithIsolatedRegions);
 
   /// Verify the dominance property of regions contained within the given
   /// Operation.
   LogicalResult verifyDominanceOfContainedRegions(Operation &op,
                                                   DominanceInfo &domInfo);
+
+  AnalysisManager *analysisManager;
 };
 } // end anonymous namespace
 
 LogicalResult OperationVerifier::verifyOpAndDominance(Operation &op) {
   SmallVector<Operation *> opsWithIsolatedRegions;
 
   // Verify the operation first, collecting any IsolatedFromAbove operations.
   if (failed(verifyOperation(op, opsWithIsolatedRegions)))
for (auto attr : op.getAttrs()) {
     if (auto *dialect = attr.first.getDialect())
       if (failed(dialect->verifyOperationAttribute(&op, attr)))
         return failure();
   }
 
   // If we can get operation info for this, check the custom hook.
   OperationName opName = op.getName();
   auto *opInfo = opName.getAbstractOperation();
-  if (opInfo && failed(opInfo->verifyInvariants(&op)))
+  if (opInfo && failed(opInfo->verifyInvariants(&op, analysisManager)))
     return failure();
 
   if (unsigned numRegions = op.getNumRegions()) {
     auto kindInterface = dyn_cast<RegionKindInterface>(op);
 
     // Verify that all child regions are ok.
     for (unsigned i = 0; i < numRegions; ++i) {
       Region &region = op.getRegion(i);

 
 //===----------------------------------------------------------------------===//
 // Entrypoint
 //===----------------------------------------------------------------------===//
 
 /// Perform (potentially expensive) checks of invariants, used to detect
 /// compiler bugs.  On error, this reports the error through the MLIRContext and
 /// returns failure.
-LogicalResult mlir::verify(Operation *op) {
-  return OperationVerifier().verifyOpAndDominance(*op);
+LogicalResult mlir::verify(Operation *op, AnalysisManager *analysisManager) {
+  return OperationVerifier(analysisManager).verifyOpAndDominance(*op);
 }

mlir/lib/Pass/Pass.cpp

if (!passFailed && verifyPasses) {
     //     subunit (and those passes aren't allowed to modify the parent).
     //
     // We run these checks in EXPENSIVE_CHECKS mode out of caution.
 #ifndef EXPENSIVE_CHECKS
     runVerifierNow = !isa<OpToOpPassAdaptor>(pass) &&
                      !pass->passState->preservedAnalyses.isAll();
 #endif
     if (runVerifierNow)
-      passFailed = failed(verify(op));
+      passFailed = failed(verify(op, &am));
   }
 
   // Instrument after the pass has run.
   if (pi) {
     if (passFailed)
       pi->runAfterPassFailed(pass, op);
     else
       pi->runAfterPass(pass, op);

mlir/test/lib/Dialect/Test/TestDialect.h

 // that do not have a respective counterpart in the main source directories.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTDIALECT_H
 #define MLIR_TESTDIALECT_H
 
 #include "TestInterfaces.h"
+#include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/DLTI/Traits.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/DerivedAttributeOpInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/AnalysisManager.h"
 
 namespace mlir {
 class DLTIDialect;
 class RewritePatternSet;
 } // namespace mlir
 
 #include "TestOpEnums.h.inc"
 #include "TestOpInterfaces.h.inc"

mlir/test/lib/Dialect/Test/TestOps.td

def TestLocationDstOp : TEST_Op<"loc_dst", [SameOperandsAndResultType]> {
   let arguments = (ins I32:$input);
   let results = (outs I32:$output);
 }
 
 //===----------------------------------------------------------------------===//
 // Test Patterns
 //===----------------------------------------------------------------------===//
 
+def OpWithAMVerifier : TEST_Op<"op_with_am_verifier"> {
+  let extraClassDeclaration = [{
+    ::mlir::LogicalResult verifyAnalyzed(::mlir::AnalysisManager *am) {
+      if (am) {
+        llvm::errs() << "has analysis manager\n";
+        llvm::errs() << &am->getAnalysis<::mlir::DataLayoutAnalysis>() << "\n";
+      } else {
+        llvm::errs() << "no analysis manager\n";
+      }
+      return ::mlir::success();
+    }
+  }];
+}
+
 def OpA : TEST_Op<"op_a"> {
   let arguments = (ins I32, I32Attr:$attr);
   let results = (outs I32);
 }
 
 def OpB : TEST_Op<"op_b"> {
   let arguments = (ins I32, I32Attr:$attr);
   let results = (outs I32);

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

includes = [
         "lib/Dialect/Test",
     ],
     deps = [
         ":TestAttrDefsIncGen",
         ":TestInterfacesIncGen",
         ":TestOpsIncGen",
         ":TestTypeDefsIncGen",
         "//llvm:Support",
+        "//mlir:Analysis",
         "//mlir:ControlFlowInterfaces",
         "//mlir:CopyOpInterface",
         "//mlir:DLTIDialect",
         "//mlir:DataLayoutInterfaces",
         "//mlir:DerivedAttributeOpInterface",
         "//mlir:Dialect",
         "//mlir:IR",
         "//mlir:InferTypeOpInterface",