[Mlir-commits] [mlir] 8a77917 - [Printing] Add flag to not print resources.
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 23 09:54:30 PDT 2023
Author: Mogball
Date: 2023-08-23T16:54:23Z
New Revision: 8a779174134d7d01c6835cb0254b8a6dbf419150
URL: https://github.com/llvm/llvm-project/commit/8a779174134d7d01c6835cb0254b8a6dbf419150
DIFF: https://github.com/llvm/llvm-project/commit/8a779174134d7d01c6835cb0254b8a6dbf419150.diff
LOG: [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
Added:
mlir/test/IR/pretty-resources-print.mlir
Modified:
mlir/include/mlir/IR/OperationSupport.h
mlir/lib/IR/AsmPrinter.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index f3a79eb52f8ec0..b33a861a92021c 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/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;
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index f626d39cb6cf27..8bae108d13cf26 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/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);
diff --git a/mlir/test/IR/pretty-resources-print.mlir b/mlir/test/IR/pretty-resources-print.mlir
new file mode 100644
index 00000000000000..ea6e9f4b28c280
--- /dev/null
+++ b/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
+ }
+ }
+#-}
More information about the Mlir-commits
mailing list