diff --git a/src/enqueue.cc b/src/enqueue.cc index 9f59b217c4..2167c6710a 100644 --- a/src/enqueue.cc +++ b/src/enqueue.cc @@ -17,6 +17,7 @@ #include "register_inline.h" #include "ce_coll.h" #include "nvtx.h" +#include "nvtx_payload_schemas.h" #include "scheduler.h" #include "compiler.h" #include "rma/rma.h" @@ -25,11 +26,67 @@ #include // PRIx64 #include #include // FLT_MAX +#include // std::extent NCCL_PARAM(L1SharedMemoryCarveout, "L1_SHARED_MEMORY_CARVEOUT", 0); NCCL_PARAM(AllgathervEnable, "ALLGATHERV_ENABLE", 1); NCCL_PARAM(SymCeThreshold, "SYM_CE_THRESHOLD", 8 * 1024 * 1024); +static const char* ncclNvtxProtoRangeName(int proto) { + switch (proto) { + case NCCL_PROTO_LL: return "NCCL Kernel Task Proto LL"; + case NCCL_PROTO_LL128: return "NCCL Kernel Task Proto LL128"; + case NCCL_PROTO_SIMPLE: return "NCCL Kernel Task Proto Simple"; + default: return "NCCL Kernel Task Proto Unknown"; + } +} + +static int ncclNvtxPushKernelTaskProtoRanges(struct ncclComm* comm, struct ncclKernelPlan* plan) { + if (ncclParamNvtxDisable()) return 0; + + constexpr uint64_t schemaId = + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_KernelTask; + static const payload_schema schema{ + NcclNvtxParamsKernelTaskSchema, + std::extent::value - 1, + schemaId, + sizeof(NcclNvtxParamsKernelTask) + }; + + int ranges = 0; + struct ncclTaskColl* task = ncclIntruQueueHead(&plan->collTaskQueue); + while (task != nullptr) { + const NcclNvtxParamsKernelTask payload = { + comm->commHash, + comm->rank, + task->count * ncclTypeSize(task->datatype), + task->func, + task->algorithm, + task->protocol, + task->datatype, + task->root, + task->nChannels + }; + nvtxPayloadData_t payloadData = {}; + nvtxEventAttributes_t attr = {}; + attr.version = NVTX_VERSION; + attr.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE; + attr.messageType = NVTX_MESSAGE_TYPE_ASCII; + attr.message.ascii = ncclNvtxProtoRangeName(task->protocol); + NVTX_PAYLOAD_EVTATTR_SET_DATA(attr, &payloadData, schemaId, &payload, sizeof(payload)); + nvtxDomainRangePushEx(nvtx3::domain::get(), &attr); + ranges++; + task = task->next; + } + return ranges; +} + +static void ncclNvtxPopKernelTaskProtoRanges(int ranges) { + while (ranges-- > 0) { + nvtxDomainRangePop(nvtx3::domain::get()); + } +} + // Returns maximum kernel stack size of all CUDA kernels ncclResult_t ncclInitKernelsForDevice(int cudaArch, int maxSharedMem, size_t* maxStackSize) { ncclResult_t result = ncclSuccess; @@ -1697,6 +1754,7 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) { } } + int nvtxTaskProtoRanges = ncclNvtxPushKernelTaskProtoRanges(comm, plan); if (!persistent && comm->sharedRes->persistentRefs) { status = CUDACLEARERROR(cudaEventQuery(comm->sharedRes->hostStream.serialEvent)); } @@ -1846,6 +1904,7 @@ ncclResult_t ncclLaunchKernel(struct ncclComm* comm, struct ncclKernelPlan* plan } do_return: + ncclNvtxPopKernelTaskProtoRanges(nvtxTaskProtoRanges); NCCLCHECK(ncclProfilerStopKernelLaunchEvent(plan)); return ret; } diff --git a/src/include/nvtx.h b/src/include/nvtx.h index ce86cf012d..656e780ac1 100644 --- a/src/include/nvtx.h +++ b/src/include/nvtx.h @@ -43,6 +43,7 @@ #define NVTX_SID_PutSignal 21 #define NVTX_SID_Signal 22 #define NVTX_SID_WaitSignal 23 +#define NVTX_SID_KernelTask 25 // When adding new schema IDs, DO NOT re-use/overlap with the enum schema ID below! // Define static schema ID for the reduction operation. diff --git a/src/include/nvtx_payload_schemas.h b/src/include/nvtx_payload_schemas.h index a4c54788d5..dfefd1e89d 100644 --- a/src/include/nvtx_payload_schemas.h +++ b/src/include/nvtx_payload_schemas.h @@ -5,8 +5,8 @@ * See LICENSE.txt for more license information *************************************************************************/ -/// Definitions of NVTX payload types and schemas used for the NVTX -/// instrumentation in init.cc and collectives.cc. +/// Definitions of NVTX payload types and schemas used for NCCL NVTX +/// instrumentation. #ifndef NVTX_PAYLOAD_SCHEMAS_H_ #define NVTX_PAYLOAD_SCHEMAS_H_ @@ -32,6 +32,11 @@ static constexpr char const* nccl_nvtxRankStr = "Rank"; static constexpr char const* nccl_nvtxNranksStr = "No. of ranks"; static constexpr char const* nccl_nvtxMsgSizeStr = "Message size [bytes]"; static constexpr char const* nccl_nvtxReductionOpStrpStr = "Reduction operation"; +static constexpr char const* nccl_nvtxCollectiveStr = "Collective"; +static constexpr char const* nccl_nvtxAlgorithmStr = "Algorithm"; +static constexpr char const* nccl_nvtxProtocolStr = "Protocol"; +static constexpr char const* nccl_nvtxDatatypeStr = "Datatype"; +static constexpr char const* nccl_nvtxChannelCountStr = "No. of channels"; NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES( NcclNvtxParamsCommInitAll, static constexpr, @@ -138,4 +143,18 @@ NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsWaitSignal, static con (int, npeers, TYPE_INT, "Number of peers"), (int, ctx, TYPE_INT, "Context ID"))) +NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsKernelTask, static constexpr, + NCCL_NVTX_PAYLOAD_ENTRIES( + (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr), + (int, rank, TYPE_INT, nccl_nvtxRankStr), + (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr), + (int, func, TYPE_INT, nccl_nvtxCollectiveStr), + (int, algorithm, TYPE_INT, nccl_nvtxAlgorithmStr), + (int, protocol, TYPE_INT, nccl_nvtxProtocolStr), + (int, datatype, TYPE_INT, nccl_nvtxDatatypeStr), + (int, root, TYPE_INT, "Root"), + (int, nChannels, TYPE_INT, nccl_nvtxChannelCountStr) + ) +) + #endif // end include guard