diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h --- a/mlir/include/mlir/IR/OperationSupport.h +++ b/mlir/include/mlir/IR/OperationSupport.h @@ -1139,6 +1139,9 @@ /// Return the size limit for printing large ElementsAttr. std::optional getLargeElementsAttrLimit() const; + /// Return the size limit in chars for printing large resources. + std::optional getLargeResourceStringLimit() const; + /// Return if debug information should be printed. bool shouldPrintDebugInfo() const; @@ -1165,6 +1168,9 @@ /// the upper limit. std::optional elementsAttrElementLimit; + /// Elide printing large resources based on size of string. + std::optional resourceStringCharLimit; + /// Print debug information. bool printDebugInfoFlag : 1; bool printDebugInfoPrettyFormFlag : 1; diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp --- 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 #include @@ -140,6 +141,11 @@ llvm::cl::desc("Elide ElementsAttrs with \"...\" that have " "more elements than the given upper limit")}; + llvm::cl::opt 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 printDebugInfoOpt{ "mlir-print-debuginfo", llvm::cl::init(false), llvm::cl::desc("Print debug info in MLIR output")}; @@ -191,6 +197,8 @@ 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 @@ return elementsAttrElementLimit; } +/// Return the size limit for printing large ElementsAttr. +std::optional OpPrintingFlags::getLargeResourceStringLimit() const { + return resourceStringCharLimit; +} + /// Return if debug information should be printed. bool OpPrintingFlags::shouldPrintDebugInfo() const { return printDebugInfoFlag; @@ -3085,11 +3098,11 @@ ~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 data, @@ -3176,23 +3189,41 @@ 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 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 --- /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 : tensor<3xi64> +"test.blob1op"() {attr = dense_resource : tensor<3xi64> } : () -> () + +// CHECK-NEXT: attr = dense_resource : tensor<3xi64> +"test.blob2op"() {attr = dense_resource : 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 + } + } +#-}