AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support fused all reduce and elementwise operations (#1729)

Browse Source 
* Support fused all reduce and elementwise operations

Add additional "acc" parameter to RCCL Replayer logs

Add flag which indicates availability of new API

* Fix Recorder json parsing

* Remove unreachable code

* Remove extra acc pointer check

* .

* Revert "[DEVICE] Adding ability to choose unroll factor at runtime (#1734)"

This reverts commit 4cadf3597c.

* Use noinline to reduce kernels linking time

* Don't use noinline for gfx942 and gfx950 to avoid perf regression

---------

Co-authored-by: AtlantaPepsi <timhu102@amd.com>
Co-authored-by: BertanDogancay <bertan.dogancay@gmail.com>

[ROCm/rccl commit: 9a4213356d]
This commit is contained in:
Wenkai Du
2025-07-23 09:04:17 -07:00
committed by GitHub
parent cbb648505a
commit caff9764d3
24 changed files with 656 additions and 231 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/misc/api_trace.cc
+20 -2
View File
@@ -153,6 +153,11 @@ ncclCommRegister_impl(const ncclComm_t comm, void* buff, size_t size, void** han
ncclResult_t
ncclCommDeregister_impl(const ncclComm_t comm, void* handle);
ncclResult_t
ncclAllReduceWithBias_impl(const void* sendbuff, void* recvbuff, size_t count,
ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm,
cudaStream_t stream, const void* acc);
namespace rccl
{
namespace
@@ -211,10 +216,11 @@ RCCL_ASSERT_OFFSET(rcclApiFuncTable, mscclRunAlgo_fn, 33);
RCCL_ASSERT_OFFSET(rcclApiFuncTable, mscclUnloadAlgo_fn, 34);
RCCL_ASSERT_OFFSET(rcclApiFuncTable, ncclCommRegister_fn, 35);
RCCL_ASSERT_OFFSET(rcclApiFuncTable, ncclCommDeregister_fn, 36);
RCCL_ASSERT_OFFSET(rcclApiFuncTable, ncclAllReduceWithBias_fn, 37);
#undef RCCL_ASSERT_OFFSET
static_assert(sizeof(rcclApiFuncTable) == compute_table_size(37),
static_assert(sizeof(rcclApiFuncTable) == compute_table_size(38),
"Update table major/step version and add a new offset assertion if this "
"fails to compile");
@@ -261,7 +267,8 @@ RcclGetFunctionTable_impl()
&mscclRunAlgo_impl,
&mscclUnloadAlgo_impl,
&ncclCommRegister_impl,
&ncclCommDeregister_impl };
&ncclCommDeregister_impl,
&ncclAllReduceWithBias_impl };
#if defined(RCCL_ROCPROFILER_REGISTER) && RCCL_ROCPROFILER_REGISTER > 0
std::array<void*, 1> table_array{ tbl };
@@ -301,6 +308,9 @@ NCCL_API(ncclResult_t, ncclAllGather, const void* sendbuff, void* recvbuff,
NCCL_API(ncclResult_t, ncclAllReduce, const void* sendbuff, void* recvbuff, size_t count,
ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, hipStream_t stream);
NCCL_API(ncclResult_t, ncclAllReduceWithBias, const void* sendbuff, void* recvbuff, size_t count,
ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, hipStream_t stream, const void* acc);
NCCL_API(ncclResult_t, ncclAllToAll, const void* sendbuff, void* recvbuff, size_t count,
ncclDataType_t datatype, ncclComm_t comm, hipStream_t stream);
@@ -411,6 +421,14 @@ ncclAllReduce(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t
datatype, op, comm, stream);
}
ncclResult_t
ncclAllReduceWithBias(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
ncclRedOp_t op, ncclComm* comm, cudaStream_t stream, const void* acc)
{
return ::rccl::RcclGetFunctionTable()->ncclAllReduceWithBias_fn(sendbuff, recvbuff, count,
datatype, op, comm, stream, acc);
}
ncclResult_t
ncclAllToAll(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
ncclComm_t comm, hipStream_t stream)
projects/rccl/src/misc/recorder.cc
+5 -4
View File
@@ -35,6 +35,7 @@ rcclApiCall::rcclApiCall(rcclCall_t type, const ncclInfo& info)://name(rcclCallS
opCount(info.comm->opCount),
sendbuff(info.sendbuff),
recvbuff(info.recvbuff),
acc(info.acc),
count(info.count),
datatype(info.datatype),
op(info.op),
@@ -69,7 +70,7 @@ std::string alloc_fmt = "%s : [returned ptr : %p, size : %zu, context : [";
std::string free_fmt = "%s : [ptr : %p, context : [";
std::string redop_fmt = "%s : [scalar : %p, datatype : %d, op : %d, residence : %d, comm : %p, context : [";
std::string redopdestroy_fmt = "%s : [op : %d, comm : %p, context : [";
std::string coll_fmt = "%s : [opCount : %lx, sendbuff : [addr : %p, base : %p, size : %zu], recvbuff : [addr : %p, base : %p, size : %zu], count : %zu, datatype : %d, op : %d, root : %d, comm : %p, nranks : %d, stream : %p, task : %d, globalrank : %d, context : [";
std::string coll_fmt = "%s : [opCount : %lx, sendbuff : [addr : %p, base : %p, size : %zu], recvbuff : [addr : %p, base : %p, size : %zu], acc : %p, count : %zu, datatype : %d, op : %d, root : %d, comm : %p, nranks : %d, stream : %p, task : %d, globalrank : %d, context : [";
Recorder::Recorder()
{
@@ -256,7 +257,7 @@ void Recorder::write(const rcclApiCall &call)
default: // collectives
len = snprintf(buffer, 4096, coll_fmt.c_str(),
rcclCallStr[call.type], call.opCount, call.sendbuff, call.sendPtrBase, call.sendPtrExtent,
call.recvbuff, call.recvPtrBase, call.recvPtrExtent, call.count, call.datatype,
call.recvbuff, call.recvPtrBase, call.recvPtrExtent, call.acc, call.count, call.datatype,
call.op, call.root, call.comm, call.nRanks, call.stream, call.nTasks, call.globalRank);
}
@@ -686,9 +687,9 @@ void parseJsonEntry(const char* entry, std::vector<rcclApiCall>& calls)
default:
assert(sscanf(str.c_str() + end + 3, (coll_fmt.substr(5) + ctxt_fmt).c_str(),
&call.opCount, &call.sendbuff, &call.sendPtrBase, &call.sendPtrExtent, &call.recvbuff, &call.recvPtrBase, &call.recvPtrExtent,
&call.count, &call.datatype, &call.op, &call.root,
&call.acc, &call.count, &call.datatype, &call.op, &call.root,
&call.comm, &call.nRanks, &call.stream, &call.nTasks, &call.globalRank, &call.timestamp, &call.tid,
&call.hipDev, &call.graphCaptured, &call.graphID) == 21);
&call.hipDev, &call.graphCaptured, &call.graphID) == 22);
}
calls.push_back(call);
}