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)


[ROCm/rocshmem commit: 92fb1abaf2]

projects/rocshmem/src/ipc/backend_ipc.cpp

@@ -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;
 }
 

projects/rocshmem/src/ipc/context_ipc_device.cpp

@@ -88,8 +88,8 @@ __device__ void IPCContext::getmem_nbi(void *dest, const void *source,
 }
 
 __device__ void IPCContext::fence() {
-  for (int i{0}; i < num_pes; i++) {
-    detail::atomic::store<int, detail::atomic::memory_scope_system>(&fence_pool[i], 1, orders_);
+  for (int i{0}, j{tinfo->pe_start}; i < tinfo->size; i++, j += tinfo->stride) {
+    detail::atomic::store<int, detail::atomic::memory_scope_system>(&fence_pool[j], 1, orders_);
   }
 }
 

projects/rocshmem/src/ipc/context_ipc_device.hpp

@@ -25,6 +25,7 @@
 
 #include "../context.hpp"
 #include "../atomic.hpp"
+#include "../team.hpp"
 
 namespace rocshmem {
 
@@ -239,6 +240,14 @@ class IPCContext : public Context {
 
   //Buffer to perform Atomic store to enforce memory ordering
   int *fence_pool{nullptr};
+
+ public:
+  //TODO(Avinash):
+  //Make tinfo private variable, it requires changes to the context
+  //creation API in backend
+
+  //Team information for the team associated with the context
+  TeamInfo *tinfo{nullptr};
 };
 
 }  // namespace rocshmem

projects/rocshmem/src/ipc/ipc_context_proxy.hpp

@@ -41,10 +41,11 @@ class IPCDefaultContextProxy {
   /*
    * Placement new the memory which is allocated by proxy_
    */
-  explicit IPCDefaultContextProxy(IPCBackend* backend) : constructed_{true} {
+  explicit IPCDefaultContextProxy(IPCBackend* backend, TeamInfo *tinfo)
+  : constructed_{true} {
     auto ctx{proxy_.get()};
     new (ctx) IPCContext(reinterpret_cast<Backend*>(backend));
-    roc_shmem_ctx_t local{ctx, nullptr};
+    roc_shmem_ctx_t local{ctx, tinfo};
     set_internal_ctx(&local);
   }