[Mlir-commits] [mlir] [mlir][OpPrintingFlags] Allow to disable ElementsAttr hex printing (PR #85766)

[Andrei Golubev]

https://github.com/andrey-golubev updated https://github.com/llvm/llvm-project/pull/85766

>From 095619831593e48d3f8afe9c5011836071d6019e Mon Sep 17 00:00:00 2001
From: "Golubev, Andrey" <andrey.golubev at intel.com>
Date: Tue, 19 Mar 2024 09:39:28 +0000
Subject: [PATCH] [mlir][OpPrintingFlags] Allow to disable ElementsAttr hex
 printing

At present, large ElementsAttr is unconditionally printed with a hex
string. This means that in IR large constant values often look like:
dense<"0x000000000004000000080000000004000000080000000..."> : tensor<10x10xi32>

Hoisting hex printing control to the user level for tooling means that one can
disable the feature and get human-readable values when necessary:
dense<[16, 32, 48, 500...]> : tensor<10x10xi32>

Note: AsmPrinterOptions::printElementsAttrWithHexIfLarger is not always
possible to be used as it requires that one exposes MLIR's command-line
options in user tooling (including an actual compiler).

Co-authored-by: Harald Rotuna <harald.razvan.rotuna at intel.com>
---
 mlir/include/mlir/IR/OperationSupport.h | 10 ++++++++++
 mlir/lib/IR/AsmPrinter.cpp              | 25 ++++++++++++++++++++++---
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index f2aa6cee840308..7d17efa7ce9b5e 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1131,6 +1131,10 @@ class OpPrintingFlags {
   /// elements.
   OpPrintingFlags &elideLargeElementsAttrs(int64_t largeElementLimit = 16);
 
+  /// Enable or disable printing of large element attributes as a hex string. By
+  /// default, printing as hex is allowed (enabled).
+  OpPrintingFlags &allowPrintingElementsAttrWithHex(bool allowHex = true);
+
   /// Enables the elision of large resources strings by omitting them from the
   /// `dialect_resources` section. The `largeResourceLimit` is used to configure
   /// what is considered to be a "large" resource by providing an upper limit to
@@ -1164,6 +1168,9 @@ class OpPrintingFlags {
   /// Return if the given ElementsAttr should be elided.
   bool shouldElideElementsAttr(ElementsAttr attr) const;
 
+  /// Return if ElementsAttr could be printed in hexadecimal form.
+  bool shouldAllowPrintingElementsAttrWithHex() const;
+
   /// Return the size limit for printing large ElementsAttr.
   std::optional<int64_t> getLargeElementsAttrLimit() const;
 
@@ -1217,6 +1224,9 @@ class OpPrintingFlags {
 
   /// Print users of values.
   bool printValueUsersFlag : 1;
+
+  /// Print ElementsAttr in hex form.
+  bool allowPrintingElementsAttrWithHexFlag : 1;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 456cf6a2c27783..693977f35110a1 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -206,12 +206,17 @@ OpPrintingFlags::OpPrintingFlags()
     : printDebugInfoFlag(false), printDebugInfoPrettyFormFlag(false),
       printGenericOpFormFlag(false), skipRegionsFlag(false),
       assumeVerifiedFlag(false), printLocalScope(false),
-      printValueUsersFlag(false) {
+      printValueUsersFlag(false), allowPrintingElementsAttrWithHexFlag(true) {
   // Initialize based upon command line options, if they are available.
   if (!clOptions.isConstructed())
     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
+  if (clOptions->printElementsAttrWithHexIfLarger.getNumOccurrences()) {
+    // Note: -1 disables the "print with hex string" feature
+    allowPrintingElementsAttrWithHexFlag =
+        (clOptions->printElementsAttrWithHexIfLarger.getValue() != -1);
+  }
   if (clOptions->elideResourceStringsIfLarger.getNumOccurrences())
     resourceStringCharLimit = clOptions->elideResourceStringsIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
@@ -233,6 +238,12 @@ OpPrintingFlags::elideLargeElementsAttrs(int64_t largeElementLimit) {
   return *this;
 }
 
+OpPrintingFlags &
+OpPrintingFlags::allowPrintingElementsAttrWithHex(bool allowHex) {
+  allowPrintingElementsAttrWithHexFlag = allowHex;
+  return *this;
+}
+
 OpPrintingFlags &
 OpPrintingFlags::elideLargeResourceString(int64_t largeResourceLimit) {
   resourceStringCharLimit = largeResourceLimit;
@@ -287,6 +298,10 @@ bool OpPrintingFlags::shouldElideElementsAttr(ElementsAttr attr) const {
          !llvm::isa<SplatElementsAttr>(attr);
 }
 
+bool OpPrintingFlags::shouldAllowPrintingElementsAttrWithHex() const {
+  return allowPrintingElementsAttrWithHexFlag;
+}
+
 /// Return the size limit for printing large ElementsAttr.
 std::optional<int64_t> OpPrintingFlags::getLargeElementsAttrLimit() const {
   return elementsAttrElementLimit;
@@ -2301,7 +2316,9 @@ void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
       printElidedElementsAttr(os);
     } else {
       os << "dense<";
-      printDenseIntOrFPElementsAttr(intOrFpEltAttr, /*allowHex=*/true);
+      printDenseIntOrFPElementsAttr(
+          intOrFpEltAttr,
+          printerFlags.shouldAllowPrintingElementsAttrWithHex());
       os << '>';
     }
 
@@ -2324,7 +2341,9 @@ void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
       if (indices.getNumElements() != 0) {
         printDenseIntOrFPElementsAttr(indices, /*allowHex=*/false);
         os << ", ";
-        printDenseElementsAttr(sparseEltAttr.getValues(), /*allowHex=*/true);
+        printDenseElementsAttr(
+            sparseEltAttr.getValues(),
+            printerFlags.shouldAllowPrintingElementsAttrWithHex());
       }
       os << '>';
     }