Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 54 additions & 59 deletions projects/rccl/cmake/DeviceLinker.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -915,12 +915,63 @@ if(ENABLE_ROCSHMEM_GIN)
)
endif()

# ===========================================================================
# rocshmem QP device bitcode for GIN GDA puts (QueuePair::put_nbi, etc.).
# Used by GIN symmetric kernels. The installed per-arch .bc files are
# arch-optimized and cannot be used with -mlink-builtin-bitcode across archs,
# so this is unoptimized IR + constmem stubs.
# ===========================================================================
set(GIN_ROCSHMEM_QP_BC_FLAG "")
set(GIN_ROCSHMEM_QP_DEPS "")
if(ENABLE_ROCSHMEM_GIN AND ROCSHMEM_SOURCE_DIR AND DL_GPU_TARGETS)
Comment thread
nusislam marked this conversation as resolved.
Resolved
list(GET DL_GPU_TARGETS 0 _bc_arch)
set(_qp_bc "${DEVICE_BUILD_DIR}/rocshmem_qp_device.bc")
find_program(_llvm_link llvm-link HINTS ${ROCM_PATH}/llvm/bin REQUIRED)
find_program(_llvm_dis llvm-dis HINTS ${ROCM_PATH}/llvm/bin REQUIRED)
find_program(_llvm_as llvm-as HINTS ${ROCM_PATH}/llvm/bin REQUIRED)

set(_cm_src "${CMAKE_SOURCE_DIR}/src/gin/gin_rocshmem_constmem.hip")
set(_cm_bc "${DEVICE_BUILD_DIR}/gin_rocshmem_constmem.bc")
set(_qp_raw "${DEVICE_BUILD_DIR}/rocshmem_qp_raw.bc")

add_custom_command(
OUTPUT ${_cm_bc}
COMMAND ${DL_CLANG}
-x hip --cuda-device-only --offload-arch=${_bc_arch}
-emit-llvm -Xclang -disable-llvm-passes
-std=c++17 -fPIC
-I${ROCSHMEM_SOURCE_DIR}/src
-I${ROCSHMEM_SOURCE_DIR}/include
-c -o ${_cm_bc} ${_cm_src}
DEPENDS ${_cm_src}
COMMENT "DL: compiling rocshmem constmem stubs to device bitcode"
VERBATIM)

add_custom_command(
OUTPUT ${_qp_bc}
COMMAND ${_llvm_link} ${_cm_bc} -o ${_qp_raw}
COMMAND ${_llvm_dis} -o ${_qp_raw}.ll ${_qp_raw}
COMMAND grep -v -E "@llvm[.]compiler[.]used|@__hip_cuid_"
Comment thread
nusislam marked this conversation as resolved.
Resolved
${_qp_raw}.ll > ${_qp_raw}.clean.ll
COMMAND ${_llvm_as} ${_qp_raw}.clean.ll -o ${_qp_bc}
DEPENDS ${_cm_bc}
COMMENT "DL: linking rocshmem QP device bitcode (with constmem, stripped)"
VERBATIM)
add_custom_target(dl_rocshmem_qp_bc DEPENDS ${_qp_bc})
if(TARGET rocshmem_static)
add_dependencies(dl_rocshmem_qp_bc rocshmem_static)
endif()
set(GIN_ROCSHMEM_QP_BC_FLAG -Xclang -mlink-builtin-bitcode -Xclang ${_qp_bc})
set(GIN_ROCSHMEM_QP_DEPS dl_rocshmem_qp_bc)
endif()

# ===========================================================================
# gin_all_reduce_sdma.cu.cpp: GIN-SDMA allreduce kernel.
#
# The defines go here because this file is filtered out of the rccl target, so
# per-source properties never apply. Leaving only SDMA on keeps ncclGinCallImpl
# on its single backend branch. GDA must stay off, as this object gets no QP bitcode.
# gfx1250 GinScatter and gfx950 two-shot both use Anvil SDMA gin.put().
# ===========================================================================
set(GIN_ALLREDUCE_SDMA_FAT_OBJ "")
if(ENABLE_ROCSHMEM_GIN)
Expand Down Expand Up @@ -1300,65 +1351,9 @@ endforeach()
# ===========================================================================
set(SYM_FAT_OBJS "")
if(GENERATE_SYM_KERNELS)
# When ENABLE_ROCSHMEM_GIN is set, GIN device templates reference rocshmem
# device symbols (QueuePair::put_nbi, atomic_add, etc.). The installed per-arch
# .bc files are arch-optimized (opt -mcpu=) and can't be used with
# -mlink-builtin-bitcode across archs. Instead, llvm-link the pre-opt
# individual source .bc files (arch-agnostic) into a minimal QP-only bitcode.
set(_sym_rocshmem_bc_flag "")
set(_sym_rocshmem_deps "")
if(ENABLE_ROCSHMEM_GIN AND ROCSHMEM_SOURCE_DIR)
# Pick the first arch's pre-opt bitcode dir (all archs produce identical
# unoptimized IR since -Xclang -disable-llvm-passes is used).
list(GET DL_GPU_TARGETS 0 _bc_arch)
set(_bc_dir "${ROCSHMEM_SOURCE_DIR}/build/bitcode/${_bc_arch}")
set(_qp_bc "${DEVICE_BUILD_DIR}/rocshmem_qp_device.bc")
find_program(_llvm_link llvm-link HINTS ${ROCM_PATH}/llvm/bin REQUIRED)
find_program(_llvm_dis llvm-dis HINTS ${ROCM_PATH}/llvm/bin REQUIRED)
find_program(_llvm_as llvm-as HINTS ${ROCM_PATH}/llvm/bin REQUIRED)

# Pipeline:
# 1. Compile gin_rocshmem_constmem.hip → device-only .bc (provides
# rocshmem::constmem and rocshmem::logd_constants definitions that
# queue_pair.bc references as external)
# 2. llvm-link QP .bc files + constmem .bc into one module
# 3. Strip @llvm.compiler.used and @__hip_cuid_ via text round-trip
# (these AMDGCN addrspace(1) appending globals clash with the
# host-side addrspace(0) equivalents in fat-object compilation)
set(_cm_src "${CMAKE_SOURCE_DIR}/src/gin/gin_rocshmem_constmem.hip")
set(_cm_bc "${DEVICE_BUILD_DIR}/gin_rocshmem_constmem.bc")
set(_qp_raw "${DEVICE_BUILD_DIR}/rocshmem_qp_raw.bc")

add_custom_command(
OUTPUT ${_cm_bc}
COMMAND ${DL_CLANG}
-x hip --cuda-device-only --offload-arch=${_bc_arch}
-emit-llvm -Xclang -disable-llvm-passes
-std=c++17 -fPIC
-I${ROCSHMEM_SOURCE_DIR}/src
-I${ROCSHMEM_SOURCE_DIR}/include
-c -o ${_cm_bc} ${_cm_src}
DEPENDS ${_cm_src}
COMMENT "DL: compiling rocshmem constmem stubs to device bitcode"
VERBATIM)

add_custom_command(
OUTPUT ${_qp_bc}
COMMAND ${_llvm_link} ${_cm_bc} -o ${_qp_raw}
COMMAND ${_llvm_dis} -o ${_qp_raw}.ll ${_qp_raw}
COMMAND grep -v -E "@llvm[.]compiler[.]used|@__hip_cuid_"
${_qp_raw}.ll > ${_qp_raw}.clean.ll
COMMAND ${_llvm_as} ${_qp_raw}.clean.ll -o ${_qp_bc}
DEPENDS ${_cm_bc}
COMMENT "DL: linking rocshmem QP device bitcode (with constmem, stripped)"
VERBATIM)
add_custom_target(dl_rocshmem_qp_bc DEPENDS ${_qp_bc})
if(TARGET rocshmem_static)
add_dependencies(dl_rocshmem_qp_bc rocshmem_static)
endif()
set(_sym_rocshmem_bc_flag -Xclang -mlink-builtin-bitcode -Xclang ${_qp_bc})
set(_sym_rocshmem_deps dl_rocshmem_qp_bc)
endif()
# GIN symmetric kernels reuse the QP bitcode generated above.
set(_sym_rocshmem_bc_flag "${GIN_ROCSHMEM_QP_BC_FLAG}")
set(_sym_rocshmem_deps "${GIN_ROCSHMEM_QP_DEPS}")
file(GLOB _sym_srcs CONFIGURE_DEPENDS "${HIPIFY_DIR}/gensrc/symmetric/*.cpp")
foreach(_sym_src IN LISTS _sym_srcs)
get_filename_component(_sym_name "${_sym_src}" NAME_WE)
Expand Down
21 changes: 15 additions & 6 deletions projects/rccl/src/algorithms/gin/gin_all_reduce.h
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@
* Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
*
* Host entry points for the GIN-SDMA AllReduce path selected by rcclSelectAllReduce
* when symmetric windows are used. By default only messages >= 256 MiB
* (GIN two-shot) take this path; smaller messages use DDA AllReduce.
* when symmetric windows are used. By default gfx950 takes GIN two-shot for
* messages >= 256 MiB; gfx1250 with Anvil SDMA takes GinScatter from 64 MiB.
* Smaller messages use DDA AllReduce.
* RCCL_GIN_ALLREDUCE_FORCE_ENABLE=1 also enables LSA one-shot (<= 4 MiB) and
* LSA two-shot ((4 MiB, 256 MiB)).
* LSA two-shot ((4 MiB, GIN floor)).
* See LICENSE.txt for license information.
******************************************************************************/

Expand All @@ -19,12 +20,14 @@
#include <stdint.h>

struct ncclComm;
struct ncclDevrWindow;

// Size and CTA constants live in gin_all_reduce_policy.h so host unit tests and
// ncclAllReduceGinSdmaEligible() share one definition.
// LSA one-shot for messages <= kGinAllReduceLsaOneShotMaxBytes (force-enable only).
// LSA two-shot for (4 MiB, 256 MiB) (force-enable only).
// GIN two-shot for messages >= kGinAllReduceGinTwoShotMinBytes (default path).
// LSA two-shot between 4 MiB and the GIN floor (force-enable only).
// gfx950 + Anvil SDMA: GIN two-shot for messages >= kGinAllReduceGinTwoShotMinBytes (256 MiB).
// gfx1250 + Anvil SDMA: GinScatter for messages >= kGinAllReduceGinScatterMinBytes (64 MiB).

// Lazily created on the first eligible AllReduce and torn down with the comm.
// Declared unconditionally: ncclComm embeds this even when ENABLE_ROCSHMEM_GIN is off.
Expand All @@ -36,10 +39,16 @@ struct ncclComm;
// would freeze at its captured value while the device counters kept advancing.
// intraGpuCtaBar is two device uint32s (arrived, sense) for the GIN two-shot intra-GPU
// CTA barrier; the pointer is stable, the words themselves advance on the device.
// scratch / scratchWin are the symmetric staging buffer for GinScatter (512 MiB of
// CE-style ping-pong incoming slots). Allocated and window-registered once on
// gfx1250 before GIN connect.
struct ncclGinAllReduceState {
bool initialized;
struct ncclDevComm devComm;
uint32_t* intraGpuCtaBar;
void* scratch;
size_t scratchBytes;
struct ncclDevrWindow* scratchWin;
};

#if defined(ENABLE_ROCSHMEM_GIN)
Expand All @@ -51,7 +60,7 @@ bool ncclAllReduceGinSdmaEligible(ncclComm* comm, const void* sendbuff, void* re
ncclDataType_t datatype, ncclRedOp_t op);

// True when this AllReduce is a GIN-SDMA candidate that the default size policy
// left for DDA (message < 256 MiB and RCCL_GIN_ALLREDUCE_FORCE_ENABLE is not 1).
// left for DDA (message below the GIN floor and RCCL_GIN_ALLREDUCE_FORCE_ENABLE is not 1).
bool ncclAllReduceGinSdmaYieldToDda(ncclComm* comm, const void* sendbuff, void* recvbuff, size_t count,
ncclDataType_t datatype, ncclRedOp_t op);

Expand Down
170 changes: 159 additions & 11 deletions projects/rccl/src/algorithms/gin/gin_all_reduce_policy.h
Original file line number Diff line number Diff line change
@@ -1,9 +1,15 @@
/*************************************************************************
* Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
*
* Size and alignment policy for GIN-SDMA AllReduce. Pure host header included by
* gin_all_reduce.h so unit tests exercise the same gates
* ncclAllReduceGinSdmaEligible() uses after the comm / buffer / gfx950 checks pass.
* Size and alignment policy for GIN-SDMA AllReduce. Included by gin_all_reduce.h
* so unit tests exercise the same gates ncclAllReduceGinSdmaEligible() uses after
* the comm / buffer / gfx950 / gfx1250 checks pass. Chunk-plan helpers are
* host+device so the GinScatter kernel splits the same way the tests do.
*
* gfx950 default: GIN two-shot from 256 MiB (LSA RS + GIN AG, Anvil SDMA).
* gfx1250 + Anvil SDMA: GinScatter from 64 MiB.
* GinScatter scratch follows CE AllReduce (ce_reduce_impl.h / ce_coll.h):
* 512 MiB of incoming ping-pong slots, reduced output in the user recvbuff.
* See LICENSE.txt for license information.
******************************************************************************/

Expand All @@ -22,9 +28,133 @@ constexpr int kGinAllReduceLsaThreadsPerCta = 512;
constexpr size_t kGinAllReduceLsaOneShotMaxBytes = 4ULL * 1024 * 1024;
constexpr size_t kGinAllReduceLsaTwoShotMidBytes = 32ULL * 1024 * 1024;
constexpr size_t kGinAllReduceGinTwoShotMinBytes = 256ULL * 1024 * 1024;
// GinScatter (CE-style local reduce) from this size inclusive on gfx1250.
constexpr size_t kGinAllReduceGinScatterMinBytes = 64ULL * 1024 * 1024;
// Put/reduce overlap starts here. Below this, ginScatter is one full scatter then
// one full reduce (or sequential slot reuse if scratch cannot hold the shard).
constexpr size_t kGinAllReduceGinScatterPipelineMinBytes = 512ULL * 1024 * 1024;
// Largest AllReduce ginScatter is sized for (2 GiB and 4 GiB included).
constexpr size_t kGinAllReduceGinScatterMaxBytes = 4ULL * 1024 * 1024 * 1024;
Comment thread
nusislam marked this conversation as resolved.
Resolved
// CE ceARTmpBuf analogue: 512 MiB of incoming ping-pong. Reduced shard is recvbuff.
constexpr size_t kGinAllReduceTwoShotScratchBytes = 512ULL * 1024 * 1024;
Comment thread
nusislam marked this conversation as resolved.
Resolved

constexpr size_t kGinAllReduceMinPutBytes = 128;

// Pipeline granule when the AllReduce is >= 512 MiB and scratch holds the shard.
// With the default 512 MiB scratch, messages above 512 MiB are staged and the
// slot (scratch / (nSlots * nRanks)) is the granule instead.
constexpr size_t kGinAllReduceGinScatterChunkBytes = 64ULL * 1024 * 1024;
Comment thread
nusislam marked this conversation as resolved.
Resolved
constexpr size_t kGinAllReduceGinScatterChunkAlign = 16;
// Matches NCCL_CE_NUM_SLOTS: ping-pong scatter staging when the shard does not
// fit in one incoming column.
constexpr int kGinAllReduceGinScatterNumSlots = 2;

#if defined(__HIPCC__) || defined(__CUDACC__)
#define GIN_AR_HD __host__ __device__ __forceinline__
#else
#define GIN_AR_HD inline
#endif

GIN_AR_HD size_t ginAllReduceGinScatterAlignDown(size_t n, size_t align) {
if (align == 0) return n;
return (n / align) * align;
}

// CE ceARTmpBuf: nSlots * nRanks * slotBytes. Reduced output is the user recvbuff.
GIN_AR_HD size_t ginAllReduceGinScatterStagedScratchBytes(size_t slotBytes, int nRanks, int nSlots) {
Comment thread
nusislam marked this conversation as resolved.
Resolved
if (nRanks <= 0 || nSlots <= 0) return 0;
return static_cast<size_t>(nSlots) * static_cast<size_t>(nRanks) * slotBytes;
}

// Largest 16-byte-aligned per-rank slot that fits in scratch with nSlots windows.
GIN_AR_HD size_t ginAllReduceGinScatterFitSlot(size_t scratchBytes, int nRanks, int nSlots) {
if (nRanks <= 0 || nSlots <= 0) return 0;
const size_t denom = static_cast<size_t>(nSlots) * static_cast<size_t>(nRanks);
return ginAllReduceGinScatterAlignDown(scratchBytes / denom, kGinAllReduceGinScatterChunkAlign);
}

// True when scratch cannot hold nRanks full incoming columns (CE pipelined path).
GIN_AR_HD bool ginAllReduceGinScatterStaged(size_t perRankBytes, int nRanks, size_t scratchBytes) {
if (nRanks <= 0) return false;
return scratchBytes < static_cast<size_t>(nRanks) * perRankBytes;
}

// Split one per-rank slice. Full-slot path: one chunk below 512 MiB, 16 MiB
// stages at 512 MiB+. Staged path (message > scratch): CE ping-pong — two slots
// and overlap at >= 512 MiB; one slot and sequential reuse below 512 MiB.
GIN_AR_HD void ginAllReduceGinScatterChunkPlan(size_t perRankBytes, int nRanks, size_t scratchBytes, int* nChunks,
size_t* uniformBytes, size_t* lastBytes, int* nPhases) {
const size_t maxChunk = kGinAllReduceGinScatterChunkBytes;
const size_t align = kGinAllReduceGinScatterChunkAlign;
const size_t totalBytes = nRanks > 0 ? perRankBytes * static_cast<size_t>(nRanks) : 0;
if (perRankBytes == 0) {
*nChunks = 1;
*uniformBytes = 0;
*lastBytes = 0;
*nPhases = 1;
return;
}
if (!ginAllReduceGinScatterStaged(perRankBytes, nRanks, scratchBytes)) {
*nPhases = 1;
if (totalBytes < kGinAllReduceGinScatterPipelineMinBytes || perRankBytes <= maxChunk) {
*nChunks = 1;
*uniformBytes = perRankBytes;
*lastBytes = perRankBytes;
return;
}
int n = static_cast<int>((perRankBytes + maxChunk - 1) / maxChunk);
Comment thread
nusislam marked this conversation as resolved.
Resolved
size_t uniform = (perRankBytes / static_cast<size_t>(n) / align) * align;
if (uniform < align) uniform = align;
*nChunks = n;
*uniformBytes = uniform;
*lastBytes = perRankBytes - uniform * static_cast<size_t>(n - 1);
return;
}

int phases = (totalBytes >= kGinAllReduceGinScatterPipelineMinBytes) ? kGinAllReduceGinScatterNumSlots : 1;
size_t slot = ginAllReduceGinScatterFitSlot(scratchBytes, nRanks, phases);
if (slot < align && phases > 1) {
Comment thread
nusislam marked this conversation as resolved.
Resolved
phases = 1;
slot = ginAllReduceGinScatterFitSlot(scratchBytes, nRanks, 1);
}
if (slot < align) {
*nChunks = 1;
*uniformBytes = perRankBytes;
*lastBytes = perRankBytes;
*nPhases = 1;
return;
}
if (perRankBytes <= slot) {
*nChunks = 1;
*uniformBytes = perRankBytes;
*lastBytes = perRankBytes;
*nPhases = 1;
return;
}
Comment thread
nusislam marked this conversation as resolved.
Resolved
int n = static_cast<int>((perRankBytes + slot - 1) / slot);
*nChunks = n;
*uniformBytes = slot;
*lastBytes = perRankBytes - slot * static_cast<size_t>(n - 1);
*nPhases = phases;
}

GIN_AR_HD size_t ginAllReduceGinScatterChunkOff(int chunk, size_t uniformBytes) {
return static_cast<size_t>(chunk) * uniformBytes;
}

GIN_AR_HD size_t ginAllReduceGinScatterChunkSize(int chunk, int nChunks, size_t uniformBytes, size_t lastBytes) {
return (chunk == nChunks - 1) ? lastBytes : uniformBytes;
}

#undef GIN_AR_HD

// Full-slot incoming columns (no reduced scratch column; CE writes the reduced
// shard to recvbuff). Used to decide whether the message fits without staging.
inline size_t ginAllReduceGinScatterScratchBytes(size_t chunkBytes, int nRanks) {
if (nRanks <= 0) return 0;
return static_cast<size_t>(nRanks) * chunkBytes;
}

// Two-shot kernels require a whole number of elements per rank and a 16-byte
// per-rank slice so vector loads stay aligned.
inline bool ginAllReduceTwoShotEligible(size_t count, size_t typeSize, int nRanks) {
Expand All @@ -40,26 +170,44 @@ inline bool ginAllReduceGinTwoShotEligible(size_t count, size_t typeSize, int nR
return chunkBytes >= kGinAllReduceMinPutBytes;
}

// ginScatter from 64 MiB when scratch holds either nRanks full incoming columns
// or at least one 16-byte-aligned staging slot per rank.
inline bool ginAllReduceGinScatterLaunch(size_t bytes, size_t chunkBytes, int nRanks, size_t scratchBytes) {
if (bytes < kGinAllReduceGinScatterMinBytes) return false;
if (nRanks <= 0) return false;
if (scratchBytes >= ginAllReduceGinScatterScratchBytes(chunkBytes, nRanks)) return true;
return ginAllReduceGinScatterFitSlot(scratchBytes, nRanks, 1) >= kGinAllReduceGinScatterChunkAlign;
Comment thread
nusislam marked this conversation as resolved.
Resolved
}

// Size-only GIN AllReduce policy. forceEnable is RCCL_GIN_ALLREDUCE_FORCE_ENABLE==1.
// Default: true only for messages >= 256 MiB that pass GIN two-shot alignment.
// Force: LSA one-shot <= 4 MiB (and >= 512 KiB), LSA two-shot in between, GIN two-shot at 256 MiB+.
inline bool ginAllReduceSizePolicyEligible(size_t count, size_t typeSize, int nRanks, bool forceEnable) {
// ginMinBytes is the default-path floor (256 MiB two-shot on gfx950, 64 MiB
// GinScatter on gfx1250). Force: LSA one-shot <= 4 MiB (and >= 512 KiB), LSA
// two-shot in between, GIN path at ginMinBytes+.
inline bool ginAllReduceSizePolicyEligible(size_t count, size_t typeSize, int nRanks, bool forceEnable,
size_t ginMinBytes = kGinAllReduceGinTwoShotMinBytes) {
const size_t bytes = count * typeSize;
if (bytes < kGinAllReduceGinTwoShotMinBytes && !forceEnable) return false;
if (bytes < ginMinBytes && !forceEnable) return false;
if (bytes < static_cast<size_t>(kGinAllReduceMinBytes)) return false;
// Inclusive 4 MiB: one-shot has no per-rank alignment requirement. Two-shot
// starts strictly above this, matching ncclAllReduceGinSdmaTyped().
if (bytes <= kGinAllReduceLsaOneShotMaxBytes) return true;
if (bytes >= kGinAllReduceGinTwoShotMinBytes) {
if (bytes >= ginMinBytes) {
return ginAllReduceGinTwoShotEligible(count, typeSize, nRanks);
}
return ginAllReduceTwoShotEligible(count, typeSize, nRanks);
}

// Smaller-than-256 MiB GIN candidates yield to DDA unless FORCE_ENABLE=1.
inline bool ginAllReduceYieldToDdaBySize(size_t count, size_t typeSize, bool forceEnable) {
// Smaller-than-ginMinBytes GIN candidates yield to DDA unless FORCE_ENABLE=1.
inline bool ginAllReduceYieldToDdaBySize(size_t count, size_t typeSize, bool forceEnable,
size_t ginMinBytes = kGinAllReduceGinTwoShotMinBytes) {
if (forceEnable) return false;
return count * typeSize < kGinAllReduceGinTwoShotMinBytes;
return count * typeSize < ginMinBytes;
}

// Host dispatch for ginAllReduceTwoShotGinScatterKernel: gfx1250 and >= 64 MiB.
// Eligibility still requires the Anvil SDMA GIN backend before launch.
inline bool ginAllReduceGinScatterDispatch(bool isGfx1250, size_t bytes) {
return isGfx1250 && bytes >= kGinAllReduceGinScatterMinBytes;
}

#endif
Loading
Loading