fix(PackConstants): don't start new storage buffer if current one is still empty (#23927)

The PackConstants pass introduces a new buffer when a constant cannot be
packed into the current buffer without exceeding the maximum size.
However, if a constant is itself larger than the limit, it gets still
added, but in a separate buffer.
This change fixes a bug in the packing logic when a constant bigger than
the limit is encountered when the current buffer is still empty. In this
case, the current buffer stayed empty and the big constant got added to
a new buffer. The empty buffer caused errors of size 0 buffers not being
supported in following passes or in the runtime. This change avoids
adding a new buffer is the current one is empty and thus avoids these
errors.

Signed-off-by: Stefan Schuermans <schuermans@roofline.ai>

compiler/src/iree/compiler/Dialect/Stream/Transforms/PackConstants.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackConstants.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackConstants.cpp
index e5a1dc7..b17c8b4 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackConstants.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackConstants.cpp
@@ -101,8 +101,13 @@
     uint64_t unpaddedLength = slice.getStorageSize();
     uint64_t paddedLength = IREE::Util::align(
         unpaddedLength, resourceConfig.getMinBufferRangeAlignment());
-    if (offset + unpaddedLength > resourceConfig.getMaxAllocationSize()) {
-      // Spilling buffer; make a new one.
+    // If unpaddedLength is bigger than resourceConfig.getMaxAllocationSize(),
+    // there no way to stay below the threshold. In this case, make sure that
+    // no new storage buffer is allocated if the current one is empty, so check
+    // current offset.
+    if (offset > 0 &&
+        offset + unpaddedLength > resourceConfig.getMaxAllocationSize()) {
+      // Spilling non-empty buffer; make a new one.
       storageBuffers.push_back({UnknownLoc::get(resourceConfig.getContext())});
       currentBuffer = &storageBuffers.back();
       offset = 0;

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/pack_constants.mlir

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/pack_constants.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/pack_constants.mlir
index bb9d934..a7002c3 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/pack_constants.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/pack_constants.mlir
@@ -137,3 +137,42 @@
   // CHECK: util.return %[[RES0]], %[[RES1]], %[[IF1]]#0
   util.return %0#0, %0#1, %0#2 : !stream.resource<constant>, !stream.resource<constant>, !stream.timepoint
 }
+
+// -----
+
+// Tests that a single constant larger than the max allocation size still gets
+// its own buffer (no infinite loop or error). A 256-byte constant exceeds the
+// 128 byte limit, but because offset is 0 (buffer is empty) it is not spilled,
+// so no empty buffer is created.
+
+#singleConstantExceedsLimitConfig = #stream.resource_config<{
+  max_allocation_size = 128,
+  min_buffer_offset_alignment = 16,
+  max_buffer_range = 1073741824,
+  min_buffer_range_alignment = 16,
+  index_bits = 32
+}>
+
+// CHECK-NOT: #composite_of_0b = #util.composite<0x{{[[:alnum:]]+}}, []>
+// CHECK: #composite_of_256b = #util.composite<256xi8, [
+// CHECK:     dense<42> : tensor<64xi32>,
+// CHECK: ]>
+
+// CHECK-LABEL: @singleConstantExceedsLimit
+util.func public @singleConstantExceedsLimit() -> (!stream.resource<constant>, !stream.timepoint)
+    attributes {stream.resources = #singleConstantExceedsLimitConfig} {
+  %c256 = arith.constant 256 : index
+
+  // CHECK-NOT: %[[RODATA:.+]] = util.buffer.constant {alignment = {{[:digit:]]+}} : index} : !util.buffer = #composite_of_0b
+  // CHECK: %[[RODATA:.+]] = util.buffer.constant {alignment = 16 : index} : !util.buffer = #composite_of_256b
+  // CHECK: %[[DID_MAP:.+]], %[[TRY_MAP:.+]] = stream.resource.try_map %[[RODATA]]
+  // CHECK: %[[IF:.+]]:2 = scf.if %[[DID_MAP]]
+  // CHECK: %[[RES:.+]] = stream.resource.subview %[[IF]]#1[%c0] : !stream.resource<constant>{%c256} -> !stream.resource<constant>{%c256}
+
+  %0:2 = stream.resource.constants :
+    !stream.resource<constant>{%c256} = dense<42> : tensor<64xi32>
+    => !stream.timepoint
+
+  // CHECK: util.return %[[RES]], %[[IF]]#0
+  util.return %0#0, %0#1 : !stream.resource<constant>, !stream.timepoint
+}