[mlir][Diagnostic] Don't store Operation arguments as a DiagnosticArgument

Summary: Diagnostics may be cached in the parallel diagnostic handler to preserve proper ordering. Storing the Operation as a DiagnosticArgument is problematic as the operation may be erased or changed before it finally gets printed.

Differential Revision: https://reviews.llvm.org/D77675

mlir/include/mlir/IR/Diagnostics.h

@@ -57,7 +57,6 @@ public:
     Attribute,
     Double,
     Integer,
-    Operation,
     String,
     Type,
     Unsigned,
@@ -84,12 +83,6 @@ public:
     return static_cast<int64_t>(opaqueVal);
   }
 
-  /// Returns this argument as an operation.
-  Operation &getAsOperation() const {
-    assert(getKind() == DiagnosticArgumentKind::Operation);
-    return *reinterpret_cast<Operation *>(opaqueVal);
-  }
-
   /// Returns this argument as a string.
   StringRef getAsString() const {
     assert(getKind() == DiagnosticArgumentKind::String);
@@ -132,14 +125,6 @@ private:
                                      sizeof(T) <= sizeof(uint64_t)>::type * = 0)
       : kind(DiagnosticArgumentKind::Unsigned), opaqueVal(uint64_t(val)) {}
 
-  // Construct from an operation reference.
-  explicit DiagnosticArgument(Operation &val) : DiagnosticArgument(&val) {}
-  explicit DiagnosticArgument(Operation *val)
-      : kind(DiagnosticArgumentKind::Operation),
-        opaqueVal(reinterpret_cast<intptr_t>(val)) {
-    assert(val && "expected valid operation");
-  }
-
   // Construct from a string reference.
   explicit DiagnosticArgument(StringRef val)
       : kind(DiagnosticArgumentKind::String), stringVal(val) {}
@@ -229,6 +214,12 @@ public:
   /// Stream in an OperationName.
   Diagnostic &operator<<(OperationName val);
 
+  /// Stream in an Operation.
+  Diagnostic &operator<<(Operation &val);
+  Diagnostic &operator<<(Operation *val) {
+    return *this << *val;
+  }
+
   /// Stream in a range.
   template <typename T> Diagnostic &operator<<(iterator_range<T> range) {
     return appendRange(range);

mlir/lib/IR/Diagnostics.cpp

@@ -30,12 +30,12 @@ using namespace mlir::detail;
 // DiagnosticArgument
 //===----------------------------------------------------------------------===//
 
-// Construct from an Attribute.
+/// Construct from an Attribute.
 DiagnosticArgument::DiagnosticArgument(Attribute attr)
     : kind(DiagnosticArgumentKind::Attribute),
       opaqueVal(reinterpret_cast<intptr_t>(attr.getAsOpaquePointer())) {}
 
-// Construct from a Type.
+/// Construct from a Type.
 DiagnosticArgument::DiagnosticArgument(Type val)
     : kind(DiagnosticArgumentKind::Type),
       opaqueVal(reinterpret_cast<intptr_t>(val.getAsOpaquePointer())) {}
@@ -65,9 +65,6 @@ void DiagnosticArgument::print(raw_ostream &os) const {
   case DiagnosticArgumentKind::Integer:
     os << getAsInteger();
     break;
-  case DiagnosticArgumentKind::Operation:
-    getAsOperation().print(os, OpPrintingFlags().useLocalScope());
-    break;
   case DiagnosticArgumentKind::String:
     os << getAsString();
     break;
@@ -125,6 +122,14 @@ Diagnostic &Diagnostic::operator<<(OperationName val) {
   return *this;
 }
 
+/// Stream in an Operation.
+Diagnostic &Diagnostic::operator<<(Operation &val) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  os << val;
+  return *this << os.str();
+}
+
 /// Outputs this diagnostic to a stream.
 void Diagnostic::print(raw_ostream &os) const {
   for (auto &arg : getArguments())