RankToAccelerator

topo.cpp

#include <iostream>

#define GPOLICY_INTERLEAVE 0
#define GPOLICY_COMPACT 1

static int map_compact(const unsigned rank_id, const unsigned num_ranks, const unsigned num_gpus,
                unsigned *gpu_id, unsigned *bin_size, unsigned *bin_id) {

  const unsigned small_bin_size = num_ranks / num_gpus;
  const unsigned big_bin_size = small_bin_size + 1;
  const unsigned big_bin_count = num_ranks % num_gpus;  

  if (rank_id < big_bin_size * big_bin_count) {
    *gpu_id = rank_id / big_bin_size;
    *bin_size = big_bin_size;
    *bin_id = rank_id % big_bin_size;
  } else {
    *gpu_id = (rank_id - big_bin_count) / small_bin_size;
    *bin_size = small_bin_size;
    *bin_id = (rank_id - big_bin_size * big_bin_count) % small_bin_size;
  }
  return 0;
}

static int map_interleave(const unsigned rank_id, const unsigned num_ranks, const unsigned num_gpus,
                unsigned *gpu_id, unsigned *bin_size, unsigned *bin_id) {

  *gpu_id = rank_id % num_gpus;
  *bin_id = rank_id / num_gpus;

  const unsigned small_bin_size = num_ranks / num_gpus;
  const unsigned big_bin_count = num_ranks % num_gpus;

  *bin_size = small_bin_size + (*gpu_id < big_bin_count);
  return 0;
}

int rankToAccelerator(unsigned rank_id, unsigned num_ranks, unsigned num_hierarchy, 
                      const unsigned * num_gpus, const unsigned *mode, 
                      unsigned *gpu_id) {

    if (rank_id >= num_ranks) return -1;
    if (num_hierarchy == 0 || num_hierarchy >= 3) return - 1;
    for (int h=0; h < num_hierarchy; h++) {
        if (num_gpus[h] == 0) return -1;
    }

    unsigned bin_size = num_ranks;
    unsigned bin_id = rank_id;

    for (int h=0; h < num_hierarchy; h++) {
        switch(mode[h]){
            case GPOLICY_INTERLEAVE :
                map_interleave(bin_id, bin_size, num_gpus[h], &gpu_id[h], &bin_size, &bin_id);
                break;
            case GPOLICY_COMPACT :
                map_compact(bin_id, bin_size, num_gpus[h], &gpu_id[h], &bin_size, &bin_id);
                break;
        }
    }

    return 0;

}

void test_compact_compact() {
    int num_ranks = 37;
    unsigned num_gpus[2] = {5,3};
    unsigned modes[2] = {GPOLICY_COMPACT, GPOLICY_COMPACT};

    unsigned gold_gpu_id_0[37] = {0,0,0,0,0,0,0,0,
                                  1,1,1,1,1,1,1,1,
                                  2,2,2,2,2,2,2,
                                  3,3,3,3,3,3,3,
                                  4,4,4,4,4,4,4};
    unsigned gold_gpu_id_1[37] = {0,0,0,1,1,1,2,2,
                                  0,0,0,1,1,1,2,2,
                                  0,0,0,1,1,2,2,
                                  0,0,0,1,1,2,2,
                                  0,0,0,1,1,2,2};
    for (int rank_id =0; rank_id < num_ranks; rank_id++)
    {
            unsigned gpu_id[2];
            rankToAccelerator(rank_id, num_ranks, 2, num_gpus, modes, gpu_id);
            assert (gpu_id[0] == gold_gpu_id_0[rank_id]);
            assert (gpu_id[1] == gold_gpu_id_1[rank_id]);

    }

}

void test_compact_interleave() {
    int num_ranks = 37;
    unsigned num_gpus[2] = {5,3};
    unsigned modes[2] = {GPOLICY_COMPACT, GPOLICY_INTERLEAVE};

    unsigned gold_gpu_id_0[37] = {0,0,0,0,0,0,0,0,
                                  1,1,1,1,1,1,1,1,
                                  2,2,2,2,2,2,2,
                                  3,3,3,3,3,3,3,
                                  4,4,4,4,4,4,4};
    unsigned gold_gpu_id_1[37] = {0,1,2,0,1,2,0,1,
                                  0,1,2,0,1,2,0,1,
                                  0,1,2,0,1,2,0,
                                  0,1,2,0,1,2,0,
                                  0,1,2,0,1,2,0};
    for (int rank_id =0; rank_id < num_ranks; rank_id++)
    {
            unsigned gpu_id[2];
            rankToAccelerator(rank_id, num_ranks, 2, num_gpus, modes, gpu_id);
            assert (gpu_id[0] == gold_gpu_id_0[rank_id]);
            assert (gpu_id[1] == gold_gpu_id_1[rank_id]);

    }
}

void test_interleave_interleave() {
    int num_ranks = 37;
    unsigned num_gpus[2] = {5,3};
    unsigned modes[2] = {GPOLICY_INTERLEAVE, GPOLICY_INTERLEAVE};

    unsigned gold_gpu_id_0[37] = {0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1};

    unsigned gold_gpu_id_1[37] = {0,0,0,0,0,
                                  1,1,1,1,1,
                                  2,2,2,2,2,
                                  0,0,0,0,0,
                                  1,1,1,1,1,
                                  2,2,2,2,2,
                                  0,0,0,0,0,
                                  1,1};

    for (int rank_id =0; rank_id < num_ranks; rank_id++)
    {
            unsigned gpu_id[2];
            rankToAccelerator(rank_id, num_ranks, 2, num_gpus, modes, gpu_id);
            assert (gpu_id[0] == gold_gpu_id_0[rank_id]);
            assert (gpu_id[1] == gold_gpu_id_1[rank_id]);

    }
}

void test_interleave_compact() {
    int num_ranks = 37;
    unsigned num_gpus[2] = {5,3};
    unsigned modes[2] = {GPOLICY_INTERLEAVE, GPOLICY_COMPACT};

    unsigned gold_gpu_id_0[37] = {0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1,2,3,4,
                                  0,1};

    unsigned gold_gpu_id_1[37] = {0,0,0,0,0,
                                  0,0,0,0,0,
                                  0,0,0,0,0,
                                  1,1,1,1,1,
                                  1,1,1,1,1,
                                  1,1,2,2,2,
                                  2,2,2,2,2,
                                  2,2};
    for (int rank_id =0; rank_id < num_ranks; rank_id++)
    {
            unsigned gpu_id[2];
            rankToAccelerator(rank_id, num_ranks, 2, num_gpus, modes, gpu_id);
            assert (gpu_id[0] == gold_gpu_id_0[rank_id]);
            assert (gpu_id[1] == gold_gpu_id_1[rank_id]);
    }
}


int main(int argc, char const *argv[])
{
    test_compact_compact();
    test_compact_interleave();
    test_interleave_interleave();
    test_interleave_compact();
    return 0;
}