[Engine] Windows HIP: keep the pinned arena's GPU aliases clear of the process's memory (#3431)

CMakeLists.txt

@@ -1626,6 +1626,12 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   target_link_libraries(pin_budget_test PRIVATE strata_core ${_strata_gpu_runtime_target})
   add_test(NAME pin_budget_test COMMAND pin_budget_test)
 
+  # #3431: the expert arena's registered ranges keep their GPU aliases clear of the process's own memory (Windows HIP; else skips)
+  add_executable(pinned_alias_test src/core/pinned_alias_test.cpp)
+  target_link_libraries(pinned_alias_test PRIVATE strata_core ${_strata_gpu_runtime_target})
+  add_test(NAME pinned_alias_test COMMAND pinned_alias_test)
+  set_tests_properties(pinned_alias_test PROPERTIES SKIP_RETURN_CODE 77)
+
   # #3350 / #3470: the conversation copy streams and events are pooled across short-lived threads (needs a device, else skips)
   add_executable(conversation_copy_streams_test src/core/conversation_copy_streams_test.cpp)
   target_include_directories(conversation_copy_streams_test PRIVATE include)

docs/AMD_HIP.md

@@ -151,6 +151,14 @@ source-checkout option for developers, not a way to start the released app.
 - **VRAM budget:** `hipMemGetInfo` on Windows does not subtract what the desktop and other programs hold. The engine lowers
   the free figure by what Windows' video memory budget for the process withholds (logged once). `STRATA_WDDM_BUDGET=0` turns
   that off.
+- **Pinned expert arena (#3431):** the HIP runtime gives each page-locked range of the arena a second address for the GPU, and
+  it takes any pointer of the process inside that range for the arena. Through engine 0.1.54 about one start in ten on an RX
+  7900 XTX had the engine thread's stack there (the same symptom was reported on an R9700): every long prompt was then read
+  wrong until a restart, with no error (the log shows `KV streaming: 100.00% of 0 block reads` for each long prompt). The
+  engine now keeps the address space below 1 TiB reserved, so its stacks and other memory are placed above the GPU addresses,
+  and it reserves each GPU address range while the arena is pinned. When a range's GPU address still covers memory in use,
+  that range and the ones after it stay unpinned and the log says so (`strata pin: ... (#3431)`); that start reads prompts more
+  slowly but correctly. `STRATA_STACK_ISOLATE=0` runs the engine as before, without the reservation below 1 TiB.
 - **WMMA prompt attention** is opt-in on Windows (`STRATA_HIP_WMMA=1`); it is the default on gfx12 on Linux.
 
 **Reporting a Windows AMD run.** Give your card, the driver version (AMD Software > System), the output of the two commands

include/strata/core/pinned.hpp

@@ -42,6 +42,13 @@ inline std::atomic<bool>& arena_uma_no_register_flag() {
 /// STRATA_ARENA_PIN_GIB to a number (STRATA_ARENA_PIN_GIB=0 registers the whole arena as before).
 bool arena_no_register();
 
+/// #3431 (Windows, HIP): the stack [lo, hi) of a thread that only waits while the arena is registered (the process's first thread
+/// when the engine runs on its own one); a GPU alias may cover it.  0, 0: none.
+void set_idle_thread_stack(uint64_t lo, uint64_t hi);
+/// #3431 (Windows, HIP): the start addresses of the address-space reservations main() keeps for the life of the process; a GPU
+/// alias inside one is clear, nothing can be placed there.
+void set_held_address_space(std::vector<uint64_t> bases);
+
 struct PinnedArena {
     void* base = nullptr;
     uint64_t capacity = 0;
@@ -88,6 +95,8 @@ struct PinnedArena {
     void register_slices(std::atomic<int>& ready);
     std::vector<uint64_t> bounds_;
     std::vector<uint64_t> slice_starts;
+    /// #3431 (Windows, HIP): the address-space reservations under the GPU aliases of the registered ranges, released after them
+    std::vector<void*> alias_claims;
     void* mapping_base = nullptr;     ///< actual mapping start; differs from base when a shared-file header exists
     uint64_t mapping_bytes = 0;       ///< bytes to release from mapping_base
     ~PinnedArena();

src/core/pinned.cu

@@ -343,6 +343,21 @@ namespace {
 bool clear_error() { (void) cudaGetLastError(); return true; }
 }  // namespace
 
+namespace {
+std::atomic<uint64_t> g_idle_stack_lo{0}, g_idle_stack_hi{0};
+std::vector<uint64_t> g_held;   // sorted; set once by main() before any thread that pins exists
+}  // namespace
+
+void set_idle_thread_stack(uint64_t lo, uint64_t hi) {
+    g_idle_stack_lo.store(lo);
+    g_idle_stack_hi.store(hi);
+}
+
+void set_held_address_space(std::vector<uint64_t> bases) {
+    std::sort(bases.begin(), bases.end());
+    g_held = std::move(bases);
+}
+
 bool arena_no_register() {
     const char* e = std::getenv("STRATA_ARENA_PIN_GIB");
     if (e != nullptr && *e != '\0') return std::string(e) == "none";   // an explicit setting wins over the engine's flag
@@ -399,6 +414,111 @@ std::vector<uint64_t> uniform_bounds(uint64_t bytes, uint64_t slice) {
     if (!b.empty()) b.push_back(b.back() + slice);
     return b;
 }
+
+#if defined(_WIN32) && defined(STRATA_USE_HIP)
+// #3431.  On Windows the HIP runtime gives every cudaHostRegister'ed range a second address for the GPU (its alias), and it
+// takes any host pointer inside that alias for the registered memory.  Unlike the addresses of cudaMalloc, the alias is not
+// reserved in the process's own address space, so a thread stack, a heap block or a mapping can be placed there.  In about one
+// engine start in ten one was (the stack of the thread that pins the arena): a cudaMemcpy from a local variable then read the
+// arena, a copy back into one wrote the arena, and every long prompt was read wrong until the process ended, with no error.
+// So the alias range is reserved here (address space only, nothing committed), and a range whose alias already covers memory
+// in use is not kept: true when the alias is clear (or there is none), with its reservations added to `claims`.  A runtime that
+// reserves the alias itself would have every range refused here, and logged.
+bool claim_alias(void* host, uint64_t n, std::vector<void*>& claims, std::string& why) {
+    void* d = nullptr;
+    if (cudaHostGetDevicePointer(&d, host, 0) != cudaSuccess || d == nullptr || d == host) {
+        (void) cudaGetLastError();
+        return true;
+    }
+    SYSTEM_INFO si;
+    GetSystemInfo(&si);
+    const uint64_t gran = si.dwAllocationGranularity;                 // where VirtualAlloc can start a reservation
+    const uint64_t top = (uint64_t) si.lpMaximumApplicationAddress;  // nothing of the process lies above
+    const uint64_t a = (uint64_t) d, end = std::min<uint64_t>(a + n, top);
+    std::vector<void*> got;
+    MEMORY_BASIC_INFORMATION mbi;
+    char msg[256] = "";
+    // the granule the alias starts in: when it is free in front of the alias too, the reservation starts at its boundary, or a
+    // small allocation could still be placed there and reach into the alias
+    uint64_t at = a;
+    const uint64_t a0 = a & ~(gran - 1);
+    if (a0 < a && a0 < end && VirtualQuery((void*) a0, &mbi, sizeof mbi) == sizeof mbi && mbi.State == MEM_FREE &&
+        (uint64_t) mbi.BaseAddress + mbi.RegionSize > a)
+        at = a0;
+    while (at < end) {
+        if (VirtualQuery((void*) at, &mbi, sizeof mbi) != sizeof mbi) {
+            std::snprintf(msg, sizeof msg, "VirtualQuery failed at %#llx", (unsigned long long) at);
+            break;
+        }
+        const uint64_t re = std::min<uint64_t>((uint64_t) mbi.BaseAddress + mbi.RegionSize, end);
+        const uint64_t idle_lo = g_idle_stack_lo.load(), idle_hi = g_idle_stack_hi.load();
+        const bool held = mbi.State == MEM_RESERVE && std::binary_search(g_held.begin(), g_held.end(), (uint64_t) mbi.AllocationBase);
+        const bool idle = mbi.State != MEM_FREE && idle_lo < idle_hi && (uint64_t) mbi.AllocationBase == idle_lo && re <= idle_hi;
+        if (held || idle) {
+            // held: main()'s reservation, kept for the life of the process.  idle: the first thread's stack, which only waits
+            at = (uint64_t) mbi.BaseAddress + mbi.RegionSize;
+            continue;
+        }
+        if (mbi.State != MEM_FREE) {
+            ULONG_PTR lo = 0, hi = 0;
+            GetCurrentThreadStackLimits(&lo, &hi);
+            const bool own_stack = at < (uint64_t) hi && (uint64_t) lo < re;
+            const char* kind = own_stack ? "this thread's stack" : mbi.Type == MEM_IMAGE ? "a module" :
+                               mbi.Type == MEM_MAPPED ? "a mapped view" : "private memory";
+            std::snprintf(msg, sizeof msg, "the GPU alias [%#llx, +%llu MiB) of the range at %p covers memory in use at %#llx (%s, from %p)",
+                          (unsigned long long) a, (unsigned long long) (n >> 20), host, (unsigned long long) at, kind,
+                          mbi.AllocationBase);
+            break;
+        }
+        // a free piece that starts inside a granule shares it with an allocation: nothing new can start there
+        const uint64_t b = (at + gran - 1) & ~(gran - 1);
+        if (b < re) {
+            void* r = VirtualAlloc((void*) b, (SIZE_T) (re - b), MEM_RESERVE, PAGE_NOACCESS);
+            if (r == nullptr) {
+                std::snprintf(msg, sizeof msg, "could not reserve the GPU alias [%#llx, +%llu MiB) of the range at %p (error %lu)",
+                              (unsigned long long) a, (unsigned long long) (n >> 20), host, GetLastError());
+                break;
+            }
+            got.push_back(r);
+        }
+        at = (uint64_t) mbi.BaseAddress + mbi.RegionSize;
+    }
+    if (msg[0] != '\0') {
+        for (void* g : got) VirtualFree(g, 0, MEM_RELEASE);
+        why = msg;
+        return false;
+    }
+    claims.insert(claims.end(), got.begin(), got.end());
+    return true;
+}
+#endif
+
+// after the ranges are unregistered
+void release_alias_claims(std::vector<void*>& claims) {
+#ifdef _WIN32
+    for (void* c : claims) VirtualFree(c, 0, MEM_RELEASE);
+#endif
+    claims.clear();
+}
+
+// pin_register for the arena's ranges.  Windows HIP: a range whose GPU alias covers memory in use is unregistered again and
+// reported as refused, so the caller's "the driver refused it" path applies to it (#3431).
+cudaError_t register_range(PinnedArena& arena, void* p, uint64_t n) {
+    const cudaError_t e = pin_register(p, (size_t) n, cudaHostRegisterPortable | cudaHostRegisterMapped);
+#if defined(_WIN32) && defined(STRATA_USE_HIP)
+    std::string why;
+    if (e == cudaSuccess && !claim_alias(p, n, arena.alias_claims, why)) {
+        (void) pin_unregister(p);
+        (void) cudaGetLastError();
+        std::fprintf(stderr, "strata pin: %s; that range stays unregistered (#3431)\n", why.c_str());
+        arena.note = "left unregistered: " + why + " (#3431); " + arena.note;
+        return cudaErrorMemoryAllocation;
+    }
+#else
+    (void) arena;
+#endif
+    return e;
+}
 }  // namespace
 
 PinnedArena::PinnedArena(uint64_t bytes, uint64_t slice) : PinnedArena(bytes, uniform_bounds(bytes, slice)) {
@@ -444,7 +564,7 @@ PinnedArena::PinnedArena(uint64_t bytes, const std::vector<uint64_t>& bounds,
         }
         const bool capped = cap > 0 && cap < bytes && bounds.size() >= 2;
         const cudaError_t e = capped ? cudaSuccess :
-            pin_register(base, (size_t) bytes, cudaHostRegisterPortable | cudaHostRegisterMapped);
+            register_range(*this, base, bytes);
         if (!capped && e == cudaSuccess) {
             note = "cudaHostRegister PORTABLE ok; " + note;
             registered_bytes = bytes;
@@ -464,7 +584,7 @@ PinnedArena::PinnedArena(uint64_t bytes, const std::vector<uint64_t>& bounds,
             for (size_t i = 0; i + 1 < bounds.size(); ++i) {
                 const uint64_t off = bounds[i], n = bounds[i + 1] - bounds[i];
                 if (limited && (off > limit || n > limit - off)) break;
-                if (pin_register((uint8_t*) base + off, (size_t) n, cudaHostRegisterPortable | cudaHostRegisterMapped) != cudaSuccess) {
+                if (register_range(*this, (uint8_t*) base + off, n) != cudaSuccess) {
                     (void) cudaGetLastError();
                     break;
                 }
@@ -572,8 +692,7 @@ PinnedArena::PinnedArena(uint64_t bytes, ByNeed, const std::vector<std::pair<uin
         if (it != done.end() && it->first < b) b = it->first;
         if (a >= b) continue;
         if (cap > 0 && total + (b - a) > cap) { capped = true; break; }
-        if (pin_register((uint8_t*) base + a, (size_t) (b - a), cudaHostRegisterPortable | cudaHostRegisterMapped) !=
-            cudaSuccess) {
+        if (register_range(*this, (uint8_t*) base + a, b - a) != cudaSuccess) {
             (void) cudaGetLastError();
             refused = 1;
             break;
@@ -628,8 +747,7 @@ void PinnedArena::register_slices(std::atomic<int>& ready) {
     int i = 0;
     for (; base != nullptr && i < n; ++i) {
         const uint64_t off = bounds_[(size_t) i], len = bounds_[(size_t) i + 1] - off;
-        if (pin_register((uint8_t*) base + off, (size_t) len, cudaHostRegisterPortable | cudaHostRegisterMapped) !=
-            cudaSuccess) {
+        if (register_range(*this, (uint8_t*) base + off, len) != cudaSuccess) {
             (void) cudaGetLastError();
             break;
         }
@@ -658,6 +776,7 @@ PinnedArena::~PinnedArena() {
     if (base && by_need) {
         for (const auto& r : locked_ranges) strata::platform::unlock_resident((uint8_t*) base + r.first, r.second);
         for (const auto& r : pin_ranges) pin_unregister((uint8_t*) base + r.first);
+        release_alias_claims(alias_claims);
         release(mapping_base ? mapping_base : base, mapping_bytes ? mapping_bytes : capacity);
         base = nullptr;
         mapping_base = nullptr;
@@ -671,6 +790,7 @@ PinnedArena::~PinnedArena() {
         } else {
             pin_unregister(base);
         }
+        release_alias_claims(alias_claims);
         release(mapping_base ? mapping_base : base, mapping_bytes ? mapping_bytes : capacity);
         base = nullptr;
         mapping_base = nullptr;

src/core/pinned_alias_test.cpp

@@ -0,0 +1,111 @@
+// pinned_alias_test (#3431, Windows HIP): the GPU alias of every range the expert arena registers is reserved in the process's
+// own address space while the arena lives, and a range whose alias covers memory already in use is not kept registered - a
+// copy through that memory must still reach the device and come back.  Needs an AMD GPU on Windows, else it skips (77).
+#include "strata/core/pinned.hpp"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+#if defined(_WIN32)
+#include <windows.h>
+#endif
+
+#if defined(_WIN32) && defined(STRATA_USE_HIP)
+static int g_fail = 0;
+#define CHECK(c)                                                              \
+    do {                                                                      \
+        if (!(c)) { std::fprintf(stderr, "FAIL %s:%d  %s\n", __FILE__, __LINE__, #c); ++g_fail; } \
+    } while (0)
+
+static bool address_free(const void* p) {
+    MEMORY_BASIC_INFORMATION m{};
+    return VirtualQuery(p, &m, sizeof m) == sizeof m && m.State == MEM_FREE;
+}
+
+static void* alias_of(void* host) {
+    void* d = nullptr;
+    if (cudaHostGetDevicePointer(&d, host, 0) != cudaSuccess) {
+        (void) cudaGetLastError();
+        return nullptr;
+    }
+    return d;
+}
+
+// 16 bytes host -> device -> host through `buf`
+static bool round_trip(uint8_t* buf) {
+    const uint32_t pat[4] = {0x11223344u, 0x55667788u, 0x99aabbccu, 0xddeeff00u};
+    void* d = nullptr;
+    if (cudaMalloc(&d, 64) != cudaSuccess) return false;
+    std::memcpy(buf, pat, sizeof pat);
+    std::memset(buf + 64, 0, sizeof pat);
+    const bool ok = cudaMemcpy(d, buf, sizeof pat, cudaMemcpyHostToDevice) == cudaSuccess &&
+                    cudaMemcpy(buf + 64, d, sizeof pat, cudaMemcpyDeviceToHost) == cudaSuccess &&
+                    std::memcmp(buf + 64, pat, sizeof pat) == 0;
+    cudaFree(d);
+    return ok;
+}
+#endif
+
+int main() {
+#if !(defined(_WIN32) && defined(STRATA_USE_HIP))
+    std::printf("pinned_alias_test: SKIP (Windows HIP only)\n");
+    return 77;
+#else
+    using namespace strata::core;
+    if (cudaFree(0) != cudaSuccess) {
+        std::printf("pinned_alias_test: SKIP (no GPU)\n");
+        return 77;
+    }
+    (void) cudaGetLastError();
+    const uint64_t slice = 64ull << 20, bytes = 4 * slice;
+    const std::vector<uint64_t> bounds{0, slice, 2 * slice, 3 * slice, 4 * slice};
+
+    // 1. while the arena lives, nothing can be placed under the alias of what it registered; afterwards the range is free again
+    void* first_alias = nullptr;
+    {
+        PinnedArena a(bytes, bounds);
+        CHECK(a.valid());
+        CHECK(a.registered_bytes == bytes);
+        first_alias = alias_of(a.base);
+        CHECK(first_alias != nullptr);
+        if (first_alias != nullptr && first_alias != a.base) {
+            CHECK(!a.alias_claims.empty());
+            CHECK(!address_free(first_alias));
+            CHECK(!address_free((uint8_t*) first_alias + bytes - 4096));
+        }
+    }
+    if (first_alias == nullptr) {
+        std::printf("pinned_alias_test: %d FAILED\n", g_fail);
+        return 1;
+    }
+    CHECK(address_free(first_alias));
+
+    // 2. memory already in use under the alias (in the engine: the stack of the thread that pins the arena).  The runtime hands
+    //    the same alias out again, so the next arena's first range would cover it: that range must stay unregistered, and a
+    //    copy through the memory must work.  Without the guard the copy reads and writes the arena instead (round trip zeros).
+    uint8_t* victim = (uint8_t*) VirtualAlloc(first_alias, 1u << 20, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    CHECK(victim != nullptr);
+    if (victim != nullptr) {
+        PinnedArena b(bytes, bounds);
+        CHECK(b.valid());
+        bool covered = false;
+        for (uint64_t off = 0; off < b.registered_bytes; off += slice) {
+            uint8_t* al = (uint8_t*) alias_of((uint8_t*) b.base + off);
+            if (al != nullptr && al <= victim && victim < al + slice) covered = true;
+        }
+        CHECK(!covered);
+        if (b.registered_bytes == bytes)
+            std::printf("pinned_alias_test: note: the runtime chose another alias this time; the refusal was not exercised\n");
+        else
+            CHECK(b.note.find("#3431") != std::string::npos);
+        CHECK(round_trip(victim));
+        VirtualFree(victim, 0, MEM_RELEASE);
+    }
+    std::printf(g_fail == 0 ? "pinned_alias_test: OK\n" : "pinned_alias_test: %d FAILED\n", g_fail);
+    return g_fail == 0 ? 0 : 1;
+#endif
+}