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

[Jacques Pienaar]

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

Enable printing without dialect wrapping.

This closely matches `AsmPrinter::printStrippedAttrOrType` implementation wise except templating component.

>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] [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>>";