ggml-hrx: a shared pool_tally pair, so the #1823 charge has device-free arms (carries 070a270ca; stacked on #277)

ggml/src/ggml-hrx/runtime/failure-status.cpp

@@ -10,12 +10,14 @@ namespace {
 // iree's status name for a pool that cannot serve the request, and HSA's for the same condition one layer down. Both appear
 // verbatim in the failures this exists for:
 //   al...: RESOURCE_EXHAUSTED; [hsa_amd_memory_pool_allocate] HSA_STATUS_ERROR_OUT_OF_RESOURCES
-constexpr const char * kResourceNeedles[] = { "RESOURCE_EXHAUSTED", "OUT_OF_RESOURCES" };
+// kPoolRefusalToken is the same condition one layer UP: the pool ledger refused before the allocator was asked.
+constexpr const char * kResourceNeedles[] = { "RESOURCE_EXHAUSTED", "OUT_OF_RESOURCES", kPoolRefusalToken };
 
 // Case-insensitive on purpose: the *classification* needles are exact because they are identifiers as iree and HSA render
 // them, but the *hint* exists to catch spellings we have not seen -- `hipErrorOutOfMemory` is real hip and mixed case, and a
 // case-sensitive match would have missed exactly the case this probe was added for (the test caught that).
-constexpr const char * kAllocationHints[] = { "memory", "alloc" };
+// "pool" is here for our own refusals: one written without kPoolRefusalToken is then at least named in the log.
+constexpr const char * kAllocationHints[] = { "memory", "alloc", "pool" };
 
 bool contains_any_of_ci(const std::string & message, const char * const * needles, size_t count) {
     std::string lowered;

ggml/src/ggml-hrx/runtime/failure-status.h

@@ -20,7 +20,16 @@ namespace ggml::hrx {
 // survive rewording of the prose around them.
 bool failure_looks_like_allocation(const std::vector<std::string> & messages);
 
-// A message that *looks* like an allocation failure (it mentions memory or allocation) but matched neither needle, so it is
+// The third identifier, and the only one this backend writes itself: the token a compute-time refusal by the pool LEDGER
+// carries in its message. The two names above come from the allocator, i.e. from a launch the guardrail admitted and the
+// device then could not serve. Since engine #1711/#1823 the ledger refuses first (the transient arena, a host staging buffer,
+// a materialized host weight), and those refusals are this backend's own prose: "refused: a 121.7 MiB host weight does not
+// fit the pool" names neither RESOURCE_EXHAUSTED nor OUT_OF_RESOURCES, so the better the accounting got, the more often an
+// out-of-memory condition went back to being "Compute error." for the client. A refusal site appends this token instead of
+// relying on its wording, so rewording a sentence cannot change the status.
+inline constexpr char kPoolRefusalToken[] = "POOL_EXHAUSTED";
+
+// A message that *looks* like an allocation failure (it mentions memory, allocation or the pool) but matched no needle, so it is
 // about to be classified as a compute failure (engine #1618 review). Returns nullptr when nothing looks suspicious. This exists
 // so the next spelling of an out-of-memory status -- `hipErrorOutOfMemory`, `HSA_STATUS_ERROR_INVALID_ALLOCATION`, whatever the
 // runtime renders tomorrow -- costs a log line rather than another day of misdiagnosis: the fallback stays the safe direction

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

@@ -5,6 +5,7 @@
 #include "hrx-interop-utils.h"
 #include "hrx_runtime.h"
 #include "hrx-1bit-safety.h"            // engine#1711: pool_reserve/pool_release
+#include "runtime/failure-status.h"     // engine#1618: kPoolRefusalToken, so a pool refusal is not a "Compute error."
 #include "transient-arena-charge.h"     // engine#1711: the same charge decision as the transient arena
 #include "runtime/host-memory-available.h"
 #include "runtime/host-memory-ternary.h"
@@ -1388,10 +1389,7 @@ struct HostWeightLease::Entry {
         // allocated uncharged and freed with nothing to release, so the DEFAULT mode (`GGML_HRX_HOST_WEIGHT_LIMIT_MB=0`
         // = per-weight residency) pierced the guardrail by however much the materialised weights held. Released here
         // because Entry owns exactly one buffer.
-        if (charged_bytes_ > 0) {
-            ggml::hrx::onebit::pool_release(charged_bytes_);
-            charged_bytes_ = 0;
-        }
+        pool_tally::release(charged_bytes_);
     }
     bool owns_buffer = false;
     size_t charged_bytes_ = 0;   // engine#1823
@@ -2019,20 +2017,14 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t             device
         // engine#1823: this is the DEFAULT mode (no GGML_HRX_HOST_WEIGHT_LIMIT_MB), and until now it was the one
         // path that allocated device memory without asking the pool. Charge the weight before the allocation, so a
         // pool that cannot hold it is a refusal naming the weight rather than an over-commit that faults later.
-        if (pool_budget > 0 && materialized_length > 0) {
-            if (!ggml::hrx::onebit::pool_reserve(pool_budget, materialized_length, "a materialized host weight")) {
-                result.status.log("refused: a %.1f MiB host weight does not fit the pool",
-                                  materialized_length / (1024.0 * 1024.0));
-                return result;
-            }
-            entry->charged_bytes_ = materialized_length;
+        if (!pool_tally::charge(pool_budget, materialized_length, entry->charged_bytes_, "a materialized host weight")) {
+            result.status.log("refused: a %.1f MiB host weight does not fit the pool (%s)",
+                              materialized_length / (1024.0 * 1024.0), kPoolRefusalToken);
+            return result;
         }
         result.status = allocate_device_buffer(device, materialized_length, entry->buffer);
         if (!result.status.success()) {
-            if (entry->charged_bytes_ > 0) {
-                ggml::hrx::onebit::pool_release(entry->charged_bytes_);
-                entry->charged_bytes_ = 0;
-            }
+            pool_tally::release(entry->charged_bytes_);
             return result;
         }
         entry->owns_buffer = true;
@@ -2174,7 +2166,7 @@ Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagin
     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);
+            refused.log("refused: a %zu-byte host staging buffer does not fit the pool (%s)", size, kPoolRefusalToken);
             return refused;
         }
         charge_taken = true;

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

@@ -3,6 +3,7 @@
 #include "dispatch/dispatch.h"
 #include "ggml.h"
 #include "runtime/host-buffer-registry.h"
+#include "hrx-1bit-safety.h"   // pool_tally (engine #1711/#1823)
 #include "runtime/host-memory-refusal.h"
 #include "status.h"
 
@@ -77,6 +78,36 @@ struct HostWeightCacheStats {
     size_t   resident_bytes   = 0;
 };
 
+// The pairing every pooled allocation in this file needs (engine #1711/#1823): charge before the allocation, return
+// exactly what was charged when the object dies. Kept as a free pair rather than repeated inline so the contract can be
+// tested with no device -- a tally that is not zeroed on release gives the bytes back twice, and that is the failure the
+// #270 review found in a move-assignment before this existed.
+namespace pool_tally {
+
+// Returns false when the pool refuses; leaves both the ledger and `charged` untouched in that case.
+inline bool charge(std::size_t budget, std::size_t bytes, std::size_t & charged, const char * what) {
+    if (budget == 0 || bytes == 0) {
+        return true;   // uncharged mode (device size unknown): nothing to reserve and nothing to release
+    }
+    if (!onebit::pool_reserve(budget, bytes, what)) {
+        return false;
+    }
+    charged = bytes;
+    return true;
+}
+
+// Releases what was charged and zeroes the tally, so a second call (a destructor reached twice by a move, say) is a no-op
+// rather than an underflow of somebody else's bytes.
+inline void release(std::size_t & charged) {
+    if (charged == 0) {
+        return;
+    }
+    onebit::pool_release(charged);
+    charged = 0;
+}
+
+}  // namespace pool_tally
+
 class HostWeightLease {
   public:
     HostWeightLease() = default;

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

@@ -2,6 +2,7 @@
 
 #include "hrx-1bit-safety.h"   // engine#1711: the pool ledger (pool_reserve/pool_release)
 #include "transient-arena-charge.h"
+#include "runtime/failure-status.h"   // engine#1618: kPoolRefusalToken
 
 #include "hrx-interop-utils.h"
 #include "hrx_runtime.h"
@@ -96,7 +97,8 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t
     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);
+            status.log("refused: the transient arena needs %zu more bytes and the pool has no room (%s)", charge.reserve_bytes,
+                       kPoolRefusalToken);
             return status;
         }
         charge_taken   = true;

tests/test-hrx-failure-status.cpp

@@ -51,6 +51,28 @@ int main() {
     REQUIRE(failure_looks_like_allocation({ "materialize host weight value 3 failed: HSA_STATUS_ERROR_OUT_OF_RESOURCES" }));
     REQUIRE(failure_status({ "x: RESOURCE_EXHAUSTED" }) == GGML_STATUS_ALLOC_FAILED);
 
+    // The pool LEDGER's own refusals (engine #1711/#1823) are the same condition one layer up, and they are this backend's
+    // prose rather than an iree or HSA name. Before kPoolRefusalToken all three were GGML_STATUS_FAILED, i.e. "Compute
+    // error." for the client -- #1618's defect again, on the path that replaced the allocator failure. The sentences are
+    // the three refusal sites' (host-memory.cpp twice, transient-arena.cpp), with the token each now appends.
+    using ggml::hrx::kPoolRefusalToken;
+    const std::string token = kPoolRefusalToken;
+    for (const std::string & refusal : {
+             "refused: a 121.7 MiB host weight does not fit the pool (" + token + ")",
+             "refused: a 8388608-byte host staging buffer does not fit the pool (" + token + ")",
+             "refused: the transient arena needs 67108864 more bytes and the pool has no room (" + token + ")",
+         }) {
+        REQUIRE(failure_looks_like_allocation({ refusal }));
+        REQUIRE(failure_status({ refusal }) == GGML_STATUS_ALLOC_FAILED);
+        REQUIRE(ggml::hrx::failure_unmatched_allocation({ refusal }) == nullptr);
+    }
+    // The token decides, not the sentence: the old wording without it stays a compute failure (so a refusal site that
+    // forgets the token is the bug, and rewording one cannot change a status) -- but it is NAMED by the warning, because
+    // "pool" is an allocation hint.
+    const std::vector<std::string> untagged = { "refused: a 121.7 MiB host weight does not fit the pool" };
+    REQUIRE(failure_status(untagged) == GGML_STATUS_FAILED);
+    REQUIRE(ggml::hrx::failure_unmatched_allocation(untagged) != nullptr);
+
     // A real compute failure must stay a compute failure: this is the direction that would be a regression, because a
     // kernel error reported as "out of memory" sends the reader to free memory that was never the problem.
     const std::vector<std::string> kernel_fault = {

tests/test-hrx-staging-ledger.cpp

@@ -94,6 +94,46 @@ int main() {
         check(pool_allocated() == base, "clearing the moved-from source does not release again", pool_allocated(), base);
     }
 
+    // 4. The pairing helper the host-weight entry calls (engine #1823): charge refuses without touching anything, release
+    //    returns exactly once, and the uncharged mode (budget 0) reserves nothing. Same contract as the staging object
+    //    above, now on the shared function, which is what makes it testable without a device.
+    {
+        using ggml::hrx::pool_tally::charge;
+        using ggml::hrx::pool_tally::release;
+        // A bystander holds 1 MiB for the whole block. Without it the ledger is at zero when the second release runs,
+        // pool_release clamps there, and "a second release is a no-op" passes for a release() that never zeroes its tally.
+        size_t bystander = 0;
+        check(charge(1ull << 40, 1u << 20, bystander, "a bystander"), "a bystander holds 1 MiB, so an over-release has bytes to take", 0, 0);
+        size_t        tally = 0;
+        const size_t  before = pool_allocated();
+        const size_t  bytes  = 6u << 20;
+
+        check(charge(1ull << 40, bytes, tally, "test weight"), "charge takes the bytes from the pool", 0, 0);
+        check(tally == bytes && pool_allocated() == before + bytes, "and records the tally and the ledger", pool_allocated(), before + bytes);
+        release(tally);
+        check(tally == 0 && pool_allocated() == before, "release returns them and zeroes the tally", pool_allocated(), before);
+        release(tally);
+        check(pool_allocated() == before, "a second release is a no-op, not an underflow of somebody else's bytes", pool_allocated(), before);
+
+        // a refusal leaves both untouched: this is the arm that matters on the default path, where the weight used to be
+        // materialised uncharged. Note `budget == size` with nothing held EXACTLY FITS and is accepted -- a first version
+        // of this arm asserted otherwise and was wrong -- so the pressure has to come from bytes already held.
+        size_t       holder = 0;
+        const size_t held   = 4u << 20;
+        check(charge(64u << 20, held, holder, "a holder"), "a holder takes 4 MiB", 0, 0);
+        tally = 0;
+        const bool refused = charge(5u << 20, bytes, tally, "test weight");   // 4 held + 6 wanted > 5 budget
+        check(!refused && tally == 0, "a refused charge leaves the tally alone", tally, 0);
+        check(pool_allocated() == before + held, "and does not touch the ledger either", pool_allocated(), before + held);
+
+        // the uncharged mode: budget 0 means "device size unknown", so nothing is reserved and nothing must be returned
+        check(charge(0, bytes, tally, "test weight") && tally == 0, "budget 0 reserves nothing and reports success", tally, 0);
+        release(holder);
+        check(pool_allocated() == before, "and the holder's release returns everything", pool_allocated(), before);
+        release(bystander);
+        check(pool_allocated() == base, "and the bystander's release leaves the ledger where the test found it", pool_allocated(), base);
+    }
+
     if (failures != 0) {
         std::printf("%d arm(s) failed\n", failures);
         return 1;