Add NCCL task protocol trace metadata

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 <cinttypes> // PRIx64
 #include <cassert>
 #include <cfloat> // FLT_MAX
+#include <type_traits> // 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<decltype(NcclNvtxParamsKernelTaskSchema)>::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<nccl_domain>(), &attr);
+    ranges++;
+    task = task->next;
+  }
+  return ranges;
+}
+
+static void ncclNvtxPopKernelTaskProtoRanges(int ranges) {
+  while (ranges-- > 0) {
+    nvtxDomainRangePop(nvtx3::domain::get<nccl_domain>());
+  }
+}
+
 // 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;
 }

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.

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