diff --git a/flagcx/adaptor/ccl/hccl_adaptor.cc b/flagcx/adaptor/ccl/hccl_adaptor.cc index 4014c7a2..abefdeb8 100644 --- a/flagcx/adaptor/ccl/hccl_adaptor.cc +++ b/flagcx/adaptor/ccl/hccl_adaptor.cc @@ -68,6 +68,7 @@ struct HcclSendRecvItemEx { }; HcclSendRecvItemEx item; std::vector sendRecvInfo; +static int groupDepth = 0; // TODO: unsupported flagcxResult_t hcclAdaptorGetVersion(int *version) { @@ -307,34 +308,55 @@ flagcxResult_t hcclAdaptorSend(const void *sendbuff, size_t count, flagcxDataType_t datatype, int peer, flagcxInnerComm_t comm, flagcxStream_t stream) { void *sendbuffptr = (void *)sendbuff; + HcclDataType hcclDatatype = (HcclDataType)f2h_datatype_map[datatype]; + if (groupDepth == 0) { + return (flagcxResult_t) + h2f_ret_map[HcclSend(sendbuffptr, count, hcclDatatype, (uint32_t)peer, + comm->base, stream->base)]; + } + item.comm = comm; item.stream = stream; - sendRecvInfo.emplace_back(HcclSendRecvItem{ - HcclSendRecvType::HCCL_SEND, sendbuffptr, count, - (HcclDataType)f2h_datatype_map[datatype], (uint32_t)peer}); + sendRecvInfo.emplace_back(HcclSendRecvItem{HcclSendRecvType::HCCL_SEND, + sendbuffptr, count, hcclDatatype, + (uint32_t)peer}); return flagcxSuccess; } flagcxResult_t hcclAdaptorRecv(void *recvbuff, size_t count, flagcxDataType_t datatype, int peer, flagcxInnerComm_t comm, flagcxStream_t stream) { - sendRecvInfo.emplace_back(HcclSendRecvItem{ - HcclSendRecvType::HCCL_RECV, recvbuff, count, - (HcclDataType)f2h_datatype_map[datatype], (uint32_t)peer}); + HcclDataType hcclDatatype = (HcclDataType)f2h_datatype_map[datatype]; + if (groupDepth == 0) { + return (flagcxResult_t) + h2f_ret_map[HcclRecv(recvbuff, count, hcclDatatype, (uint32_t)peer, + comm->base, stream->base)]; + } + + item.comm = comm; + item.stream = stream; + sendRecvInfo.emplace_back(HcclSendRecvItem{HcclSendRecvType::HCCL_RECV, + recvbuff, count, hcclDatatype, + (uint32_t)peer}); return flagcxSuccess; } flagcxResult_t hcclAdaptorGroupStart() { - sendRecvInfo.clear(); + if (groupDepth == 0) { + sendRecvInfo.clear(); + } + groupDepth++; return flagcxSuccess; } flagcxResult_t hcclAdaptorGroupEnd() { - uint32_t itemNum = 0; - itemNum = sendRecvInfo.size(); - if (itemNum > 0) { - return (flagcxResult_t)h2f_ret_map[HcclBatchSendRecv( - sendRecvInfo.data(), itemNum, item.comm->base, item.stream->base)]; + groupDepth--; + if (groupDepth == 0) { + uint32_t itemNum = sendRecvInfo.size(); + if (itemNum > 0) { + return (flagcxResult_t)h2f_ret_map[HcclBatchSendRecv( + sendRecvInfo.data(), itemNum, item.comm->base, item.stream->base)]; + } } return flagcxSuccess; }