SPLITCOMM design fix in src/misc/msccl (#1715)

* Fix TOC-TOU in mcclInit

* Improving vector resize thread safety

* Initial commit rank to comm change

* Removing unwanted include header changes

* Updated CHANGELOG.md

* Update CHANGELOG.md

Co-authored-by: Jeffrey Novotny <jnovotny@amd.com>

---------

Co-authored-by: Jeffrey Novotny <jnovotny@amd.com>

src/include/api_trace.h

@@ -139,7 +139,7 @@ typedef ncclResult_t (*ncclMemAlloc_fn_t)(void** ptr, size_t size);
 typedef ncclResult_t (*ncclMemFree_fn_t)(void* ptr);
 
 typedef ncclResult_t (*mscclLoadAlgo_fn_t)(const char*        mscclAlgoFilePath,
-                                           mscclAlgoHandle_t* mscclAlgoHandle, int rank);
+                                           mscclAlgoHandle_t* mscclAlgoHandle, const ncclComm_t comm);
 
 typedef ncclResult_t (*mscclRunAlgo_fn_t)(
     const void* sendBuff, const size_t sendCounts[], const size_t sDisPls[],

src/include/msccl/msccl_lifecycle.h

@@ -17,7 +17,15 @@ void mscclSetIsCallerFlag();
 void mscclClearIsCallerFlag();
 bool mscclIsCaller();
 
-bool mscclAvailable(int rank = -1);
+/**
+ * @brief mscclAvailable() is used to determine if msccl functionality is avaliable
+ * @param comm is an optional rccl communicator, if provided uses the mscclStatus
+ * from a global map<comm -> mscclStatus> to determine if msccl is available. If not available
+ * in the map, this invocations inserts a new key value pair in the global map.
+ * If comm == nullptr, on the first invocation it initializes a static thread local variable 
+ * mscclStatus and uses the same object in subsequent calls from same thread if comm is null ptr
+ */
+bool mscclAvailable(const ncclComm_t comm = nullptr);
 
 ncclResult_t mscclSchedulerInit(ncclComm_t comm, int* numChannelsRequired);
 
@@ -33,7 +41,7 @@ ncclResult_t mscclEnqueueCheck(
 
 ncclResult_t mscclGroupEnd();
 
-ncclResult_t mscclTeardown(int rank);
+ncclResult_t mscclTeardown(const ncclComm_t comm);
 
 size_t mscclKernMaxLocalSize();
 

src/include/msccl/msccl_setup.h

@@ -11,13 +11,13 @@
 #include "comm.h"
 #include "msccl/msccl_struct.h"
 
-ncclResult_t mscclGetCaptureStatus(int rank, hipStream_t stream);
+ncclResult_t mscclGetCaptureStatus(const ncclComm_t comm, hipStream_t stream);
 
 ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream);
 
-ncclResult_t mscclSetupSyncFlags(int rank, hipStream_t stream);
+ncclResult_t mscclSetupSyncFlags(const ncclComm_t comm, hipStream_t stream);
 
-ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm);
+ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo,const ncclComm_t comm);
 
 ncclResult_t mscclSetupCount(struct mscclAlgo* hostAlgo, ncclComm_t comm, size_t count, ncclDataType_t dataType);
 

src/include/msccl/msccl_status.h

@@ -8,15 +8,15 @@
 
 #include "msccl/msccl_struct.h"
 
-bool mscclInitialized(int rank);
+bool mscclInitialized(const ncclComm_t comm);
 
-void mscclSetInitialized(int rank, bool initialized = true);
+void mscclSetInitialized(const ncclComm_t comm, bool initialized = true);
 
-void mscclRemoveRank(int rank);
+void mscclRemoveRank(const ncclComm_t comm);
 
-mscclStatus& mscclGetStatus(int rank);
+mscclStatus& mscclGetStatus(const ncclComm_t comm);
 
-mscclSavedProxyArgs& mscclGetSavedProxyArgs(int rank);
+mscclSavedProxyArgs& mscclGetSavedProxyArgs(const ncclComm_t comm);
 
 mscclThreadLocalStatus& mscclGetThreadLocalStatus();