[mlir] Add metadata to Diagnostic. (#99398)

Summary:
Add metadata to Diagnostic. 

Motivation: we have a use case where we want to do some filtering in our
customized Diagnostic Handler based on some customized info that is not
`location` or `severity` or `diagnostic arguments` that are member
variables of `Diagnostic`. Specifically, we want to add a unique ID to
the `Diagnostic` for the handler to filter in a compiler pass that emits
errors in async tasks with multithreading and the diagnostic handling is
associated to the task.

This patch adds a field of `metadata` to `mlir::Diagnostics` as a
general solution. `metadata` is of type `SmallVector<DiagnosticArgument,
0>` to save memory size and reuse existing `DiagnosticArgument` for
metadata type.

Test Plan: 

Reviewers: 

Subscribers: 

Tasks: 

Tags: 


Differential Revision: https://phabricator.intern.facebook.com/D60250607

mlir/docs/Diagnostics.md

@@ -77,14 +77,19 @@ InFlightDiagnostic Operation::emitOpError();
 ## Diagnostic
 
 A `Diagnostic` in MLIR contains all of the necessary information for reporting a
-message to the user. A `Diagnostic` essentially boils down to three main
+message to the user. A `Diagnostic` essentially boils down to four main
 components:
 
 *   [Source Location](#source-locations)
 *   Severity Level
     -   Error, Note, Remark, Warning
 *   Diagnostic Arguments
     -   The diagnostic arguments are used when constructing the output message.
+*   Metadata
+    -   Some additional information attached that can be used to identify 
+        this diagnostic other than source location and severity level 
+        (e.g. for diagnostic handlers to do some filtering). 
+        Metadata is not part of the output message.
 
 ### Appending arguments
 
@@ -143,6 +148,11 @@ op->emitError("...").attachNote(noteLoc) << "...";
 op->emitError("...").attachNote() << "...";
 ```
 
+### Managing Metadata
+Metadata is a mutable vector of DiagnosticArguments. 
+It can be accessed and modified as a vector. 
+
+
 ## InFlight Diagnostic
 
 Now that [Diagnostics](#diagnostic) have been explained, we introduce the

mlir/include/mlir/IR/Diagnostics.h

@@ -271,6 +271,9 @@ class Diagnostic {
     return failure();
   }
 
+  /// Returns the current list of diagnostic metadata.
+  SmallVectorImpl<DiagnosticArgument> &getMetadata() { return metadata; }
+
 private:
   Diagnostic(const Diagnostic &rhs) = delete;
   Diagnostic &operator=(const Diagnostic &rhs) = delete;
@@ -290,6 +293,9 @@ class Diagnostic {
 
   /// A list of attached notes.
   NoteVector notes;
+
+  /// A list of metadata attached to this Diagnostic.
+  SmallVector<DiagnosticArgument, 0> metadata;
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, const Diagnostic &diag) {

mlir/test/IR/diagnostic-handler-metadata.mlir

@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(test-diagnostic-metadata))" -verify-diagnostics -o - 2>&1 | FileCheck %s
+// COM: This test verifies that diagnostic handler can filter the diagnostic based on its metadata
+// COM: whether to emit the errors.
+
+// CHECK-LABEL: Test 'test'
+func.func @test() {
+  // expected-error @+1 {{test diagnostic metadata}}
+  "test.emit_error"() {
+    // CHECK: attr = "emit_error"
+    attr = "emit_error"
+  } : () -> ()
+
+  "test.do_not_emit_error"() {
+    // CHECK: attr = "do_not_emit_error"
+    attr = "do_not_emit_error"
+  } : () -> ()
+
+  return
+}

mlir/test/lib/IR/CMakeLists.txt

@@ -6,6 +6,7 @@ add_mlir_library(MLIRTestIR
   TestBuiltinDistinctAttributes.cpp
   TestClone.cpp
   TestDiagnostics.cpp
+  TestDiagnosticsMetadata.cpp
   TestDominance.cpp
   TestFunc.cpp
   TestInterfaces.cpp

mlir/test/lib/IR/TestDiagnosticsMetadata.cpp

@@ -0,0 +1,65 @@
+//===- TestDiagnosticsMetadata.cpp - Test Diagnostic Metatdata ------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for constructing and resolving dominance
+// information.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/SourceMgr.h"
+
+using namespace mlir;
+
+namespace {
+struct TestDiagnosticMetadataPass
+    : public PassWrapper<TestDiagnosticMetadataPass,
+                         InterfacePass<SymbolOpInterface>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDiagnosticMetadataPass)
+
+  StringRef getArgument() const final { return "test-diagnostic-metadata"; }
+  StringRef getDescription() const final { return "Test diagnostic metadata."; }
+  TestDiagnosticMetadataPass() = default;
+  TestDiagnosticMetadataPass(const TestDiagnosticMetadataPass &) {}
+
+  void runOnOperation() override {
+    llvm::errs() << "Test '" << getOperation().getName() << "'\n";
+
+    // Build a diagnostic handler that has filtering capabilities.
+    ScopedDiagnosticHandler handler(&getContext(), [](mlir::Diagnostic &diag) {
+      return mlir::success(
+          llvm::none_of(diag.getMetadata(), [](mlir::DiagnosticArgument &arg) {
+            return arg.getKind() == mlir::DiagnosticArgument::
+                                        DiagnosticArgumentKind::String &&
+                   arg.getAsString().contains("hello");
+          }));
+    });
+
+    // Emit a diagnostic for every operation with a valid loc.
+    getOperation()->walk([&](Operation *op) {
+      if (StringAttr strAttr = op->getAttrOfType<StringAttr>("attr")) {
+        if (strAttr.getValue() == "emit_error")
+          emitError(op->getLoc(), "test diagnostic metadata")
+              .getUnderlyingDiagnostic()
+              ->getMetadata()
+              .push_back(DiagnosticArgument("hello"));
+      }
+    });
+  }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestDiagnosticsMetadataPass() {
+  PassRegistration<TestDiagnosticMetadataPass>{};
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -92,6 +92,7 @@ void registerTestDataLayoutQuery();
 void registerTestDeadCodeAnalysisPass();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
+void registerTestDiagnosticsMetadataPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestEmulateNarrowTypePass();
@@ -226,6 +227,7 @@ void registerTestPasses() {
   mlir::test::registerTestDeadCodeAnalysisPass();
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDiagnosticsPass();
+  mlir::test::registerTestDiagnosticsMetadataPass();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestEmulateNarrowTypePass();