[mlir][Printer] Honor --mlir-elide-resource-strings-if-larger when serializing of large blobs

Project: llvm/llvm-project | Language: C++ | Concepts: memory-management, algorithmic-optimization | Difficulty: mid
Link: https://github.com/llvm/llvm-project/pull/213920

The problem

When printing MLIR IR with resource blobs (e.g., large binary data), the printer supports an option (--mlir-elide-resource-strings-if-larger) to elide (skip printing) resource strings that are too large. However, the code always fully serializes (e.g., hex-encodes) the resource data first, and only checks the length afterward. For large blobs, this wastes significant time and memory encoding data that will be discarded.

Why it's non-trivial

A naive fix might simply check the size of the original data, but the actual serialized size (e.g., after hex encoding and formatting) can be much larger and isn't always trivial to compute. For some resource types, the serialized size is not known in advance. The fix must avoid unnecessary work for large blobs, but still correctly handle all resource types and preserve output for boundary cases.

The approach

1. Extend the PrintFn callback to accept a size hint:

-    using PrintFn = function_ref<void(StringRef, ValueFn)>;
+    // `sizeHint` is the exact number of characters `valueFn` will write, or -1
+    // if unknown, so the char limit can be applied before paying the cost of
+    // invoking `valueFn` (e.g. hex-encoding a large blob).
+    using PrintFn = function_ref<void(StringRef, ValueFn, int64_t sizeHint)>;

This changes the callback signature so that, when building a resource, the builder can optionally provide the exact output size (if known).

2. Compute the size hint for blobs, pass it to the callback:

-      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"
-           << llvm::toHex(StringRef(reinterpret_cast<char *>(&dataAlignmentLE),
-                                    sizeof(dataAlignment)))
-           << llvm::toHex(StringRef(data.data(), data.size())) << "\"";
-      });
+      // Two hex chars per byte of the alignment word and the data, plus the
+      // `"0x`/`"` wrapping; exact, so the limit can be checked pre-encoding.
+      int64_t sizeHint = 2 * int64_t(sizeof(dataAlignment) + data.size()) + 4;
+      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"
+               << llvm::toHex(
+                      StringRef(reinterpret_cast<char *>(&dataAlignmentLE),
+                                sizeof(dataAlignment)))
+               << llvm::toHex(StringRef(data.data(), data.size())) << "\"";
+          },
+          sizeHint);

For blobs, the code now computes the exact number of characters the output will take, so it can avoid hex-encoding if it's too large.

3. In the actual print function, skip serialization if the size hint is over the limit:

-    auto printFn = [&](StringRef key, ResourceBuilder::ValueFn valueFn) {
+    auto printFn = [&](StringRef key, ResourceBuilder::ValueFn valueFn,
+                       int64_t sizeHint) {
       checkAddMetadataDict();
 
       std::string resourceStr;
       if (charLimit) {
         if (charLimit.value() == 0)
           return;
 
+        // Skip serializing entirely if the exact size already exceeds the
+        // limit, e.g. hex-encoding a large blob.
+        if (sizeHint >= 0 && uint64_t(sizeHint) > charLimit.value()) {
+          LLVM_DEBUG(llvm::dbgs() << "eliding resource '" << key
+                                  << "' without materializing (sizeHint="
+                                  << sizeHint << ")\n");
+          return;
+        }

If the size hint is known and over the limit, the expensive serialization is skipped entirely, saving time and memory.

4. Tests for boundary cases were added:

+// blob3's exact serialized size (quotes + "0x" + hex(alignment) + hex(data)) is 14 chars;
+// these two RUN lines check the off-by-one boundary of the sizeHint-based elision.
+// RUN: mlir-opt %s --mlir-elide-resource-strings-if-larger=14| FileCheck %s --check-prefix=BOUND14
+// RUN: mlir-opt %s --mlir-elide-resource-strings-if-larger=13| FileCheck %s --check-prefix=BOUND13
...
+// At the exact boundary (limit == blob3's exact size) blob3 must still be printed,
+// since the sizeHint fast-path only elides when size is strictly greater than the limit.
+// BOUND14:      {-#
+// BOUND14-NEXT:   dialect_resources: {
+// BOUND14-NEXT:     builtin: {
+// BOUND14-NEXT:       blob3: "0x0800000001"
...
+// One below the boundary, blob3 is elided and no dialect_resources dict is emitted.
+// Note: the escaped `string` entry is also exactly 14 chars, so it is elided here too.
+// BOUND13:      {-#
+// BOUND13-NEXT:   external_resources: {
+// BOUND13-NEXT:     external: {
+// BOUND13-NEXT:       "backslash\\tab\09": true

The tests verify that the new logic correctly handles the boundary between elided and non-elided output.

From the review thread

No notable review discussion on this PR.

Takeaways

• Avoid expensive computation if you can determine up front that the result will be discarded.
• When optimizing for memory or performance, consider whether you can cheaply compute the size or cost of an operation before performing it.
• When changing logic that depends on output size or limits, add tests for boundary conditions (e.g., exactly at the limit, just below, just above).