使用MPI计算并行排名

TMPI_Rank(
    void *send_data,
    void *recv_data,
    MPI_Datatype datatype,
    MPI_Comm comm)

// 为进程0的TMPI_Rank收集数字。为MPI的数据类型分配空间
// 对进程0返回 void * 指向的缓冲区
// 对所有其他进程返回NULL
void *gather_numbers_to_root(void *number, MPI_Datatype datatype,
                             MPI_Comm comm) {
  int comm_rank, comm_size;
  MPI_Comm_rank(comm, &comm_rank);
  MPI_Comm_size(comm, &comm_size);

  // 在根进程上分配一个数组
  // 数组大小取决于所用的MPI数据类型
  int datatype_size;
  MPI_Type_size(datatype, &datatype_size);
  void *gathered_numbers;
  if (comm_rank == 0) {
    gathered_numbers = malloc(datatype_size * comm_size);
  }

  // 在根进程上收集所有数字
  MPI_Gather(number, 1, datatype, gathered_numbers, 1,
             datatype, 0, comm);

  return gathered_numbers;
}

// 保存进程在通讯器中的次序（rank）和对应数字
// 该结构体用于数组排序，
// 并同时完整保留所属进程信息

typedef struct {
  int comm_rank;
  union {
    float f;
    int i;
  } number;
} CommRankNumber;

// 这个函数在根进程上对收集到的数字排序
// 返回一个数组，数组按进程在通讯器中的次序排序
// 注意 - 该函数只在根进程上运行

int *get_ranks(void *gathered_numbers, int gathered_number_count,
               MPI_Datatype datatype) {
  int datatype_size;
  MPI_Type_size(datatype, &datatype_size);

  // 将收集到的数字数组转换为CommRankNumbers数组
  // 这允许我们在排序的同时，完整保留数字所属进程的信息
  
  CommRankNumber *comm_rank_numbers = malloc(
    gathered_number_count * sizeof(CommRankNumber));
  int i;
  for (i = 0; i < gathered_number_count; i++) {
    comm_rank_numbers[i].comm_rank = i;
    memcpy(&(comm_rank_numbers[i].number),
           gathered_numbers + (i * datatype_size),
           datatype_size);
  }

  // 根据数据类型对comm_rank_numbers排序
  if (datatype == MPI_FLOAT) {
    qsort(comm_rank_numbers, gathered_number_count,
          sizeof(CommRankNumber), &compare_float_comm_rank_number);
  } else {
    qsort(comm_rank_numbers, gathered_number_count,
          sizeof(CommRankNumber), &compare_int_comm_rank_number);
  }

  // 现在comm_rank_numbers是排好序的，下面生成一个数组，
  // 包含每个进程的排名，数组第i个元素是进程i的数字的排名
  
  int *ranks = (int *)malloc(sizeof(int) * gathered_number_count);
  for (i = 0; i < gathered_number_count; i++) {
    ranks[comm_rank_numbers[i].comm_rank] = i;
  }

  // 清理并返回排名数组
  free(comm_rank_numbers);
  return ranks;
}

// 获取send_data的排名, 类型为datatype
// 排名用recv_data返回，类型为datatype
int TMPI_Rank(void *send_data, void *recv_data, MPI_Datatype datatype,
             MPI_Comm comm) {
  // 首先检查基本情况 - 此函数只支持MPI_INT和MPI_FLOAT

  if (datatype != MPI_INT && datatype != MPI_FLOAT) {
    return MPI_ERR_TYPE;
  }

  int comm_size, comm_rank;
  MPI_Comm_size(comm, &comm_size);
  MPI_Comm_rank(comm, &comm_rank);

  // 为了计算排名，必须将数字收集到一个进程中
  // 对数字排序, 然后将排名结果分散传回
  // 首先在进程0上收集数字
  void *gathered_numbers = gather_numbers_to_root(send_data, datatype,
                                                  comm);

  // 获取每个进程的次序（rank)
  int *ranks = NULL;
  if (comm_rank == 0) {
    ranks = get_ranks(gathered_numbers, comm_size, datatype);
  }

  // 分散发回排名结果
  MPI_Scatter(ranks, 1, MPI_INT, recv_data, 1, MPI_INT, 0, comm);

  // 清理
  if (comm_rank == 0) {
    free(gathered_numbers);
    free(ranks);
  }
}