[mlir] implement verify-only-specified-diagnostics

Author: makslevental

mlir/examples/transform-opt/mlir-transform-opt.cpp

@@ -43,6 +43,12 @@ struct MlirTransformOptCLOptions {
       cl::desc("Check that emitted diagnostics match expected-* lines "
                "on the corresponding line"),
       cl::init(false)};
+  cl::opt<bool> verifyOnlyExpectedDiagnostics{
+      "verify-only-expected-diagnostics",
+      cl::desc("Check that emitted diagnostics match only specified expected-* "
+               "lines "
+               "on the corresponding line"),
+      cl::init(false)};
 
   cl::opt<std::string> payloadFilename{cl::Positional, cl::desc("<input file>"),
                                        cl::init("-")};
@@ -103,11 +109,13 @@ class DiagnosticHandlerWrapper {
   /// Constructs the diagnostic handler of the specified kind of the given
   /// source manager and context.
   DiagnosticHandlerWrapper(Kind kind, llvm::SourceMgr &mgr,
-                           mlir::MLIRContext *context) {
+                           mlir::MLIRContext *context,
+                           bool verifyOnlyExpectedDiagnostics = false) {
     if (kind == Kind::EmitDiagnostics)
       handler = new mlir::SourceMgrDiagnosticHandler(mgr, context);
     else
-      handler = new mlir::SourceMgrDiagnosticVerifierHandler(mgr, context);
+      handler = new mlir::SourceMgrDiagnosticVerifierHandler(
+          mgr, context, verifyOnlyExpectedDiagnostics);
   }
 
   /// This object is non-copyable but movable.
@@ -150,8 +158,10 @@ class TransformSourceMgr {
 public:
   /// Constructs the source manager indicating whether diagnostic messages will
   /// be verified later on.
-  explicit TransformSourceMgr(bool verifyDiagnostics)
-      : verifyDiagnostics(verifyDiagnostics) {}
+  explicit TransformSourceMgr(bool verifyDiagnostics,
+                              bool verifyOnlyExpectedDiagnostics)
+      : verifyDiagnostics(verifyDiagnostics),
+        verifyOnlyExpectedDiagnostics(verifyOnlyExpectedDiagnostics) {}
 
   /// Deconstructs the source manager. Note that `checkResults` must have been
   /// called on this instance before deconstructing it.
@@ -179,7 +189,8 @@ class TransformSourceMgr {
     // verification needs to happen and store it.
     if (verifyDiagnostics) {
       diagHandlers.emplace_back(
-          DiagnosticHandlerWrapper::Kind::VerifyDiagnostics, mgr, &context);
+          DiagnosticHandlerWrapper::Kind::VerifyDiagnostics, mgr, &context,
+          verifyOnlyExpectedDiagnostics);
     } else {
       diagHandlers.emplace_back(DiagnosticHandlerWrapper::Kind::EmitDiagnostics,
                                 mgr, &context);
@@ -205,6 +216,9 @@ class TransformSourceMgr {
 private:
   /// Indicates whether diagnostic message verification is requested.
   const bool verifyDiagnostics;
+  /// Indicates whether *only specified* diagnostic message verification is
+  /// requested.
+  const bool verifyOnlyExpectedDiagnostics;
 
   /// Indicates that diagnostic message verification has taken place, and the
   /// deconstruction is therefore safe.
@@ -248,7 +262,8 @@ static llvm::LogicalResult processPayloadBuffer(
   context.allowUnregisteredDialects(clOptions->allowUnregisteredDialects);
   mlir::ParserConfig config(&context);
   TransformSourceMgr sourceMgr(
-      /*verifyDiagnostics=*/clOptions->verifyDiagnostics);
+      /*verifyDiagnostics=*/clOptions->verifyDiagnostics,
+      clOptions->verifyOnlyExpectedDiagnostics);
 
   // Parse the input buffer that will be used as transform payload.
   mlir::OwningOpRef<mlir::Operation *> payloadRoot =

mlir/include/mlir/IR/Diagnostics.h

@@ -626,9 +626,12 @@ struct SourceMgrDiagnosticVerifierHandlerImpl;
 /// corresponding line of the source file.
 class SourceMgrDiagnosticVerifierHandler : public SourceMgrDiagnosticHandler {
 public:
-  SourceMgrDiagnosticVerifierHandler(llvm::SourceMgr &srcMgr, MLIRContext *ctx,
-                                     raw_ostream &out);
-  SourceMgrDiagnosticVerifierHandler(llvm::SourceMgr &srcMgr, MLIRContext *ctx);
+  SourceMgrDiagnosticVerifierHandler(
+      llvm::SourceMgr &srcMgr, MLIRContext *ctx, raw_ostream &out,
+      bool verifyOnlyExpectedDiagnostics = false);
+  SourceMgrDiagnosticVerifierHandler(
+      llvm::SourceMgr &srcMgr, MLIRContext *ctx,
+      bool verifyOnlyExpectedDiagnostics = false);
   ~SourceMgrDiagnosticVerifierHandler();
 
   /// Returns the status of the handler and verifies that all expected
@@ -644,6 +647,10 @@ class SourceMgrDiagnosticVerifierHandler : public SourceMgrDiagnosticHandler {
   void process(LocationAttr loc, StringRef msg, DiagnosticSeverity kind);
 
   std::unique_ptr<detail::SourceMgrDiagnosticVerifierHandlerImpl> impl;
+
+  /// Set whether to check that emitted diagnostics match *only specified*
+  /// `expected-*` lines on the corresponding line.
+  bool verifyOnlyExpectedDiagnostics = false;
 };
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

@@ -191,6 +191,16 @@ class MlirOptMainConfig {
   }
   bool shouldVerifyDiagnostics() const { return verifyDiagnosticsFlag; }
 
+  /// Set whether to check that emitted diagnostics match *only specified*
+  /// `expected-*` lines on the corresponding line.
+  MlirOptMainConfig &verifyOnlyExpectedDiagnostics(bool verify) {
+    verifyOnlyExpectedDiagnosticsFlag = verify;
+    return *this;
+  }
+  bool shouldVerifyOnlyExpectedDiagnostics() const {
+    return verifyOnlyExpectedDiagnosticsFlag;
+  }
+
   /// Set whether to run the verifier after each transformation pass.
   MlirOptMainConfig &verifyPasses(bool verify) {
     verifyPassesFlag = verify;
@@ -277,6 +287,9 @@ class MlirOptMainConfig {
   /// Set whether to check that emitted diagnostics match `expected-*` lines on
   /// the corresponding line. This is meant for implementing diagnostic tests.
   bool verifyDiagnosticsFlag = false;
+  /// Set whether to check that emitted diagnostics match *only specified*
+  /// `expected-*` lines on the corresponding line.
+  bool verifyOnlyExpectedDiagnosticsFlag = false;
 
   /// Run the verifier after each transformation pass.
   bool verifyPassesFlag = true;

mlir/lib/IR/Diagnostics.cpp

@@ -803,9 +803,11 @@ SourceMgrDiagnosticVerifierHandlerImpl::computeExpectedDiags(
 }
 
 SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
-    llvm::SourceMgr &srcMgr, MLIRContext *ctx, raw_ostream &out)
+    llvm::SourceMgr &srcMgr, MLIRContext *ctx, raw_ostream &out,
+    bool verifyOnlyExpectedDiagnostics)
     : SourceMgrDiagnosticHandler(srcMgr, ctx, out),
-      impl(new SourceMgrDiagnosticVerifierHandlerImpl()) {
+      impl(new SourceMgrDiagnosticVerifierHandlerImpl()),
+      verifyOnlyExpectedDiagnostics(verifyOnlyExpectedDiagnostics) {
   // Compute the expected diagnostics for each of the current files in the
   // source manager.
   for (unsigned i = 0, e = mgr.getNumBuffers(); i != e; ++i)
@@ -823,8 +825,10 @@ SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
 }
 
 SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
-    llvm::SourceMgr &srcMgr, MLIRContext *ctx)
-    : SourceMgrDiagnosticVerifierHandler(srcMgr, ctx, llvm::errs()) {}
+    llvm::SourceMgr &srcMgr, MLIRContext *ctx,
+    bool verifyOnlyExpectedDiagnostics)
+    : SourceMgrDiagnosticVerifierHandler(srcMgr, ctx, llvm::errs(),
+                                         verifyOnlyExpectedDiagnostics) {}
 
 SourceMgrDiagnosticVerifierHandler::~SourceMgrDiagnosticVerifierHandler() {
   // Ensure that all expected diagnostics were handled.
@@ -898,6 +902,9 @@ void SourceMgrDiagnosticVerifierHandler::process(LocationAttr loc,
     }
   }
 
+  if (verifyOnlyExpectedDiagnostics)
+    return;
+
   // Otherwise, emit an error for the near miss.
   if (nearMiss)
     mgr.PrintMessage(os, nearMiss->fileLoc, llvm::SourceMgr::DK_Error,

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

@@ -168,6 +168,14 @@ struct MlirOptMainConfigCLOptions : public MlirOptMainConfig {
         cl::desc("Check that emitted diagnostics match "
                  "expected-* lines on the corresponding line"),
         cl::location(verifyDiagnosticsFlag), cl::init(false));
+    static cl::opt<bool, /*ExternalStorage=*/true>
+        verifyOnlyExpectedDiagnostics{
+            "verify-only-expected-diagnostics",
+            cl::desc("Check that emitted diagnostics match only specified "
+                     "expected-* "
+                     "lines "
+                     "on the corresponding line"),
+            cl::location(verifyOnlyExpectedDiagnosticsFlag), cl::init(false)};
 
     static cl::opt<bool, /*ExternalStorage=*/true> verifyPasses(
         "verify-each",
@@ -537,7 +545,8 @@ static LogicalResult processBuffer(raw_ostream &os,
     return performActions(os, sourceMgr, &context, config);
   }
 
-  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(*sourceMgr, &context);
+  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(
+      *sourceMgr, &context, config.shouldVerifyOnlyExpectedDiagnostics());
 
   // Do any processing requested by command line flags.  We don't care whether
   // these actions succeed or fail, we only care what diagnostics they produce

mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp

@@ -58,7 +58,8 @@ LogicalResult mlir::mlirTranslateMain(int argc, char **argv,
 
   static llvm::cl::opt<bool> allowUnregisteredDialects(
       "allow-unregistered-dialect",
-      llvm::cl::desc("Allow operation with no registered dialects (discouraged: testing only!)"),
+      llvm::cl::desc("Allow operation with no registered dialects "
+                     "(discouraged: testing only!)"),
       llvm::cl::init(false));
 
   static llvm::cl::opt<std::string> inputSplitMarker{
@@ -77,6 +78,13 @@ LogicalResult mlir::mlirTranslateMain(int argc, char **argv,
       llvm::cl::desc("Check that emitted diagnostics match "
                      "expected-* lines on the corresponding line"),
       llvm::cl::init(false));
+  static llvm::cl::opt<bool> verifyOnlyExpectedDiagnostics{
+      "verify-only-expected-diagnostics",
+      llvm::cl::desc(
+          "Check that emitted diagnostics match only specified expected-* "
+          "lines "
+          "on the corresponding line"),
+      llvm::cl::init(false)};
 
   static llvm::cl::opt<bool> errorDiagnosticsOnly(
       "error-diagnostics-only",
@@ -158,8 +166,8 @@ LogicalResult mlir::mlirTranslateMain(int argc, char **argv,
         // translation failed (in most cases, it is expected to fail) and we do
         // not filter non-error diagnostics even if `errorDiagnosticsOnly` is
         // set. Instead, we check if the diagnostics were produced as expected.
-        SourceMgrDiagnosticVerifierHandler sourceMgrHandler(*sourceMgr,
-                                                            &context);
+        SourceMgrDiagnosticVerifierHandler sourceMgrHandler(
+            *sourceMgr, &context, verifyOnlyExpectedDiagnostics);
         (void)(*translationRequested)(sourceMgr, os, &context);
         result = sourceMgrHandler.verify();
       } else if (errorDiagnosticsOnly) {