Merge remote-tracking branch 'nccl/master' into develop

src/include/comm.h

@@ -174,8 +174,8 @@ struct ncclComm {
   ncclCollNet_t* ncclCollNet;
   void* bootstrap;
   // Bitmasks for ncclTransportP2pSetup
-  uint32_t* connectSend;
-  uint32_t* connectRecv;
+  uint64_t* connectSend;
+  uint64_t* connectRecv;
 
   int rank;    // my rank in the communicator
   int nRanks;  // number of GPUs in communicator

src/include/cudawrap.h

@@ -8,6 +8,8 @@
 #define NCCL_CUDAWRAP_H_
 
 #include <cuda.h>
+#include <cuda_runtime.h>
+#include "checks.h"
 
 #if CUDART_VERSION >= 11030
 #include <cudaTypedefs.h>
@@ -83,6 +85,18 @@ DECLARE_CUDA_PFN_EXTERN(cuDriverGetVersion, 2020);
 DECLARE_CUDA_PFN_EXTERN(cuGetProcAddress, 11030);
 
 
-ncclResult_t cudaLibraryInit(void);
+ncclResult_t ncclCudaLibraryInit(void);
+
+extern int ncclCudaDriverVersionCache;
+
+inline ncclResult_t ncclCudaDriverVersion(int* driver) {
+  int version = __atomic_load_n(&ncclCudaDriverVersionCache, __ATOMIC_RELAXED);
+  if (version == -1) {
+    CUDACHECK(cudaDriverGetVersion(&version));
+    __atomic_store_n(&ncclCudaDriverVersionCache, version, __ATOMIC_RELAXED);
+  }
+  *driver = version;
+  return ncclSuccess;
+}
 
 #endif

src/include/graph.h

@@ -37,6 +37,7 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int ne
 #define MAX_XGMI_INTER_GPUS 4
 ncclResult_t ncclTopoGetIntraNetDev(struct ncclTopoSystem* system, int rank, struct ncclTopoGraph* graph, int channelId, int type, int* dev);
 ncclResult_t ncclTopoGetLinkType(struct ncclTopoSystem* system, int cudaDev1, int cudaDev2, bool* isXGMI, int maxInter=MAX_XGMI_INTER_GPUS, int nInter=0, int *inter=nullptr);
+ncclResult_t ncclTopoNeedFlush(struct ncclTopoSystem* system, int64_t busId, int* flush);
 ncclResult_t ncclTopoCheckNet(struct ncclTopoSystem* system, int64_t id1, int64_t id2, int* net);
 int ncclPxnDisable(struct ncclComm* comm);
 ncclResult_t ncclTopoGetPxnRanks(struct ncclComm* comm, int** intermediateRanks, int* nranks);
@@ -109,7 +110,7 @@ struct ncclTopoRanks {
 };
 
 ncclResult_t ncclTopoPreset(struct ncclComm* comm,
-    struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph,
+    struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph, struct ncclTopoGraph* collNetGraph,
     struct ncclTopoRanks* topoRanks);
 
 ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePatterns,

src/include/nvmlwrap.h

@@ -107,6 +107,75 @@ typedef enum nvmlGpuP2PCapsIndex_enum
     NVML_P2P_CAPS_INDEX_UNKNOWN
 } nvmlGpuP2PCapsIndex_t;
 
+/**
+ * Represents the type for sample value returned
+ */
+typedef enum nvmlValueType_enum
+{
+    NVML_VALUE_TYPE_DOUBLE = 0,
+    NVML_VALUE_TYPE_UNSIGNED_INT = 1,
+    NVML_VALUE_TYPE_UNSIGNED_LONG = 2,
+    NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3,
+    NVML_VALUE_TYPE_SIGNED_LONG_LONG = 4,
+
+    // Keep this last
+    NVML_VALUE_TYPE_COUNT
+}nvmlValueType_t;
+
+
+/**
+ * Union to represent different types of Value
+ */
+typedef union nvmlValue_st
+{
+    double dVal;                    //!< If the value is double
+    unsigned int uiVal;             //!< If the value is unsigned int
+    unsigned long ulVal;            //!< If the value is unsigned long
+    unsigned long long ullVal;      //!< If the value is unsigned long long
+    signed long long sllVal;        //!< If the value is signed long long
+}nvmlValue_t;
+
+/**
+ * Field Identifiers.
+ *
+ * All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change.
+ */
+
+/* NVLink Speed */
+#define NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON 90  //!< Common NVLink Speed in MBps for active links
+#define NVML_FI_DEV_NVLINK_LINK_COUNT        91  //!< Number of NVLinks present on the device
+
+/**
+ * Remote device NVLink ID
+ *
+ * Link ID needs to be specified in the scopeId field in nvmlFieldValue_t.
+ */
+#define NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID     146 //!< Remote device NVLink ID
+
+/**
+ * NVSwitch: connected NVLink count
+ */
+#define NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT   147  //!< Number of NVLinks connected to NVSwitch
+
+#define NVML_FI_DEV_NVLINK_GET_SPEED                  164
+#define NVML_FI_DEV_NVLINK_GET_STATE                  165
+#define NVML_FI_DEV_NVLINK_GET_VERSION                166
+#define NVML_FI_MAX 167 //!< One greater than the largest field ID defined above
+
+/**
+ * Information for a Field Value Sample
+ */
+typedef struct nvmlFieldValue_st
+{
+    unsigned int fieldId;       //!< ID of the NVML field to retrieve. This must be set before any call that uses this struct. See the constants starting with NVML_FI_ above.
+    unsigned int scopeId;       //!< Scope ID can represent data used by NVML depending on fieldId's context. For example, for NVLink throughput counter data, scopeId can represent linkId.
+    long long timestamp;        //!< CPU Timestamp of this value in microseconds since 1970
+    long long latencyUsec;      //!< How long this field value took to update (in usec) within NVML. This may be averaged across several fields that are serviced by the same driver call.
+    nvmlValueType_t valueType;  //!< Type of the value stored in value
+    nvmlReturn_t nvmlReturn;    //!< Return code for retrieving this value. This must be checked before looking at value, as value is undefined if nvmlReturn != NVML_SUCCESS
+    nvmlValue_t value;          //!< Value for this field. This is only valid if nvmlReturn == NVML_SUCCESS
+} nvmlFieldValue_t;
+
 /* End of nvml.h */
 #endif // NCCL_NVML_DIRECT
 
@@ -135,4 +204,6 @@ ncclResult_t ncclNvmlDeviceGetNvLinkRemotePciInfo(nvmlDevice_t device, unsigned
 ncclResult_t ncclNvmlDeviceGetNvLinkCapability(nvmlDevice_t device, unsigned int link, nvmlNvLinkCapability_t capability, unsigned int *capResult);
 ncclResult_t ncclNvmlDeviceGetCudaComputeCapability(nvmlDevice_t device, int* major, int* minor);
 ncclResult_t ncclNvmlDeviceGetP2PStatus(nvmlDevice_t device1, nvmlDevice_t device2, nvmlGpuP2PCapsIndex_t p2pIndex, nvmlGpuP2PStatus_t* p2pStatus);
+ncclResult_t ncclNvmlDeviceGetFieldValues(nvmlDevice_t device, int valuesCount, nvmlFieldValue_t *values);
+
 #endif // End include guard

src/include/strongstream.h

@@ -18,11 +18,11 @@
 struct ncclCudaGraph {
 #if CUDART_VERSION >= 11030
   cudaGraph_t graph;
-  uint64_t graphId;
+  unsigned long long graphId;
 #endif
 };
 
-inline struct ncclCudaGraph ncclCudaGraphNull() {
+inline struct ncclCudaGraph ncclCudaGraphNone() {
   struct ncclCudaGraph tmp;
   #if CUDART_VERSION >= 11030
     tmp.graph = nullptr;
@@ -50,7 +50,6 @@ inline bool ncclCudaGraphSame(struct ncclCudaGraph a, struct ncclCudaGraph b) {
 ncclResult_t ncclCudaGetCapturingGraph(struct ncclCudaGraph* graph, hipStream_t stream);
 ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t fn, void* arg);
 
-
 /* ncclStrongStream: An abstraction over CUDA streams that do not lose their
  * identity while being captured. Regular streams have the deficiency that the
  * captured form of a stream in one graph launch has no relation to the
@@ -58,37 +57,30 @@ ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t
  * streams unfit for the use of serializing access to a persistent resource.
  * Strong streams have been introduced to address this need.
  *
- * Constraints of using strong streams:
+ * - All updates to a strong stream must be enclosed by a Acquire/Release pair.
  *
- * - Operations that enqueue work to the strong stream need to be enclosed by
- *   ncclStrongStream[Acquire/Release] pairs. Acquire/release act like fences,
- *   the strong stream is not stateful so there is no harm in redundant acquire
- *   or releases.
+ * - The Acquire, Release, and all updates take a ncclCudaGraph parameter
+ *   indicating the currently capturing graph (or none). This parameter must be
+ *   the same for the entire sequence of {Acquire; ...; Release}.
  *
  * - An {Acquire; ...; Release} sequence must not be concurrent with any
  *   other operations against the strong stream including graph launches which
  *   reference this stream.
- *
- * - All strong stream functions take a "graph" parameter which must reference
- *   the currently capturing graph, or null if none.
  */
 struct ncclStrongStream;
 
 ncclResult_t ncclStrongStreamConstruct(struct ncclStrongStream* ss);
 ncclResult_t ncclStrongStreamDestruct(struct ncclStrongStream* ss);
 
-// Has this strong stream ever been captured in a graph.
-bool ncclStrongStreamEverCaptured(struct ncclStrongStream* ss);
-
 // Acquire-fence the strong stream.
 ncclResult_t ncclStrongStreamAcquire(
   struct ncclCudaGraph graph, struct ncclStrongStream* ss
 );
 
 // Acquire-fence the strong stream assuming no graph is capturing. This permits
-// the caller to enqueue directly to the `ss->stream` member using native CUDA
-// calls. Strong stream must be released via:
-//   ncclStrongStreamRelease(ncclCudaGraphNull(), graphRefs, ss);
+// the caller to enqueue directly to the `ss->cudaStream` member using native CUDA
+// calls. Strong stream still must be released via:
+//   ncclStrongStreamRelease(ncclCudaGraphNone(), ss);
 ncclResult_t ncclStrongStreamAcquireUncaptured(struct ncclStrongStream* ss);
 
 // Release-fence of the strong stream.
@@ -104,6 +96,7 @@ ncclResult_t ncclStrongStreamLaunchKernel(
   struct ncclCudaGraph graph, struct ncclStrongStream* ss,
   void* fn, dim3 grid, dim3 block, void** args, size_t sharedMemBytes
 );
+
 // Cause `a` to wait for the current state `b`. Both `a` and `b` must be acquired.
 ncclResult_t ncclStrongStreamWaitStream(
   struct ncclCudaGraph graph, struct ncclStrongStream* a, struct ncclStrongStream* b
@@ -122,21 +115,23 @@ ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss);
 
 ////////////////////////////////////////////////////////////////////////////////
 
+struct ncclStrongStreamGraph; // internal to ncclStrongStream
+
 struct ncclStrongStream {
-  hipStream_t stream;
-  hipEvent_t event;
-  #if CUDART_VERSION >= 11030
-  cudaGraphNode_t node; // null if never captured, otherwise never null again
-  uint64_t graphId:63, eventIsLagging:1;
-  #endif
+  // Used when not graph capturing.
+  hipStream_t cudaStream;
+#if CUDART_VERSION >= 11030
+  // The event used to establish order between graphs and streams. During acquire
+  // this event is waited on, during release it is recorded to.
+  cudaEvent_t serialEvent;
+  // This stream ever appeared in a graph capture.
+  bool everCaptured;
+  // Tracks whether serialEvent needs to be recorded to upon Release().
+  bool serialEventNeedsRecord;
+  struct ncclStrongStreamGraph* graphHead;
+#else
+  hipEvent_t scratchEvent;
+#endif
 };
 
-inline bool ncclStrongStreamEverCaptured(struct ncclStrongStream* ss) {
-  #if CUDART_VERSION >= 11030
-    return ss->node != nullptr;
-  #else
-    return false;
-  #endif
-}
-
 #endif