Update backend to use provided MPI communicator during library initialization (#79)

* Update backend to use provided MPI communicator during library initialization, default to `MPI_COMM_WORLD`

* Update `rocshmem_my_pe` and `rocshmem_n_pes` host APIs
   - Return values from backend if initialized; otherwise, fallback to MPI_Singleton.
This commit is contained in:
Avinash Kethineedi
2025-04-14 09:18:57 -05:00
committed by GitHub
parent 0fd628458c
commit 05755847f5
10 changed files with 79 additions and 58 deletions
+15 -4
View File
@@ -197,13 +197,24 @@ rocshmem_ctx_t ROCSHMEM_HOST_CTX_DEFAULT;
}
[[maybe_unused]] __host__ int rocshmem_my_pe() {
MPIInitSingleton *s = s->GetInstance();
return s->get_rank();
if(backend == nullptr) {
MPIInitSingleton *s = s->GetInstance();
return s->get_rank();
}
else
{
return backend->getMyPE();
}
}
[[maybe_unused]] __host__ int rocshmem_n_pes() {
MPIInitSingleton *s = s->GetInstance();
return s->get_nprocs();
if(backend == nullptr) {
MPIInitSingleton *s = s->GetInstance();
return s->get_nprocs();
}
else {
return backend->getNumPEs();
}
}
[[maybe_unused]] __host__ void *rocshmem_malloc(size_t size) {