Add team information to the context

* Update roc_shmem_ctx_fence API to use team-relative PE numbering
* Update backend to populate team_opaque member of ROC_SHMEM_CTX_DEFAULT (used to store information about the team wrt TEAM_WORLD)
This commit is contained in:
avinashkethineedi
2024-10-04 17:56:15 +00:00
parent 979aed105a
commit 92fb1abaf2
8 changed files with 33 additions and 14 deletions
+5 -2
View File
@@ -82,10 +82,11 @@ IPCBackend::IPCBackend(MPI_Comm comm)
allocate_atomic_region(&bp->atomic_ret, MAX_NUM_BLOCKS);
default_context_proxy_ = IPCDefaultContextProxyT(this);
setup_team_world();
TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world;
default_context_proxy_ = IPCDefaultContextProxyT(this, tinfo);
roc_shmem_collective_init();
setup_fence_buffer();
@@ -143,6 +144,8 @@ __device__ bool IPCBackend::create_ctx(int64_t options, roc_shmem_ctx_t *ctx) {
ctx_ = pop_result.value;
ctx->ctx_opaque = ctx_;
ctx_->tinfo = reinterpret_cast<TeamInfo *>(ctx->team_opaque);
return true;
}