[Printing] Add flag to not print resources.

[AndrewZhaoLuo]

Summary

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.

[Mogball]

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

15

You could just have the resource references be dropped entirely. They are not required to be present in the IR.

[AndrewZhaoLuo]

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

15

If you feel super strongly about this I can do this. My thought around this was:

1. didn't know you could drop the resources
2. wanted to keep printing as close as possible to original text as possible

[Mogball]

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

15

Dropping them would make more sense to me.

[AndrewZhaoLuo]

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

15

IMO it may be sometimes helpful to have them anyway.

For example, you might store small "things" in resources which might want to know the values of for debugging.

See the use of "true" in the test which is not elided.

[Mogball]

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

15

I'm not suggesting that all resources be dropped, but that instead of replacing them with a garbage value that the elided ones be completely dropped from the IR. All resource references are potentially nullable which means this keeps the IR valid and roundtrippable.

[AndrewZhaoLuo]

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

15

Ah I see, I think that does make sense and is better. I guess the only argument against is it may be helpful to know what resources were defined in the original IR.

I have made your suggested changes.

[Mogball]

mlir/lib/IR/AsmPrinter.cpp

3193

As a minor optimization, can you skip the string buffer if charLimit is not defined?

[AndrewZhaoLuo]

Add path to not use string buffer if flag not there

[AndrewZhaoLuo]

Done

[jpienaar]

mlir/lib/IR/AsmPrinter.cpp

3105

Why was this change needed?

[AndrewZhaoLuo]

mlir/lib/IR/AsmPrinter.cpp

3105

The old implementation didn't used the passed in ostream.

printEscapedString isn't in ostream so I just manually do the equivalent here.

[Mogball]

mlir/lib/IR/AsmPrinter.cpp

3105

printEscapedString does more than just wrap in quotes

[AndrewZhaoLuo]

mlir/lib/IR/AsmPrinter.cpp

3105

Ah darn you are right. Probably glossed it over. I'll fix and add a test later.

https://github.com/llvm/llvm-project/blob/8a779174134d7d01c6835cb0254b8a6dbf419150/mlir/lib/IR/AsmPrinter.cpp#L2673

[AndrewZhaoLuo]

mlir/lib/IR/AsmPrinter.cpp

3105

https://reviews.llvm.org/D158700 Fix. Sorry about that.

mlir/include/mlir/IR/OperationSupport.h

   /// 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<int64_t> getLargeResourceStringLimit() const;
+
   /// Return if debug information should be printed.
   bool shouldPrintDebugInfo() const;
 

   /// the upper limit.
   std::optional<int64_t> elementsAttrElementLimit;
 
+  /// Elide printing large resources based on size of string.
+  std::optional<int64_t> resourceStringCharLimit;
+
   /// Print debug information.
   bool printDebugInfoFlag : 1;
   bool printDebugInfoPrettyFormFlag : 1;

mlir/lib/IR/AsmPrinter.cpp

 #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>

       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")};

     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
+  if (clOptions->elideResourceStringsIfLarger.getNumOccurrences())
+    resourceStringCharLimit = clOptions->elideResourceStringsIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
   printDebugInfoPrettyFormFlag = clOptions->printPrettyDebugInfoOpt;
   printGenericOpFormFlag = clOptions->printGenericOpFormOpt;

   return elementsAttrElementLimit;
 }
 
+/// Return the size limit for printing large ElementsAttr.
+std::optional<int64_t> OpPrintingFlags::getLargeResourceStringLimit() const {
+  return resourceStringCharLimit;
+}
+
 /// Return if debug information should be printed.
 bool OpPrintingFlags::shouldPrintDebugInfo() const {
   return printDebugInfoFlag;

       }
 
       os << "      " << key << ": ";
-      valueFn(os);
+
+      // Print only if the character limit is met
+      std::optional<int64_t> charLimit =
+          printerFlags.getLargeResourceStringLimit();
+      if (charLimit.has_value()) {
+        std::string s;
+        llvm::raw_string_ostream ss(s);
+        valueFn(ss);
+        if (s.size() > charLimit)
+          os << "\"__elided__\"";
+        else
+          os << s;
+      } else {
+        valueFn(os);
+      }
     };
     ResourceBuilder entryBuilder(*this, printFn);
     provider.buildResources(op, providerArgs..., entryBuilder);

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

+// 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:   dialect_resources: {
+// CHECK-NEXT:     builtin: {
+// CHECK-NEXT:       blob1: "__elided__",
+// CHECK-NEXT:       blob2: "__elided__"
+// CHECK-NEXT:     }
+// CHECK-NEXT:   },
+// CHECK-NEXT:   external_resources: {
+// CHECK-NEXT:     external: {
+// CHECK-NEXT:       blob: "__elided__",
+// 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
+    }
+  }
+#-}