diff --git a/docs/Doxyfile b/docs/Doxyfile index d17aa6eb..1709aa5c 100644 --- a/docs/Doxyfile +++ b/docs/Doxyfile @@ -2188,6 +2188,7 @@ INCLUDE_FILE_PATTERNS = # This tag requires that the tag ENABLE_PREPROCESSING is set to YES. PREDEFINED = __CUDACC__ \ + MSCCLPP_BULK_AVAILABLE=1 \ MSCCLPP_DEVICE_COMPILE \ MSCCLPP_DEVICE_CUDA \ MSCCLPP_DEVICE_HIP \ diff --git a/docs/cpp_api.rst b/docs/cpp_api.rst index a7ebaaf9..f9bb91c5 100644 --- a/docs/cpp_api.rst +++ b/docs/cpp_api.rst @@ -292,6 +292,27 @@ FIFO Device Interfaces Device Utilities ~~~~~~~~~~~~~~~~ +.. doxygendefine:: MSCCLPP_BULK_AVAILABLE + +.. doxygenstruct:: mscclpp::BulkBarrier + :members: + +.. doxygenenum:: mscclpp::BulkRedOp + +.. doxygenfunction:: mscclpp::bulkLoad + +.. doxygenfunction:: mscclpp::bulkStore + +.. doxygenfunction:: mscclpp::bulkReduceStore + +.. doxygenfunction:: mscclpp::bulkStoreCommit + +.. doxygenfunction:: mscclpp::bulkStoreWait + +.. doxygenfunction:: mscclpp::bulkStoreWaitSource + +.. doxygenfunction:: mscclpp::bulkFence + .. doxygenstruct:: mscclpp::DeviceSemaphore :members: diff --git a/include/mscclpp/bulk_device.hpp b/include/mscclpp/bulk_device.hpp new file mode 100644 index 00000000..1377db38 --- /dev/null +++ b/include/mscclpp/bulk_device.hpp @@ -0,0 +1,278 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +#ifndef MSCCLPP_BULK_DEVICE_HPP_ +#define MSCCLPP_BULK_DEVICE_HPP_ + +#include +#include + +#include "assert_device.hpp" +#include "device.hpp" +#include "poll_device.hpp" + +/// 1 if bulk asynchronous copy is available on the current device compilation target, 0 otherwise. +/// The declarations in this header exist only where this is 1, so call sites must be guarded by it. +#if defined(MSCCLPP_DEVICE_CUDA) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +#define MSCCLPP_BULK_AVAILABLE 1 +#else +#define MSCCLPP_BULK_AVAILABLE 0 +#endif + +// Included whenever the CUDA toolkit is present, not only where bulk copy is available, so that the +// element types accepted by bulkReduceStore() can be named from host code too. +#if defined(MSCCLPP_DEVICE_CUDA) +#include +#endif // defined(MSCCLPP_DEVICE_CUDA) + +namespace mscclpp { + +#if MSCCLPP_BULK_AVAILABLE + +namespace detail { // NOLINT + +template +constexpr bool bulkDependentFalse = false; + +/// Bulk copies require 16-byte aligned addresses and a size that is a multiple of 16 bytes. +MSCCLPP_DEVICE_INLINE bool bulkAligned(const void* ptr) { return (reinterpret_cast(ptr) & 15) == 0; } + +} // namespace detail + +struct BulkBarrier; + +MSCCLPP_DEVICE_INLINE void bulkFence(); +MSCCLPP_DEVICE_INLINE void bulkLoad(void* dstShared, const void* srcGlobal, uint32_t bytes, BulkBarrier& barrier); + +#endif // MSCCLPP_BULK_AVAILABLE + +/// Completion barrier for bulk loads. +/// +/// A bulk load completes asynchronously; its completion is tracked by counting transferred bytes +/// against this barrier. The barrier must live in shared memory and alternates between two phases, +/// so one barrier serves an unbounded number of batches without re-initialization. +/// +/// Per batch the caller declares the bytes to wait for with expect(), signals participation with +/// arrive(), and blocks with wait(); arriveAndExpect() fuses the common single-issuer case. The +/// phase starts at 0 and wait() flips it, so the caller keeps the phase in a register and passes it +/// back on the next batch. +/// +/// @warning Every expect() for a batch must be issued before the arrival that completes the arrival +/// count. Arriving while the outstanding byte count is momentarily zero completes the phase +/// immediately, and wait() then returns over tiles that have not been filled. Issuing a single +/// arriveAndExpect() carrying the batch total avoids this entirely and is the recommended form. +/// +/// The storage is declared on every build target, including host compilation, so that host code can +/// size the shared memory that holds barriers. The operations exist only where MSCCLPP_BULK_AVAILABLE +/// is 1. +struct BulkBarrier { +#if MSCCLPP_BULK_AVAILABLE + /// Initialize the barrier and publish it to the asynchronous proxy. Called by a single thread + /// before any other method, and not while a bulk load tracked by this barrier is in flight. + /// @param arriveCount Number of arrivals that complete a phase. + MSCCLPP_DEVICE_INLINE void init(uint32_t arriveCount = 1) { + relaxedInit(arriveCount); + bulkFence(); + } + + /// Initialize the barrier without publishing it to the asynchronous proxy. + /// + /// The initialization is not visible to bulk copies until a bulkFence() runs on the initializing + /// thread. Use this to initialize several barriers under a single fence; otherwise use init(). + /// @param arriveCount Number of arrivals that complete a phase. + MSCCLPP_DEVICE_INLINE void relaxedInit(uint32_t arriveCount = 1) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(addr()), "r"(arriveCount)); + } + + /// Invalidate the barrier, releasing the underlying hardware state. Called by a single thread + /// before the shared memory holding this barrier is reused for another purpose. + MSCCLPP_DEVICE_INLINE void invalidate() { asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"(addr())); } + + /// Add to the bytes the current phase waits for, without arriving. + /// @param bytes Bytes expected to arrive. + MSCCLPP_DEVICE_INLINE void expect(uint32_t bytes) { + asm volatile("mbarrier.expect_tx.shared::cta.b64 [%0], %1;" ::"r"(addr()), "r"(bytes)); + } + + /// Signal one arrival on the current phase. + MSCCLPP_DEVICE_INLINE void arrive() { asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(addr())); } + + /// Signal one arrival and add to the bytes the current phase waits for. Equivalent to expect() + /// followed by arrive(), issued as a single instruction. + /// @param bytes Bytes expected to arrive. + MSCCLPP_DEVICE_INLINE void arriveAndExpect(uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(addr()), "r"(bytes)); + } + + /// Check whether the given phase has completed, without blocking or advancing the phase. + /// @param phase The phase to check. + /// @return True if the phase has completed. + MSCCLPP_DEVICE_INLINE bool poll(uint32_t phase) { + uint32_t done; + asm volatile("{.reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2; selp.u32 %0, 1, 0, p;}" + : "=r"(done) + : "r"(addr()), "r"(phase & 1u)); + return done != 0; + } + + /// Wait for the given phase to complete and advance @p phase to the next one. + /// + /// Does not order the delivered data against generic shared memory reads; run bulkFence() on each + /// consuming thread, or a thread barrier that orders it against one that did, before reading the + /// loaded tiles. + /// + /// @param phase The phase to wait for, flipped on return. Starts at 0 and is kept by the caller. + /// @param maxSpinCount The maximum number of spins before asserting. Never assert if negative. + MSCCLPP_DEVICE_INLINE void wait(uint32_t& phase, int64_t maxSpinCount = 10000000) { + POLL_MAYBE_JAILBREAK(!poll(phase), maxSpinCount); + phase ^= 1u; + } + + private: + friend MSCCLPP_DEVICE_INLINE void bulkLoad(void* dstShared, const void* srcGlobal, uint32_t bytes, + BulkBarrier& barrier); + + MSCCLPP_DEVICE_INLINE uint32_t addr() const { return static_cast(__cvta_generic_to_shared(&mbar_)); } +#endif // MSCCLPP_BULK_AVAILABLE + + private: + // Marked maybe_unused because the operations that read it exist only where MSCCLPP_BULK_AVAILABLE + // is 1, while the storage is declared everywhere so host code can size shared memory. + [[maybe_unused]] alignas(8) uint64_t mbar_; +}; + +#if MSCCLPP_BULK_AVAILABLE + +/// Issue an asynchronous bulk load from global memory to shared memory. +/// +/// Returns immediately; completion is tracked by @p barrier, which must expect at least @p bytes for +/// the phase being waited on. Issued by a single thread. Loaded data is visible to generic shared +/// memory reads only after the barrier's phase completes and bulkFence() has run. +/// +/// @param dstShared Destination in shared memory. 16-byte aligned; 128 bytes is faster. +/// @param srcGlobal Source in global memory, local or peer-mapped. 16-byte aligned. +/// @param bytes Bytes to load. A multiple of 16. +/// @param barrier Barrier tracking completion of this load. +MSCCLPP_DEVICE_INLINE void bulkLoad(void* dstShared, const void* srcGlobal, uint32_t bytes, BulkBarrier& barrier) { + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(dstShared), "bulkLoad destination is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(srcGlobal), "bulkLoad source is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE((bytes & 15) == 0, "bulkLoad size is not a multiple of 16 bytes"); + asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(dstShared))), + "l"(srcGlobal), "r"(bytes), "r"(barrier.addr()) + : "memory"); +} + +/// Issue an asynchronous bulk store from shared memory to global memory. +/// +/// Returns immediately; the store joins the calling thread's open bulk group, which is closed by +/// bulkStoreCommit() and drained by bulkStoreWait() or bulkStoreWaitSource(). Issued by a single +/// thread. Data written to @p srcShared by generic shared memory writes is visible to the store only +/// after bulkFence(). +/// +/// @param dstGlobal Destination in global memory, local or peer-mapped. 16-byte aligned. +/// @param srcShared Source in shared memory. 16-byte aligned. +/// @param bytes Bytes to store. A multiple of 16. +MSCCLPP_DEVICE_INLINE void bulkStore(void* dstGlobal, const void* srcShared, uint32_t bytes) { + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(dstGlobal), "bulkStore destination is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(srcShared), "bulkStore source is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE((bytes & 15) == 0, "bulkStore size is not a multiple of 16 bytes"); + asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(static_cast(__cvta_generic_to_shared(srcShared))), "r"(bytes) + : "memory"); +} + +/// Reduction operation applied by bulkReduceStore(). +enum class BulkRedOp { Add }; + +/// Issue an asynchronous bulk reduction from shared memory into global memory. +/// +/// Accumulates @p srcShared into @p dstGlobal elementwise, in the copy engine rather than on the +/// SM, so the destination is never read back across the interconnect. Works on peer-mapped +/// destinations, which makes it a remote accumulate. Otherwise behaves exactly like bulkStore(): +/// returns immediately, joins the calling thread's open bulk group, and requires a preceding +/// bulkFence() to make generic writes to @p srcShared visible. +/// +/// Measured on H200 at roughly 90% of the bulkStore() rate for the same payload from a single +/// issuing thread, a roughly 10% throughput cost in that microbenchmark. The cost visible to a +/// workload depends on its transfer overlap and bottlenecks. +/// +/// PTX ISA 9.3 specifies each element-wise reduction as atomic and defaults an omitted scope to +/// `.relaxed.sys`. Earlier PTX documentation, including that shipped with CUDA 12.9 and 13.0, +/// specifies `.relaxed.gpu` semantics and does not guarantee atomicity across devices. Because this +/// API supports those toolkits, concurrent reductions from several devices into the same peer address +/// must be treated as empirical rather than portable; exact results were observed on H200 and GB200, +/// but callers must confirm the behavior on their target toolchain and hardware. +/// +/// @tparam T Source and destination element type. Currently `float`, `__nv_bfloat16`, and `uint32_t`. +/// The instruction does not perform mixed-precision conversion: `__nv_bfloat16` accumulates in bf16; +/// for fp32 accumulation, convert the source tile and use a float destination. Other types are +/// rejected at compile time rather than silently mapped, because the underlying instruction accepts +/// only certain operation and type combinations. +/// @tparam Op Reduction operation. +/// @param dstGlobal Destination in global memory, local or peer-mapped. 16-byte aligned. +/// @param srcShared Source in shared memory. 16-byte aligned. +/// @param bytes Bytes to reduce. A multiple of 16. +template +MSCCLPP_DEVICE_INLINE void bulkReduceStore(void* dstGlobal, const void* srcShared, uint32_t bytes) { + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(dstGlobal), "bulkReduceStore destination is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(srcShared), "bulkReduceStore source is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE((bytes & 15) == 0, "bulkReduceStore size is not a multiple of 16 bytes"); + const uint32_t src = static_cast(__cvta_generic_to_shared(srcShared)); + if constexpr (Op == BulkRedOp::Add && std::is_same_v) { + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(src), "r"(bytes) + : "memory"); + } else if constexpr (Op == BulkRedOp::Add && std::is_same_v) { + // The instruction requires an explicit .noftz for bf16 addition. + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(src), "r"(bytes) + : "memory"); + } else if constexpr (Op == BulkRedOp::Add && std::is_same_v) { + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.u32 [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(src), "r"(bytes) + : "memory"); + } else { + static_assert(detail::bulkDependentFalse, "Unsupported bulk reduction type or operation"); + } +} + +/// Close the calling thread's open bulk group, committing every bulkStore() and bulkReduceStore() +/// issued since the previous bulkStoreCommit(). +MSCCLPP_DEVICE_INLINE void bulkStoreCommit() { asm volatile("cp.async.bulk.commit_group;" ::: "memory"); } + +/// Wait until at most @p PendingGroups committed bulk groups have yet to complete. +/// +/// Completion means the data has reached the destination. Wait for this before signaling a peer that +/// the data is available. To reuse the source shared memory only, bulkStoreWaitSource() is cheaper. +/// @tparam PendingGroups Number of most recently committed groups allowed to remain outstanding. +template +MSCCLPP_DEVICE_INLINE void bulkStoreWait() { + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(PendingGroups) : "memory"); +} + +/// Wait until at most @p PendingGroups committed bulk groups may still read their source. +/// +/// Weaker than bulkStoreWait(): it guarantees only that the source shared memory can be overwritten, +/// not that the data has reached the destination. This is what a double-buffered store pipeline +/// needs between refills, and it lets stores stay in flight while the next tile is being staged. +/// @tparam PendingGroups Number of most recently committed groups allowed to still read their source. +template +MSCCLPP_DEVICE_INLINE void bulkStoreWaitSource() { + asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(PendingGroups) : "memory"); +} + +/// Order the calling thread's generic shared memory accesses against bulk copies. +/// +/// Bulk copies reach shared memory through a separate proxy, so they are not ordered against generic +/// accesses by default. Run this after a bulk load completes and before reading its destination, +/// after writing a source and before issuing a bulk store from it, and after initializing a barrier +/// with relaxedInit(). Every thread that touches the shared memory must run it, unless a thread +/// barrier such as __syncthreads() already orders that thread against one that did. +MSCCLPP_DEVICE_INLINE void bulkFence() { asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); } + +#endif // MSCCLPP_BULK_AVAILABLE + +} // namespace mscclpp + +#endif // MSCCLPP_BULK_DEVICE_HPP_ diff --git a/include/mscclpp/gpu_utils.hpp b/include/mscclpp/gpu_utils.hpp index 82fa3ec0..b705f883 100644 --- a/include/mscclpp/gpu_utils.hpp +++ b/include/mscclpp/gpu_utils.hpp @@ -312,6 +312,11 @@ inline void gpuMemset(void* ptr, int value, size_t bytes) { detail::gpuMemset(pt /// @return True if NVLink SHARP (NVLS) is supported, false otherwise. bool isNvlsSupported(); +/// Check if bulk asynchronous copy is supported by the current device. +/// +/// @return True if bulk asynchronous copy is supported, false otherwise. +bool isBulkSupported(); + /// Check if ptr is allocaed by cuMemMap. /// @param ptr The pointer to check. /// @return True if the pointer is allocated by cuMemMap, false otherwise. diff --git a/python/csrc/gpu_utils_py.cpp b/python/csrc/gpu_utils_py.cpp index d6527502..c05e5cc3 100644 --- a/python/csrc/gpu_utils_py.cpp +++ b/python/csrc/gpu_utils_py.cpp @@ -113,6 +113,7 @@ static nb::capsule toDlpack(GpuBuffer buffer, std::string dataType, std::v void register_gpu_utils(nb::module_& m) { m.def("is_nvls_supported", &isNvlsSupported); + m.def("is_bulk_supported", &isBulkSupported); nb::enum_(m, "CppGpuBufferGranularity") .value("MultiCastMinimum", GpuBufferGranularity::MultiCastMinimum) diff --git a/python/mscclpp/__init__.py b/python/mscclpp/__init__.py index 09408171..55004eeb 100644 --- a/python/mscclpp/__init__.py +++ b/python/mscclpp/__init__.py @@ -52,6 +52,7 @@ CppRawGpuBuffer as RawGpuBuffer, CppReduceOp as ReduceOp, env, + is_bulk_supported, is_nvls_supported, cpp_npkit as npkit, ) @@ -87,6 +88,7 @@ "ReduceOp", "env", "version", + "is_bulk_supported", "is_nvls_supported", "alloc_shared_physical_cuda", "npkit", diff --git a/src/core/gpu_utils.cc b/src/core/gpu_utils.cc index 7ea838c6..cd780a50 100644 --- a/src/core/gpu_utils.cc +++ b/src/core/gpu_utils.cc @@ -303,6 +303,23 @@ bool isNvlsSupported() { return false; } +bool isBulkSupported() { + [[maybe_unused]] static bool result = false; + [[maybe_unused]] static bool isChecked = false; +#if !defined(MSCCLPP_USE_ROCM) + if (!isChecked) { + int deviceId; + int major; + MSCCLPP_CUDATHROW(cudaGetDevice(&deviceId)); + MSCCLPP_CUDATHROW(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, deviceId)); + result = (major >= 9); + isChecked = true; + } + return result; +#endif + return false; +} + bool isCuMemMapAllocated([[maybe_unused]] void* ptr) { #if defined(MSCCLPP_USE_ROCM) return false; diff --git a/test/mp_unit/CMakeLists.txt b/test/mp_unit/CMakeLists.txt index 7d64145d..ff85a250 100644 --- a/test/mp_unit/CMakeLists.txt +++ b/test/mp_unit/CMakeLists.txt @@ -8,6 +8,7 @@ target_sources(mp_unit_tests PRIVATE communicator_tests.cu port_channel_tests.cu memory_channel_tests.cu + bulk_pattern_tests.cu semaphore_perf_tests.cu switch_channel_tests.cu switch_channel_perf_tests.cu diff --git a/test/mp_unit/bulk_pattern_tests.cu b/test/mp_unit/bulk_pattern_tests.cu new file mode 100644 index 00000000..2fc276e4 --- /dev/null +++ b/test/mp_unit/bulk_pattern_tests.cu @@ -0,0 +1,425 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// Bulk-copy usage patterns, shaped after the expert-parallel dispatch and combine kernels. +// +// These are deliberately simplified: fixed token counts, one contribution per rank pair, and plain +// arithmetic instead of quantization or top-k routing. What they preserve is the communication +// shape, which is what the API has to support: +// +// Dispatch local buffer -> shared (transform) -> peer buffer. Push, staged through shared +// memory, pipelined so a tile is refilled while its store is still in flight. +// Combine every peer's buffer -> shared -> reduce -> local output. Pull, many sources into +// one barrier, double-buffered across chunks. +// CombineReduce the same reduction expressed as a push, with the copy engine accumulating +// directly into the peer's memory, so nothing is ever read back. +// +// Note where the channels are and are not used. Bulk data moves through raw peer pointers, exactly +// as the expert-parallel kernels do; the channels carry only signal/wait. That split is the reason +// the bulk primitives are pointer-based rather than channel methods. + +#include +#include +#include +#include +#include + +#include "mp_unit_tests.hpp" + +void BulkPatternTest::SetUp() { + if (gEnv->nRanksPerNode < 2) { + SKIP_TEST(); + } + setNumRanksToUse(gEnv->nRanksPerNode); + CommunicatorTestBase::SetUp(); + rank = gEnv->rank; + worldSize = numRanksToUse; +} + +void BulkPatternTest::TearDown() { + syncHandles.clear(); + syncChannels.clear(); + peerPools.clear(); + remotePoolMemories.clear(); + CommunicatorTestBase::TearDown(); +} + +void BulkPatternTest::setupPeerPools(void* pool, size_t poolBytes) { + const mscclpp::TransportFlags transport = mscclpp::Transport::CudaIpc; + mscclpp::RegisteredMemory localMem = communicator->registerMemory(pool, poolBytes, transport); + + std::vector> connFutures(worldSize); + std::vector> memFutures(worldSize); + for (int r = 0; r < worldSize; ++r) { + if (r == rank) continue; + connFutures[r] = communicator->connect(mscclpp::Transport::CudaIpc, r); + communicator->sendMemory(localMem, r); + memFutures[r] = communicator->recvMemory(r); + } + + peerPools.assign(worldSize, nullptr); + peerPools[rank] = pool; + for (int r = 0; r < worldSize; ++r) { + if (r == rank) continue; + mscclpp::RegisteredMemory remote = memFutures[r].get(); + peerPools[r] = remote.data(); + remotePoolMemories.push_back(remote); + syncChannels.emplace_back(communicator->buildSemaphore(connFutures[r].get(), r).get()); + } + for (const auto& chan : syncChannels) syncHandles.push_back(mscclpp::deviceHandle(chan)); + registeredMemories.push_back(localMem); +} + +#if defined(MSCCLPP_DEVICE_CUDA) + +namespace { + +constexpr int kNumTokens = 16; +constexpr int kHidden = 1024; // floats per token +constexpr uint32_t kTokenBytes = kHidden * sizeof(float); // 4 KB +constexpr uint32_t kChunkBytes = 1024; // staging granularity +constexpr int kChunksPerToken = kTokenBytes / kChunkBytes; +constexpr int kChunkFloats = kChunkBytes / sizeof(float); +constexpr int kStages = 2; + +// Distinct, exactly representable value for (source rank, destination rank, token, element). +MSCCLPP_HOST_DEVICE_INLINE float payload(int src, int dst, int token, int elem) { + return (float)(src * 1000 + dst * 100 + token * 10 + (elem % 10)); +} + +// Offset in floats of the [src][token] row within a receive pool. +MSCCLPP_HOST_DEVICE_INLINE int64_t rowOffset(int src, int token) { + return ((int64_t)src * kNumTokens + token) * kHidden; +} + +} // namespace + +// Signal every peer, then wait on every peer. +// +// Launched as its own kernel rather than folded into the pattern kernels: the ranks have to agree on +// a point that every block has passed, and kernel completion already provides exactly that. Doing it +// inside a kernel would need a cooperative grid launch. +__global__ void kernelPeerBarrier(mscclpp::BaseMemoryChannelDeviceHandle* chans, int nPeers) { + if ((int)threadIdx.x < nPeers) { + chans[threadIdx.x].signal(); + chans[threadIdx.x].wait(); + } +} + +namespace { + +// Peer pool pointers and synchronization handles, staged in device memory for the kernels. +struct DeviceState { + std::shared_ptr> chans; + std::shared_ptr pools; + int nPeers; +}; + +DeviceState uploadPeerState(const std::vector>& handles, + const std::vector& pools) { + DeviceState state; + state.nPeers = (int)handles.size(); + state.chans = mscclpp::detail::gpuCallocShared>(handles.size()); + MSCCLPP_CUDATHROW(cudaMemcpy(state.chans.get(), handles.data(), + handles.size() * sizeof(DeviceHandle), + cudaMemcpyHostToDevice)); + state.pools = mscclpp::detail::gpuCallocShared(pools.size()); + MSCCLPP_CUDATHROW(cudaMemcpy(state.pools.get(), pools.data(), pools.size() * sizeof(void*), cudaMemcpyHostToDevice)); + return state; +} + +} // namespace + +// ------------------------------------------------------------------------------------------------ +// Dispatch: push each local token to every peer, staged and transformed in shared memory. +// +// Per chunk the leader loads into a tile, the block transforms it in place, and the leader pushes it +// out. bulkStoreWaitSource() lets the next load start while the store is still travelling, which is +// the reason that entry point exists separately from bulkStoreWait(). +// ------------------------------------------------------------------------------------------------ +__global__ void kernelDispatch([[maybe_unused]] const float* localTokens, [[maybe_unused]] void** peerPools, + [[maybe_unused]] int rank, [[maybe_unused]] int worldSize) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) uint8_t tile[kChunkBytes]; + __shared__ mscclpp::BulkBarrier barrier; + + if (threadIdx.x == 0) barrier.init(); + __syncthreads(); + + uint32_t phase = 0; // initialized once; wait() advances it on every chunk + bool storePending = false; + + // One block per (destination, token) pair, strided over the grid. + const int totalRows = worldSize * kNumTokens; + for (int row = blockIdx.x; row < totalRows; row += gridDim.x) { + const int dst = row / kNumTokens; + const int token = row % kNumTokens; + const float* srcRow = localTokens + (int64_t)token * kHidden; + auto* dstRow = reinterpret_cast(peerPools[dst]) + rowOffset(rank, token) * sizeof(float); + + for (int c = 0; c < kChunksPerToken; ++c) { + const uint32_t off = c * kChunkBytes; + if (threadIdx.x == 0) { + // The previous chunk's store has released the tile by now; make sure before refilling. + if (storePending) mscclpp::bulkStoreWaitSource<0>(); + barrier.arriveAndExpect(kChunkBytes); + mscclpp::bulkLoad(tile, reinterpret_cast(srcRow) + off, kChunkBytes, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); // loaded data -> generic reads + + // Stand-in for the quantize/scale step a real dispatch performs while the token is staged. + float* staged = reinterpret_cast(tile); + for (int i = threadIdx.x; i < kChunkFloats; i += blockDim.x) staged[i] += (float)(dst * 100); + + __syncthreads(); + mscclpp::bulkFence(); // generic writes -> the store's read of the tile + if (threadIdx.x == 0) { + mscclpp::bulkStore(dstRow + off, tile, kChunkBytes); + mscclpp::bulkStoreCommit(); + storePending = true; + } + __syncthreads(); + } + } + + // Every push must have landed before any peer is told the data is ready. + if (threadIdx.x == 0) { + mscclpp::bulkStoreWait<0>(); + barrier.invalidate(); + } +#endif +} + +TEST(BulkPatternTest, Dispatch) { + if (gEnv->rank >= numRanksToUse) return; + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + + const size_t poolFloats = (size_t)worldSize * kNumTokens * kHidden; + std::shared_ptr pool = mscclpp::GpuBuffer(poolFloats).memory(); + std::shared_ptr tokens = mscclpp::GpuBuffer((size_t)kNumTokens * kHidden).memory(); + MSCCLPP_CUDATHROW(cudaMemset(pool.get(), 0, poolFloats * sizeof(float))); + + // Local tokens carry payload(rank, 0, token, elem); the kernel adds dst * 100 while staged, so the + // value that lands on rank dst is payload(rank, dst, token, elem). + std::vector host((size_t)kNumTokens * kHidden); + for (int t = 0; t < kNumTokens; ++t) + for (int h = 0; h < kHidden; ++h) host[(size_t)t * kHidden + h] = payload(gEnv->rank, 0, t, h); + MSCCLPP_CUDATHROW(cudaMemcpy(tokens.get(), host.data(), host.size() * sizeof(float), cudaMemcpyHostToDevice)); + + setupPeerPools(pool.get(), poolFloats * sizeof(float)); + + DeviceState dev = uploadPeerState(syncHandles, peerPools); + + kernelDispatch<<>>(tokens.get(), dev.pools.get(), rank, worldSize); + // Kernel completion is the grid-wide point after which every push has landed; only then is it safe + // to tell the peers, and to read what they pushed here. + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out(poolFloats); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), pool.get(), poolFloats * sizeof(float), cudaMemcpyDeviceToHost)); + for (int src = 0; src < worldSize; ++src) { + for (int t = 0; t < kNumTokens; ++t) { + for (int h = 0; h < kHidden; h += 97) { // sparse check keeps the assertion count sane + EXPECT_EQ(out[rowOffset(src, t) + h], payload(src, rank, t, h)); + } + } + } +} + +// ------------------------------------------------------------------------------------------------ +// Combine: pull this rank's row from every peer into shared memory and reduce it. +// +// This is the pattern the barrier design exists for. All contributors of a chunk are issued against +// one barrier and waited on once, and two stages let the next chunk's loads be issued before the +// current one is reduced. The two barriers are set up with relaxedInit() under a single bulkFence() +// rather than paying a fence each. +// ------------------------------------------------------------------------------------------------ +__global__ void kernelCombine([[maybe_unused]] float* output, [[maybe_unused]] void** peerPools, + [[maybe_unused]] int rank, [[maybe_unused]] int worldSize) { +#if MSCCLPP_BULK_AVAILABLE + extern __shared__ __align__(128) uint8_t shared[]; + // Layout: [stage][contributor][chunk] tiles, then the per-stage barriers. + uint8_t* tiles = shared; + auto* barriers = reinterpret_cast(shared + (size_t)kStages * worldSize * kChunkBytes); + auto tile = [&](int stage, int contributor) { + return tiles + ((size_t)stage * worldSize + contributor) * kChunkBytes; + }; + + if (threadIdx.x == 0) { + for (int s = 0; s < kStages; ++s) barriers[s].relaxedInit(); + mscclpp::bulkFence(); // one fence publishes both barriers to the async proxy + } + __syncthreads(); + + uint32_t phase[kStages] = {0, 0}; + + auto issue = [&](int stage, int token, int chunk) { + barriers[stage].arriveAndExpect(kChunkBytes * worldSize); // whole batch, before any arrival lands + for (int src = 0; src < worldSize; ++src) { + const auto* row = reinterpret_cast(peerPools[src]) + rowOffset(rank, token) * sizeof(float); + mscclpp::bulkLoad(tile(stage, src), row + (size_t)chunk * kChunkBytes, kChunkBytes, barriers[stage]); + } + }; + + for (int token = blockIdx.x; token < kNumTokens; token += gridDim.x) { + if (threadIdx.x == 0) issue(0, token, 0); + for (int c = 0; c < kChunksPerToken; ++c) { + const int stage = c % kStages; + if (threadIdx.x == 0) { + if (c + 1 < kChunksPerToken) issue((c + 1) % kStages, token, c + 1); // prefetch next chunk + barriers[stage].wait(phase[stage]); + } + __syncthreads(); + mscclpp::bulkFence(); + + float* out = output + (int64_t)token * kHidden + (int64_t)c * kChunkFloats; + for (int i = threadIdx.x; i < kChunkFloats; i += blockDim.x) { + float acc = 0.0f; + for (int src = 0; src < worldSize; ++src) acc += reinterpret_cast(tile(stage, src))[i]; + out[i] = acc; + } + __syncthreads(); // this stage's tiles are consumed before they are reissued + } + } + + if (threadIdx.x == 0) { + for (int s = 0; s < kStages; ++s) barriers[s].invalidate(); + } +#endif +} + +TEST(BulkPatternTest, Combine) { + if (gEnv->rank >= numRanksToUse) return; + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + + const size_t poolFloats = (size_t)worldSize * kNumTokens * kHidden; + std::shared_ptr pool = mscclpp::GpuBuffer(poolFloats).memory(); + std::shared_ptr output = mscclpp::GpuBuffer((size_t)kNumTokens * kHidden).memory(); + + // Each rank publishes, for every destination, the contribution that destination will pull. + std::vector host(poolFloats); + for (int dst = 0; dst < worldSize; ++dst) + for (int t = 0; t < kNumTokens; ++t) + for (int h = 0; h < kHidden; ++h) host[rowOffset(dst, t) + h] = payload(gEnv->rank, dst, t, h); + MSCCLPP_CUDATHROW(cudaMemcpy(pool.get(), host.data(), host.size() * sizeof(float), cudaMemcpyHostToDevice)); + MSCCLPP_CUDATHROW(cudaMemset(output.get(), 0, (size_t)kNumTokens * kHidden * sizeof(float))); + + setupPeerPools(pool.get(), poolFloats * sizeof(float)); + + DeviceState dev = uploadPeerState(syncHandles, peerPools); + + const size_t sharedBytes = (size_t)kStages * worldSize * kChunkBytes + kStages * sizeof(mscclpp::BulkBarrier); + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // every pool is filled before any pull + kernelCombine<<>>(output.get(), dev.pools.get(), rank, worldSize); + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // no rank tears down while peers read + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out((size_t)kNumTokens * kHidden); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), output.get(), out.size() * sizeof(float), cudaMemcpyDeviceToHost)); + for (int t = 0; t < kNumTokens; ++t) { + for (int h = 0; h < kHidden; h += 97) { + float expected = 0.0f; + for (int src = 0; src < worldSize; ++src) expected += payload(src, rank, t, h); + EXPECT_EQ(out[(size_t)t * kHidden + h], expected); + } + } +} + +// ------------------------------------------------------------------------------------------------ +// CombineReduce: the same reduction, expressed as a push. +// +// Instead of every rank pulling and summing, each rank accumulates its own contribution straight +// into the destination's accumulator with bulkReduceStore(). The copy engine performs the add at the +// destination, so the accumulator is never read back across the link and no rank stages anyone +// else's data. Cross-device atomicity is guaranteed only by PTX ISA 9.3's system-scope semantics; +// with earlier supported toolkits this pattern exercises empirically observed behavior rather than a +// portable guarantee. The reverse link direction stays idle. +// ------------------------------------------------------------------------------------------------ +__global__ void kernelCombineReduce([[maybe_unused]] const float* contribution, [[maybe_unused]] void** peerPools, + [[maybe_unused]] int rank, [[maybe_unused]] int worldSize) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) uint8_t tile[kChunkBytes]; + __shared__ mscclpp::BulkBarrier barrier; + + if (threadIdx.x == 0) barrier.init(); + __syncthreads(); + uint32_t phase = 0; + + const int totalRows = worldSize * kNumTokens; + for (int row = blockIdx.x; row < totalRows; row += gridDim.x) { + const int dst = row / kNumTokens; + const int token = row % kNumTokens; + const auto* srcRow = reinterpret_cast(contribution + rowOffset(dst, token)); + // Every rank accumulates into the same [rank-independent] row of the destination. + auto* dstRow = reinterpret_cast(peerPools[dst]) + rowOffset(0, token) * sizeof(float); + + for (int c = 0; c < kChunksPerToken; ++c) { + const uint32_t off = c * kChunkBytes; + if (threadIdx.x == 0) { + barrier.arriveAndExpect(kChunkBytes); + mscclpp::bulkLoad(tile, srcRow + off, kChunkBytes, barrier); + barrier.wait(phase); + mscclpp::bulkFence(); // staged data -> the reduction's read of the tile + mscclpp::bulkReduceStore(dstRow + off, tile, kChunkBytes); + mscclpp::bulkStoreCommit(); + mscclpp::bulkStoreWaitSource<0>(); // the tile can be refilled once its source read completes + } + __syncthreads(); + } + } + + if (threadIdx.x == 0) { + mscclpp::bulkStoreWait<0>(); // every accumulate must land before the kernel completes + barrier.invalidate(); + } +#endif +} + +TEST(BulkPatternTest, CombineReduce) { + if (gEnv->rank >= numRanksToUse) return; + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + + const size_t poolFloats = (size_t)worldSize * kNumTokens * kHidden; + std::shared_ptr pool = mscclpp::GpuBuffer(poolFloats).memory(); + std::shared_ptr contribution = mscclpp::GpuBuffer(poolFloats).memory(); + MSCCLPP_CUDATHROW(cudaMemset(pool.get(), 0, poolFloats * sizeof(float))); // accumulator starts at zero + + std::vector host(poolFloats); + for (int dst = 0; dst < worldSize; ++dst) + for (int t = 0; t < kNumTokens; ++t) + for (int h = 0; h < kHidden; ++h) host[rowOffset(dst, t) + h] = payload(gEnv->rank, dst, t, h); + MSCCLPP_CUDATHROW(cudaMemcpy(contribution.get(), host.data(), host.size() * sizeof(float), cudaMemcpyHostToDevice)); + + setupPeerPools(pool.get(), poolFloats * sizeof(float)); + + DeviceState dev = uploadPeerState(syncHandles, peerPools); + + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // every accumulator is zeroed first + kernelCombineReduce<<>>(contribution.get(), dev.pools.get(), rank, worldSize); + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // every contribution has landed + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out(poolFloats); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), pool.get(), poolFloats * sizeof(float), cudaMemcpyDeviceToHost)); + for (int t = 0; t < kNumTokens; ++t) { + for (int h = 0; h < kHidden; h += 97) { + float expected = 0.0f; + for (int src = 0; src < worldSize; ++src) expected += payload(src, rank, t, h); + EXPECT_EQ(out[rowOffset(0, t) + h], expected); + } + } +} + +#endif // defined(MSCCLPP_DEVICE_CUDA) diff --git a/test/mp_unit/mp_unit_tests.hpp b/test/mp_unit/mp_unit_tests.hpp index c5a42b5f..8654ccc9 100644 --- a/test/mp_unit/mp_unit_tests.hpp +++ b/test/mp_unit/mp_unit_tests.hpp @@ -187,6 +187,30 @@ class MemoryChannelOneToOneTest : public CommunicatorTestBase { std::unordered_map> memorySemaphores; }; +/// Fixture for the bulk-copy usage patterns, shaped after the expert-parallel dispatch and combine +/// kernels: a full mesh where every rank holds a receive pool addressable by every peer, channels +/// used only for synchronization, and the peer pool pointers handed to the kernel as a plain array. +class BulkPatternTest : public CommunicatorTestBase { + protected: + void SetUp() override; + void TearDown() override; + + /// Register @p pool on every rank, exchange it around a full mesh, and build one synchronization + /// channel per peer. Fills peerPools and syncHandles. + void setupPeerPools(void* pool, size_t poolBytes); + + int worldSize = 0; + int rank = 0; + /// Receive pool base address of each rank, indexed by rank. The local entry is the local pool. + /// This mirrors the expert-parallel kernels, which address peers by raw pointer and never route + /// bulk data through a channel. + std::vector peerPools; + /// One channel per peer, used only for signal/wait. + std::vector syncChannels; + std::vector> syncHandles; + std::vector remotePoolMemories; +}; + class SemaphorePerfTest : public CommunicatorTestBase { protected: void SetUp() override; diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt index 2547278b..e6d09d2d 100644 --- a/test/unit/CMakeLists.txt +++ b/test/unit/CMakeLists.txt @@ -17,4 +17,5 @@ target_sources(unit_tests PRIVATE gpu_data_types_tests.cu local_channel_tests.cu gpu_ipc_mem_tests.cc + bulk_tests.cu ) diff --git a/test/unit/bulk_tests.cu b/test/unit/bulk_tests.cu new file mode 100644 index 00000000..73c92a82 --- /dev/null +++ b/test/unit/bulk_tests.cu @@ -0,0 +1,233 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#include +#include +#include +#include + +#include "../framework.hpp" + +#if defined(MSCCLPP_DEVICE_CUDA) + +constexpr uint32_t kTile = 4096; +constexpr uint32_t kElems = kTile / sizeof(int); + +// Load one tile into shared memory, then copy it out. +__global__ void kernelBulkLoad([[maybe_unused]] const int* src, [[maybe_unused]] int* dst) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) int tile[kElems]; + __shared__ mscclpp::BulkBarrier barrier; + + uint32_t phase = 0; + if (threadIdx.x == 0) { + barrier.init(); + barrier.arriveAndExpect(kTile); + mscclpp::bulkLoad(tile, src, kTile, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); + + for (uint32_t i = threadIdx.x; i < kElems; i += blockDim.x) dst[i] = tile[i]; + + __syncthreads(); + if (threadIdx.x == 0) barrier.invalidate(); +#endif +} + +// Gather NumSrc tiles into one barrier, then reduce them. +template +__global__ void kernelBulkGather([[maybe_unused]] const int* src, [[maybe_unused]] int* dst) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) int tiles[NumSrc][kElems]; + __shared__ mscclpp::BulkBarrier barrier; + + uint32_t phase = 0; + if (threadIdx.x == 0) { + barrier.init(); + barrier.arriveAndExpect(kTile * NumSrc); // whole batch total, before any arrival completes + for (int s = 0; s < NumSrc; ++s) mscclpp::bulkLoad(tiles[s], src + s * kElems, kTile, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); + + for (uint32_t i = threadIdx.x; i < kElems; i += blockDim.x) { + int acc = 0; + for (int s = 0; s < NumSrc; ++s) acc += tiles[s][i]; + dst[i] = acc; + } + + __syncthreads(); + if (threadIdx.x == 0) barrier.invalidate(); +#endif +} + +// Reuse one barrier across NumChunks phases, staging each chunk in and storing it back out. +template +__global__ void kernelBulkRoundTrip([[maybe_unused]] const int* src, [[maybe_unused]] int* dst) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) int tile[kElems]; + __shared__ mscclpp::BulkBarrier barrier; + + if (threadIdx.x == 0) barrier.init(); + __syncthreads(); + + uint32_t phase = 0; // one barrier, no re-initialization; wait() advances the phase + for (int c = 0; c < NumChunks; ++c) { + if (threadIdx.x == 0) { + barrier.arriveAndExpect(kTile); + mscclpp::bulkLoad(tile, src + c * kElems, kTile, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); + + for (uint32_t i = threadIdx.x; i < kElems; i += blockDim.x) tile[i] += 1; + + __syncthreads(); + mscclpp::bulkFence(); + if (threadIdx.x == 0) { + mscclpp::bulkStore(dst + c * kElems, tile, kTile); + mscclpp::bulkStoreCommit(); + mscclpp::bulkStoreWaitSource(); // tile refillable; the store may still be in flight + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + mscclpp::bulkStoreWait(); // every store has landed + barrier.invalidate(); + } +#endif +} + +// Accumulate a staged tile into a seeded destination with the copy engine. +template +__global__ void kernelBulkReduce([[maybe_unused]] T* dst, [[maybe_unused]] T addend, [[maybe_unused]] uint32_t count) { +#if MSCCLPP_BULK_AVAILABLE + extern __shared__ __align__(128) uint8_t raw[]; + T* tile = reinterpret_cast(raw); + for (uint32_t i = threadIdx.x; i < count; i += blockDim.x) tile[i] = addend; + __syncthreads(); + mscclpp::bulkFence(); + if (threadIdx.x == 0) { + mscclpp::bulkReduceStore(dst, tile, count * sizeof(T)); + mscclpp::bulkStoreCommit(); + mscclpp::bulkStoreWait(); + } +#endif +} + +class BulkTestData { + public: + BulkTestData(int numElems) + : src_(mscclpp::GpuBuffer(numElems).memory()), + dst_(mscclpp::GpuBuffer(numElems).memory()), + numElems_(numElems) { + std::vector host(numElems); + for (int i = 0; i < numElems; ++i) host[i] = i + 1; + MSCCLPP_CUDATHROW(cudaMemcpy(src_.get(), host.data(), numElems * sizeof(int), cudaMemcpyHostToDevice)); + MSCCLPP_CUDATHROW(cudaMemset(dst_.get(), 0, numElems * sizeof(int))); + } + + std::vector result() { + std::vector out(numElems_); + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), dst_.get(), out.size() * sizeof(int), cudaMemcpyDeviceToHost)); + return out; + } + + int* src() { return src_.get(); } + int* dst() { return dst_.get(); } + + private: + std::shared_ptr src_; + std::shared_ptr dst_; + int numElems_; +}; + +TEST(BulkTest, Load) { + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + BulkTestData f(kElems); + kernelBulkLoad<<<1, 256>>>(f.src(), f.dst()); + std::vector out = f.result(); + for (uint32_t i = 0; i < kElems; ++i) EXPECT_EQ(out[i], (int)i + 1); +} + +TEST(BulkTest, Gather) { + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + constexpr int kNumSrc = 4; + BulkTestData f(kElems * kNumSrc); + kernelBulkGather<<<1, 256>>>(f.src(), f.dst()); + std::vector out = f.result(); + for (uint32_t i = 0; i < kElems; ++i) { + int expected = 0; + for (int s = 0; s < kNumSrc; ++s) expected += s * kElems + i + 1; + EXPECT_EQ(out[i], expected); + } +} + +TEST(BulkTest, RoundTrip) { + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + constexpr int kNumChunks = 8; + BulkTestData f(kElems * kNumChunks); + kernelBulkRoundTrip<<<1, 256>>>(f.src(), f.dst()); + std::vector out = f.result(); + for (uint32_t i = 0; i < kElems * kNumChunks; ++i) EXPECT_EQ(out[i], (int)i + 2); +} + +// Each reduction type accumulates into a destination seeded with a known base, so the test +// distinguishes an accumulate from an overwrite. +template +static void reduceStoreTest(float base, float addend) { + constexpr uint32_t kCount = 1024; + std::shared_ptr dst = mscclpp::GpuBuffer(kCount).memory(); + std::vector host(kCount, static_cast(base)); + MSCCLPP_CUDATHROW(cudaMemcpy(dst.get(), host.data(), kCount * sizeof(T), cudaMemcpyHostToDevice)); + + kernelBulkReduce<<<1, 256, kCount * sizeof(T)>>>(dst.get(), static_cast(addend), kCount); + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out(kCount); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), dst.get(), kCount * sizeof(T), cudaMemcpyDeviceToHost)); + for (uint32_t i = 0; i < kCount; ++i) { + EXPECT_EQ(static_cast(out[i]), base + addend); + } +} + +TEST(BulkTest, ReduceStoreFloat) { + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + reduceStoreTest(10.0f, 1.5f); +} + +TEST(BulkTest, ReduceStoreBf16) { + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + reduceStoreTest<__nv_bfloat16>(10.0f, 2.0f); +} + +TEST(BulkTest, ReduceStoreUint32) { + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + reduceStoreTest(10.0f, 3.0f); +} + +#endif // defined(MSCCLPP_DEVICE_CUDA)