From 71af1b366d0e474a6da4c3fb0cec5b8d379c7623 Mon Sep 17 00:00:00 2001
From: avinashkethineedi <avinash.kethineedi@amd.com>
Date: Wed, 5 Feb 2025 03:41:40 +0000
Subject: [PATCH] Fix `rocshmem_ctx_my_pe` and `rocshmem_ctx_n_pes` APIs to
 return PE numbering and size relative to the team in a team-specific context.

[ROCm/rocshmem commit: e311400d15d423a7a3aa292c77e93c96bb240f51]
---
 projects/rocshmem/src/rocshmem_gpu.cpp | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/projects/rocshmem/src/rocshmem_gpu.cpp b/projects/rocshmem/src/rocshmem_gpu.cpp
index 0e4ce22508..78350deb92 100644
--- a/projects/rocshmem/src/rocshmem_gpu.cpp
+++ b/projects/rocshmem/src/rocshmem_gpu.cpp
@@ -606,7 +606,8 @@ __device__ void rocshmem_wg_team_sync(rocshmem_team_t team) {
 __device__ int rocshmem_ctx_n_pes(rocshmem_ctx_t ctx) {
   GPU_DPRINTF("Function: rocshmem_n_pes\n");
 
-  return get_internal_ctx(ctx)->num_pes;
+  TeamInfo *tinfo = reinterpret_cast<TeamInfo *>(ctx.team_opaque);
+  return tinfo->size;
 }
 
 __device__ int rocshmem_n_pes() {
@@ -616,7 +617,21 @@ __device__ int rocshmem_n_pes() {
 __device__ int rocshmem_ctx_my_pe(rocshmem_ctx_t ctx) {
   GPU_DPRINTF("Function: rocshmem_ctx_my_pe\n");
 
-  return get_internal_ctx(ctx)->my_pe;
+  TeamInfo *tinfo = reinterpret_cast<TeamInfo *>(ctx.team_opaque);
+  int my_pe{get_internal_ctx(ctx)->my_pe};
+  int pe_start{tinfo->pe_start};
+  int stride{tinfo->stride};
+  int size{tinfo->size};
+
+  int translated_pe = (my_pe - pe_start) / stride;
+
+  if ((my_pe < pe_start) ||
+     ((my_pe - pe_start) % stride) ||
+     (translated_pe >= size)) {
+    translated_pe = -1;
+  }
+
+  return translated_pe;
 }
 
 __device__ int rocshmem_my_pe() {