main.cc 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. #include <mpi.h>
  2. #include <string>
  3. #include <cstring>
  4. #include <functional>
  5. #include <iostream>
  6. int main(int argc, char *argv[]) {
  7. int rank, size;
  8. char proc_name[2048];
  9. int proc_name_len;
  10. MPI_Comm node, head;
  11. MPI_Init(&argc, &argv);
  12. MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  13. MPI_Comm_size(MPI_COMM_WORLD, &size);
  14. // node comm
  15. MPI_Get_processor_name(proc_name, &proc_name_len);
  16. std::string proc_name_str(proc_name, proc_name_len);
  17. int proc_name_hash = std::hash<std::string>{}(proc_name_str);
  18. // std::cout << proc_name_str << ", " << proc_name_hash << std::endl;
  19. MPI_Comm_split(MPI_COMM_WORLD, proc_name_hash, -1, &node);
  20. int node_rank, node_size;
  21. MPI_Comm_rank(node, &node_rank);
  22. MPI_Comm_size(node, &node_size);
  23. // head comm
  24. int head_rank, head_size;
  25. int headornot = node_rank == 0;
  26. MPI_Comm_split(MPI_COMM_WORLD, headornot, -1, &head);
  27. MPI_Comm_rank(head, &head_rank);
  28. MPI_Comm_size(head, &head_size);
  29. //std::cout << "rank:" << rank << ", head_rank:" << head_rank << ", head_size:" << head_size << std::endl;
  30. char message[2048] = "default string";
  31. printf("rank:%d, is_head:%d, node:%s, node_rank:%d, msg:%s\n", rank, headornot, proc_name, node_rank, message);
  32. // bcast in head
  33. if (head_rank == 0 && headornot) {
  34. strcpy(message, "hello from root");
  35. }
  36. if (headornot) {
  37. MPI_Bcast((void*) message, 16, MPI_CHAR, 0, head);
  38. }
  39. // bcast in node
  40. MPI_Bcast((void*) message, 16, MPI_CHAR, 0, node);
  41. printf("rank:%d, is_head:%d, node:%s, node_rank:%d, msg:%s\n", rank, headornot, proc_name, node_rank, message);
  42. MPI_Finalize();
  43. }