MSCCL: Improve executor and integrate scheduler (#694)

* MSCCL: improve executor and add scheduler for testing * Use external scheduler * Fix cmake error * Address comments * Fix thread safe issue * Make MSCCL lifecycle APIs thread safe * Make MSCCL internal scheduler aware of topology hint * Revise error message [ROCm/rccl commit: e3b2342f39]
2023-03-15 05:34:25 +08:00
parent 8fdc4795fd
commit f7f669e7f0
25 changed files with 36884 additions and 231 deletions
@@ -100,4 +100,6 @@ static ncclResult_t mscclXmlFindTag(struct mscclXml* xml, const char* tagName, s

 ncclResult_t mscclGetAlgoFromXmlFile(const char* xmlGraphFile, struct mscclAlgo* algo, int rank);

+ncclResult_t mscclGetAlgoMetaFromXmlFile(const char* xmlGraphFile, struct mscclAlgoMeta* algoMeta);
+
 #endif
@@ -0,0 +1,52 @@
+/*************************************************************************
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT License.
+ ************************************************************************/
+
+#ifndef MSCCL_SCHEDULER_H_
+#define MSCCL_SCHEDULER_H_
+
+typedef enum { mscclFuncReduce             =  0,
+               mscclFuncBroadcast          =  1,
+               mscclFuncAllReduce          =  2,
+               mscclFuncReduceScatter      =  3,
+               mscclFuncAllGather          =  4,
+               mscclFuncSend               =  5,
+               mscclFuncRecv               =  6,
+               mscclFuncGather             =  7,
+               mscclFuncScatter            =  8,
+               mscclFuncAllToAll           =  9,
+               mscclFuncAllToAllv          =  10,
+               mscclNumFuncs               =  11 } mscclFunc_t;
+
+struct mscclSchedulerParam {
+  const void* sendBuff;
+  const size_t* sendCounts;
+  const size_t* sDisPls;
+  void* recvBuff;
+  const size_t* recvCounts;
+  const size_t* rDisPls;
+  size_t count;
+  ncclDataType_t dataType;
+  int root;
+  int peer;
+  ncclRedOp_t op;
+  mscclFunc_t func;
+  int rank;
+  int nRanks;
+  bool scheduled;
+  mscclAlgoHandle_t handle;
+};
+
+typedef struct {
+  // Name of the scheduler (mainly for logs)
+  const char* name;
+  // Load all algorithms
+  ncclResult_t (*init)();
+  // Select an algorithm
+  ncclResult_t (*selectAlgo)(struct mscclSchedulerParam* param);
+  // Unload all algorithms
+  ncclResult_t (*teardown)();
+} mscclSchedulerInterface;
+
+#endif
@@ -10,4 +10,6 @@

 mscclStatus& mscclGetStatus();

+mscclThreadLocalStatus& mscclGetThreadLocalStatus();
+
 #endif
@@ -11,8 +11,9 @@
 #include <set>
 #include <vector>
 #include "devcomm.h"
+#include "msccl/msccl_scheduler.h"

-#define MSCCL_MAX_NUM_STEPS 256
+#define MSCCL_MAX_NUM_STEPS 64
 #define MSCCL_MAX_NUM_THREAD_BLOCKS_PER_CHANNEL 32
 #define MSCCL_MAX_NUM_THREAD_BLOCKS (MSCCL_MAX_NUM_THREAD_BLOCKS_PER_CHANNEL * MAXCHANNELS)
 #define MSCCL_MAX_COUNT 72 // max concurrent number of msccl chunk transmission
@@ -35,19 +36,6 @@
 #define MSCCL_LOCAL_COPY 6
 #define MSCCL_REDUCE 7

-typedef enum { mscclFuncReduce             =  0,
-               mscclFuncBroadcast          =  1,
-               mscclFuncAllReduce          =  2,
-               mscclFuncReduceScatter      =  3,
-               mscclFuncAllGather          =  4,
-               mscclFuncSend               =  5,
-               mscclFuncRecv               =  6,
-               mscclFuncGather             =  7,
-               mscclFuncScatter            =  8,
-               mscclFuncAllToAll           =  9,
-               mscclFuncAllToAllv          =  10,
-               mscclNumFuncs               =  11 } mscclFunc_t;
-
 struct mscclTransmission {
  int16_t dependencePointer; // index to the first dependence
  int16_t numDependencies; // dependencePointer+numDependencies indicate the last dependence
@@ -98,6 +86,27 @@ struct mscclChannelInfo {
  int nRecvPeers;
 };

+struct mscclAlgoMeta {
+  // Path to algorithm file
+  std::string filePath;
+  // number of chunks of input/output in each MSCCL algorithm loop
+  int nChunksPerLoop;
+  // number of ranks required by this algorithm
+  int nRanks;
+  // need to times nRanks for all-gather, reduce-scatter and all-to-all
+  int sizeMultiplier;
+  // MSCCL function type
+  mscclFunc_t func;
+  // Min message size allowed for this algorithm.
+  int64_t minBytes;
+  // Max message size allowed for this algorithm, 0 for no limit.
+  int64_t maxBytes;
+  // Whether this algorithm is suitable for in-place.
+  bool inPlace;
+  // Whether this algorithm is suitable for out-of-place.
+  bool outOfPlace;
+};
+
 struct mscclAlgo {
  // number of chunks of input/output in each MSCCL algorithm loop
  int nChunksPerLoop;
@@ -141,29 +150,23 @@ enum mscclGroupStatus {
  mscclGroupUnsupportedOp
 };

-struct mscclSchedulerParam {
-  const void* sendBuff;
-  const size_t* sendCounts;
+struct mscclSavedSchedulerParam {
+  struct mscclSchedulerParam p;
  std::vector<size_t> savedSendCounts;
-  const size_t* sDisPls;
  std::vector<size_t> savedSDisPls;
-  void* recvBuff;
-  const size_t* recvCounts;
  std::vector<size_t> savedRecvCounts;
-  const size_t* rDisPls;
  std::vector<size_t> savedRDisPls;
-  size_t count;
-  ncclDataType_t dataType;
-  int root;
-  int peer;
-  ncclRedOp_t op;
-  mscclFunc_t func;
-  bool scheduled;
-  mscclAlgoHandle_t handle;
  ncclComm_t comm;
  hipStream_t stream;
 };

+struct mscclThreadLocalStatus {
+  bool mscclIsCallerFlag;
+  mscclGroupStatus groupStatus;
+  int groupDepth;
+  std::vector<struct mscclSavedSchedulerParam> savedSchedulerParams;
+};
+
 struct mscclStatus {
  std::vector<mscclAlgoHandle_t> freeAlgoHandles;
  std::map<mscclAlgoHandle_t, mscclAlgo *> hostAlgos;
@@ -177,13 +180,15 @@ struct mscclStatus {
  int sliceSteps;
  int chunkSize;
  int chunkEffectiveSize;
-  int rank;
  uint32_t workIndex;
  uint32_t maxAllowedCount;
  ncclDataType_t dataType;
-  mscclGroupStatus groupStatus;
-  int groupDepth;
-  std::vector<struct mscclSchedulerParam> savedSchedulerParams;
+  std::map<ncclComm_t, std::set<mscclAlgoHandle_t>> connectedAlgos;
+  hipStream_t lastStream;
+  void* mscclSchedulerLib;
+  mscclSchedulerInterface* mscclSchedulerPtr;
+  std::vector<mscclAlgoMeta> algoMetas;
+  std::vector<std::map<int, mscclAlgoHandle_t>> rankToAlgoHandles;
 };

 struct alignas(16) mscclWork {