[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 dba43c58790457f27ecefbbfe411f913ae8f6ba6 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 | 17 ++++++++++
 mlir/lib/IR/AsmPrinter.cpp              | 43 +++++++++++++------------
 2 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 90e63ff8fcb38f0..2c1c490aac49b80 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1136,6 +1136,13 @@ class OpPrintingFlags {
   /// elements.
   OpPrintingFlags &elideLargeElementsAttrs(int64_t largeElementLimit = 16);
 
+  /// Enables the printing of large element attributes with a hex string. The
+  /// `largeElementLimit` is used to configure what is considered to be a
+  /// "large" ElementsAttr by providing an upper limit to the number of
+  /// elements. Use -1 to disable the hex printing.
+  OpPrintingFlags &
+  printLargeElementsAttrWithHex(int64_t largeElementLimit = 100);
+
   /// 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
@@ -1169,9 +1176,15 @@ class OpPrintingFlags {
   /// Return if the given ElementsAttr should be elided.
   bool shouldElideElementsAttr(ElementsAttr attr) const;
 
+  /// Return if the given ElementsAttr should be printed as hex string.
+  bool shouldPrintElementsAttrWithHex(ElementsAttr attr) const;
+
   /// Return the size limit for printing large ElementsAttr.
   std::optional<int64_t> getLargeElementsAttrLimit() const;
 
+  /// Return the size limit for printing large ElementsAttr as hex string.
+  int64_t getLargeElementsAttrHexLimit() const;
+
   /// Return the size limit in chars for printing large resources.
   std::optional<uint64_t> getLargeResourceStringLimit() const;
 
@@ -1204,6 +1217,10 @@ class OpPrintingFlags {
   /// Elide printing large resources based on size of string.
   std::optional<uint64_t> resourceStringCharLimit;
 
+  /// Print large element attributes with hex strings if the number of elements
+  /// is larger than the upper limit.
+  int64_t elementsAttrHexElementLimit = 100;
+
   /// Print debug information.
   bool printDebugInfoFlag : 1;
   bool printDebugInfoPrettyFormFlag : 1;
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 456cf6a2c27783e..e915b97d9ff17b1 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -212,6 +212,9 @@ OpPrintingFlags::OpPrintingFlags()
     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
+  if (clOptions->printElementsAttrWithHexIfLarger.getNumOccurrences())
+    elementsAttrHexElementLimit =
+        clOptions->printElementsAttrWithHexIfLarger.getValue();
   if (clOptions->elideResourceStringsIfLarger.getNumOccurrences())
     resourceStringCharLimit = clOptions->elideResourceStringsIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
@@ -233,6 +236,12 @@ OpPrintingFlags::elideLargeElementsAttrs(int64_t largeElementLimit) {
   return *this;
 }
 
+OpPrintingFlags &
+OpPrintingFlags::printLargeElementsAttrWithHex(int64_t largeElementLimit) {
+  elementsAttrHexElementLimit = largeElementLimit;
+  return *this;
+}
+
 OpPrintingFlags &
 OpPrintingFlags::elideLargeResourceString(int64_t largeResourceLimit) {
   resourceStringCharLimit = largeResourceLimit;
@@ -287,11 +296,24 @@ bool OpPrintingFlags::shouldElideElementsAttr(ElementsAttr attr) const {
          !llvm::isa<SplatElementsAttr>(attr);
 }
 
+/// Return if the given ElementsAttr should be printed as hex string.
+bool OpPrintingFlags::shouldPrintElementsAttrWithHex(ElementsAttr attr) const {
+  // -1 is used to disable hex printing.
+  return (elementsAttrHexElementLimit != -1) &&
+         (elementsAttrHexElementLimit < int64_t(attr.getNumElements())) &&
+         !llvm::isa<SplatElementsAttr>(attr);
+}
+
 /// Return the size limit for printing large ElementsAttr.
 std::optional<int64_t> OpPrintingFlags::getLargeElementsAttrLimit() const {
   return elementsAttrElementLimit;
 }
 
+/// Return the size limit for printing large ElementsAttr as hex string.
+int64_t OpPrintingFlags::getLargeElementsAttrHexLimit() const {
+  return elementsAttrHexElementLimit;
+}
+
 /// Return the size limit for printing large ElementsAttr.
 std::optional<uint64_t> OpPrintingFlags::getLargeResourceStringLimit() const {
   return resourceStringCharLimit;
@@ -328,23 +350,6 @@ bool OpPrintingFlags::shouldPrintValueUsers() const {
   return printValueUsersFlag;
 }
 
-/// Returns true if an ElementsAttr with the given number of elements should be
-/// printed with hex.
-static bool shouldPrintElementsAttrWithHex(int64_t numElements) {
-  // Check to see if a command line option was provided for the limit.
-  if (clOptions.isConstructed()) {
-    if (clOptions->printElementsAttrWithHexIfLarger.getNumOccurrences()) {
-      // -1 is used to disable hex printing.
-      if (clOptions->printElementsAttrWithHexIfLarger == -1)
-        return false;
-      return numElements > clOptions->printElementsAttrWithHexIfLarger;
-    }
-  }
-
-  // Otherwise, default to printing with hex if the number of elements is >100.
-  return numElements > 100;
-}
-
 //===----------------------------------------------------------------------===//
 // NewLineCounter
 //===----------------------------------------------------------------------===//
@@ -2435,9 +2440,7 @@ void AsmPrinter::Impl::printDenseIntOrFPElementsAttr(
   auto elementType = type.getElementType();
 
   // Check to see if we should format this attribute as a hex string.
-  auto numElements = type.getNumElements();
-  if (!attr.isSplat() && allowHex &&
-      shouldPrintElementsAttrWithHex(numElements)) {
+  if (allowHex && printerFlags.shouldPrintElementsAttrWithHex(attr)) {
     ArrayRef<char> rawData = attr.getRawData();
     if (llvm::endianness::native == llvm::endianness::big) {
       // Convert endianess in big-endian(BE) machines. `rawData` is BE in BE