[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

16

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

16

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

16

Dropping them would make more sense to me.

[AndrewZhaoLuo]

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

16

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

16

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

16

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

public:
   OpPrintingFlags &printValueUsers();
 
   /// Return if the given ElementsAttr should be elided.
   bool shouldElideElementsAttr(ElementsAttr attr) const;
 
   /// 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;
 
   /// Return if debug information should be printed in the pretty form.
   bool shouldPrintDebugInfoPrettyForm() const;
 
   /// Return if operations should be printed in the generic form.
   bool shouldPrintGenericOpForm() const;
public:
   /// Return if the printer should print users of values.
   bool shouldPrintValueUsers() const;
 
 private:
   /// Elide large elements attributes if the number of elements is larger than
   /// 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;
 
   /// Print operations in the generic form.
   bool printGenericOpFormFlag : 1;
 
   /// Always skip Regions.

mlir/lib/IR/AsmPrinter.cpp

 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Endian.h"
 #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>
 
 using namespace mlir;
 using namespace mlir::detail;
 
 #define DEBUG_TYPE "mlir-asm-printer"
llvm::cl::opt<int64_t> printElementsAttrWithHexIfLarger{
           "Print DenseElementsAttrs with a hex string that have "
           "more elements than the given upper limit (use -1 to disable)")};
 
   llvm::cl::opt<unsigned> elideElementsAttrIfLarger{
       "mlir-elide-elementsattrs-if-larger",
       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")};
 
   llvm::cl::opt<bool> printPrettyDebugInfoOpt{
       "mlir-pretty-debuginfo", llvm::cl::init(false),
       llvm::cl::desc("Print pretty debug info in MLIR output")};
 
: printDebugInfoFlag(false), printDebugInfoPrettyFormFlag(false),
       printGenericOpFormFlag(false), skipRegionsFlag(false),
       assumeVerifiedFlag(false), printLocalScope(false),
       printValueUsersFlag(false) {
   // Initialize based upon command line options, if they are available.
   if (!clOptions.isConstructed())
     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
+  if (clOptions->elideResourceStringsIfLarger.getNumOccurrences())
+    resourceStringCharLimit = clOptions->elideResourceStringsIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
   printDebugInfoPrettyFormFlag = clOptions->printPrettyDebugInfoOpt;
   printGenericOpFormFlag = clOptions->printGenericOpFormOpt;
   assumeVerifiedFlag = clOptions->assumeVerifiedOpt;
   printLocalScope = clOptions->printLocalScopeOpt;
   printValueUsersFlag = clOptions->printValueUsers;
 }
 
return elementsAttrElementLimit &&
          !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.
+std::optional<uint64_t> OpPrintingFlags::getLargeResourceStringLimit() const {
+  return resourceStringCharLimit;
+}
+
 /// Return if debug information should be printed.
 bool OpPrintingFlags::shouldPrintDebugInfo() const {
   return printDebugInfoFlag;
 }
 
 /// Return if debug information should be printed in the pretty form.
 bool OpPrintingFlags::shouldPrintDebugInfoPrettyForm() const {
   return printDebugInfoPrettyFormFlag;
public:
     using ValueFn = function_ref<void(raw_ostream &)>;
     using PrintFn = function_ref<void(StringRef, ValueFn)>;
 
     ResourceBuilder(OperationPrinter &p, PrintFn printFn)
         : p(p), printFn(printFn) {}
     ~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,
                    uint32_t dataAlignment) final {
       printFn(key, [&](raw_ostream &os) {
         // Store the blob in a hex string containing the alignment and the data.
         llvm::support::ulittle32_t dataAlignmentLE(dataAlignment);
         os << "\"0x"
void OperationPrinter::printResourceFileMetadata(
   bool needResourceComma = false;
   bool needEntryComma = false;
   auto processProvider = [&](StringRef dictName, StringRef name, auto &provider,
                              auto &&...providerArgs) {
     bool hadEntry = false;
     auto printFn = [&](StringRef key, ResourceBuilder::ValueFn valueFn) {
       checkAddMetadataDict();
 
+      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;
         }
+      };
 
+      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 {
+        printFormatting();
         os << "      " << key << ": ";
         valueFn(os);
+      }
     };
     ResourceBuilder entryBuilder(*this, printFn);
     provider.buildResources(op, providerArgs..., entryBuilder);
 
     needEntryComma |= hadEntry;
     if (hadEntry)
       os << newLine << "    }";
   };

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:   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
+    }
+  }
+#-}