ggml-hrx: host staging buffers are charged to the pool too (engine #1711)

ggml/src/ggml-hrx/runtime/command-program-executor.cpp

@@ -757,7 +757,8 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec
         }
 
         HostStagingBuffer staging;
-        Status            allocation_status = allocate_host_staging_buffer(context.device, binding.length, staging);
+        Status            allocation_status = allocate_host_staging_buffer(context.device, binding.length, staging,
+                                                         context.pool_budget);
         if (!allocation_status.success()) {
             status.log("allocate host staging for value %d failed", binding.value.value);
             status.append(allocation_status);

ggml/src/ggml-hrx/runtime/host-memory.cpp

@@ -2108,6 +2108,11 @@ HostStagingBuffer & HostStagingBuffer::operator=(HostStagingBuffer && other) noe
     length          = other.length;
     upload          = other.upload;
     download        = other.download;
+    // engine#1711: the charge travels with the buffer. Dropping it leaves the destination holding bytes it is not
+    // charged for (its clear() returns nothing) while the source keeps a tally for a buffer it no longer owns (its
+    // destructor releases bytes it does not hold) -- a ledger that disagrees with the memory in both objects at once.
+    charged_bytes_  = other.charged_bytes_;
+    other.charged_bytes_ = 0;
     other.buffer    = nullptr;
     other.host_data = nullptr;
     other.value     = -1;
@@ -2122,6 +2127,10 @@ void HostStagingBuffer::clear() {
         hrx_buffer_release(buffer);
         buffer = nullptr;
     }
+    if (charged_bytes_ > 0) {
+        ggml::hrx::onebit::pool_release(charged_bytes_);   // engine#1711
+        charged_bytes_ = 0;
+    }
     host_data = nullptr;
     source_host_buffer = HostBufferRef{};
     value     = -1;
@@ -2130,11 +2139,30 @@ void HostStagingBuffer::clear() {
     download  = false;
 }
 
-Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging) {
+Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging,
+                                    size_t pool_budget) {
     staging.clear();
+    // engine#1711: staging buffers are device-visible memory too. Charged before the allocation (a refusal must leave
+    // the previous state intact -- clear() above has already released the old charge, so a refusal here just means no
+    // staging buffer) and given back in clear(), which every path reaches: the destructor, the move-assignment, and the
+    // next allocate.
+    bool charge_taken = false;
+    if (size > 0 && pool_budget > 0) {
+        if (!ggml::hrx::onebit::pool_reserve(pool_budget, size, "a host staging buffer")) {
+            Status refused;
+            refused.log("refused: a %zu-byte host staging buffer does not fit the pool", size);
+            return refused;
+        }
+        charge_taken = true;
+    }
     Status status = allocate_device_buffer(device, size, staging.buffer);
     if (status.success()) {
         staging.length = size;
+        if (charge_taken) {
+            staging.charged_bytes_ = size;
+        }
+    } else if (charge_taken) {
+        ggml::hrx::onebit::pool_release(size);
     }
     return status;
 }

ggml/src/ggml-hrx/runtime/host-memory.h

@@ -249,6 +249,8 @@ class HostWeightExecutionScope {
 
 struct HostStagingBuffer {
     HostStagingBuffer() = default;
+    // engine#1711: bytes this staging buffer holds against the pool ledger. The object owns exactly one buffer, so the
+    // charge is taken on allocate and given back in clear() -- no delta arithmetic, unlike the two arenas.
     ~HostStagingBuffer();
 
     HostStagingBuffer(HostStagingBuffer && other) noexcept;
@@ -262,13 +264,15 @@ struct HostStagingBuffer {
     HostBufferRef source_host_buffer;
     int32_t      value     = -1;
     size_t       length    = 0;
+    size_t       charged_bytes_ = 0;   // engine#1711: pool bytes held for this buffer
     bool         upload    = false;
     bool         download  = false;
 
     void clear();
 };
 
-Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging);
+Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging,
+                                    size_t pool_budget = 0);
 Status allocate_mapped_host_staging_buffer(hrx_device_t device,
                                            size_t       size,
                                            hrx_buffer_t & buffer,

ggml/src/ggml-hrx/runtime/transient-arena-charge.h

@@ -43,9 +43,15 @@ inline ArenaChargePlan plan_arena_charge(std::size_t required, std::size_t capac
     if (required <= capacity) {
         return plan;  // no growth, nothing to charge and nothing to give back
     }
-    plan.reserve_bytes = required - capacity;
+    // REPLACE, not transfer: charge the whole new allocation and give back the superseded one. The delta form --
+    // `reserve(required - capacity)` together with `release(capacity)` -- subtracts the old capacity TWICE, so the
+    // ledger settles at `required - capacity` instead of `required`, under-counting by exactly the arena it is meant to
+    // account for; the guardrail then over-admits. The invariant every caller relies on is "after a growth the ledger
+    // equals the arena's live capacity", which is what clear() releases.
+    plan.reserve_bytes = required;
     // A first allocation has nothing to release; a growth releases the superseded buffer's charge only when that
-    // buffer is actually freed. Reserved-and-retired bytes stay charged (see this file's header).
+    // buffer is actually freed. A RETIRED buffer stays resident, so its charge stays too: under the replace model the
+    // ledger then ends at old + new, which is exactly the memory the arena holds in that mode.
     plan.release_bytes = (!retires_old && capacity > 0) ? capacity : 0;
     return plan;
 }

tests/CMakeLists.txt

@@ -381,6 +381,12 @@ if (TARGET ggml-hrx)
     add_executable(test-hrx-transient-arena-charge test-hrx-transient-arena-charge.cpp)
     target_include_directories(test-hrx-transient-arena-charge PRIVATE ../ggml/src/ggml-hrx)
     add_test(NAME test-hrx-transient-arena-charge COMMAND test-hrx-transient-arena-charge)
+    # the host staging buffer's charge: the ledger is the only thing under test, no device and no libhrx (engine #1711)
+    add_executable(test-hrx-staging-ledger test-hrx-staging-ledger.cpp
+                   ../ggml/src/ggml-hrx/hrx-1bit-safety.cpp)
+    target_include_directories(test-hrx-staging-ledger PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx)
+    target_link_libraries(test-hrx-staging-ledger PRIVATE ggml)
+    add_test(NAME test-hrx-staging-ledger COMMAND test-hrx-staging-ledger)
 
     # MemAvailable parsing for the unbounded-residency warning: pure, no device and no /proc (engine #669)
     add_executable(test-hrx-host-memory-available test-hrx-host-memory-available.cpp

tests/test-hrx-staging-ledger.cpp

@@ -0,0 +1,103 @@
+// Copyright 2026 bong-water-water-bong
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// The host staging buffer's pool charge (engine#1711). Unlike the two arenas this one carries no delta -- the object
+// owns exactly one buffer -- so the whole contract is that clear() gives the bytes back EXACTLY once. clear() is reached
+// by the destructor, by move-assignment and by the next allocate, so a tally that survived a release would hand the same
+// bytes back two or three times and inflate the budget for everything afterwards.
+//
+// No device: the ledger is the only thing under test, which is why this does not need libhrx.
+
+#include "hrx-1bit-safety.h"
+#include "runtime/host-memory.h"
+
+#include <cstdio>
+#include <utility>
+
+using ggml::hrx::HostStagingBuffer;
+
+static int failures = 0;
+
+static void check(bool ok, const char * what, size_t got, size_t want) {
+    if (!ok) {
+        std::printf("FAIL %s (ledger %zu, want %zu)\n", what, got, want);
+        ++failures;
+    } else {
+        std::printf("ok   %s\n", what);
+    }
+}
+
+int main() {
+    using ggml::hrx::onebit::pool_reserve;
+    using ggml::hrx::onebit::pool_allocated;   // the ledger (NOT pool_reserve_bytes, which is the 16 GiB headroom reserve)
+    using ggml::hrx::onebit::pool_release;
+
+    const size_t base = pool_allocated();
+
+    // 1. clear() on a buffer that holds nothing must not touch the ledger: this is the state of every default-constructed
+    //    staging buffer and of every one already cleared, so an unconditional release here would underflow everything.
+    {
+        HostStagingBuffer empty;
+        empty.clear();
+        check(pool_allocated() == base, "clear() on an empty staging buffer leaves the ledger alone", pool_allocated(), base);
+        empty.clear();
+        check(pool_allocated() == base, "and calling it twice does not either", pool_allocated(), base);
+    }
+
+    // 2. A held charge is returned exactly once. The buffer is simulated: reserve the bytes, record them on the object,
+    //    clear it, and the ledger must be back where it started.
+    {
+        HostStagingBuffer held;
+        const size_t     bytes = 8u << 20;
+        // budget 1 TiB: pool_reserve() returns true WITHOUT touching the ledger when budget == 0, so a real budget is
+        // needed for the bytes to actually be charged
+        const bool reserved = pool_reserve(1ull << 40, bytes, "a host staging buffer");
+        check(reserved, "the pool accepts an 8 MiB staging request", 0, 0);
+        held.charged_bytes_ = bytes;
+        check(pool_allocated() == base + bytes, "the ledger carries the staging bytes", pool_allocated(), base + bytes);
+        held.clear();
+        check(pool_allocated() == base, "clear() returns exactly what was charged", pool_allocated(), base);
+        held.clear();
+        check(pool_allocated() == base, "a second clear() does not release it again", pool_allocated(), base);
+        (void) pool_release;  // the release under test happens inside clear()
+    }
+
+    // 3. The MOVE path carries the tally (the defect review of #270 found): the destination ends up holding the
+    //    buffer, so it must hold the charge; the source ends up owning nothing, so it must hold no tally. Dropping
+    //    either half leaves a ledger that disagrees with the memory in both objects at once.
+    {
+        HostStagingBuffer from, to;
+        const size_t     bytes = 4u << 20;
+        const bool       reserved = pool_reserve(1ull << 40, bytes, "a host staging buffer");
+        check(reserved, "the pool accepts a 4 MiB staging request for the move arm", 0, 0);
+        from.charged_bytes_ = bytes;
+        const size_t before = pool_allocated();
+        to = std::move(from);
+        check(pool_allocated() == before, "a move does not touch the ledger", pool_allocated(), before);
+        check(to.charged_bytes_ == bytes, "the destination carries the charge", to.charged_bytes_, bytes);
+        check(from.charged_bytes_ == 0, "the source keeps no tally for a buffer it no longer owns", from.charged_bytes_, 0);
+        to.clear();
+        check(pool_allocated() == base, "clearing the destination returns the bytes exactly once", pool_allocated(), base);
+        from.clear();                                     // must be a no-op: it holds nothing and is charged nothing
+        check(pool_allocated() == base, "clearing the moved-from source does not release again", pool_allocated(), base);
+    }
+
+    if (failures != 0) {
+        std::printf("%d arm(s) failed\n", failures);
+        return 1;
+    }
+    std::printf("all arms passed\n");
+    return 0;
+}

tests/test-hrx-transient-arena-charge.cpp

@@ -27,6 +27,15 @@ using ggml::hrx::plan_arena_charge;
 
 static int failures = 0;
 
+static void check_sizes(bool ok, const char * what, std::size_t got, std::size_t want) {
+    if (!ok) {
+        std::printf("FAIL %s (%zu, want %zu)\n", what, got, want);
+        ++failures;
+    } else {
+        std::printf("ok   %s\n", what);
+    }
+}
+
 static void check(bool ok, const char * what, ArenaChargePlan got, ArenaChargePlan want) {
     if (!ok) {
         std::printf("FAIL %s: reserve %zu want %zu, release %zu want %zu\n", what, got.reserve_bytes, want.reserve_bytes,
@@ -46,14 +55,14 @@ int main() {
     // 2. Growth, buffer freed: charge the delta, give back the superseded allocation's charge. The net charge after this
     //    is the new size - which is what the ledger must show, not the sum of every size the arena ever had.
     got = plan_arena_charge(16u << 20, 8u << 20, false);
-    check(got.reserve_bytes == (8u << 20) && got.release_bytes == (8u << 20), "growth charges the delta and releases the old", got,
-          {8u << 20, 8u << 20});
+    check(got.reserve_bytes == (16u << 20) && got.release_bytes == (8u << 20),
+          "growth charges the WHOLE new arena and releases the old", got, {16u << 20, 8u << 20});
 
     // 3. Growth in diagnostic mode: the old buffer is RETIRED and stays resident, so its charge stays too. This is the arm
     //    a naive "release the old capacity" version fails: it would release memory that is still on the device.
     got = plan_arena_charge(16u << 20, 8u << 20, true);
-    check(got.reserve_bytes == (8u << 20) && got.release_bytes == 0, "diagnostic growth keeps the retired buffer charged", got,
-          {8u << 20, 0});
+    check(got.reserve_bytes == (16u << 20) && got.release_bytes == 0,
+          "diagnostic growth keeps the retired buffer charged (ledger ends at old + new)", got, {16u << 20, 0});
 
     // 4. No growth: nothing charged, nothing released - the common case, and the one that must not perturb the ledger.
     got = plan_arena_charge(8u << 20, 8u << 20, false);
@@ -74,8 +83,39 @@ int main() {
           {0, 0});
     // and a grow past the budget releases the superseded arena's charge in the same step (the old buffer is freed there)
     got = plan_arena_charge(8192u << 20, 4096u << 20, false);
-    check(got.reserve_bytes == (4096u << 20) && got.release_bytes == (4096u << 20), "host-weight arena: growing past the budget swaps charge for charge", got,
-          {4096u << 20, 4096u << 20});
+    check(got.reserve_bytes == (8192u << 20) && got.release_bytes == (4096u << 20),
+          "host-weight arena: growing past the budget charges the new arena and releases the old", got,
+          {8192u << 20, 4096u << 20});
+
+    // 9. THE INVARIANT the #270 review named, walked over a sequence rather than one step: after every growth the
+    //    ledger must equal the arena's live capacity. A delta-shaped plan with a paired release passes a single-step
+    //    check and fails this one on the second growth, settling at new - old.
+    {
+        size_t ledger = 0, capacity = 0;
+        bool   ok     = true;
+        for (size_t step = 0; step < 5; ++step) {
+            const size_t          required = (1u << 20) << step;   // 1, 2, 4, 8, 16 MiB
+            const ArenaChargePlan plan     = plan_arena_charge(required, capacity, false);
+            ledger += plan.reserve_bytes - plan.release_bytes;
+            capacity = required;
+            if (ledger != capacity) {
+                ok = false;
+            }
+        }
+        check_sizes(ok, "ledger equals the live capacity after every growth (not new - old)", ledger, capacity);
+        size_t retired_ledger = 0, retired_held = 0;
+        ok = true;
+        for (size_t step = 0; step < 4; ++step) {
+            const size_t          required = (1u << 20) << step;
+            const ArenaChargePlan plan     = plan_arena_charge(required, retired_held, true);
+            retired_ledger += plan.reserve_bytes - plan.release_bytes;
+            retired_held += required;                              // every superseded arena stays resident
+            if (retired_ledger != retired_held) {
+                ok = false;
+            }
+        }
+        check_sizes(ok, "retired growth accumulates old + new (they are all still on the device)", retired_ledger, retired_held);
+    }
 
     if (failures != 0) {
         std::printf("%d arm(s) failed\n", failures);