Use a map to host scratch buffers (#1004)

* Use a map to host scratch buffers * Address review feedbacks. Deliberately keep mscclSetupScratch function. [ROCm/rccl commit: 293f0fb752]
2023-12-05 13:15:28 -06:00
parent 403a91137c
commit 33aa8b67be
3 changed files with 16 additions and 14 deletions
@@ -205,8 +205,7 @@ struct mscclStatus {
  std::map<mscclAlgoHandle_t, mscclAlgo *> hostAlgos;
  std::map<mscclAlgoHandle_t, mscclAlgo *> devAlgos;
  struct mscclFlag* syncFlags;
-  void *scratchBuffer;
-  uint64_t scratchBufferSize;
+  std::map<size_t, void *> scratchBuffers;
  size_t nBytes;
  int stepSize;
  int chunkSteps;