ggml-hrx: the transient arena is charged to the pool before it grows (engine #1711)

ggml/src/ggml-hrx/ggml-hrx.cpp

@@ -380,7 +380,11 @@ static const ggml_backend_buffer_i buffer_i = {
 // engine #1502 (#251) calls this from buffer_alloc, eight hundred lines above its definition, with no declaration in
 // scope - so the translation unit does not compile ("'device_free_total' was not declared in this scope"). One forward
 // declaration is enough; the definition and its callers are unchanged.
-static void device_free_total(ggml_backend_hrx_device_context * device_ctx, size_t & free_bytes, size_t & total_bytes);
+// The declaration must match the definition's type: it is `bool` (it reports whether the driver counters were
+// available). Two merges each added a forward declaration - #252's above, #240's `bool` one further down - and the
+// `void`/`bool` pair is an ambiguous redeclaration, so the translation unit fails to compile at the `bool` declaration
+// and the pin builds for nobody. One word, and every lane that builds this pin is unblocked.
+static bool device_free_total(ggml_backend_hrx_device_context * device_ctx, size_t & free_bytes, size_t & total_bytes);
 
 static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_t size) {
     ggml::hrx::onebit::poison_selftest();  // no-op unless GGML_HRX_POISON_SELFTEST=1

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

@@ -387,7 +387,8 @@ static Status ensure_transient_arena(const CommandProgramExecutionContext & cont
     if (commands.transients.arena_size == 0) {
         return status;
     }
-    status = context.transient_arena->ensure_capacity(context.device, context.stream, commands.transients.arena_size);
+    status = context.transient_arena->ensure_capacity(context.device, context.stream, commands.transients.arena_size,
+                                                            context.pool_budget);
     if (!status.success()) {
         return status;
     }
@@ -1912,7 +1913,8 @@ bool bind_and_execute_prepared_command_program(const CommandProgramExecutionCont
     TransientArena::AllocationLease lease = context.transient_arena->acquire_allocation_lease();
     {
         const HostGapScope scope("prep.transient");
-        status = lease.ensure_capacity(context.device, context.stream, commands.transients.arena_size);
+        status = lease.ensure_capacity(context.device, context.stream, commands.transients.arena_size,
+                                                            context.pool_budget);
     }
     if (!status.success()) {
         GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status));
@@ -2011,7 +2013,8 @@ RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph(
         {
             const HostGapScope scope("replay.transient");
             lease  = context.transient_arena->acquire_allocation_lease();
-            status = lease.ensure_capacity(context.device, context.stream, commands.transients.arena_size);
+            status = lease.ensure_capacity(context.device, context.stream, commands.transients.arena_size,
+                                                            context.pool_budget);
             if (status.success()) {
                 transient_allocation = lease.current_allocation();
                 if (!bind_prepared_command_program_transients(commands, transient_allocation, prepared)) {

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

@@ -65,6 +65,10 @@ struct CommandProgramExecutionContext {
     HostWeightCache *       host_weights       = nullptr;
     HostBufferRegistry *    host_buffers       = nullptr;
     GraphReplayStreamState * graph_replay_state = nullptr;
+    // engine#1711: the device's pool budget, so the transient arena can be charged before it grows. 0 means the device
+    // size is unknown, in which case the arena keeps the old uncharged behaviour (and pool_reserve() would not have
+    // added to the ledger either, so a paired release would unbalance it).
+    size_t                  pool_budget        = 0;
 };
 
 struct PreparedCommandBinding {

ggml/src/ggml-hrx/runtime/graph-executor.cpp

@@ -147,6 +147,7 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const {
         &context_.host_weights,
         &context_.device->host_buffers,
         &context_.graph_replay_state,
+        context_.device->memory_budget,
     };
     PreparedCommandProgramCacheExecutionResult execution;
     {

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

@@ -0,0 +1,53 @@
+// 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.
+//
+// What the transient arena owes the pool ledger when it grows (engine #1711).
+//
+// The guardrail used to be consulted only by the buffer-type allocator, so the arena - the scratch a launch grows into -
+// was allocated uncharged, and "a launch that needs more is refused here" (#680) held for tensors only. Charging it is
+// two numbers, and the second one is where a naive version goes wrong: when the superseded allocation is RELEASED its
+// charge goes with it, but in diagnostic mode it is RETIRED (diagnostic_retired_buffers_) and stays resident, so its
+// charge has to stay too. Releasing it there under-reports by exactly the buffers kept around to diagnose memory
+// problems, i.e. the ledger stays comfortable while the device is full.
+//
+// Pure and header-only so the decision can be tested with no device and no libhrx (the same reason
+// host-memory-refusal.h lives apart from HostWeightCache).
+
+#pragma once
+
+#include <cstddef>
+
+namespace ggml::hrx {
+
+struct ArenaChargePlan {
+    std::size_t reserve_bytes = 0;  // to take from the pool BEFORE the old allocation is touched
+    std::size_t release_bytes = 0;  // to return once the old allocation is gone (0 when it is retired, not freed)
+};
+
+// required: the arena size the launch needs; capacity: the arena's current allocation (0 for the first one);
+// retires_old: the arena's diagnostic mode, which keeps superseded buffers alive.
+inline ArenaChargePlan plan_arena_charge(std::size_t required, std::size_t capacity, bool retires_old) {
+    ArenaChargePlan plan;
+    if (required <= capacity) {
+        return plan;  // no growth, nothing to charge and nothing to give back
+    }
+    plan.reserve_bytes = required - capacity;
+    // 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).
+    plan.release_bytes = (!retires_old && capacity > 0) ? capacity : 0;
+    return plan;
+}
+
+}  // namespace ggml::hrx

ggml/src/ggml-hrx/runtime/transient-arena.cpp

@@ -1,5 +1,8 @@
 #include "transient-arena.h"
 
+#include "hrx-1bit-safety.h"   // engine#1711: the pool ledger (pool_reserve/pool_release)
+#include "transient-arena-charge.h"
+
 #include "hrx-interop-utils.h"
 #include "hrx_runtime.h"
 
@@ -14,13 +17,14 @@ TransientArena::AllocationLease::AllocationLease(TransientArena & arena, std::un
 
 Status TransientArena::AllocationLease::ensure_capacity(hrx_device_t device,
                                                         hrx_stream_t stream,
-                                                        size_t       required_size) {
+                                                        size_t       required_size,
+                                                        size_t       pool_budget) {
     Status status;
     if (arena_ == nullptr || !lock_.owns_lock()) {
         status.log("missing transient arena allocation lease");
         return status;
     }
-    return arena_->ensure_capacity_locked(device, stream, required_size);
+    return arena_->ensure_capacity_locked(device, stream, required_size, pool_budget);
 }
 
 TransientArenaAllocationRef TransientArena::AllocationLease::current_allocation() const {
@@ -44,6 +48,10 @@ void TransientArena::clear() {
         hrx_buffer_release(buffer);
     }
     diagnostic_retired_buffers_.clear();
+    if (charged_bytes_ > 0) {
+        ggml::hrx::onebit::pool_release(charged_bytes_);   // engine#1711: retired buffers were charged too
+        charged_bytes_ = 0;
+    }
     allocation_capacity_ = 0;
     allocation_id_       = kInvalidTransientArenaAllocationId;
 }
@@ -56,16 +64,18 @@ uint64_t TransientArena::next_allocation_id() {
     return id;
 }
 
-Status TransientArena::ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size) {
+Status TransientArena::ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size,
+                                      size_t pool_budget) {
     std::lock_guard<std::mutex> lock(mutex_);
-    return ensure_capacity_locked(device, stream, required_size);
+    return ensure_capacity_locked(device, stream, required_size, pool_budget);
 }
 
 TransientArena::AllocationLease TransientArena::acquire_allocation_lease() {
     return AllocationLease(*this, std::unique_lock<std::mutex>(mutex_));
 }
 
-Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size) {
+Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size,
+                                               size_t pool_budget) {
     Status status;
     const bool diagnostic_fresh = std::getenv("GGML_HRX_DIAGNOSTIC_FRESH_TRANSIENT_ARENA") != nullptr;
     if (required_size == 0 || (!diagnostic_fresh && allocation_capacity_ >= required_size)) {
@@ -79,6 +89,19 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t
         status.log("missing HRX stream for transient arena allocation");
         return status;
     }
+    // engine#1711: the guardrail used to be consulted only by the buffer-type allocator, so the arena grew uncharged and a
+    // launch that ran out of pool memory faulted instead of being refused (#680's promise covered tensors only). The charge
+    // is taken BEFORE the old buffer is released: a refusal must leave the arena with the capacity it already had.
+    const ArenaChargePlan charge = plan_arena_charge(required_size, allocation_capacity_, diagnostic_fresh);
+    bool                  charge_taken = false;
+    if (charge.reserve_bytes > 0 && pool_budget > 0) {
+        if (!ggml::hrx::onebit::pool_reserve(pool_budget, charge.reserve_bytes, "the transient arena")) {
+            status.log("refused: the transient arena needs %zu more bytes and the pool has no room", charge.reserve_bytes);
+            return status;
+        }
+        charge_taken   = true;
+        charged_bytes_ += charge.reserve_bytes;
+    }
     if (buffer_ != nullptr) {
         if (diagnostic_fresh) {
             diagnostic_retired_buffers_.push_back(buffer_);
@@ -88,6 +111,12 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t
             return status;
         }
         hrx_buffer_release(buffer_);
+            // the superseded allocation is gone from the device, so its charge goes with it; in diagnostic mode the
+            // buffer was retired above and stays resident, so its charge stays too.
+            if (charge.release_bytes > 0) {
+                ggml::hrx::onebit::pool_release(charge.release_bytes);
+                charged_bytes_ -= charged_bytes_ >= charge.release_bytes ? charge.release_bytes : charged_bytes_;
+            }
         }
         buffer_              = nullptr;
         allocation_capacity_ = 0;
@@ -104,6 +133,10 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t
     if (ErrorResult error = take_status(
             hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, required_size, &allocation))) {
         status.log("allocate transient arena: %s", error->c_str());
+        if (charge_taken) {
+            ggml::hrx::onebit::pool_release(charge.reserve_bytes);
+            charged_bytes_ -= charged_bytes_ >= charge.reserve_bytes ? charge.reserve_bytes : charged_bytes_;
+        }
         return status;
     }
 

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

@@ -24,7 +24,8 @@ class TransientArena {
         AllocationLease(const AllocationLease &)             = delete;
         AllocationLease & operator=(const AllocationLease &) = delete;
 
-        Status                      ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size);
+        Status                      ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size,
+                                                      size_t pool_budget = 0);
         TransientArenaAllocationRef current_allocation() const;
 
       private:
@@ -42,7 +43,8 @@ class TransientArena {
     TransientArena(const TransientArena &)             = delete;
     TransientArena & operator=(const TransientArena &) = delete;
 
-    Status          ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size);
+    Status          ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size,
+                                    size_t pool_budget = 0);
     AllocationLease acquire_allocation_lease();
     void            clear();
 
@@ -54,12 +56,17 @@ class TransientArena {
 
   private:
     uint64_t                    next_allocation_id();
-    Status                      ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size);
+    Status                      ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size,
+                                                           size_t pool_budget);
     TransientArenaAllocationRef current_allocation_locked() const;
 
     mutable std::mutex mutex_;
     hrx_buffer_t       buffer_              = nullptr;
     std::vector<hrx_buffer_t> diagnostic_retired_buffers_;
+    // Bytes this arena currently holds against the pool ledger (engine#1711). It is NOT allocation_capacity_ in
+    // diagnostic mode, where superseded buffers are retired and stay resident: the tally covers retired + current, so
+    // releasing allocation_capacity_ there would under-report exactly the mode used to diagnose out-of-memory.
+    size_t                     charged_bytes_ = 0;
     size_t             allocation_capacity_ = 0;
     uint64_t           allocation_id_       = kInvalidTransientArenaAllocationId;
     uint64_t           next_allocation_id_  = 1;

tests/CMakeLists.txt

@@ -364,6 +364,10 @@ if (TARGET ggml-hrx)
     add_executable(test-hrx-arena-refusal test-hrx-arena-refusal.cpp)
     target_include_directories(test-hrx-arena-refusal PRIVATE ../ggml/src/ggml-hrx)
     add_test(NAME test-hrx-arena-refusal COMMAND test-hrx-arena-refusal)
+    # what the transient arena owes the pool when it grows: pure helper, no device and no HRX link (engine #1711)
+    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)
 
     # 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-transient-arena-charge.cpp

@@ -0,0 +1,72 @@
+// 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.
+//
+// What the transient arena owes the pool ledger when it grows (engine #1711). No device and no libhrx: the decision is a
+// pure function (runtime/transient-arena-charge.h) so every corner can be pinned here, including the one that makes a
+// naive charge under-report - diagnostic mode retires the superseded buffer instead of releasing it, so its charge has
+// to stay. The arms are the four states the growth path can be in.
+
+#include "runtime/transient-arena-charge.h"
+
+#include <cstdio>
+
+using ggml::hrx::ArenaChargePlan;
+using ggml::hrx::plan_arena_charge;
+
+static int failures = 0;
+
+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,
+                    got.release_bytes, want.release_bytes);
+        ++failures;
+    } else {
+        std::printf("ok   %s\n", what);
+    }
+}
+
+int main() {
+    // 1. First allocation: charge all of it, nothing to give back.
+    ArenaChargePlan got = plan_arena_charge(8u << 20, 0, false);
+    check(got.reserve_bytes == (8u << 20) && got.release_bytes == 0, "first allocation charges its whole size", got,
+          {8u << 20, 0});
+
+    // 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});
+
+    // 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});
+
+    // 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);
+    check(got.reserve_bytes == 0 && got.release_bytes == 0, "a request that fits charges nothing", got, {0, 0});
+
+    // 5. A smaller request than the current capacity: same as no growth (the arena never shrinks on its own).
+    got = plan_arena_charge(1u << 20, 8u << 20, true);
+    check(got.reserve_bytes == 0 && got.release_bytes == 0, "a smaller request charges nothing", got, {0, 0});
+
+    if (failures != 0) {
+        std::printf("%d arm(s) failed\n", failures);
+        return 1;
+    }
+    std::printf("all arms passed\n");
+    return 0;
+}