Clique kernel support (#295) (#15)

* Adding experimental clique-based kernels (opt-in only) Co-authored-by: Stanley Tsang <stanley.tsang@amd.com> Co-authored-by: Gilbert Lee <gilbert.lee@amd.com> Co-authored-by: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Co-authored-by: Stanley Tsang <stanley.tsang@amd.com> Co-authored-by: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> [ROCm/rccl commit: 3e62ceddc5]
2021-01-28 09:45:01 -07:00
Commit c981e76efe
@@ -1,6 +1,6 @@
 /*************************************************************************
 * Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2020 Advanced Micro Devices, Inc. All rights reserved.
+ * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
 *
 * See LICENSE.txt for license information
 ************************************************************************/
@@ -9,6 +9,8 @@
 #include "argcheck.h"
 #include "coll_net.h"
 #include "graph/topo.h"
+#include <hip/hip_runtime.h>
+#include <hip/hip_ext.h>

 // Only generate inline kernels for LL
 #define NCCL_FUNC5(func, algo, redop, dtype) \
@@ -126,6 +128,10 @@ static ncclResult_t setupLaunch(struct ncclComm* comm, hipLaunchParams* params)
    STORE(&channel->workFifo[(channel->workFifoTail-1)%NCCL_MAX_OPS].elems[0].active, 2);
  }

+  { // [RCCL] Wait for any clique-based collectives
+    NCCLCHECK(comm->cliqueManager->WaitForPointers());
+  } // [/RCCL]
+
  // Find the first operation, choose the kernel accordingly and pass it
  // as the first argument.
  struct ncclChannel* c0 = comm->channels;
@@ -224,7 +230,8 @@ ncclResult_t ncclBarrierEnqueueWait(ncclComm_t comm) {
        (comm->launchMode == ncclComm::GROUP && comm->groupCudaStream) ? "/Stream" : "");
  }

-
+  hipEvent_t startEvent;
+  hipEvent_t stopEvent;
  if (comm->launchMode == ncclComm::PARALLEL) {
    hipLaunchKernelGGL(((void (*)(struct ncclWorkElem))params->func), params->gridDim, params->blockDim, params->sharedMem, params->stream, **((struct ncclWorkElem**)params->args));
  } else {
@@ -283,6 +290,7 @@ static ncclResult_t getAlgoInfo(struct ncclInfo* info) {
  info->algorithm = -1;
  info->protocol = -1;
  int nAlgos = NCCL_NUM_ALGORITHMS;
+
  // Check collNet support
  int collNetTypeSupport = 0;
  if (info->comm->collNetSupport)
@@ -321,7 +329,7 @@ static ncclResult_t getAlgoInfo(struct ncclInfo* info) {
  }
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
 #else
-  if (info->protocol == NCCL_PROTO_SIMPLE) nt += WARP_SIZE; // Extra warp for sync 
+  if (info->protocol == NCCL_PROTO_SIMPLE) nt += WARP_SIZE; // Extra warp for sync
  if (info->protocol == NCCL_PROTO_SIMPLE && info->algorithm == NCCL_ALGO_TREE) nt += WARP_SIZE;
 #endif
  info->nChannels = nc;
@@ -401,6 +409,33 @@ static ncclResult_t computeColl(struct ncclInfo* info /* input */, struct ncclWo

  work->funcIndex = FUNC_INDEX(info->coll, info->op, info->datatype, info->algorithm, info->protocol);

+  { // [RCCL] Check for clique-based kernel support
+    if (info->comm->cliqueManager->IsSupported(info->coll,
+                                               info->count,
+                                               info->datatype,
+                                               info->op))
+    {
+      // Declare the input / output pointers being used (to exchange via IPC with other ranks)
+      NCCLCHECK(info->comm->cliqueManager->DeclarePointers(info->comm->opCount,
+                                                           info->sendbuff,
+                                                           info->recvbuff));
+
+
+      info->algorithm = NCCL_ALGO_RING;
+      info->protocol = NCCL_PROTO_CLIQUE;
+      // Determine the number of channels to use for clique-kernel
+      NCCLCHECK(info->comm->cliqueManager->GetNumChannelsToUse(info->coll,
+                                                               info->count,
+                                                               info->datatype,
+                                                               info->op,
+                                                               info->comm->nChannels,
+                                                               &work->clique.nChannels));
+      work->clique.count = info->count;
+      work->funcIndex = FUNC_INDEX(info->coll, info->op, info->datatype, info->algorithm, info->protocol);
+      return ncclSuccess;
+    }
+  } // [RCCL]
+
  int stepSize   = info->comm->buffSizes[info->protocol]/NCCL_STEPS;
  int chunkSteps = (info->protocol == NCCL_PROTO_SIMPLE && info->algorithm == NCCL_ALGO_RING) ? info->chunkSteps : 1;
  int sliceSteps = (info->protocol == NCCL_PROTO_SIMPLE && info->algorithm == NCCL_ALGO_RING) ? info->sliceSteps : 1;
@@ -487,7 +522,7 @@ ncclResult_t ncclSaveKernel(struct ncclInfo* info) {

  info->comm->myParams->blockDim.x = std::max<unsigned>(info->comm->myParams->blockDim.x, info->nThreads);

-  int nChannels = work.coll.nChannels;
+  int nChannels = (info->coll == ncclFuncAllToAllv) ? work.a2av.nChannels : work.coll.nChannels;
  int nSubChannels = (info->pattern == ncclPatternCollTreeUp || info->pattern == ncclPatternCollTreeDown) ? 2 : 1;

  for (int bid=0; bid<nChannels*nSubChannels; bid++) {
@@ -509,6 +544,13 @@ ncclResult_t ncclSaveKernel(struct ncclInfo* info) {
    } else {
      work.coll.bid = bid % nChannels;
    }
+
+    // [RCCL] Setup pointers to where all the input/output pointers will be
+    if (info->protocol == NCCL_PROTO_CLIQUE) {
+      NCCLCHECK(info->comm->cliqueManager->SetCliqueArgs(&work));
+    }
+    // [/RCCL]
+
    struct ncclWork* w;
    NCCLCHECK(getNextOp(channel, &w, &work));
    if (info->coll == ncclFuncAllToAllv) {
@@ -687,6 +729,7 @@ ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
    } else {
      NCCLCHECKGOTO(ncclSaveAsyncColl(info), ret, end);
    }
+
 end:
    if (savedDev != -1) CUDACHECK(hipSetDevice(savedDev));
    ncclAsyncErrCheck(ret);