[mlir][remarks] Add support for Attributes in Remark Args (#170882)

Extend Remark::Arg to optionally store the original Attribute when
constructed from one. This allows custom remark streamers to access the
underlying attribute for type-specific handling, rather than being
limited to the string representation.

For example, a streamer processing an ArrayAttr argument can iterate
over its elements, or a streamer can check the specific attribute type
to apply custom formatting.

Changes:
- Add std::optional<Attribute> member to Remark::Arg
- Update Arg(StringRef, Attribute) constructor to store the attribute
- Add hasAttribute() and getAttribute() accessors
- Add unit test for Arg with ArrayAttr

Author: razvanlupusoru

mlir/include/mlir/IR/Remarks.h

@@ -99,18 +99,30 @@ class Remark {
   }
 
   // Remark argument that is a key-value pair that can be printed as machine
-  // parsable args.
+  // parsable args. For Attribute arguments, the original attribute is also
+  // stored to allow custom streamers to handle them specially.
   struct Arg {
     std::string key;
     std::string val;
+    /// Optional attribute storage for Attribute-based args. Allows streamers
+    /// to access the original attribute for custom handling.
+    std::optional<Attribute> attr;
+
     Arg(llvm::StringRef m) : key("Remark"), val(m) {}
     Arg(llvm::StringRef k, llvm::StringRef v) : key(k), val(v) {}
     Arg(llvm::StringRef k, std::string v) : key(k), val(std::move(v)) {}
     Arg(llvm::StringRef k, const char *v) : Arg(k, llvm::StringRef(v)) {}
     Arg(llvm::StringRef k, Value v);
     Arg(llvm::StringRef k, Type t);
+    Arg(llvm::StringRef k, Attribute a);
     Arg(llvm::StringRef k, bool b) : key(k), val(b ? "true" : "false") {}
 
+    /// Check if this arg has an associated attribute.
+    bool hasAttribute() const { return attr.has_value(); }
+
+    /// Get the attribute if present.
+    Attribute getAttribute() const { return attr.value_or(Attribute()); }
+
     // One constructor for all arithmetic types except bool.
     template <typename T, typename = std::enable_if_t<std::is_arithmetic_v<T> &&
                                                       !std::is_same_v<T, bool>>>

mlir/lib/IR/Remarks.cpp

@@ -31,6 +31,11 @@ Remark::Arg::Arg(llvm::StringRef k, Type t) : key(k) {
   os << t;
 }
 
+Remark::Arg::Arg(llvm::StringRef k, Attribute a) : key(k), attr(a) {
+  llvm::raw_string_ostream os(val);
+  os << a;
+}
+
 void Remark::insert(llvm::StringRef s) { args.emplace_back(s); }
 void Remark::insert(Arg a) { args.push_back(std::move(a)); }
 

mlir/unittests/IR/RemarkTest.cpp

@@ -6,6 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Remarks.h"
@@ -377,4 +379,35 @@ TEST(Remark, TestRemarkFinal) {
   EXPECT_NE(errOut.find(pass3Msg), std::string::npos); // shown
   EXPECT_NE(errOut.find(pass4Msg), std::string::npos); // shown
 }
+
+TEST(Remark, TestArgWithAttribute) {
+  MLIRContext context;
+
+  SmallVector<Attribute> elements;
+  elements.push_back(IntegerAttr::get(IntegerType::get(&context, 32), 1));
+  elements.push_back(IntegerAttr::get(IntegerType::get(&context, 32), 2));
+  elements.push_back(IntegerAttr::get(IntegerType::get(&context, 32), 3));
+  ArrayAttr arrayAttr = ArrayAttr::get(&context, elements);
+  remark::detail::Remark::Arg argWithArray("Values", arrayAttr);
+
+  // Verify the attribute is stored
+  EXPECT_TRUE(argWithArray.hasAttribute());
+  EXPECT_EQ(argWithArray.getAttribute(), arrayAttr);
+
+  // Ensure it can be retrieved as an ArrayAttr.
+  auto retrievedAttr = dyn_cast<ArrayAttr>(argWithArray.getAttribute());
+  EXPECT_TRUE(retrievedAttr);
+  EXPECT_EQ(retrievedAttr.size(), 3u);
+  EXPECT_EQ(cast<IntegerAttr>(retrievedAttr[0]).getInt(), 1);
+  EXPECT_EQ(cast<IntegerAttr>(retrievedAttr[1]).getInt(), 2);
+  EXPECT_EQ(cast<IntegerAttr>(retrievedAttr[2]).getInt(), 3);
+
+  // Create an Arg without an Attribute (string-based)
+  remark::detail::Remark::Arg argWithoutAttr("Key", "Value");
+
+  // Verify no attribute is stored
+  EXPECT_FALSE(argWithoutAttr.hasAttribute());
+  EXPECT_FALSE(argWithoutAttr.getAttribute()); // Returns null Attribute
+  EXPECT_EQ(argWithoutAttr.val, "Value");
+}
 } // namespace