[rocshmem] Expose CTX Array (#4121)

AIROCSHMEM-8

Author: Yiltan

projects/rocshmem/include/rocshmem/rocshmem_common.hpp

@@ -133,6 +133,8 @@ extern "C" __device__  rocshmem_ctx_t __attribute__((visibility("default"))) ROC
  * valid context.
  */
 extern __constant__ rocshmem_ctx_t ROCSHMEM_CTX_INVALID;
+
+extern __constant__  rocshmem_ctx_t *rocshmem_ctx_array;
 /**
  * Used internally to set default context.
  */
@@ -180,7 +182,7 @@ const rocshmem_team_t ROCSHMEM_TEAM_INVALID = nullptr;
 using rocshmem_uniqueid_t = std::array<uint8_t, ROCSHMEM_UNIQUE_ID_BYTES>;
 
 /**
- * @brief Data structure used for attribute based 
+ * @brief Data structure used for attribute based
  *        initialization
  */
 struct rocshmem_init_attr_t  {

projects/rocshmem/src/gda/backend_gda.cpp

@@ -178,9 +178,42 @@ void GDABackend::setup_ctxs() {
     new (&ctx_array[i]) GDAContext(this, i + 1, gda_provider);
     ctx_free_list.get()->push_back(ctx_array + i);
   }
+
+  rocshmem_ctx_t *rocshmem_ctx_array_device = nullptr;
+  rocshmem_ctx_t *rocshmem_ctx_array_ptr = nullptr;
+  size_t ctx_array_size = sizeof(rocshmem_ctx_t) * envvar::max_num_contexts;
+
+  CHECK_HIP(hipMalloc((void**)&rocshmem_ctx_array_device, ctx_array_size));
+
+  for (size_t i = 0; i < envvar::max_num_contexts; i++) {
+    rocshmem_ctx_array_device[i].ctx_opaque  = &ctx_array[i];
+    rocshmem_ctx_array_device[i].team_opaque = team_tracker.get_team_world()->tinfo_wrt_world;
+  }
+
+  CHECK_HIP(hipGetSymbolAddress(reinterpret_cast<void**>(&rocshmem_ctx_array_ptr),
+                                HIP_SYMBOL(rocshmem_ctx_array)));
+
+  CHECK_HIP(hipMemcpy(rocshmem_ctx_array_ptr,
+                      &rocshmem_ctx_array_device,
+                      sizeof(rocshmem_ctx*),
+                      hipMemcpyDefault));
 }
 
 void GDABackend::cleanup_ctxs() {
+  /* Free ctx array */
+  rocshmem_ctx_t *rocshmem_ctx_array_ptr = nullptr;
+  rocshmem_ctx_t *rocshmem_ctx_array_device = nullptr;
+
+  CHECK_HIP(hipGetSymbolAddress(reinterpret_cast<void**>(&rocshmem_ctx_array_ptr),
+                                HIP_SYMBOL(rocshmem_ctx_array)));
+
+  CHECK_HIP(hipMemcpy(&rocshmem_ctx_array_device,
+                      rocshmem_ctx_array_ptr,
+                      sizeof(rocshmem_ctx*),
+                      hipMemcpyDefault));
+
+  CHECK_HIP(hipFree(rocshmem_ctx_array_device));
+
   ctx_free_list.~FreeListProxy();
   for (size_t i = 0; i < envvar::max_num_contexts; i++) {
     ctx_array[i].~GDAContext();

projects/rocshmem/src/rocshmem_gpu.cpp

@@ -77,6 +77,8 @@ namespace rocshmem {
 __device__  rocshmem_ctx_t
 __attribute__((visibility("default"))) ROCSHMEM_CTX_DEFAULT{};
 
+__constant__  rocshmem_ctx_t *rocshmem_ctx_array;
+
 __constant__ Backend *device_backend_proxy;
 
 __constant__ constmem_t constmem;