andraghetti · Nov 6, 2026 · Nov 5, 2026 · Nov 6, 2026
@@ -551,12 +551,13 @@ or direct GPU access via the mapped flag. At capture time:
 
 At playback:
 
-- `playback_hipHostRegister` allocates a fresh 64-byte-aligned host buffer (`_aligned_malloc`
-  on Windows, `posix_memalign` on Linux), restores the blob snapshot into it, calls
-  `hipHostRegister` on the new buffer, and records `(recorded_ptr → live_ptr)` in `alloc_map`
+- `playback_hipHostRegister` takes the recorded address's block in `host_reg_blocks` (the
+  same block again when it is large enough), restores the blob snapshot into it, calls
+  `hipHostRegister` on it, and records `(recorded_ptr → live_ptr)` in `alloc_map`
   (for kernel-arg pointer translation) and `host_reg_bufs` (for cleanup at Unregister).
-- `playback_hipHostUnregister` translates the pointer, calls `hipHostUnregister`, frees the
-  backing buffer, and removes both map entries.
+- `playback_hipHostUnregister` translates the pointer, calls `hipHostUnregister`, and removes
+  both map entries. The block stays, charged to the host budget, until the context is
+  destroyed, since a graph node may still point into it.
 
 ### Pinned Host Snapshots
 

@@ -3342,12 +3342,12 @@ hipError_t playback_hipMemPoolCreate(PlaybackContext& ctx, const uint8_t* pl) {
 // An archive can record a host allocation again at an address that still holds
 // one: a second pass replays the same events, and a crafted archive can repeat
 // the address without a free. The entry being replaced keeps a live pinned
-// buffer, so free it, and give its charge back, before the new one is charged.
-// Otherwise the old buffer would stay allocated with no charge against it.
+// buffer, so release it before the new one is charged: a hipHostMalloc buffer
+// is freed and its charge given back; a hipHostRegister block is unregistered
+// and kept in host_reg_blocks for the next registration at this address.
 // A buffer the runtime refuses to release (hipHostFree or hipHostUnregister
-// fails, as it does under a capture on another thread) is left alone, buffer
-// and charge both: the charge goes back only once the runtime has let go.
-// Freeing a registration would leave it aimed at freed heap.
+// fails, as it does under a capture on another thread) is left alone and
+// keeps its charge: the charge goes back only once the runtime has let go.
 static void free_replaced_host_alloc(PlaybackContext& ctx, uint64_t rec) {
     std::unique_lock host_free(ctx.host_free_mu);
     void* host_ptr = nullptr;
@@ -3374,13 +3374,11 @@ static void free_replaced_host_alloc(PlaybackContext& ctx, uint64_t rec) {
         return;
     }
     void* buf = reg_buf ? reg_buf : host_ptr;
-    if (hipHostUnregister(buf) != hipSuccess) return;
-    ctx.host_budget.release(size);
-#ifdef _WIN32
-    _aligned_free(buf);
-#else
-    free(buf);
-#endif
+    if (hipHostUnregister(buf) == hipSuccess) return;
+    // Still registered: the next registration here gets a new block.
+    std::unique_lock lk(ctx.map_mutex);
+    auto& block = ctx.host_reg_blocks[rec];
+    if (block.cur.get() == buf) block.retired.push_back(std::move(block.cur));
 }
 
 hipError_t playback_hipHostMalloc(PlaybackContext& ctx, const uint8_t* pl) {
@@ -3450,11 +3448,12 @@ hipError_t playback_hipFreeHost(PlaybackContext& ctx, const uint8_t* pl) {
 // Manual playback: hipHostRegister / hipHostUnregister
 // ---------------------------------------------------------------------------
 // hipHostRegister recorded a snapshot of the host memory as a blob.
-// At replay we allocate a fresh host buffer (malloc), restore the blob
-// into it, call hipHostRegister on it, and track the (recorded -> live)
+// At replay we take the recorded address's block in host_reg_blocks, restore
+// the blob into it, call hipHostRegister on it, and track the (recorded -> live)
 // mapping so kernel-arg pointer translations work. A registration recorded
 // without its blob is refused.
-// hipHostUnregister unregisters, frees the backing buffer, and removes the entry.
+// hipHostUnregister unregisters and removes the entry. The block stays until
+// the context is destroyed, since a graph node may still point into it.
 
 hipError_t playback_hipHostRegister(PlaybackContext& ctx, const uint8_t* pl) {
     const auto* a = reinterpret_cast<const hrr_args_hipHostRegister*>(pl);
@@ -3463,20 +3462,9 @@ hipError_t playback_hipHostRegister(PlaybackContext& ctx, const uint8_t* pl) {
     free_replaced_host_alloc(ctx, a->hostPtr);
     if (host_contents_left_out("hipHostRegister", sz, a->blob_hash_lo, a->blob_hash_hi))
         return hipErrorNotFound;
-    if (!ctx.charge_host(sz, "hipHostRegister backing buffer"))
-        return hipErrorMemoryAllocation;
-
-    // Allocate backing host buffer aligned to 64 bytes (page-register friendly).
-    void* buf = nullptr;
-#ifdef _WIN32
-    buf = _aligned_malloc(sz, 64);
-#else
-    if (posix_memalign(&buf, 64, sz) != 0) buf = nullptr;
-#endif
-    if (!buf) {
-        ctx.host_budget.release(sz);
-        return hipErrorMemoryAllocation;
-    }
+    void* buf = ctx.host_block(ctx.host_reg_blocks, a->hostPtr, sz,
+                               "hipHostRegister backing buffer");
+    if (!buf) return hipErrorMemoryAllocation;
 
     // Restore snapshot into the buffer.
     size_t blob_sz = 0;
@@ -3488,11 +3476,6 @@ hipError_t playback_hipHostRegister(PlaybackContext& ctx, const uint8_t* pl) {
     if (!blob && ctx.fatal_error.load(std::memory_order_acquire)) {
         // The blob was refused (cap or allocation), not merely absent:
         // do not register zeros in its place.
-#ifdef _WIN32
-        _aligned_free(buf);
-#else
-        free(buf);
-#endif
         return hipErrorUnknown;
     }
 
@@ -3501,13 +3484,6 @@ hipError_t playback_hipHostRegister(PlaybackContext& ctx, const uint8_t* pl) {
         ctx.record_alloc(a->hostPtr, buf, sz, AllocKind::HostRegister);
         std::unique_lock lk(ctx.map_mutex);
         ctx.host_reg_bufs[a->hostPtr] = buf;
-    } else {
-        ctx.host_budget.release(sz);
-#ifdef _WIN32
-        _aligned_free(buf);
-#else
-        free(buf);
-#endif
     }
     return r;
 }
@@ -3516,9 +3492,8 @@ hipError_t playback_hipHostUnregister(PlaybackContext& ctx, const uint8_t* pl) {
     std::unique_lock host_free(ctx.host_free_mu);
     const auto* a = reinterpret_cast<const hrr_args_hipHostUnregister*>(pl);
 
-    // Leave the maps and the backing buffer alone until the runtime call
-    // succeeds. A failed unregister did not release the registration, so
-    // freeing the buffer would leave alloc_map aimed at freed heap.
+    // Leave the maps alone until the runtime call succeeds: a failed
+    // unregister did not release the registration.
     void* buf = nullptr;
     {
         std::shared_lock lk(ctx.map_mutex);
@@ -3537,13 +3512,6 @@ hipError_t playback_hipHostUnregister(PlaybackContext& ctx, const uint8_t* pl) {
         ctx.host_reg_bufs.erase(a->hostPtr);
     }
     ctx.remove_alloc(a->hostPtr);
-    if (buf) {
-#ifdef _WIN32
-        _aligned_free(buf);
-#else
-        free(buf);
-#endif
-    }
     return r;
 }
 

@@ -10,6 +10,7 @@
 #include <algorithm>
 #include <iterator>
 #include <memory>
+#include <new>
 #include <map>
 #include <set>
 #include <unordered_map>
@@ -199,8 +200,9 @@ inline thread_local HrrUnmappedNote hrr_unmapped_note;
 enum class AllocKind : uint8_t {
     Device,        // hipMalloc / hipMallocManaged / hipMallocPitch -> hipFree
     HostMalloc,    // hipHostMalloc / hipMallocHost                 -> hipHostFree
-    HostRegister,  // hipHostRegister backing buffer  -> hipHostUnregister + free
-                   //   (released via host_reg_bufs; skipped in the alloc_map loop)
+    HostRegister,  // hipHostRegister backing buffer  -> hipHostUnregister
+                   //   (via host_reg_bufs; skipped in the alloc_map loop; the
+                   //   block is freed with the context)
     DevicePtrAlias // hipHostGetDevicePointer result  -> not separately freed
                    //   (alias into an already-tracked pinned host allocation)
 };
@@ -445,8 +447,8 @@ struct PlaybackContext {
     // Appended under map_mutex (unique_lock); no lock needed to read (single writer).
     std::vector<hipEvent_t> owned_timing_events;
 
-    // Backing buffers for hipHostRegister replay: recorded ptr -> malloc'd host buffer.
-    // Needed so hipHostUnregister can call hipHostUnregister then free the buffer.
+    // Registered hipHostRegister replay buffers: recorded ptr -> block in
+    // host_reg_blocks. Needed so hipHostUnregister can find the buffer.
     // Guarded by map_mutex.
     std::unordered_map<uint64_t, void*> host_reg_bufs;
 
@@ -886,10 +888,9 @@ struct PlaybackContext {
         std::unique_lock lk(map_mutex);
         auto it = alloc_map.find(rec);
         if (it == alloc_map.end()) return;
-        // Host allocations were charged to host_budget before they were made.
-        if (it->second.kind == AllocKind::HostMalloc ||
-            it->second.kind == AllocKind::HostRegister)
-            host_budget.release(it->second.size);
+        // Charged before it was made. A hipHostRegister block stays charged
+        // in host_reg_blocks.
+        if (it->second.kind == AllocKind::HostMalloc) host_budget.release(it->second.size);
         alloc_map.erase(it);
     }
     // True if an allocation is already tracked under this recorded address.
@@ -902,33 +903,56 @@ struct PlaybackContext {
     // moved or freed before teardown.
     // Keyed by the recorded host address so the same recorded destination
     // reuses one block until a later event asks for more room.
-    struct LandingBuffer {
-        std::unique_ptr<uint8_t[]> cur;   // the block new nodes are given
+    // Blocks start on a 64-byte cache line, as hipHostRegister asks of
+    // registered memory, so two blocks never share one.
+    struct HostBlockFree {
+        void operator()(uint8_t* p) const { ::operator delete[](p, std::align_val_t{64}); }
+    };
+    using HostBlockPtr = std::unique_ptr<uint8_t[], HostBlockFree>;
+    struct HostBlock {
+        HostBlockPtr cur;   // the block new nodes are given
         size_t size = 0;
         // Blocks an earlier node may still point into, outgrown but not freed.
-        std::vector<std::unique_ptr<uint8_t[]>> retired;
+        std::vector<HostBlockPtr> retired;
     };
-    std::unordered_map<uint64_t, LandingBuffer> host_landing_buffers;
-
-    // A recorded key of 0 is a null host pointer the capturing process passed,
-    // and the call it belonged to failed there. Handing back a real buffer
-    // would turn that into a success at replay, so the sentinel is preserved
-    // and the caller's null check reports it.
-    void* host_landing_buffer(uint64_t rec, size_t sz) {
-        if (!rec) return nullptr;
+    std::unordered_map<uint64_t, HostBlock> host_landing_buffers;
+    // hipHostRegister backing buffers, kept the same way: a graph node may
+    // still point into one after hipHostUnregister.
+    std::unordered_map<uint64_t, HostBlock> host_reg_blocks;
+
+    // The block for a recorded host address, at least sz bytes, zero-filled
+    // when new. Null when the allocation fails, and with the replay ended
+    // when the budget refuses a new block.
+    void* host_block(std::unordered_map<uint64_t, HostBlock>& blocks, uint64_t rec,
+                     size_t sz, const char* what) {
         std::unique_lock lk(map_mutex);
-        auto& buf = host_landing_buffers[rec];
+        auto& buf = blocks[rec];
         if (!buf.cur || buf.size < sz) {
             // Charged whole, since the outgrown block stays resident.
-            if (!charge_host(sz, "graph host operand buffer")) return nullptr;
-            auto next = std::make_unique<uint8_t[]>(sz ? sz : 1);
+            if (!charge_host(sz, what)) return nullptr;
+            HostBlockPtr next(new (std::align_val_t{64}, std::nothrow) uint8_t[sz ? sz : 1]());
+            if (!next) {
+                host_budget.release(sz);
+                fprintf(stderr, "[HRR] no %zu-byte %s for recorded host address 0x%llx\n", sz,
+                        what, static_cast<unsigned long long>(rec));
+                return nullptr;
+            }
             if (buf.cur) buf.retired.push_back(std::move(buf.cur));
             buf.cur = std::move(next);
             buf.size = sz;
         }
         return buf.cur.get();
     }
 
+    // A recorded key of 0 is a null host pointer the capturing process passed,
+    // and the call it belonged to failed there. Handing back a real buffer
+    // would turn that into a success at replay, so the sentinel is preserved
+    // and the caller's null check reports it.
+    void* host_landing_buffer(uint64_t rec, size_t sz) {
+        if (!rec) return nullptr;
+        return host_block(host_landing_buffers, rec, sz, "graph host operand buffer");
+    }
+
     // ---- Symbol registration and lookup ----
     void record_symbol(uint64_t rec_host_var, const std::string& name,
                        void* live, size_t sz) {

@@ -1945,7 +1945,8 @@ int main(int argc, char** argv) {
     ctx.kernel_filter.clear();
     if (!run_pass(ctx, archive, thread_events, use_mt, /*log=*/false)) {
       fprintf(stderr, "[HRR] Fatal error in warm-up pass — aborting\n");
-      return 1;
+      // As after the timed pass: queued work may still use ctx's host blocks.
+      std::exit(1);
     }
     ctx.fatal_error.store(false, std::memory_order_release);
     ctx.kernel_filter    = filter;
@@ -1987,7 +1988,7 @@ int main(int argc, char** argv) {
     // matching API. Dispatch on AllocEntry::kind:
     //   Device         -> hipFree
     //   HostMalloc     -> hipHostFree
-    //   HostRegister   -> released below via host_reg_bufs (hipHostUnregister+free)
+    //   HostRegister   -> unregistered below via host_reg_bufs
     //   DevicePtrAlias -> not separately freed (alias into a pinned host alloc)
     for (auto& [rec, entry] : ctx.alloc_map) {
       switch (entry.kind) {
@@ -1998,17 +1999,10 @@ int main(int argc, char** argv) {
       }
     }
     // Captures routinely end mid-stream with hipHostRegister'd buffers still live;
-    // playback_hipHostUnregister never ran for them, so unregister + free each
-    // remaining backing buffer here to avoid leaking both the pinned registration
-    // and the malloc'd buffer every run.
+    // playback_hipHostUnregister never ran for them, so unregister each here.
+    // The blocks are freed with the context.
     for (auto& [rec, buf] : ctx.host_reg_bufs) {
-      if (!buf) continue;
-      (void)hipHostUnregister(buf);
-#ifdef _WIN32
-      _aligned_free(buf);
-#else
-      free(buf);
-#endif
+      if (buf) (void)hipHostUnregister(buf);
     }
     ctx.host_reg_bufs.clear();
     for (auto& [rec, str]   : ctx.stream_map)  (void)hipStreamDestroy(str);
@@ -2041,7 +2035,9 @@ int main(int argc, char** argv) {
       return 2;
     }
     fprintf(stderr, "[HRR] Replay aborted due to fatal HIP error — exiting\n");
-    return 1;
+    // The pass ended without its device sync, so work still queued may use
+    // ctx's host blocks: exit without destroying ctx.
+    std::exit(1);
   }
 
   // ---------------------------------------------------------------------------

@@ -38,6 +38,10 @@ hipError_t playback_hipMallocHost(PlaybackContext& ctx, const uint8_t* pl);
 hipError_t playback_hipHostRegister(PlaybackContext& ctx, const uint8_t* pl);
 hipError_t playback_hipFreeHost(PlaybackContext& ctx, const uint8_t* pl);
 
+// AddressSanitizer aborts on a failed allocation unless told to return null,
+// which the allocation-failure case below needs.
+extern "C" const char* __asan_default_options() { return "allocator_may_return_null=1"; }
+
 namespace fs = std::filesystem;
 
 namespace {
@@ -263,9 +267,10 @@ TEST_CASE("Unit_HRR_HostBudget_DispatchTableUsesTheChargingHandlers", "[hrr]") {
  *   - The old charge goes back only if the runtime let go of the buffer. A
  *     hipHostMalloc entry's buffer is freed (a null live pointer frees
  *     cleanly where a runtime is up, and fails without one, which the charge
- *     must then follow); a failing hipHostFree keeps the charge, as a failing
- *     hipHostUnregister does for a hipHostRegister entry, and the buffer is
- *     left alone. The seeds are not pinned, so the failures need no GPU.
+ *     must then follow); a failing hipHostFree keeps the charge, and the
+ *     buffer is left alone. A hipHostRegister entry's charge stays either way,
+ *     and a failing hipHostUnregister leaves its buffer alone too. The seeds
+ *     are not pinned, so the failures need no GPU.
  */
 TEST_CASE("Unit_HRR_HostBudget_DispatchFreesReplacedHostAllocation", "[hrr]") {
   auto run = [](uint32_t api, auto& rec, AllocKind kind, void* live, size_t expect_used) {
@@ -362,4 +367,38 @@ TEST_CASE("Unit_HRR_HostBudget_LandingBufferPointerSurvivesGrowth", "[hrr]") {
   CHECK(n1[15] == 0x44);
 }
 
+/**
+ * Test Description
+ * ----------------
+ *   - Host blocks start on a 64-byte cache line, so two small blocks never
+ *     share one (hipHostRegister's alignment advice).
+ */
+TEST_CASE("Unit_HRR_HostBudget_HostBlocksStartOnCacheLine", "[hrr]") {
+  PlaybackContext ctx;
+  for (uint64_t rec = 1; rec <= 8; ++rec) {
+    void* p = ctx.host_landing_buffer(rec, 8);
+    REQUIRE(p != nullptr);
+    CHECK(reinterpret_cast<uintptr_t>(p) % 64 == 0);
+  }
+}
+
+/**
+ * Test Description
+ * ----------------
+ *   - A hipHostRegister the budget allows but the allocator cannot back is
+ *     refused with hipErrorMemoryAllocation, and its charge is given back.
+ */
+TEST_CASE("Unit_HRR_HostBudget_HostRegisterAllocationFailureRefuses", "[hrr]") {
+  constexpr uint64_t kHuge = uint64_t{1} << 62;  // more than any address space
+  PlaybackContext ctx;
+  ctx.host_budget.set_cap(kHuge);
+
+  hrr_test::ExactRecord<hrr_args_hipHostRegister> reg;
+  reg.get()->hostPtr = 0x6000;
+  reg.get()->sizeBytes = kHuge;
+  reg.get()->blob_hash_lo = 1;
+  CHECK(playback_hipHostRegister(ctx, reg.bytes()) == hipErrorMemoryAllocation);
+  CHECK(ctx.host_budget.used() == 0);
+}
+
 /** @} */