From 22a977e730148080f0774071c525dce69c506406 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Thu, 30 Mar 2023 06:29:38 +0000
Subject: [PATCH 01/23] init

---
 Makefile                 |   4 +-
 tests/allgather_test3.cu |  93 +++++++++++++++++
 tests/common.cu          | 212 +++++++++++++++++++++++++++++++++++++++
 tests/common.h           | 152 +++++++++++++++++++++++++++-
 4 files changed, 457 insertions(+), 4 deletions(-)
 create mode 100644 tests/allgather_test3.cu
 create mode 100644 tests/common.cu

diff --git a/Makefile b/Makefile
index 8a2fd327..2793b7bb 100644
--- a/Makefile
+++ b/Makefile
@@ -129,7 +129,7 @@ LIBSONAME := $(LIBNAME).$(MSCCLPP_MAJOR)
 LIBTARGET := $(BUILDDIR)/$(LIBDIR)/$(LIBNAME).$(MSCCLPP_MAJOR).$(MSCCLPP_MINOR).$(MSCCLPP_PATCH)
 
 TESTSDIR  := tests
-TESTSSRCS := $(addprefix $(TESTSDIR)/,bootstrap_test.cc allgather_test.cu)
+TESTSSRCS := $(addprefix $(TESTSDIR)/,bootstrap_test.cc allgather_test.cu common.cu)
 TESTSOBJS := $(patsubst %.cc,%.o,$(TESTSSRCS)) $(patsubst %.cu,%.o,$(TESTSSRCS))
 TESTSOBJTARGETS := $(TESTSOBJS:%=$(BUILDDIR)/$(OBJDIR)/%)
 TESTSBINS       := $(patsubst %.o,$(BUILDDIR)/$(BINDIR)/%,$(TESTSOBJS))
@@ -179,7 +179,7 @@ $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/%.o: $(TESTSDIR)/%.cc $(INCTARGETS)
 # Compile .cu tests
 $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/%.o: $(TESTSDIR)/%.cu $(INCTARGETS)
 	@mkdir -p $(@D)
-	$(NVCC) -o $@ -I$(BUILDDIR)/$(INCDIR) $(MPI_INC) $(NVCUFLAGS) -c $< $(MPI_MACRO)
+	$(NVCC) -o $@ -I$(BUILDDIR)/$(INCDIR) $(MPI_INC) $(NVCUFLAGS) $(INCLUDE) -c $< $(MPI_MACRO)
 
 # Test bins
 $(BUILDDIR)/$(BINDIR)/%: $(BUILDDIR)/$(OBJDIR)/%.o $(LIBTARGET)
diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
new file mode 100644
index 00000000..ce75864b
--- /dev/null
+++ b/tests/allgather_test3.cu
@@ -0,0 +1,93 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include "cuda_runtime.h"
+#include "common.h"
+
+#define ALIGN 4
+
+void AllGatherGetCollByteCount(size_t *sendcount, size_t *recvcount, size_t *paramcount, size_t *sendInplaceOffset, size_t *recvInplaceOffset, size_t count, int nranks) {
+  size_t base = (count/(ALIGN*nranks))*ALIGN;
+  *sendcount = base;
+  *recvcount = base*nranks;
+  *sendInplaceOffset = base;
+  *recvInplaceOffset = 0;
+  *paramcount = base;
+}
+
+testResult_t AllGatherInitData(struct threadArgs* args, ncclDataType_t type, ncclRedOp_t op, int root, int rep, int in_place) {
+  size_t sendcount = args->sendBytes / wordSize(type);
+  size_t recvcount = args->expectedBytes / wordSize(type);
+  int nranks = args->nProcs*args->nThreads*args->nGpus;
+
+  for (int i=0; i<args->nGpus; i++) {
+    CUDACHECK(cudaSetDevice(args->gpus[i]));
+    int rank = ((args->proc*args->nThreads + args->thread)*args->nGpus + i);
+    CUDACHECK(cudaMemset(args->recvbuffs[i], 0, args->expectedBytes));
+    void* data = in_place ? ((char*)args->recvbuffs[i])+rank*args->sendBytes : args->sendbuffs[i];
+    TESTCHECK(InitData(data, sendcount, 0, type, ncclSum, 33*rep + rank, 1, 0));
+    for (int j=0; j<nranks; j++) {
+      TESTCHECK(InitData((char*)args->expected[i] + args->sendBytes*j, sendcount, 0, type, ncclSum, 33*rep + j, 1, 0));
+    }
+    CUDACHECK(cudaDeviceSynchronize());
+  }
+  return testSuccess;
+}
+
+void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks) {
+  double baseBw = (double)(count * typesize * nranks) / 1.0E9 / sec;
+
+  *algBw = baseBw;
+  double factor = ((double)(nranks - 1))/((double)nranks);
+  *busBw = baseBw * factor;
+}
+
+testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count, ncclDataType_t type, ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream) {
+  NCCLCHECK(ncclAllGather(sendbuff, recvbuff, count, type, comm, stream));
+  return testSuccess;
+}
+
+struct testColl allGatherTest = {
+  "AllGather",
+  AllGatherGetCollByteCount,
+  AllGatherInitData,
+  AllGatherGetBw,
+  AllGatherRunColl
+};
+
+void AllGatherGetBuffSize(size_t *sendcount, size_t *recvcount, size_t count, int nranks) {
+  size_t paramcount, sendInplaceOffset, recvInplaceOffset;
+  AllGatherGetCollByteCount(sendcount, recvcount, &paramcount, &sendInplaceOffset, &recvInplaceOffset, count, nranks);
+}
+
+testResult_t AllGatherRunTest(struct threadArgs* args, ncclDataType_t type, const char* typeName, ncclRedOp_t op, const char* opName) {
+  args->collTest = &allGatherTest;
+  ncclDataType_t *run_types;
+  const char **run_typenames;
+  int type_count;
+
+  if ((int)type != -1) {
+    type_count = 1;
+    run_types = &type;
+    run_typenames = &typeName;
+  } else {
+    type_count = test_typenum;
+    run_types = test_types;
+    run_typenames = test_typenames;
+  }
+
+  for (int i=0; i<type_count; i++) {
+    TESTCHECK(TimeTest(args, run_types[i], run_typenames[i], (ncclRedOp_t)0, "none", -1));
+  }
+  return testSuccess;
+}
+
+struct testEngine allGatherEngine = {
+  AllGatherGetBuffSize,
+  AllGatherRunTest
+};
+
+#pragma weak mcclppTestEngine=allGatherEngine
\ No newline at end of file
diff --git a/tests/common.cu b/tests/common.cu
new file mode 100644
index 00000000..beffee0a
--- /dev/null
+++ b/tests/common.cu
@@ -0,0 +1,212 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#define MSCCLPP_USE_MPI_FOR_TESTS 1
+
+#include "common.h"
+#include "cuda.h"
+
+#include <cstdio>
+#include <iostream>
+#include <type_traits>
+
+#include <getopt.h>
+#include <libgen.h>
+#include <pthread.h>
+
+int is_main_proc = 0;
+thread_local int is_main_thread = 0;
+
+// Command line parameter defaults
+static int nThreads = 1;
+static int nGpus = 1;
+static size_t minBytes = 32*1024*1024;
+static size_t maxBytes = 32*1024*1024;
+static size_t stepBytes = 1*1024*1024;
+static size_t stepFactor = 1;
+static int datacheck = 1;
+static int warmup_iters = 5;
+static int iters = 20;
+static int timeout = 0;
+static int report_cputime = 0;
+// Report average iteration time: (0=RANK0,1=AVG,2=MIN,3=MAX)
+static int average = 1;
+
+#define NUM_BLOCKS 32
+
+static double parsesize(const char *value) {
+    long long int units;
+    double size;
+    char size_lit;
+
+    int count = sscanf(value, "%lf %1s", &size, &size_lit);
+
+    switch (count) {
+    case 2:
+      switch (size_lit) {
+      case 'G':
+      case 'g':
+        units = 1024*1024*1024;
+        break;
+      case 'M':
+      case 'm':
+        units = 1024*1024;
+        break;
+      case 'K':
+      case 'k':
+        units = 1024;
+        break;
+      default:
+        return -1.0;
+      };
+      break;
+    case 1:
+      units = 1;
+      break;
+    default:
+      return -1.0;
+    }
+
+    return size * units;
+}
+
+testResult_t run(); // Main function
+
+int main(int argc, char* argv[]) {
+  // Make sure everyline is flushed so that we see the progress of the test
+  setlinebuf(stdout);
+
+  // Parse args
+  double parsed;
+  int longindex;
+  static struct option longopts[] = {
+    {"nthreads", required_argument, 0, 't'},
+    {"ngpus", required_argument, 0, 'g'},
+    {"minbytes", required_argument, 0, 'b'},
+    {"maxbytes", required_argument, 0, 'e'},
+    {"stepbytes", required_argument, 0, 'i'},
+    {"stepfactor", required_argument, 0, 'f'},
+    {"iters", required_argument, 0, 'n'},
+    {"agg_iters", required_argument, 0, 'm'},
+    {"warmup_iters", required_argument, 0, 'w'},
+    {"check", required_argument, 0, 'c'},
+    {"timeout", required_argument, 0, 'T'},
+    {"report_cputime", required_argument, 0, 'C'},
+    {"average", required_argument, 0, 'a'},
+    {"help", no_argument, 0, 'h'},
+    {}
+  };
+
+  while(1) {
+    int c;
+    c = getopt_long(argc, argv, "t:g:b:e:i:f:n:w:c:o:T:h:C:a:", longopts, &longindex);
+
+    if (c == -1)
+      break;
+
+    switch(c) {
+      case 't':
+        nThreads = strtol(optarg, NULL, 0);
+        break;
+      case 'g':
+        nGpus = strtol(optarg, NULL, 0);
+        break;
+      case 'b':
+        parsed = parsesize(optarg);
+        if (parsed < 0) {
+          fprintf(stderr, "invalid size specified for 'minbytes'\n");
+          return -1;
+        }
+        minBytes = (size_t)parsed;
+        break;
+      case 'e':
+        parsed = parsesize(optarg);
+        if (parsed < 0) {
+          fprintf(stderr, "invalid size specified for 'maxbytes'\n");
+          return -1;
+        }
+        maxBytes = (size_t)parsed;
+        break;
+      case 'i':
+        stepBytes = strtol(optarg, NULL, 0);
+        break;
+      case 'f':
+        stepFactor = strtol(optarg, NULL, 0);
+        break;
+      case 'n':
+        iters = (int)strtol(optarg, NULL, 0);
+        break;
+      case 'w':
+        warmup_iters = (int)strtol(optarg, NULL, 0);
+        break;
+      case 'c':
+        datacheck = (int)strtol(optarg, NULL, 0);
+        break;
+      case 'T':
+        timeout = strtol(optarg, NULL, 0);
+        break;
+      case 'C':
+        report_cputime = strtol(optarg, NULL, 0);
+        break;
+      case 'a':
+        average = (int)strtol(optarg, NULL, 0);
+        break;
+      case 'h':
+      default:
+        if (c != 'h') printf("invalid option '%c'\n", c);
+        printf("USAGE: %s \n\t"
+            "[-t,--nthreads <num threads>] \n\t"
+            "[-g,--ngpus <gpus per thread>] \n\t"
+            "[-b,--minbytes <min size in bytes>] \n\t"
+            "[-e,--maxbytes <max size in bytes>] \n\t"
+            "[-i,--stepbytes <increment size>] \n\t"
+            "[-f,--stepfactor <increment factor>] \n\t"
+            "[-n,--iters <iteration count>] \n\t"
+            "[-w,--warmup_iters <warmup iteration count>] \n\t"
+            "[-c,--check <0/1>] \n\t"
+            "[-T,--timeout <time in seconds>] \n\t"
+            "[-C,--report_cputime <0/1>] \n\t"
+            "[-a,--average <0/1/2/3> report average iteration time <0=RANK0/1=AVG/2=MIN/3=MAX>] \n\t"
+            "[-h,--help]\n",
+          basename(argv[0]));
+        return 0;
+    }
+  }
+  if (minBytes > maxBytes) {
+    fprintf(stderr, "invalid sizes for 'minbytes' and 'maxbytes': %llu > %llu\n",
+           (unsigned long long)minBytes,
+           (unsigned long long)maxBytes);
+    return -1;
+  }
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+  MPI_Init(&argc, &argv);
+#endif
+  TESTCHECK(run());
+  return 0;
+}
+
+testResult_t run() {
+  int worldSize = 1, rank = 0;
+  int ranksPerNode = 0;
+
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+  MPI_Comm_size(MPI_COMM_WORLD, &worldSize);
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm shmcomm;
+  MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &shmcomm);
+  MPI_Comm_size(shmcomm, &ranksPerNode);
+  MPI_Comm_free(&shmcomm);
+#endif
+  is_main_thread = is_main_proc = (rank == 0) ? 1 : 0;
+    is_main_thread = is_main_proc = (rank == 0) ? 1 : 0;
+
+    PRINT("# nThread %d nGpus %d minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d\n",
+          nThreads, nGpus, minBytes, maxBytes, (stepFactor > 1) ? stepFactor : stepBytes,
+          (stepFactor > 1) ? "factor" : "bytes", warmup_iters, iters, datacheck);
+    PRINT("#\n");
+    PRINT("# Using devices\n");
+    return testSuccess;
+}
\ No newline at end of file
diff --git a/tests/common.h b/tests/common.h
index 11b5527f..970aa8af 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -1,13 +1,157 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
 #ifndef MSCCLPP_TESTS_COMMON_H_
 #define MSCCLPP_TESTS_COMMON_H_
 
-#include <stdio.h>
-#include <stdlib.h>
+#include <cstdio>
+#include <cstdlib>
+
+#include <unistd.h>
 
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
 #include "mpi.h"
 #endif // MSCCLPP_USE_MPI_FOR_TESTS
 
+// Relay errors up and trace
+#define TESTCHECK(cmd)                                                                                                 \
+  do {                                                                                                                 \
+    testResult_t r = cmd;                                                                                              \
+    if (r != testSuccess) {                                                                                            \
+      char hostname[1024];                                                                                             \
+      getHostName(hostname, 1024);                                                                                     \
+      printf(" .. %s pid %d: Test failure %s:%d\n", hostname, getpid(), __FILE__, __LINE__);                           \
+      return r;                                                                                                        \
+    }                                                                                                                  \
+  } while (0)
+
+typedef enum {
+  testSuccess = 0,
+  testInternalError = 1,
+  testCudaError = 2,
+  testMcclppError = 3,
+  testTimeout = 4,
+  testNumResults = 5
+} testResult_t;
+
+struct testColl {
+  const char name[20];
+  void (*getCollByteCount)(
+      size_t *sendcount, size_t *recvcount, size_t *paramcount,
+      size_t *sendInplaceOffset, size_t *recvInplaceOffset,
+      size_t count, int nranks);
+  testResult_t (*initData)(struct threadArgs* args, int root, int rep, int in_place);
+  void (*getBw)(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks);
+  testResult_t (*runColl)(void* sendbuff, void* recvbuff, size_t count);
+};
+
+struct testEngine
+{
+  void (*getBuffSize)(size_t* sendcount, size_t* recvcount, size_t count, int nranks);
+  testResult_t (*runTest)(struct threadArgs* args, int root, const char* typeName, const char* opName);
+};
+
+extern struct testEngine mscclppTestEngine;
+
+struct threadArgs
+{
+  size_t nbytes;
+  size_t minbytes;
+  size_t maxbytes;
+  size_t stepbytes;
+  size_t stepfactor;
+
+  int totalProcs;
+  int nProcs;
+  int proc;
+  int nThreads;
+  int thread;
+  int nGpus;
+  int* gpus;
+  int localRank;
+  void** sendbuffs;
+  size_t sendBytes;
+  size_t sendInplaceOffset;
+  void** recvbuffs;
+  size_t recvInplaceOffset;
+  cudaStream_t* streams;
+
+  void** expected;
+  size_t expectedBytes;
+  int* errors;
+  double* bw;
+  int* bw_count;
+
+  int reportErrors;
+
+  struct testColl* collTest;
+};
+
+typedef testResult_t (*threadFunc_t)(struct threadArgs* args);
+struct testThread
+{
+  pthread_t thread;
+  threadFunc_t func;
+  struct threadArgs args;
+  testResult_t ret;
+};
+
+static void getHostName(char* hostname, int maxlen)
+{
+  gethostname(hostname, maxlen);
+  for (int i = 0; i < maxlen; i++) {
+    if (hostname[i] == '.') {
+      hostname[i] = '\0';
+      return;
+    }
+  }
+}
+
+// getHash and getHostHash are duplicated with utils.h
+// since libmscclpp not export this function, copy here to avoid link issue
+static uint64_t getHash(const char* string, size_t n) {
+  // Based on DJB2a, result = result * 33 ^ char
+  uint64_t result = 5381;
+  for (size_t c = 0; c < n; c++) {
+    result = ((result << 5) + result) ^ string[c];
+  }
+  return result;
+}
+
+/* Generate a hash of the unique identifying string for this host
+ * that will be unique for both bare-metal and container instances
+ * Equivalent of a hash of;
+ *
+ * $(hostname)$(cat /proc/sys/kernel/random/boot_id)
+ *
+ */
+#define HOSTID_FILE "/proc/sys/kernel/random/boot_id"
+static uint64_t getHostHash(const char* hostname) {
+  char hostHash[1024];
+
+  // Fall back is the hostname if something fails
+  (void) strncpy(hostHash, hostname, sizeof(hostHash));
+  int offset = strlen(hostHash);
+
+  FILE *file = fopen(HOSTID_FILE, "r");
+  if (file != NULL) {
+    char *p;
+    if (fscanf(file, "%ms", &p) == 1) {
+        strncpy(hostHash+offset, p, sizeof(hostHash)-offset-1);
+        free(p);
+    }
+  }
+  fclose(file);
+
+  // Make sure the string is terminated
+  hostHash[sizeof(hostHash)-1]='\0';
+
+  return getHash(hostHash, strlen(hostHash));
+}
+
 void print_usage(const char* prog)
 {
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
@@ -43,4 +187,8 @@ void parse_arguments(int argc, const char* argv[], const char** ip_port, int* ra
 #endif
 }
 
+#define PRINT                                                                                                          \
+  if (is_main_thread)                                                                                                  \
+  printf
+
 #endif // MSCCLPP_TESTS_COMMON_H_
\ No newline at end of file

From 20b644b88d81a7310c00ba3d6f9eb0520c0a9142 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Thu, 30 Mar 2023 09:13:28 +0000
Subject: [PATCH 02/23] port more code from nccl-test

---
 tests/common.cu | 525 ++++++++++++++++++++++++++++++++++++++++++++++--
 tests/common.h  |  75 +++----
 2 files changed, 539 insertions(+), 61 deletions(-)

diff --git a/tests/common.cu b/tests/common.cu
index beffee0a..5f6bf014 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -8,9 +8,11 @@
 
 #include "common.h"
 #include "cuda.h"
+#include "mscclpp.h"
 
 #include <cstdio>
 #include <iostream>
+#include <string>
 #include <type_traits>
 
 #include <getopt.h>
@@ -34,6 +36,7 @@ static int timeout = 0;
 static int report_cputime = 0;
 // Report average iteration time: (0=RANK0,1=AVG,2=MIN,3=MAX)
 static int average = 1;
+static std::string ip_port;
 
 #define NUM_BLOCKS 32
 
@@ -73,6 +76,334 @@ static double parsesize(const char *value) {
     return size * units;
 }
 
+void Barrier(struct threadArgs *args) {
+  thread_local int epoch = 0;
+  static pthread_mutex_t lock[2] = {PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER};
+  static pthread_cond_t cond[2] = {PTHREAD_COND_INITIALIZER, PTHREAD_COND_INITIALIZER};
+  static int counter[2] = {0, 0};
+
+  pthread_mutex_lock(&lock[epoch]);
+  if(++counter[epoch] == args->nThreads)
+    pthread_cond_broadcast(&cond[epoch]);
+
+  if(args->thread+1 == args->nThreads) {
+    while(counter[epoch] != args->nThreads)
+      pthread_cond_wait(&cond[epoch], &lock[epoch]);
+    #ifdef MPI_SUPPORT
+      MPI_Barrier(MPI_COMM_WORLD);
+    #endif
+    counter[epoch] = 0;
+    pthread_cond_broadcast(&cond[epoch]);
+  }
+  else {
+    while(counter[epoch] != 0)
+      pthread_cond_wait(&cond[epoch], &lock[epoch]);
+  }
+  pthread_mutex_unlock(&lock[epoch]);
+  epoch ^= 1;
+}
+
+testResult_t AllocateBuffs(void **sendbuff, size_t sendBytes, void **recvbuff, size_t recvBytes, void **expected, size_t nbytes) {
+    CUDACHECK(cudaMalloc(sendbuff, nbytes));
+    CUDACHECK(cudaMalloc(recvbuff, nbytes));
+    if (datacheck) CUDACHECK(cudaMalloc(expected, recvBytes));
+    return testSuccess;
+}
+
+testResult_t startColl(struct threadArgs* args, ncclDataType_t type, ncclRedOp_t opIndex, int root, int in_place, int iter) {
+  size_t count = args->nbytes / wordSize(type);
+
+  // Try to change offset for each iteration so that we avoid cache effects and catch race conditions in ptrExchange
+  size_t totalnbytes = max(args->sendBytes, args->expectedBytes);
+  size_t steps = totalnbytes ? args->maxbytes / totalnbytes : 1;
+  size_t shift = totalnbytes * (iter % steps);
+
+  if (args->nGpus > 1) NCCLCHECK(ncclGroupStart());
+  for (int i = 0; i < args->nGpus; i++) {
+#ifndef NCCL_MAJOR
+    CUDACHECK(cudaSetDevice(args->gpus[i]));
+#endif
+    int rank = ((args->proc*args->nThreads + args->thread)*args->nGpus + i);
+    char* recvBuff = ((char*)args->recvbuffs[i]) + shift;
+    char* sendBuff = ((char*)args->sendbuffs[i]) + shift;
+    ncclRedOp_t op;
+
+    if(opIndex < ncclNumOps) {
+      op = opIndex;
+    }
+    #if NCCL_VERSION_CODE >= NCCL_VERSION(2,11,0)
+    else {
+      union {
+        int8_t i8; uint8_t u8; int32_t i32; uint32_t u32; int64_t i64; uint64_t u64;
+        half f16; float f32; double f64;
+        #if defined(__CUDA_BF16_TYPES_EXIST__)
+        __nv_bfloat16 bf16;
+        #endif
+      };
+      switch(type) {
+      case ncclInt8: i8 = ncclVerifiablePremulScalar<int8_t>(rank); break;
+      case ncclUint8: u8 = ncclVerifiablePremulScalar<uint8_t>(rank); break;
+      case ncclInt32: i32 = ncclVerifiablePremulScalar<int32_t>(rank); break;
+      case ncclUint32: u32 = ncclVerifiablePremulScalar<uint32_t>(rank); break;
+      case ncclInt64: i64 = ncclVerifiablePremulScalar<int64_t>(rank); break;
+      case ncclUint64: u64 = ncclVerifiablePremulScalar<uint64_t>(rank); break;
+      case ncclFloat16: f16 = ncclVerifiablePremulScalar<half>(rank); break;
+      case ncclFloat32: f32 = ncclVerifiablePremulScalar<float>(rank); break;
+      case ncclFloat64: f64 = ncclVerifiablePremulScalar<double>(rank); break;
+      #if defined(__CUDA_BF16_TYPES_EXIST__)
+      case ncclBfloat16: bf16 = ncclVerifiablePremulScalar<__nv_bfloat16>(rank); break;
+      #endif
+      }
+      NCCLCHECK(ncclRedOpCreatePreMulSum(&op, &u64, type, ncclScalarHostImmediate, args->comms[i]));
+    }
+    #endif
+
+    TESTCHECK(args->collTest->runColl(
+          (void*)(in_place ? recvBuff + args->sendInplaceOffset*rank : sendBuff),
+          (void*)(in_place ? recvBuff + args->recvInplaceOffset*rank : recvBuff),
+        count, type, op, root, args->comms[i], args->streams[i]));
+
+    #if NCCL_VERSION_CODE >= NCCL_VERSION(2,11,0)
+    if(opIndex >= ncclNumOps) {
+      NCCLCHECK(ncclRedOpDestroy(op, args->comms[i]));
+    }
+    #endif
+  }
+  if (args->nGpus > 1) NCCLCHECK(ncclGroupEnd());
+
+  if (blocking_coll) {
+    // Complete op before returning
+    TESTCHECK(testStreamSynchronize(args->nGpus, args->streams, args->comms));
+  }
+  if (blocking_coll) Barrier(args);
+  return testSuccess;
+}
+
+testResult_t completeColl(struct threadArgs* args) {
+  if (blocking_coll) return testSuccess;
+
+  TESTCHECK(testStreamSynchronize(args->nGpus, args->streams, args->comms));
+  return testSuccess;
+}
+
+testResult_t BenchTime(struct threadArgs* args, ncclDataType_t type, ncclRedOp_t op, int root, int in_place) {
+  size_t count = args->nbytes / wordSize(type);
+  if (datacheck) {
+    // Initialize sendbuffs, recvbuffs and expected
+    TESTCHECK(args->collTest->initData(args, type, op, root, 99, in_place));
+  }
+
+  // Sync
+  TESTCHECK(startColl(args, type, op, root, in_place, 0));
+  TESTCHECK(completeColl(args));
+
+  Barrier(args);
+
+#if CUDART_VERSION >= 11030
+  cudaGraph_t graphs[args->nGpus];
+  cudaGraphExec_t graphExec[args->nGpus];
+  if (cudaGraphLaunches >= 1) {
+    // Begin cuda graph capture
+    for (int i=0; i<args->nGpus; i++) {
+      // Thread local mdoe is needed for:
+      // - Multi-thread mode: where graph capture and instantiation can happen concurrently across threads
+      // - P2P pre-connect: when there is no warm-up, P2P pre-connect is done during graph capture.
+      //   Since pre-connect calls cudaMalloc, we cannot use global capture mode
+      CUDACHECK(cudaStreamBeginCapture(args->streams[i], cudaStreamCaptureModeThreadLocal));
+    }
+  }
+#endif
+
+  // Performance Benchmark
+  timer tim;
+  for (int iter = 0; iter < iters; iter++) {
+    if (agg_iters>1) NCCLCHECK(ncclGroupStart());
+    for (int aiter = 0; aiter < agg_iters; aiter++) {
+      TESTCHECK(startColl(args, type, op, root, in_place, iter*agg_iters+aiter));
+    }
+    if (agg_iters>1) NCCLCHECK(ncclGroupEnd());
+  }
+
+#if CUDART_VERSION >= 11030
+  if (cudaGraphLaunches >= 1) {
+    // End cuda graph capture
+    for (int i=0; i<args->nGpus; i++) {
+      CUDACHECK(cudaStreamEndCapture(args->streams[i], graphs+i));
+    }
+    // Instantiate cuda graph
+    for (int i=0; i<args->nGpus; i++) {
+      CUDACHECK(cudaGraphInstantiate(graphExec+i, graphs[i], NULL, NULL, 0));
+    }
+    // Resync CPU, restart timing, launch cuda graph
+    Barrier(args);
+    tim.reset();
+    for (int l=0; l<cudaGraphLaunches; l++) {
+      for (int i=0; i<args->nGpus; i++) {
+        CUDACHECK(cudaGraphLaunch(graphExec[i], args->streams[i]));
+      }
+    }
+  }
+#endif
+
+  double cputimeSec = tim.elapsed()/(iters*agg_iters);
+  TESTCHECK(completeColl(args));
+
+  double deltaSec = tim.elapsed();
+  deltaSec = deltaSec/(iters*agg_iters);
+  if (cudaGraphLaunches >= 1) deltaSec = deltaSec/cudaGraphLaunches;
+  Allreduce(args, &deltaSec, average);
+
+#if CUDART_VERSION >= 11030
+  if (cudaGraphLaunches >= 1) {
+    //destroy cuda graph
+    for (int i=0; i<args->nGpus; i++) {
+      CUDACHECK(cudaGraphExecDestroy(graphExec[i]));
+      CUDACHECK(cudaGraphDestroy(graphs[i]));
+    }
+  }
+#endif
+
+  double algBw, busBw;
+  args->collTest->getBw(count, wordSize(type), deltaSec, &algBw, &busBw, args->nProcs*args->nThreads*args->nGpus);
+
+  Barrier(args);
+
+  int64_t wrongElts = 0;
+  static __thread int rep = 0;
+  rep++;
+  if (datacheck) {
+      // Initialize sendbuffs, recvbuffs and expected
+      TESTCHECK(args->collTest->initData(args, type, op, root, rep, in_place));
+
+#if CUDART_VERSION >= 11030
+      if (cudaGraphLaunches >= 1) {
+        // Begin cuda graph capture for data check
+        for (int i=0; i<args->nGpus; i++) {
+          CUDACHECK(cudaStreamBeginCapture(args->streams[i], args->nThreads > 1 ? cudaStreamCaptureModeThreadLocal : cudaStreamCaptureModeGlobal));
+        }
+      }
+#endif
+
+      //test validation in single itertion, should ideally be included into the multi-iteration run
+      TESTCHECK(startColl(args, type, op, root, in_place, 0));
+
+#if CUDART_VERSION >= 11030
+      if (cudaGraphLaunches >= 1) {
+        // End cuda graph capture
+        for (int i=0; i<args->nGpus; i++) {
+          CUDACHECK(cudaStreamEndCapture(args->streams[i], graphs+i));
+        }
+        // Instantiate cuda graph
+        for (int i=0; i<args->nGpus; i++) {
+          CUDACHECK(cudaGraphInstantiate(graphExec+i, graphs[i], NULL, NULL, 0));
+        }
+        // Launch cuda graph
+        for (int i=0; i<args->nGpus; i++) {
+          CUDACHECK(cudaGraphLaunch(graphExec[i], args->streams[i]));
+        }
+      }
+#endif
+
+      TESTCHECK(completeColl(args));
+
+#if CUDART_VERSION >= 11030
+      if (cudaGraphLaunches >= 1) {
+        //destroy cuda graph
+        for (int i=0; i<args->nGpus; i++) {
+          CUDACHECK(cudaGraphExecDestroy(graphExec[i]));
+          CUDACHECK(cudaGraphDestroy(graphs[i]));
+        }
+      }
+#endif
+
+      TESTCHECK(CheckData(args, type, op, root, in_place, &wrongElts));
+
+      //aggregate delta from all threads and procs
+      long long wrongElts1 = wrongElts;
+      Allreduce(args, &wrongElts1, /*sum*/4);
+      wrongElts = wrongElts1;
+  }
+
+  double timeUsec = (report_cputime ? cputimeSec : deltaSec)*1.0E6;
+  char timeStr[100];
+  if (timeUsec >= 10000.0) {
+    sprintf(timeStr, "%7.0f", timeUsec);
+  } else if (timeUsec >= 100.0) {
+    sprintf(timeStr, "%7.1f", timeUsec);
+  } else {
+    sprintf(timeStr, "%7.2f", timeUsec);
+  }
+  if (args->reportErrors) {
+    PRINT("  %7s  %6.2f  %6.2f  %5g", timeStr, algBw, busBw, (double)wrongElts);
+  } else {
+    PRINT("  %7s  %6.2f  %6.2f  %5s", timeStr, algBw, busBw, "N/A");
+  }
+
+  args->bw[0] += busBw;
+  args->bw_count[0]++;
+  return testSuccess;
+}
+
+void setupArgs(size_t size, ncclDataType_t type, struct threadArgs* args) {
+  int nranks = args->nProcs*args->nGpus*args->nThreads;
+  size_t count, sendCount, recvCount, paramCount, sendInplaceOffset, recvInplaceOffset;
+
+  count = size / wordSize(type);
+  args->collTest->getCollByteCount(&sendCount, &recvCount, &paramCount, &sendInplaceOffset, &recvInplaceOffset, (size_t)count, (size_t)nranks);
+
+  args->nbytes = paramCount * wordSize(type);
+  args->sendBytes = sendCount * wordSize(type);
+  args->expectedBytes = recvCount * wordSize(type);
+  args->sendInplaceOffset = sendInplaceOffset * wordSize(type);
+  args->recvInplaceOffset = recvInplaceOffset * wordSize(type);
+}
+
+testResult_t TimeTest(struct threadArgs* args, ncclDataType_t type, const char* typeName, ncclRedOp_t op, const char* opName, int root) {
+  // Sync to avoid first-call timeout
+  Barrier(args);
+
+  // Warm-up for large size
+  setupArgs(args->maxbytes, type, args);
+  for (int iter = 0; iter < warmup_iters; iter++) {
+    TESTCHECK(startColl(args, type, op, root, 0, iter));
+  }
+  TESTCHECK(completeColl(args));
+
+  // Warm-up for small size
+  setupArgs(args->minbytes, type, args);
+  for (int iter = 0; iter < warmup_iters; iter++) {
+    TESTCHECK(startColl(args, type, op, root, 0, iter));
+  }
+  TESTCHECK(completeColl(args));
+
+  // Benchmark
+  for (size_t size = args->minbytes; size<=args->maxbytes; size = ((args->stepfactor > 1) ? size*args->stepfactor : size+args->stepbytes)) {
+      setupArgs(size, type, args);
+      char rootName[100];
+      sprintf(rootName, "%6i", root);
+      PRINT("%12li  %12li  %8s  %6s  %6s", max(args->sendBytes, args->expectedBytes), args->nbytes / wordSize(type), typeName, opName, rootName);
+      TESTCHECK(BenchTime(args, type, op, root, 0));
+      TESTCHECK(BenchTime(args, type, op, root, 1));
+      PRINT("\n");
+  }
+  return testSuccess;
+}
+
+testResult_t threadRunTests(struct threadArgs* args) {
+  TESTCHECK(mscclppTestEngine.runTest(args, test_typenames[nccltype], (ncclRedOp_t)ncclop, test_opnames[ncclop]));
+  return testSuccess;
+}
+
+void* threadLauncher(void* thread_) {
+  struct testThread* thread = (struct testThread*)thread_;
+  thread->ret = thread->func(&thread->args);
+  return NULL;
+}
+testResult_t threadLaunch(struct testThread* thread) {
+  pthread_create(&thread->thread, NULL, threadLauncher, thread);
+  return testSuccess;
+}
+
 testResult_t run(); // Main function
 
 int main(int argc, char* argv[]) {
@@ -90,19 +421,19 @@ int main(int argc, char* argv[]) {
     {"stepbytes", required_argument, 0, 'i'},
     {"stepfactor", required_argument, 0, 'f'},
     {"iters", required_argument, 0, 'n'},
-    {"agg_iters", required_argument, 0, 'm'},
     {"warmup_iters", required_argument, 0, 'w'},
     {"check", required_argument, 0, 'c'},
     {"timeout", required_argument, 0, 'T'},
     {"report_cputime", required_argument, 0, 'C'},
     {"average", required_argument, 0, 'a'},
+    {"ip_port", required_argument, 0, 'P'},
     {"help", no_argument, 0, 'h'},
     {}
   };
 
   while(1) {
     int c;
-    c = getopt_long(argc, argv, "t:g:b:e:i:f:n:w:c:o:T:h:C:a:", longopts, &longindex);
+    c = getopt_long(argc, argv, "t:g:b:e:i:f:n:w:c:T:C:a:P:h:", longopts, &longindex);
 
     if (c == -1)
       break;
@@ -154,6 +485,9 @@ int main(int argc, char* argv[]) {
       case 'a':
         average = (int)strtol(optarg, NULL, 0);
         break;
+      case 'P':
+        ip_port = optarg;
+        break;
       case 'h':
       default:
         if (c != 'h') printf("invalid option '%c'\n", c);
@@ -170,15 +504,19 @@ int main(int argc, char* argv[]) {
             "[-T,--timeout <time in seconds>] \n\t"
             "[-C,--report_cputime <0/1>] \n\t"
             "[-a,--average <0/1/2/3> report average iteration time <0=RANK0/1=AVG/2=MIN/3=MAX>] \n\t"
+            "[-P,--ip_port <ip port for bootstrap>] \n\t"
             "[-h,--help]\n",
           basename(argv[0]));
         return 0;
     }
   }
   if (minBytes > maxBytes) {
-    fprintf(stderr, "invalid sizes for 'minbytes' and 'maxbytes': %llu > %llu\n",
-           (unsigned long long)minBytes,
-           (unsigned long long)maxBytes);
+    fprintf(stderr, "invalid sizes for 'minbytes' and 'maxbytes': %llu > %llu\n", (unsigned long long)minBytes,
+            (unsigned long long)maxBytes);
+    return -1;
+  }
+  if (ip_port.empty()) {
+    fprintf(stderr, "--ip_port is required'\n");
     return -1;
   }
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
@@ -189,24 +527,177 @@ int main(int argc, char* argv[]) {
 }
 
 testResult_t run() {
-  int worldSize = 1, rank = 0;
-  int ranksPerNode = 0;
+  int totalProcs = 1, proc = 0;
+  int ranksPerNode = 0, localRank = 0;
+  char hostname[1024];
+  getHostName(hostname, 1024);
 
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  MPI_Comm_size(MPI_COMM_WORLD, &worldSize);
-  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm_size(MPI_COMM_WORLD, &totalProcs);
+  MPI_Comm_rank(MPI_COMM_WORLD, &proc);
   MPI_Comm shmcomm;
   MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &shmcomm);
   MPI_Comm_size(shmcomm, &ranksPerNode);
   MPI_Comm_free(&shmcomm);
+  localRank = proc % ranksPerNode;
 #endif
-  is_main_thread = is_main_proc = (rank == 0) ? 1 : 0;
-    is_main_thread = is_main_proc = (rank == 0) ? 1 : 0;
+  is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
+  is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
 
-    PRINT("# nThread %d nGpus %d minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d\n",
-          nThreads, nGpus, minBytes, maxBytes, (stepFactor > 1) ? stepFactor : stepBytes,
-          (stepFactor > 1) ? "factor" : "bytes", warmup_iters, iters, datacheck);
-    PRINT("#\n");
-    PRINT("# Using devices\n");
-    return testSuccess;
+  PRINT("# nThread %d nGpus %d minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip "
+        "port: %s\n",
+        nThreads, nGpus, minBytes, maxBytes, (stepFactor > 1) ? stepFactor : stepBytes,
+        (stepFactor > 1) ? "factor" : "bytes", warmup_iters, iters, datacheck, ip_port.c_str());
+  PRINT("#\n");
+  PRINT("# Using devices\n");
+
+#define MAX_LINE 2048
+  char line[MAX_LINE];
+  int len = 0;
+  size_t maxMem = ~0;
+  for (int i = 0; i < nThreads * nGpus; i++) {
+    int cudaDev = (localRank * nThreads * nGpus) + i;
+    int rank = proc * nThreads * nGpus + i;
+    cudaDeviceProp prop;
+    CUDACHECK(cudaGetDeviceProperties(&prop, cudaDev));
+    len += snprintf(line + len, MAX_LINE - len, "#  Rank %2d Pid %6d on %10s device %2d [0x%02x] %s\n", rank, getpid(),
+                    hostname, cudaDev, prop.pciBusID, prop.name);
+    maxMem = std::min(maxMem, prop.totalGlobalMem);
+  }
+
+#if MSCCLPP_USE_MPI_FOR_TESTS
+  char* lines = (proc == 0) ? (char*)malloc(totalProcs * MAX_LINE) : NULL;
+  // Gather all output in rank order to root (0)
+  MPI_Gather(line, MAX_LINE, MPI_BYTE, lines, MAX_LINE, MPI_BYTE, 0, MPI_COMM_WORLD);
+  if (proc == 0) {
+    for (int p = 0; p < totalProcs; p++)
+        PRINT("%s", lines + MAX_LINE * p);
+    free(lines);
+  }
+  MPI_Allreduce(MPI_IN_PLACE, &maxMem, 1, MPI_LONG, MPI_MIN, MPI_COMM_WORLD);
+#else
+  PRINT("%s", line);
+#endif
+  // We need sendbuff, recvbuff, expected (when datacheck enabled), plus 1G for the rest.
+  size_t memMaxBytes = (maxMem - (1<<30)) / (datacheck ? 3 : 2);
+  if (maxBytes > memMaxBytes) {
+    maxBytes = memMaxBytes;
+    if (proc == 0) printf("#\n# Reducing maxBytes to %ld due to memory limitation\n", maxBytes);
+  }
+
+  int gpus[nGpus*nThreads];
+  cudaStream_t streams[nGpus*nThreads];
+  void* sendbuffs[nGpus*nThreads];
+  void* recvbuffs[nGpus*nThreads];
+  void* expected[nGpus*nThreads];
+  size_t sendBytes = 10, recvBytes = 10;
+
+  // mscclppTestEngine.getBuffSize(&sendBytes, &recvBytes, (size_t)maxBytes, (size_t)totalProcs*nGpus*nThreads);
+
+  for (int i = 0; i < nGpus * nThreads; i++) {
+    gpus[i] = (localRank * nThreads * nGpus) + i;
+    CUDACHECK(cudaSetDevice(gpus[i]));
+    TESTCHECK(AllocateBuffs(sendbuffs + i, sendBytes, recvbuffs + i, recvBytes, expected + i, (size_t)maxBytes));
+    CUDACHECK(cudaStreamCreateWithFlags(streams + i, cudaStreamNonBlocking));
+  }
+  PRINT("#\n");
+  PRINT("# Initializing MSCCL++\n");
+
+  mscclppComm_t* comms = (mscclppComm_t*)malloc(sizeof(mscclppComm_t)*nThreads*nGpus);
+  for (int i = 0; i < nGpus * nThreads; i++) {
+    MSCCLPPCHECK(
+      mscclppCommInitRank(comms + i, totalProcs * nGpus * nThreads, ip_port.c_str(), localRank * nThreads * nGpus + i));
+  }
+
+  int errors[nThreads];
+  double bw[nThreads];
+  double* delta;
+  CUDACHECK(cudaHostAlloc(&delta, sizeof(double) * nThreads * NUM_BLOCKS, cudaHostAllocPortable | cudaHostAllocMapped));
+  int bw_count[nThreads];
+  for (int t = 0; t < nThreads; t++) {
+    bw[t] = 0.0;
+    errors[t] = bw_count[t] = 0;
+  }
+
+  fflush(stdout);
+
+  const char* timeStr = report_cputime ? "cputime" : "time";
+  PRINT("#\n");
+  PRINT("# %10s  %12s  %8s  %6s  %6s           out-of-place                       in-place          \n", "", "", "", "",
+        "");
+  PRINT("# %10s  %12s  %8s  %6s  %6s  %7s  %6s  %6s %6s  %7s  %6s  %6s %6s\n", "size", "count", "type", "redop", "root",
+        timeStr, "algbw", "busbw", "#wrong", timeStr, "algbw", "busbw", "#wrong");
+  PRINT("# %10s  %12s  %8s  %6s  %6s  %7s  %6s  %6s  %5s  %7s  %6s  %6s  %5s\n", "(B)", "(elements)", "", "", "",
+        "(us)", "(GB/s)", "(GB/s)", "", "(us)", "(GB/s)", "(GB/s)", "");
+
+  struct testThread threads[nThreads];
+  memset(threads, 0, sizeof(struct testThread) * nThreads);
+
+  for (int t = nThreads - 1; t >= 0; t--) {
+    threads[t].args.minbytes = minBytes;
+    threads[t].args.maxbytes = maxBytes;
+    threads[t].args.stepbytes = stepBytes;
+    threads[t].args.stepfactor = stepFactor;
+    threads[t].args.localRank = localRank;
+
+    threads[t].args.totalProcs = totalProcs;
+    threads[t].args.proc = proc;
+    threads[t].args.nThreads = nThreads;
+    threads[t].args.thread = t;
+    threads[t].args.nGpus = nGpus;
+    threads[t].args.gpus = gpus + t * nGpus;
+    threads[t].args.sendbuffs = sendbuffs + t * nGpus;
+    threads[t].args.recvbuffs = recvbuffs + t * nGpus;
+    threads[t].args.expected = expected + t * nGpus;
+    threads[t].args.comms = comms + t * nGpus;
+    threads[t].args.streams = streams + t * nGpus;
+
+    threads[t].args.errors = errors + t;
+    threads[t].args.bw = bw + t;
+    threads[t].args.bw_count = bw_count + t;
+
+    threads[t].args.reportErrors = datacheck;
+
+    threads[t].func = threadRunTests;
+    if (t) {
+        TESTCHECK(threadLaunch(threads + t));
+    } else {
+        TESTCHECK(threads[t].func(&threads[t].args));
+    }
+  }
+
+  // Wait for other threads and accumulate stats and errors
+  for (int t = nThreads - 1; t >= 0; t--) {
+    if (t)
+        pthread_join(threads[t].thread, NULL);
+    TESTCHECK(threads[t].ret);
+    if (t) {
+        errors[0] += errors[t];
+        bw[0] += bw[t];
+        bw_count[0] += bw_count[t];
+    }
+  }
+
+  for (int i = 0; i < nGpus * nThreads; i++) {
+    MSCCLPPCHECK(mscclppCommDestroy(comms[i]));
+  }
+  free(comms);
+
+  // Free off CUDA allocated memory
+  for (int i=0; i<nGpus*nThreads; i++) {
+    if (sendbuffs[i]) CUDACHECK(cudaFree((char*)sendbuffs[i]));
+    if (recvbuffs[i]) CUDACHECK(cudaFree((char*)recvbuffs[i]));
+    if (datacheck) CUDACHECK(cudaFree(expected[i]));
+  }
+  CUDACHECK(cudaFreeHost(delta));
+
+  bw[0] /= bw_count[0];
+
+  PRINT("# Out of bounds values : %d %s\n", errors[0], errors[0] ? "FAILED" : "OK");
+  PRINT("#\n");
+
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+  MPI_Finalize();
+#endif
+  return testSuccess;
 }
\ No newline at end of file
diff --git a/tests/common.h b/tests/common.h
index 970aa8af..82f2af16 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -7,15 +7,41 @@
 #ifndef MSCCLPP_TESTS_COMMON_H_
 #define MSCCLPP_TESTS_COMMON_H_
 
+#include "mscclpp.h"
+
 #include <cstdio>
 #include <cstdlib>
 
 #include <unistd.h>
 
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
-#include "mpi.h"
+#include <mpi.h>
 #endif // MSCCLPP_USE_MPI_FOR_TESTS
 
+#define CUDACHECK(cmd) do {                         \
+  cudaError_t err = cmd;                            \
+  if( err != cudaSuccess ) {                        \
+    char hostname[1024];                            \
+    getHostName(hostname, 1024);                    \
+    printf("%s: Test CUDA failure %s:%d '%s'\n",    \
+         hostname,                                  \
+        __FILE__,__LINE__,cudaGetErrorString(err)); \
+    return testCudaError;                           \
+  }                                                 \
+} while(0)
+
+// Propagate errors up
+#define MSCCLPPCHECK(cmd)                                                                                              \
+  do {                                                                                                                 \
+    mscclppResult_t res = cmd;                                                                                         \
+    if (res != mscclppSuccess && res != mscclppInProgress) {                                                           \
+      char hostname[1024];                                                                                             \
+      getHostName(hostname, 1024);                                                                                     \
+      printf("%s: Failure at %s:%d -> %s\n", hostname, __FILE__, __LINE__, mscclppGetErrorString(res));                \
+      return testMcclppError;                                                                                          \
+    }                                                                                                                  \
+  } while (0);
+
 // Relay errors up and trace
 #define TESTCHECK(cmd)                                                                                                 \
   do {                                                                                                                 \
@@ -65,7 +91,6 @@ struct threadArgs
   size_t stepfactor;
 
   int totalProcs;
-  int nProcs;
   int proc;
   int nThreads;
   int thread;
@@ -77,6 +102,7 @@ struct threadArgs
   size_t sendInplaceOffset;
   void** recvbuffs;
   size_t recvInplaceOffset;
+  mscclppComm_t* comms;
   cudaStream_t* streams;
 
   void** expected;
@@ -99,6 +125,9 @@ struct testThread
   testResult_t ret;
 };
 
+// Provided by common.cu
+extern testResult_t TimeTest(struct threadArgs* args, ncclDataType_t type, const char* typeName, ncclRedOp_t op,  const char* opName, int root);
+
 static void getHostName(char* hostname, int maxlen)
 {
   gethostname(hostname, maxlen);
@@ -110,48 +139,6 @@ static void getHostName(char* hostname, int maxlen)
   }
 }
 
-// getHash and getHostHash are duplicated with utils.h
-// since libmscclpp not export this function, copy here to avoid link issue
-static uint64_t getHash(const char* string, size_t n) {
-  // Based on DJB2a, result = result * 33 ^ char
-  uint64_t result = 5381;
-  for (size_t c = 0; c < n; c++) {
-    result = ((result << 5) + result) ^ string[c];
-  }
-  return result;
-}
-
-/* Generate a hash of the unique identifying string for this host
- * that will be unique for both bare-metal and container instances
- * Equivalent of a hash of;
- *
- * $(hostname)$(cat /proc/sys/kernel/random/boot_id)
- *
- */
-#define HOSTID_FILE "/proc/sys/kernel/random/boot_id"
-static uint64_t getHostHash(const char* hostname) {
-  char hostHash[1024];
-
-  // Fall back is the hostname if something fails
-  (void) strncpy(hostHash, hostname, sizeof(hostHash));
-  int offset = strlen(hostHash);
-
-  FILE *file = fopen(HOSTID_FILE, "r");
-  if (file != NULL) {
-    char *p;
-    if (fscanf(file, "%ms", &p) == 1) {
-        strncpy(hostHash+offset, p, sizeof(hostHash)-offset-1);
-        free(p);
-    }
-  }
-  fclose(file);
-
-  // Make sure the string is terminated
-  hostHash[sizeof(hostHash)-1]='\0';
-
-  return getHash(hostHash, strlen(hostHash));
-}
-
 void print_usage(const char* prog)
 {
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS

From c7b3d0b0fd6c173424fed757c133e8e158b5645a Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 31 Mar 2023 03:58:09 +0000
Subject: [PATCH 03/23] update

---
 tests/allgather_test3.cu |  61 ++---
 tests/common.cu          | 513 +++++++++++++++------------------------
 tests/common.h           |   7 +-
 tests/timer.cc           |  28 +++
 tests/timer.h            |  15 ++
 5 files changed, 272 insertions(+), 352 deletions(-)
 create mode 100644 tests/timer.cc
 create mode 100644 tests/timer.h

diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
index ce75864b..5aad077c 100644
--- a/tests/allgather_test3.cu
+++ b/tests/allgather_test3.cu
@@ -9,31 +9,33 @@
 
 #define ALIGN 4
 
-void AllGatherGetCollByteCount(size_t *sendcount, size_t *recvcount, size_t *paramcount, size_t *sendInplaceOffset, size_t *recvInplaceOffset, size_t count, int nranks) {
-  size_t base = (count/(ALIGN*nranks))*ALIGN;
+void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* paramcount, size_t* sendInplaceOffset,
+                               size_t* recvInplaceOffset, size_t count, int nranks)
+{
+  size_t base = (count / (ALIGN * nranks)) * ALIGN;
   *sendcount = base;
-  *recvcount = base*nranks;
+  *recvcount = base * nranks;
   *sendInplaceOffset = base;
   *recvInplaceOffset = 0;
   *paramcount = base;
 }
 
-testResult_t AllGatherInitData(struct threadArgs* args, ncclDataType_t type, ncclRedOp_t op, int root, int rep, int in_place) {
-  size_t sendcount = args->sendBytes / wordSize(type);
-  size_t recvcount = args->expectedBytes / wordSize(type);
-  int nranks = args->nProcs*args->nThreads*args->nGpus;
+testResult_t AllGatherInitData(struct threadArgs* args, int in_place) {
+  size_t sendcount = args->sendBytes;
+  size_t recvcount = args->expectedBytes;
+  int nranks = args->totalProcs;
+
+  CUDACHECK(cudaSetDevice(args->gpus[0]));
+  int rank = args->proc;
+  CUDACHECK(cudaMemset(args->recvbuffs[0], 0, args->expectedBytes));
+  void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
+  // TESTCHECK(InitData(data, sendcount, 0, type, ncclSum, 33 * rep + rank, 1, 0));
+  // for (int j = 0; j < nranks; j++) {
+  //   TESTCHECK(
+  //     InitData((char*)args->expected[0] + args->sendBytes * j, sendcount, 0, type, ncclSum, 33 * rep + j, 1, 0));
+  // }
+  CUDACHECK(cudaDeviceSynchronize());
 
-  for (int i=0; i<args->nGpus; i++) {
-    CUDACHECK(cudaSetDevice(args->gpus[i]));
-    int rank = ((args->proc*args->nThreads + args->thread)*args->nGpus + i);
-    CUDACHECK(cudaMemset(args->recvbuffs[i], 0, args->expectedBytes));
-    void* data = in_place ? ((char*)args->recvbuffs[i])+rank*args->sendBytes : args->sendbuffs[i];
-    TESTCHECK(InitData(data, sendcount, 0, type, ncclSum, 33*rep + rank, 1, 0));
-    for (int j=0; j<nranks; j++) {
-      TESTCHECK(InitData((char*)args->expected[i] + args->sendBytes*j, sendcount, 0, type, ncclSum, 33*rep + j, 1, 0));
-    }
-    CUDACHECK(cudaDeviceSynchronize());
-  }
   return testSuccess;
 }
 
@@ -45,8 +47,8 @@ void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, doubl
   *busBw = baseBw * factor;
 }
 
-testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count, ncclDataType_t type, ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream) {
-  NCCLCHECK(ncclAllGather(sendbuff, recvbuff, count, type, comm, stream));
+testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count) {
+  // NCCLCHECK(ncclAllGather(sendbuff, recvbuff, count, type, comm, stream));
   return testSuccess;
 }
 
@@ -63,25 +65,10 @@ void AllGatherGetBuffSize(size_t *sendcount, size_t *recvcount, size_t count, in
   AllGatherGetCollByteCount(sendcount, recvcount, &paramcount, &sendInplaceOffset, &recvInplaceOffset, count, nranks);
 }
 
-testResult_t AllGatherRunTest(struct threadArgs* args, ncclDataType_t type, const char* typeName, ncclRedOp_t op, const char* opName) {
+testResult_t AllGatherRunTest(struct threadArgs* args) {
   args->collTest = &allGatherTest;
-  ncclDataType_t *run_types;
-  const char **run_typenames;
-  int type_count;
 
-  if ((int)type != -1) {
-    type_count = 1;
-    run_types = &type;
-    run_typenames = &typeName;
-  } else {
-    type_count = test_typenum;
-    run_types = test_types;
-    run_typenames = test_typenames;
-  }
-
-  for (int i=0; i<type_count; i++) {
-    TESTCHECK(TimeTest(args, run_types[i], run_typenames[i], (ncclRedOp_t)0, "none", -1));
-  }
+  TESTCHECK(TimeTest(args));
   return testSuccess;
 }
 
diff --git a/tests/common.cu b/tests/common.cu
index 5f6bf014..5074312b 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -9,6 +9,7 @@
 #include "common.h"
 #include "cuda.h"
 #include "mscclpp.h"
+#include "timer.h"
 
 #include <cstdio>
 #include <iostream>
@@ -23,8 +24,6 @@ int is_main_proc = 0;
 thread_local int is_main_thread = 0;
 
 // Command line parameter defaults
-static int nThreads = 1;
-static int nGpus = 1;
 static size_t minBytes = 32*1024*1024;
 static size_t maxBytes = 32*1024*1024;
 static size_t stepBytes = 1*1024*1024;
@@ -110,220 +109,165 @@ testResult_t AllocateBuffs(void **sendbuff, size_t sendBytes, void **recvbuff, s
     return testSuccess;
 }
 
-testResult_t startColl(struct threadArgs* args, ncclDataType_t type, ncclRedOp_t opIndex, int root, int in_place, int iter) {
-  size_t count = args->nbytes / wordSize(type);
+testResult_t startColl(struct threadArgs* args, int in_place, int iter) {
+  size_t count = args->nbytes;
 
   // Try to change offset for each iteration so that we avoid cache effects and catch race conditions in ptrExchange
   size_t totalnbytes = max(args->sendBytes, args->expectedBytes);
   size_t steps = totalnbytes ? args->maxbytes / totalnbytes : 1;
   size_t shift = totalnbytes * (iter % steps);
 
-  if (args->nGpus > 1) NCCLCHECK(ncclGroupStart());
   for (int i = 0; i < args->nGpus; i++) {
-#ifndef NCCL_MAJOR
-    CUDACHECK(cudaSetDevice(args->gpus[i]));
-#endif
     int rank = ((args->proc*args->nThreads + args->thread)*args->nGpus + i);
     char* recvBuff = ((char*)args->recvbuffs[i]) + shift;
     char* sendBuff = ((char*)args->sendbuffs[i]) + shift;
-    ncclRedOp_t op;
 
-    if(opIndex < ncclNumOps) {
-      op = opIndex;
-    }
-    #if NCCL_VERSION_CODE >= NCCL_VERSION(2,11,0)
-    else {
-      union {
-        int8_t i8; uint8_t u8; int32_t i32; uint32_t u32; int64_t i64; uint64_t u64;
-        half f16; float f32; double f64;
-        #if defined(__CUDA_BF16_TYPES_EXIST__)
-        __nv_bfloat16 bf16;
-        #endif
-      };
-      switch(type) {
-      case ncclInt8: i8 = ncclVerifiablePremulScalar<int8_t>(rank); break;
-      case ncclUint8: u8 = ncclVerifiablePremulScalar<uint8_t>(rank); break;
-      case ncclInt32: i32 = ncclVerifiablePremulScalar<int32_t>(rank); break;
-      case ncclUint32: u32 = ncclVerifiablePremulScalar<uint32_t>(rank); break;
-      case ncclInt64: i64 = ncclVerifiablePremulScalar<int64_t>(rank); break;
-      case ncclUint64: u64 = ncclVerifiablePremulScalar<uint64_t>(rank); break;
-      case ncclFloat16: f16 = ncclVerifiablePremulScalar<half>(rank); break;
-      case ncclFloat32: f32 = ncclVerifiablePremulScalar<float>(rank); break;
-      case ncclFloat64: f64 = ncclVerifiablePremulScalar<double>(rank); break;
-      #if defined(__CUDA_BF16_TYPES_EXIST__)
-      case ncclBfloat16: bf16 = ncclVerifiablePremulScalar<__nv_bfloat16>(rank); break;
-      #endif
-      }
-      NCCLCHECK(ncclRedOpCreatePreMulSum(&op, &u64, type, ncclScalarHostImmediate, args->comms[i]));
-    }
-    #endif
-
-    TESTCHECK(args->collTest->runColl(
-          (void*)(in_place ? recvBuff + args->sendInplaceOffset*rank : sendBuff),
-          (void*)(in_place ? recvBuff + args->recvInplaceOffset*rank : recvBuff),
-        count, type, op, root, args->comms[i], args->streams[i]));
-
-    #if NCCL_VERSION_CODE >= NCCL_VERSION(2,11,0)
-    if(opIndex >= ncclNumOps) {
-      NCCLCHECK(ncclRedOpDestroy(op, args->comms[i]));
-    }
-    #endif
+    TESTCHECK(args->collTest->runColl((void*)(in_place ? recvBuff + args->sendInplaceOffset * rank : sendBuff),
+                                      (void*)(in_place ? recvBuff + args->recvInplaceOffset * rank : recvBuff), count));
   }
-  if (args->nGpus > 1) NCCLCHECK(ncclGroupEnd());
-
-  if (blocking_coll) {
-    // Complete op before returning
-    TESTCHECK(testStreamSynchronize(args->nGpus, args->streams, args->comms));
-  }
-  if (blocking_coll) Barrier(args);
   return testSuccess;
 }
 
 testResult_t completeColl(struct threadArgs* args) {
-  if (blocking_coll) return testSuccess;
-
-  TESTCHECK(testStreamSynchronize(args->nGpus, args->streams, args->comms));
+  TESTCHECK(testStreamSynchronize(args->nGpus, args->streams));
   return testSuccess;
 }
 
-testResult_t BenchTime(struct threadArgs* args, ncclDataType_t type, ncclRedOp_t op, int root, int in_place) {
-  size_t count = args->nbytes / wordSize(type);
-  if (datacheck) {
-    // Initialize sendbuffs, recvbuffs and expected
-    TESTCHECK(args->collTest->initData(args, type, op, root, 99, in_place));
+testResult_t testStreamSynchronize(int ngpus, cudaStream_t* streams)
+{
+  cudaError_t cudaErr;
+  int remaining = ngpus;
+  int* done = (int*)malloc(sizeof(int) * ngpus);
+  memset(done, 0, sizeof(int) * ngpus);
+  timer tim;
+
+  while (remaining) {
+    int idle = 1;
+    for (int i = 0; i < ngpus; i++) {
+      if (done[i])
+        continue;
+
+      cudaErr = cudaStreamQuery(streams[i]);
+      if (cudaErr == cudaSuccess) {
+        done[i] = 1;
+        remaining--;
+        idle = 0;
+        continue;
+      }
+
+      if (cudaErr != cudaErrorNotReady)
+        CUDACHECK(cudaErr);
+
+      // We might want to let other threads (including NCCL threads) use the CPU.
+      if (idle)
+        sched_yield();
+    }
+    free(done);
+    return testSuccess;
+  }
+}
+
+// Inter-thread/process barrier+allreduce. The quality of the return value
+// for average=0 (which means broadcast from rank=0) is dubious. The returned
+// value will actually be the result of process-local broadcast from the local thread=0.
+template<typename T>
+void Allreduce(struct threadArgs* args, T* value, int average)
+{
+  thread_local int epoch = 0;
+  static pthread_mutex_t lock[2] = {PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER};
+  static pthread_cond_t cond[2] = {PTHREAD_COND_INITIALIZER, PTHREAD_COND_INITIALIZER};
+  static T accumulator[2];
+  static int counter[2] = {0, 0};
+
+  pthread_mutex_lock(&lock[epoch]);
+  if (counter[epoch] == 0) {
+    if (average != 0 || args->thread == 0)
+      accumulator[epoch] = *value;
+  } else {
+    switch (average) {
+    case /*r0*/ 0:
+      if (args->thread == 0)
+        accumulator[epoch] = *value;
+      break;
+    case /*avg*/ 1:
+      accumulator[epoch] += *value;
+      break;
+    case /*min*/ 2:
+      accumulator[epoch] = std::min<T>(accumulator[epoch], *value);
+      break;
+    case /*max*/ 3:
+      accumulator[epoch] = std::max<T>(accumulator[epoch], *value);
+      break;
+    case /*sum*/ 4:
+      accumulator[epoch] += *value;
+      break;
+    }
   }
 
+  if (++counter[epoch] == args->nThreads)
+    pthread_cond_broadcast(&cond[epoch]);
+
+  if (args->thread + 1 == args->nThreads) {
+    while (counter[epoch] != args->nThreads)
+      pthread_cond_wait(&cond[epoch], &lock[epoch]);
+
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+    if (average != 0) {
+      static_assert(std::is_same<T, long long>::value || std::is_same<T, double>::value,
+                    "Allreduce<T> only for T in {long long, double}");
+      MPI_Datatype ty = std::is_same<T, long long>::value ? MPI_LONG_LONG
+                        : std::is_same<T, double>::value  ? MPI_DOUBLE
+                                                          : MPI_Datatype();
+      MPI_Op op = average == 1   ? MPI_SUM
+                  : average == 2 ? MPI_MIN
+                  : average == 3 ? MPI_MAX
+                  : average == 4 ? MPI_SUM
+                                 : MPI_Op();
+      MPI_Allreduce(MPI_IN_PLACE, (void*)&accumulator[epoch], 1, ty, op, MPI_COMM_WORLD);
+    }
+#endif
+
+    if (average == 1)
+      accumulator[epoch] /= args->totalProcs * args->nThreads;
+    counter[epoch] = 0;
+    pthread_cond_broadcast(&cond[epoch]);
+  } else {
+    while (counter[epoch] != 0)
+      pthread_cond_wait(&cond[epoch], &lock[epoch]);
+  }
+  pthread_mutex_unlock(&lock[epoch]);
+
+  *value = accumulator[epoch];
+  epoch ^= 1;
+}
+
+testResult_t BenchTime(struct threadArgs* args, int in_place) {
+  size_t count = args->nbytes;
+
   // Sync
-  TESTCHECK(startColl(args, type, op, root, in_place, 0));
+  TESTCHECK(startColl(args, in_place, 0));
   TESTCHECK(completeColl(args));
 
   Barrier(args);
 
-#if CUDART_VERSION >= 11030
-  cudaGraph_t graphs[args->nGpus];
-  cudaGraphExec_t graphExec[args->nGpus];
-  if (cudaGraphLaunches >= 1) {
-    // Begin cuda graph capture
-    for (int i=0; i<args->nGpus; i++) {
-      // Thread local mdoe is needed for:
-      // - Multi-thread mode: where graph capture and instantiation can happen concurrently across threads
-      // - P2P pre-connect: when there is no warm-up, P2P pre-connect is done during graph capture.
-      //   Since pre-connect calls cudaMalloc, we cannot use global capture mode
-      CUDACHECK(cudaStreamBeginCapture(args->streams[i], cudaStreamCaptureModeThreadLocal));
-    }
-  }
-#endif
-
   // Performance Benchmark
   timer tim;
   for (int iter = 0; iter < iters; iter++) {
-    if (agg_iters>1) NCCLCHECK(ncclGroupStart());
-    for (int aiter = 0; aiter < agg_iters; aiter++) {
-      TESTCHECK(startColl(args, type, op, root, in_place, iter*agg_iters+aiter));
-    }
-    if (agg_iters>1) NCCLCHECK(ncclGroupEnd());
+    TESTCHECK(startColl(args, in_place, iter));
   }
 
-#if CUDART_VERSION >= 11030
-  if (cudaGraphLaunches >= 1) {
-    // End cuda graph capture
-    for (int i=0; i<args->nGpus; i++) {
-      CUDACHECK(cudaStreamEndCapture(args->streams[i], graphs+i));
-    }
-    // Instantiate cuda graph
-    for (int i=0; i<args->nGpus; i++) {
-      CUDACHECK(cudaGraphInstantiate(graphExec+i, graphs[i], NULL, NULL, 0));
-    }
-    // Resync CPU, restart timing, launch cuda graph
-    Barrier(args);
-    tim.reset();
-    for (int l=0; l<cudaGraphLaunches; l++) {
-      for (int i=0; i<args->nGpus; i++) {
-        CUDACHECK(cudaGraphLaunch(graphExec[i], args->streams[i]));
-      }
-    }
-  }
-#endif
-
-  double cputimeSec = tim.elapsed()/(iters*agg_iters);
+  double cputimeSec = tim.elapsed()/(iters);
   TESTCHECK(completeColl(args));
 
   double deltaSec = tim.elapsed();
-  deltaSec = deltaSec/(iters*agg_iters);
-  if (cudaGraphLaunches >= 1) deltaSec = deltaSec/cudaGraphLaunches;
+  deltaSec = deltaSec/(iters);
   Allreduce(args, &deltaSec, average);
 
-#if CUDART_VERSION >= 11030
-  if (cudaGraphLaunches >= 1) {
-    //destroy cuda graph
-    for (int i=0; i<args->nGpus; i++) {
-      CUDACHECK(cudaGraphExecDestroy(graphExec[i]));
-      CUDACHECK(cudaGraphDestroy(graphs[i]));
-    }
-  }
-#endif
-
   double algBw, busBw;
-  args->collTest->getBw(count, wordSize(type), deltaSec, &algBw, &busBw, args->nProcs*args->nThreads*args->nGpus);
+  args->collTest->getBw(count, 1, deltaSec, &algBw, &busBw, args->totalProcs * args->nThreads * args->nGpus);
 
   Barrier(args);
 
-  int64_t wrongElts = 0;
-  static __thread int rep = 0;
-  rep++;
-  if (datacheck) {
-      // Initialize sendbuffs, recvbuffs and expected
-      TESTCHECK(args->collTest->initData(args, type, op, root, rep, in_place));
-
-#if CUDART_VERSION >= 11030
-      if (cudaGraphLaunches >= 1) {
-        // Begin cuda graph capture for data check
-        for (int i=0; i<args->nGpus; i++) {
-          CUDACHECK(cudaStreamBeginCapture(args->streams[i], args->nThreads > 1 ? cudaStreamCaptureModeThreadLocal : cudaStreamCaptureModeGlobal));
-        }
-      }
-#endif
-
-      //test validation in single itertion, should ideally be included into the multi-iteration run
-      TESTCHECK(startColl(args, type, op, root, in_place, 0));
-
-#if CUDART_VERSION >= 11030
-      if (cudaGraphLaunches >= 1) {
-        // End cuda graph capture
-        for (int i=0; i<args->nGpus; i++) {
-          CUDACHECK(cudaStreamEndCapture(args->streams[i], graphs+i));
-        }
-        // Instantiate cuda graph
-        for (int i=0; i<args->nGpus; i++) {
-          CUDACHECK(cudaGraphInstantiate(graphExec+i, graphs[i], NULL, NULL, 0));
-        }
-        // Launch cuda graph
-        for (int i=0; i<args->nGpus; i++) {
-          CUDACHECK(cudaGraphLaunch(graphExec[i], args->streams[i]));
-        }
-      }
-#endif
-
-      TESTCHECK(completeColl(args));
-
-#if CUDART_VERSION >= 11030
-      if (cudaGraphLaunches >= 1) {
-        //destroy cuda graph
-        for (int i=0; i<args->nGpus; i++) {
-          CUDACHECK(cudaGraphExecDestroy(graphExec[i]));
-          CUDACHECK(cudaGraphDestroy(graphs[i]));
-        }
-      }
-#endif
-
-      TESTCHECK(CheckData(args, type, op, root, in_place, &wrongElts));
-
-      //aggregate delta from all threads and procs
-      long long wrongElts1 = wrongElts;
-      Allreduce(args, &wrongElts1, /*sum*/4);
-      wrongElts = wrongElts1;
-  }
-
   double timeUsec = (report_cputime ? cputimeSec : deltaSec)*1.0E6;
   char timeStr[100];
   if (timeUsec >= 10000.0) {
@@ -333,74 +277,60 @@ testResult_t BenchTime(struct threadArgs* args, ncclDataType_t type, ncclRedOp_t
   } else {
     sprintf(timeStr, "%7.2f", timeUsec);
   }
-  if (args->reportErrors) {
-    PRINT("  %7s  %6.2f  %6.2f  %5g", timeStr, algBw, busBw, (double)wrongElts);
-  } else {
-    PRINT("  %7s  %6.2f  %6.2f  %5s", timeStr, algBw, busBw, "N/A");
-  }
+  PRINT("  %7s  %6.2f  %6.2f  %5s", timeStr, algBw, busBw, "N/A");
 
   args->bw[0] += busBw;
   args->bw_count[0]++;
   return testSuccess;
 }
 
-void setupArgs(size_t size, ncclDataType_t type, struct threadArgs* args) {
-  int nranks = args->nProcs*args->nGpus*args->nThreads;
+void setupArgs(size_t size, struct threadArgs* args) {
+  int nranks = args->totalProcs*args->nGpus*args->nThreads;
   size_t count, sendCount, recvCount, paramCount, sendInplaceOffset, recvInplaceOffset;
 
-  count = size / wordSize(type);
+  // TODO: support more data types
+  int typeSize = sizeof(char);
+  count = size / typeSize;
   args->collTest->getCollByteCount(&sendCount, &recvCount, &paramCount, &sendInplaceOffset, &recvInplaceOffset, (size_t)count, (size_t)nranks);
 
-  args->nbytes = paramCount * wordSize(type);
-  args->sendBytes = sendCount * wordSize(type);
-  args->expectedBytes = recvCount * wordSize(type);
-  args->sendInplaceOffset = sendInplaceOffset * wordSize(type);
-  args->recvInplaceOffset = recvInplaceOffset * wordSize(type);
+  args->nbytes = paramCount * typeSize;
+  args->sendBytes = sendCount * typeSize;
+  args->expectedBytes = recvCount * typeSize;
+  args->sendInplaceOffset = sendInplaceOffset * typeSize;
+  args->recvInplaceOffset = recvInplaceOffset * typeSize;
 }
 
-testResult_t TimeTest(struct threadArgs* args, ncclDataType_t type, const char* typeName, ncclRedOp_t op, const char* opName, int root) {
+testResult_t TimeTest(struct threadArgs* args) {
   // Sync to avoid first-call timeout
   Barrier(args);
 
   // Warm-up for large size
-  setupArgs(args->maxbytes, type, args);
+  setupArgs(args->maxbytes, args);
   for (int iter = 0; iter < warmup_iters; iter++) {
-    TESTCHECK(startColl(args, type, op, root, 0, iter));
+    TESTCHECK(startColl(args, 0, iter));
   }
   TESTCHECK(completeColl(args));
 
   // Warm-up for small size
-  setupArgs(args->minbytes, type, args);
+  setupArgs(args->minbytes, args);
   for (int iter = 0; iter < warmup_iters; iter++) {
-    TESTCHECK(startColl(args, type, op, root, 0, iter));
+    TESTCHECK(startColl(args, 0, iter));
   }
   TESTCHECK(completeColl(args));
 
   // Benchmark
   for (size_t size = args->minbytes; size<=args->maxbytes; size = ((args->stepfactor > 1) ? size*args->stepfactor : size+args->stepbytes)) {
-      setupArgs(size, type, args);
-      char rootName[100];
-      sprintf(rootName, "%6i", root);
-      PRINT("%12li  %12li  %8s  %6s  %6s", max(args->sendBytes, args->expectedBytes), args->nbytes / wordSize(type), typeName, opName, rootName);
-      TESTCHECK(BenchTime(args, type, op, root, 0));
-      TESTCHECK(BenchTime(args, type, op, root, 1));
+      setupArgs(size, args);
+      PRINT("%12li  %12li", max(args->sendBytes, args->expectedBytes), args->nbytes);
+      TESTCHECK(BenchTime(args, 0));
+      TESTCHECK(BenchTime(args, 1));
       PRINT("\n");
   }
   return testSuccess;
 }
 
 testResult_t threadRunTests(struct threadArgs* args) {
-  TESTCHECK(mscclppTestEngine.runTest(args, test_typenames[nccltype], (ncclRedOp_t)ncclop, test_opnames[ncclop]));
-  return testSuccess;
-}
-
-void* threadLauncher(void* thread_) {
-  struct testThread* thread = (struct testThread*)thread_;
-  thread->ret = thread->func(&thread->args);
-  return NULL;
-}
-testResult_t threadLaunch(struct testThread* thread) {
-  pthread_create(&thread->thread, NULL, threadLauncher, thread);
+  TESTCHECK(mscclppTestEngine.runTest(args));
   return testSuccess;
 }
 
@@ -414,8 +344,6 @@ int main(int argc, char* argv[]) {
   double parsed;
   int longindex;
   static struct option longopts[] = {
-    {"nthreads", required_argument, 0, 't'},
-    {"ngpus", required_argument, 0, 'g'},
     {"minbytes", required_argument, 0, 'b'},
     {"maxbytes", required_argument, 0, 'e'},
     {"stepbytes", required_argument, 0, 'i'},
@@ -433,18 +361,12 @@ int main(int argc, char* argv[]) {
 
   while(1) {
     int c;
-    c = getopt_long(argc, argv, "t:g:b:e:i:f:n:w:c:T:C:a:P:h:", longopts, &longindex);
+    c = getopt_long(argc, argv, "b:e:i:f:n:w:c:T:C:a:P:h:", longopts, &longindex);
 
     if (c == -1)
       break;
 
     switch(c) {
-      case 't':
-        nThreads = strtol(optarg, NULL, 0);
-        break;
-      case 'g':
-        nGpus = strtol(optarg, NULL, 0);
-        break;
       case 'b':
         parsed = parsesize(optarg);
         if (parsed < 0) {
@@ -492,8 +414,6 @@ int main(int argc, char* argv[]) {
       default:
         if (c != 'h') printf("invalid option '%c'\n", c);
         printf("USAGE: %s \n\t"
-            "[-t,--nthreads <num threads>] \n\t"
-            "[-g,--ngpus <gpus per thread>] \n\t"
             "[-b,--minbytes <min size in bytes>] \n\t"
             "[-e,--maxbytes <max size in bytes>] \n\t"
             "[-i,--stepbytes <increment size>] \n\t"
@@ -544,10 +464,9 @@ testResult_t run() {
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
 
-  PRINT("# nThread %d nGpus %d minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip "
-        "port: %s\n",
-        nThreads, nGpus, minBytes, maxBytes, (stepFactor > 1) ? stepFactor : stepBytes,
-        (stepFactor > 1) ? "factor" : "bytes", warmup_iters, iters, datacheck, ip_port.c_str());
+  PRINT("# minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip port: %s\n", minBytes,
+        maxBytes, (stepFactor > 1) ? stepFactor : stepBytes, (stepFactor > 1) ? "factor" : "bytes", warmup_iters, iters,
+        datacheck, ip_port.c_str());
   PRINT("#\n");
   PRINT("# Using devices\n");
 
@@ -555,15 +474,14 @@ testResult_t run() {
   char line[MAX_LINE];
   int len = 0;
   size_t maxMem = ~0;
-  for (int i = 0; i < nThreads * nGpus; i++) {
-    int cudaDev = (localRank * nThreads * nGpus) + i;
-    int rank = proc * nThreads * nGpus + i;
-    cudaDeviceProp prop;
-    CUDACHECK(cudaGetDeviceProperties(&prop, cudaDev));
-    len += snprintf(line + len, MAX_LINE - len, "#  Rank %2d Pid %6d on %10s device %2d [0x%02x] %s\n", rank, getpid(),
-                    hostname, cudaDev, prop.pciBusID, prop.name);
-    maxMem = std::min(maxMem, prop.totalGlobalMem);
-  }
+
+  int cudaDev = localRank;
+  int rank = proc;
+  cudaDeviceProp prop;
+  CUDACHECK(cudaGetDeviceProperties(&prop, cudaDev));
+  len += snprintf(line + len, MAX_LINE - len, "#  Rank %2d Pid %6d on %10s device %2d [0x%02x] %s\n", rank, getpid(),
+                  hostname, cudaDev, prop.pciBusID, prop.name);
+  maxMem = std::min(maxMem, prop.totalGlobalMem);
 
 #if MSCCLPP_USE_MPI_FOR_TESTS
   char* lines = (proc == 0) ? (char*)malloc(totalProcs * MAX_LINE) : NULL;
@@ -585,39 +503,29 @@ testResult_t run() {
     if (proc == 0) printf("#\n# Reducing maxBytes to %ld due to memory limitation\n", maxBytes);
   }
 
-  int gpus[nGpus*nThreads];
-  cudaStream_t streams[nGpus*nThreads];
-  void* sendbuffs[nGpus*nThreads];
-  void* recvbuffs[nGpus*nThreads];
-  void* expected[nGpus*nThreads];
+  int gpu = cudaDev;
+  cudaStream_t stream;
+  void* sendbuff;
+  void* recvbuff;
+  void* expected;
   size_t sendBytes = 10, recvBytes = 10;
 
-  // mscclppTestEngine.getBuffSize(&sendBytes, &recvBytes, (size_t)maxBytes, (size_t)totalProcs*nGpus*nThreads);
+  mscclppTestEngine.getBuffSize(&sendBytes, &recvBytes, (size_t)maxBytes, (size_t)totalProcs);
 
-  for (int i = 0; i < nGpus * nThreads; i++) {
-    gpus[i] = (localRank * nThreads * nGpus) + i;
-    CUDACHECK(cudaSetDevice(gpus[i]));
-    TESTCHECK(AllocateBuffs(sendbuffs + i, sendBytes, recvbuffs + i, recvBytes, expected + i, (size_t)maxBytes));
-    CUDACHECK(cudaStreamCreateWithFlags(streams + i, cudaStreamNonBlocking));
-  }
+  CUDACHECK(cudaSetDevice(gpu));
+  TESTCHECK(AllocateBuffs(&sendbuff, sendBytes, &recvbuff, recvBytes, &expected, (size_t)maxBytes));
+  CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
   PRINT("#\n");
   PRINT("# Initializing MSCCL++\n");
 
-  mscclppComm_t* comms = (mscclppComm_t*)malloc(sizeof(mscclppComm_t)*nThreads*nGpus);
-  for (int i = 0; i < nGpus * nThreads; i++) {
-    MSCCLPPCHECK(
-      mscclppCommInitRank(comms + i, totalProcs * nGpus * nThreads, ip_port.c_str(), localRank * nThreads * nGpus + i));
-  }
+  mscclppComm_t comms;
+  MSCCLPPCHECK(mscclppCommInitRank(&comms, totalProcs, ip_port.c_str(), localRank));
 
-  int errors[nThreads];
-  double bw[nThreads];
+  int error = 0;
+  double bw = 0.0;
   double* delta;
-  CUDACHECK(cudaHostAlloc(&delta, sizeof(double) * nThreads * NUM_BLOCKS, cudaHostAllocPortable | cudaHostAllocMapped));
-  int bw_count[nThreads];
-  for (int t = 0; t < nThreads; t++) {
-    bw[t] = 0.0;
-    errors[t] = bw_count[t] = 0;
-  }
+  CUDACHECK(cudaHostAlloc(&delta, sizeof(double) * NUM_BLOCKS, cudaHostAllocPortable | cudaHostAllocMapped));
+  int bw_count = 0;
 
   fflush(stdout);
 
@@ -630,70 +538,51 @@ testResult_t run() {
   PRINT("# %10s  %12s  %8s  %6s  %6s  %7s  %6s  %6s  %5s  %7s  %6s  %6s  %5s\n", "(B)", "(elements)", "", "", "",
         "(us)", "(GB/s)", "(GB/s)", "", "(us)", "(GB/s)", "(GB/s)", "");
 
-  struct testThread threads[nThreads];
-  memset(threads, 0, sizeof(struct testThread) * nThreads);
+  struct testThread thread = {0};
 
-  for (int t = nThreads - 1; t >= 0; t--) {
-    threads[t].args.minbytes = minBytes;
-    threads[t].args.maxbytes = maxBytes;
-    threads[t].args.stepbytes = stepBytes;
-    threads[t].args.stepfactor = stepFactor;
-    threads[t].args.localRank = localRank;
+  thread.args.minbytes = minBytes;
+  thread.args.maxbytes = maxBytes;
+  thread.args.stepbytes = stepBytes;
+  thread.args.stepfactor = stepFactor;
+  thread.args.localRank = localRank;
 
-    threads[t].args.totalProcs = totalProcs;
-    threads[t].args.proc = proc;
-    threads[t].args.nThreads = nThreads;
-    threads[t].args.thread = t;
-    threads[t].args.nGpus = nGpus;
-    threads[t].args.gpus = gpus + t * nGpus;
-    threads[t].args.sendbuffs = sendbuffs + t * nGpus;
-    threads[t].args.recvbuffs = recvbuffs + t * nGpus;
-    threads[t].args.expected = expected + t * nGpus;
-    threads[t].args.comms = comms + t * nGpus;
-    threads[t].args.streams = streams + t * nGpus;
+  thread.args.totalProcs = totalProcs;
+  thread.args.proc = proc;
+  thread.args.nThreads = 1;
+  thread.args.thread = 0;
+  thread.args.nGpus = 1;
+  thread.args.gpus = &gpu;
+  thread.args.sendbuffs = &sendbuff;
+  thread.args.recvbuffs = &recvbuff;
+  thread.args.expected = &expected;
+  thread.args.comms = &comms;
+  thread.args.streams = &stream;
 
-    threads[t].args.errors = errors + t;
-    threads[t].args.bw = bw + t;
-    threads[t].args.bw_count = bw_count + t;
+  thread.args.errors = &error;
+  thread.args.bw = &bw;
+  thread.args.bw_count = &bw_count;
 
-    threads[t].args.reportErrors = datacheck;
+  thread.args.reportErrors = datacheck;
 
-    threads[t].func = threadRunTests;
-    if (t) {
-        TESTCHECK(threadLaunch(threads + t));
-    } else {
-        TESTCHECK(threads[t].func(&threads[t].args));
-    }
-  }
+  thread.func = threadRunTests;
+  TESTCHECK(thread.func(&thread.args));
 
   // Wait for other threads and accumulate stats and errors
-  for (int t = nThreads - 1; t >= 0; t--) {
-    if (t)
-        pthread_join(threads[t].thread, NULL);
-    TESTCHECK(threads[t].ret);
-    if (t) {
-        errors[0] += errors[t];
-        bw[0] += bw[t];
-        bw_count[0] += bw_count[t];
-    }
-  }
-
-  for (int i = 0; i < nGpus * nThreads; i++) {
-    MSCCLPPCHECK(mscclppCommDestroy(comms[i]));
-  }
-  free(comms);
+  TESTCHECK(thread.ret);
+  MSCCLPPCHECK(mscclppCommDestroy(comms));
 
   // Free off CUDA allocated memory
-  for (int i=0; i<nGpus*nThreads; i++) {
-    if (sendbuffs[i]) CUDACHECK(cudaFree((char*)sendbuffs[i]));
-    if (recvbuffs[i]) CUDACHECK(cudaFree((char*)recvbuffs[i]));
-    if (datacheck) CUDACHECK(cudaFree(expected[i]));
-  }
+  if (sendbuff)
+    CUDACHECK(cudaFree((char*)sendbuff));
+  if (recvbuff)
+    CUDACHECK(cudaFree((char*)recvbuff));
+  if (datacheck)
+    CUDACHECK(cudaFree(expected));
   CUDACHECK(cudaFreeHost(delta));
 
-  bw[0] /= bw_count[0];
+  bw /= bw_count;
 
-  PRINT("# Out of bounds values : %d %s\n", errors[0], errors[0] ? "FAILED" : "OK");
+  PRINT("# Out of bounds values : %d %s\n", error, error ? "FAILED" : "OK");
   PRINT("#\n");
 
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
diff --git a/tests/common.h b/tests/common.h
index 82f2af16..92aa645c 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -69,7 +69,7 @@ struct testColl {
       size_t *sendcount, size_t *recvcount, size_t *paramcount,
       size_t *sendInplaceOffset, size_t *recvInplaceOffset,
       size_t count, int nranks);
-  testResult_t (*initData)(struct threadArgs* args, int root, int rep, int in_place);
+  testResult_t (*initData)(struct threadArgs* args, int in_place);
   void (*getBw)(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks);
   testResult_t (*runColl)(void* sendbuff, void* recvbuff, size_t count);
 };
@@ -77,7 +77,8 @@ struct testColl {
 struct testEngine
 {
   void (*getBuffSize)(size_t* sendcount, size_t* recvcount, size_t count, int nranks);
-  testResult_t (*runTest)(struct threadArgs* args, int root, const char* typeName, const char* opName);
+  // We can add more parameters for other communication primitives
+  testResult_t (*runTest)(struct threadArgs* args);
 };
 
 extern struct testEngine mscclppTestEngine;
@@ -126,7 +127,7 @@ struct testThread
 };
 
 // Provided by common.cu
-extern testResult_t TimeTest(struct threadArgs* args, ncclDataType_t type, const char* typeName, ncclRedOp_t op,  const char* opName, int root);
+extern testResult_t TimeTest(struct threadArgs* args);
 
 static void getHostName(char* hostname, int maxlen)
 {
diff --git a/tests/timer.cc b/tests/timer.cc
new file mode 100644
index 00000000..eb510d9a
--- /dev/null
+++ b/tests/timer.cc
@@ -0,0 +1,28 @@
+#include "timer.h"
+
+// Make sure to compile this translation unit with the host compiler and not
+// nvcc, lest you hit an internal compiler error (ICE) with GCC 10.3.0
+#include <chrono>
+
+namespace {
+  std::uint64_t now() {
+    using clock = std::chrono::steady_clock;
+    return std::chrono::duration_cast<std::chrono::nanoseconds>(clock::now().time_since_epoch()).count();
+  }
+}
+
+timer::timer() {
+  t0 = now();
+}
+
+double timer::elapsed() const {
+  std::uint64_t t1 = now();
+  return 1.e-9*(t1 - t0);
+}
+
+double timer::reset() {
+  std::uint64_t t1 = now();
+  double ans = 1.e-9*(t1 - t0);
+  t0 = t1;
+  return ans;
+}
\ No newline at end of file
diff --git a/tests/timer.h b/tests/timer.h
new file mode 100644
index 00000000..6edfdd4c
--- /dev/null
+++ b/tests/timer.h
@@ -0,0 +1,15 @@
+#ifndef _408319ecdd5b47b28bf8f511c4fdf816
+#define _408319ecdd5b47b28bf8f511c4fdf816
+
+#include <cstdint>
+
+// Can't include <chrono> because of bug with gcc 10.3.0
+class timer {
+  std::uint64_t t0;
+public:
+  timer();
+  double elapsed() const;
+  double reset();
+};
+
+#endif
\ No newline at end of file

From 98020f5b524c45f4fefbcafcbfb98f5a9f8b6f91 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 31 Mar 2023 07:10:43 +0000
Subject: [PATCH 04/23] update

---
 Makefile                 |  8 +++++++-
 tests/allgather_test3.cu |  2 +-
 tests/common.cu          | 32 ++++++++++++++++----------------
 tests/common.h           |  4 ++--
 4 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/Makefile b/Makefile
index 2793b7bb..55e51dc2 100644
--- a/Makefile
+++ b/Makefile
@@ -129,10 +129,11 @@ LIBSONAME := $(LIBNAME).$(MSCCLPP_MAJOR)
 LIBTARGET := $(BUILDDIR)/$(LIBDIR)/$(LIBNAME).$(MSCCLPP_MAJOR).$(MSCCLPP_MINOR).$(MSCCLPP_PATCH)
 
 TESTSDIR  := tests
-TESTSSRCS := $(addprefix $(TESTSDIR)/,bootstrap_test.cc allgather_test.cu common.cu)
+TESTSSRCS := $(addprefix $(TESTSDIR)/,bootstrap_test.cc allgather_test.cu)
 TESTSOBJS := $(patsubst %.cc,%.o,$(TESTSSRCS)) $(patsubst %.cu,%.o,$(TESTSSRCS))
 TESTSOBJTARGETS := $(TESTSOBJS:%=$(BUILDDIR)/$(OBJDIR)/%)
 TESTSBINS       := $(patsubst %.o,$(BUILDDIR)/$(BINDIR)/%,$(TESTSOBJS))
+TESTSPERFBIN := $(BUILDDIR)/$(BINDIR)/allgather_test3
 
 INCLUDE := -Isrc -Isrc/include
 
@@ -146,6 +147,8 @@ lib: $(LIBOBJTARGETS) $(INCTARGETS) $(LIBTARGET)
 
 tests: $(TESTSBINS)
 
+allgather_perf: $(TESTSPERFBIN)
+
 cpplint:
 	clang-format-12 -style=file --verbose --Werror --dry-run $(CPPSOURCES)
 
@@ -172,6 +175,9 @@ $(LIBTARGET): $(LIBOBJTARGETS)
 	ln -sf $(LIBTARGET) $(BUILDDIR)/$(LIBDIR)/$(LIBSONAME)
 
 # Compile .cc tests
+$(TESTSPERFBIN): $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/allgather_test3.o $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/common.o $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/timer.o
+	$(NVCC) -o $@ $^ $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp
+
 $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/%.o: $(TESTSDIR)/%.cc $(INCTARGETS)
 	@mkdir -p $(@D)
 	$(CXX) -o $@ -I$(BUILDDIR)/$(INCDIR) $(MPI_INC) $(CXXFLAGS) -c $< $(MPI_MACRO)
diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
index 5aad077c..adf9426b 100644
--- a/tests/allgather_test3.cu
+++ b/tests/allgather_test3.cu
@@ -77,4 +77,4 @@ struct testEngine allGatherEngine = {
   AllGatherRunTest
 };
 
-#pragma weak mcclppTestEngine=allGatherEngine
\ No newline at end of file
+#pragma weak mscclppTestEngine=allGatherEngine
\ No newline at end of file
diff --git a/tests/common.cu b/tests/common.cu
index 5074312b..a4ffde38 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -128,11 +128,6 @@ testResult_t startColl(struct threadArgs* args, int in_place, int iter) {
   return testSuccess;
 }
 
-testResult_t completeColl(struct threadArgs* args) {
-  TESTCHECK(testStreamSynchronize(args->nGpus, args->streams));
-  return testSuccess;
-}
-
 testResult_t testStreamSynchronize(int ngpus, cudaStream_t* streams)
 {
   cudaError_t cudaErr;
@@ -157,14 +152,19 @@ testResult_t testStreamSynchronize(int ngpus, cudaStream_t* streams)
 
       if (cudaErr != cudaErrorNotReady)
         CUDACHECK(cudaErr);
-
-      // We might want to let other threads (including NCCL threads) use the CPU.
-      if (idle)
-        sched_yield();
     }
-    free(done);
-    return testSuccess;
+
+    // We might want to let other threads (including NCCL threads) use the CPU.
+    if (idle)
+      sched_yield();
   }
+  free(done);
+  return testSuccess;
+}
+
+testResult_t completeColl(struct threadArgs* args) {
+  TESTCHECK(testStreamSynchronize(args->nGpus, args->streams));
+  return testSuccess;
 }
 
 // Inter-thread/process barrier+allreduce. The quality of the return value
@@ -531,12 +531,12 @@ testResult_t run() {
 
   const char* timeStr = report_cputime ? "cputime" : "time";
   PRINT("#\n");
-  PRINT("# %10s  %12s  %8s  %6s  %6s           out-of-place                       in-place          \n", "", "", "", "",
+  PRINT("# %10s  %12s           out-of-place                       in-place          \n", "",
         "");
-  PRINT("# %10s  %12s  %8s  %6s  %6s  %7s  %6s  %6s %6s  %7s  %6s  %6s %6s\n", "size", "count", "type", "redop", "root",
-        timeStr, "algbw", "busbw", "#wrong", timeStr, "algbw", "busbw", "#wrong");
-  PRINT("# %10s  %12s  %8s  %6s  %6s  %7s  %6s  %6s  %5s  %7s  %6s  %6s  %5s\n", "(B)", "(elements)", "", "", "",
-        "(us)", "(GB/s)", "(GB/s)", "", "(us)", "(GB/s)", "(GB/s)", "");
+  PRINT("# %10s  %12s  %7s  %6s  %6s  %6s  %7s  %6s  %6s  %6s\n", "size", "count", timeStr, "algbw", "busbw", "#wrong",
+        timeStr, "algbw", "busbw", "#wrong");
+  PRINT("# %10s  %12s  %7s  %6s  %6s  %5s  %7s  %6s  %6s  %5s\n", "(B)", "(elements)", "(us)", "(GB/s)", "(GB/s)", "",
+        "(us)", "(GB/s)", "(GB/s)", "");
 
   struct testThread thread = {0};
 
diff --git a/tests/common.h b/tests/common.h
index 92aa645c..5e0dca5e 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -140,7 +140,7 @@ static void getHostName(char* hostname, int maxlen)
   }
 }
 
-void print_usage(const char* prog)
+inline void print_usage(const char* prog)
 {
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
   printf("usage: %s IP:PORT [rank nranks]\n", prog);
@@ -149,7 +149,7 @@ void print_usage(const char* prog)
 #endif
 }
 
-void parse_arguments(int argc, const char* argv[], const char** ip_port, int* rank, int* world_size)
+inline void parse_arguments(int argc, const char* argv[], const char** ip_port, int* rank, int* world_size)
 {
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
   if (argc != 2 && argc != 4) {

From a3adc99b5c63220a0b91ab87fa44edb6dfe2885f Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 31 Mar 2023 15:10:51 +0000
Subject: [PATCH 05/23] more code

---
 tests/allgather_test3.cu | 97 +++++++++++++++++++++++++++++++---------
 tests/common.cu          | 53 +++++++++++++++++++---
 tests/common.h           | 24 +++++-----
 3 files changed, 138 insertions(+), 36 deletions(-)

diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
index adf9426b..53628004 100644
--- a/tests/allgather_test3.cu
+++ b/tests/allgather_test3.cu
@@ -4,11 +4,66 @@
  * See LICENSE.txt for license information
  ************************************************************************/
 
-#include "cuda_runtime.h"
 #include "common.h"
+#include "cuda_runtime.h"
+
+#include <string>
 
 #define ALIGN 4
 
+__device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, int nelemsPerGPU)
+{
+  // this allgather is really simple and implemented as an alltoall
+
+  // this thread's role is a sender role
+  // put your data asynchronously
+  devConn.put(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
+  // make sure everyone is put their data before some thread randomly blocks everyone else in signal
+  __syncthreads();
+  // push with flag and sync to make sure the data is received
+  devConn.signal();
+
+  // this thread's role is a receiver role. wait on the semaphore to make sure the data is ready
+  devConn.wait();
+}
+
+__device__ void allgather1(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, int nelemsPerGPU)
+{
+  // this allgather algorithm works as follows:
+  // Step 1: GPU rank i sends data to GPU rank (i+1) % world_size
+  // Step 2: GPU rank i waits for data from GPU rank (i+2) % world_size
+  // ...
+  // This order is much better for DMA engine for NVLinks
+
+  for (int i = 1; i < world_size; i++) {
+    __syncthreads();
+    if (remoteRank != ((rank + i) % world_size))
+      continue;
+    // put your data to GPU (rank+i) % world_size and signal all in one call
+    devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
+  }
+  // all connections wait for the signal from the sender
+  devConn.wait();
+}
+
+__global__ void kernel(int rank, int world_size, int nelemsPerGPU, int kernel)
+{
+  // only use a single thread from each warp
+  if (threadIdx.x % 32 != 0)
+    return;
+
+  // find the mapping between remoteRank and devConns
+  int warpId = threadIdx.x / 32;
+  int remoteRank = (warpId < rank) ? warpId : warpId + 1;
+  // Each warp is responsible for one of the remote ranks
+  mscclppDevConn_t devConn = constDevConns[warpId];
+
+  if (kernel == 0)
+    allgather0(devConn, rank, world_size, remoteRank, nelemsPerGPU);
+  else if (kernel == 1)
+    allgather1(devConn, rank, world_size, remoteRank, nelemsPerGPU);
+}
+
 void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* paramcount, size_t* sendInplaceOffset,
                                size_t* recvInplaceOffset, size_t count, int nranks)
 {
@@ -29,11 +84,19 @@ testResult_t AllGatherInitData(struct threadArgs* args, int in_place) {
   int rank = args->proc;
   CUDACHECK(cudaMemset(args->recvbuffs[0], 0, args->expectedBytes));
   void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
-  // TESTCHECK(InitData(data, sendcount, 0, type, ncclSum, 33 * rep + rank, 1, 0));
-  // for (int j = 0; j < nranks; j++) {
-  //   TESTCHECK(
-  //     InitData((char*)args->expected[0] + args->sendBytes * j, sendcount, 0, type, ncclSum, 33 * rep + j, 1, 0));
-  // }
+
+  int* dataHost = new int[recvcount];
+  for (int i = 0; i < static_cast<int>(recvcount); i++) {
+    int val = i + 1;
+    if (i / args->ranksPerNode == rank) {
+      dataHost[i] = val;
+    } else {
+      dataHost[i] = 0;
+    }
+  }
+  CUDACHECK(cudaMemcpy(args->recvbuffs[0], dataHost, recvcount, cudaMemcpyHostToDevice));
+  delete dataHost;
+  // TODO: need to init expected data here
   CUDACHECK(cudaDeviceSynchronize());
 
   return testSuccess;
@@ -47,34 +110,28 @@ void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, doubl
   *busBw = baseBw * factor;
 }
 
-testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count) {
+testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count)
+{
   // NCCLCHECK(ncclAllGather(sendbuff, recvbuff, count, type, comm, stream));
   return testSuccess;
 }
 
-struct testColl allGatherTest = {
-  "AllGather",
-  AllGatherGetCollByteCount,
-  AllGatherInitData,
-  AllGatherGetBw,
-  AllGatherRunColl
-};
+struct testColl allGatherTest = {"AllGather", AllGatherGetCollByteCount, AllGatherInitData, AllGatherGetBw,
+                                 AllGatherRunColl};
 
 void AllGatherGetBuffSize(size_t *sendcount, size_t *recvcount, size_t count, int nranks) {
   size_t paramcount, sendInplaceOffset, recvInplaceOffset;
   AllGatherGetCollByteCount(sendcount, recvcount, &paramcount, &sendInplaceOffset, &recvInplaceOffset, count, nranks);
 }
 
-testResult_t AllGatherRunTest(struct threadArgs* args) {
+testResult_t AllGatherRunTest(struct threadArgs* args)
+{
   args->collTest = &allGatherTest;
 
   TESTCHECK(TimeTest(args));
   return testSuccess;
 }
 
-struct testEngine allGatherEngine = {
-  AllGatherGetBuffSize,
-  AllGatherRunTest
-};
+struct testEngine allGatherEngine = {AllGatherGetBuffSize, AllGatherRunTest};
 
-#pragma weak mscclppTestEngine=allGatherEngine
\ No newline at end of file
+#pragma weak mscclppTestEngine = allGatherEngine
diff --git a/tests/common.cu b/tests/common.cu
index a4ffde38..20b7e9b3 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -88,7 +88,7 @@ void Barrier(struct threadArgs *args) {
   if(args->thread+1 == args->nThreads) {
     while(counter[epoch] != args->nThreads)
       pthread_cond_wait(&cond[epoch], &lock[epoch]);
-    #ifdef MPI_SUPPORT
+    #ifdef MSCCLPP_USE_MPI_FOR_TESTS
       MPI_Barrier(MPI_COMM_WORLD);
     #endif
     counter[epoch] = 0;
@@ -244,6 +244,7 @@ void Allreduce(struct threadArgs* args, T* value, int average)
 testResult_t BenchTime(struct threadArgs* args, int in_place) {
   size_t count = args->nbytes;
 
+  TESTCHECK(args->collTest->initData(args, in_place));
   // Sync
   TESTCHECK(startColl(args, in_place, 0));
   TESTCHECK(completeColl(args));
@@ -322,15 +323,56 @@ testResult_t TimeTest(struct threadArgs* args) {
   for (size_t size = args->minbytes; size<=args->maxbytes; size = ((args->stepfactor > 1) ? size*args->stepfactor : size+args->stepbytes)) {
       setupArgs(size, args);
       PRINT("%12li  %12li", max(args->sendBytes, args->expectedBytes), args->nbytes);
-      TESTCHECK(BenchTime(args, 0));
+      // Don't support out-of-place for now
+      // TESTCHECK(BenchTime(args, 0));
       TESTCHECK(BenchTime(args, 1));
       PRINT("\n");
   }
   return testSuccess;
 }
 
-testResult_t threadRunTests(struct threadArgs* args) {
+testResult_t setupMscclppConnections(int rank, int ranksPerNode, int worldSize, mscclppComm_t comm, void* dataDst,
+                                        size_t dataSize)
+{
+  int thisNode = rank / ranksPerNode;
+  int localRank = rank % ranksPerNode;
+  std::string ibDevStr = "mlx5_ib" + std::to_string(localRank);
+
+  for (int r = 0; r < worldSize; ++r) {
+    if (r == rank)
+      continue;
+    mscclppTransport_t transportType;
+    const char* ibDev = ibDevStr.c_str();
+    if (r / ranksPerNode == thisNode) {
+      ibDev = NULL;
+      transportType = mscclppTransportP2P;
+    } else {
+      transportType = mscclppTransportIB;
+    }
+    // Connect with all other ranks
+    MSCCLPPCHECK(mscclppConnect(comm, r, 0, dataDst, dataSize, transportType, ibDev));
+  }
+
+  MSCCLPPCHECK(mscclppConnectionSetup(comm));
+
+  mscclppDevConn_t* devConns;
+  int nCons;
+  MSCCLPPCHECK(mscclppGetAllDeviceConnections(comm, &devConns, &nCons));
+
+  CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * nCons));
+
+  return testSuccess;
+}
+
+testResult_t threadRunTests(struct threadArgs* args)
+{
+  TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->ranksPerNode, args->comms[0],
+                                    args->recvbuffs[0], args->expectedBytes));
+  PRINT("Setting up the connection in MSCCL++\n");
+  MSCCLPPCHECK(mscclppProxyLaunch(args->comms[0]));
   TESTCHECK(mscclppTestEngine.runTest(args));
+  PRINT("Stopping MSCCL++ proxy threads\n");
+  MSCCLPPCHECK(mscclppProxyStop(args->comms[0]));
   return testSuccess;
 }
 
@@ -508,7 +550,7 @@ testResult_t run() {
   void* sendbuff;
   void* recvbuff;
   void* expected;
-  size_t sendBytes = 10, recvBytes = 10;
+  size_t sendBytes, recvBytes;
 
   mscclppTestEngine.getBuffSize(&sendBytes, &recvBytes, (size_t)maxBytes, (size_t)totalProcs);
 
@@ -531,7 +573,7 @@ testResult_t run() {
 
   const char* timeStr = report_cputime ? "cputime" : "time";
   PRINT("#\n");
-  PRINT("# %10s  %12s           out-of-place                       in-place          \n", "",
+  PRINT("# %10s  %12s           in-place                       out-of-place          \n", "",
         "");
   PRINT("# %10s  %12s  %7s  %6s  %6s  %6s  %7s  %6s  %6s  %6s\n", "size", "count", timeStr, "algbw", "busbw", "#wrong",
         timeStr, "algbw", "busbw", "#wrong");
@@ -545,6 +587,7 @@ testResult_t run() {
   thread.args.stepbytes = stepBytes;
   thread.args.stepfactor = stepFactor;
   thread.args.localRank = localRank;
+  thread.args.ranksPerNode = ranksPerNode;
 
   thread.args.totalProcs = totalProcs;
   thread.args.proc = proc;
diff --git a/tests/common.h b/tests/common.h
index 5e0dca5e..655f2d85 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -18,17 +18,18 @@
 #include <mpi.h>
 #endif // MSCCLPP_USE_MPI_FOR_TESTS
 
-#define CUDACHECK(cmd) do {                         \
-  cudaError_t err = cmd;                            \
-  if( err != cudaSuccess ) {                        \
-    char hostname[1024];                            \
-    getHostName(hostname, 1024);                    \
-    printf("%s: Test CUDA failure %s:%d '%s'\n",    \
-         hostname,                                  \
-        __FILE__,__LINE__,cudaGetErrorString(err)); \
-    return testCudaError;                           \
-  }                                                 \
-} while(0)
+__constant__ mscclppDevConn_t constDevConns[16];
+
+#define CUDACHECK(cmd)                                                                                                 \
+  do {                                                                                                                 \
+    cudaError_t err = cmd;                                                                                             \
+    if (err != cudaSuccess) {                                                                                          \
+      char hostname[1024];                                                                                             \
+      getHostName(hostname, 1024);                                                                                     \
+      printf("%s: Test CUDA failure %s:%d '%s'\n", hostname, __FILE__, __LINE__, cudaGetErrorString(err));             \
+      return testCudaError;                                                                                            \
+    }                                                                                                                  \
+  } while (0)
 
 // Propagate errors up
 #define MSCCLPPCHECK(cmd)                                                                                              \
@@ -98,6 +99,7 @@ struct threadArgs
   int nGpus;
   int* gpus;
   int localRank;
+  int ranksPerNode;
   void** sendbuffs;
   size_t sendBytes;
   size_t sendInplaceOffset;

From 9c1be4c03a685f1df8c8a455075f53e0a374420c Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Sat, 1 Apr 2023 05:10:14 +0000
Subject: [PATCH 06/23] fix

---
 tests/allgather_test3.cu | 15 +++++++++-----
 tests/common.cu          | 43 +++++++++++++++++++++++-----------------
 tests/common.h           |  4 +---
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
index 53628004..1b5d28fe 100644
--- a/tests/allgather_test3.cu
+++ b/tests/allgather_test3.cu
@@ -4,12 +4,14 @@
  * See LICENSE.txt for license information
  ************************************************************************/
 
+#include "comm.h"
 #include "common.h"
-#include "cuda_runtime.h"
 
+#include <cuda_runtime.h>
 #include <string>
 
 #define ALIGN 4
+__constant__ mscclppDevConn_t constDevConns[16];
 
 __device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, int nelemsPerGPU)
 {
@@ -98,7 +100,6 @@ testResult_t AllGatherInitData(struct threadArgs* args, int in_place) {
   delete dataHost;
   // TODO: need to init expected data here
   CUDACHECK(cudaDeviceSynchronize());
-
   return testSuccess;
 }
 
@@ -110,9 +111,10 @@ void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, doubl
   *busBw = baseBw * factor;
 }
 
-testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count)
+testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count, mscclppComm_t comm, cudaStream_t stream)
 {
-  // NCCLCHECK(ncclAllGather(sendbuff, recvbuff, count, type, comm, stream));
+  int worldSize = comm->nRanks;
+  kernel<<<1, 32 * (worldSize - 1), 0, stream>>>(comm->rank , worldSize, count / sizeof(int), 0);
   return testSuccess;
 }
 
@@ -127,7 +129,10 @@ void AllGatherGetBuffSize(size_t *sendcount, size_t *recvcount, size_t count, in
 testResult_t AllGatherRunTest(struct threadArgs* args)
 {
   args->collTest = &allGatherTest;
-
+    mscclppDevConn_t* devConns;
+  int nCons;
+  MSCCLPPCHECK(mscclppGetAllDeviceConnections(args->comms[0], &devConns, &nCons));
+  CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * nCons));
   TESTCHECK(TimeTest(args));
   return testSuccess;
 }
diff --git a/tests/common.cu b/tests/common.cu
index 20b7e9b3..7d849bd2 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -123,7 +123,8 @@ testResult_t startColl(struct threadArgs* args, int in_place, int iter) {
     char* sendBuff = ((char*)args->sendbuffs[i]) + shift;
 
     TESTCHECK(args->collTest->runColl((void*)(in_place ? recvBuff + args->sendInplaceOffset * rank : sendBuff),
-                                      (void*)(in_place ? recvBuff + args->recvInplaceOffset * rank : recvBuff), count));
+                                      (void*)(in_place ? recvBuff + args->recvInplaceOffset * rank : recvBuff), count,
+                                      args->comms[0], args->streams[i]));
   }
   return testSuccess;
 }
@@ -252,10 +253,20 @@ testResult_t BenchTime(struct threadArgs* args, int in_place) {
   Barrier(args);
 
   // Performance Benchmark
+  cudaGraph_t graph;
+  cudaGraphExec_t graphExec;
+  CUDACHECK(cudaStreamBeginCapture(args->streams[0], cudaStreamCaptureModeGlobal));
   timer tim;
   for (int iter = 0; iter < iters; iter++) {
     TESTCHECK(startColl(args, in_place, iter));
   }
+  CUDACHECK(cudaStreamEndCapture(args->streams[0], &graph));
+  CUDACHECK(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+
+  // Launch the graph
+  Barrier(args);
+  tim.reset();
+  CUDACHECK(cudaGraphLaunch(graphExec, args->streams[0]));
 
   double cputimeSec = tim.elapsed()/(iters);
   TESTCHECK(completeColl(args));
@@ -264,6 +275,9 @@ testResult_t BenchTime(struct threadArgs* args, int in_place) {
   deltaSec = deltaSec/(iters);
   Allreduce(args, &deltaSec, average);
 
+  CUDACHECK(cudaGraphExecDestroy(graphExec));
+  CUDACHECK(cudaGraphDestroy(graph));
+
   double algBw, busBw;
   args->collTest->getBw(count, 1, deltaSec, &algBw, &busBw, args->totalProcs * args->nThreads * args->nGpus);
 
@@ -307,18 +321,26 @@ testResult_t TimeTest(struct threadArgs* args) {
 
   // Warm-up for large size
   setupArgs(args->maxbytes, args);
+  TESTCHECK(args->collTest->initData(args, 1));
   for (int iter = 0; iter < warmup_iters; iter++) {
-    TESTCHECK(startColl(args, 0, iter));
+    TESTCHECK(startColl(args, 1, iter));
   }
   TESTCHECK(completeColl(args));
 
   // Warm-up for small size
   setupArgs(args->minbytes, args);
   for (int iter = 0; iter < warmup_iters; iter++) {
-    TESTCHECK(startColl(args, 0, iter));
+    TESTCHECK(startColl(args, 1, iter));
   }
   TESTCHECK(completeColl(args));
 
+  PRINT("#\n");
+  PRINT("# %10s  %12s           in-place                       out-of-place          \n", "",
+        "");
+  PRINT("# %10s  %12s  %7s  %6s  %6s  %6s  %7s  %6s  %6s  %6s\n", "size", "count", "time", "algbw", "busbw", "#wrong",
+        "time", "algbw", "busbw", "#wrong");
+  PRINT("# %10s  %12s  %7s  %6s  %6s  %5s  %7s  %6s  %6s  %5s\n", "(B)", "(elements)", "(us)", "(GB/s)", "(GB/s)", "",
+        "(us)", "(GB/s)", "(GB/s)", "");
   // Benchmark
   for (size_t size = args->minbytes; size<=args->maxbytes; size = ((args->stepfactor > 1) ? size*args->stepfactor : size+args->stepbytes)) {
       setupArgs(size, args);
@@ -355,12 +377,6 @@ testResult_t setupMscclppConnections(int rank, int ranksPerNode, int worldSize,
 
   MSCCLPPCHECK(mscclppConnectionSetup(comm));
 
-  mscclppDevConn_t* devConns;
-  int nCons;
-  MSCCLPPCHECK(mscclppGetAllDeviceConnections(comm, &devConns, &nCons));
-
-  CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * nCons));
-
   return testSuccess;
 }
 
@@ -571,15 +587,6 @@ testResult_t run() {
 
   fflush(stdout);
 
-  const char* timeStr = report_cputime ? "cputime" : "time";
-  PRINT("#\n");
-  PRINT("# %10s  %12s           in-place                       out-of-place          \n", "",
-        "");
-  PRINT("# %10s  %12s  %7s  %6s  %6s  %6s  %7s  %6s  %6s  %6s\n", "size", "count", timeStr, "algbw", "busbw", "#wrong",
-        timeStr, "algbw", "busbw", "#wrong");
-  PRINT("# %10s  %12s  %7s  %6s  %6s  %5s  %7s  %6s  %6s  %5s\n", "(B)", "(elements)", "(us)", "(GB/s)", "(GB/s)", "",
-        "(us)", "(GB/s)", "(GB/s)", "");
-
   struct testThread thread = {0};
 
   thread.args.minbytes = minBytes;
diff --git a/tests/common.h b/tests/common.h
index 655f2d85..7ff3a340 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -18,8 +18,6 @@
 #include <mpi.h>
 #endif // MSCCLPP_USE_MPI_FOR_TESTS
 
-__constant__ mscclppDevConn_t constDevConns[16];
-
 #define CUDACHECK(cmd)                                                                                                 \
   do {                                                                                                                 \
     cudaError_t err = cmd;                                                                                             \
@@ -72,7 +70,7 @@ struct testColl {
       size_t count, int nranks);
   testResult_t (*initData)(struct threadArgs* args, int in_place);
   void (*getBw)(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks);
-  testResult_t (*runColl)(void* sendbuff, void* recvbuff, size_t count);
+  testResult_t (*runColl)(void* sendbuff, void* recvbuff, size_t count, mscclppComm_t comm, cudaStream_t stream);
 };
 
 struct testEngine

From f02ba52b0add8c5a7689e48740b4dc0a8aa216a0 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Sat, 1 Apr 2023 07:10:20 +0000
Subject: [PATCH 07/23] fix

---
 tests/common.cu | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/common.cu b/tests/common.cu
index 7d849bd2..fc34a1e5 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -353,7 +353,7 @@ testResult_t TimeTest(struct threadArgs* args) {
   return testSuccess;
 }
 
-testResult_t setupMscclppConnections(int rank, int ranksPerNode, int worldSize, mscclppComm_t comm, void* dataDst,
+testResult_t setupMscclppConnections(int rank, int worldSize, int ranksPerNode, mscclppComm_t comm, void* dataDst,
                                         size_t dataSize)
 {
   int thisNode = rank / ranksPerNode;
@@ -382,9 +382,10 @@ testResult_t setupMscclppConnections(int rank, int ranksPerNode, int worldSize,
 
 testResult_t threadRunTests(struct threadArgs* args)
 {
+  PRINT("# Setting up the connection in MSCCL++\n");
   TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->ranksPerNode, args->comms[0],
-                                    args->recvbuffs[0], args->expectedBytes));
-  PRINT("Setting up the connection in MSCCL++\n");
+                                    args->recvbuffs[0], args->maxbytes));
+  PRINT("# Launching MSCCL++ proxy threads\n");
   MSCCLPPCHECK(mscclppProxyLaunch(args->comms[0]));
   TESTCHECK(mscclppTestEngine.runTest(args));
   PRINT("Stopping MSCCL++ proxy threads\n");
@@ -577,7 +578,7 @@ testResult_t run() {
   PRINT("# Initializing MSCCL++\n");
 
   mscclppComm_t comms;
-  MSCCLPPCHECK(mscclppCommInitRank(&comms, totalProcs, ip_port.c_str(), localRank));
+  MSCCLPPCHECK(mscclppCommInitRank(&comms, totalProcs, ip_port.c_str(), rank));
 
   int error = 0;
   double bw = 0.0;

From 212f92cfba7744a3ded7919955ff9502f02a3dbf Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Sun, 2 Apr 2023 11:13:43 +0000
Subject: [PATCH 08/23] fix

---
 tests/allgather_test3.cu | 135 ++++++++++++++++++++++++++++++---------
 tests/common.cu          | 116 ++++++++++++++++-----------------
 tests/common.h           |  10 +--
 3 files changed, 169 insertions(+), 92 deletions(-)

diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
index 1b5d28fe..56f08cd3 100644
--- a/tests/allgather_test3.cu
+++ b/tests/allgather_test3.cu
@@ -13,47 +13,121 @@
 #define ALIGN 4
 __constant__ mscclppDevConn_t constDevConns[16];
 
-__device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, int nelemsPerGPU)
+__device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, size_t nelemsPerGPU)
 {
   // this allgather is really simple and implemented as an alltoall
 
   // this thread's role is a sender role
   // put your data asynchronously
-  devConn.put(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
+  if (threadIdx.x % 32 != 0)
+    devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
   // make sure everyone is put their data before some thread randomly blocks everyone else in signal
   __syncthreads();
   // push with flag and sync to make sure the data is received
-  devConn.signal();
+  if (threadIdx.x % 32 != 0)
+    devConn.flush();
 
   // this thread's role is a receiver role. wait on the semaphore to make sure the data is ready
-  devConn.wait();
+  if (threadIdx.x % 32 != 0)
+    devConn.wait();
 }
 
-__device__ void allgather1(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, int nelemsPerGPU)
+__device__ void localAllGather(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
+                               uint64_t offset, uint64_t size)
 {
   // this allgather algorithm works as follows:
-  // Step 1: GPU rank i sends data to GPU rank (i+1) % world_size
-  // Step 2: GPU rank i waits for data from GPU rank (i+2) % world_size
+  // Step 1: GPU rank i sends data to GPU rank (i+1) % nranksPerNode
+  // and waits for data from GPU rank (i-1) % nranksPerNode
+  // Step 2: GPU rank i sends data to GPU rank (i+2) % nranksPerNode
   // ...
   // This order is much better for DMA engine for NVLinks
-
-  for (int i = 1; i < world_size; i++) {
-    __syncthreads();
-    if (remoteRank != ((rank + i) % world_size))
-      continue;
-    // put your data to GPU (rank+i) % world_size and signal all in one call
-    devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
+  for (int i = 1; i < nranksPerNode; i++) {
+    if ((remoteRank % nranksPerNode) == ((rank + i) % nranksPerNode)) {
+      // put your data to GPU (rank+i) % nranksPerNode and signal in one call
+      if ((threadIdx.x % 32) == 0)
+        devConn.putWithSignalAndFlush(offset, size);
+    }
+    // wait for the data from GPU (rank-i) % nranksPerNode to arrive
+    if ((remoteRank % nranksPerNode) == ((rank - i + nranksPerNode) % nranksPerNode)) {
+      if ((threadIdx.x % 32) == 0)
+        devConn.wait();
+    }
+    asm volatile("bar.sync %0, %1;" ::"r"(11), "r"((nranksPerNode - 1) * 32) : "memory");
   }
-  // all connections wait for the signal from the sender
-  devConn.wait();
 }
 
-__global__ void kernel(int rank, int world_size, int nelemsPerGPU, int kernel)
+__device__ void allgather1(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
+                           size_t nelemsPerGPU)
 {
-  // only use a single thread from each warp
-  if (threadIdx.x % 32 != 0)
-    return;
+  localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, rank * nelemsPerGPU * sizeof(int),
+                 nelemsPerGPU * sizeof(int));
+}
 
+__device__ void allgather2(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
+                           size_t nelemsPerGPU)
+{
+  // this allgather is a pipelined and hierarchical one and only works for two nodes
+  // it is implemented as follows:
+  // Step 1: each node does a local allgather and concurrently,
+  // local GPU i exchange (piplineSize-1)/pipelineSize portion of their data with
+  // its cross-node neighbor (local GPU i on the other node) via IB
+  // Step 2: each node does a local allgather again with the data just received from its
+  // cross-node neighbor in step 1, and concurrently, exchange the rest of the data with
+  // its cross-node neighbor
+  // Step 3: each node does a local allgather for the last time with the rest of the data
+
+  int pipelineSize = 3;
+
+  // Step 1
+  // local allgather
+  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
+    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, rank * nelemsPerGPU * sizeof(int),
+                   nelemsPerGPU * sizeof(int));
+  }
+  // cross-node exchange
+  if (remoteRank % nranksPerNode == rank % nranksPerNode) {
+    // opposite side
+    if ((threadIdx.x % 32) == 0)
+      devConn.putWithSignalAndFlush(rank * nelemsPerGPU * sizeof(int),
+                                    (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize * sizeof(int));
+    if ((threadIdx.x % 32) == 0)
+      devConn.wait();
+  }
+
+  __syncthreads();
+
+  // Step 2
+  // local allgather
+  int otherNghr = (rank + nranksPerNode) % world_size;
+  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
+    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, otherNghr * nelemsPerGPU * sizeof(int),
+                   (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize * sizeof(int));
+  }
+
+  // cross-node exchange
+  if (remoteRank % nranksPerNode == rank % nranksPerNode) {
+    // opposite side
+    if ((threadIdx.x % 32) == 0)
+      devConn.putWithSignalAndFlush((rank * nelemsPerGPU + (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize) *
+                                      sizeof(int),
+                                    nelemsPerGPU / pipelineSize * sizeof(int));
+    if ((threadIdx.x % 32) == 0)
+      devConn.wait();
+  }
+
+  __syncthreads();
+
+  // Step 3
+  // local allgather
+  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
+    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank,
+                   (otherNghr * nelemsPerGPU + (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize) * sizeof(int),
+                   nelemsPerGPU / pipelineSize * sizeof(int));
+  }
+}
+
+__global__ void kernel(int rank, int world_size, int nranksPerNode, size_t nelemsPerGPU, int kernel)
+{
   // find the mapping between remoteRank and devConns
   int warpId = threadIdx.x / 32;
   int remoteRank = (warpId < rank) ? warpId : warpId + 1;
@@ -63,7 +137,9 @@ __global__ void kernel(int rank, int world_size, int nelemsPerGPU, int kernel)
   if (kernel == 0)
     allgather0(devConn, rank, world_size, remoteRank, nelemsPerGPU);
   else if (kernel == 1)
-    allgather1(devConn, rank, world_size, remoteRank, nelemsPerGPU);
+    allgather1(devConn, rank, world_size, nranksPerNode, remoteRank, nelemsPerGPU);
+  else if (kernel == 2)
+    allgather2(devConn, rank, world_size, nranksPerNode, remoteRank, nelemsPerGPU);
 }
 
 void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* paramcount, size_t* sendInplaceOffset,
@@ -78,19 +154,19 @@ void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* par
 }
 
 testResult_t AllGatherInitData(struct threadArgs* args, int in_place) {
-  size_t sendcount = args->sendBytes;
+  // size_t sendcount = args->sendBytes;
   size_t recvcount = args->expectedBytes;
-  int nranks = args->totalProcs;
+  // int nranks = args->totalProcs;
 
   CUDACHECK(cudaSetDevice(args->gpus[0]));
   int rank = args->proc;
   CUDACHECK(cudaMemset(args->recvbuffs[0], 0, args->expectedBytes));
-  void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
+  // void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
 
   int* dataHost = new int[recvcount];
   for (int i = 0; i < static_cast<int>(recvcount); i++) {
     int val = i + 1;
-    if (i / args->ranksPerNode == rank) {
+    if (i / args->nranksPerNode == rank) {
       dataHost[i] = val;
     } else {
       dataHost[i] = 0;
@@ -111,10 +187,11 @@ void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, doubl
   *busBw = baseBw * factor;
 }
 
-testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, size_t count, mscclppComm_t comm, cudaStream_t stream)
+testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, int nranksPerNode, size_t count, mscclppComm_t comm,
+                              cudaStream_t stream)
 {
   int worldSize = comm->nRanks;
-  kernel<<<1, 32 * (worldSize - 1), 0, stream>>>(comm->rank , worldSize, count / sizeof(int), 0);
+  kernel<<<1, 32 * (worldSize - 1), 0, stream>>>(comm->rank, worldSize, nranksPerNode, count / sizeof(int), 1);
   return testSuccess;
 }
 
@@ -129,9 +206,9 @@ void AllGatherGetBuffSize(size_t *sendcount, size_t *recvcount, size_t count, in
 testResult_t AllGatherRunTest(struct threadArgs* args)
 {
   args->collTest = &allGatherTest;
-    mscclppDevConn_t* devConns;
+  mscclppDevConn_t* devConns;
   int nCons;
-  MSCCLPPCHECK(mscclppGetAllDeviceConnections(args->comms[0], &devConns, &nCons));
+  MSCCLPPCHECK(mscclppGetAllDeviceConnections(args->comm, &devConns, &nCons));
   CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * nCons));
   TESTCHECK(TimeTest(args));
   return testSuccess;
diff --git a/tests/common.cu b/tests/common.cu
index fc34a1e5..4343c181 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -29,13 +29,14 @@ static size_t maxBytes = 32*1024*1024;
 static size_t stepBytes = 1*1024*1024;
 static size_t stepFactor = 1;
 static int datacheck = 1;
-static int warmup_iters = 5;
+static int warmup_iters = 10;
 static int iters = 20;
 static int timeout = 0;
 static int report_cputime = 0;
 // Report average iteration time: (0=RANK0,1=AVG,2=MIN,3=MAX)
 static int average = 1;
 static std::string ip_port;
+static int cudaGraphLaunches = 10;
 
 #define NUM_BLOCKS 32
 
@@ -117,54 +118,46 @@ testResult_t startColl(struct threadArgs* args, int in_place, int iter) {
   size_t steps = totalnbytes ? args->maxbytes / totalnbytes : 1;
   size_t shift = totalnbytes * (iter % steps);
 
-  for (int i = 0; i < args->nGpus; i++) {
-    int rank = ((args->proc*args->nThreads + args->thread)*args->nGpus + i);
-    char* recvBuff = ((char*)args->recvbuffs[i]) + shift;
-    char* sendBuff = ((char*)args->sendbuffs[i]) + shift;
+  int rank = ((args->proc * args->nThreads + args->thread) * args->nGpus);
+  char* recvBuff = ((char*)args->recvbuffs[0]) + shift;
+  char* sendBuff = ((char*)args->sendbuffs[0]) + shift;
 
-    TESTCHECK(args->collTest->runColl((void*)(in_place ? recvBuff + args->sendInplaceOffset * rank : sendBuff),
-                                      (void*)(in_place ? recvBuff + args->recvInplaceOffset * rank : recvBuff), count,
-                                      args->comms[0], args->streams[i]));
-  }
+  TESTCHECK(args->collTest->runColl((void*)(in_place ? recvBuff + args->sendInplaceOffset * rank : sendBuff),
+                                    (void*)(in_place ? recvBuff + args->recvInplaceOffset * rank : recvBuff),
+                                    args->nranksPerNode, count, args->comm, args->stream));
   return testSuccess;
 }
 
-testResult_t testStreamSynchronize(int ngpus, cudaStream_t* streams)
+testResult_t testStreamSynchronize(cudaStream_t stream)
 {
   cudaError_t cudaErr;
-  int remaining = ngpus;
-  int* done = (int*)malloc(sizeof(int) * ngpus);
-  memset(done, 0, sizeof(int) * ngpus);
   timer tim;
 
-  while (remaining) {
-    int idle = 1;
-    for (int i = 0; i < ngpus; i++) {
-      if (done[i])
-        continue;
-
-      cudaErr = cudaStreamQuery(streams[i]);
-      if (cudaErr == cudaSuccess) {
-        done[i] = 1;
-        remaining--;
-        idle = 0;
-        continue;
-      }
-
-      if (cudaErr != cudaErrorNotReady)
-        CUDACHECK(cudaErr);
+  while (true) {
+    cudaErr = cudaStreamQuery(stream);
+    if (cudaErr == cudaSuccess) {
+      break;
     }
 
-    // We might want to let other threads (including NCCL threads) use the CPU.
-    if (idle)
-      sched_yield();
+    if (cudaErr != cudaErrorNotReady)
+      CUDACHECK(cudaErr);
+
+    double delta = tim.elapsed();
+    if (delta > timeout && timeout > 0) {
+      char hostname[1024];
+      getHostName(hostname, 1024);
+      printf("%s: Test timeout (%ds) %s:%d\n", hostname, timeout, __FILE__, __LINE__);
+      return testTimeout;
+    }
+
+    // We might want to let other threads (including MSCCLPP threads) use the CPU.
+    sched_yield();
   }
-  free(done);
   return testSuccess;
 }
 
 testResult_t completeColl(struct threadArgs* args) {
-  TESTCHECK(testStreamSynchronize(args->nGpus, args->streams));
+  TESTCHECK(testStreamSynchronize(args->stream));
   return testSuccess;
 }
 
@@ -255,24 +248,26 @@ testResult_t BenchTime(struct threadArgs* args, int in_place) {
   // Performance Benchmark
   cudaGraph_t graph;
   cudaGraphExec_t graphExec;
-  CUDACHECK(cudaStreamBeginCapture(args->streams[0], cudaStreamCaptureModeGlobal));
+  CUDACHECK(cudaStreamBeginCapture(args->stream, cudaStreamCaptureModeGlobal));
   timer tim;
   for (int iter = 0; iter < iters; iter++) {
     TESTCHECK(startColl(args, in_place, iter));
   }
-  CUDACHECK(cudaStreamEndCapture(args->streams[0], &graph));
+  CUDACHECK(cudaStreamEndCapture(args->stream, &graph));
   CUDACHECK(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
 
   // Launch the graph
   Barrier(args);
   tim.reset();
-  CUDACHECK(cudaGraphLaunch(graphExec, args->streams[0]));
+  for (int l = 0; l < cudaGraphLaunches; ++l) {
+    CUDACHECK(cudaGraphLaunch(graphExec, args->stream));
+  }
 
   double cputimeSec = tim.elapsed()/(iters);
   TESTCHECK(completeColl(args));
 
   double deltaSec = tim.elapsed();
-  deltaSec = deltaSec/(iters);
+  deltaSec = deltaSec/(iters)/(cudaGraphLaunches);
   Allreduce(args, &deltaSec, average);
 
   CUDACHECK(cudaGraphExecDestroy(graphExec));
@@ -383,13 +378,13 @@ testResult_t setupMscclppConnections(int rank, int worldSize, int ranksPerNode,
 testResult_t threadRunTests(struct threadArgs* args)
 {
   PRINT("# Setting up the connection in MSCCL++\n");
-  TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->ranksPerNode, args->comms[0],
+  TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->nranksPerNode, args->comm,
                                     args->recvbuffs[0], args->maxbytes));
   PRINT("# Launching MSCCL++ proxy threads\n");
-  MSCCLPPCHECK(mscclppProxyLaunch(args->comms[0]));
+  MSCCLPPCHECK(mscclppProxyLaunch(args->comm));
   TESTCHECK(mscclppTestEngine.runTest(args));
   PRINT("Stopping MSCCL++ proxy threads\n");
-  MSCCLPPCHECK(mscclppProxyStop(args->comms[0]));
+  MSCCLPPCHECK(mscclppProxyStop(args->comm));
   return testSuccess;
 }
 
@@ -411,6 +406,7 @@ int main(int argc, char* argv[]) {
     {"warmup_iters", required_argument, 0, 'w'},
     {"check", required_argument, 0, 'c'},
     {"timeout", required_argument, 0, 'T'},
+    {"cudagraph", required_argument, 0, 'G'},
     {"report_cputime", required_argument, 0, 'C'},
     {"average", required_argument, 0, 'a'},
     {"ip_port", required_argument, 0, 'P'},
@@ -460,6 +456,9 @@ int main(int argc, char* argv[]) {
       case 'T':
         timeout = strtol(optarg, NULL, 0);
         break;
+      case 'G':
+        cudaGraphLaunches = strtol(optarg, NULL, 0);
+        break;
       case 'C':
         report_cputime = strtol(optarg, NULL, 0);
         break;
@@ -481,6 +480,7 @@ int main(int argc, char* argv[]) {
             "[-w,--warmup_iters <warmup iteration count>] \n\t"
             "[-c,--check <0/1>] \n\t"
             "[-T,--timeout <time in seconds>] \n\t"
+            "[-G,--cudagraph <num graph launches>] \n\t"
             "[-C,--report_cputime <0/1>] \n\t"
             "[-a,--average <0/1/2/3> report average iteration time <0=RANK0/1=AVG/2=MIN/3=MAX>] \n\t"
             "[-P,--ip_port <ip port for bootstrap>] \n\t"
@@ -507,7 +507,7 @@ int main(int argc, char* argv[]) {
 
 testResult_t run() {
   int totalProcs = 1, proc = 0;
-  int ranksPerNode = 0, localRank = 0;
+  int nranksPerNode = 0, localRank = 0;
   char hostname[1024];
   getHostName(hostname, 1024);
 
@@ -516,16 +516,16 @@ testResult_t run() {
   MPI_Comm_rank(MPI_COMM_WORLD, &proc);
   MPI_Comm shmcomm;
   MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &shmcomm);
-  MPI_Comm_size(shmcomm, &ranksPerNode);
+  MPI_Comm_size(shmcomm, &nranksPerNode);
   MPI_Comm_free(&shmcomm);
-  localRank = proc % ranksPerNode;
+  localRank = proc % nranksPerNode;
 #endif
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
 
-  PRINT("# minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip port: %s\n", minBytes,
-        maxBytes, (stepFactor > 1) ? stepFactor : stepBytes, (stepFactor > 1) ? "factor" : "bytes", warmup_iters, iters,
-        datacheck, ip_port.c_str());
+  PRINT("# minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip port: %s graph: %d\n",
+        minBytes, maxBytes, (stepFactor > 1) ? stepFactor : stepBytes, (stepFactor > 1) ? "factor" : "bytes",
+        warmup_iters, iters, datacheck, ip_port.c_str(), cudaGraphLaunches);
   PRINT("#\n");
   PRINT("# Using devices\n");
 
@@ -537,9 +537,11 @@ testResult_t run() {
   int cudaDev = localRank;
   int rank = proc;
   cudaDeviceProp prop;
+  char busIdChar[] = "00000000:00:00.0";
   CUDACHECK(cudaGetDeviceProperties(&prop, cudaDev));
-  len += snprintf(line + len, MAX_LINE - len, "#  Rank %2d Pid %6d on %10s device %2d [0x%02x] %s\n", rank, getpid(),
-                  hostname, cudaDev, prop.pciBusID, prop.name);
+  CUDACHECK(cudaDeviceGetPCIBusId(busIdChar, sizeof(busIdChar), cudaDev));
+  len += snprintf(line + len, MAX_LINE - len, "#  Rank %2d Pid %6d on %10s device %2d [%s] %s\n", rank, getpid(),
+                  hostname, cudaDev, busIdChar, prop.name);
   maxMem = std::min(maxMem, prop.totalGlobalMem);
 
 #if MSCCLPP_USE_MPI_FOR_TESTS
@@ -577,8 +579,8 @@ testResult_t run() {
   PRINT("#\n");
   PRINT("# Initializing MSCCL++\n");
 
-  mscclppComm_t comms;
-  MSCCLPPCHECK(mscclppCommInitRank(&comms, totalProcs, ip_port.c_str(), rank));
+  mscclppComm_t comm;
+  MSCCLPPCHECK(mscclppCommInitRank(&comm, totalProcs, ip_port.c_str(), rank));
 
   int error = 0;
   double bw = 0.0;
@@ -588,14 +590,14 @@ testResult_t run() {
 
   fflush(stdout);
 
-  struct testThread thread = {0};
+  struct testThread thread;
 
   thread.args.minbytes = minBytes;
   thread.args.maxbytes = maxBytes;
   thread.args.stepbytes = stepBytes;
   thread.args.stepfactor = stepFactor;
   thread.args.localRank = localRank;
-  thread.args.ranksPerNode = ranksPerNode;
+  thread.args.nranksPerNode = nranksPerNode;
 
   thread.args.totalProcs = totalProcs;
   thread.args.proc = proc;
@@ -606,8 +608,8 @@ testResult_t run() {
   thread.args.sendbuffs = &sendbuff;
   thread.args.recvbuffs = &recvbuff;
   thread.args.expected = &expected;
-  thread.args.comms = &comms;
-  thread.args.streams = &stream;
+  thread.args.comm = comm;
+  thread.args.stream = stream;
 
   thread.args.errors = &error;
   thread.args.bw = &bw;
@@ -618,9 +620,7 @@ testResult_t run() {
   thread.func = threadRunTests;
   TESTCHECK(thread.func(&thread.args));
 
-  // Wait for other threads and accumulate stats and errors
-  TESTCHECK(thread.ret);
-  MSCCLPPCHECK(mscclppCommDestroy(comms));
+  MSCCLPPCHECK(mscclppCommDestroy(comm));
 
   // Free off CUDA allocated memory
   if (sendbuff)
diff --git a/tests/common.h b/tests/common.h
index 7ff3a340..b7e2c226 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -70,7 +70,8 @@ struct testColl {
       size_t count, int nranks);
   testResult_t (*initData)(struct threadArgs* args, int in_place);
   void (*getBw)(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks);
-  testResult_t (*runColl)(void* sendbuff, void* recvbuff, size_t count, mscclppComm_t comm, cudaStream_t stream);
+  testResult_t (*runColl)(void* sendbuff, void* recvbuff, int nranksPerNode, size_t count, mscclppComm_t comm,
+                          cudaStream_t stream);
 };
 
 struct testEngine
@@ -97,14 +98,14 @@ struct threadArgs
   int nGpus;
   int* gpus;
   int localRank;
-  int ranksPerNode;
+  int nranksPerNode;
   void** sendbuffs;
   size_t sendBytes;
   size_t sendInplaceOffset;
   void** recvbuffs;
   size_t recvInplaceOffset;
-  mscclppComm_t* comms;
-  cudaStream_t* streams;
+  mscclppComm_t comm;
+  cudaStream_t stream;
 
   void** expected;
   size_t expectedBytes;
@@ -123,7 +124,6 @@ struct testThread
   pthread_t thread;
   threadFunc_t func;
   struct threadArgs args;
-  testResult_t ret;
 };
 
 // Provided by common.cu

From 617f39daf13f69412d7a79d8a4ebf6501bab3a08 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Mon, 3 Apr 2023 03:45:36 +0000
Subject: [PATCH 09/23] change makefile

---
 Makefile        | 19 ++++++++++++-------
 tests/common.cu |  4 ++--
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/Makefile b/Makefile
index 10fc9f3e..a46c691a 100644
--- a/Makefile
+++ b/Makefile
@@ -133,21 +133,24 @@ TESTSSRCS := $(addprefix $(TESTSDIR)/,bootstrap_test.cc allgather_test.cu)
 TESTSOBJS := $(patsubst %.cc,%.o,$(TESTSSRCS)) $(patsubst %.cu,%.o,$(TESTSSRCS))
 TESTSOBJTARGETS := $(TESTSOBJS:%=$(BUILDDIR)/$(OBJDIR)/%)
 TESTSBINS       := $(patsubst %.o,$(BUILDDIR)/$(BINDIR)/%,$(TESTSOBJS))
-TESTSPERFBIN := $(BUILDDIR)/$(BINDIR)/allgather_test3
+
+MSCLLPPTESTSOBJSDIR:= $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)
+MSCLLPPTESTBINFILESLIST := allgather_test3
+MSCLLPPTESTBINS         := $(MSCLLPPTESTBINFILESLIST:%=$(BUILDDIR)/$(BINDIR)/$(TESTSDIR)/%_perf)
 
 INCLUDE := -Isrc -Isrc/include
 
-.PHONY: all build lib tests clean
+.PHONY: all build lib tests mscclpp_test clean
 
 all: build
 
-build: lib tests
+build: lib tests mscclpp_test
 
 lib: $(LIBOBJTARGETS) $(INCTARGETS) $(LIBTARGET)
 
 tests: $(TESTSBINS)
 
-allgather_perf: $(TESTSPERFBIN)
+mscclpp_test: $(LIBTARGET) $(MSCLLPPTESTBINS)
 
 cpplint:
 	clang-format-12 -style=file --verbose --Werror --dry-run $(CPPSOURCES)
@@ -175,9 +178,6 @@ $(LIBTARGET): $(LIBOBJTARGETS)
 	ln -sf $(LIBTARGET) $(BUILDDIR)/$(LIBDIR)/$(LIBSONAME)
 
 # Compile .cc tests
-$(TESTSPERFBIN): $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/allgather_test3.o $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/common.o $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/timer.o
-	$(NVCC) -o $@ $^ $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp
-
 $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)/%.o: $(TESTSDIR)/%.cc $(INCTARGETS)
 	@mkdir -p $(@D)
 	$(CXX) -o $@ -I$(BUILDDIR)/$(INCDIR) $(MPI_INC) $(CXXFLAGS) -c $< $(MPI_MACRO)
@@ -192,5 +192,10 @@ $(BUILDDIR)/$(BINDIR)/%: $(BUILDDIR)/$(OBJDIR)/%.o $(LIBTARGET)
 	@mkdir -p $(@D)
 	$(NVCC) -o $@ $< $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp
 
+# Comile mscclpp_test
+$(BUILDDIR)/$(BINDIR)/$(TESTSDIR)/%_perf: $(MSCLLPPTESTSOBJSDIR)/%.o $(MSCLLPPTESTSOBJSDIR)/common.o $(MSCLLPPTESTSOBJSDIR)/timer.o
+	@mkdir -p $(@D)
+	$(NVCC) -o $@ $^ $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp
+
 clean:
 	rm -rf $(BUILDDIR)
diff --git a/tests/common.cu b/tests/common.cu
index 4343c181..4c819c98 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -30,13 +30,13 @@ static size_t stepBytes = 1*1024*1024;
 static size_t stepFactor = 1;
 static int datacheck = 1;
 static int warmup_iters = 10;
-static int iters = 20;
+static int iters = 100;
 static int timeout = 0;
 static int report_cputime = 0;
 // Report average iteration time: (0=RANK0,1=AVG,2=MIN,3=MAX)
 static int average = 1;
 static std::string ip_port;
-static int cudaGraphLaunches = 10;
+static int cudaGraphLaunches = 15;
 
 #define NUM_BLOCKS 32
 

From d3c7cc721afb22905dc66088eb5cf9b25b23befe Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Mon, 3 Apr 2023 04:04:16 +0000
Subject: [PATCH 10/23] fix lint

---
 tests/allgather_test3.cu |  11 +-
 tests/common.cu          | 325 ++++++++++++++++++++-------------------
 2 files changed, 175 insertions(+), 161 deletions(-)

diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
index 56f08cd3..3c72209d 100644
--- a/tests/allgather_test3.cu
+++ b/tests/allgather_test3.cu
@@ -153,7 +153,8 @@ void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* par
   *paramcount = base;
 }
 
-testResult_t AllGatherInitData(struct threadArgs* args, int in_place) {
+testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
+{
   // size_t sendcount = args->sendBytes;
   size_t recvcount = args->expectedBytes;
   // int nranks = args->totalProcs;
@@ -179,11 +180,12 @@ testResult_t AllGatherInitData(struct threadArgs* args, int in_place) {
   return testSuccess;
 }
 
-void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks) {
+void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks)
+{
   double baseBw = (double)(count * typesize * nranks) / 1.0E9 / sec;
 
   *algBw = baseBw;
-  double factor = ((double)(nranks - 1))/((double)nranks);
+  double factor = ((double)(nranks - 1)) / ((double)nranks);
   *busBw = baseBw * factor;
 }
 
@@ -198,7 +200,8 @@ testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, int nranksPerNode,
 struct testColl allGatherTest = {"AllGather", AllGatherGetCollByteCount, AllGatherInitData, AllGatherGetBw,
                                  AllGatherRunColl};
 
-void AllGatherGetBuffSize(size_t *sendcount, size_t *recvcount, size_t count, int nranks) {
+void AllGatherGetBuffSize(size_t* sendcount, size_t* recvcount, size_t count, int nranks)
+{
   size_t paramcount, sendInplaceOffset, recvInplaceOffset;
   AllGatherGetCollByteCount(sendcount, recvcount, &paramcount, &sendInplaceOffset, &recvInplaceOffset, count, nranks);
 }
diff --git a/tests/common.cu b/tests/common.cu
index 4c819c98..58381a78 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -24,9 +24,9 @@ int is_main_proc = 0;
 thread_local int is_main_thread = 0;
 
 // Command line parameter defaults
-static size_t minBytes = 32*1024*1024;
-static size_t maxBytes = 32*1024*1024;
-static size_t stepBytes = 1*1024*1024;
+static size_t minBytes = 32 * 1024 * 1024;
+static size_t maxBytes = 32 * 1024 * 1024;
+static size_t stepBytes = 1 * 1024 * 1024;
 static size_t stepFactor = 1;
 static int datacheck = 1;
 static int warmup_iters = 10;
@@ -40,77 +40,82 @@ static int cudaGraphLaunches = 15;
 
 #define NUM_BLOCKS 32
 
-static double parsesize(const char *value) {
-    long long int units;
-    double size;
-    char size_lit;
+static double parsesize(const char* value)
+{
+  long long int units;
+  double size;
+  char size_lit;
 
-    int count = sscanf(value, "%lf %1s", &size, &size_lit);
+  int count = sscanf(value, "%lf %1s", &size, &size_lit);
 
-    switch (count) {
-    case 2:
-      switch (size_lit) {
-      case 'G':
-      case 'g':
-        units = 1024*1024*1024;
-        break;
-      case 'M':
-      case 'm':
-        units = 1024*1024;
-        break;
-      case 'K':
-      case 'k':
-        units = 1024;
-        break;
-      default:
-        return -1.0;
-      };
+  switch (count) {
+  case 2:
+    switch (size_lit) {
+    case 'G':
+    case 'g':
+      units = 1024 * 1024 * 1024;
       break;
-    case 1:
-      units = 1;
+    case 'M':
+    case 'm':
+      units = 1024 * 1024;
+      break;
+    case 'K':
+    case 'k':
+      units = 1024;
       break;
     default:
       return -1.0;
-    }
+    };
+    break;
+  case 1:
+    units = 1;
+    break;
+  default:
+    return -1.0;
+  }
 
-    return size * units;
+  return size * units;
 }
 
-void Barrier(struct threadArgs *args) {
+void Barrier(struct threadArgs* args)
+{
   thread_local int epoch = 0;
   static pthread_mutex_t lock[2] = {PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER};
   static pthread_cond_t cond[2] = {PTHREAD_COND_INITIALIZER, PTHREAD_COND_INITIALIZER};
   static int counter[2] = {0, 0};
 
   pthread_mutex_lock(&lock[epoch]);
-  if(++counter[epoch] == args->nThreads)
+  if (++counter[epoch] == args->nThreads)
     pthread_cond_broadcast(&cond[epoch]);
 
-  if(args->thread+1 == args->nThreads) {
-    while(counter[epoch] != args->nThreads)
+  if (args->thread + 1 == args->nThreads) {
+    while (counter[epoch] != args->nThreads)
       pthread_cond_wait(&cond[epoch], &lock[epoch]);
-    #ifdef MSCCLPP_USE_MPI_FOR_TESTS
-      MPI_Barrier(MPI_COMM_WORLD);
-    #endif
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+    MPI_Barrier(MPI_COMM_WORLD);
+#endif
     counter[epoch] = 0;
     pthread_cond_broadcast(&cond[epoch]);
-  }
-  else {
-    while(counter[epoch] != 0)
+  } else {
+    while (counter[epoch] != 0)
       pthread_cond_wait(&cond[epoch], &lock[epoch]);
   }
   pthread_mutex_unlock(&lock[epoch]);
   epoch ^= 1;
 }
 
-testResult_t AllocateBuffs(void **sendbuff, size_t sendBytes, void **recvbuff, size_t recvBytes, void **expected, size_t nbytes) {
-    CUDACHECK(cudaMalloc(sendbuff, nbytes));
-    CUDACHECK(cudaMalloc(recvbuff, nbytes));
-    if (datacheck) CUDACHECK(cudaMalloc(expected, recvBytes));
-    return testSuccess;
+testResult_t AllocateBuffs(void** sendbuff, size_t sendBytes, void** recvbuff, size_t recvBytes, void** expected,
+                           size_t nbytes)
+{
+  CUDACHECK(cudaMalloc(sendbuff, nbytes));
+  CUDACHECK(cudaMalloc(recvbuff, nbytes));
+  if (datacheck)
+    CUDACHECK(cudaMalloc(expected, recvBytes));
+  return testSuccess;
 }
 
-testResult_t startColl(struct threadArgs* args, int in_place, int iter) {
+testResult_t startColl(struct threadArgs* args, int in_place, int iter)
+{
   size_t count = args->nbytes;
 
   // Try to change offset for each iteration so that we avoid cache effects and catch race conditions in ptrExchange
@@ -156,7 +161,8 @@ testResult_t testStreamSynchronize(cudaStream_t stream)
   return testSuccess;
 }
 
-testResult_t completeColl(struct threadArgs* args) {
+testResult_t completeColl(struct threadArgs* args)
+{
   TESTCHECK(testStreamSynchronize(args->stream));
   return testSuccess;
 }
@@ -164,8 +170,7 @@ testResult_t completeColl(struct threadArgs* args) {
 // Inter-thread/process barrier+allreduce. The quality of the return value
 // for average=0 (which means broadcast from rank=0) is dubious. The returned
 // value will actually be the result of process-local broadcast from the local thread=0.
-template<typename T>
-void Allreduce(struct threadArgs* args, T* value, int average)
+template <typename T> void Allreduce(struct threadArgs* args, T* value, int average)
 {
   thread_local int epoch = 0;
   static pthread_mutex_t lock[2] = {PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER};
@@ -235,7 +240,8 @@ void Allreduce(struct threadArgs* args, T* value, int average)
   epoch ^= 1;
 }
 
-testResult_t BenchTime(struct threadArgs* args, int in_place) {
+testResult_t BenchTime(struct threadArgs* args, int in_place)
+{
   size_t count = args->nbytes;
 
   TESTCHECK(args->collTest->initData(args, in_place));
@@ -263,11 +269,11 @@ testResult_t BenchTime(struct threadArgs* args, int in_place) {
     CUDACHECK(cudaGraphLaunch(graphExec, args->stream));
   }
 
-  double cputimeSec = tim.elapsed()/(iters);
+  double cputimeSec = tim.elapsed() / (iters);
   TESTCHECK(completeColl(args));
 
   double deltaSec = tim.elapsed();
-  deltaSec = deltaSec/(iters)/(cudaGraphLaunches);
+  deltaSec = deltaSec / (iters) / (cudaGraphLaunches);
   Allreduce(args, &deltaSec, average);
 
   CUDACHECK(cudaGraphExecDestroy(graphExec));
@@ -278,7 +284,7 @@ testResult_t BenchTime(struct threadArgs* args, int in_place) {
 
   Barrier(args);
 
-  double timeUsec = (report_cputime ? cputimeSec : deltaSec)*1.0E6;
+  double timeUsec = (report_cputime ? cputimeSec : deltaSec) * 1.0E6;
   char timeStr[100];
   if (timeUsec >= 10000.0) {
     sprintf(timeStr, "%7.0f", timeUsec);
@@ -294,14 +300,16 @@ testResult_t BenchTime(struct threadArgs* args, int in_place) {
   return testSuccess;
 }
 
-void setupArgs(size_t size, struct threadArgs* args) {
-  int nranks = args->totalProcs*args->nGpus*args->nThreads;
+void setupArgs(size_t size, struct threadArgs* args)
+{
+  int nranks = args->totalProcs * args->nGpus * args->nThreads;
   size_t count, sendCount, recvCount, paramCount, sendInplaceOffset, recvInplaceOffset;
 
   // TODO: support more data types
   int typeSize = sizeof(char);
   count = size / typeSize;
-  args->collTest->getCollByteCount(&sendCount, &recvCount, &paramCount, &sendInplaceOffset, &recvInplaceOffset, (size_t)count, (size_t)nranks);
+  args->collTest->getCollByteCount(&sendCount, &recvCount, &paramCount, &sendInplaceOffset, &recvInplaceOffset,
+                                   (size_t)count, (size_t)nranks);
 
   args->nbytes = paramCount * typeSize;
   args->sendBytes = sendCount * typeSize;
@@ -310,7 +318,8 @@ void setupArgs(size_t size, struct threadArgs* args) {
   args->recvInplaceOffset = recvInplaceOffset * typeSize;
 }
 
-testResult_t TimeTest(struct threadArgs* args) {
+testResult_t TimeTest(struct threadArgs* args)
+{
   // Sync to avoid first-call timeout
   Barrier(args);
 
@@ -330,26 +339,26 @@ testResult_t TimeTest(struct threadArgs* args) {
   TESTCHECK(completeColl(args));
 
   PRINT("#\n");
-  PRINT("# %10s  %12s           in-place                       out-of-place          \n", "",
-        "");
+  PRINT("# %10s  %12s           in-place                       out-of-place          \n", "", "");
   PRINT("# %10s  %12s  %7s  %6s  %6s  %6s  %7s  %6s  %6s  %6s\n", "size", "count", "time", "algbw", "busbw", "#wrong",
         "time", "algbw", "busbw", "#wrong");
   PRINT("# %10s  %12s  %7s  %6s  %6s  %5s  %7s  %6s  %6s  %5s\n", "(B)", "(elements)", "(us)", "(GB/s)", "(GB/s)", "",
         "(us)", "(GB/s)", "(GB/s)", "");
   // Benchmark
-  for (size_t size = args->minbytes; size<=args->maxbytes; size = ((args->stepfactor > 1) ? size*args->stepfactor : size+args->stepbytes)) {
-      setupArgs(size, args);
-      PRINT("%12li  %12li", max(args->sendBytes, args->expectedBytes), args->nbytes);
-      // Don't support out-of-place for now
-      // TESTCHECK(BenchTime(args, 0));
-      TESTCHECK(BenchTime(args, 1));
-      PRINT("\n");
+  for (size_t size = args->minbytes; size <= args->maxbytes;
+       size = ((args->stepfactor > 1) ? size * args->stepfactor : size + args->stepbytes)) {
+    setupArgs(size, args);
+    PRINT("%12li  %12li", max(args->sendBytes, args->expectedBytes), args->nbytes);
+    // Don't support out-of-place for now
+    // TESTCHECK(BenchTime(args, 0));
+    TESTCHECK(BenchTime(args, 1));
+    PRINT("\n");
   }
   return testSuccess;
 }
 
 testResult_t setupMscclppConnections(int rank, int worldSize, int ranksPerNode, mscclppComm_t comm, void* dataDst,
-                                        size_t dataSize)
+                                     size_t dataSize)
 {
   int thisNode = rank / ranksPerNode;
   int localRank = rank % ranksPerNode;
@@ -378,8 +387,8 @@ testResult_t setupMscclppConnections(int rank, int worldSize, int ranksPerNode,
 testResult_t threadRunTests(struct threadArgs* args)
 {
   PRINT("# Setting up the connection in MSCCL++\n");
-  TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->nranksPerNode, args->comm,
-                                    args->recvbuffs[0], args->maxbytes));
+  TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->nranksPerNode, args->comm, args->recvbuffs[0],
+                                    args->maxbytes));
   PRINT("# Launching MSCCL++ proxy threads\n");
   MSCCLPPCHECK(mscclppProxyLaunch(args->comm));
   TESTCHECK(mscclppTestEngine.runTest(args));
@@ -390,103 +399,103 @@ testResult_t threadRunTests(struct threadArgs* args)
 
 testResult_t run(); // Main function
 
-int main(int argc, char* argv[]) {
+int main(int argc, char* argv[])
+{
   // Make sure everyline is flushed so that we see the progress of the test
   setlinebuf(stdout);
 
   // Parse args
   double parsed;
   int longindex;
-  static struct option longopts[] = {
-    {"minbytes", required_argument, 0, 'b'},
-    {"maxbytes", required_argument, 0, 'e'},
-    {"stepbytes", required_argument, 0, 'i'},
-    {"stepfactor", required_argument, 0, 'f'},
-    {"iters", required_argument, 0, 'n'},
-    {"warmup_iters", required_argument, 0, 'w'},
-    {"check", required_argument, 0, 'c'},
-    {"timeout", required_argument, 0, 'T'},
-    {"cudagraph", required_argument, 0, 'G'},
-    {"report_cputime", required_argument, 0, 'C'},
-    {"average", required_argument, 0, 'a'},
-    {"ip_port", required_argument, 0, 'P'},
-    {"help", no_argument, 0, 'h'},
-    {}
-  };
+  static struct option longopts[] = {{"minbytes", required_argument, 0, 'b'},
+                                     {"maxbytes", required_argument, 0, 'e'},
+                                     {"stepbytes", required_argument, 0, 'i'},
+                                     {"stepfactor", required_argument, 0, 'f'},
+                                     {"iters", required_argument, 0, 'n'},
+                                     {"warmup_iters", required_argument, 0, 'w'},
+                                     {"check", required_argument, 0, 'c'},
+                                     {"timeout", required_argument, 0, 'T'},
+                                     {"cudagraph", required_argument, 0, 'G'},
+                                     {"report_cputime", required_argument, 0, 'C'},
+                                     {"average", required_argument, 0, 'a'},
+                                     {"ip_port", required_argument, 0, 'P'},
+                                     {"help", no_argument, 0, 'h'},
+                                     {}};
 
-  while(1) {
+  while (1) {
     int c;
     c = getopt_long(argc, argv, "b:e:i:f:n:w:c:T:C:a:P:h:", longopts, &longindex);
 
     if (c == -1)
       break;
 
-    switch(c) {
-      case 'b':
-        parsed = parsesize(optarg);
-        if (parsed < 0) {
-          fprintf(stderr, "invalid size specified for 'minbytes'\n");
-          return -1;
-        }
-        minBytes = (size_t)parsed;
-        break;
-      case 'e':
-        parsed = parsesize(optarg);
-        if (parsed < 0) {
-          fprintf(stderr, "invalid size specified for 'maxbytes'\n");
-          return -1;
-        }
-        maxBytes = (size_t)parsed;
-        break;
-      case 'i':
-        stepBytes = strtol(optarg, NULL, 0);
-        break;
-      case 'f':
-        stepFactor = strtol(optarg, NULL, 0);
-        break;
-      case 'n':
-        iters = (int)strtol(optarg, NULL, 0);
-        break;
-      case 'w':
-        warmup_iters = (int)strtol(optarg, NULL, 0);
-        break;
-      case 'c':
-        datacheck = (int)strtol(optarg, NULL, 0);
-        break;
-      case 'T':
-        timeout = strtol(optarg, NULL, 0);
-        break;
-      case 'G':
-        cudaGraphLaunches = strtol(optarg, NULL, 0);
-        break;
-      case 'C':
-        report_cputime = strtol(optarg, NULL, 0);
-        break;
-      case 'a':
-        average = (int)strtol(optarg, NULL, 0);
-        break;
-      case 'P':
-        ip_port = optarg;
-        break;
-      case 'h':
-      default:
-        if (c != 'h') printf("invalid option '%c'\n", c);
-        printf("USAGE: %s \n\t"
-            "[-b,--minbytes <min size in bytes>] \n\t"
-            "[-e,--maxbytes <max size in bytes>] \n\t"
-            "[-i,--stepbytes <increment size>] \n\t"
-            "[-f,--stepfactor <increment factor>] \n\t"
-            "[-n,--iters <iteration count>] \n\t"
-            "[-w,--warmup_iters <warmup iteration count>] \n\t"
-            "[-c,--check <0/1>] \n\t"
-            "[-T,--timeout <time in seconds>] \n\t"
-            "[-G,--cudagraph <num graph launches>] \n\t"
-            "[-C,--report_cputime <0/1>] \n\t"
-            "[-a,--average <0/1/2/3> report average iteration time <0=RANK0/1=AVG/2=MIN/3=MAX>] \n\t"
-            "[-P,--ip_port <ip port for bootstrap>] \n\t"
-            "[-h,--help]\n",
-          basename(argv[0]));
-        return 0;
+    switch (c) {
+    case 'b':
+      parsed = parsesize(optarg);
+      if (parsed < 0) {
+        fprintf(stderr, "invalid size specified for 'minbytes'\n");
+        return -1;
+      }
+      minBytes = (size_t)parsed;
+      break;
+    case 'e':
+      parsed = parsesize(optarg);
+      if (parsed < 0) {
+        fprintf(stderr, "invalid size specified for 'maxbytes'\n");
+        return -1;
+      }
+      maxBytes = (size_t)parsed;
+      break;
+    case 'i':
+      stepBytes = strtol(optarg, NULL, 0);
+      break;
+    case 'f':
+      stepFactor = strtol(optarg, NULL, 0);
+      break;
+    case 'n':
+      iters = (int)strtol(optarg, NULL, 0);
+      break;
+    case 'w':
+      warmup_iters = (int)strtol(optarg, NULL, 0);
+      break;
+    case 'c':
+      datacheck = (int)strtol(optarg, NULL, 0);
+      break;
+    case 'T':
+      timeout = strtol(optarg, NULL, 0);
+      break;
+    case 'G':
+      cudaGraphLaunches = strtol(optarg, NULL, 0);
+      break;
+    case 'C':
+      report_cputime = strtol(optarg, NULL, 0);
+      break;
+    case 'a':
+      average = (int)strtol(optarg, NULL, 0);
+      break;
+    case 'P':
+      ip_port = optarg;
+      break;
+    case 'h':
+    default:
+      if (c != 'h')
+        printf("invalid option '%c'\n", c);
+      printf("USAGE: %s \n\t"
+             "[-b,--minbytes <min size in bytes>] \n\t"
+             "[-e,--maxbytes <max size in bytes>] \n\t"
+             "[-i,--stepbytes <increment size>] \n\t"
+             "[-f,--stepfactor <increment factor>] \n\t"
+             "[-n,--iters <iteration count>] \n\t"
+             "[-w,--warmup_iters <warmup iteration count>] \n\t"
+             "[-c,--check <0/1>] \n\t"
+             "[-T,--timeout <time in seconds>] \n\t"
+             "[-G,--cudagraph <num graph launches>] \n\t"
+             "[-C,--report_cputime <0/1>] \n\t"
+             "[-a,--average <0/1/2/3> report average iteration time <0=RANK0/1=AVG/2=MIN/3=MAX>] \n\t"
+             "[-P,--ip_port <ip port for bootstrap>] \n\t"
+             "[-h,--help]\n",
+             basename(argv[0]));
+      return 0;
     }
   }
   if (minBytes > maxBytes) {
@@ -505,7 +514,8 @@ int main(int argc, char* argv[]) {
   return 0;
 }
 
-testResult_t run() {
+testResult_t run()
+{
   int totalProcs = 1, proc = 0;
   int nranksPerNode = 0, localRank = 0;
   char hostname[1024];
@@ -550,7 +560,7 @@ testResult_t run() {
   MPI_Gather(line, MAX_LINE, MPI_BYTE, lines, MAX_LINE, MPI_BYTE, 0, MPI_COMM_WORLD);
   if (proc == 0) {
     for (int p = 0; p < totalProcs; p++)
-        PRINT("%s", lines + MAX_LINE * p);
+      PRINT("%s", lines + MAX_LINE * p);
     free(lines);
   }
   MPI_Allreduce(MPI_IN_PLACE, &maxMem, 1, MPI_LONG, MPI_MIN, MPI_COMM_WORLD);
@@ -558,10 +568,11 @@ testResult_t run() {
   PRINT("%s", line);
 #endif
   // We need sendbuff, recvbuff, expected (when datacheck enabled), plus 1G for the rest.
-  size_t memMaxBytes = (maxMem - (1<<30)) / (datacheck ? 3 : 2);
+  size_t memMaxBytes = (maxMem - (1 << 30)) / (datacheck ? 3 : 2);
   if (maxBytes > memMaxBytes) {
     maxBytes = memMaxBytes;
-    if (proc == 0) printf("#\n# Reducing maxBytes to %ld due to memory limitation\n", maxBytes);
+    if (proc == 0)
+      printf("#\n# Reducing maxBytes to %ld due to memory limitation\n", maxBytes);
   }
 
   int gpu = cudaDev;

From 4de3dba8182cd37dab6c4452f7f8a5ebe2e34ad9 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Mon, 3 Apr 2023 06:14:41 +0000
Subject: [PATCH 11/23] fix lint

---
 tests/common.h | 12 ++++++------
 tests/timer.cc | 22 +++++++++++++---------
 tests/timer.h  |  4 +++-
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/tests/common.h b/tests/common.h
index b7e2c226..1fffcbac 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -53,7 +53,8 @@
     }                                                                                                                  \
   } while (0)
 
-typedef enum {
+typedef enum
+{
   testSuccess = 0,
   testInternalError = 1,
   testCudaError = 2,
@@ -62,12 +63,11 @@ typedef enum {
   testNumResults = 5
 } testResult_t;
 
-struct testColl {
+struct testColl
+{
   const char name[20];
-  void (*getCollByteCount)(
-      size_t *sendcount, size_t *recvcount, size_t *paramcount,
-      size_t *sendInplaceOffset, size_t *recvInplaceOffset,
-      size_t count, int nranks);
+  void (*getCollByteCount)(size_t* sendcount, size_t* recvcount, size_t* paramcount, size_t* sendInplaceOffset,
+                           size_t* recvInplaceOffset, size_t count, int nranks);
   testResult_t (*initData)(struct threadArgs* args, int in_place);
   void (*getBw)(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks);
   testResult_t (*runColl)(void* sendbuff, void* recvbuff, int nranksPerNode, size_t count, mscclppComm_t comm,
diff --git a/tests/timer.cc b/tests/timer.cc
index eb510d9a..28477ae8 100644
--- a/tests/timer.cc
+++ b/tests/timer.cc
@@ -5,24 +5,28 @@
 #include <chrono>
 
 namespace {
-  std::uint64_t now() {
-    using clock = std::chrono::steady_clock;
-    return std::chrono::duration_cast<std::chrono::nanoseconds>(clock::now().time_since_epoch()).count();
-  }
+std::uint64_t now()
+{
+  using clock = std::chrono::steady_clock;
+  return std::chrono::duration_cast<std::chrono::nanoseconds>(clock::now().time_since_epoch()).count();
 }
+} // namespace
 
-timer::timer() {
+timer::timer()
+{
   t0 = now();
 }
 
-double timer::elapsed() const {
+double timer::elapsed() const
+{
   std::uint64_t t1 = now();
-  return 1.e-9*(t1 - t0);
+  return 1.e-9 * (t1 - t0);
 }
 
-double timer::reset() {
+double timer::reset()
+{
   std::uint64_t t1 = now();
-  double ans = 1.e-9*(t1 - t0);
+  double ans = 1.e-9 * (t1 - t0);
   t0 = t1;
   return ans;
 }
\ No newline at end of file
diff --git a/tests/timer.h b/tests/timer.h
index 6edfdd4c..c30b264c 100644
--- a/tests/timer.h
+++ b/tests/timer.h
@@ -4,8 +4,10 @@
 #include <cstdint>
 
 // Can't include <chrono> because of bug with gcc 10.3.0
-class timer {
+class timer
+{
   std::uint64_t t0;
+
 public:
   timer();
   double elapsed() const;

From 7d6dad226b7e2936a32ed44af3608b0b9be2844d Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Mon, 3 Apr 2023 09:02:58 +0000
Subject: [PATCH 12/23] fix

---
 tests/common.cu | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/common.cu b/tests/common.cu
index 58381a78..d7ee6314 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -4,8 +4,6 @@
  * See LICENSE.txt for license information
  ************************************************************************/
 
-#define MSCCLPP_USE_MPI_FOR_TESTS 1
-
 #include "common.h"
 #include "cuda.h"
 #include "mscclpp.h"

From 69a49c189f078c880a3ea599b2da02d924196ff5 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 4 Apr 2023 07:01:21 +0000
Subject: [PATCH 13/23] Add correctness check

---
 Makefile                 |  2 +-
 tests/allgather_test3.cu | 19 +++++++-----
 tests/common.cu          | 63 +++++++++++++++++++++++++++++++++++++---
 3 files changed, 71 insertions(+), 13 deletions(-)

diff --git a/Makefile b/Makefile
index a46c691a..6349bad8 100644
--- a/Makefile
+++ b/Makefile
@@ -150,7 +150,7 @@ lib: $(LIBOBJTARGETS) $(INCTARGETS) $(LIBTARGET)
 
 tests: $(TESTSBINS)
 
-mscclpp_test: $(LIBTARGET) $(MSCLLPPTESTBINS)
+mscclpp-test: $(LIBTARGET) $(MSCLLPPTESTBINS)
 
 cpplint:
 	clang-format-12 -style=file --verbose --Werror --dry-run $(CPPSOURCES)
diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
index 3c72209d..1f16fba9 100644
--- a/tests/allgather_test3.cu
+++ b/tests/allgather_test3.cu
@@ -19,16 +19,16 @@ __device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, i
 
   // this thread's role is a sender role
   // put your data asynchronously
-  if (threadIdx.x % 32 != 0)
+  if (threadIdx.x % 32 == 0)
     devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
   // make sure everyone is put their data before some thread randomly blocks everyone else in signal
   __syncthreads();
   // push with flag and sync to make sure the data is received
-  if (threadIdx.x % 32 != 0)
+  if (threadIdx.x % 32 == 0)
     devConn.flush();
 
   // this thread's role is a receiver role. wait on the semaphore to make sure the data is ready
-  if (threadIdx.x % 32 != 0)
+  if (threadIdx.x % 32 == 0)
     devConn.wait();
 }
 
@@ -155,8 +155,8 @@ void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* par
 
 testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
 {
-  // size_t sendcount = args->sendBytes;
-  size_t recvcount = args->expectedBytes;
+  size_t sendcount = args->sendBytes / sizeof(int);
+  size_t recvcount = args->expectedBytes / sizeof(int);
   // int nranks = args->totalProcs;
 
   CUDACHECK(cudaSetDevice(args->gpus[0]));
@@ -165,17 +165,20 @@ testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
   // void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
 
   int* dataHost = new int[recvcount];
-  for (int i = 0; i < static_cast<int>(recvcount); i++) {
+  for (size_t i = 0; i < recvcount; i++) {
     int val = i + 1;
-    if (i / args->nranksPerNode == rank) {
+    if (i / sendcount == (size_t)rank) {
       dataHost[i] = val;
     } else {
       dataHost[i] = 0;
     }
   }
   CUDACHECK(cudaMemcpy(args->recvbuffs[0], dataHost, recvcount, cudaMemcpyHostToDevice));
+  for (int i = 0; i < static_cast<int>(recvcount); i++) {
+    dataHost[i] = i + 1;
+  }
+  CUDACHECK(cudaMemcpy(args->expected[0], dataHost, recvcount, cudaMemcpyHostToDevice));
   delete dataHost;
-  // TODO: need to init expected data here
   CUDACHECK(cudaDeviceSynchronize());
   return testSuccess;
 }
diff --git a/tests/common.cu b/tests/common.cu
index d7ee6314..9fc45545 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -238,6 +238,29 @@ template <typename T> void Allreduce(struct threadArgs* args, T* value, int aver
   epoch ^= 1;
 }
 
+testResult_t CheckData(struct threadArgs* args, int in_place, int64_t *wrongElts) {
+  if (in_place == 0) {
+    return testInternalError;
+  }
+  size_t count = args->expectedBytes / sizeof(int);
+
+  int* dataHostRecv = new int[count];
+  int* dataHostExpected = new int[count];
+  CUDACHECK(cudaMemcpy(dataHostRecv, args->recvbuffs[0], args->expectedBytes, cudaMemcpyDeviceToHost));
+  CUDACHECK(cudaMemcpy(dataHostExpected, args->expected[0], args->expectedBytes, cudaMemcpyDeviceToHost));
+
+  for (size_t i = 0; i < count; i++) {
+    if (dataHostRecv[i] != dataHostExpected[i]) {
+      *wrongElts += 1;
+    }
+  }
+
+  if (args->reportErrors && *wrongElts) {
+    (*args->errors)++;
+  }
+  return testSuccess;
+}
+
 testResult_t BenchTime(struct threadArgs* args, int in_place)
 {
   size_t count = args->nbytes;
@@ -279,9 +302,37 @@ testResult_t BenchTime(struct threadArgs* args, int in_place)
 
   double algBw, busBw;
   args->collTest->getBw(count, 1, deltaSec, &algBw, &busBw, args->totalProcs * args->nThreads * args->nGpus);
-
   Barrier(args);
 
+  int64_t wrongElts = 0;
+  if (datacheck) {
+    // Initialize sendbuffs, recvbuffs and expected
+    TESTCHECK(args->collTest->initData(args, in_place));
+    // Begin cuda graph capture for data check
+    CUDACHECK(cudaStreamBeginCapture(args->stream, cudaStreamCaptureModeGlobal));
+    // test validation in single itertion, should ideally be included into the multi-iteration run
+    TESTCHECK(startColl(args, in_place, 0));
+    // End cuda graph capture
+    CUDACHECK(cudaStreamEndCapture(args->stream, &graph));
+    // Instantiate cuda graph
+    CUDACHECK(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+    // Launch cuda graph
+    CUDACHECK(cudaGraphLaunch(graphExec, args->stream));
+
+    TESTCHECK(completeColl(args));
+
+    // destroy cuda graph
+    CUDACHECK(cudaGraphExecDestroy(graphExec));
+    CUDACHECK(cudaGraphDestroy(graph));
+
+    TESTCHECK(CheckData(args, in_place, &wrongElts));
+
+    // aggregate delta from all threads and procs
+    long long wrongElts1 = wrongElts;
+    Allreduce(args, &wrongElts1, /*sum*/ 4);
+    wrongElts = wrongElts1;
+  }
+
   double timeUsec = (report_cputime ? cputimeSec : deltaSec) * 1.0E6;
   char timeStr[100];
   if (timeUsec >= 10000.0) {
@@ -291,7 +342,11 @@ testResult_t BenchTime(struct threadArgs* args, int in_place)
   } else {
     sprintf(timeStr, "%7.2f", timeUsec);
   }
-  PRINT("  %7s  %6.2f  %6.2f  %5s", timeStr, algBw, busBw, "N/A");
+  if (args->reportErrors) {
+    PRINT("  %7s  %6.2f  %6.2f  %5g", timeStr, algBw, busBw, (double)wrongElts);
+  } else {
+    PRINT("  %7s  %6.2f  %6.2f  %5s", timeStr, algBw, busBw, "N/A");
+  }
 
   args->bw[0] += busBw;
   args->bw_count[0]++;
@@ -304,7 +359,7 @@ void setupArgs(size_t size, struct threadArgs* args)
   size_t count, sendCount, recvCount, paramCount, sendInplaceOffset, recvInplaceOffset;
 
   // TODO: support more data types
-  int typeSize = sizeof(char);
+  int typeSize = sizeof(int);
   count = size / typeSize;
   args->collTest->getCollByteCount(&sendCount, &recvCount, &paramCount, &sendInplaceOffset, &recvInplaceOffset,
                                    (size_t)count, (size_t)nranks);
@@ -346,7 +401,7 @@ testResult_t TimeTest(struct threadArgs* args)
   for (size_t size = args->minbytes; size <= args->maxbytes;
        size = ((args->stepfactor > 1) ? size * args->stepfactor : size + args->stepbytes)) {
     setupArgs(size, args);
-    PRINT("%12li  %12li", max(args->sendBytes, args->expectedBytes), args->nbytes);
+    PRINT("%12li  %12li", max(args->sendBytes, args->expectedBytes), args->nbytes / sizeof(int));
     // Don't support out-of-place for now
     // TESTCHECK(BenchTime(args, 0));
     TESTCHECK(BenchTime(args, 1));

From 5648ac5d7dd797e4b1069533b8b14a79230d179e Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 4 Apr 2023 07:04:12 +0000
Subject: [PATCH 14/23] fix lint

---
 tests/common.cu | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/common.cu b/tests/common.cu
index 9fc45545..b1f98388 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -238,7 +238,8 @@ template <typename T> void Allreduce(struct threadArgs* args, T* value, int aver
   epoch ^= 1;
 }
 
-testResult_t CheckData(struct threadArgs* args, int in_place, int64_t *wrongElts) {
+testResult_t CheckData(struct threadArgs* args, int in_place, int64_t* wrongElts)
+{
   if (in_place == 0) {
     return testInternalError;
   }

From 674d30a813d87712326214d2c8b5a7fcc5d6d4b6 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 4 Apr 2023 08:37:16 +0000
Subject: [PATCH 15/23] minor fix

---
 Makefile        | 2 +-
 tests/common.cu | 6 +++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/Makefile b/Makefile
index 6349bad8..3fdb3ba3 100644
--- a/Makefile
+++ b/Makefile
@@ -144,7 +144,7 @@ INCLUDE := -Isrc -Isrc/include
 
 all: build
 
-build: lib tests mscclpp_test
+build: lib tests mscclpp-test
 
 lib: $(LIBOBJTARGETS) $(INCTARGETS) $(LIBTARGET)
 
diff --git a/tests/common.cu b/tests/common.cu
index b1f98388..3e11cb44 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -478,7 +478,7 @@ int main(int argc, char* argv[])
 
   while (1) {
     int c;
-    c = getopt_long(argc, argv, "b:e:i:f:n:w:c:T:C:a:P:h:", longopts, &longindex);
+    c = getopt_long(argc, argv, "b:e:i:f:n:w:c:T:G:C:a:P:h:", longopts, &longindex);
 
     if (c == -1)
       break;
@@ -520,6 +520,10 @@ int main(int argc, char* argv[])
       break;
     case 'G':
       cudaGraphLaunches = strtol(optarg, NULL, 0);
+      if (cudaGraphLaunches <= 0) {
+        fprintf(stderr, "invalid number for 'cudaGraphLaunches'\n");
+        return -1;
+      }
       break;
     case 'C':
       report_cputime = strtol(optarg, NULL, 0);

From e82a75c132c093e90d9f12cec94b9c0415e3c7ba Mon Sep 17 00:00:00 2001
From: Saeed Maleki <saemal@microsoft.com>
Date: Thu, 6 Apr 2023 21:58:48 +0000
Subject: [PATCH 16/23] typo fix

---
 Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index 3fdb3ba3..6c3577a1 100644
--- a/Makefile
+++ b/Makefile
@@ -192,7 +192,7 @@ $(BUILDDIR)/$(BINDIR)/%: $(BUILDDIR)/$(OBJDIR)/%.o $(LIBTARGET)
 	@mkdir -p $(@D)
 	$(NVCC) -o $@ $< $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp
 
-# Comile mscclpp_test
+# Compile mscclpp_test
 $(BUILDDIR)/$(BINDIR)/$(TESTSDIR)/%_perf: $(MSCLLPPTESTSOBJSDIR)/%.o $(MSCLLPPTESTSOBJSDIR)/common.o $(MSCLLPPTESTSOBJSDIR)/timer.o
 	@mkdir -p $(@D)
 	$(NVCC) -o $@ $^ $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp

From 1b2db68e93db9433275920fffe8aefa5bc84d859 Mon Sep 17 00:00:00 2001
From: Saeed Maleki <saemal@microsoft.com>
Date: Thu, 6 Apr 2023 23:55:11 +0000
Subject: [PATCH 17/23] filename changes

---
 Makefile                           |   4 +-
 tests/allgather_test.cu            | 420 +++++-------------------
 tests/allgather_test2.cu           | 295 -----------------
 tests/allgather_test3.cu           | 225 -------------
 tests/allgather_test_standalone.cu | 501 +++++++++++++++++++++++++++++
 5 files changed, 575 insertions(+), 870 deletions(-)
 delete mode 100644 tests/allgather_test2.cu
 delete mode 100644 tests/allgather_test3.cu
 create mode 100644 tests/allgather_test_standalone.cu

diff --git a/Makefile b/Makefile
index 6c3577a1..692ae56b 100644
--- a/Makefile
+++ b/Makefile
@@ -129,13 +129,13 @@ LIBSONAME := $(LIBNAME).$(MSCCLPP_MAJOR)
 LIBTARGET := $(BUILDDIR)/$(LIBDIR)/$(LIBNAME).$(MSCCLPP_MAJOR).$(MSCCLPP_MINOR).$(MSCCLPP_PATCH)
 
 TESTSDIR  := tests
-TESTSSRCS := $(addprefix $(TESTSDIR)/,bootstrap_test.cc allgather_test.cu)
+TESTSSRCS := $(addprefix $(TESTSDIR)/,bootstrap_test.cc allgather_test_standalone.cu)
 TESTSOBJS := $(patsubst %.cc,%.o,$(TESTSSRCS)) $(patsubst %.cu,%.o,$(TESTSSRCS))
 TESTSOBJTARGETS := $(TESTSOBJS:%=$(BUILDDIR)/$(OBJDIR)/%)
 TESTSBINS       := $(patsubst %.o,$(BUILDDIR)/$(BINDIR)/%,$(TESTSOBJS))
 
 MSCLLPPTESTSOBJSDIR:= $(BUILDDIR)/$(OBJDIR)/$(TESTSDIR)
-MSCLLPPTESTBINFILESLIST := allgather_test3
+MSCLLPPTESTBINFILESLIST := allgather_test
 MSCLLPPTESTBINS         := $(MSCLLPPTESTBINFILESLIST:%=$(BUILDDIR)/$(BINDIR)/$(TESTSDIR)/%_perf)
 
 INCLUDE := -Isrc -Isrc/include
diff --git a/tests/allgather_test.cu b/tests/allgather_test.cu
index 05aa3304..1f16fba9 100644
--- a/tests/allgather_test.cu
+++ b/tests/allgather_test.cu
@@ -1,50 +1,16 @@
-#include "mscclpp.h"
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
 
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-#include "mpi.h"
-#endif // MSCCLPP_USE_MPI_FOR_TESTS
-#include <iostream>
-#include <stdio.h>
-#include <stdlib.h>
+#include "comm.h"
+#include "common.h"
+
+#include <cuda_runtime.h>
 #include <string>
-#include <unistd.h>
-#include <unordered_map>
-
-static int nranksPerNode = 8;
-
-// Propagate errors up
-
-#define MSCCLPPCHECK(call)                                                                                             \
-  do {                                                                                                                 \
-    mscclppResult_t res = call;                                                                                        \
-    if (res != mscclppSuccess && res != mscclppInProgress) {                                                           \
-      /* Print the back trace*/                                                                                        \
-      printf("Failure at %s:%d -> %s\n", __FILE__, __LINE__, mscclppGetErrorString(res));                              \
-      return res;                                                                                                      \
-    }                                                                                                                  \
-  } while (0)
-
-// Check CUDA RT calls
-#define CUDACHECK(cmd)                                                                                                 \
-  do {                                                                                                                 \
-    cudaError_t err = cmd;                                                                                             \
-    if (err != cudaSuccess) {                                                                                          \
-      printf("%s:%d Cuda failure '%s'\n", __FILE__, __LINE__, cudaGetErrorString(err));                                \
-      exit(EXIT_FAILURE);                                                                                              \
-    }                                                                                                                  \
-  } while (false)
-
-// Measure current time in second.
-static double getTime(void)
-{
-  struct timespec tspec;
-  if (clock_gettime(CLOCK_MONOTONIC, &tspec) == -1) {
-    printf("clock_gettime failed\n");
-    exit(EXIT_FAILURE);
-  }
-  return (tspec.tv_nsec / 1.0e9) + tspec.tv_sec;
-}
 
+#define ALIGN 4
 __constant__ mscclppDevConn_t constDevConns[16];
 
 __device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, size_t nelemsPerGPU)
@@ -53,16 +19,16 @@ __device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, i
 
   // this thread's role is a sender role
   // put your data asynchronously
-  if ((threadIdx.x % 32) == 0)
+  if (threadIdx.x % 32 == 0)
     devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
   // make sure everyone is put their data before some thread randomly blocks everyone else in signal
   __syncthreads();
   // push with flag and sync to make sure the data is received
-  if ((threadIdx.x % 32) == 0)
+  if (threadIdx.x % 32 == 0)
     devConn.flush();
 
   // this thread's role is a receiver role. wait on the semaphore to make sure the data is ready
-  if ((threadIdx.x % 32) == 0)
+  if (threadIdx.x % 32 == 0)
     devConn.wait();
 }
 
@@ -176,326 +142,84 @@ __global__ void kernel(int rank, int world_size, int nranksPerNode, size_t nelem
     allgather2(devConn, rank, world_size, nranksPerNode, remoteRank, nelemsPerGPU);
 }
 
-int rankToLocalRank(int rank)
+void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* paramcount, size_t* sendInplaceOffset,
+                               size_t* recvInplaceOffset, size_t count, int nranks)
 {
-  return rank % nranksPerNode;
+  size_t base = (count / (ALIGN * nranks)) * ALIGN;
+  *sendcount = base;
+  *recvcount = base * nranks;
+  *sendInplaceOffset = base;
+  *recvInplaceOffset = 0;
+  *paramcount = base;
 }
 
-int rankToNode(int rank)
+testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
 {
-  return rank / nranksPerNode;
-}
+  size_t sendcount = args->sendBytes / sizeof(int);
+  size_t recvcount = args->expectedBytes / sizeof(int);
+  // int nranks = args->totalProcs;
 
-void print_usage(const char* prog)
-{
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  printf("usage: %s IP:PORT [rank nranks]\n", prog);
-#else
-  printf("usage: %s IP:PORT rank nranks\n", prog);
-#endif
-}
+  CUDACHECK(cudaSetDevice(args->gpus[0]));
+  int rank = args->proc;
+  CUDACHECK(cudaMemset(args->recvbuffs[0], 0, args->expectedBytes));
+  // void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
 
-void initializeAndAllocateAllGatherData(int rank, int world_size, size_t dataSize, size_t nelemsPerGPU, int** data_h,
-                                        int** data_d)
-{
-  CUDACHECK(cudaMalloc(data_d, dataSize));
-  CUDACHECK(cudaMemset(*data_d, 0, dataSize));
-
-  *data_h = new int[nelemsPerGPU * world_size];
-  for (size_t i = 0; i < nelemsPerGPU * world_size; i++) {
+  int* dataHost = new int[recvcount];
+  for (size_t i = 0; i < recvcount; i++) {
     int val = i + 1;
-    if (i / nelemsPerGPU == (size_t)rank) {
-      (*data_h)[i] = val;
+    if (i / sendcount == (size_t)rank) {
+      dataHost[i] = val;
     } else {
-      (*data_h)[i] = 0;
+      dataHost[i] = 0;
     }
   }
-  CUDACHECK(cudaMemcpy(*data_d, *data_h, dataSize, cudaMemcpyHostToDevice));
+  CUDACHECK(cudaMemcpy(args->recvbuffs[0], dataHost, recvcount, cudaMemcpyHostToDevice));
+  for (int i = 0; i < static_cast<int>(recvcount); i++) {
+    dataHost[i] = i + 1;
+  }
+  CUDACHECK(cudaMemcpy(args->expected[0], dataHost, recvcount, cudaMemcpyHostToDevice));
+  delete dataHost;
+  CUDACHECK(cudaDeviceSynchronize());
+  return testSuccess;
 }
 
-mscclppResult_t setupMscclppConnections(int rank, int world_size, mscclppComm_t comm, int* data_d, size_t dataSize)
+void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks)
 {
-  int thisNode = rankToNode(rank);
-  int cudaNum = rankToLocalRank(rank);
-  std::string ibDevStr = "mlx5_ib" + std::to_string(cudaNum);
+  double baseBw = (double)(count * typesize * nranks) / 1.0E9 / sec;
 
-  for (int r = 0; r < world_size; ++r) {
-    if (r == rank)
-      continue;
-    mscclppTransport_t transportType;
-    const char* ibDev = ibDevStr.c_str();
-    if (rankToNode(r) == thisNode) {
-      ibDev = NULL;
-      transportType = mscclppTransportP2P;
-    } else {
-      transportType = mscclppTransportIB;
-    }
-    // Connect with all other ranks
-    MSCCLPPCHECK(mscclppConnect(comm, r, 0, data_d, dataSize, transportType, ibDev));
-  }
+  *algBw = baseBw;
+  double factor = ((double)(nranks - 1)) / ((double)nranks);
+  *busBw = baseBw * factor;
+}
 
-  MSCCLPPCHECK(mscclppConnectionSetup(comm));
+testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, int nranksPerNode, size_t count, mscclppComm_t comm,
+                              cudaStream_t stream)
+{
+  int worldSize = comm->nRanks;
+  kernel<<<1, 32 * (worldSize - 1), 0, stream>>>(comm->rank, worldSize, nranksPerNode, count / sizeof(int), 1);
+  return testSuccess;
+}
 
+struct testColl allGatherTest = {"AllGather", AllGatherGetCollByteCount, AllGatherInitData, AllGatherGetBw,
+                                 AllGatherRunColl};
+
+void AllGatherGetBuffSize(size_t* sendcount, size_t* recvcount, size_t count, int nranks)
+{
+  size_t paramcount, sendInplaceOffset, recvInplaceOffset;
+  AllGatherGetCollByteCount(sendcount, recvcount, &paramcount, &sendInplaceOffset, &recvInplaceOffset, count, nranks);
+}
+
+testResult_t AllGatherRunTest(struct threadArgs* args)
+{
+  args->collTest = &allGatherTest;
   mscclppDevConn_t* devConns;
   int nCons;
-  MSCCLPPCHECK(mscclppGetAllDeviceConnections(comm, &devConns, &nCons));
-
+  MSCCLPPCHECK(mscclppGetAllDeviceConnections(args->comm, &devConns, &nCons));
   CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * nCons));
-
-  return mscclppSuccess;
+  TESTCHECK(TimeTest(args));
+  return testSuccess;
 }
 
-void printUsage(const char* prog, bool isMpi)
-{
-  if (isMpi) {
-    std::string st = "you are using MPI for this test\n";
-    st += "two possilbe usages are:\n";
-    st += "> " + std::string(prog) + "\n";
-    st += "or\n";
-    st += "> " + std::string(prog) + " -ip_port [ip:port]\n";
-    printf("%s", st.c_str());
-  } else {
-    std::string st = "you are NOT using MPI for this test\n";
-    st += "the only possible usage:\n";
-    st += "> " + std::string(prog) + " -ip_port [ip:port] -rank [rank] -nranks [nranks]\n";
-    printf("%s", st.c_str());
-  }
-}
+struct testEngine allGatherEngine = {AllGatherGetBuffSize, AllGatherRunTest};
 
-std::unordered_map<std::string, std::string> parseArgs(int argc, const char* argv[], bool isMpi)
-{
-  std::unordered_map<std::string, std::string> options;
-
-  for (int i = 1; i < argc; i++) {
-    std::string arg = argv[i];
-    if (arg == "-rankspernode") {
-      if (isMpi) {
-        fprintf(stderr, "Error: -rankspernode should not be specified with MPI.\n");
-        exit(-1);
-      }
-      if (i + 1 < argc) {
-        options["rankspernode"] = argv[++i];
-      } else {
-        fprintf(stderr, "Error: -rankspernode option requires an argument.\n");
-        ;
-        exit(-1);
-      }
-    } else if (arg == "-kernel") {
-      if (i + 1 < argc) {
-        options["kernel"] = argv[++i];
-      } else {
-        fprintf(stderr, "Error: -kernel option requires an argument.\n");
-        exit(-1);
-      }
-    } else if (arg == "-ip_port") {
-      if (i + 1 < argc) {
-        options["ip_port"] = argv[++i];
-      } else {
-        fprintf(stderr, "Error: -ip_port option requires an argument.\n");
-        exit(-1);
-      }
-    } else if (arg == "-rank") {
-      if (isMpi) {
-        fprintf(stderr, "Error: -rank should not be specified with MPI.\n");
-        exit(-1);
-      }
-      if (i + 1 < argc) {
-        options["rank"] = argv[++i];
-      } else {
-        fprintf(stderr, "Error: -ip_port option requires an argument.\n");
-        exit(-1);
-      }
-    } else if (arg == "-nranks") {
-      if (isMpi) {
-        fprintf(stderr, "Error: -nranks should not be specified with MPI.\n");
-        exit(-1);
-      }
-      if (i + 1 < argc) {
-        options["nranks"] = argv[++i];
-      } else {
-        fprintf(stderr, "Error: -ip_port option requires an argument.\n");
-        exit(-1);
-      }
-    } else if (arg == "-datasize") {
-      if (i + 1 < argc) {
-        options["datasize"] = argv[++i];
-      } else {
-        fprintf(stderr, "Error: -datasize option requires an argument.\n");
-        exit(-1);
-      }
-    } else if (arg == "-help" || arg == "-h") {
-      printUsage(argv[0], isMpi);
-      exit(0);
-    } else {
-      fprintf(stderr, "Error: Unknown option %s\n", argv[i]);
-      exit(-1);
-    }
-  }
-  return options;
-}
-
-int main(int argc, const char* argv[])
-{
-  bool isMpi = false;
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  isMpi = true;
-#endif
-
-  auto parsedArgs = parseArgs(argc, argv, isMpi);
-
-  int rank;
-  int world_size;
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  MPI_Init(NULL, NULL);
-  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-  MPI_Comm_size(MPI_COMM_WORLD, &world_size);
-  // get the local number of nodes with MPI
-  MPI_Comm shmcomm;
-  MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &shmcomm);
-  int shmrank;
-  MPI_Comm_size(shmcomm, &shmrank);
-  nranksPerNode = shmrank;
-  MPI_Comm_free(&shmcomm);
-#else
-  if (parsedArgs.find("rank") == parsedArgs.end() || parsedArgs.find("nranks") == parsedArgs.end()) {
-    printUsage(argv[0], isMpi);
-    exit(-1);
-  }
-  rank = std::stoi(parsedArgs["rank"]);
-  world_size = std::stoi(parsedArgs["nranks"]);
-  if (parsedArgs.find("rankspernode") == parsedArgs.end()) {
-    printUsage(argv[0], isMpi);
-    exit(-1);
-  }
-  nranksPerNode = std::stoi(parsedArgs["rankspernode"]);
-#endif
-  int kernelNum = 0;
-  if (parsedArgs.find("kernel") != parsedArgs.end()) {
-    kernelNum = std::stoi(parsedArgs["kernel"]);
-  }
-  char* ip_port = NULL;
-  if (parsedArgs.find("ip_port") == parsedArgs.end()) {
-    printUsage(argv[0], isMpi);
-    exit(-1);
-  }
-  ip_port = (char*)parsedArgs["ip_port"].c_str();
-
-  int thisNode = rankToNode(rank);
-  int cudaNum = rankToLocalRank(rank);
-  CUDACHECK(cudaSetDevice(cudaNum));
-
-  if (rank == 0)
-    printf("Initializing MSCCL++\n");
-  mscclppComm_t comm;
-  MSCCLPPCHECK(mscclppCommInitRank(&comm, world_size, ip_port, rank));
-
-  int* data_d;
-  int* data_h;
-  size_t dataSize = 1024 * 1024 * 1024;
-  if (parsedArgs.find("datasize") != parsedArgs.end()) {
-    dataSize = std::stoul(parsedArgs["datasize"]);
-  }
-  size_t nelemsPerGPU = dataSize / sizeof(int) / world_size;
-
-  if (rank == 0)
-    printf("Initializing data for allgather test\n");
-  initializeAndAllocateAllGatherData(rank, world_size, dataSize, nelemsPerGPU, &data_h, &data_d);
-
-  if (rank == 0)
-    printf("Setting up the connection in MSCCL++\n");
-  MSCCLPPCHECK(setupMscclppConnections(rank, world_size, comm, data_d, dataSize));
-
-  if (rank == 0)
-    printf("Launching MSCCL++ proxy threads\n");
-  MSCCLPPCHECK(mscclppProxyLaunch(comm));
-
-  if (rank == 0)
-    printf("Testing the correctness of AllGather implementation\n");
-  cudaStream_t stream;
-  CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
-  CUDACHECK(cudaDeviceSynchronize());
-  kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size, nranksPerNode, nelemsPerGPU, kernelNum);
-  CUDACHECK(cudaDeviceSynchronize());
-  CUDACHECK(cudaMemcpy(data_h, data_d, dataSize, cudaMemcpyDeviceToHost));
-
-  for (size_t i = 0; i < nelemsPerGPU * world_size; i++) {
-    int val = i + 1;
-    if (data_h[i] != val) {
-      printf("oh uh! data_h[%ld] (%d) != val (%d)\n", i, data_h[i], val);
-      break;
-    }
-  }
-  int tmp[16];
-  // A simple barrier
-  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-  if (rank == 0)
-    printf("Successfully checked the correctness\n");
-
-  // Perf test
-  int iterwithoutcudagraph = 10;
-  if (rank == 0)
-    printf("Running %d iterations of the kernel without CUDA graph\n", iterwithoutcudagraph);
-  CUDACHECK(cudaStreamSynchronize(stream));
-  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-  for (int i = 0; i < iterwithoutcudagraph; ++i) {
-    kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size, nranksPerNode, nelemsPerGPU, kernelNum);
-  }
-  CUDACHECK(cudaStreamSynchronize(stream));
-  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-
-  // cudaGraph Capture
-  int cudagraphiter = 10;
-  if (rank == 0)
-    printf("Capturing %d iterations of the kernel in a CUDA graph\n", cudagraphiter);
-  cudaGraph_t graph;
-  cudaGraphExec_t instance;
-  cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
-  for (int i = 0; i < cudagraphiter; ++i) {
-    kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size, nranksPerNode, nelemsPerGPU, kernelNum);
-  }
-  cudaStreamEndCapture(stream, &graph);
-  cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);
-
-  int cudagraphwarmup = 10;
-  if (rank == 0)
-    printf("Warming up %d iterations of the CUDA graph with %d iterations of the kernel\n", cudagraphwarmup,
-           cudagraphiter);
-  for (int i = 0; i < cudagraphwarmup; ++i) {
-    cudaGraphLaunch(instance, stream);
-  }
-  CUDACHECK(cudaStreamSynchronize(stream));
-
-  // measure runtime
-  int cudagraphlaunch = 10;
-  if (rank == 0)
-    printf("Running %d iterations of the CUDA graph with %d iterations of the kernel\n", cudagraphlaunch,
-           cudagraphiter);
-  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-  double t0, t1, ms, time_in_us;
-  t0 = getTime();
-  for (int i = 0; i < cudagraphlaunch; ++i) {
-    cudaGraphLaunch(instance, stream);
-  }
-  CUDACHECK(cudaStreamSynchronize(stream));
-
-  t1 = getTime();
-  ms = (t1 - t0) * 1000.0;
-  time_in_us = ms * 1000. / (float)cudagraphlaunch / (float)cudagraphiter;
-  printf("Rank %d report: size %lu time: %f us/iter algBW %f GBps\n", rank, dataSize, time_in_us,
-         (double)(dataSize) / 1e9 / (time_in_us / 1e6));
-  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-
-  if (rank == 0)
-    printf("Stopping MSCCL++ proxy threads\n");
-  MSCCLPPCHECK(mscclppProxyStop(comm));
-
-  if (rank == 0)
-    printf("Destroying MSCCL++ communicator\n");
-  MSCCLPPCHECK(mscclppCommDestroy(comm));
-  printf("Rank %d succeeded!\n", rank);
-
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  MPI_Finalize();
-#endif
-  return 0;
-}
+#pragma weak mscclppTestEngine = allGatherEngine
diff --git a/tests/allgather_test2.cu b/tests/allgather_test2.cu
deleted file mode 100644
index 047869fb..00000000
--- a/tests/allgather_test2.cu
+++ /dev/null
@@ -1,295 +0,0 @@
-#include "mscclpp.h"
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-#include "mpi.h"
-#endif // MSCCLPP_USE_MPI_FOR_TESTS
-#include <stdio.h>
-#include <stdlib.h>
-#include <string>
-#include <unistd.h>
-
-#define RANKS_PER_NODE 8
-
-// Check CUDA RT calls
-#define CUDACHECK(cmd)                                                                                                 \
-  do {                                                                                                                 \
-    cudaError_t err = cmd;                                                                                             \
-    if (err != cudaSuccess) {                                                                                          \
-      printf("%s:%d Cuda failure '%s'\n", __FILE__, __LINE__, cudaGetErrorString(err));                                \
-      exit(EXIT_FAILURE);                                                                                              \
-    }                                                                                                                  \
-  } while (false)
-
-// Measure current time in second.
-static double getTime(void)
-{
-  struct timespec tspec;
-  if (clock_gettime(CLOCK_MONOTONIC, &tspec) == -1) {
-    printf("clock_gettime failed\n");
-    exit(EXIT_FAILURE);
-  }
-  return (tspec.tv_nsec / 1.0e9) + tspec.tv_sec;
-}
-
-__constant__ mscclppDevConn_t constDevConns[16];
-
-__global__ void kernel(int rank, int world_size, size_t nelemsPerGPU)
-{
-  if (threadIdx.x % 32 != 0)
-    return;
-
-  int warpId = threadIdx.x / 32;
-  bool isIB = false;
-  if (warpId >= world_size - 1)
-    isIB = true;
-  if (isIB)
-    warpId = warpId - (world_size - 1);
-  int remoteRank = (warpId < rank) ? warpId : warpId + 1;
-  mscclppDevConn_t devConn = constDevConns[remoteRank];
-  if (isIB)
-    devConn = constDevConns[remoteRank + world_size];
-
-    // Each warp receives data from different ranks
-#if 1
-
-  // Trigger sending data, flag and synchronize after
-  devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
-
-  devConn.wait();
-
-#else
-  for (int i = 1; i < world_size; i++) {
-    __syncthreads();
-    if (remoteRank != ((rank + i) % world_size))
-      continue;
-
-    // Trigger sending data, flag and synchronize after
-    size_t ibPortion = nelemsPerGPU / 12; // nelemsPerGPU/12;
-    if (isIB)
-      devConn.fifo.setTrigger(trig, mscclppFlag | mscclppData | mscclppSync,
-                              rank * nelemsPerGPU * sizeof(int) + (nelemsPerGPU - ibPortion) * sizeof(int),
-                              rank * nelemsPerGPU * sizeof(int) + (nelemsPerGPU - ibPortion) * sizeof(int),
-                              ibPortion * sizeof(int));
-    else
-      devConn.fifo.setTrigger(trig, mscclppFlag | mscclppData | mscclppSync, rank * nelemsPerGPU * sizeof(int),
-                              rank * nelemsPerGPU * sizeof(int), (nelemsPerGPU - ibPortion) * sizeof(int));
-    // Wait on the request to make sure it is safe to reuse buffer and flag
-    auto req = devConn.fifo.putWithSignal(dataOffset, dataSize);
-    devConn.fifo.sync(req);
-  }
-  // Wait for receiving data from remote rank
-  while (*proxyFlag == baseFlag)
-    ;
-#endif
-}
-
-int rankToLocalRank(int rank)
-{
-  return rank % RANKS_PER_NODE;
-}
-
-int rankToNode(int rank)
-{
-  return rank / RANKS_PER_NODE;
-}
-
-int cudaNumToIbNum(int cudaNum)
-{
-  int ibNum;
-  if (cudaNum == 0) {
-    ibNum = 0;
-  } else if (cudaNum == 1) {
-    ibNum = 4;
-  } else if (cudaNum == 2) {
-    ibNum = 1;
-  } else if (cudaNum == 3) {
-    ibNum = 5;
-  } else if (cudaNum == 4) {
-    ibNum = 2;
-  } else if (cudaNum == 5) {
-    ibNum = 6;
-  } else if (cudaNum == 6) {
-    ibNum = 3;
-  } else if (cudaNum == 7) {
-    ibNum = 7;
-  } else {
-    printf("Invalid cudaNum: %d\n", cudaNum);
-    exit(EXIT_FAILURE);
-  }
-  return ibNum;
-}
-
-void print_usage(const char* prog)
-{
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  printf("usage: %s IP:PORT [rank nranks]\n", prog);
-#else
-  printf("usage: %s IP:PORT rank nranks\n", prog);
-#endif
-}
-
-int main(int argc, const char* argv[])
-{
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  if (argc != 2 && argc != 4) {
-    print_usage(argv[0]);
-    return -1;
-  }
-  const char* ip_port = argv[1];
-  int rank;
-  int world_size;
-  if (argc == 4) {
-    rank = atoi(argv[2]);
-    world_size = atoi(argv[3]);
-  } else {
-    MPI_Init(NULL, NULL);
-    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-    MPI_Comm_size(MPI_COMM_WORLD, &world_size);
-  }
-#else
-  if (argc != 4) {
-    print_usage(argv[0]);
-    return -1;
-  }
-  const char* ip_port = argv[1];
-  int rank = atoi(argv[2]);
-  int world_size = atoi(argv[3]);
-#endif
-  int localRank = rankToLocalRank(rank);
-  int thisNode = rankToNode(rank);
-  int cudaNum = localRank;
-  int ibNum = cudaNumToIbNum(cudaNum);
-
-  CUDACHECK(cudaSetDevice(cudaNum));
-  std::string ibDevStr = "mlx5_ib" + std::to_string(localRank);
-
-  mscclppComm_t comm;
-  MSCCLPPCHECK(mscclppCommInitRank(&comm, world_size, rank, ip_port));
-
-  int* data_d;
-  uint64_t* flag_d;
-  size_t data_size = 1536 * 1024 * 1024;
-  size_t nelemsPerGPU = data_size / sizeof(int) / world_size;
-  CUDACHECK(cudaMalloc(&data_d, data_size));
-  CUDACHECK(cudaMalloc(&flag_d, sizeof(uint64_t)));
-  CUDACHECK(cudaMemset(data_d, 0, data_size));
-  CUDACHECK(cudaMemset(flag_d, 0, sizeof(uint64_t)));
-
-  int* data_h = new int[nelemsPerGPU * world_size];
-  for (int i = 0; i < nelemsPerGPU * world_size; i++) {
-    size_t val = i + 1;
-    if (i / nelemsPerGPU == rank) {
-      data_h[i] = val;
-    } else {
-      data_h[i] = 0;
-    }
-  }
-  CUDACHECK(cudaMemcpy(data_d, data_h, data_size, cudaMemcpyHostToDevice));
-
-  mscclppDevConn_t devConns[16];
-  for (int r = 0; r < world_size; ++r) {
-    if (r == rank)
-      continue;
-    mscclppTransport_t transportType;
-    const char* ibDev = NULL;
-    transportType = mscclppTransportP2P;
-    // Connect with all other ranks
-    MSCCLPPCHECK(mscclppConnect(comm, &devConns[r], r, 0, data_d, data_size, flag_d, transportType, ibDev));
-  }
-  for (int r = 0; r < world_size; ++r) {
-    if (r == rank)
-      continue;
-    mscclppTransport_t transportType;
-    const char* ibDev = ibDevStr.c_str();
-    transportType = mscclppTransportIB;
-    // Connect with all other ranks
-    MSCCLPPCHECK(
-      mscclppConnect(comm, &devConns[r + world_size], r, 0, data_d, data_size, flag_d, transportType, ibDev));
-  }
-
-  MSCCLPPCHECK(mscclppConnectionSetup(comm));
-
-  MSCCLPPCHECK(mscclppProxyLaunch(comm));
-
-  CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * 2 * world_size));
-
-  cudaStream_t stream;
-  CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
-
-  CUDACHECK(cudaDeviceSynchronize());
-  kernel<<<1, 32 * 2 * (world_size - 1), 0, stream>>>(rank, world_size, nelemsPerGPU);
-  CUDACHECK(cudaDeviceSynchronize());
-  CUDACHECK(cudaMemcpy(data_h, data_d, data_size, cudaMemcpyDeviceToHost));
-  CUDACHECK(cudaDeviceSynchronize());
-
-  for (size_t i = 0; i < nelemsPerGPU * world_size; i++) {
-    int val = i + 1;
-    if (data_h[i] != val) {
-      printf("oh uh things went wrong! data_h[%d] (%d) != val (%d)\n", i, data_h[i], val);
-      break;
-    }
-  }
-  int tmp[16];
-  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-
-  //   // Perf test
-  //   cudaEvent_t ev_start;
-  //   cudaEvent_t ev_end;
-  //   CUDACHECK(cudaEventCreate(&ev_start));
-  //   CUDACHECK(cudaEventCreate(&ev_end));
-
-  // warm up
-  // int warmupiter = 1000;
-  // for (int i = 0; i < warmupiter; ++i) {
-  //   kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size, nelemsPerGPU);
-  // }
-  // CUDACHECK(cudaDeviceSynchronize());
-  // MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-
-  // cudaGraph Capture
-  cudaGraph_t graph;
-  cudaGraphExec_t instance;
-  cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
-  int cudagraphiter = 10;
-  for (int i = 0; i < cudagraphiter; ++i) {
-    kernel<<<1, 32 * 2 * (world_size - 1), 0, stream>>>(rank, world_size, nelemsPerGPU);
-  }
-  cudaStreamEndCapture(stream, &graph);
-  cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);
-
-  int cudagraphwarmup = 10;
-  for (int i = 0; i < cudagraphwarmup; ++i) {
-    cudaGraphLaunch(instance, stream);
-  }
-  CUDACHECK(cudaStreamSynchronize(stream));
-
-  // measure runtime
-  //  CUDACHECK(cudaEventRecord(ev_start, stream));
-  double t0 = getTime();
-  int cudagraphlaunch = 10;
-  for (int i = 0; i < cudagraphlaunch; ++i) {
-    // kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size);
-    cudaGraphLaunch(instance, stream);
-  }
-  //  CUDACHECK(cudaEventRecord(ev_end, stream));
-  CUDACHECK(cudaStreamSynchronize(stream));
-
-  double t1 = getTime();
-  float ms = (t1 - t0) * 1000.0;
-  //  CUDACHECK(cudaEventElapsedTime(&ms, ev_start, ev_end));
-  double time_in_us = ms * 1000. / (float)cudagraphlaunch / (float)cudagraphiter;
-  printf("rank: %d, time: %f us/iter algBW %f\n", rank, time_in_us,
-         (double)(data_size) / 1024. / 1024. / 1024. / (time_in_us / 1e6));
-
-  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
-  MSCCLPPCHECK(mscclppProxyStop(comm));
-
-  MSCCLPPCHECK(mscclppCommDestroy(comm));
-
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  if (argc == 2) {
-    MPI_Finalize();
-  }
-#endif
-  printf("Succeeded! %d\n", rank);
-  return 0;
-}
diff --git a/tests/allgather_test3.cu b/tests/allgather_test3.cu
deleted file mode 100644
index 1f16fba9..00000000
--- a/tests/allgather_test3.cu
+++ /dev/null
@@ -1,225 +0,0 @@
-/*************************************************************************
- * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
- *
- * See LICENSE.txt for license information
- ************************************************************************/
-
-#include "comm.h"
-#include "common.h"
-
-#include <cuda_runtime.h>
-#include <string>
-
-#define ALIGN 4
-__constant__ mscclppDevConn_t constDevConns[16];
-
-__device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, size_t nelemsPerGPU)
-{
-  // this allgather is really simple and implemented as an alltoall
-
-  // this thread's role is a sender role
-  // put your data asynchronously
-  if (threadIdx.x % 32 == 0)
-    devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
-  // make sure everyone is put their data before some thread randomly blocks everyone else in signal
-  __syncthreads();
-  // push with flag and sync to make sure the data is received
-  if (threadIdx.x % 32 == 0)
-    devConn.flush();
-
-  // this thread's role is a receiver role. wait on the semaphore to make sure the data is ready
-  if (threadIdx.x % 32 == 0)
-    devConn.wait();
-}
-
-__device__ void localAllGather(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
-                               uint64_t offset, uint64_t size)
-{
-  // this allgather algorithm works as follows:
-  // Step 1: GPU rank i sends data to GPU rank (i+1) % nranksPerNode
-  // and waits for data from GPU rank (i-1) % nranksPerNode
-  // Step 2: GPU rank i sends data to GPU rank (i+2) % nranksPerNode
-  // ...
-  // This order is much better for DMA engine for NVLinks
-  for (int i = 1; i < nranksPerNode; i++) {
-    if ((remoteRank % nranksPerNode) == ((rank + i) % nranksPerNode)) {
-      // put your data to GPU (rank+i) % nranksPerNode and signal in one call
-      if ((threadIdx.x % 32) == 0)
-        devConn.putWithSignalAndFlush(offset, size);
-    }
-    // wait for the data from GPU (rank-i) % nranksPerNode to arrive
-    if ((remoteRank % nranksPerNode) == ((rank - i + nranksPerNode) % nranksPerNode)) {
-      if ((threadIdx.x % 32) == 0)
-        devConn.wait();
-    }
-    asm volatile("bar.sync %0, %1;" ::"r"(11), "r"((nranksPerNode - 1) * 32) : "memory");
-  }
-}
-
-__device__ void allgather1(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
-                           size_t nelemsPerGPU)
-{
-  localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, rank * nelemsPerGPU * sizeof(int),
-                 nelemsPerGPU * sizeof(int));
-}
-
-__device__ void allgather2(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
-                           size_t nelemsPerGPU)
-{
-  // this allgather is a pipelined and hierarchical one and only works for two nodes
-  // it is implemented as follows:
-  // Step 1: each node does a local allgather and concurrently,
-  // local GPU i exchange (piplineSize-1)/pipelineSize portion of their data with
-  // its cross-node neighbor (local GPU i on the other node) via IB
-  // Step 2: each node does a local allgather again with the data just received from its
-  // cross-node neighbor in step 1, and concurrently, exchange the rest of the data with
-  // its cross-node neighbor
-  // Step 3: each node does a local allgather for the last time with the rest of the data
-
-  int pipelineSize = 3;
-
-  // Step 1
-  // local allgather
-  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
-    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, rank * nelemsPerGPU * sizeof(int),
-                   nelemsPerGPU * sizeof(int));
-  }
-  // cross-node exchange
-  if (remoteRank % nranksPerNode == rank % nranksPerNode) {
-    // opposite side
-    if ((threadIdx.x % 32) == 0)
-      devConn.putWithSignalAndFlush(rank * nelemsPerGPU * sizeof(int),
-                                    (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize * sizeof(int));
-    if ((threadIdx.x % 32) == 0)
-      devConn.wait();
-  }
-
-  __syncthreads();
-
-  // Step 2
-  // local allgather
-  int otherNghr = (rank + nranksPerNode) % world_size;
-  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
-    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, otherNghr * nelemsPerGPU * sizeof(int),
-                   (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize * sizeof(int));
-  }
-
-  // cross-node exchange
-  if (remoteRank % nranksPerNode == rank % nranksPerNode) {
-    // opposite side
-    if ((threadIdx.x % 32) == 0)
-      devConn.putWithSignalAndFlush((rank * nelemsPerGPU + (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize) *
-                                      sizeof(int),
-                                    nelemsPerGPU / pipelineSize * sizeof(int));
-    if ((threadIdx.x % 32) == 0)
-      devConn.wait();
-  }
-
-  __syncthreads();
-
-  // Step 3
-  // local allgather
-  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
-    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank,
-                   (otherNghr * nelemsPerGPU + (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize) * sizeof(int),
-                   nelemsPerGPU / pipelineSize * sizeof(int));
-  }
-}
-
-__global__ void kernel(int rank, int world_size, int nranksPerNode, size_t nelemsPerGPU, int kernel)
-{
-  // find the mapping between remoteRank and devConns
-  int warpId = threadIdx.x / 32;
-  int remoteRank = (warpId < rank) ? warpId : warpId + 1;
-  // Each warp is responsible for one of the remote ranks
-  mscclppDevConn_t devConn = constDevConns[warpId];
-
-  if (kernel == 0)
-    allgather0(devConn, rank, world_size, remoteRank, nelemsPerGPU);
-  else if (kernel == 1)
-    allgather1(devConn, rank, world_size, nranksPerNode, remoteRank, nelemsPerGPU);
-  else if (kernel == 2)
-    allgather2(devConn, rank, world_size, nranksPerNode, remoteRank, nelemsPerGPU);
-}
-
-void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* paramcount, size_t* sendInplaceOffset,
-                               size_t* recvInplaceOffset, size_t count, int nranks)
-{
-  size_t base = (count / (ALIGN * nranks)) * ALIGN;
-  *sendcount = base;
-  *recvcount = base * nranks;
-  *sendInplaceOffset = base;
-  *recvInplaceOffset = 0;
-  *paramcount = base;
-}
-
-testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
-{
-  size_t sendcount = args->sendBytes / sizeof(int);
-  size_t recvcount = args->expectedBytes / sizeof(int);
-  // int nranks = args->totalProcs;
-
-  CUDACHECK(cudaSetDevice(args->gpus[0]));
-  int rank = args->proc;
-  CUDACHECK(cudaMemset(args->recvbuffs[0], 0, args->expectedBytes));
-  // void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
-
-  int* dataHost = new int[recvcount];
-  for (size_t i = 0; i < recvcount; i++) {
-    int val = i + 1;
-    if (i / sendcount == (size_t)rank) {
-      dataHost[i] = val;
-    } else {
-      dataHost[i] = 0;
-    }
-  }
-  CUDACHECK(cudaMemcpy(args->recvbuffs[0], dataHost, recvcount, cudaMemcpyHostToDevice));
-  for (int i = 0; i < static_cast<int>(recvcount); i++) {
-    dataHost[i] = i + 1;
-  }
-  CUDACHECK(cudaMemcpy(args->expected[0], dataHost, recvcount, cudaMemcpyHostToDevice));
-  delete dataHost;
-  CUDACHECK(cudaDeviceSynchronize());
-  return testSuccess;
-}
-
-void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks)
-{
-  double baseBw = (double)(count * typesize * nranks) / 1.0E9 / sec;
-
-  *algBw = baseBw;
-  double factor = ((double)(nranks - 1)) / ((double)nranks);
-  *busBw = baseBw * factor;
-}
-
-testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, int nranksPerNode, size_t count, mscclppComm_t comm,
-                              cudaStream_t stream)
-{
-  int worldSize = comm->nRanks;
-  kernel<<<1, 32 * (worldSize - 1), 0, stream>>>(comm->rank, worldSize, nranksPerNode, count / sizeof(int), 1);
-  return testSuccess;
-}
-
-struct testColl allGatherTest = {"AllGather", AllGatherGetCollByteCount, AllGatherInitData, AllGatherGetBw,
-                                 AllGatherRunColl};
-
-void AllGatherGetBuffSize(size_t* sendcount, size_t* recvcount, size_t count, int nranks)
-{
-  size_t paramcount, sendInplaceOffset, recvInplaceOffset;
-  AllGatherGetCollByteCount(sendcount, recvcount, &paramcount, &sendInplaceOffset, &recvInplaceOffset, count, nranks);
-}
-
-testResult_t AllGatherRunTest(struct threadArgs* args)
-{
-  args->collTest = &allGatherTest;
-  mscclppDevConn_t* devConns;
-  int nCons;
-  MSCCLPPCHECK(mscclppGetAllDeviceConnections(args->comm, &devConns, &nCons));
-  CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * nCons));
-  TESTCHECK(TimeTest(args));
-  return testSuccess;
-}
-
-struct testEngine allGatherEngine = {AllGatherGetBuffSize, AllGatherRunTest};
-
-#pragma weak mscclppTestEngine = allGatherEngine
diff --git a/tests/allgather_test_standalone.cu b/tests/allgather_test_standalone.cu
new file mode 100644
index 00000000..05aa3304
--- /dev/null
+++ b/tests/allgather_test_standalone.cu
@@ -0,0 +1,501 @@
+#include "mscclpp.h"
+
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+#include "mpi.h"
+#endif // MSCCLPP_USE_MPI_FOR_TESTS
+#include <iostream>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string>
+#include <unistd.h>
+#include <unordered_map>
+
+static int nranksPerNode = 8;
+
+// Propagate errors up
+
+#define MSCCLPPCHECK(call)                                                                                             \
+  do {                                                                                                                 \
+    mscclppResult_t res = call;                                                                                        \
+    if (res != mscclppSuccess && res != mscclppInProgress) {                                                           \
+      /* Print the back trace*/                                                                                        \
+      printf("Failure at %s:%d -> %s\n", __FILE__, __LINE__, mscclppGetErrorString(res));                              \
+      return res;                                                                                                      \
+    }                                                                                                                  \
+  } while (0)
+
+// Check CUDA RT calls
+#define CUDACHECK(cmd)                                                                                                 \
+  do {                                                                                                                 \
+    cudaError_t err = cmd;                                                                                             \
+    if (err != cudaSuccess) {                                                                                          \
+      printf("%s:%d Cuda failure '%s'\n", __FILE__, __LINE__, cudaGetErrorString(err));                                \
+      exit(EXIT_FAILURE);                                                                                              \
+    }                                                                                                                  \
+  } while (false)
+
+// Measure current time in second.
+static double getTime(void)
+{
+  struct timespec tspec;
+  if (clock_gettime(CLOCK_MONOTONIC, &tspec) == -1) {
+    printf("clock_gettime failed\n");
+    exit(EXIT_FAILURE);
+  }
+  return (tspec.tv_nsec / 1.0e9) + tspec.tv_sec;
+}
+
+__constant__ mscclppDevConn_t constDevConns[16];
+
+__device__ void allgather0(mscclppDevConn_t devConn, int rank, int world_size, int remoteRank, size_t nelemsPerGPU)
+{
+  // this allgather is really simple and implemented as an alltoall
+
+  // this thread's role is a sender role
+  // put your data asynchronously
+  if ((threadIdx.x % 32) == 0)
+    devConn.putWithSignal(rank * nelemsPerGPU * sizeof(int), nelemsPerGPU * sizeof(int));
+  // make sure everyone is put their data before some thread randomly blocks everyone else in signal
+  __syncthreads();
+  // push with flag and sync to make sure the data is received
+  if ((threadIdx.x % 32) == 0)
+    devConn.flush();
+
+  // this thread's role is a receiver role. wait on the semaphore to make sure the data is ready
+  if ((threadIdx.x % 32) == 0)
+    devConn.wait();
+}
+
+__device__ void localAllGather(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
+                               uint64_t offset, uint64_t size)
+{
+  // this allgather algorithm works as follows:
+  // Step 1: GPU rank i sends data to GPU rank (i+1) % nranksPerNode
+  // and waits for data from GPU rank (i-1) % nranksPerNode
+  // Step 2: GPU rank i sends data to GPU rank (i+2) % nranksPerNode
+  // ...
+  // This order is much better for DMA engine for NVLinks
+  for (int i = 1; i < nranksPerNode; i++) {
+    if ((remoteRank % nranksPerNode) == ((rank + i) % nranksPerNode)) {
+      // put your data to GPU (rank+i) % nranksPerNode and signal in one call
+      if ((threadIdx.x % 32) == 0)
+        devConn.putWithSignalAndFlush(offset, size);
+    }
+    // wait for the data from GPU (rank-i) % nranksPerNode to arrive
+    if ((remoteRank % nranksPerNode) == ((rank - i + nranksPerNode) % nranksPerNode)) {
+      if ((threadIdx.x % 32) == 0)
+        devConn.wait();
+    }
+    asm volatile("bar.sync %0, %1;" ::"r"(11), "r"((nranksPerNode - 1) * 32) : "memory");
+  }
+}
+
+__device__ void allgather1(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
+                           size_t nelemsPerGPU)
+{
+  localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, rank * nelemsPerGPU * sizeof(int),
+                 nelemsPerGPU * sizeof(int));
+}
+
+__device__ void allgather2(mscclppDevConn_t devConn, int rank, int world_size, int nranksPerNode, int remoteRank,
+                           size_t nelemsPerGPU)
+{
+  // this allgather is a pipelined and hierarchical one and only works for two nodes
+  // it is implemented as follows:
+  // Step 1: each node does a local allgather and concurrently,
+  // local GPU i exchange (piplineSize-1)/pipelineSize portion of their data with
+  // its cross-node neighbor (local GPU i on the other node) via IB
+  // Step 2: each node does a local allgather again with the data just received from its
+  // cross-node neighbor in step 1, and concurrently, exchange the rest of the data with
+  // its cross-node neighbor
+  // Step 3: each node does a local allgather for the last time with the rest of the data
+
+  int pipelineSize = 3;
+
+  // Step 1
+  // local allgather
+  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
+    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, rank * nelemsPerGPU * sizeof(int),
+                   nelemsPerGPU * sizeof(int));
+  }
+  // cross-node exchange
+  if (remoteRank % nranksPerNode == rank % nranksPerNode) {
+    // opposite side
+    if ((threadIdx.x % 32) == 0)
+      devConn.putWithSignalAndFlush(rank * nelemsPerGPU * sizeof(int),
+                                    (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize * sizeof(int));
+    if ((threadIdx.x % 32) == 0)
+      devConn.wait();
+  }
+
+  __syncthreads();
+
+  // Step 2
+  // local allgather
+  int otherNghr = (rank + nranksPerNode) % world_size;
+  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
+    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank, otherNghr * nelemsPerGPU * sizeof(int),
+                   (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize * sizeof(int));
+  }
+
+  // cross-node exchange
+  if (remoteRank % nranksPerNode == rank % nranksPerNode) {
+    // opposite side
+    if ((threadIdx.x % 32) == 0)
+      devConn.putWithSignalAndFlush((rank * nelemsPerGPU + (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize) *
+                                      sizeof(int),
+                                    nelemsPerGPU / pipelineSize * sizeof(int));
+    if ((threadIdx.x % 32) == 0)
+      devConn.wait();
+  }
+
+  __syncthreads();
+
+  // Step 3
+  // local allgather
+  if (remoteRank / nranksPerNode == rank / nranksPerNode) {
+    localAllGather(devConn, rank, world_size, nranksPerNode, remoteRank,
+                   (otherNghr * nelemsPerGPU + (nelemsPerGPU * (pipelineSize - 1)) / pipelineSize) * sizeof(int),
+                   nelemsPerGPU / pipelineSize * sizeof(int));
+  }
+}
+
+__global__ void kernel(int rank, int world_size, int nranksPerNode, size_t nelemsPerGPU, int kernel)
+{
+  // find the mapping between remoteRank and devConns
+  int warpId = threadIdx.x / 32;
+  int remoteRank = (warpId < rank) ? warpId : warpId + 1;
+  // Each warp is responsible for one of the remote ranks
+  mscclppDevConn_t devConn = constDevConns[warpId];
+
+  if (kernel == 0)
+    allgather0(devConn, rank, world_size, remoteRank, nelemsPerGPU);
+  else if (kernel == 1)
+    allgather1(devConn, rank, world_size, nranksPerNode, remoteRank, nelemsPerGPU);
+  else if (kernel == 2)
+    allgather2(devConn, rank, world_size, nranksPerNode, remoteRank, nelemsPerGPU);
+}
+
+int rankToLocalRank(int rank)
+{
+  return rank % nranksPerNode;
+}
+
+int rankToNode(int rank)
+{
+  return rank / nranksPerNode;
+}
+
+void print_usage(const char* prog)
+{
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+  printf("usage: %s IP:PORT [rank nranks]\n", prog);
+#else
+  printf("usage: %s IP:PORT rank nranks\n", prog);
+#endif
+}
+
+void initializeAndAllocateAllGatherData(int rank, int world_size, size_t dataSize, size_t nelemsPerGPU, int** data_h,
+                                        int** data_d)
+{
+  CUDACHECK(cudaMalloc(data_d, dataSize));
+  CUDACHECK(cudaMemset(*data_d, 0, dataSize));
+
+  *data_h = new int[nelemsPerGPU * world_size];
+  for (size_t i = 0; i < nelemsPerGPU * world_size; i++) {
+    int val = i + 1;
+    if (i / nelemsPerGPU == (size_t)rank) {
+      (*data_h)[i] = val;
+    } else {
+      (*data_h)[i] = 0;
+    }
+  }
+  CUDACHECK(cudaMemcpy(*data_d, *data_h, dataSize, cudaMemcpyHostToDevice));
+}
+
+mscclppResult_t setupMscclppConnections(int rank, int world_size, mscclppComm_t comm, int* data_d, size_t dataSize)
+{
+  int thisNode = rankToNode(rank);
+  int cudaNum = rankToLocalRank(rank);
+  std::string ibDevStr = "mlx5_ib" + std::to_string(cudaNum);
+
+  for (int r = 0; r < world_size; ++r) {
+    if (r == rank)
+      continue;
+    mscclppTransport_t transportType;
+    const char* ibDev = ibDevStr.c_str();
+    if (rankToNode(r) == thisNode) {
+      ibDev = NULL;
+      transportType = mscclppTransportP2P;
+    } else {
+      transportType = mscclppTransportIB;
+    }
+    // Connect with all other ranks
+    MSCCLPPCHECK(mscclppConnect(comm, r, 0, data_d, dataSize, transportType, ibDev));
+  }
+
+  MSCCLPPCHECK(mscclppConnectionSetup(comm));
+
+  mscclppDevConn_t* devConns;
+  int nCons;
+  MSCCLPPCHECK(mscclppGetAllDeviceConnections(comm, &devConns, &nCons));
+
+  CUDACHECK(cudaMemcpyToSymbol(constDevConns, devConns, sizeof(mscclppDevConn_t) * nCons));
+
+  return mscclppSuccess;
+}
+
+void printUsage(const char* prog, bool isMpi)
+{
+  if (isMpi) {
+    std::string st = "you are using MPI for this test\n";
+    st += "two possilbe usages are:\n";
+    st += "> " + std::string(prog) + "\n";
+    st += "or\n";
+    st += "> " + std::string(prog) + " -ip_port [ip:port]\n";
+    printf("%s", st.c_str());
+  } else {
+    std::string st = "you are NOT using MPI for this test\n";
+    st += "the only possible usage:\n";
+    st += "> " + std::string(prog) + " -ip_port [ip:port] -rank [rank] -nranks [nranks]\n";
+    printf("%s", st.c_str());
+  }
+}
+
+std::unordered_map<std::string, std::string> parseArgs(int argc, const char* argv[], bool isMpi)
+{
+  std::unordered_map<std::string, std::string> options;
+
+  for (int i = 1; i < argc; i++) {
+    std::string arg = argv[i];
+    if (arg == "-rankspernode") {
+      if (isMpi) {
+        fprintf(stderr, "Error: -rankspernode should not be specified with MPI.\n");
+        exit(-1);
+      }
+      if (i + 1 < argc) {
+        options["rankspernode"] = argv[++i];
+      } else {
+        fprintf(stderr, "Error: -rankspernode option requires an argument.\n");
+        ;
+        exit(-1);
+      }
+    } else if (arg == "-kernel") {
+      if (i + 1 < argc) {
+        options["kernel"] = argv[++i];
+      } else {
+        fprintf(stderr, "Error: -kernel option requires an argument.\n");
+        exit(-1);
+      }
+    } else if (arg == "-ip_port") {
+      if (i + 1 < argc) {
+        options["ip_port"] = argv[++i];
+      } else {
+        fprintf(stderr, "Error: -ip_port option requires an argument.\n");
+        exit(-1);
+      }
+    } else if (arg == "-rank") {
+      if (isMpi) {
+        fprintf(stderr, "Error: -rank should not be specified with MPI.\n");
+        exit(-1);
+      }
+      if (i + 1 < argc) {
+        options["rank"] = argv[++i];
+      } else {
+        fprintf(stderr, "Error: -ip_port option requires an argument.\n");
+        exit(-1);
+      }
+    } else if (arg == "-nranks") {
+      if (isMpi) {
+        fprintf(stderr, "Error: -nranks should not be specified with MPI.\n");
+        exit(-1);
+      }
+      if (i + 1 < argc) {
+        options["nranks"] = argv[++i];
+      } else {
+        fprintf(stderr, "Error: -ip_port option requires an argument.\n");
+        exit(-1);
+      }
+    } else if (arg == "-datasize") {
+      if (i + 1 < argc) {
+        options["datasize"] = argv[++i];
+      } else {
+        fprintf(stderr, "Error: -datasize option requires an argument.\n");
+        exit(-1);
+      }
+    } else if (arg == "-help" || arg == "-h") {
+      printUsage(argv[0], isMpi);
+      exit(0);
+    } else {
+      fprintf(stderr, "Error: Unknown option %s\n", argv[i]);
+      exit(-1);
+    }
+  }
+  return options;
+}
+
+int main(int argc, const char* argv[])
+{
+  bool isMpi = false;
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+  isMpi = true;
+#endif
+
+  auto parsedArgs = parseArgs(argc, argv, isMpi);
+
+  int rank;
+  int world_size;
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+  MPI_Init(NULL, NULL);
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm_size(MPI_COMM_WORLD, &world_size);
+  // get the local number of nodes with MPI
+  MPI_Comm shmcomm;
+  MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &shmcomm);
+  int shmrank;
+  MPI_Comm_size(shmcomm, &shmrank);
+  nranksPerNode = shmrank;
+  MPI_Comm_free(&shmcomm);
+#else
+  if (parsedArgs.find("rank") == parsedArgs.end() || parsedArgs.find("nranks") == parsedArgs.end()) {
+    printUsage(argv[0], isMpi);
+    exit(-1);
+  }
+  rank = std::stoi(parsedArgs["rank"]);
+  world_size = std::stoi(parsedArgs["nranks"]);
+  if (parsedArgs.find("rankspernode") == parsedArgs.end()) {
+    printUsage(argv[0], isMpi);
+    exit(-1);
+  }
+  nranksPerNode = std::stoi(parsedArgs["rankspernode"]);
+#endif
+  int kernelNum = 0;
+  if (parsedArgs.find("kernel") != parsedArgs.end()) {
+    kernelNum = std::stoi(parsedArgs["kernel"]);
+  }
+  char* ip_port = NULL;
+  if (parsedArgs.find("ip_port") == parsedArgs.end()) {
+    printUsage(argv[0], isMpi);
+    exit(-1);
+  }
+  ip_port = (char*)parsedArgs["ip_port"].c_str();
+
+  int thisNode = rankToNode(rank);
+  int cudaNum = rankToLocalRank(rank);
+  CUDACHECK(cudaSetDevice(cudaNum));
+
+  if (rank == 0)
+    printf("Initializing MSCCL++\n");
+  mscclppComm_t comm;
+  MSCCLPPCHECK(mscclppCommInitRank(&comm, world_size, ip_port, rank));
+
+  int* data_d;
+  int* data_h;
+  size_t dataSize = 1024 * 1024 * 1024;
+  if (parsedArgs.find("datasize") != parsedArgs.end()) {
+    dataSize = std::stoul(parsedArgs["datasize"]);
+  }
+  size_t nelemsPerGPU = dataSize / sizeof(int) / world_size;
+
+  if (rank == 0)
+    printf("Initializing data for allgather test\n");
+  initializeAndAllocateAllGatherData(rank, world_size, dataSize, nelemsPerGPU, &data_h, &data_d);
+
+  if (rank == 0)
+    printf("Setting up the connection in MSCCL++\n");
+  MSCCLPPCHECK(setupMscclppConnections(rank, world_size, comm, data_d, dataSize));
+
+  if (rank == 0)
+    printf("Launching MSCCL++ proxy threads\n");
+  MSCCLPPCHECK(mscclppProxyLaunch(comm));
+
+  if (rank == 0)
+    printf("Testing the correctness of AllGather implementation\n");
+  cudaStream_t stream;
+  CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+  CUDACHECK(cudaDeviceSynchronize());
+  kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size, nranksPerNode, nelemsPerGPU, kernelNum);
+  CUDACHECK(cudaDeviceSynchronize());
+  CUDACHECK(cudaMemcpy(data_h, data_d, dataSize, cudaMemcpyDeviceToHost));
+
+  for (size_t i = 0; i < nelemsPerGPU * world_size; i++) {
+    int val = i + 1;
+    if (data_h[i] != val) {
+      printf("oh uh! data_h[%ld] (%d) != val (%d)\n", i, data_h[i], val);
+      break;
+    }
+  }
+  int tmp[16];
+  // A simple barrier
+  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
+  if (rank == 0)
+    printf("Successfully checked the correctness\n");
+
+  // Perf test
+  int iterwithoutcudagraph = 10;
+  if (rank == 0)
+    printf("Running %d iterations of the kernel without CUDA graph\n", iterwithoutcudagraph);
+  CUDACHECK(cudaStreamSynchronize(stream));
+  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
+  for (int i = 0; i < iterwithoutcudagraph; ++i) {
+    kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size, nranksPerNode, nelemsPerGPU, kernelNum);
+  }
+  CUDACHECK(cudaStreamSynchronize(stream));
+  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
+
+  // cudaGraph Capture
+  int cudagraphiter = 10;
+  if (rank == 0)
+    printf("Capturing %d iterations of the kernel in a CUDA graph\n", cudagraphiter);
+  cudaGraph_t graph;
+  cudaGraphExec_t instance;
+  cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
+  for (int i = 0; i < cudagraphiter; ++i) {
+    kernel<<<1, 32 * (world_size - 1), 0, stream>>>(rank, world_size, nranksPerNode, nelemsPerGPU, kernelNum);
+  }
+  cudaStreamEndCapture(stream, &graph);
+  cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);
+
+  int cudagraphwarmup = 10;
+  if (rank == 0)
+    printf("Warming up %d iterations of the CUDA graph with %d iterations of the kernel\n", cudagraphwarmup,
+           cudagraphiter);
+  for (int i = 0; i < cudagraphwarmup; ++i) {
+    cudaGraphLaunch(instance, stream);
+  }
+  CUDACHECK(cudaStreamSynchronize(stream));
+
+  // measure runtime
+  int cudagraphlaunch = 10;
+  if (rank == 0)
+    printf("Running %d iterations of the CUDA graph with %d iterations of the kernel\n", cudagraphlaunch,
+           cudagraphiter);
+  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
+  double t0, t1, ms, time_in_us;
+  t0 = getTime();
+  for (int i = 0; i < cudagraphlaunch; ++i) {
+    cudaGraphLaunch(instance, stream);
+  }
+  CUDACHECK(cudaStreamSynchronize(stream));
+
+  t1 = getTime();
+  ms = (t1 - t0) * 1000.0;
+  time_in_us = ms * 1000. / (float)cudagraphlaunch / (float)cudagraphiter;
+  printf("Rank %d report: size %lu time: %f us/iter algBW %f GBps\n", rank, dataSize, time_in_us,
+         (double)(dataSize) / 1e9 / (time_in_us / 1e6));
+  MSCCLPPCHECK(mscclppBootstrapAllGather(comm, tmp, sizeof(int)));
+
+  if (rank == 0)
+    printf("Stopping MSCCL++ proxy threads\n");
+  MSCCLPPCHECK(mscclppProxyStop(comm));
+
+  if (rank == 0)
+    printf("Destroying MSCCL++ communicator\n");
+  MSCCLPPCHECK(mscclppCommDestroy(comm));
+  printf("Rank %d succeeded!\n", rank);
+
+#ifdef MSCCLPP_USE_MPI_FOR_TESTS
+  MPI_Finalize();
+#endif
+  return 0;
+}

From bf472ff8648098b5803d0fa27412ac831a3b3282 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 7 Apr 2023 03:37:47 +0000
Subject: [PATCH 18/23] Fix bug & remove pthread related code

---
 Makefile                |   2 +-
 tests/allgather_test.cu |  10 +--
 tests/common.cu         | 190 +++++++++++++++++-----------------------
 tests/common.h          |  28 +-----
 tests/timer.cc          |  32 -------
 tests/timer.h           |  17 ----
 6 files changed, 86 insertions(+), 193 deletions(-)
 delete mode 100644 tests/timer.cc
 delete mode 100644 tests/timer.h

diff --git a/Makefile b/Makefile
index 692ae56b..ae1343e9 100644
--- a/Makefile
+++ b/Makefile
@@ -193,7 +193,7 @@ $(BUILDDIR)/$(BINDIR)/%: $(BUILDDIR)/$(OBJDIR)/%.o $(LIBTARGET)
 	$(NVCC) -o $@ $< $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp
 
 # Compile mscclpp_test
-$(BUILDDIR)/$(BINDIR)/$(TESTSDIR)/%_perf: $(MSCLLPPTESTSOBJSDIR)/%.o $(MSCLLPPTESTSOBJSDIR)/common.o $(MSCLLPPTESTSOBJSDIR)/timer.o
+$(BUILDDIR)/$(BINDIR)/$(TESTSDIR)/%_perf: $(MSCLLPPTESTSOBJSDIR)/%.o $(MSCLLPPTESTSOBJSDIR)/common.o
 	@mkdir -p $(@D)
 	$(NVCC) -o $@ $^ $(MPI_LDFLAGS) -L$(BUILDDIR)/$(LIBDIR) -lmscclpp
 
diff --git a/tests/allgather_test.cu b/tests/allgather_test.cu
index 1f16fba9..8354e69a 100644
--- a/tests/allgather_test.cu
+++ b/tests/allgather_test.cu
@@ -1,9 +1,3 @@
-/*************************************************************************
- * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
- *
- * See LICENSE.txt for license information
- ************************************************************************/
-
 #include "comm.h"
 #include "common.h"
 
@@ -173,11 +167,11 @@ testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
       dataHost[i] = 0;
     }
   }
-  CUDACHECK(cudaMemcpy(args->recvbuffs[0], dataHost, recvcount, cudaMemcpyHostToDevice));
+  CUDACHECK(cudaMemcpy(args->recvbuffs[0], dataHost, recvcount * sizeof(int), cudaMemcpyHostToDevice));
   for (int i = 0; i < static_cast<int>(recvcount); i++) {
     dataHost[i] = i + 1;
   }
-  CUDACHECK(cudaMemcpy(args->expected[0], dataHost, recvcount, cudaMemcpyHostToDevice));
+  CUDACHECK(cudaMemcpy(args->expected[0], dataHost, recvcount * sizeof(int), cudaMemcpyHostToDevice));
   delete dataHost;
   CUDACHECK(cudaDeviceSynchronize());
   return testSuccess;
diff --git a/tests/common.cu b/tests/common.cu
index 3e11cb44..6722f57c 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -7,8 +7,9 @@
 #include "common.h"
 #include "cuda.h"
 #include "mscclpp.h"
-#include "timer.h"
 
+#include <chrono>
+#include <cstdint>
 #include <cstdio>
 #include <iostream>
 #include <string>
@@ -16,29 +17,45 @@
 
 #include <getopt.h>
 #include <libgen.h>
-#include <pthread.h>
+
+#define NUM_BLOCKS 32
 
 int is_main_proc = 0;
 thread_local int is_main_thread = 0;
 
+namespace {
+class timer
+{
+  std::uint64_t t0;
+
+public:
+  timer();
+  double elapsed() const;
+  double reset();
+};
+
+std::uint64_t now()
+{
+  using clock = std::chrono::steady_clock;
+  return std::chrono::duration_cast<std::chrono::nanoseconds>(clock::now().time_since_epoch()).count();
+}
+
 // Command line parameter defaults
-static size_t minBytes = 32 * 1024 * 1024;
-static size_t maxBytes = 32 * 1024 * 1024;
-static size_t stepBytes = 1 * 1024 * 1024;
-static size_t stepFactor = 1;
-static int datacheck = 1;
-static int warmup_iters = 10;
-static int iters = 100;
-static int timeout = 0;
-static int report_cputime = 0;
+size_t minBytes = 32 * 1024 * 1024;
+size_t maxBytes = 32 * 1024 * 1024;
+size_t stepBytes = 1 * 1024 * 1024;
+size_t stepFactor = 1;
+int datacheck = 1;
+int warmup_iters = 10;
+int iters = 100;
+int timeout = 0;
+int report_cputime = 0;
 // Report average iteration time: (0=RANK0,1=AVG,2=MIN,3=MAX)
-static int average = 1;
-static std::string ip_port;
-static int cudaGraphLaunches = 15;
+int average = 1;
+std::string ip_port;
+int cudaGraphLaunches = 15;
 
-#define NUM_BLOCKS 32
-
-static double parsesize(const char* value)
+double parsesize(const char* value)
 {
   long long int units;
   double size;
@@ -75,31 +92,32 @@ static double parsesize(const char* value)
   return size * units;
 }
 
-void Barrier(struct threadArgs* args)
+inline testResult_t Barrier(struct threadArgs* args)
 {
-  thread_local int epoch = 0;
-  static pthread_mutex_t lock[2] = {PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER};
-  static pthread_cond_t cond[2] = {PTHREAD_COND_INITIALIZER, PTHREAD_COND_INITIALIZER};
-  static int counter[2] = {0, 0};
+  int tmp[16];
+  // A simple barrier
+  MSCCLPPCHECK(mscclppBootstrapAllGather(args->comm, tmp, sizeof(int)));
+  return testSuccess;
+}
+} // namespace
 
-  pthread_mutex_lock(&lock[epoch]);
-  if (++counter[epoch] == args->nThreads)
-    pthread_cond_broadcast(&cond[epoch]);
+timer::timer()
+{
+  t0 = now();
+}
 
-  if (args->thread + 1 == args->nThreads) {
-    while (counter[epoch] != args->nThreads)
-      pthread_cond_wait(&cond[epoch], &lock[epoch]);
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-    MPI_Barrier(MPI_COMM_WORLD);
-#endif
-    counter[epoch] = 0;
-    pthread_cond_broadcast(&cond[epoch]);
-  } else {
-    while (counter[epoch] != 0)
-      pthread_cond_wait(&cond[epoch], &lock[epoch]);
-  }
-  pthread_mutex_unlock(&lock[epoch]);
-  epoch ^= 1;
+double timer::elapsed() const
+{
+  std::uint64_t t1 = now();
+  return 1.e-9 * (t1 - t0);
+}
+
+double timer::reset()
+{
+  std::uint64_t t1 = now();
+  double ans = 1.e-9 * (t1 - t0);
+  t0 = t1;
+  return ans;
 }
 
 testResult_t AllocateBuffs(void** sendbuff, size_t sendBytes, void** recvbuff, size_t recvBytes, void** expected,
@@ -165,77 +183,33 @@ testResult_t completeColl(struct threadArgs* args)
   return testSuccess;
 }
 
-// Inter-thread/process barrier+allreduce. The quality of the return value
-// for average=0 (which means broadcast from rank=0) is dubious. The returned
-// value will actually be the result of process-local broadcast from the local thread=0.
+// Inter process barrier+allreduce. The quality of the return value
+// for average=0 is just value itself.
+// Inter process barrier+allreduce. The quality of the return value
+// for average=0 is just value itself.
 template <typename T> void Allreduce(struct threadArgs* args, T* value, int average)
 {
-  thread_local int epoch = 0;
-  static pthread_mutex_t lock[2] = {PTHREAD_MUTEX_INITIALIZER, PTHREAD_MUTEX_INITIALIZER};
-  static pthread_cond_t cond[2] = {PTHREAD_COND_INITIALIZER, PTHREAD_COND_INITIALIZER};
-  static T accumulator[2];
-  static int counter[2] = {0, 0};
-
-  pthread_mutex_lock(&lock[epoch]);
-  if (counter[epoch] == 0) {
-    if (average != 0 || args->thread == 0)
-      accumulator[epoch] = *value;
-  } else {
-    switch (average) {
-    case /*r0*/ 0:
-      if (args->thread == 0)
-        accumulator[epoch] = *value;
-      break;
-    case /*avg*/ 1:
-      accumulator[epoch] += *value;
-      break;
-    case /*min*/ 2:
-      accumulator[epoch] = std::min<T>(accumulator[epoch], *value);
-      break;
-    case /*max*/ 3:
-      accumulator[epoch] = std::max<T>(accumulator[epoch], *value);
-      break;
-    case /*sum*/ 4:
-      accumulator[epoch] += *value;
-      break;
-    }
-  }
-
-  if (++counter[epoch] == args->nThreads)
-    pthread_cond_broadcast(&cond[epoch]);
-
-  if (args->thread + 1 == args->nThreads) {
-    while (counter[epoch] != args->nThreads)
-      pthread_cond_wait(&cond[epoch], &lock[epoch]);
+  T accumulator = *value;
 
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
-    if (average != 0) {
-      static_assert(std::is_same<T, long long>::value || std::is_same<T, double>::value,
-                    "Allreduce<T> only for T in {long long, double}");
-      MPI_Datatype ty = std::is_same<T, long long>::value ? MPI_LONG_LONG
-                        : std::is_same<T, double>::value  ? MPI_DOUBLE
-                                                          : MPI_Datatype();
-      MPI_Op op = average == 1   ? MPI_SUM
-                  : average == 2 ? MPI_MIN
-                  : average == 3 ? MPI_MAX
-                  : average == 4 ? MPI_SUM
-                                 : MPI_Op();
-      MPI_Allreduce(MPI_IN_PLACE, (void*)&accumulator[epoch], 1, ty, op, MPI_COMM_WORLD);
-    }
+  if (average != 0) {
+    static_assert(std::is_same<T, long long>::value || std::is_same<T, double>::value,
+                  "Allreduce<T> only for T in {long long, double}");
+    MPI_Datatype ty = std::is_same<T, long long>::value ? MPI_LONG_LONG
+                      : std::is_same<T, double>::value  ? MPI_DOUBLE
+                                                        : MPI_Datatype();
+    MPI_Op op = average == 1   ? MPI_SUM
+                : average == 2 ? MPI_MIN
+                : average == 3 ? MPI_MAX
+                : average == 4 ? MPI_SUM
+                               : MPI_Op();
+    MPI_Allreduce(MPI_IN_PLACE, (void*)&accumulator, 1, ty, op, MPI_COMM_WORLD);
+  }
 #endif
 
-    if (average == 1)
-      accumulator[epoch] /= args->totalProcs * args->nThreads;
-    counter[epoch] = 0;
-    pthread_cond_broadcast(&cond[epoch]);
-  } else {
-    while (counter[epoch] != 0)
-      pthread_cond_wait(&cond[epoch], &lock[epoch]);
-  }
-  pthread_mutex_unlock(&lock[epoch]);
-
-  *value = accumulator[epoch];
-  epoch ^= 1;
+  if (average == 1)
+    accumulator /= args->totalProcs;
+  *value = accumulator;
 }
 
 testResult_t CheckData(struct threadArgs* args, int in_place, int64_t* wrongElts)
@@ -271,7 +245,7 @@ testResult_t BenchTime(struct threadArgs* args, int in_place)
   TESTCHECK(startColl(args, in_place, 0));
   TESTCHECK(completeColl(args));
 
-  Barrier(args);
+  TESTCHECK(Barrier(args));
 
   // Performance Benchmark
   cudaGraph_t graph;
@@ -285,7 +259,7 @@ testResult_t BenchTime(struct threadArgs* args, int in_place)
   CUDACHECK(cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
 
   // Launch the graph
-  Barrier(args);
+  TESTCHECK(Barrier(args));
   tim.reset();
   for (int l = 0; l < cudaGraphLaunches; ++l) {
     CUDACHECK(cudaGraphLaunch(graphExec, args->stream));
@@ -303,7 +277,7 @@ testResult_t BenchTime(struct threadArgs* args, int in_place)
 
   double algBw, busBw;
   args->collTest->getBw(count, 1, deltaSec, &algBw, &busBw, args->totalProcs * args->nThreads * args->nGpus);
-  Barrier(args);
+  TESTCHECK(Barrier(args));
 
   int64_t wrongElts = 0;
   if (datacheck) {
@@ -375,7 +349,7 @@ void setupArgs(size_t size, struct threadArgs* args)
 testResult_t TimeTest(struct threadArgs* args)
 {
   // Sync to avoid first-call timeout
-  Barrier(args);
+  TESTCHECK(Barrier(args));
 
   // Warm-up for large size
   setupArgs(args->maxbytes, args);
diff --git a/tests/common.h b/tests/common.h
index 1fffcbac..7d0e9267 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -9,6 +9,7 @@
 
 #include "mscclpp.h"
 
+#include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 
@@ -121,7 +122,6 @@ struct threadArgs
 typedef testResult_t (*threadFunc_t)(struct threadArgs* args);
 struct testThread
 {
-  pthread_t thread;
   threadFunc_t func;
   struct threadArgs args;
 };
@@ -149,32 +149,6 @@ inline void print_usage(const char* prog)
 #endif
 }
 
-inline void parse_arguments(int argc, const char* argv[], const char** ip_port, int* rank, int* world_size)
-{
-#ifdef MSCCLPP_USE_MPI_FOR_TESTS
-  if (argc != 2 && argc != 4) {
-    print_usage(argv[0]);
-    exit(-1);
-  }
-  *ip_port = argv[1];
-  if (argc == 4) {
-    *rank = atoi(argv[2]);
-    *world_size = atoi(argv[3]);
-  } else {
-    MPI_Comm_rank(MPI_COMM_WORLD, rank);
-    MPI_Comm_size(MPI_COMM_WORLD, world_size);
-  }
-#else
-  if (argc != 4) {
-    print_usage(argv[0]);
-    exit(-1);
-  }
-  *ip_port = argv[1];
-  *rank = atoi(argv[2]);
-  *world_size = atoi(argv[3]);
-#endif
-}
-
 #define PRINT                                                                                                          \
   if (is_main_thread)                                                                                                  \
   printf
diff --git a/tests/timer.cc b/tests/timer.cc
deleted file mode 100644
index 28477ae8..00000000
--- a/tests/timer.cc
+++ /dev/null
@@ -1,32 +0,0 @@
-#include "timer.h"
-
-// Make sure to compile this translation unit with the host compiler and not
-// nvcc, lest you hit an internal compiler error (ICE) with GCC 10.3.0
-#include <chrono>
-
-namespace {
-std::uint64_t now()
-{
-  using clock = std::chrono::steady_clock;
-  return std::chrono::duration_cast<std::chrono::nanoseconds>(clock::now().time_since_epoch()).count();
-}
-} // namespace
-
-timer::timer()
-{
-  t0 = now();
-}
-
-double timer::elapsed() const
-{
-  std::uint64_t t1 = now();
-  return 1.e-9 * (t1 - t0);
-}
-
-double timer::reset()
-{
-  std::uint64_t t1 = now();
-  double ans = 1.e-9 * (t1 - t0);
-  t0 = t1;
-  return ans;
-}
\ No newline at end of file
diff --git a/tests/timer.h b/tests/timer.h
deleted file mode 100644
index c30b264c..00000000
--- a/tests/timer.h
+++ /dev/null
@@ -1,17 +0,0 @@
-#ifndef _408319ecdd5b47b28bf8f511c4fdf816
-#define _408319ecdd5b47b28bf8f511c4fdf816
-
-#include <cstdint>
-
-// Can't include <chrono> because of bug with gcc 10.3.0
-class timer
-{
-  std::uint64_t t0;
-
-public:
-  timer();
-  double elapsed() const;
-  double reset();
-};
-
-#endif
\ No newline at end of file

From e303a532ff7b14830c25519309e85be9c6922dae Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 7 Apr 2023 05:42:21 +0000
Subject: [PATCH 19/23] address comments

---
 tests/allgather_test.cu |  16 +++---
 tests/common.cu         | 110 ++++++++++++++++++++--------------------
 tests/common.h          |  36 +++++++------
 3 files changed, 80 insertions(+), 82 deletions(-)

diff --git a/tests/allgather_test.cu b/tests/allgather_test.cu
index 8354e69a..416f2e66 100644
--- a/tests/allgather_test.cu
+++ b/tests/allgather_test.cu
@@ -147,15 +147,15 @@ void AllGatherGetCollByteCount(size_t* sendcount, size_t* recvcount, size_t* par
   *paramcount = base;
 }
 
-testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
+testResult_t AllGatherInitData(struct testArgs* args, int in_place)
 {
   size_t sendcount = args->sendBytes / sizeof(int);
   size_t recvcount = args->expectedBytes / sizeof(int);
   // int nranks = args->totalProcs;
 
-  CUDACHECK(cudaSetDevice(args->gpus[0]));
+  CUDACHECK(cudaSetDevice(args->gpuNum));
   int rank = args->proc;
-  CUDACHECK(cudaMemset(args->recvbuffs[0], 0, args->expectedBytes));
+  CUDACHECK(cudaMemset(args->recvbuff, 0, args->expectedBytes));
   // void* data = in_place ? ((char*)args->recvbuffs[0]) + rank * args->sendBytes : args->sendbuffs[0];
 
   int* dataHost = new int[recvcount];
@@ -167,11 +167,11 @@ testResult_t AllGatherInitData(struct threadArgs* args, int in_place)
       dataHost[i] = 0;
     }
   }
-  CUDACHECK(cudaMemcpy(args->recvbuffs[0], dataHost, recvcount * sizeof(int), cudaMemcpyHostToDevice));
+  CUDACHECK(cudaMemcpy(args->recvbuff, dataHost, recvcount * sizeof(int), cudaMemcpyHostToDevice));
   for (int i = 0; i < static_cast<int>(recvcount); i++) {
     dataHost[i] = i + 1;
   }
-  CUDACHECK(cudaMemcpy(args->expected[0], dataHost, recvcount * sizeof(int), cudaMemcpyHostToDevice));
+  CUDACHECK(cudaMemcpy(args->expected, dataHost, recvcount * sizeof(int), cudaMemcpyHostToDevice));
   delete dataHost;
   CUDACHECK(cudaDeviceSynchronize());
   return testSuccess;
@@ -187,10 +187,10 @@ void AllGatherGetBw(size_t count, int typesize, double sec, double* algBw, doubl
 }
 
 testResult_t AllGatherRunColl(void* sendbuff, void* recvbuff, int nranksPerNode, size_t count, mscclppComm_t comm,
-                              cudaStream_t stream)
+                              cudaStream_t stream, int kernel_num)
 {
   int worldSize = comm->nRanks;
-  kernel<<<1, 32 * (worldSize - 1), 0, stream>>>(comm->rank, worldSize, nranksPerNode, count / sizeof(int), 1);
+  kernel<<<1, 32 * (worldSize - 1), 0, stream>>>(comm->rank, worldSize, nranksPerNode, count / sizeof(int), kernel_num);
   return testSuccess;
 }
 
@@ -203,7 +203,7 @@ void AllGatherGetBuffSize(size_t* sendcount, size_t* recvcount, size_t count, in
   AllGatherGetCollByteCount(sendcount, recvcount, &paramcount, &sendInplaceOffset, &recvInplaceOffset, count, nranks);
 }
 
-testResult_t AllGatherRunTest(struct threadArgs* args)
+testResult_t AllGatherRunTest(struct testArgs* args)
 {
   args->collTest = &allGatherTest;
   mscclppDevConn_t* devConns;
diff --git a/tests/common.cu b/tests/common.cu
index 6722f57c..20baac0f 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -53,6 +53,7 @@ int report_cputime = 0;
 // Report average iteration time: (0=RANK0,1=AVG,2=MIN,3=MAX)
 int average = 1;
 std::string ip_port;
+int kernel_num = 0;
 int cudaGraphLaunches = 15;
 
 double parsesize(const char* value)
@@ -92,7 +93,7 @@ double parsesize(const char* value)
   return size * units;
 }
 
-inline testResult_t Barrier(struct threadArgs* args)
+inline testResult_t Barrier(struct testArgs* args)
 {
   int tmp[16];
   // A simple barrier
@@ -130,7 +131,7 @@ testResult_t AllocateBuffs(void** sendbuff, size_t sendBytes, void** recvbuff, s
   return testSuccess;
 }
 
-testResult_t startColl(struct threadArgs* args, int in_place, int iter)
+testResult_t startColl(struct testArgs* args, int in_place, int iter)
 {
   size_t count = args->nbytes;
 
@@ -139,13 +140,13 @@ testResult_t startColl(struct threadArgs* args, int in_place, int iter)
   size_t steps = totalnbytes ? args->maxbytes / totalnbytes : 1;
   size_t shift = totalnbytes * (iter % steps);
 
-  int rank = ((args->proc * args->nThreads + args->thread) * args->nGpus);
-  char* recvBuff = ((char*)args->recvbuffs[0]) + shift;
-  char* sendBuff = ((char*)args->sendbuffs[0]) + shift;
+  int rank = args->proc;
+  char* recvBuff = ((char*)args->recvbuff) + shift;
+  char* sendBuff = ((char*)args->sendbuff) + shift;
 
   TESTCHECK(args->collTest->runColl((void*)(in_place ? recvBuff + args->sendInplaceOffset * rank : sendBuff),
                                     (void*)(in_place ? recvBuff + args->recvInplaceOffset * rank : recvBuff),
-                                    args->nranksPerNode, count, args->comm, args->stream));
+                                    args->nranksPerNode, count, args->comm, args->stream, args->kernel_num));
   return testSuccess;
 }
 
@@ -177,7 +178,7 @@ testResult_t testStreamSynchronize(cudaStream_t stream)
   return testSuccess;
 }
 
-testResult_t completeColl(struct threadArgs* args)
+testResult_t completeColl(struct testArgs* args)
 {
   TESTCHECK(testStreamSynchronize(args->stream));
   return testSuccess;
@@ -187,7 +188,7 @@ testResult_t completeColl(struct threadArgs* args)
 // for average=0 is just value itself.
 // Inter process barrier+allreduce. The quality of the return value
 // for average=0 is just value itself.
-template <typename T> void Allreduce(struct threadArgs* args, T* value, int average)
+template <typename T> void Allreduce(struct testArgs* args, T* value, int average)
 {
   T accumulator = *value;
 
@@ -212,7 +213,7 @@ template <typename T> void Allreduce(struct threadArgs* args, T* value, int aver
   *value = accumulator;
 }
 
-testResult_t CheckData(struct threadArgs* args, int in_place, int64_t* wrongElts)
+testResult_t CheckData(struct testArgs* args, int in_place, int64_t* wrongElts)
 {
   if (in_place == 0) {
     return testInternalError;
@@ -221,8 +222,8 @@ testResult_t CheckData(struct threadArgs* args, int in_place, int64_t* wrongElts
 
   int* dataHostRecv = new int[count];
   int* dataHostExpected = new int[count];
-  CUDACHECK(cudaMemcpy(dataHostRecv, args->recvbuffs[0], args->expectedBytes, cudaMemcpyDeviceToHost));
-  CUDACHECK(cudaMemcpy(dataHostExpected, args->expected[0], args->expectedBytes, cudaMemcpyDeviceToHost));
+  CUDACHECK(cudaMemcpy(dataHostRecv, args->recvbuff, args->expectedBytes, cudaMemcpyDeviceToHost));
+  CUDACHECK(cudaMemcpy(dataHostExpected, args->expected, args->expectedBytes, cudaMemcpyDeviceToHost));
 
   for (size_t i = 0; i < count; i++) {
     if (dataHostRecv[i] != dataHostExpected[i]) {
@@ -231,12 +232,12 @@ testResult_t CheckData(struct threadArgs* args, int in_place, int64_t* wrongElts
   }
 
   if (args->reportErrors && *wrongElts) {
-    (*args->errors)++;
+    (args->error)++;
   }
   return testSuccess;
 }
 
-testResult_t BenchTime(struct threadArgs* args, int in_place)
+testResult_t BenchTime(struct testArgs* args, int in_place)
 {
   size_t count = args->nbytes;
 
@@ -276,7 +277,7 @@ testResult_t BenchTime(struct threadArgs* args, int in_place)
   CUDACHECK(cudaGraphDestroy(graph));
 
   double algBw, busBw;
-  args->collTest->getBw(count, 1, deltaSec, &algBw, &busBw, args->totalProcs * args->nThreads * args->nGpus);
+  args->collTest->getBw(count, 1, deltaSec, &algBw, &busBw, args->totalProcs);
   TESTCHECK(Barrier(args));
 
   int64_t wrongElts = 0;
@@ -323,14 +324,14 @@ testResult_t BenchTime(struct threadArgs* args, int in_place)
     PRINT("  %7s  %6.2f  %6.2f  %5s", timeStr, algBw, busBw, "N/A");
   }
 
-  args->bw[0] += busBw;
-  args->bw_count[0]++;
+  args->bw += busBw;
+  args->bw_count++;
   return testSuccess;
 }
 
-void setupArgs(size_t size, struct threadArgs* args)
+void setupArgs(size_t size, struct testArgs* args)
 {
-  int nranks = args->totalProcs * args->nGpus * args->nThreads;
+  int nranks = args->totalProcs;
   size_t count, sendCount, recvCount, paramCount, sendInplaceOffset, recvInplaceOffset;
 
   // TODO: support more data types
@@ -346,7 +347,7 @@ void setupArgs(size_t size, struct threadArgs* args)
   args->recvInplaceOffset = recvInplaceOffset * typeSize;
 }
 
-testResult_t TimeTest(struct threadArgs* args)
+testResult_t TimeTest(struct testArgs* args)
 {
   // Sync to avoid first-call timeout
   TESTCHECK(Barrier(args));
@@ -412,10 +413,10 @@ testResult_t setupMscclppConnections(int rank, int worldSize, int ranksPerNode,
   return testSuccess;
 }
 
-testResult_t threadRunTests(struct threadArgs* args)
+testResult_t runTests(struct testArgs* args)
 {
   PRINT("# Setting up the connection in MSCCL++\n");
-  TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->nranksPerNode, args->comm, args->recvbuffs[0],
+  TESTCHECK(setupMscclppConnections(args->proc, args->totalProcs, args->nranksPerNode, args->comm, args->recvbuff,
                                     args->maxbytes));
   PRINT("# Launching MSCCL++ proxy threads\n");
   MSCCLPPCHECK(mscclppProxyLaunch(args->comm));
@@ -447,12 +448,13 @@ int main(int argc, char* argv[])
                                      {"report_cputime", required_argument, 0, 'C'},
                                      {"average", required_argument, 0, 'a'},
                                      {"ip_port", required_argument, 0, 'P'},
+                                     {"kernel_num", required_argument, 0, 'k'},
                                      {"help", no_argument, 0, 'h'},
                                      {}};
 
   while (1) {
     int c;
-    c = getopt_long(argc, argv, "b:e:i:f:n:w:c:T:G:C:a:P:h:", longopts, &longindex);
+    c = getopt_long(argc, argv, "b:e:i:f:n:w:c:T:G:C:a:P:k:h:", longopts, &longindex);
 
     if (c == -1)
       break;
@@ -508,6 +510,9 @@ int main(int argc, char* argv[])
     case 'P':
       ip_port = optarg;
       break;
+    case 'k':
+      kernel_num = (int)strtol(optarg, NULL, 0);
+      break;
     case 'h':
     default:
       if (c != 'h')
@@ -525,6 +530,7 @@ int main(int argc, char* argv[])
              "[-C,--report_cputime <0/1>] \n\t"
              "[-a,--average <0/1/2/3> report average iteration time <0=RANK0/1=AVG/2=MIN/3=MAX>] \n\t"
              "[-P,--ip_port <ip port for bootstrap>] \n\t"
+             "[-k,--kernel_num <kernel number of commnication primitive>] \n\t"
              "[-h,--help]\n",
              basename(argv[0]));
       return 0;
@@ -565,9 +571,10 @@ testResult_t run()
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
 
-  PRINT("# minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip port: %s graph: %d\n",
+  PRINT("# minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip port: %s graph: %d, "
+        "kernel num: %d\n",
         minBytes, maxBytes, (stepFactor > 1) ? stepFactor : stepBytes, (stepFactor > 1) ? "factor" : "bytes",
-        warmup_iters, iters, datacheck, ip_port.c_str(), cudaGraphLaunches);
+        warmup_iters, iters, datacheck, ip_port.c_str(), cudaGraphLaunches, kernel_num);
   PRINT("#\n");
   PRINT("# Using devices\n");
 
@@ -607,7 +614,6 @@ testResult_t run()
       printf("#\n# Reducing maxBytes to %ld due to memory limitation\n", maxBytes);
   }
 
-  int gpu = cudaDev;
   cudaStream_t stream;
   void* sendbuff;
   void* recvbuff;
@@ -616,7 +622,7 @@ testResult_t run()
 
   mscclppTestEngine.getBuffSize(&sendBytes, &recvBytes, (size_t)maxBytes, (size_t)totalProcs);
 
-  CUDACHECK(cudaSetDevice(gpu));
+  CUDACHECK(cudaSetDevice(cudaDev));
   TESTCHECK(AllocateBuffs(&sendbuff, sendBytes, &recvbuff, recvBytes, &expected, (size_t)maxBytes));
   CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
   PRINT("#\n");
@@ -625,43 +631,38 @@ testResult_t run()
   mscclppComm_t comm;
   MSCCLPPCHECK(mscclppCommInitRank(&comm, totalProcs, ip_port.c_str(), rank));
 
-  int error = 0;
-  double bw = 0.0;
   double* delta;
   CUDACHECK(cudaHostAlloc(&delta, sizeof(double) * NUM_BLOCKS, cudaHostAllocPortable | cudaHostAllocMapped));
-  int bw_count = 0;
 
   fflush(stdout);
 
-  struct testThread thread;
+  struct testWorker worker;
 
-  thread.args.minbytes = minBytes;
-  thread.args.maxbytes = maxBytes;
-  thread.args.stepbytes = stepBytes;
-  thread.args.stepfactor = stepFactor;
-  thread.args.localRank = localRank;
-  thread.args.nranksPerNode = nranksPerNode;
+  worker.args.minbytes = minBytes;
+  worker.args.maxbytes = maxBytes;
+  worker.args.stepbytes = stepBytes;
+  worker.args.stepfactor = stepFactor;
+  worker.args.localRank = localRank;
+  worker.args.nranksPerNode = nranksPerNode;
 
-  thread.args.totalProcs = totalProcs;
-  thread.args.proc = proc;
-  thread.args.nThreads = 1;
-  thread.args.thread = 0;
-  thread.args.nGpus = 1;
-  thread.args.gpus = &gpu;
-  thread.args.sendbuffs = &sendbuff;
-  thread.args.recvbuffs = &recvbuff;
-  thread.args.expected = &expected;
-  thread.args.comm = comm;
-  thread.args.stream = stream;
+  worker.args.totalProcs = totalProcs;
+  worker.args.proc = proc;
+  worker.args.gpuNum = cudaDev;
+  worker.args.kernel_num = kernel_num;
+  worker.args.sendbuff = sendbuff;
+  worker.args.recvbuff = recvbuff;
+  worker.args.expected = expected;
+  worker.args.comm = comm;
+  worker.args.stream = stream;
 
-  thread.args.errors = &error;
-  thread.args.bw = &bw;
-  thread.args.bw_count = &bw_count;
+  worker.args.error = 0;
+  worker.args.bw = 0.0;
+  worker.args.bw_count = 0;
 
-  thread.args.reportErrors = datacheck;
+  worker.args.reportErrors = datacheck;
 
-  thread.func = threadRunTests;
-  TESTCHECK(thread.func(&thread.args));
+  worker.func = runTests;
+  TESTCHECK(worker.func(&worker.args));
 
   MSCCLPPCHECK(mscclppCommDestroy(comm));
 
@@ -674,8 +675,7 @@ testResult_t run()
     CUDACHECK(cudaFree(expected));
   CUDACHECK(cudaFreeHost(delta));
 
-  bw /= bw_count;
-
+  int error = worker.args.error;
   PRINT("# Out of bounds values : %d %s\n", error, error ? "FAILED" : "OK");
   PRINT("#\n");
 
diff --git a/tests/common.h b/tests/common.h
index 7d0e9267..39766de0 100644
--- a/tests/common.h
+++ b/tests/common.h
@@ -69,22 +69,22 @@ struct testColl
   const char name[20];
   void (*getCollByteCount)(size_t* sendcount, size_t* recvcount, size_t* paramcount, size_t* sendInplaceOffset,
                            size_t* recvInplaceOffset, size_t count, int nranks);
-  testResult_t (*initData)(struct threadArgs* args, int in_place);
+  testResult_t (*initData)(struct testArgs* args, int in_place);
   void (*getBw)(size_t count, int typesize, double sec, double* algBw, double* busBw, int nranks);
   testResult_t (*runColl)(void* sendbuff, void* recvbuff, int nranksPerNode, size_t count, mscclppComm_t comm,
-                          cudaStream_t stream);
+                          cudaStream_t stream, int kernel_num);
 };
 
 struct testEngine
 {
   void (*getBuffSize)(size_t* sendcount, size_t* recvcount, size_t count, int nranks);
   // We can add more parameters for other communication primitives
-  testResult_t (*runTest)(struct threadArgs* args);
+  testResult_t (*runTest)(struct testArgs* args);
 };
 
 extern struct testEngine mscclppTestEngine;
 
-struct threadArgs
+struct testArgs
 {
   size_t nbytes;
   size_t minbytes;
@@ -94,40 +94,38 @@ struct threadArgs
 
   int totalProcs;
   int proc;
-  int nThreads;
-  int thread;
-  int nGpus;
-  int* gpus;
+  int gpuNum;
   int localRank;
   int nranksPerNode;
-  void** sendbuffs;
+  int kernel_num;
+  void* sendbuff;
   size_t sendBytes;
   size_t sendInplaceOffset;
-  void** recvbuffs;
+  void* recvbuff;
   size_t recvInplaceOffset;
   mscclppComm_t comm;
   cudaStream_t stream;
 
-  void** expected;
+  void* expected;
   size_t expectedBytes;
-  int* errors;
-  double* bw;
-  int* bw_count;
+  int error;
+  double bw;
+  int bw_count;
 
   int reportErrors;
 
   struct testColl* collTest;
 };
 
-typedef testResult_t (*threadFunc_t)(struct threadArgs* args);
-struct testThread
+typedef testResult_t (*entryFunc_t)(struct testArgs* args);
+struct testWorker
 {
-  threadFunc_t func;
-  struct threadArgs args;
+  entryFunc_t func;
+  struct testArgs args;
 };
 
 // Provided by common.cu
-extern testResult_t TimeTest(struct threadArgs* args);
+extern testResult_t TimeTest(struct testArgs* args);
 
 static void getHostName(char* hostname, int maxlen)
 {

From ec83a27e834a2239999bfba0691b0c0e5644cb35 Mon Sep 17 00:00:00 2001
From: Saeed Maleki <saemal@microsoft.com>
Date: Sat, 8 Apr 2023 01:57:22 +0000
Subject: [PATCH 20/23] wip

---
 src/bootstrap/socket.cc | 49 ++++++++++++++++++++++++-----------------
 src/include/alloc.h     |  2 --
 src/include/utils.h     | 10 ++++-----
 src/utils.cc            | 10 +++++++++
 tests/allgather_test.cu |  1 +
 5 files changed, 44 insertions(+), 28 deletions(-)

diff --git a/src/bootstrap/socket.cc b/src/bootstrap/socket.cc
index 487f38cb..cfa795b3 100644
--- a/src/bootstrap/socket.cc
+++ b/src/bootstrap/socket.cc
@@ -452,9 +452,12 @@ mscclppResult_t mscclppSocketGetAddr(struct mscclppSocket* sock, union mscclppSo
 
 static mscclppResult_t socketTryAccept(struct mscclppSocket* sock)
 {
-  static time_t initTime = -1;
-  if (initTime == -1)
-    initTime = clockNano() / 1e9;
+  static bool timeInitialized = false;
+  static mscclppTime_t initTime;
+  if (!timeInitialized){
+    timeInitialized = true;
+    initTime = getClock();
+  }
 
   mscclppConfig* config = mscclppConfig::getInstance();
   time_t acceptTimeout = config->getBootstrapConnectionTimeoutConfig();
@@ -462,14 +465,14 @@ static mscclppResult_t socketTryAccept(struct mscclppSocket* sock)
   sock->fd = accept(sock->acceptFd, &sock->addr.sa, &socklen);
   if (sock->fd != -1) {
     sock->state = mscclppSocketStateAccepted;
-    initTime = -1;
+    timeInitialized = false;
   } else if (errno != EAGAIN && errno != EWOULDBLOCK) {
     WARN("socketTryAccept: get errno %d that is not EAGAIN or EWOULDBLOCK", errno);
-    initTime = -1;
+    timeInitialized = false;
     return mscclppSystemError;
-  } else if ((clockNano() / 1e9) - initTime > acceptTimeout) {
+  } else if (elapsedClock(getClock(), initTime) > acceptTimeout) {
     WARN("socketTryAccept: exceeded timeout (%ld) sec", acceptTimeout);
-    initTime = -1;
+    timeInitialized = false;
     return mscclppRemoteError;
   } else {
     usleep(SLEEP_INT);
@@ -513,10 +516,13 @@ static mscclppResult_t socketFinalizeAccept(struct mscclppSocket* sock)
 
 static mscclppResult_t socketStartConnect(struct mscclppSocket* sock)
 {
-  static time_t initTime = -1;
-  if (initTime == -1) {
-    initTime = clockNano() / 1e9;
+  static bool timeInitialized = false;
+  static mscclppTime_t initTime;
+  if (!timeInitialized){
+    timeInitialized = true;
+    initTime = getClock();
   }
+
   mscclppConfig* config = mscclppConfig::getInstance();
   time_t acceptTimeout = config->getBootstrapConnectionTimeoutConfig();
 
@@ -524,16 +530,16 @@ static mscclppResult_t socketStartConnect(struct mscclppSocket* sock)
   int ret = connect(sock->fd, &sock->addr.sa, sock->salen);
   if (ret == 0) {
     sock->state = mscclppSocketStateConnected;
-    initTime = -1;
+    timeInitialized = false;
     return mscclppSuccess;
   } else if (errno == EINPROGRESS) {
     sock->state = mscclppSocketStateConnectPolling;
     return mscclppSuccess;
   } else if (errno == ECONNREFUSED || errno == ETIMEDOUT) {
-    if ((clockNano() / 1e9) - initTime > acceptTimeout) {
+    if (elapsedClock(getClock(), initTime) > acceptTimeout) {
       WARN("socketStartConnect: exceeded timeout (%ld) sec", acceptTimeout);
       sock->state = mscclppSocketStateError;
-      initTime = -1;
+      timeInitialized = false;
       return mscclppRemoteError;
     }
     usleep(SLEEP_INT);
@@ -544,17 +550,20 @@ static mscclppResult_t socketStartConnect(struct mscclppSocket* sock)
     char line[SOCKET_NAME_MAXLEN + 1];
     sock->state = mscclppSocketStateError;
     WARN("socketStartConnect: Connect to %s failed : %s", mscclppSocketToString(&sock->addr, line), strerror(errno));
-    initTime = -1;
+    timeInitialized = false;
     return mscclppSystemError;
   }
 }
 
 static mscclppResult_t socketPollConnect(struct mscclppSocket* sock)
 {
-  static time_t initTime = -1;
-  if (initTime == -1) {
-    initTime = clockNano() / 1e9;
+  static bool timeInitialized = false;
+  static mscclppTime_t initTime;
+  if (!timeInitialized){
+    timeInitialized = true;
+    initTime = getClock();
   }
+
   mscclppConfig* config = mscclppConfig::getInstance();
   time_t acceptTimeout = config->getBootstrapConnectionTimeoutConfig();
 
@@ -567,7 +576,7 @@ static mscclppResult_t socketPollConnect(struct mscclppSocket* sock)
   pfd.events = POLLOUT;
   SYSCHECK(ret = poll(&pfd, 1, timeout), "poll");
   if (ret == 0) {
-    initTime = -1;
+    timeInitialized = false;
     return mscclppSuccess;
   }
 
@@ -576,10 +585,10 @@ static mscclppResult_t socketPollConnect(struct mscclppSocket* sock)
   SYSCHECK(getsockopt(sock->fd, SOL_SOCKET, SO_ERROR, (void*)&ret, &rlen), "getsockopt");
 
   if (ret == 0) {
-    initTime = -1;
+    timeInitialized = false;
     sock->state = mscclppSocketStateConnected;
   } else if (ret == ECONNREFUSED || ret == ETIMEDOUT) {
-    if ((clockNano() / 1e9) - initTime > acceptTimeout) {
+    if (elapsedClock(getClock(), initTime) > acceptTimeout) {
       WARN("socketPollConnect: exceeded timeout (%ld) sec", acceptTimeout);
       sock->state = mscclppSocketStateError;
       return mscclppRemoteError;
diff --git a/src/include/alloc.h b/src/include/alloc.h
index 65e953e1..496af197 100644
--- a/src/include/alloc.h
+++ b/src/include/alloc.h
@@ -16,8 +16,6 @@
 #include <sys/mman.h>
 #include <unistd.h>
 
-uint64_t clockNano(); // from utils.h with which we have a circular dependency
-
 template <typename T> mscclppResult_t mscclppCudaHostCallocDebug(T** ptr, size_t nelem, const char* filefunc, int line)
 {
   mscclppResult_t result = mscclppSuccess;
diff --git a/src/include/utils.h b/src/include/utils.h
index f3eef1a9..bcf50d9b 100644
--- a/src/include/utils.h
+++ b/src/include/utils.h
@@ -16,6 +16,7 @@
 #include <stdint.h>
 #include <time.h>
 
+
 // int mscclppCudaCompCap();
 
 // PCI Bus ID <-> int64 conversion functions
@@ -48,12 +49,9 @@ static long log2i(long n)
   return l;
 }
 
-inline uint64_t clockNano()
-{
-  struct timespec ts;
-  clock_gettime(CLOCK_MONOTONIC, &ts);
-  return uint64_t(ts.tv_sec) * 1000 * 1000 * 1000 + ts.tv_nsec;
-}
+typedef std::chrono::steady_clock::time_point mscclppTime_t;
+mscclppTime_t getClock();
+int64_t elapsedClock(mscclppTime_t start, mscclppTime_t end);
 
 /* get any bytes of random data from /dev/urandom, return 0 if it succeeds; else
  * return -1 */
diff --git a/src/utils.cc b/src/utils.cc
index 9bb01284..f0a73303 100644
--- a/src/utils.cc
+++ b/src/utils.cc
@@ -270,3 +270,13 @@ mscclppResult_t setNumaState(mscclppNumaState state)
   numa_bind(state);
   return mscclppSuccess;
 }
+
+inline mscclppTime_t getClock()
+{
+  return std::chrono::steady_clock::now();
+}
+
+inline int64_t elapsedClock(mscclppTime_t start, mscclppTime_t end)
+{
+  return std::chrono::duration_cast<std::chrono::seconds>(end - start).count();
+}
diff --git a/tests/allgather_test.cu b/tests/allgather_test.cu
index 416f2e66..647debef 100644
--- a/tests/allgather_test.cu
+++ b/tests/allgather_test.cu
@@ -174,6 +174,7 @@ testResult_t AllGatherInitData(struct testArgs* args, int in_place)
   CUDACHECK(cudaMemcpy(args->expected, dataHost, recvcount * sizeof(int), cudaMemcpyHostToDevice));
   delete dataHost;
   CUDACHECK(cudaDeviceSynchronize());
+  MSCCLPPCHECK(mscclppBootstrapBarrier(args->comm));
   return testSuccess;
 }
 

From 7e091d3806345427948922236666ade78a41d1b4 Mon Sep 17 00:00:00 2001
From: Saeed Maleki <saemal@microsoft.com>
Date: Sat, 8 Apr 2023 06:22:12 +0000
Subject: [PATCH 21/23] all good now

---
 src/utils.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/utils.cc b/src/utils.cc
index f0a73303..c0766765 100644
--- a/src/utils.cc
+++ b/src/utils.cc
@@ -271,12 +271,12 @@ mscclppResult_t setNumaState(mscclppNumaState state)
   return mscclppSuccess;
 }
 
-inline mscclppTime_t getClock()
+mscclppTime_t getClock()
 {
   return std::chrono::steady_clock::now();
 }
 
-inline int64_t elapsedClock(mscclppTime_t start, mscclppTime_t end)
+int64_t elapsedClock(mscclppTime_t start, mscclppTime_t end)
 {
   return std::chrono::duration_cast<std::chrono::seconds>(end - start).count();
 }

From 5ca72120652156af4d6b4200a87e042040557770 Mon Sep 17 00:00:00 2001
From: Saeed Maleki <saemal@microsoft.com>
Date: Sat, 8 Apr 2023 06:29:54 +0000
Subject: [PATCH 22/23] removed ip_port

---
 tests/common.cu | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/tests/common.cu b/tests/common.cu
index 20baac0f..69b3f284 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -52,7 +52,6 @@ int timeout = 0;
 int report_cputime = 0;
 // Report average iteration time: (0=RANK0,1=AVG,2=MIN,3=MAX)
 int average = 1;
-std::string ip_port;
 int kernel_num = 0;
 int cudaGraphLaunches = 15;
 
@@ -447,7 +446,6 @@ int main(int argc, char* argv[])
                                      {"cudagraph", required_argument, 0, 'G'},
                                      {"report_cputime", required_argument, 0, 'C'},
                                      {"average", required_argument, 0, 'a'},
-                                     {"ip_port", required_argument, 0, 'P'},
                                      {"kernel_num", required_argument, 0, 'k'},
                                      {"help", no_argument, 0, 'h'},
                                      {}};
@@ -507,9 +505,6 @@ int main(int argc, char* argv[])
     case 'a':
       average = (int)strtol(optarg, NULL, 0);
       break;
-    case 'P':
-      ip_port = optarg;
-      break;
     case 'k':
       kernel_num = (int)strtol(optarg, NULL, 0);
       break;
@@ -529,7 +524,6 @@ int main(int argc, char* argv[])
              "[-G,--cudagraph <num graph launches>] \n\t"
              "[-C,--report_cputime <0/1>] \n\t"
              "[-a,--average <0/1/2/3> report average iteration time <0=RANK0/1=AVG/2=MIN/3=MAX>] \n\t"
-             "[-P,--ip_port <ip port for bootstrap>] \n\t"
              "[-k,--kernel_num <kernel number of commnication primitive>] \n\t"
              "[-h,--help]\n",
              basename(argv[0]));
@@ -541,10 +535,6 @@ int main(int argc, char* argv[])
             (unsigned long long)maxBytes);
     return -1;
   }
-  if (ip_port.empty()) {
-    fprintf(stderr, "--ip_port is required'\n");
-    return -1;
-  }
 #ifdef MSCCLPP_USE_MPI_FOR_TESTS
   MPI_Init(&argc, &argv);
 #endif
@@ -571,10 +561,10 @@ testResult_t run()
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
   is_main_thread = is_main_proc = (proc == 0) ? 1 : 0;
 
-  PRINT("# minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d ip port: %s graph: %d, "
+  PRINT("# minBytes %ld maxBytes %ld step: %ld(%s) warmup iters: %d iters: %d validation: %d graph: %d, "
         "kernel num: %d\n",
         minBytes, maxBytes, (stepFactor > 1) ? stepFactor : stepBytes, (stepFactor > 1) ? "factor" : "bytes",
-        warmup_iters, iters, datacheck, ip_port.c_str(), cudaGraphLaunches, kernel_num);
+        warmup_iters, iters, datacheck, cudaGraphLaunches, kernel_num);
   PRINT("#\n");
   PRINT("# Using devices\n");
 
@@ -628,8 +618,11 @@ testResult_t run()
   PRINT("#\n");
   PRINT("# Initializing MSCCL++\n");
 
+  mscclppUniqueId mscclppId;
+  if (proc == 0) MSCCLPPCHECK(mscclppGetUniqueId(&mscclppId));
+  MPI_Bcast((void*)&mscclppId, sizeof(mscclppId), MPI_BYTE, 0, MPI_COMM_WORLD);
   mscclppComm_t comm;
-  MSCCLPPCHECK(mscclppCommInitRank(&comm, totalProcs, ip_port.c_str(), rank));
+  MSCCLPPCHECK(mscclppCommInitRankFromId(&comm, totalProcs, mscclppId, rank));
 
   double* delta;
   CUDACHECK(cudaHostAlloc(&delta, sizeof(double) * NUM_BLOCKS, cudaHostAllocPortable | cudaHostAllocMapped));

From f3f53a41486570aab2669f8f946a0036c045b7f7 Mon Sep 17 00:00:00 2001
From: Saeed Maleki <saemal@microsoft.com>
Date: Sat, 8 Apr 2023 06:32:57 +0000
Subject: [PATCH 23/23] lint

---
 src/bootstrap/socket.cc | 6 +++---
 src/include/utils.h     | 1 -
 tests/common.cu         | 3 ++-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/bootstrap/socket.cc b/src/bootstrap/socket.cc
index cfa795b3..b3998d91 100644
--- a/src/bootstrap/socket.cc
+++ b/src/bootstrap/socket.cc
@@ -454,7 +454,7 @@ static mscclppResult_t socketTryAccept(struct mscclppSocket* sock)
 {
   static bool timeInitialized = false;
   static mscclppTime_t initTime;
-  if (!timeInitialized){
+  if (!timeInitialized) {
     timeInitialized = true;
     initTime = getClock();
   }
@@ -518,7 +518,7 @@ static mscclppResult_t socketStartConnect(struct mscclppSocket* sock)
 {
   static bool timeInitialized = false;
   static mscclppTime_t initTime;
-  if (!timeInitialized){
+  if (!timeInitialized) {
     timeInitialized = true;
     initTime = getClock();
   }
@@ -559,7 +559,7 @@ static mscclppResult_t socketPollConnect(struct mscclppSocket* sock)
 {
   static bool timeInitialized = false;
   static mscclppTime_t initTime;
-  if (!timeInitialized){
+  if (!timeInitialized) {
     timeInitialized = true;
     initTime = getClock();
   }
diff --git a/src/include/utils.h b/src/include/utils.h
index bcf50d9b..3eff9842 100644
--- a/src/include/utils.h
+++ b/src/include/utils.h
@@ -16,7 +16,6 @@
 #include <stdint.h>
 #include <time.h>
 
-
 // int mscclppCudaCompCap();
 
 // PCI Bus ID <-> int64 conversion functions
diff --git a/tests/common.cu b/tests/common.cu
index 69b3f284..77c3223f 100644
--- a/tests/common.cu
+++ b/tests/common.cu
@@ -619,7 +619,8 @@ testResult_t run()
   PRINT("# Initializing MSCCL++\n");
 
   mscclppUniqueId mscclppId;
-  if (proc == 0) MSCCLPPCHECK(mscclppGetUniqueId(&mscclppId));
+  if (proc == 0)
+    MSCCLPPCHECK(mscclppGetUniqueId(&mscclppId));
   MPI_Bcast((void*)&mscclppId, sizeof(mscclppId), MPI_BYTE, 0, MPI_COMM_WORLD);
   mscclppComm_t comm;
   MSCCLPPCHECK(mscclppCommInitRankFromId(&comm, totalProcs, mscclppId, rank));