From a57541b477a8c23fd6c5d697c88ae1cf5d3e73e5 Mon Sep 17 00:00:00 2001
From: shaoxunli <446032208@qq.com>
Date: Thu, 30 Oct 2025 09:31:02 +0800
Subject: [PATCH] add shmemx_barrier_on_stream and shmemx_barrier_all_on_stream

---
 include/host/shmem_host_sync.h | 13 +++++++++++++
 src/host/sync/shmemi_sync.cpp  |  4 ++--
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/include/host/shmem_host_sync.h b/include/host/shmem_host_sync.h
index 6bd4d7b7..b8b24f0e 100644
--- a/include/host/shmem_host_sync.h
+++ b/include/host/shmem_host_sync.h
@@ -43,6 +43,19 @@ SHMEM_HOST_API uint64_t shmemx_get_ffts_config();
  */
 SHMEM_HOST_API void shmem_handle_wait(shmem_handle_t handle, aclrtStream stream);
 
+/**
+ * @brief The shmemx_barrier_on_stream is a collective synchronization routine over a team.
+ * @param tid              [in] team to do barrier
+ * @param stream           [in] copy used stream (use default stream if stream == NULL)
+ */
+SHMEM_HOST_API void shmemx_barrier_on_stream(shmem_team_t tid, aclrtStream stream);
+
+/**
+ * @brief The shmemx_barrier_all_on_stream routine is a mechanism for synchronizing all PEs at once.
+ * @param stream           [in] copy used stream (use default stream if stream == NULL)
+ */
+SHMEM_HOST_API void shmemx_barrier_all_on_stream(aclrtStream stream);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/host/sync/shmemi_sync.cpp b/src/host/sync/shmemi_sync.cpp
index 7bca1332..1312a267 100644
--- a/src/host/sync/shmemi_sync.cpp
+++ b/src/host/sync/shmemi_sync.cpp
@@ -47,12 +47,12 @@ void shmem_barrier_all()
     shmem_barrier(SHMEM_TEAM_WORLD);
 }
 
-void shmem_barrier_on_stream(shmem_team_t tid, aclrtStream stream)
+void shmemx_barrier_on_stream(shmem_team_t tid, aclrtStream stream)
 {
     shmemi_barrier_on_stream(tid, stream);
 }
 
-void shmem_barrier_all_on_stream(aclrtStream stream)
+void shmemx_barrier_all_on_stream(aclrtStream stream)
 {
     shmemi_barrier_on_stream(SHMEM_TEAM_WORLD, stream);
 }
-- 
Gitee