diff --git a/src/umbp/distributed/transfer/hbm_copy_engine.cpp b/src/umbp/distributed/transfer/hbm_copy_engine.cpp index a72cc204b..baec7e040 100644 --- a/src/umbp/distributed/transfer/hbm_copy_engine.cpp +++ b/src/umbp/distributed/transfer/hbm_copy_engine.cpp @@ -264,26 +264,55 @@ void LogPlanGeometry(const char* path, size_t plan_index, const TransferPlan& pl } } -// One non-blocking stream per device, per thread. Submit runs on PoolClient's -// executor threads, so a shared stream would serialize copies that the -// hipMemcpy path ran independently. -std::map& GatherStreams() { - static thread_local std::map* streams = new std::map(); - return *streams; -} +// Non-blocking gather streams, pooled per device and lent to one gather at a +// time, so concurrent copies still run on separate streams. They are not tied +// to the calling thread: the standalone server calls Submit from gRPC's sync +// pool, which retires and spawns threads under load, and every HIP stream holds +// ~6.4 MiB of device memory for its kernarg pool. Kept per thread, they leaked +// until the GPU was full and stream creation failed (which HIP 7.2 turns into a +// segfault). A device's pool grows only to its peak number of concurrent +// gathers. Never destroyed: worker threads can still be copying during static +// teardown. +class GatherStreamPool { + public: + static GatherStreamPool& Instance() { + static GatherStreamPool* pool = new GatherStreamPool(); + return *pool; + } -hipStream_t GatherStream(int device_id) { - auto& streams = GatherStreams(); - auto it = streams.find(device_id); - if (it != streams.end()) return it->second; - hipStream_t stream = nullptr; - if (hipStreamCreateWithFlags(&stream, hipStreamNonBlocking) != hipSuccess) { - (void)hipGetLastError(); - return nullptr; + // The current device must already be `device_id`. nullptr if no stream is + // free and a new one cannot be created. + hipStream_t Acquire(int device_id) { + { + std::lock_guard lock(mutex_); + auto& idle = idle_[device_id]; + if (!idle.empty()) { + hipStream_t stream = idle.back(); + idle.pop_back(); + return stream; + } + } + hipStream_t stream = nullptr; + if (hipStreamCreateWithFlags(&stream, hipStreamNonBlocking) != hipSuccess) { + (void)hipGetLastError(); + return nullptr; + } + created_.fetch_add(1, std::memory_order_relaxed); + return stream; } - streams.emplace(device_id, stream); - return stream; -} + + void Release(int device_id, hipStream_t stream) { + std::lock_guard lock(mutex_); + idle_[device_id].push_back(stream); + } + + uint64_t Created() const { return created_.load(std::memory_order_relaxed); } + + private: + std::mutex mutex_; + std::map> idle_; + std::atomic created_{0}; +}; // Latches on the first gather failure and never tries again. A kernel that // cannot launch here will not start working later, and the fallback is always @@ -398,6 +427,8 @@ void* HbmCopyEngine::HostRegionDeviceAddress(const void* ptr, size_t size, int d return nullptr; } +uint64_t HbmCopyEngine::GatherStreamsCreated() { return GatherStreamPool::Instance().Created(); } + // Run the gather kernel over every eligible segment in the batch, and report // which plans it completed. // @@ -558,7 +589,7 @@ std::vector HbmCopyEngine::GatherEligiblePlans(const std::vector HbmCopyEngine::GatherEligiblePlans(const std::vector Submit(std::vector plans) override; + // Gather streams created since process start, across all engines. They are + // pooled per device and outlive the thread that first used them, so this + // tracks the peak number of gathers in flight at once, not the number of + // threads that have called Submit. Lets a test assert the reuse. + static uint64_t GatherStreamsCreated(); + private: // Why a plan did not take the gather kernel. The hot path never computes // this: it is filled in only when debug mode is on (UMBP_HBM_COPY_DEBUG), so diff --git a/src/umbp/local/tiers/device_gather.hip b/src/umbp/local/tiers/device_gather.hip index 5ba665d43..9372511f5 100644 --- a/src/umbp/local/tiers/device_gather.hip +++ b/src/umbp/local/tiers/device_gather.hip @@ -26,8 +26,9 @@ #include #include #include -#include +#include #include +#include #include "mori/utils/mori_log.hpp" #include "umbp/common/device_gather.h" @@ -71,25 +72,31 @@ __global__ void GatherFragmentsKernel(const DeviceGatherFragment* __restrict__ f for (size_t i = index; i < fragment.bytes; i += stride) dst[i] = src[i]; } -// Per-thread, per-device descriptor staging. +// Per-stream descriptor staging. // -// Per-thread because two threads serving the same device share one stream (see -// DeviceStream), so a shared staging buffer would be a data race even though -// the stream itself is ordered. Never freed, for the same reason DeviceStream -// never destroys its streams: releasing HIP resources from a thread or static -// destructor races with runtime teardown. The count is bounded by the server's -// worker threads times the devices they touch. +// The caller holds `stream` exclusively from the upload to the synchronize +// (see LaunchDeviceGather), so a stream's buffer is never in use twice at once, +// and the count is bounded by the caller's streams rather than by however many +// threads have ever called in. Never freed, for the same reason the streams are +// never destroyed: releasing HIP resources from a static destructor races with +// runtime teardown. struct DescriptorStaging { DeviceGatherFragment* host = nullptr; DeviceGatherFragment* device = nullptr; size_t capacity = 0; }; -DescriptorStaging* StagingFor(int device_id, size_t count) { - static thread_local std::map* staging_by_device = - new std::map(); +DescriptorStaging* StagingFor(hipStream_t stream, int device_id, size_t count) { + static std::mutex* staging_mutex = new std::mutex(); + static auto* staging_by_stream = new std::unordered_map(); - DescriptorStaging& staging = (*staging_by_device)[device_id]; + DescriptorStaging* staging_entry = nullptr; + { + // References into an unordered_map stay valid across rehashing. + std::lock_guard lock(*staging_mutex); + staging_entry = &(*staging_by_stream)[stream]; + } + DescriptorStaging& staging = *staging_entry; if (staging.capacity >= count) return &staging; // Grow with headroom so a slowly increasing batch size does not reallocate @@ -154,7 +161,7 @@ bool LaunchDeviceGather(const DeviceGatherFragment* fragments, size_t count, int (chunks + kThreadsPerBlock - 1) / kThreadsPerBlock, 1, kMaxBlocksPerFragment)); if (count > std::numeric_limits::max() / blocks_per_fragment) return false; - DescriptorStaging* staging = StagingFor(device_id, count); + DescriptorStaging* staging = StagingFor(stream, device_id, count); if (staging == nullptr) return false; std::memcpy(staging->host, fragments, count * sizeof(DeviceGatherFragment)); diff --git a/tests/cpp/umbp/distributed/test_transfer_engine.cpp b/tests/cpp/umbp/distributed/test_transfer_engine.cpp index 54e42aa42..592a520ff 100644 --- a/tests/cpp/umbp/distributed/test_transfer_engine.cpp +++ b/tests/cpp/umbp/distributed/test_transfer_engine.cpp @@ -26,6 +26,7 @@ #include #include +#include #include #include @@ -34,9 +35,11 @@ #include #include #include +#include #include #include "umbp/common/aligned_buffer.h" +#include "umbp/common/device_gather.h" #include "umbp/common/host_registration.h" #include "umbp/distributed/transfer/composite_transfer_engine.h" #include "umbp/distributed/transfer/hbm_copy_engine.h" @@ -292,6 +295,59 @@ TEST(HbmCopyEngine, BuilderSplitsAtHostRegistrationPieces) { ExpectPlanContract(items, planned); } +// Gathers submitted from a succession of short-lived threads, the way the +// standalone server's gRPC pool retires and spawns threads under load, share a +// pooled stream. A stream per thread grew until HIP could not create another. +TEST(HbmCopyEngine, GatherStreamsOutliveTheirThreads) { + int devices = 0; + if (hipGetDeviceCount(&devices) != hipSuccess || devices == 0) GTEST_SKIP() << "no GPU"; + if (!DeviceGatherEnabled()) GTEST_SKIP() << "gather kernel disabled"; + + constexpr size_t kBytes = 1 << 20; + constexpr size_t kFragment = 4096; + constexpr int kThreads = 16; + AlignedBuffer host(kBytes); + for (size_t i = 0; i < kBytes; ++i) host.data()[i] = static_cast(i * 131 + 7); + HbmCopyEngine engine; + engine.AddHostGatherRegion(host.data(), host.size()); + void* device = nullptr; + ASSERT_EQ(hipMalloc(&device, kBytes), hipSuccess); + ASSERT_EQ(hipMemset(device, 0, kBytes), hipSuccess); + const auto src = TransferRef::HostBytes(host.data(), host.size()); + const auto dst = TransferRef::HostBytes(device, kBytes, mori::io::MemoryLocationType::GPU, 0); + + const uint64_t launches_before = DeviceGatherLaunchCount(); + const uint64_t streams_before = HbmCopyEngine::GatherStreamsCreated(); + for (int t = 0; t < kThreads; ++t) { + std::thread([&, t] { + // Two fragments far apart, so they stay two segments: one is hipMemcpy's case. + const uint64_t first = t * 2 * kFragment; + const uint64_t second = kBytes / 2 + first; + auto planned = engine.Plan({MakeItem(src, first, dst, first, kFragment, 1), + MakeItem(src, second, dst, second, kFragment, 2)}); + std::vector failures; + if (auto handle = engine.Submit(std::move(planned.plans))) handle->Wait(&failures); + EXPECT_TRUE(failures.empty()); + }).join(); + } + if (DeviceGatherLaunchCount() == launches_before) { + (void)hipFree(device); + GTEST_SKIP() << "host region not gatherable on this device"; + } + EXPECT_EQ(DeviceGatherLaunchCount(), launches_before + kThreads); + EXPECT_LE(HbmCopyEngine::GatherStreamsCreated() - streams_before, 1u); + + std::vector copied(kBytes); + ASSERT_EQ(hipMemcpy(copied.data(), device, kBytes, hipMemcpyDeviceToHost), hipSuccess); + for (int t = 0; t < kThreads; ++t) { + for (const size_t offset : {t * 2 * kFragment, kBytes / 2 + t * 2 * kFragment}) { + EXPECT_EQ(std::memcmp(copied.data() + offset, host.data() + offset, kFragment), 0) + << "thread " << t << " offset " << offset; + } + } + (void)hipFree(device); +} + // --------------------------------------------------------------------------- // MoriIoEngine (planning only: the engine is never Init()ed) // ---------------------------------------------------------------------------