Merge pull request #1662 from BertanDogancay/2.25

[SYNC] 2.25.1-1 [ROCm/rccl commit: 590ad6acc2]
2025-05-06 09:39:09 -04:00
commit c75ebd9147
@@ -228,14 +228,13 @@ static inline ncclResult_t ncclCuMemFreeAddr(void *ptr) {
  return result;
 }

-static inline ncclResult_t ncclCuMemAlloc(void **ptr, CUmemGenericAllocationHandle *handlep, size_t size) {
+static inline ncclResult_t ncclCuMemAlloc(void **ptr, CUmemGenericAllocationHandle *handlep, CUmemAllocationHandleType type, size_t size) {
  ncclResult_t result = ncclSuccess;
  size_t granularity = 0;
  CUdevice currentDev;
  CUmemAllocationProp prop = {};
  CUmemAccessDesc accessDesc = {};
  CUmemGenericAllocationHandle handle;
-  CUmemAllocationHandleType type = ncclCuMemHandleType;
  int cudaDev;
  int flag = 0;
  CUDACHECK(cudaGetDevice(&cudaDev));
@@ -284,7 +283,7 @@ static inline ncclResult_t ncclCuMemFree(void *ptr) {

 extern int ncclCuMemEnable();

-static inline ncclResult_t ncclCuMemAlloc(void **ptr, void *handlep, size_t size) {
+static inline ncclResult_t ncclCuMemAlloc(void **ptr, void *handlep, int type, size_t size) {
  WARN("CUMEM not supported prior to CUDA 11.3");
  return ncclInternalError;
 }
@@ -70,7 +70,7 @@ union ncclLLFifoLine {
 #define WARP_SIZE warpSize
 #define MAXCHANNELS 128
 #define CHANNEL_LIMIT 16
-#define NCCL_MAX_LOCAL_RANKS 64
+#define NCCL_MAX_LOCAL_RANKS 72
 #define NCCL_MAX_NTHREADS 256
 #define NCCL_MIN_NTHREADS (4*WARP_SIZE)
 #define NCCL_SIMPLE_MAX_NTHREADS NCCL_MAX_NTHREADS
@@ -204,8 +204,6 @@ struct ncclNvls {
  int down;
  int treeUp;
  int treeDown[NCCL_MAX_NVLS_TREE_ARITY];
-  int node;
-  int nNodes;
 };

 #if __CUDA_ARCH__ >= 900
@@ -17,7 +17,7 @@
 #define NCCL_SIMPLE_ALIGNMENT (WARP_SIZE * 8LL * 16LL)
 #define NCCL_BYTES_ALIGNMENT 16

-ncclResult_t ncclInitKernelsForDevice(int cudaArch, size_t* maxStackSize);
+ncclResult_t ncclInitKernelsForDevice(int cudaArch, int maxSharedMem, size_t* maxStackSize);
 ncclResult_t ncclEnqueueCheck(struct ncclInfo* info);
 ncclResult_t ncclLaunchPrepare(struct ncclComm* comm);
 ncclResult_t ncclLaunchKernelBefore_NoUncapturedCuda(struct ncclComm* comm, struct ncclKernelPlan* plan);
@@ -29,7 +29,8 @@ void ncclTopoFree(struct ncclTopoSystem* system);
 ncclResult_t ncclTopoTrimSystem(struct ncclTopoSystem* system, struct ncclComm* comm);
 ncclResult_t ncclTopoComputeP2pChannels(struct ncclComm* comm);
 ncclResult_t ncclTopoGetNvbGpus(struct ncclTopoSystem* system, int rank, int* nranks, int** ranks);
-int ncclTopoPathAllNVLink(struct ncclTopoSystem* system);
+ncclResult_t ncclTopoPathAllNVLink(struct ncclTopoSystem* system, int* allNvLink);
+
 ncclResult_t ncclTopoComputeCommCPU(struct ncclComm* comm);

 // Query topology
@@ -0,0 +1,15 @@
+/*************************************************************************
+ * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_MNNVL_H_
+#define NCCL_MNNVL_H_
+
+#include "nccl.h"
+#include "comm.h"
+
+ncclResult_t ncclMnnvlCheck(struct ncclComm* comm);
+
+#endif
@@ -36,19 +36,22 @@
 #define NVTX_SID_CommInitRankConfig   16 // same schema as NVTX_SID_CommInitRank
 #define NVTX_SID_CommInitRankScalable 17 // same schema as NVTX_SID_CommInitRank
 #define NVTX_SID_CommSplit            18
+#define NVTX_SID_CommFinalize         19

 // Define static schema ID for the reduction operation.
-#define NVTX_PAYLOAD_ENTRY_NCCL_REDOP 14 + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START
+#define NVTX_PAYLOAD_ENTRY_NCCL_REDOP 19 + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START

 extern const nvtxDomainHandle_t ncclNvtxDomainHandle;

 struct nccl_domain{static constexpr char const* name{"NCCL"};};

+/// @brief Register an NVTX payload schema for static-size payloads.
 class payload_schema {
-public:
-  explicit payload_schema(const nvtxPayloadSchemaEntry_t entries[], size_t numEntries, const uint64_t schemaId, const char* schemaName = nullptr) noexcept
+ public:
+  explicit payload_schema(const nvtxPayloadSchemaEntry_t entries[], size_t numEntries,
+    const uint64_t schemaId, const size_t size) noexcept
  {
-    schema_attr.name = schemaName;
+    schema_attr.payloadStaticSize = size;
    schema_attr.entries = entries;
    schema_attr.numEntries = numEntries;
    schema_attr.schemaId = schemaId;
@@ -69,33 +72,105 @@ private:
    NVTX_PAYLOAD_SCHEMA_ATTR_NUM_ENTRIES |
    NVTX_PAYLOAD_SCHEMA_ATTR_STATIC_SIZE |
    NVTX_PAYLOAD_SCHEMA_ATTR_SCHEMA_ID,
-    nullptr,
+    nullptr, /* schema name is not needed */
    NVTX_PAYLOAD_SCHEMA_TYPE_STATIC,
    NVTX_PAYLOAD_SCHEMA_FLAG_NONE,
    nullptr, 0, 0, 0, 0, nullptr};
 };

+// Convenience macro to give the payload parameters a scope.
+#define NVTX3_PAYLOAD(...) __VA_ARGS__
+
 // Create NVTX push/pop range with parameters
-// @param name of the operation (see `NVTX_SID_*`)
-// @param N  schema name
-// @param S  schema (entries)
-// @param P  payload (struct)
+// @param N NCCL API name without the `nccl` prefix.
+// @param T name of the used NVTX payload schema without "Schema" suffix.
+// @param P payload parameters/entries
 #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
-#define NVTX3_FUNC_WITH_PARAMS(ID, S, P) \
-  nvtxPayloadData_t nvtx3_bpl__[] = { \
-    {NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##ID, sizeof(P), &(P)}}; \
-  roctx_scoped_range_in const roctx_range__{S, nvtx3_bpl__, std::extent<decltype(S)>::value, "RCCL_" #ID};
+#define NVTX3_FUNC_WITH_PARAMS(N, T, P) \
+  constexpr uint64_t schemaId = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
+  static const payload_schema schema{T##Schema, std::extent<decltype(T##Schema)>::value - 1, \
+    schemaId, sizeof(T)}; \
+  const T _payload = {P}; \
+  nvtxPayloadData_t nvtx3_bpl__[] = {{schemaId, sizeof(_payload), &_payload}}; \
+  roctx_scoped_range_in const roctx_range__{T##Schema, nvtx3_bpl__, std::extent<decltype(T##Schema)>::value - 1, "RCCL_" #N};
 #else
-#define NVTX3_FUNC_WITH_PARAMS(ID, S, P) \
-  static const payload_schema schema{S, std::extent<decltype(S)>::value, \
-    NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##ID, #ID}; \
+#define NVTX3_FUNC_WITH_PARAMS(N, T, P) \
+  constexpr uint64_t schemaId = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
+  static const payload_schema schema{T##Schema, std::extent<decltype(T##Schema)>::value - 1, \
+    schemaId, sizeof(T)}; \
  static ::nvtx3::v1::registered_string_in<nccl_domain> const nvtx3_func_name__{__func__}; \
-  nvtxPayloadData_t nvtx3_bpl__[] = { \
-    {NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##ID, sizeof(P), &(P)}}; \
+  const T _payload = {P}; \
+  nvtxPayloadData_t nvtx3_bpl__[] = {{schemaId, sizeof(_payload), &_payload}}; \
  ::nvtx3::v1::event_attributes const nvtx3_func_attr__{nvtx3_func_name__, nvtx3_bpl__}; \
  ::nvtx3::v1::scoped_range_in<nccl_domain> const nvtx3_range__{nvtx3_func_attr__};
 #endif

+/// @brief Creates an NVTX range with extended payload using the RAII pattern.
+/// @tparam PayloadType Data type of the payload.
+template <typename PayloadType>
+class ncclNvtxRange {
+ public:
+  explicit ncclNvtxRange(const nvtxEventAttributes_t* evtAttr) noexcept {
+    nvtxDomainRangePushEx(nvtx3::domain::get<nccl_domain>(), evtAttr);
+  }
+
+  ~ncclNvtxRange() noexcept {
+    if (payloadData.payload) {
+      nvtxRangePopPayload(nvtx3::domain::get<nccl_domain>(), &payloadData, 1);
+    } else {
+      nvtxDomainRangePop(nvtx3::domain::get<nccl_domain>());
+    }
+  }
+
+  void setPayloadData(const uint64_t schemaId) noexcept
+  {
+    payloadData = {schemaId, sizeof(PayloadType), &payload};
+  }
+
+  ncclNvtxRange() = delete;
+  ncclNvtxRange(ncclNvtxRange const&) = default;
+  ncclNvtxRange& operator=(ncclNvtxRange const&) = default;
+  ncclNvtxRange(ncclNvtxRange&&) = default;
+  ncclNvtxRange& operator=(ncclNvtxRange&&) = default;
+
+  // Holds the payload data.
+  PayloadType payload{};
+
+  nvtxPayloadData_t payloadData = {NVTX_PAYLOAD_ENTRY_TYPE_INVALID, 0, NULL};
+};
+
+// Create an NVTX range with the function name as the range name. Use RAII pattern.
+// @param T Type ID of the NVTX payload (pointer for variable-size payloads).
+#define NVTX3_RANGE(T) \
+  static ::nvtx3::v1::registered_string_in<nccl_domain> const nvtx3_func_name__{__func__}; \
+  ::nvtx3::v1::event_attributes const nvtx3_func_attr__{nvtx3_func_name__}; \
+  ncclNvtxRange<T> nvtx3_range__{nvtx3_func_attr__.get()};
+
+// Add static-size payload to the NVTX range created with `NVTX3_RANGE()`,
+// which must be in this or an outer scope.
+// @param N NCCL API name without the `nccl` prefix.
+// @param S name of the used NVTX payload schema.
+// @param P payload parameters/entries
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+#define NVTX3_RANGE_ADD_PAYLOAD(N, S, P) do { \
+  constexpr uint64_t schema_id = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
+  static const payload_schema schema{S, std::extent<decltype(S)>::value - 1, schema_id, \
+    sizeof(nvtx3_range__.payload)}; \
+  nvtx3_range__.payload = {P}; \
+  nvtx3_range__.setPayloadData(schema_id); \
+  nvtxPayloadData_t nvtx3_bpl__[] = {{schema_id, sizeof(nvtx3_range__.payloadData), &nvtx3_range__.payloadData}}; \
+  roctx_scoped_range_in const roctx_range__{S, nvtx3_bpl__, std::extent<decltype(S)>::value - 1, "RCCL_" #N}; \
+} while (0)
+#else
+#define NVTX3_RANGE_ADD_PAYLOAD(N, S, P) do { \
+  constexpr uint64_t schema_id = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
+  static const payload_schema schema{S, std::extent<decltype(S)>::value - 1, schema_id, \
+    sizeof(nvtx3_range__.payload)}; \
+  nvtx3_range__.payload = {P}; \
+  nvtx3_range__.setPayloadData(schema_id); \
+} while (0)
+#endif
+
 extern void initNvtxRegisteredEnums();

 #endif
@@ -11,7 +11,7 @@

 /* This is just an empty marker (for readability), which can be omitted. */
 /* TODO: Fix issue with trailing comma at end of entry list. */
-#define NVTX_PAYLOAD_ENTRIES
+#define NCCL_NVTX_PAYLOAD_ENTRIES


 /**
@@ -32,7 +32,7 @@
 *
 * Example:
 *  NVTX_DEFINE_SCHEMA_FOR_STRUCT(your_struct, "SchemaName",
- *      NVTX_PAYLOAD_ENTRIES(
+ *      NCCL_NVTX_PAYLOAD_ENTRIES(
 *          (index, TYPE_INT, "integer value"),
 *          (dpfloat, TYPE_DOUBLE, "fp64 value"),
 *          (text, TYPE_CSTRING, "text", NULL, 24)
@@ -80,7 +80,7 @@
 *
 * Example:
 *  NVTX_DEFINE_STRUCT_WITH_SCHEMA(your_struct_name, "Your schema name",
- *      NVTX_PAYLOAD_ENTRIES(
+ *      NCCL_NVTX_PAYLOAD_ENTRIES(
 *          (int, index, TYPE_INT, "integer value"),
 *          (double, dpfloat, TYPE_DOUBLE, "fp64 value"),
 *          (const char, (text, 24), TYPE_CSTRING, "text", NULL, 24)
@@ -0,0 +1,178 @@
+/*************************************************************************
+ * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+/// Definitions of NVTX payload types and schemas used for the NVTX
+/// instrumentation in init.cc and collectives.cc.
+
+#ifndef NVTX_PAYLOAD_SCHEMAS_H_
+#define NVTX_PAYLOAD_SCHEMAS_H_
+
+
+#include "nccl.h"
+#include "nvtx3/nvToolsExtPayload.h"
+#include "nvtx3/nvToolsExtPayloadHelper.h"
+
+/**
+ * \brief Define a C struct together with the matching schema entries.
+ *
+ * Does the same as `NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA`, but without creating the
+ * schema attributes. (Remove this helper when it is available in the NVTX headers.)
+ */
+#define NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(struct_id, prefix, entries) \
+  _NVTX_PAYLOAD_TYPEDEF_STRUCT(struct_id, _NVTX_PAYLOAD_PASS_THROUGH entries) \
+  prefix _NVTX_PAYLOAD_SCHEMA_INIT_ENTRIES(struct_id, _NVTX_PAYLOAD_PASS_THROUGH entries)
+
+// C strings used as NVTX payload entry names.
+static constexpr char const* nccl_nvtxCommStr = "NCCL communicator ID";
+static constexpr char const* nccl_nvtxCudaDevStr = "CUDA device";
+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_nvtxMsgSizeSendStr = "Message size [bytes] (Send)";
+static constexpr char const* nccl_nvtxMsgSizeRecvStr = "Message size [bytes] (Recv)";
+static constexpr char const* nccl_nvtxReductionOpStrpStr = "Reduction operation";
+static constexpr char const* nccl_nvtxDataTypeStr = "Data type";
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsCommInitAll, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, commhash, TYPE_UINT64, nccl_nvtxCommStr),
+    (int, ndev, TYPE_INT, "No. of devices")
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsCommInitRank, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, newcomm, TYPE_UINT64, nccl_nvtxCommStr),
+    (int, nranks, TYPE_INT, nccl_nvtxNranksStr),
+    (int, myrank, TYPE_INT, nccl_nvtxRankStr),
+    (int, cudaDev, TYPE_INT, nccl_nvtxCudaDevStr)
+  )
+)
+// The typedef and payload schema for ncclCommInitRank is also used for,
+// ncclCommInitRankConfig, ncclCommInitRankScalable, ncclCommDestroy, and ncclCommAbort.
+typedef NcclNvtxParamsCommInitRank NcclNvtxParamsCommInitRankConfig;
+typedef NcclNvtxParamsCommInitRank NcclNvtxParamsCommInitRankScalable;
+typedef NcclNvtxParamsCommInitRank NcclNvtxParamsCommAbort;
+typedef NcclNvtxParamsCommInitRank NcclNvtxParamsCommDestroy;
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsCommSplit, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, newcomm, TYPE_UINT64, nccl_nvtxCommStr),
+    (uint64_t, parentcomm, TYPE_UINT64, "Parent NCCL communicator ID"),
+    (int, nranks, TYPE_INT, nccl_nvtxNranksStr),
+    (int, myrank, TYPE_INT, nccl_nvtxRankStr),
+    (int, cudaDev, TYPE_INT, nccl_nvtxCudaDevStr),
+    (int, color, TYPE_INT, "Color"),
+    (int, key, TYPE_INT, "Key")
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsCommFinalize, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsAllGather, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsAllReduce, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (ncclRedOp_t, op, NCCL_REDOP, nccl_nvtxReductionOpStrpStr),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsAllToAll, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsAllToAllv, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, sendBytes, TYPE_SIZE, nccl_nvtxMsgSizeSendStr),
+    (size_t, recvBytes, TYPE_SIZE, nccl_nvtxMsgSizeRecvStr),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsBroadcast, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (int, root, TYPE_INT, "Root"),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsGather, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (int, root, TYPE_INT, "Root"),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsReduce, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (int, root, TYPE_INT, "Root"),
+    (ncclRedOp_t, op, NCCL_REDOP, nccl_nvtxReductionOpStrpStr),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsReduceScatter, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (ncclRedOp_t, op, NCCL_REDOP, nccl_nvtxReductionOpStrpStr),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsScatter, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (int, root, TYPE_INT, "Root"),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+// Used in NCCL APIs `ncclSend` and `ncclRecv`.
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsSendRecv, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (int, peer, TYPE_INT, "Peer rank"),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsMSCCL, static constexpr,
+  NCCL_NVTX_PAYLOAD_ENTRIES(
+    (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr),
+    (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr),
+    (ncclRedOp_t, op, NCCL_REDOP, nccl_nvtxReductionOpStrpStr),
+    (ncclDataType_t, datatype, TYPE_DATATYPE, nccl_nvtxDataTypeStr)
+  )
+)
+
+#endif // end include guard
@@ -13,7 +13,10 @@ struct nccl_domain{static constexpr char const* name{"NCCL"};};

 #define NVTX3_FUNC_RANGE_IN(domain)
 #define nvtxNameOsThreadA(syscall, thread)
-#define NVTX3_FUNC_WITH_PARAMS(ID, S, P)
+#define NVTX3_FUNC_WITH_PARAMS(N, T, P)
+#define NVTX3_PAYLOAD(...) __VA_ARGS__
+#define NVTX3_RANGE(T)
+#define NVTX3_RANGE_ADD_PAYLOAD(N, S, P)

 #define NVTX_PAYLOAD_ENTRY_NCCL_REDOP 11

@@ -380,6 +380,8 @@ ncclResult_t ncclProxyStart(struct ncclComm* comm);
 ncclResult_t ncclProxyInit(struct ncclComm* comm, struct ncclSocket* sock, union ncclSocketAddress* peerAddresses, uint64_t *peerAddressesUDS);
 ncclResult_t ncclProxyCreate(struct ncclComm* comm);
 ncclResult_t ncclProxyConnect(struct ncclComm* comm, int transport, int send, int proxyRank, struct ncclProxyConnector* proxyConn);
+
+// NB: ncclProxyMsgTypeStr[] in proxy.cc needs to match
 enum ncclProxyMsgType {
  ncclProxyMsgInit = 1,
  ncclProxyMsgSharedInit = 2,
@@ -70,6 +70,9 @@ DECLARE_ROCM_PFN_EXTERN(hsa_status_string);
 extern int ncclCuMemEnable();
 extern int ncclCuMemHostEnable();

+// Handle type used for cuMemCreate()
+extern CUmemAllocationHandleType ncclCuMemHandleType;
+
 ncclResult_t rocmLibraryInit(void);

 extern bool ncclCudaLaunchBlocking; // initialized by ncclCudaLibraryInit()