[Printing] Add flag to not print resources.

Often times, large weights for ML models will be stored as resources in MLIR. It is sometimes advantageous to control whether to print these resources for debugging purposes. For example, some models contain very big weights with millions of characters in printed size, which may slow down whatever text editor you are using.

This diff adds a flag which allows users to disable printing resources in these scenarios.

Reviewed By: Mogball

Differential Revision: https://reviews.llvm.org/D157928

Author: Mogball

mlir/include/mlir/IR/OperationSupport.h

@@ -1139,6 +1139,9 @@ class OpPrintingFlags {
   /// Return the size limit for printing large ElementsAttr.
   std::optional<int64_t> getLargeElementsAttrLimit() const;
 
+  /// Return the size limit in chars for printing large resources.
+  std::optional<uint64_t> getLargeResourceStringLimit() const;
+
   /// Return if debug information should be printed.
   bool shouldPrintDebugInfo() const;
 
@@ -1165,6 +1168,9 @@ class OpPrintingFlags {
   /// the upper limit.
   std::optional<int64_t> elementsAttrElementLimit;
 
+  /// Elide printing large resources based on size of string.
+  std::optional<uint64_t> resourceStringCharLimit;
+
   /// Print debug information.
   bool printDebugInfoFlag : 1;
   bool printDebugInfoPrettyFormFlag : 1;

mlir/lib/IR/AsmPrinter.cpp

@@ -43,6 +43,7 @@
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/Threading.h"
+#include "llvm/Support/raw_ostream.h"
 
 #include <optional>
 #include <tuple>
@@ -140,6 +141,11 @@ struct AsmPrinterOptions {
       llvm::cl::desc("Elide ElementsAttrs with \"...\" that have "
                      "more elements than the given upper limit")};
 
+  llvm::cl::opt<unsigned> elideResourceStringsIfLarger{
+      "mlir-elide-resource-strings-if-larger",
+      llvm::cl::desc(
+          "Elide printing value of resources if string is too long in chars.")};
+
   llvm::cl::opt<bool> printDebugInfoOpt{
       "mlir-print-debuginfo", llvm::cl::init(false),
       llvm::cl::desc("Print debug info in MLIR output")};
@@ -191,6 +197,8 @@ OpPrintingFlags::OpPrintingFlags()
     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
+  if (clOptions->elideResourceStringsIfLarger.getNumOccurrences())
+    resourceStringCharLimit = clOptions->elideResourceStringsIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
   printDebugInfoPrettyFormFlag = clOptions->printPrettyDebugInfoOpt;
   printGenericOpFormFlag = clOptions->printGenericOpFormOpt;
@@ -262,6 +270,11 @@ std::optional<int64_t> OpPrintingFlags::getLargeElementsAttrLimit() const {
   return elementsAttrElementLimit;
 }
 
+/// Return the size limit for printing large ElementsAttr.
+std::optional<uint64_t> OpPrintingFlags::getLargeResourceStringLimit() const {
+  return resourceStringCharLimit;
+}
+
 /// Return if debug information should be printed.
 bool OpPrintingFlags::shouldPrintDebugInfo() const {
   return printDebugInfoFlag;
@@ -3085,11 +3098,11 @@ class OperationPrinter : public AsmPrinter::Impl, private OpAsmPrinter {
     ~ResourceBuilder() override = default;
 
     void buildBool(StringRef key, bool data) final {
-      printFn(key, [&](raw_ostream &os) { p.os << (data ? "true" : "false"); });
+      printFn(key, [&](raw_ostream &os) { os << (data ? "true" : "false"); });
     }
 
     void buildString(StringRef key, StringRef data) final {
-      printFn(key, [&](raw_ostream &os) { p.printEscapedString(data); });
+      printFn(key, [&](raw_ostream &os) { os << "\"" << data << "\""; });
     }
 
     void buildBlob(StringRef key, ArrayRef<char> data,
@@ -3176,23 +3189,41 @@ void OperationPrinter::printResourceFileMetadata(
     auto printFn = [&](StringRef key, ResourceBuilder::ValueFn valueFn) {
       checkAddMetadataDict();
 
-      // Emit the top-level resource entry if we haven't yet.
-      if (!std::exchange(hadResource, true)) {
-        if (needResourceComma)
-          os << "," << newLine;
-        os << "  " << dictName << "_resources: {" << newLine;
-      }
-      // Emit the parent resource entry if we haven't yet.
-      if (!std::exchange(hadEntry, true)) {
-        if (needEntryComma)
+      auto printFormatting = [&]() {
+        // Emit the top-level resource entry if we haven't yet.
+        if (!std::exchange(hadResource, true)) {
+          if (needResourceComma)
+            os << "," << newLine;
+          os << "  " << dictName << "_resources: {" << newLine;
+        }
+        // Emit the parent resource entry if we haven't yet.
+        if (!std::exchange(hadEntry, true)) {
+          if (needEntryComma)
+            os << "," << newLine;
+          os << "    " << name << ": {" << newLine;
+        } else {
           os << "," << newLine;
-        os << "    " << name << ": {" << newLine;
+        }
+      };
+
+      std::optional<uint64_t> charLimit =
+          printerFlags.getLargeResourceStringLimit();
+      if (charLimit.has_value()) {
+        std::string resourceStr;
+        llvm::raw_string_ostream ss(resourceStr);
+        valueFn(ss);
+
+        // Only print entry if it's string is small enough
+        if (resourceStr.size() > charLimit.value())
+          return;
+
+        printFormatting();
+        os << "      " << key << ": " << resourceStr;
       } else {
-        os << "," << newLine;
+        printFormatting();
+        os << "      " << key << ": ";
+        valueFn(os);
       }
-
-      os << "      " << key << ": ";
-      valueFn(os);
     };
     ResourceBuilder entryBuilder(*this, printFn);
     provider.buildResources(op, providerArgs..., entryBuilder);

mlir/test/IR/pretty-resources-print.mlir

@@ -0,0 +1,41 @@
+// Check printing with --mlir-elide-resource-strings-if-larger elides printing large resources
+
+// RUN: mlir-opt %s --mlir-elide-resource-strings-if-larger=10| FileCheck %s
+
+// To ensure we print the resource keys, have reference to them
+// CHECK: attr = dense_resource<blob1> : tensor<3xi64>
+"test.blob1op"() {attr = dense_resource<blob1> : tensor<3xi64> } : () -> ()
+
+// CHECK-NEXT: attr = dense_resource<blob2> : tensor<3xi64>
+"test.blob2op"() {attr = dense_resource<blob2> : tensor<3xi64> } : () -> ()
+
+// CHECK:      {-#
+// CHECK-NEXT:   external_resources: {
+// CHECK-NEXT:     external: {
+// CHECK-NEXT:       bool: true,
+// CHECK-NEXT:       string: "string"
+// CHECK-NEXT:     },
+// CHECK-NEXT:     other_stuff: {
+// CHECK-NEXT:       bool: true
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: #-}
+
+{-#
+  dialect_resources: {
+    builtin: {
+      blob1: "0x08000000010000000000000002000000000000000300000000000000",
+      blob2: "0x08000000040000000000000005000000000000000600000000000000"
+    }
+  },
+  external_resources: {
+    external: {
+      blob: "0x08000000010000000000000002000000000000000300000000000000",
+      bool: true,
+      string: "string"
+    },
+    other_stuff: {
+      bool: true
+    }
+  }
+#-}