[Mlir-commits] [mlir] [mlir] Attribute add printStripped (PR #78008)

[Jacques Pienaar]

https://github.com/jpienaar updated https://github.com/llvm/llvm-project/pull/78008

>From f6836d96ec499e042e90485a951497f4e808f893 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Fri, 12 Jan 2024 16:37:48 -0800
Subject: [PATCH 1/3] [mlir] Attribute add printStripped

Enable printing without dialect wrapping.

This closely matches `AsmPrinter::printStrippedAttrOrType`.
---
 mlir/include/mlir/IR/Attributes.h |  4 ++++
 mlir/lib/IR/AsmPrinter.cpp        | 31 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/mlir/include/mlir/IR/Attributes.h b/mlir/include/mlir/IR/Attributes.h
index f433363e5dedec..cc0cee6a31183c 100644
--- a/mlir/include/mlir/IR/Attributes.h
+++ b/mlir/include/mlir/IR/Attributes.h
@@ -78,6 +78,10 @@ class Attribute {
   void print(raw_ostream &os, AsmState &state, bool elideType = false) const;
   void dump() const;
 
+  /// Print the attribute without dialect wrapping.
+  void printStripped(raw_ostream &os) const;
+  void printStripped(raw_ostream &os, AsmState &state) const;
+
   /// Get an opaque pointer to the attribute.
   const void *getAsOpaquePointer() const { return impl; }
   /// Construct an attribute from the opaque pointer representation.
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 8fe8c78efecf9f..df2c85556406e9 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -3746,6 +3746,37 @@ void Attribute::dump() const {
   llvm::errs() << "\n";
 }
 
+void Attribute::printStripped(raw_ostream &os, AsmState &state) const {
+  if (!*this) {
+    os << "<<NULL ATTRIBUTE>>";
+    return;
+  }
+
+  AsmPrinter::Impl subPrinter(os, state.getImpl());
+  if (succeeded(subPrinter.printAlias(*this)))
+    return;
+
+  auto &dialect = this->getDialect();
+  uint64_t posPrior = os.tell();
+  DialectAsmPrinter printer(subPrinter);
+  dialect.printAttribute(*this, printer);
+  if (posPrior != os.tell())
+    return;
+
+  // Fallback to printing with prefix if the above failed to write anything
+  // to the output stream.
+  print(os, state);
+}
+void Attribute::printStripped(raw_ostream &os) const {
+  if (!*this) {
+    os << "<<NULL ATTRIBUTE>>";
+    return;
+  }
+
+  AsmState state(getContext());
+  printStripped(os, state);
+}
+
 void Type::print(raw_ostream &os) const {
   if (!*this) {
     os << "<<NULL TYPE>>";

>From c1d6205828652ba658821ae2dc8738133ed8d5af Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Tue, 16 Jan 2024 04:42:12 +0000
Subject: [PATCH 2/3] Add test

---
 mlir/unittests/IR/AttributeTest.cpp | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index 20167825d4ca82..438634fb99e16a 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -483,4 +483,22 @@ TEST(CopyCountAttr, CopyCount) {
 #endif
 }
 
+// Test stripped printing using test dialect attribute.
+TEST(CopyCountAttr, PrintStripped) {
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+  // Doesn't matter which dialect attribute is used, just chose TestCopyCount
+  // given proximity.
+  test::CopyCount::counter = 0;
+  test::CopyCount copyCount("hello");
+  auto res = test::TestCopyCountAttr::get(&context, std::move(copyCount));
+
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  os << "|" << res << "|";
+  res.printStripped(os << "[");
+  os << "]";
+  EXPECT_EQ(os.str(), "|#test.copy_count<hello>|[copy_count<hello>]");
+}
+
 } // namespace

>From 8f8392313b01941c883d8bc622bf119a5fce1833 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Tue, 16 Jan 2024 04:55:23 +0000
Subject: [PATCH 3/3] Update test

---
 mlir/unittests/IR/AttributeTest.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index 438634fb99e16a..e72bfe9d82e7cf 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -491,7 +491,7 @@ TEST(CopyCountAttr, PrintStripped) {
   // given proximity.
   test::CopyCount::counter = 0;
   test::CopyCount copyCount("hello");
-  auto res = test::TestCopyCountAttr::get(&context, std::move(copyCount));
+  Attribute res = test::TestCopyCountAttr::get(&context, std::move(copyCount));
 
   std::string str;
   llvm::raw_string_ostream os(str);