stress test for the infiniband network

mpi-infiniband-stress.c

#include "stdio.h"
#include "stdlib.h"
#include "mpi.h"

#define ITERATION_COUNT 8192
#define BUFFER_SIZE 268435456

int main(int argc, char **argv) {
  MPI_Init(&argc, &argv);

  int rankCount, myRank;
  MPI_Comm_rank(MPI_COMM_WORLD, &myRank);
  MPI_Comm_size(MPI_COMM_WORLD, &rankCount);

  unsigned *refBuf = malloc(BUFFER_SIZE * sizeof(unsigned));
  unsigned *rxBuf = malloc(BUFFER_SIZE * sizeof(unsigned));
  unsigned *txBuf = malloc(BUFFER_SIZE * sizeof(unsigned));

  srand(myRank);
  for(int i = 0; i < BUFFER_SIZE; i++) {
    txBuf[i] = refBuf[i] = rand();
  }

  int txRank = (myRank + 1) % rankCount;
  int rxRank = myRank - 1;
  if(rxRank == -1) rxRank = rankCount - 1;

  for(int i = 0; i < ITERATION_COUNT; i++) {
    double startTime = MPI_Wtime();

    MPI_Request sendRequest; MPI_Status sendStatus, recvStatus;
    MPI_Isend(txBuf, BUFFER_SIZE, MPI_UNSIGNED, txRank, i, MPI_COMM_WORLD, &sendRequest);
    MPI_Recv(rxBuf, BUFFER_SIZE, MPI_UNSIGNED, rxRank, i, MPI_COMM_WORLD, &recvStatus);
    MPI_Wait(&sendRequest, &sendStatus);

    unsigned *tmp = txBuf; txBuf = rxBuf; rxBuf = tmp;

    double time = MPI_Wtime() - startTime;
    unsigned speed = BUFFER_SIZE * 2 * sizeof(unsigned) / ((1024 * 1024) * time);
    printf("Rank %d, loop %d, net bandwidth: %dMB/s\n", myRank, i, speed);
  }

  MPI_Finalize();
  return 0;
}