Harden large CUDA host staging allocations

csrc/instant_tensor/dl_binding/cuda_binding.hpp

@@ -18,6 +18,11 @@
 #define cudaHostRegisterIoMemory 0x04
 #define cudaHostRegisterReadOnly 0x08
 
+#define cudaHostAllocDefault 0x00
+#define cudaHostAllocPortable 0x01
+#define cudaHostAllocMapped 0x02
+#define cudaHostAllocWriteCombined 0x04
+
 namespace instanttensor {
 namespace cuda_binding {
 
@@ -54,6 +59,9 @@ inline cudaError_t (*cudaEventSynchronize_fn)(cudaEvent_t event) = nullptr;
 inline cudaError_t (*cudaStreamWaitEvent_fn)(cudaStream_t stream, cudaEvent_t event, unsigned int flags) = nullptr;
 inline cudaError_t (*cudaHostRegister_fn)(void* ptr, size_t size, unsigned int flags) = nullptr;
 inline cudaError_t (*cudaHostUnregister_fn)(void* ptr) = nullptr;
+inline cudaError_t (*cudaHostAlloc_fn)(void** ptr, size_t size, unsigned int flags) = nullptr;
+inline cudaError_t (*cudaFreeHost_fn)(void* ptr) = nullptr;
+inline cudaError_t (*cudaGetLastError_fn)() = nullptr;
 inline const char* (*cudaGetErrorString_fn)(cudaError_t error) = nullptr;
 
 inline bool init() {
@@ -95,6 +103,9 @@ inline bool init() {
     cudaStreamWaitEvent_fn = resolve<decltype(cudaStreamWaitEvent_fn)>(lib_handle, {"cudaStreamWaitEvent", "hipStreamWaitEvent"});
     cudaHostRegister_fn = resolve<decltype(cudaHostRegister_fn)>(lib_handle, {"cudaHostRegister", "hipHostRegister"});
     cudaHostUnregister_fn = resolve<decltype(cudaHostUnregister_fn)>(lib_handle, {"cudaHostUnregister", "hipHostUnregister"});
+    cudaHostAlloc_fn = resolve<decltype(cudaHostAlloc_fn)>(lib_handle, {"cudaHostAlloc", "hipHostMalloc"});
+    cudaFreeHost_fn = resolve<decltype(cudaFreeHost_fn)>(lib_handle, {"cudaFreeHost", "hipHostFree"});
+    cudaGetLastError_fn = resolve<decltype(cudaGetLastError_fn)>(lib_handle, {"cudaGetLastError", "hipGetLastError"});
     cudaGetErrorString_fn = resolve<decltype(cudaGetErrorString_fn)>(lib_handle, {"cudaGetErrorString", "hipGetErrorString"});
 
     return true;
@@ -137,9 +148,18 @@ inline cudaError_t cudaHostRegister(void* ptr, size_t size, unsigned int flags)
 inline cudaError_t cudaHostUnregister(void* ptr) {
     return cudaHostUnregister_fn(ptr);
 }
+inline cudaError_t cudaHostAlloc(void** ptr, size_t size, unsigned int flags) {
+    return cudaHostAlloc_fn(ptr, size, flags);
+}
+inline cudaError_t cudaFreeHost(void* ptr) {
+    return cudaFreeHost_fn(ptr);
+}
+inline cudaError_t cudaGetLastError() {
+    return cudaGetLastError_fn();
+}
 inline const char* cudaGetErrorString(cudaError_t error) {
     return cudaGetErrorString_fn(error);
 }
 
 } // namespace cuda_binding
-}
+}

csrc/instant_tensor/host_registration.hpp

@@ -0,0 +1,143 @@
+#pragma once
+
+#include <algorithm>
+#include <cstddef>
+#include <exception>
+#include <stdexcept>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace instanttensor {
+
+struct HostRegistrationRange {
+    void* ptr;
+    size_t size;
+};
+
+struct HostRegistration {
+    std::vector<HostRegistrationRange> ranges;
+    int whole_buffer_error = 0;
+};
+
+struct HostBufferAllocation {
+    void* ptr = nullptr;
+    HostRegistration registration;
+    bool runtime_allocated = false;
+    std::string registration_failure;
+};
+
+template <typename RegisterFn, typename UnregisterFn, typename ClearErrorFn>
+HostRegistration register_host_buffer(
+    void* ptr,
+    size_t size,
+    unsigned int flags,
+    size_t segment_size,
+    RegisterFn&& register_fn,
+    UnregisterFn&& unregister_fn,
+    ClearErrorFn&& clear_error_fn
+) {
+    if (ptr == nullptr || size == 0 || segment_size == 0) {
+        throw std::invalid_argument("Host registration requires non-empty storage and segments");
+    }
+
+    HostRegistration registration;
+    int result = register_fn(ptr, size, flags);
+    if (result == 0) {
+        registration.ranges.push_back({ptr, size});
+        return registration;
+    }
+
+    registration.whole_buffer_error = result;
+    clear_error_fn();
+    auto* base = static_cast<char*>(ptr);
+    for (size_t offset = 0; offset < size; offset += segment_size) {
+        const size_t current_size = std::min(segment_size, size - offset);
+        void* current_ptr = base + offset;
+        result = register_fn(current_ptr, current_size, flags);
+        if (result == 0) {
+            registration.ranges.push_back({current_ptr, current_size});
+            continue;
+        }
+
+        clear_error_fn();
+        for (auto it = registration.ranges.rbegin(); it != registration.ranges.rend(); ++it) {
+            unregister_fn(it->ptr);
+        }
+        throw std::runtime_error(
+            "Host registration failed for the whole buffer (code "
+            + std::to_string(registration.whole_buffer_error)
+            + ") and segment at offset " + std::to_string(offset)
+            + " (code " + std::to_string(result) + ")"
+        );
+    }
+    return registration;
+}
+
+template <
+    typename AlignedAllocFn,
+    typename FreeFn,
+    typename RegisterFn,
+    typename UnregisterFn,
+    typename ClearErrorFn,
+    typename RuntimeAllocFn>
+HostBufferAllocation allocate_registered_host_buffer(
+    size_t size,
+    size_t alignment,
+    unsigned int register_flags,
+    unsigned int runtime_alloc_flags,
+    size_t segment_size,
+    AlignedAllocFn&& aligned_alloc_fn,
+    FreeFn&& free_fn,
+    RegisterFn&& register_fn,
+    UnregisterFn&& unregister_fn,
+    ClearErrorFn&& clear_error_fn,
+    RuntimeAllocFn&& runtime_alloc_fn
+) {
+    if (size == 0 || alignment == 0) {
+        throw std::invalid_argument("Host allocation requires non-empty aligned storage");
+    }
+
+    HostBufferAllocation allocation;
+    allocation.ptr = aligned_alloc_fn(alignment, size);
+    if (allocation.ptr == nullptr) {
+        throw std::runtime_error("Failed to allocate aligned host storage");
+    }
+
+    try {
+        allocation.registration = register_host_buffer(
+            allocation.ptr,
+            size,
+            register_flags,
+            segment_size,
+            std::forward<RegisterFn>(register_fn),
+            std::forward<UnregisterFn>(unregister_fn),
+            std::forward<ClearErrorFn>(clear_error_fn)
+        );
+        return allocation;
+    } catch (const std::exception& error) {
+        allocation.registration_failure = error.what();
+    }
+
+    free_fn(allocation.ptr);
+    allocation.ptr = nullptr;
+    clear_error_fn();
+
+    const int result = runtime_alloc_fn(
+        &allocation.ptr,
+        size,
+        runtime_alloc_flags
+    );
+    if (result != 0 || allocation.ptr == nullptr) {
+        throw std::runtime_error(
+            allocation.registration_failure
+            + "; runtime pinned allocation failed (code "
+            + std::to_string(result) + ")"
+        );
+    }
+
+    allocation.runtime_allocated = true;
+    return allocation;
+}
+
+} // namespace instanttensor

csrc/loader_common.cpp

@@ -1,8 +1,13 @@
 #include <instant_tensor/loader.hpp>
+#include <instant_tensor/host_registration.hpp>
 
 
 namespace instanttensor {
 
+namespace {
+constexpr size_t HOST_REGISTER_SEGMENT_SIZE = 256ULL * 1024 * 1024;
+}
+
 
 int Loader::next_executor_request_id() {
     int request_id = this->executor_request_id;
@@ -95,22 +100,72 @@ void Loader::init_buffer() {
         size_t host_buffer_size = inflight_host_buffer_size;
         this->host_buffer_entry = host_buffer_cache->get(host_buffer_size);
         if (this->host_buffer_entry.ptr == NULL) {
-            // aligned_alloc + cudaHostRegister is faster than cudaHostAlloc
-            this->host_buffer_entry.ptr = aligned_alloc(this->thread_alignment, host_buffer_size);
-            if (this->host_buffer_entry.ptr == NULL) {
-                throw std::runtime_error("Failed to allocate host buffer: " + std::string(strerror(errno)));
+            // Some CUDA driver/kernel combinations reject a multi-GiB
+            // cudaHostRegister call even though the same storage can be pinned
+            // in smaller adjacent ranges. A few driver/kernel combinations
+            // reject registration itself; cudaHostAlloc remains the portable
+            // pinned-memory fallback for those hosts.
+            const unsigned int runtime_alloc_flags =
+                this->cuda_host_register_flags
+                & (cudaHostAllocPortable | cudaHostAllocMapped);
+            HostBufferAllocation allocation = allocate_registered_host_buffer(
+                host_buffer_size,
+                this->thread_alignment,
+                this->cuda_host_register_flags,
+                runtime_alloc_flags,
+                HOST_REGISTER_SEGMENT_SIZE,
+                [](size_t alignment, size_t size) {
+                    return aligned_alloc(alignment, size);
+                },
+                [](void* ptr) { free(ptr); },
+                [](void* ptr, size_t size, unsigned int flags) {
+                    return static_cast<int>(cudaHostRegister(ptr, size, flags));
+                },
+                [](void* ptr) {
+                    return static_cast<int>(cudaHostUnregister(ptr));
+                },
+                []() { cudaGetLastError(); },
+                [](void** ptr, size_t size, unsigned int flags) {
+                    return static_cast<int>(cudaHostAlloc(ptr, size, flags));
+                }
+            );
+            this->host_buffer_entry.ptr = allocation.ptr;
+            if (allocation.runtime_allocated) {
+                fprintf(
+                    stderr,
+                    "[InstantTensor] cudaHostRegister paths failed for a %.2f "
+                    "GiB buffer; using runtime-allocated pinned memory. %s\n",
+                    host_buffer_size / static_cast<double>(1ULL << 30),
+                    allocation.registration_failure.c_str()
+                );
+            } else if (allocation.registration.whole_buffer_error != 0) {
+                fprintf(
+                    stderr,
+                    "[InstantTensor] cudaHostRegister rejected a %.2f GiB buffer "
+                    "(code %d); registered it as %zu segments instead.\n",
+                    host_buffer_size / static_cast<double>(1ULL << 30),
+                    allocation.registration.whole_buffer_error,
+                    allocation.registration.ranges.size()
+                );
             }
 
-            // NOTE: cudaHostRegisterReadOnly is not used since the host buffer is writable for the CPU
-            CUDA_CHECK(cudaHostRegister(this->host_buffer_entry.ptr, host_buffer_size, this->cuda_host_register_flags));
-
             this->host_buffer_entry.size = host_buffer_size;
-            this->host_buffer_entry.deleter = [=](void *ptr) {
+            this->host_buffer_entry.deleter = [=](void* ptr) {
                 if (this->backend == Backend::URING || this->backend == Backend::URING_BUFFERED) {
                     this->deregister_host_buffer_uring();
                 }
-                CUDA_CHECK(cudaHostUnregister(ptr));
-                free(ptr);
+                if (allocation.runtime_allocated) {
+                    CUDA_CHECK(cudaFreeHost(ptr));
+                } else {
+                    for (
+                        auto it = allocation.registration.ranges.rbegin();
+                        it != allocation.registration.ranges.rend();
+                        ++it
+                    ) {
+                        CUDA_CHECK(cudaHostUnregister(it->ptr));
+                    }
+                    free(ptr);
+                }
             };
         }
         this->host_buffer = (char*)this->host_buffer_entry.ptr;