diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 705dcc54155a..b45dd2abe0d6 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/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 diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 3485ff955fcf..67f2332596c3 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/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)) { diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h index db5ae5b43051..163570c2d9ef 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/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 { diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index 6f0184fdd861..7b9b8c4dcf02 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/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; { diff --git a/ggml/src/ggml-hrx/runtime/transient-arena-charge.h b/ggml/src/ggml-hrx/runtime/transient-arena-charge.h new file mode 100644 index 000000000000..be87e6388371 --- /dev/null +++ b/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 + +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 diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.cpp b/ggml/src/ggml-hrx/runtime/transient-arena.cpp index a12f89e0ab38..b4a3099bb349 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena.cpp +++ b/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 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(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; } diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.h b/ggml/src/ggml-hrx/runtime/transient-arena.h index a50639962c7c..546de324caf2 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena.h +++ b/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 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; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c592e6b41034..ea05b5bd0fd2 100644 --- a/tests/CMakeLists.txt +++ b/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 diff --git a/tests/test-hrx-transient-arena-charge.cpp b/tests/test-hrx-transient-arena-charge.cpp new file mode 100644 index 000000000000..6e5d508ba5da --- /dev/null +++ b/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 + +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; +}