From 39c5be18b1f8205c15cab789c0df0423ae3b41b5 Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Mon, 15 Jul 2024 10:18:20 -0500
Subject: [PATCH 1/6] Remove CACHED_HEAP option

---
 CMakeLists.txt                                    | 1 -
 README.md                                         | 2 +-
 cmake/config.h.in                                 | 3 +--
 scripts/build_configs/rc_single                   | 1 -
 scripts/build_configs/rc_single_single_node       | 2 --
 scripts/build_configs/rc_single_single_node_debug | 1 -
 src/gpu_ib/connection.cpp                         | 2 +-
 src/memory/heap_type.hpp                          | 2 +-
 src/util.hpp                                      | 2 +-
 9 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index ee7f42fb12..8603737308 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -60,7 +60,6 @@ option(USE_IPC "Enable IPC support (using HIP)" OFF)
 option(USE_THREADS "Enable workgroup threads to share network queues" OFF)
 option(USE_WF_COAL "Enable wavefront message coalescing" OFF)
 option(USE_COHERENT_HEAP "Enable support for coherent systems" OFF)
-option(USE_CACHED_HEAP "Enable support for cached systems" OFF)
 option(USE_MANAGED_HEAP "Enable managed memory" OFF)
 option(USE_HOST_HEAP "Enable host memory using malloc/free" OFF)
 option(USE_HIP_HOST_HEAP "Enable host memory using hip api" OFF)
diff --git a/README.md b/README.md
index c974fd7052..ba671ff51a 100644
--- a/README.md
+++ b/README.md
@@ -45,7 +45,7 @@ ROC_SHMEM base requirements:
     *  May work with other versions, but not tested
 * AMD GFX9 GPUs (e.g.: MI25, Vega 56, Vega 64, MI50, MI60, MI100, Radeon VII)
 * AMD MI200 GPUs: To enable the support on MI200, please configure the library
- with USE_CACHED_HEAP
+ with USE_COHERENT_HEAP
 * ROCm-aware MPI as described in
   [Building the Dependencies](#building-the-dependencies)
 * InfiniBand adaptor compatable with ROCm RDMA technology
diff --git a/cmake/config.h.in b/cmake/config.h.in
index 6ec07c2853..f409c10ba3 100644
--- a/cmake/config.h.in
+++ b/cmake/config.h.in
@@ -7,10 +7,9 @@
 #cmakedefine USE_SHARED_CTX
 #cmakedefine USE_WF_COAL
 #cmakedefine USE_COHERENT_HEAP
-#cmakedefine USE_CACHED_HEAP
 #cmakedefine USE_MANAGED_HEAP
 #cmakedefine USE_HOST_HEAP
 #cmakedefine USE_HIP_HOST_HEAP
 #cmakedefine USE_FUNC_CALL
 #cmakedefine USE_SINGLE_NODE
-#cmakedefine USE_HOST_SIDE_HDP_FLUSH
\ No newline at end of file
+#cmakedefine USE_HOST_SIDE_HDP_FLUSH
diff --git a/scripts/build_configs/rc_single b/scripts/build_configs/rc_single
index 31eb7283c6..aceca7dbed 100755
--- a/scripts/build_configs/rc_single
+++ b/scripts/build_configs/rc_single
@@ -20,7 +20,6 @@ cmake \
     -DUSE_DC=OFF \
     -DUSE_IPC=OFF \
     -DUSE_COHERENT_HEAP=OFF \
-    -DUSE_CACHED_HEAP=OFF \
     -DUSE_THREADS=OFF \
     -DUSE_WF_COAL=OFF \
     $src_path
diff --git a/scripts/build_configs/rc_single_single_node b/scripts/build_configs/rc_single_single_node
index 8bcf9fcb0d..a5223c379f 100755
--- a/scripts/build_configs/rc_single_single_node
+++ b/scripts/build_configs/rc_single_single_node
@@ -20,12 +20,10 @@ cmake \
     -DUSE_DC=OFF \
     -DUSE_IPC=ON \
     -DUSE_COHERENT_HEAP=OFF \
-    -DUSE_CACHED_HEAP=OFF \
     -DUSE_THREADS=OFF \
     -DUSE_WF_COAL=OFF \
     -DUSE_SINGLE_NODE=ON \
     -DUSE_HOST_SIDE_HDP_FLUSH=ON\
-    -DROCM_PATH="/opt/rocm-5.4.2/"\
     $src_path
 cmake --build . --parallel 8
 cmake --install .
diff --git a/scripts/build_configs/rc_single_single_node_debug b/scripts/build_configs/rc_single_single_node_debug
index b535f1b853..c4e2477da9 100755
--- a/scripts/build_configs/rc_single_single_node_debug
+++ b/scripts/build_configs/rc_single_single_node_debug
@@ -20,7 +20,6 @@ cmake \
     -DUSE_DC=OFF \
     -DUSE_IPC=ON \
     -DUSE_COHERENT_HEAP=OFF \
-    -DUSE_CACHED_HEAP=OFF \
     -DUSE_THREADS=OFF \
     -DUSE_WF_COAL=OFF \
     -DUSE_SINGLE_NODE=ON \
diff --git a/src/gpu_ib/connection.cpp b/src/gpu_ib/connection.cpp
index 631f72f324..b6b2ae4dcd 100644
--- a/src/gpu_ib/connection.cpp
+++ b/src/gpu_ib/connection.cpp
@@ -256,7 +256,7 @@ void* Connection::buf_alloc([[maybe_unused]] struct ibv_pd* pd,
   if (use_gpu_mem) {
     void* dev_ptr;
     if (coherent_cq == 1) {
-#if defined USE_COHERENT_HEAP || defined USE_CACHED_HEAP
+#if defined USE_COHERENT_HEAP
       CHECK_HIP(hipMalloc(reinterpret_cast<void**>(&dev_ptr), size));
 #else
   #ifdef HIP_SUPPORTS_MALLOC_UNCACHED
diff --git a/src/memory/heap_type.hpp b/src/memory/heap_type.hpp
index 0008317498..3f48edd11e 100644
--- a/src/memory/heap_type.hpp
+++ b/src/memory/heap_type.hpp
@@ -40,7 +40,7 @@ namespace rocshmem {
 
 #if defined USE_MANAGED_HEAP
 using HEAP_T = HeapMemory<HIPAllocatorManaged>;
-#elif defined USE_COHERENT_HEAP || defined USE_CACHED_HEAP
+#elif defined USE_COHERENT_HEAP
 using HEAP_T = HeapMemory<HIPAllocator>;
 #elif defined USE_HOST_HEAP
 using HEAP_T = HeapMemory<HostAllocator>;
diff --git a/src/util.hpp b/src/util.hpp
index 0d4dad806a..1d5cd9f9a6 100644
--- a/src/util.hpp
+++ b/src/util.hpp
@@ -126,7 +126,7 @@ extern const int gpu_clock_freq_mhz;
 __device__ __forceinline__ void __roc_inv() { asm volatile("buffer_wbinvl1;"); }
 
 __device__ __forceinline__ void __roc_flush() {
-#if defined USE_COHERENT_HEAP || defined USE_CACHED_HEAP
+#if defined USE_COHERENT_HEAP
 #if __gfx90a__
   asm volatile("s_dcache_wb;");
   asm volatile("buffer_wbl2;");

From 23f2983214ec9ab8215364cb63f366f68d7735e1 Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Mon, 15 Jul 2024 10:38:38 -0500
Subject: [PATCH 2/6] Create new ro_ipc config

---
 scripts/build_configs/ro_ipc | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100755 scripts/build_configs/ro_ipc

diff --git a/scripts/build_configs/ro_ipc b/scripts/build_configs/ro_ipc
new file mode 100755
index 0000000000..d57404b88c
--- /dev/null
+++ b/scripts/build_configs/ro_ipc
@@ -0,0 +1,29 @@
+#!/bin/bash
+# Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
+
+if [ -z $1 ]
+then
+  install_path=~/rocshmem
+else
+  install_path=$1
+fi
+
+src_path=$(dirname "$(realpath $0)")/../../
+
+cmake \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_INSTALL_PREFIX=$install_path \
+    -DCMAKE_VERBOSE_MAKEFILE=OFF \
+    -DDEBUG=OFF \
+    -DPROFILE=OFF \
+    -DUSE_GPU_IB=OFF \
+    -DUSE_DC=OFF \
+    -DUSE_IPC=ON \
+    -DUSE_COHERENT_HEAP=ON \
+    -DUSE_THREADS=OFF \
+    -DUSE_WF_COAL=OFF \
+    -DUSE_SINGLE_NODE=ON \
+    -DUSE_HOST_SIDE_HDP_FLUSH=OFF\
+    $src_path
+cmake --build . --parallel 8
+cmake --install .

From ca36cc249e92bc38429fbb8cc6a410e1b071d4c8 Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Thu, 11 Jul 2024 08:37:52 -0700
Subject: [PATCH 3/6] Update functional_tests driver.sh script

---
 scripts/functional_tests/driver.sh | 785 +++++++++++++++--------------
 1 file changed, 414 insertions(+), 371 deletions(-)

diff --git a/scripts/functional_tests/driver.sh b/scripts/functional_tests/driver.sh
index a7a3e869e5..5adb033166 100755
--- a/scripts/functional_tests/driver.sh
+++ b/scripts/functional_tests/driver.sh
@@ -22,417 +22,460 @@
 
 if [ $# -eq 0 ] ; then
     echo "This script must be run with at least 2 arguments."
-    echo 'Usage: ${0} argument1 argument2 [argument3] [argument4] [argument5]'
+    echo 'Usage: ${0} argument1 argument2 [argument3]'
     echo "  argument1 : path to the tester driver"
     echo "  argument2 : test type to run, e.g put"
     echo "  argument3 : directory to put the output logs"
-    echo "  argument4 : enable gdb debug"
-    echo "  argument5 : shmem context type"
     exit 1
 fi
 
-shm_ctx=4  # run tests with SHMEM_CONTEXT_WG_PRIVATE
-if [ $# -eq 5 ] ; then
-    shm_ctx=$5
-fi
-
-ENABLE_DBG=false
-if [ $# -ge 4 ] ; then
-    ENABLE_DBG=$4
-fi
-
-if [ "$ENABLE_DBG" = true ]
-then
-    gdb_cmd="xterm -e gdb -x gdbscript --args"
-else
-    gdb_cmd=""
-fi
-
-echo "Test Name ${2} (with shmem context: ${shm_ctx})"
+echo "Test Name ${2}"
 
 check() {
     if [ $? -ne 0 ]
     then
         echo "Failed $1" >&2
-        exit 1
     fi
 }
 
 case $2 in
-    *"single_thread")
-        echo "get"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 0 -x ${shm_ctx} > $3/get.log
-        check get
-        echo "get_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 0 -x ${shm_ctx} > $3/get_th.log
-        check get_th
-        echo "get_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 1 -x ${shm_ctx} > $3/get_nbi.log
-        check get_nbi
-        echo "get_nbi_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 1 -x ${shm_ctx} > $3/get_nbi_th.log
-        check get_nbi_th
-        echo "put"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 2 -x ${shm_ctx} > $3/put.log
-        check put
-        echo "put_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 2 -x ${shm_ctx} > $3/put_th.log
-        check put_th
-        echo "put_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 3 -x ${shm_ctx} > $3/put_nbi.log
-        check put_nbi
-        echo "put_nbi_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 3 -x ${shm_ctx} > $3/put_nbi_th.log
-        check put_nbi_th
-#        echo "team_ctx_infra"
-#        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 42 -x ${shm_ctx} > $3/team_ctx_infra.log
-#        check team_ctx_infra
-        echo "team_ctx_put_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 41 -x ${shm_ctx} > $3/team_ctx_put_nbi.log
-        check team_ctx_put_nbi
-        echo "team_ctx_put_nbi_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 41 -x ${shm_ctx} > $3/team_ctx_put_nbi_th.log
-        check team_ctx_put_nbi_th
-        echo "team_ctx_get_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 39 -x ${shm_ctx} > $3/team_ctx_get_nbi.log
-        check team_ctx_get_nbi
-        echo "team_ctx_get_nbi_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 39 -x ${shm_ctx} > $3/team_ctx_get_nbi_th.log
-        check team_ctx_get_nbi_th
-        echo "reduction"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 5 -x ${shm_ctx} > $3/reduction.log
-        check reduction
-        echo "amo_fadd"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 6 -x ${shm_ctx} > $3/amo_fadd.log
-        check amo_fadd
-        echo "amo_fadd_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 6 -x ${shm_ctx} > $3/amo_fadd_th.log
-        check amo_fadd_th
-        echo "amo_finc"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 7 -x ${shm_ctx} > $3/amo_finc.log
-        check amo_finc
-        echo "amo_finc_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 7 -x ${shm_ctx} > $3/amo_finc_th.log
-        check amo_finc_th
-        echo "amo_fetch"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 8 -x ${shm_ctx} > $3/amo_fetch.log
-        check amo_fetch
-        echo "amo_fetch_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 8 -x ${shm_ctx} > $3/amo_fetch_th.log
-        check amo_fetch_th
-        echo "amo_fcswap"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 9 -x ${shm_ctx} > $3/amo_fcswap.log
-        check amo_fcswap
-        echo "amo_fcswap_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 9 -x ${shm_ctx} > $3/amo_fcswap_th.log
-        check amo_fcswap_th
-        echo "amo_add"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 10 -x ${shm_ctx} > $3/amo_add.log
-        check amo_add
-        echo "amo_add_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 10 -x ${shm_ctx} > $3/amo_add_th.log
-        check amo_add_th
-        echo "amo_set"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 44 -x ${shm_ctx} > $3/amo_set.log
-        check amo_set
-        echo "amo_set_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 44 -x ${shm_ctx} > $3/amo_set_th.log
-        check amo_set_th
-        echo "amo_inc"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 11 -x ${shm_ctx} > $3/amo_inc.log
-        check amo_inc
-        echo "amo_inc_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 11 -x ${shm_ctx} > $3/amo_inc_th.log
-        check amo_inc_th
-        echo "init"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 13 -x ${shm_ctx} > $3/init.log
-        check init
-        echo "ping_pong"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 14 -x ${shm_ctx} > $3/ping_pong.log
-        check ping_pong
-        echo "ping_pong_th"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1024 -w 25 -s 32768 -a 14 -x ${shm_ctx} > $3/ping_pong_th.log
-        check ping_pong_th
-        echo "barrier_all"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 17 -x ${shm_ctx} > $3/barrier_all.log
-        check barrier_all
-        echo "sync_all"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 18 -x ${shm_ctx} > $3/sync_all.log
-        check sync_all
-        echo "sync"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 19 -x ${shm_ctx} > $3/sync.log
-        check sync
-        echo "alltoall"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 23 -x ${shm_ctx} > $3/alltoall.log
-        check alltoall
-        echo "fcollect"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 22 -x ${shm_ctx} > $3/fcollect.log
-        check fcollect
+    ###########################################################################
+    ############################## SERIAL TESTS ###############################
+    ###########################################################################
+    *"serial")
+        echo "get_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 0 > $3/get_n2_w1_z1_1MB.log
+        check get_n2_w1_z1_1MB
+        echo "getnbi_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 1 > $3/getnbi_n2_w1_z1_1MB.log
+        check getnbi_n2_w1_z1_1MB
+        echo "put_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 2 > $3/put_n2_w1_z1_1MB.log
+        check put_n2_w1_z1_1MB
+        echo "putnbi_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 3 > $3/putnbi_n2_w1_z1_1MB.log
+        check putnbi_n2_w1_z1_1MB
+        echo "amofadd_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 6 > $3/amofadd_n2_w1_z1.log
+        check amofadd_n2_w1_z1
+        echo "amofinc_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 7 > $3/amofinc_n2_w1_z1.log
+        check amofinc_n2_w1_z1
+        echo "amofetch_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 8 > $3/amofetch_n2_w1_z1.log
+        check amofetch_n2_w1_z1
+        echo "amofcswap_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 9 > $3/amofcswap_n2_w1_z1.log
+        check amofcswap_n2_w1_z1
+        echo "amoadd_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 10 > $3/amoadd_n2_w1_z1.log
+        check amoadd_n2_w1_z1
+        echo "amoinc_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 11 > $3/amoinc_n2_w1_z1.log
+        check amoinc_n2_w1_z1
+        echo "pingpong_n2_w1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -a 14 > $3/pingpong_n2_w1.log
+        check pingpong_n2_w1
+        echo "amoset_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 44 > $3/amoset_n2_w1_z1.log
+        check amoset_n2_w1_z1
         ;;
-    *"multi_thread")
-        echo "get"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 0 -x ${shm_ctx} > $3/get.log
-        check get
-        echo "get_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 1 -x ${shm_ctx} > $3/get_nbi.log
-        check get_nbi
-        echo "put"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 2 -x ${shm_ctx} > $3/put.log
-        check put
-        echo "put_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 3 -x ${shm_ctx} > $3/put_nbi.log
-        check put_nbi
-#        echo "team_ctx_infra"
-#        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 42 -x ${shm_ctx} > $3/team_ctx_infra.log
-#        check team_ctx_infra
-        echo "team_ctx_put_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 41 -x ${shm_ctx} > $3/team_ctx_put_nbi.log
-        check team_ctx_put_nbi
-        echo "team_ctx_get_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 39 -x ${shm_ctx} > $3/team_ctx_get_nbi.log
-        check team_ctx_get_nbi
-        echo "get_swarm"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 4 -x ${shm_ctx} > $3/get_swarm.log
-        check get_swarm
-        echo "reduction"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 5 -x ${shm_ctx} > $3/reduction.log
-        check reduction
-        echo "amo_fadd"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 6 -x ${shm_ctx} > $3/amo_fadd.log
-        check amo_fadd
-        echo "amo_finc"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 7 -x ${shm_ctx} > $3/amo_finc.log
-        check amo_finc
-        echo "amo_fetch"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 8 -x ${shm_ctx} > $3/amo_fetch.log
-        check amo_fetch
-        echo "amo_fcswap"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 9 -x ${shm_ctx} > $3/amo_fcswap.log
-        check amo_fcswap
-        echo "amo_add"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 10 -x ${shm_ctx} > $3/amo_add.log
-        check amo_add
-        echo "amo_set"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 44 -x ${shm_ctx} > $3/amo_set.log
-        check amo_set
-        echo "amo_inc"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 11 -x ${shm_ctx} > $3/amo_inc.log
-        check amo_inc
-        echo "init"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 13 -x ${shm_ctx} > $3/init.log
-        check init
-        echo "ping_pong"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 14 -x ${shm_ctx} > $3/ping_pong.log
-        check ping_pong
-        echo "barrier_all"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 17 -x ${shm_ctx} > $3/barrier_all.log
-        check barrier_all
-        echo "sync_all"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 18 -x ${shm_ctx} > $3/sync_all.log
-        check sync_all
-        echo "sync"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 19 -x ${shm_ctx} > $3/sync.log
-        check sync
-        echo "alltoall"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 23 -x ${shm_ctx} > $3/alltoall.log
-        check alltoall
-        echo "fcollect"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 22 -x ${shm_ctx} > $3/fcollect.log
-        check fcollect
+
+    ###########################################################################
+    ############################### SHORT TESTS ###############################
+    ###########################################################################
+    *"short")
+        echo "get_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 0 > $3/get_n2_w16_z128_8B.log
+        check get_n2_w16_z128_8B
+        echo "getnbi_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 1 > $3/getnbi_n2_w16_z128_8B.log
+        check getnbi_n2_w16_z128_8B
+        echo "put_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 2 > $3/put_n2_w16_z128_8B.log
+        check put_n2_w16_z128_8B
+        echo "putnbi_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 3 > $3/putnbi_n2_w16_z128_8B.log
+        check putnbi_n2_w16_z128_8B
+        echo "amofadd_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 6 > $3/amofadd_n2_w8_z1.log
+        check amofadd_n2_w8_z1
+        echo "amofinc_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 7 > $3/amofinc_n2_w8_z1.log
+        check amofinc_n2_w8_z1
+        echo "amofetch_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 8 > $3/amofetch_n2_w8_z1.log
+        check amofetch_n2_w8_z1
+        echo "amofcswap_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 9 > $3/amofcswap_n2_w8_z1.log
+        check amofcswap_n2_w8_z1
+        echo "amoadd_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 10 > $3/amoadd_n2_w8_z1.log
+        check amoadd_n2_w8_z1
+        echo "amoinc_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 11 > $3/amoinc_n2_w8_z1.log
+        check amoinc_n2_w8_z1
+        echo "pingpong_n2_w1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -a 14 > $3/pingpong_n2_w1.log
+        check pingpong_n2_w1
+        echo "amoset_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 44 > $3/amoset_n2_w8_z1.log
+        check amoset_n2_w8_z1
         ;;
-    *"ro")
-        echo "get"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 0 -x ${shm_ctx} > $3/get.log
-        check get
-        echo "get_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 1 -x ${shm_ctx} > $3/get_nbi.log
-        check get_nbi
-        echo "put"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 2 -x ${shm_ctx} > $3/put.log
-        check put
-        echo "put_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 3 -x ${shm_ctx} > $3/put_nbi.log
-        check put_nbi
-#        echo "team_ctx_infra"
-#        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 42 -x ${shm_ctx} > $3/team_ctx_infra.log
-#        check team_ctx_infra
-        echo "team_ctx_put_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 41 -x ${shm_ctx} > $3/team_ctx_put_nbi.log
-        check team_ctx_put_nbi
-        echo "team_ctx_get_nbi"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 39 -x ${shm_ctx} > $3/team_ctx_get_nbi.log
-        check team_ctx_get_nbi
-#        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 4 -x ${shm_ctx} > $3/get_swarm.log
-#        check get_swarm
-        #mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 5 -x ${shm_ctx} > $3/reduction.log
-        echo "amo_fadd"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 6 -x ${shm_ctx} > $3/amo_fadd.log
-        check amo_fadd
-        echo "amo_finc"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 7 -x ${shm_ctx} > $3/amo_finc.log
-        check amo_finc
-        echo "amo_fetch"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 8 -x ${shm_ctx} > $3/amo_fetch.log
-        check amo_fetch
-        echo "amo_fcswap"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 9 -x ${shm_ctx} > $3/amo_fcswap.log
-        check amo_fcswap
-        echo "amo_add"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 10 -x ${shm_ctx} > $3/amo_add.log
-        check amo_add
-        echo "amo_set"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 44 -x ${shm_ctx} > $3/amo_set.log
-        check amo_set
-        echo "amo_swap"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 45 -x ${shm_ctx} > $3/amo_swap.log
-        check amo_swap
-        echo "amo_fetch_and"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 46 -x ${shm_ctx} > $3/amo_fetch_and.log
-        check amo_fetch_and
-        echo "amo_and"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 49 -x ${shm_ctx} > $3/amo_and.log
-        check amo_and
-        echo "amo_fetch_or"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 47 -x ${shm_ctx} > $3/amo_fetch_or.log
-        check amo_fetch_or
-        echo "amo_or"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 50 -x ${shm_ctx} > $3/amo_or.log
-        check amo_or
-        echo "amo_fetch_xor"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 47 -x ${shm_ctx} > $3/amo_fetch_xor.log
-        check amo_fetch_xor
-        echo "amo_xor"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 50 -x ${shm_ctx} > $3/amo_xor.log
-        check amo_xor
-        echo "amo_inc"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 11 -x ${shm_ctx} > $3/amo_inc.log
-        check amo_inc
-        echo "init"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 13 -x ${shm_ctx} > $3/init.log
-        check init
-        echo "ping_pong"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 14 -x ${shm_ctx} > $3/ping_pong.log
-        check ping_pong
-        echo "barrier_all"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 17 -x ${shm_ctx} > $3/barrier_all.log
-        check barrier_all
-        echo "sync_all"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 18 -x ${shm_ctx} > $3/sync_all.log
-        check sync_all
-        echo "sync"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 19 -x ${shm_ctx} > $3/sync.log
-        check sync
-        echo "alltoall"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 23 -x ${shm_ctx} > $3/alltoall.log
-        check alltoall
-        echo "fcollect"
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 22 -x ${shm_ctx} > $3/fcollect.log
-        check fcollect
-        ;;
-    *"team_ctx_get")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 38 -x ${shm_ctx}
-        ;;
-    *"team_ctx_get_nbi")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 39 -x ${shm_ctx}
-        ;;
-    *"team_ctx_put")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 40 -x ${shm_ctx}
-        ;;
-    *"team_ctx_put_nbi")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 41 -x ${shm_ctx}
+
+    ###########################################################################
+    ############################# EXHAUSTIVE TESTS ############################
+    ###########################################################################
+    *"exhaustive")
+        ############################### GET ###################################
+        echo "get_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 0 > $3/get_n2_w1_z1_1MB.log
+        check get_n2_w1_z1_1MB
+        echo "get_n2_w1_z1024_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -s 512 -a 0 > $3/get_n2_w1_z1024_512B.log
+        check get_n2_w1_z1024_512B
+        echo "get_n2_w8_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -s 1048576 -a 0 > $3/get_n2_w8_z1_1MB.log
+        check get_n2_w8_z1_1MB
+        echo "get_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 0 > $3/get_n2_w16_z128_8B.log
+        check get_n2_w16_z128_8B
+        echo "get_n2_w32_z256_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 256 -s 512 -a 0 > $3/get_n2_w32_z256_512B.log
+        check get_n2_w32_z256_512B
+        echo "get_n2_w64_z1024_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=64 mpirun -np 2 $1 -w 64 -z 1024 -s 8 -a 0 > $3/get_n2_w64_z1024_8B.log
+        check get_n2_w64_z1024_8B
+        ############################### GETNBI ################################
+        echo "getnbi_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 1 > $3/getnbi_n2_w1_z1_1MB.log
+        check getnbi_n2_w1_z1_1MB
+        echo "getnbi_n2_w1_z1024_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -s 512 -a 1 > $3/getnbi_n2_w1_z1024_512B.log
+        check getnbi_n2_w1_z1024_512B
+        echo "getnbi_n2_w8_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -s 1048576 -a 1 > $3/getnbi_n2_w8_z1_1MB.log
+        check getnbi_n2_w8_z1_1MB
+        echo "getnbi_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 1 > $3/getnbi_n2_w16_z128_8B.log
+        check getnbi_n2_w16_z128_8B
+        echo "getnbi_n2_w32_z256_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 256 -s 512 -a 1 > $3/getnbi_n2_w32_z256_512B.log
+        check getnbi_n2_w32_z256_512B
+        echo "getnbi_n2_w64_z1024_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=64 mpirun -np 2 $1 -w 64 -z 1024 -s 8 -a 1 > $3/getnbi_n2_w64_z1024_8B.log
+        check getnbi_n2_w64_z1024_8B
+        ############################### PUT ###################################
+        echo "put_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 2 > $3/put_n2_w1_z1_1MB.log
+        check put_n2_w1_z1_1MB
+        echo "put_n2_w1_z1024_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -s 512 -a 2 > $3/put_n2_w1_z1024_512B.log
+        check put_n2_w1_z1024_512B
+        echo "put_n2_w8_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -s 1048576 -a 2 > $3/put_n2_w8_z1_1MB.log
+        check put_n2_w8_z1_1MB
+        echo "put_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 2 > $3/put_n2_w16_z128_8B.log
+        check put_n2_w16_z128_8B
+        echo "put_n2_w32_z256_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 256 -s 512 -a 2 > $3/put_n2_w32_z256_512B.log
+        check put_n2_w32_z256_512B
+        echo "put_n2_w64_z1024_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=64 mpirun -np 2 $1 -w 64 -z 1024 -s 8 -a 2 > $3/put_n2_w64_z1024_8B.log
+        check put_n2_w64_z1024_8B
+        ############################### PUTNBI ################################
+        echo "putnbi_n2_w1_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 3 > $3/putnbi_n2_w1_z1_1MB.log
+        check putnbi_n2_w1_z1_1MB
+        echo "putnbi_n2_w1_z1024_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -s 512 -a 3 > $3/putnbi_n2_w1_z1024_512B.log
+        check putnbi_n2_w1_z1024_512B
+        echo "putnbi_n2_w8_z1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -s 1048576 -a 3 > $3/putnbi_n2_w8_z1_1MB.log
+        check putnbi_n2_w8_z1_1MB
+        echo "putnbi_n2_w16_z128_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=16 mpirun -np 2 $1 -w 16 -z 128 -s 8 -a 3 > $3/putnbi_n2_w16_z128_8B.log
+        check putnbi_n2_w16_z128_8B
+        echo "putnbi_n2_w32_z256_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 256 -s 512 -a 3 > $3/putnbi_n2_w32_z256_512B.log
+        check putnbi_n2_w32_z256_512B
+        echo "putnbi_n2_w64_z1024_8B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=64 mpirun -np 2 $1 -w 64 -z 1024 -s 8 -a 3 > $3/putnbi_n2_w64_z1024_8B.log
+        check putnbi_n2_w64_z1024_8B
+        ############################# REDUCTION ##############################
+        echo "reduction_n2_w1_z1_32K"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -s 32768 -a 5 > $3/reduction_n2_w1_z1_32K.log
+        check reduction_n2_w1_z1_32K
+        echo "reduction_n2_w8_z1_32K"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -s 32768 -a 5 > $3/reduction_n2_w8_z1_32K.log
+        check reduction_n2_w8_z1_32K
+        echo "reduction_n2_w32_z1_32K"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 1 -s 32768 -a 5 > $3/reduction_n2_w32_z1_32K.log
+        check reduction_n2_w32_z1_32K
+        ############################## AMOFADD ###############################
+        echo "amofadd_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 6 > $3/amofadd_n2_w1_z1.log
+        check amofadd_n2_w1_z1
+        echo "amofadd_n2_w1_z1024"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -a 6 > $3/amofadd_n2_w1_z1024.log
+        check amofadd_n2_w1_z1024
+        echo "amofadd_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 6 > $3/amofadd_n2_w8_z1.log
+        check amofadd_n2_w8_z1
+        echo "amofadd_n2_w32_z128"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 128 -a 6 > $3/amofadd_n2_w32_z128.log
+        check amofadd_n2_w32_z128
+        ############################## AMOFINC ###############################
+        echo "amofinc_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 7 > $3/amofinc_n2_w1_z1.log
+        check amofinc_n2_w1_z1
+        echo "amofinc_n2_w1_z1024"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -a 7 > $3/amofinc_n2_w1_z1024.log
+        check amofinc_n2_w1_z1024
+        echo "amofinc_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 7 > $3/amofinc_n2_w8_z1.log
+        check amofinc_n2_w8_z1
+        echo "amofinc_n2_w32_z128"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 128 -a 7 > $3/amofinc_n2_w32_z128.log
+        check amofinc_n2_w32_z128
+        ############################ AMOFETCH ################################
+        echo "amofetch_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 8 > $3/amofetch_n2_w1_z1.log
+        check amofetch_n2_w1_z1
+        echo "amofetch_n2_w1_z1024"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -a 8 > $3/amofetch_n2_w1_z1024.log
+        check amofetch_n2_w1_z1024
+        echo "amofetch_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 8 > $3/amofetch_n2_w8_z1.log
+        check amofetch_n2_w8_z1
+        echo "amofetch_n2_w32_z128"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 128 -a 8 > $3/amofetch_n2_w32_z128.log
+        check amofetch_n2_w32_z128
+        ########################### AMOFCSWAP ################################
+        echo "amofcswap_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 9 > $3/amofcswap_n2_w1_z1.log
+        check amofcswap_n2_w1_z1
+        echo "amofcswap_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 9 > $3/amofcswap_n2_w8_z1.log
+        check amofcswap_n2_w8_z1
+        echo "amofcswap_n2_w32_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 1 -a 9 > $3/amofcswap_n2_w32_z1.log
+        check amofcswap_n2_w32_z1
+        ############################# AMOADD ################################
+        echo "amoadd_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 10 > $3/amoadd_n2_w1_z1.log
+        check amoadd_n2_w1_z1
+        echo "amoadd_n2_w1_z1024"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -a 10 > $3/amoadd_n2_w1_z1024.log
+        check amoadd_n2_w1_z1024
+        echo "amoadd_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 10 > $3/amoadd_n2_w8_z1.log
+        check amoadd_n2_w8_z1
+        echo "amoadd_n2_w32_z128"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 128 -a 10 > $3/amoadd_n2_w32_z128.log
+        check amoadd_n2_w32_z128
+        ############################# AMOINC ################################
+        echo "amoinc_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 11 > $3/amoinc_n2_w1_z1.log
+        check amoinc_n2_w1_z1
+        echo "amoinc_n2_w1_z1024"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1024 -a 11 > $3/amoinc_n2_w1_z1024.log
+        check amoinc_n2_w1_z1024
+        echo "amoinc_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 11 > $3/amoinc_n2_w8_z1.log
+        check amoinc_n2_w8_z1
+        echo "amoinc_n2_w32_z128"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 128 -a 11 > $3/amoinc_n2_w32_z128.log
+        check amoinc_n2_w32_z128
+        ############################## INIT #################################
+        echo "init_n2"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -a 13 > $3/init_n2.log
+        check init_n2
+        ########################### PINGPONG ################################
+        echo "pingpong_n2_w1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -a 14 > $3/pingpong_n2_w1.log
+        check pingpong_n2_w1
+        echo "pingpong_n2_w8"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -a 14 > $3/pingpong_n2_w8.log
+        check pingpong_n2_w8
+        echo "pingpong_n2_w32"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -a 14 > $3/pingpong_n2_w32.log
+        check pingpong_n2_w32
+        ############################ BARRIER ################################
+        echo "barrier_n2_w1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -a 17 > $3/barrier_n2_w1.log
+        check barrier_n2_w1
+        echo "barrier_n2_w8"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -a 17 > $3/barrier_n2_w8.log
+        check barrier_n2_w8
+        echo "barrier_n2_w32"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -a 17 > $3/barrier_n2_w32.log
+        check barrier_n2_w32
+        ############################ SYNCALL ################################
+        echo "syncall_n2_w1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -a 18 > $3/syncall_n2_w1.log
+        check syncall_n2_w1
+        echo "syncall_n2_w8"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -a 18 > $3/syncall_n2_w8.log
+        check syncall_n2_w8
+        echo "syncall_n2_w32"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -a 18 > $3/syncall_n2_w32.log
+        check syncall_n2_w32
+        ############################# SYNC ##################################
+        echo "sync_n2_w1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -a 19 > $3/sync_n2_w1.log
+        check sync_n2_w1
+        echo "sync_n2_w8"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -a 19 > $3/sync_n2_w8.log
+        check sync_n2_w8
+        echo "sync_n2_w32"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -a 19 > $3/sync_n2_w32.log
+        check sync_n2_w32
+        ########################### FCOLLECT ################################
+        echo "fcollect_n2_w1_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -s 512 -a 22 > $3/fcollect_n2_w1_512B.log
+        check fcollect_n2_w1_512B
+        echo "fcollect_n2_w8_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -s 512 -a 22 > $3/fcollect_n2_w8_512B.log
+        check fcollect_n2_w8_512B
+        echo "fcollect_n2_w32_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -s 512 -a 22 > $3/fcollect_n2_w32_512B.log
+        check fcollect_n2_w32_512B
+        ########################### ALLTOALL ################################
+        echo "alltoall_n2_w1_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -s 512 -a 23 > $3/alltoall_n2_w1_512B.log
+        check alltoall_n2_w1_512B
+        echo "alltoall_n2_w8_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -s 512 -a 23 > $3/alltoall_n2_w8_512B.log
+        check alltoall_n2_w8_512B
+        echo "alltoall_n2_w32_512B"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -s 512 -a 23 > $3/alltoall_n2_w32_512B.log
+        check alltoall_n2_w32_512B
+        ########################## TEAMGETNBI ###############################
+        echo "teamgetnbi_n2_w1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -s 1048576 -a 39 > $3/teamgetnbi_n2_w1_1MB.log
+        check teamgetnbi_n2_w1_1MB
+        ########################## TEAMPUTNBI ###############################
+        echo "teamputnbi_n2_w1_1MB"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -s 1048576 -a 41 > $3/teamputnbi_n2_w1_1MB.log
+        check teamputnbi_n2_w1_1MB
+        ############################ AMOSET #################################
+        echo "amoset_n2_w1_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=1 mpirun -np 2 $1 -w 1 -z 1 -a 44 > $3/amoset_n2_w1_z1.log
+        check amoset_n2_w1_z1
+        echo "amoset_n2_w8_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=8 mpirun -np 2 $1 -w 8 -z 1 -a 44 > $3/amoset_n2_w8_z1.log
+        check amoset_n2_w8_z1
+        echo "amoset_n2_w32_z1"
+        ROC_SHMEM_MAX_NUM_CONTEXTS=32 mpirun -np 2 $1 -w 32 -z 1 -a 44 > $3/amoset_n2_w32_z1.log
+        check amoset_n2_w32_z1
         ;;
+
+    ###########################################################################
+    ############################# INDIVIDUAL TESTS ############################
+    ###########################################################################
     *"get")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 0 -x ${shm_ctx}
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 0
         ;;
-    *"get_nbi")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 1 -x ${shm_ctx}
+    *"getnbi")
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 1
         ;;
     *"put")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 2 -x ${shm_ctx}
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 2
         ;;
-    *"put_nbi")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 3 -x ${shm_ctx}
-        ;;
-    *"get_swarm")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 4 -x ${shm_ctx}
-        ;;
-    *"team_reduction")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 37 -x ${shm_ctx}
+    *"putnbi")
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 3
         ;;
     *"reduction")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 5 -x ${shm_ctx}
+        mpirun -np 2 $1 -w 1 -z 1 -s 32768 -a 5
         ;;
-    *"amo_fadd")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 6 -x ${shm_ctx}
+    *"amofadd")
+        mpirun -np 2 $1 -w 1 -z 1 -a 6
         ;;
-    *"amo_finc")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 7 -x ${shm_ctx}
+    *"amofinc")
+        mpirun -np 2 $1 -w 1 -z 1 -a 7
         ;;
-    *"amo_fetch")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 8 -x ${shm_ctx}
+    *"amofetch")
+        mpirun -np 2 $1 -w 1 -z 1 -a 8
         ;;
-    *"amo_fcswap")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 9 -x ${shm_ctx}
+    *"amofcswap")
+        mpirun -np 2 $1 -w 1 -z 1 -a 9
         ;;
-    *"amo_add")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 10 -x ${shm_ctx}
+    *"amoadd")
+        mpirun -np 2 $1 -w 1 -z 1 -a 10
         ;;
-    *"amo_set")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 44 -x ${shm_ctx}
-        ;;
-    *"amo_swap")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 45 -x ${shm_ctx}
-        ;;
-    *"amo_fetch_and")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 46 -x ${shm_ctx}
-        ;;
-    *"amo_and")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 49 -x ${shm_ctx}
-        ;;
-    *"amo_fetch_or")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 47 -x ${shm_ctx}
-        ;;
-    *"amo_or")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 50 -x ${shm_ctx}
-        ;;
-    *"amo_fetch_xor")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 48 -x ${shm_ctx}
-        ;;
-    *"amo_xor")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 51 -x ${shm_ctx}
-        ;;
-    *"amo_inc")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 11 -x ${shm_ctx}
+    *"amoinc")
+        mpirun -np 2 $1 -w 1 -z 1 -a 11
         ;;
     *"init")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 13 -x ${shm_ctx}
+        mpirun -np 2 $1 -a 13
         ;;
-    *"ping_pong")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 32768 -a 14 -x ${shm_ctx}
+    *"pingpong")
+        mpirun -np 2 $1 -w 1 -z 1 -a 14
         ;;
-    *"team_broadcast")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 36 -x ${shm_ctx}
+    *"barrier")
+        mpirun -np 2 $1 -w 1 -z 1 -a 17
         ;;
-    *"alltoall")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 23 -x ${shm_ctx}
-    ;;
-    *"fcollect")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 22 -x ${shm_ctx}
-    ;;
-    *"broadcast")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 512 -a 20 -x ${shm_ctx}
-        ;;
-    *"barrier_all")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 17 -x ${shm_ctx}
-        ;;
-    *"sync_all")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 18 -x ${shm_ctx}
+    *"syncall")
+        mpirun -np 2 $1 -w 1 -z 1 -a 18
         ;;
     *"sync")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 19 -x ${shm_ctx}
+        mpirun -np 2 $1 -w 1 -z 1 -s 8 -a 19
+        ;;
+    *"broadcast")
+        mpirun -np 2 $1 -w 1 -z 1 -s 32768 -a 20
+        ;;
+    *"fcollect")
+        mpirun -np 2 $1 -w 1 -z 1 -s 32768 -a 22
+        ;;
+    *"alltoall")
+        mpirun -np 2 $1 -w 1 -z 1 -s 32768 -a 23
+        ;;
+    *"team_broadcast")
+        mpirun -np 2 $1 -w 1 -z 1 -s 32768 -a 36
+        ;;
+    *"team_reduction")
+        mpirun -np 2 $1 -w 1 -z 1 -s 32768 -a 37
+        ;;
+    *"team_get")
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 38
+        ;;
+    *"team_getnbi")
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 39
+        ;;
+    *"team_put")
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 40
+        ;;
+    *"team_putnbi")
+        mpirun -np 2 $1 -w 1 -z 1 -s 1048576 -a 41
         ;;
     *"ctx_infra")
-        mpirun -np 2 ${gdb_cmd} $1 -t 1 -w 1 -s 8 -a 42 -x ${shm_ctx}
+        mpirun -np 2 $1 -w 1 -z 1 -a 42
+        ;;
+    *"amoset")
+        mpirun -np 2 $1 -w 1 -z 1 -a 44
+        ;;
+    *"amoswap")
+        mpirun -np 2 $1 -w 1 -z 1 -a 45
+        ;;
+    *"amofetchand")
+        mpirun -np 2 $1 -w 1 -z 1 -a 46
+        ;;
+    *"amofetchor")
+        mpirun -np 2 $1 -w 1 -z 1 -a 47
+        ;;
+    *"amofetchxor")
+        mpirun -np 2 $1 -w 1 -z 1 -a 48
+        ;;
+    *"amoand")
+        mpirun -np 2 $1 -w 1 -z 1 -a 49
+        ;;
+    *"amoor")
+        mpirun -np 2 $1 -w 1 -z 1 -a 50
+        ;;
+    *"amoxor")
+        mpirun -np 2 $1 -w 1 -z 1 -a 51
         ;;
     *)
         echo "UNKNOWN TEST TYPE: $2"

From 73303ca2d26658bb41a63224fea5a41545753fad Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Tue, 30 Jul 2024 12:56:32 -0500
Subject: [PATCH 4/6] Move inline assembly into arch defines blocks

---
 src/reverse_offload/context_ro_device.cpp |  8 ++++--
 src/util.cpp                              |  2 ++
 src/util.hpp                              | 30 ++++++++++++++++++++---
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/reverse_offload/context_ro_device.cpp b/src/reverse_offload/context_ro_device.cpp
index 6182ce0917..6c6ea8f809 100644
--- a/src/reverse_offload/context_ro_device.cpp
+++ b/src/reverse_offload/context_ro_device.cpp
@@ -391,20 +391,24 @@ __device__ bool enough_space(BlockHandle *h, uint64_t required) {
 
 __device__ void refresh_volatile_dwordx2(volatile uint64_t *assigned_value,
                                          volatile uint64_t *read_value) {
-  __asm__ volatile(
+#if __gfx90a__
+  asm volatile(
     "global_load_dwordx2 %0 %1 off glc slc\n "
     "s_waitcnt vmcnt(0)"
     : "=v"(*assigned_value)
     : "v"(read_value));
+#endif
 }
 
 __device__ void refresh_volatile_sbyte(volatile int *assigned_value,
                                        volatile char *read_value) {
-  __asm__ volatile(
+#if __gfx90a__
+  asm volatile(
     "global_load_sbyte %0 %1 off glc slc\n "
     "s_waitcnt vmcnt(0)"
     : "=v"(*assigned_value)
     : "v"(read_value));
+#endif
 }
 
 __device__ void acquire_lock(BlockHandle *handle) {
diff --git a/src/util.cpp b/src/util.cpp
index 11122d7b16..c68bfc55c1 100644
--- a/src/util.cpp
+++ b/src/util.cpp
@@ -42,10 +42,12 @@ std::vector<device_agent_t> cpu_agents;
 
 __device__ uint64_t __read_clock() {
   uint64_t clock{};
+#if __gfx90a__
   asm volatile(
       "s_memrealtime %0\n"
       "s_waitcnt lgkmcnt(0)\n"
       : "=s"(clock));
+#endif
   return clock;
 }
 
diff --git a/src/util.hpp b/src/util.hpp
index 1d5cd9f9a6..1fa9330509 100644
--- a/src/util.hpp
+++ b/src/util.hpp
@@ -43,11 +43,13 @@ namespace rocshmem {
 
 __device__ __forceinline__ int uncached_load_ubyte(uint8_t* src) {
   int ret;
-  __asm__ volatile(
+#if __gfx90a__
+  asm volatile(
       "global_load_ubyte %0 %1 off glc slc \n"
       "s_waitcnt vmcnt(0)"
       : "=v"(ret)
       : "v"(src));
+#endif
   return ret;
 }
 
@@ -62,18 +64,22 @@ NOWARN(-Wdeprecated-volatile,
     T ret;
     switch (sizeof(T)) {
       case 4:
-        __asm__ volatile(
+#if __gfx90a__
+        asm volatile(
             "global_load_dword %0 %1 off glc slc \n"
             "s_waitcnt vmcnt(0)"
             : "=v"(ret)
             : "v"(src));
+#endif
         break;
       case 8:
-        __asm__ volatile(
+#if __gfx90a__
+        asm volatile(
             "global_load_dwordx2 %0 %1 off glc slc \n"
             "s_waitcnt vmcnt(0)"
             : "=v"(ret)
             : "v"(src));
+#endif
         break;
       default:
         break;
@@ -123,7 +129,13 @@ NOWARN(-Wdeprecated-volatile,
 extern const int gpu_clock_freq_mhz;
 
 /* Device-side internal functions */
-__device__ __forceinline__ void __roc_inv() { asm volatile("buffer_wbinvl1;"); }
+__device__ __forceinline__ void __roc_inv() {
+#if defined USE_COHERENT_HEAP
+#if __gfx90a__
+  asm volatile("buffer_wbinvl1;");
+#endif
+#endif
+}
 
 __device__ __forceinline__ void __roc_flush() {
 #if defined USE_COHERENT_HEAP
@@ -131,6 +143,10 @@ __device__ __forceinline__ void __roc_flush() {
   asm volatile("s_dcache_wb;");
   asm volatile("buffer_wbl2;");
 #endif
+#if __gfx90a__
+  asm volatile("s_dcache_wb;");
+  asm volatile("buffer_wbl2;");
+#endif
 #endif
 }
 
@@ -227,17 +243,23 @@ __device__ __forceinline__ void store_asm(uint8_t* val, uint8_t* dst,
   switch (size) {
     case 2: {
       int16_t val16{*(reinterpret_cast<int16_t*>(val))};
+#if __gfx90a__
       asm volatile("flat_store_short %0 %1 glc slc" : : "v"(dst), "v"(val16));
+#endif
       break;
     }
     case 4: {
       int32_t val32{*(reinterpret_cast<int32_t*>(val))};
+#if __gfx90a__
       asm volatile("flat_store_dword %0 %1 glc slc" : : "v"(dst), "v"(val32));
+#endif
       break;
     }
     case 8: {
       int64_t val64{*(reinterpret_cast<int64_t*>(val))};
+#if __gfx90a__
       asm volatile("flat_store_dwordx2 %0 %1 glc slc" : : "v"(dst), "v"(val64));
+#endif
       break;
     }
     default:

From 862ef5713f0aebbba5292fc6bd23d6ab33eb1b14 Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Tue, 30 Jul 2024 14:53:19 -0500
Subject: [PATCH 5/6] Move inline asm into separate file

---
 src/assembly.hpp                          | 236 ++++++++++++++++++++++
 src/reverse_offload/context_ro_device.cpp |  22 --
 src/util.cpp                              |  11 -
 src/util.hpp                              | 110 +---------
 4 files changed, 237 insertions(+), 142 deletions(-)
 create mode 100644 src/assembly.hpp

diff --git a/src/assembly.hpp b/src/assembly.hpp
new file mode 100644
index 0000000000..8d69249561
--- /dev/null
+++ b/src/assembly.hpp
@@ -0,0 +1,236 @@
+/******************************************************************************
+ * Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ *****************************************************************************/
+
+#ifndef LIBRARY_SRC_ASSEMBLY_HPP_
+#define LIBRARY_SRC_ASSEMBLY_HPP_
+
+#include <hip/hip_runtime.h>
+#include <hsa/hsa.h>
+#include <hsa/hsa_ext_amd.h>
+
+namespace rocshmem {
+
+#define DO_PRAGMA(x) _Pragma(#x)
+#define NOWARN(warnoption, ...)                 \
+  DO_PRAGMA(GCC diagnostic push)                \
+  DO_PRAGMA(GCC diagnostic ignored #warnoption) \
+  __VA_ARGS__                                   \
+  DO_PRAGMA(GCC diagnostic pop)
+
+#define SFENCE() asm volatile("sfence" ::: "memory")
+
+__device__ __forceinline__ int uncached_load_ubyte(uint8_t* src) {
+  int ret;
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+  asm volatile(
+      "global_load_ubyte %0 %1 off glc slc \n"
+      "s_waitcnt vmcnt(0)"
+      : "=v"(ret)
+      : "v"(src));
+#endif
+#if __gfx940__
+#endif
+  return ret;
+}
+
+__device__ __forceinline__ void refresh_volatile_sbyte(volatile int *assigned_value,
+                                                       volatile char *read_value) {
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+  asm volatile(
+    "global_load_sbyte %0 %1 off glc slc\n "
+    "s_waitcnt vmcnt(0)"
+    : "=v"(*assigned_value)
+    : "v"(read_value));
+#endif
+#if __gfx940__
+#endif
+}
+
+__device__ __forceinline__ void refresh_volatile_dwordx2(volatile uint64_t *assigned_value,
+                                                         volatile uint64_t *read_value) {
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+  asm volatile(
+    "global_load_dwordx2 %0 %1 off glc slc\n "
+    "s_waitcnt vmcnt(0)"
+    : "=v"(*assigned_value)
+    : "v"(read_value));
+#endif
+#if __gfx940__
+#endif
+}
+
+/* Ignore the warning about deprecated volatile.
+ * The only usage of volatile is to force the compiler to generate
+ * the assembly instruction. If volatile is omitted, the compiler
+ * will NOT generate the non-temporal load or the waitcnt.
+ */
+// clang-format off
+NOWARN(-Wdeprecated-volatile,
+  template <typename T> __device__ __forceinline__ T uncached_load(T* src) {
+    T ret;
+    switch (sizeof(T)) {
+      case 4:
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+        asm volatile(
+            "global_load_dword %0 %1 off glc slc \n"
+            "s_waitcnt vmcnt(0)"
+            : "=v"(ret)
+            : "v"(src));
+#endif
+#if __gfx940__
+#endif
+        break;
+      case 8:
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+        asm volatile(
+            "global_load_dwordx2 %0 %1 off glc slc \n"
+            "s_waitcnt vmcnt(0)"
+            : "=v"(ret)
+            : "v"(src));
+#endif
+#if __gfx940__
+#endif
+        break;
+      default:
+        break;
+    }
+    return ret;
+  }
+)
+// clang-format on
+
+__device__ __forceinline__ void __roc_inv() {
+#if defined USE_COHERENT_HEAP
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+  asm volatile("buffer_wbinvl1;");
+#endif
+#if __gfx940__
+#endif
+#endif
+}
+
+__device__ __forceinline__ void __roc_flush() {
+#if defined USE_COHERENT_HEAP
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+  asm volatile("s_dcache_wb;");
+  asm volatile("buffer_wbl2;");
+#endif
+#if __gfx940__
+#endif
+#endif
+}
+
+__device__ __forceinline__ void store_asm(uint8_t* val, uint8_t* dst,
+                                          int size) {
+  switch (size) {
+    case 2: {
+      int16_t val16{*(reinterpret_cast<int16_t*>(val))};
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+      asm volatile("flat_store_short %0 %1 glc slc" : : "v"(dst), "v"(val16));
+#endif
+#if __gfx940__
+#endif
+      break;
+    }
+    case 4: {
+      int32_t val32{*(reinterpret_cast<int32_t*>(val))};
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+      asm volatile("flat_store_dword %0 %1 glc slc" : : "v"(dst), "v"(val32));
+#endif
+#if __gfx940__
+#endif
+      break;
+    }
+    case 8: {
+      int64_t val64{*(reinterpret_cast<int64_t*>(val))};
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+      asm volatile("flat_store_dwordx2 %0 %1 glc slc" : : "v"(dst), "v"(val64));
+#endif
+#if __gfx940__
+#endif
+      break;
+    }
+    default:
+      break;
+  }
+}
+
+__device__ __forceinline__ uint64_t __read_clock() {
+  uint64_t clock{};
+#if __gfx906__
+#endif
+#if __gfx908__
+#endif
+#if __gfx90a__
+  asm volatile(
+      "s_memrealtime %0\n"
+      "s_waitcnt lgkmcnt(0)\n"
+      : "=s"(clock));
+#endif
+#if __gfx940__
+#endif
+  return clock;
+}
+
+}  // namespace rocshmem
+
+#endif  // LIBRARY_SRC_ASSEMBLY_HPP_
diff --git a/src/reverse_offload/context_ro_device.cpp b/src/reverse_offload/context_ro_device.cpp
index 6c6ea8f809..8e933e9777 100644
--- a/src/reverse_offload/context_ro_device.cpp
+++ b/src/reverse_offload/context_ro_device.cpp
@@ -389,28 +389,6 @@ __device__ bool enough_space(BlockHandle *h, uint64_t required) {
   return (h->queue_size - (h->write_index - h->read_index)) >= required;
 }
 
-__device__ void refresh_volatile_dwordx2(volatile uint64_t *assigned_value,
-                                         volatile uint64_t *read_value) {
-#if __gfx90a__
-  asm volatile(
-    "global_load_dwordx2 %0 %1 off glc slc\n "
-    "s_waitcnt vmcnt(0)"
-    : "=v"(*assigned_value)
-    : "v"(read_value));
-#endif
-}
-
-__device__ void refresh_volatile_sbyte(volatile int *assigned_value,
-                                       volatile char *read_value) {
-#if __gfx90a__
-  asm volatile(
-    "global_load_sbyte %0 %1 off glc slc\n "
-    "s_waitcnt vmcnt(0)"
-    : "=v"(*assigned_value)
-    : "v"(read_value));
-#endif
-}
-
 __device__ void acquire_lock(BlockHandle *handle) {
   while(atomicCAS((uint64_t *)&handle->lock, 0, 1) == 1) ;
 }
diff --git a/src/util.cpp b/src/util.cpp
index c68bfc55c1..7318b14208 100644
--- a/src/util.cpp
+++ b/src/util.cpp
@@ -40,17 +40,6 @@ typedef struct device_agent {
 std::vector<device_agent_t> gpu_agents;
 std::vector<device_agent_t> cpu_agents;
 
-__device__ uint64_t __read_clock() {
-  uint64_t clock{};
-#if __gfx90a__
-  asm volatile(
-      "s_memrealtime %0\n"
-      "s_waitcnt lgkmcnt(0)\n"
-      : "=s"(clock));
-#endif
-  return clock;
-}
-
 hsa_status_t rocm_hsa_amd_memory_pool_callback(
     hsa_amd_memory_pool_t memory_pool, void* data) {
   hsa_amd_memory_pool_global_flag_t pool_flag{};
diff --git a/src/util.hpp b/src/util.hpp
index 1fa9330509..c967ec1a9c 100644
--- a/src/util.hpp
+++ b/src/util.hpp
@@ -29,66 +29,12 @@
 
 #include <cstdio>
 
+#include "assembly.hpp"
 #include "config.h"  // NOLINT(build/include_subdir)
 #include "constants.hpp"
 
 namespace rocshmem {
 
-#define DO_PRAGMA(x) _Pragma(#x)
-#define NOWARN(warnoption, ...)                 \
-  DO_PRAGMA(GCC diagnostic push)                \
-  DO_PRAGMA(GCC diagnostic ignored #warnoption) \
-  __VA_ARGS__                                   \
-  DO_PRAGMA(GCC diagnostic pop)
-
-__device__ __forceinline__ int uncached_load_ubyte(uint8_t* src) {
-  int ret;
-#if __gfx90a__
-  asm volatile(
-      "global_load_ubyte %0 %1 off glc slc \n"
-      "s_waitcnt vmcnt(0)"
-      : "=v"(ret)
-      : "v"(src));
-#endif
-  return ret;
-}
-
-/* Ignore the warning about deprecated volatile.
- * The only usage of volatile is to force the compiler to generate
- * the assembly instruction. If volatile is omitted, the compiler
- * will NOT generate the non-temporal load or the waitcnt.
- */
-// clang-format off
-NOWARN(-Wdeprecated-volatile,
-  template <typename T> __device__ __forceinline__ T uncached_load(T* src) {
-    T ret;
-    switch (sizeof(T)) {
-      case 4:
-#if __gfx90a__
-        asm volatile(
-            "global_load_dword %0 %1 off glc slc \n"
-            "s_waitcnt vmcnt(0)"
-            : "=v"(ret)
-            : "v"(src));
-#endif
-        break;
-      case 8:
-#if __gfx90a__
-        asm volatile(
-            "global_load_dwordx2 %0 %1 off glc slc \n"
-            "s_waitcnt vmcnt(0)"
-            : "=v"(ret)
-            : "v"(src));
-#endif
-        break;
-      default:
-        break;
-    }
-    return ret;
-  }
-)
-// clang-format on
-
 #define LOAD(VAR) __atomic_load_n((VAR), __ATOMIC_SEQ_CST)
 #define STORE(DST, SRC) __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST)
 
@@ -102,8 +48,6 @@ NOWARN(-Wdeprecated-volatile,
     }                                                                         \
   }
 
-#define SFENCE() asm volatile("sfence" ::: "memory")
-
 #ifdef DEBUG
 #define DPRINTF(...)     \
   do {                   \
@@ -129,29 +73,6 @@ NOWARN(-Wdeprecated-volatile,
 extern const int gpu_clock_freq_mhz;
 
 /* Device-side internal functions */
-__device__ __forceinline__ void __roc_inv() {
-#if defined USE_COHERENT_HEAP
-#if __gfx90a__
-  asm volatile("buffer_wbinvl1;");
-#endif
-#endif
-}
-
-__device__ __forceinline__ void __roc_flush() {
-#if defined USE_COHERENT_HEAP
-#if __gfx90a__
-  asm volatile("s_dcache_wb;");
-  asm volatile("buffer_wbl2;");
-#endif
-#if __gfx90a__
-  asm volatile("s_dcache_wb;");
-  asm volatile("buffer_wbl2;");
-#endif
-#endif
-}
-
-__device__ uint64_t __read_clock();
-
 __device__ __forceinline__ uint32_t lowerID() {
   return __ffsll(__ballot(1)) - 1;
 }
@@ -238,35 +159,6 @@ __device__ void gpu_dprintf(const char* fmt, const Args&... args) {
   }
 }
 
-__device__ __forceinline__ void store_asm(uint8_t* val, uint8_t* dst,
-                                          int size) {
-  switch (size) {
-    case 2: {
-      int16_t val16{*(reinterpret_cast<int16_t*>(val))};
-#if __gfx90a__
-      asm volatile("flat_store_short %0 %1 glc slc" : : "v"(dst), "v"(val16));
-#endif
-      break;
-    }
-    case 4: {
-      int32_t val32{*(reinterpret_cast<int32_t*>(val))};
-#if __gfx90a__
-      asm volatile("flat_store_dword %0 %1 glc slc" : : "v"(dst), "v"(val32));
-#endif
-      break;
-    }
-    case 8: {
-      int64_t val64{*(reinterpret_cast<int64_t*>(val))};
-#if __gfx90a__
-      asm volatile("flat_store_dwordx2 %0 %1 glc slc" : : "v"(dst), "v"(val64));
-#endif
-      break;
-    }
-    default:
-      break;
-  }
-}
-
 __device__ __forceinline__ void memcpy(void* dst, void* src, size_t size) {
   uint8_t* dst_bytes{static_cast<uint8_t*>(dst)};
   uint8_t* src_bytes{static_cast<uint8_t*>(src)};

From bb6e48dded1a6088c2bd0404faf9f059cac80e90 Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Wed, 31 Jul 2024 14:14:24 -0500
Subject: [PATCH 6/6] Add assembly for gfx94x

---
 src/assembly.hpp | 133 ++++++++++++++++++++++++++++++-----------------
 1 file changed, 84 insertions(+), 49 deletions(-)

diff --git a/src/assembly.hpp b/src/assembly.hpp
index 8d69249561..076600067e 100644
--- a/src/assembly.hpp
+++ b/src/assembly.hpp
@@ -40,53 +40,68 @@ namespace rocshmem {
 
 __device__ __forceinline__ int uncached_load_ubyte(uint8_t* src) {
   int ret;
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
   asm volatile(
       "global_load_ubyte %0 %1 off glc slc \n"
       "s_waitcnt vmcnt(0)"
       : "=v"(ret)
       : "v"(src));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+  asm volatile(
+      "global_load_ubyte %0 %1 off sc0 sc1 \n"
+      "s_waitcnt vmcnt(0)"
+      : "=v"(ret)
+      : "v"(src));
 #endif
   return ret;
 }
 
 __device__ __forceinline__ void refresh_volatile_sbyte(volatile int *assigned_value,
                                                        volatile char *read_value) {
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
   asm volatile(
     "global_load_sbyte %0 %1 off glc slc\n "
     "s_waitcnt vmcnt(0)"
     : "=v"(*assigned_value)
     : "v"(read_value));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+  asm volatile(
+    "global_load_sbyte %0 %1 off sc0 sc1\n "
+    "s_waitcnt vmcnt(0)"
+    : "=v"(*assigned_value)
+    : "v"(read_value));
 #endif
 }
 
 __device__ __forceinline__ void refresh_volatile_dwordx2(volatile uint64_t *assigned_value,
                                                          volatile uint64_t *read_value) {
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
   asm volatile(
     "global_load_dwordx2 %0 %1 off glc slc\n "
     "s_waitcnt vmcnt(0)"
     : "=v"(*assigned_value)
     : "v"(read_value));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+  asm volatile(
+    "global_load_dwordx2 %0 %1 off sc0 sc1\n "
+    "s_waitcnt vmcnt(0)"
+    : "=v"(*assigned_value)
+    : "v"(read_value));
 #endif
 }
 
@@ -101,33 +116,43 @@ NOWARN(-Wdeprecated-volatile,
     T ret;
     switch (sizeof(T)) {
       case 4:
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
         asm volatile(
             "global_load_dword %0 %1 off glc slc \n"
             "s_waitcnt vmcnt(0)"
             : "=v"(ret)
             : "v"(src));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+        asm volatile(
+            "global_load_dword %0 %1 off sc0 sc1 \n"
+            "s_waitcnt vmcnt(0)"
+            : "=v"(ret)
+            : "v"(src));
 #endif
         break;
+#if defined(__gfx906__)
+#endif
+#if defined(__gfx908__)
+#endif
+#if defined(__gfx90a__)
       case 8:
-#if __gfx906__
-#endif
-#if __gfx908__
-#endif
-#if __gfx90a__
         asm volatile(
             "global_load_dwordx2 %0 %1 off glc slc \n"
             "s_waitcnt vmcnt(0)"
             : "=v"(ret)
             : "v"(src));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+        asm volatile(
+            "global_load_dwordx2 %0 %1 off sc0 sc1 \n"
+            "s_waitcnt vmcnt(0)"
+            : "=v"(ret)
+            : "v"(src));
 #endif
         break;
       default:
@@ -140,29 +165,32 @@ NOWARN(-Wdeprecated-volatile,
 
 __device__ __forceinline__ void __roc_inv() {
 #if defined USE_COHERENT_HEAP
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
-  asm volatile("buffer_wbinvl1;");
+#if defined(__gfx90a__)
+//  asm volatile("buffer_wbinvl1;");
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+//  asm volatile("buffer_inv sc0 sc1;");
 #endif
 #endif
 }
 
 __device__ __forceinline__ void __roc_flush() {
 #if defined USE_COHERENT_HEAP
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
-  asm volatile("s_dcache_wb;");
-  asm volatile("buffer_wbl2;");
+#if defined(__gfx90a__)
+//  asm volatile("s_dcache_wb;");
+//  asm volatile("buffer_wbl2;");
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+//  asm volatile("s_dcache_wb;");
+//  asm volatile("buffer_wbl2;");
 #endif
 #endif
 }
@@ -172,40 +200,43 @@ __device__ __forceinline__ void store_asm(uint8_t* val, uint8_t* dst,
   switch (size) {
     case 2: {
       int16_t val16{*(reinterpret_cast<int16_t*>(val))};
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
       asm volatile("flat_store_short %0 %1 glc slc" : : "v"(dst), "v"(val16));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+      asm volatile("flat_store_short %0 %1 sc0 sc1" : : "v"(dst), "v"(val16));
 #endif
       break;
     }
     case 4: {
       int32_t val32{*(reinterpret_cast<int32_t*>(val))};
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
       asm volatile("flat_store_dword %0 %1 glc slc" : : "v"(dst), "v"(val32));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+      asm volatile("flat_store_dword %0 %1 sc0 sc1" : : "v"(dst), "v"(val32));
 #endif
       break;
     }
     case 8: {
       int64_t val64{*(reinterpret_cast<int64_t*>(val))};
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
       asm volatile("flat_store_dwordx2 %0 %1 glc slc" : : "v"(dst), "v"(val64));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+      asm volatile("flat_store_dwordx2 %0 %1 sc0 sc1" : : "v"(dst), "v"(val64));
 #endif
       break;
     }
@@ -216,17 +247,21 @@ __device__ __forceinline__ void store_asm(uint8_t* val, uint8_t* dst,
 
 __device__ __forceinline__ uint64_t __read_clock() {
   uint64_t clock{};
-#if __gfx906__
+#if defined(__gfx906__)
 #endif
-#if __gfx908__
+#if defined(__gfx908__)
 #endif
-#if __gfx90a__
+#if defined(__gfx90a__)
   asm volatile(
       "s_memrealtime %0\n"
       "s_waitcnt lgkmcnt(0)\n"
       : "=s"(clock));
 #endif
-#if __gfx940__
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+  asm volatile(
+      "s_memrealtime %0\n"
+      "s_waitcnt lgkmcnt(0)\n"
+      : "=s"(clock));
 #endif
   return clock;
 }