rccl/repos/community-x86_64/test.cpp

   1 #include <stdio.h>
   2 #include <hip/hip_runtime.h>
   3 #include <rccl/rccl.h>
   4 #include <math.h>
   5
   6 #define NDEV 1
   7
   8 int main(int argc, char* argv[])
   9 {
  10   ncclComm_t comms[NDEV];
  11
  12   int size = 100;
  13   int devs[NDEV];
  14   for(int i = 0; i < NDEV; i++){
  15       devs[i] = i;
  16   }
  17
  18   float* hsend = (float*) malloc(size * sizeof *hsend);
  19   float* hrecv = (float*) malloc(size * sizeof *hrecv);
  20
  21   for (int i = 0; i < size; i++) {
  22     hsend[i] = 1.0f;
  23     hrecv[i] = 0.0f;
  24   }
  25
  26   float* sendbuff[NDEV];
  27   float* recvbuff[NDEV];
  28   hipStream_t s[NDEV];
  29   for (int i = 0; i < NDEV; i++) {
  30     hipSetDevice(i);
  31     hipMalloc(sendbuff + i, size * sizeof(float));
  32     hipMalloc(recvbuff + i, size * sizeof(float));
  33     hipMemcpy(sendbuff[i], hsend, size * sizeof *hsend, hipMemcpyHostToDevice);
  34     hipMemcpy(recvbuff[i], hrecv, size * sizeof *hrecv, hipMemcpyHostToDevice);
  35     hipStreamCreate(s+i);
  36     hipDeviceSynchronize();
  37   }
  38
  39   ncclCommInitAll(comms, NDEV, devs);
  40
  41   ncclGroupStart();
  42   for (int i = 0; i < NDEV; i++)
  43     ncclAllReduce(&sendbuff[i][0], &recvbuff[i][0], size, ncclFloat, ncclSum,
  44         comms[i], s[i]);
  45   ncclGroupEnd();
  46
  47
  48   for (int i = 0; i < NDEV; i++) {
  49     hipSetDevice(i);
  50     hipStreamSynchronize(s[i]);
  51   }
  52
  53   hipSetDevice(0);
  54   hipMemcpy(hrecv, recvbuff[0], size * sizeof *hrecv, hipMemcpyDeviceToHost);
  55   hipDeviceSynchronize();
  56   for (int i = 0; i < size; i++) {
  57     if(fabsf(NDEV * hsend[i] - hrecv[i]) > 0.001f){
  58         printf("At entry %i\n", i);
  59         printf("Expected %f\n", hsend[i]);
  60         printf("Received %f\n", hrecv[i]);
  61         return EXIT_FAILURE;
  62     }
  63   }
  64   printf("TESTS PASSED!\n");
  65
  66   for (int i = 0; i < NDEV; i++) {
  67     hipSetDevice(i);
  68     hipFree(sendbuff[i]);
  69     hipFree(recvbuff[i]);
  70   }
  71
  72   for(int i = 0; i < NDEV; ++i){
  73       ncclCommDestroy(comms[i]);
  74   }
  75
  76   free(hsend);
  77   free(hrecv);
  78 }