Enable multi-threading for MSCCL (#1203)

MSCCL can now run in a multi-threaded configuration. To test in the unit tests, added the ENABLE_OPENMP compile definition flag and the --openmp-test-enable flag to the unit test build script. To activate, set the environment variables UT_MULTITHREADED=1 and UT_PROCESS_MASK=1. Set Jenkins to use this mode.

src/include/msccl/msccl_lifecycle.h

@@ -17,7 +17,7 @@ void mscclSetIsCallerFlag();
 void mscclClearIsCallerFlag();
 bool mscclIsCaller();
 
-bool mscclAvailable();
+bool mscclAvailable(int rank = -1);
 
 ncclResult_t mscclSchedulerInit(ncclComm_t comm, int* numChannelsRequired);
 
@@ -33,7 +33,7 @@ ncclResult_t mscclEnqueueCheck(
 
 ncclResult_t mscclGroupEnd();
 
-ncclResult_t mscclTeardown();
+ncclResult_t mscclTeardown(int rank);
 
 size_t mscclKernMaxLocalSize();
 

src/include/msccl/msccl_setup.h

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

src/include/msccl/msccl_status.h

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