Feat: add mori_shmem_barrier_all to wrapper

Author: jhchouuu

include/mori/shmem/shmem_device_api_wrapper.hpp

@@ -44,6 +44,9 @@ __device__ __attribute__((visibility("default"))) void mori_shmem_fence_thread_p
 __device__ __attribute__((visibility("default"))) void mori_shmem_fence_thread_pe_qp(int pe,
                                                                                      int qpId);
 
+__device__ __attribute__((visibility("default"))) void mori_shmem_barrier_all_thread();
+__device__ __attribute__((visibility("default"))) void mori_shmem_barrier_all_block();
+
 // ============================================================================
 // PutNbi APIs - Thread Scope (Address-based only)
 // ============================================================================

setup.py

@@ -86,6 +86,7 @@ def build_extension(self, ext: Extension) -> None:
         build_type = os.environ.get("CMAKE_BUILD_TYPE", "Release")
         unroll_value = os.environ.get("WARP_ACCUM_UNROLL", "1")
         use_bnxt = os.environ.get("USE_BNXT", "OFF")
+        build_shmem_device_wrapper = os.environ.get("BUILD_SHMEM_DEVICE_WRAPPER", "ON")
         use_ionic = os.environ.get("USE_IONIC", "OFF")
         enable_profiler = os.environ.get("ENABLE_PROFILER", "OFF")
         enable_debug_printf = os.environ.get("ENABLE_DEBUG_PRINTF", "OFF")
@@ -97,6 +98,7 @@ def build_extension(self, ext: Extension) -> None:
                 f"-DCMAKE_BUILD_TYPE={build_type}",
                 f"-DWARP_ACCUM_UNROLL={unroll_value}",
                 f"-DUSE_BNXT={use_bnxt}",
+                f"-DBUILD_SHMEM_DEVICE_WRAPPER={build_shmem_device_wrapper}",
                 f"-DUSE_IONIC={use_ionic}",
                 f"-DENABLE_DEBUG_PRINTF={enable_debug_printf}",
                 f"-DGPU_TARGETS={gpu_archs}",

src/shmem/CMakeLists.txt

@@ -4,6 +4,7 @@ option(BUILD_SHMEM_DEVICE_WRAPPER "Build SHMEM device API wrapper" OFF)
 
 if(BUILD_SHMEM_DEVICE_WRAPPER)
   add_library(mori_shmem init.cpp memory.cpp shmem_device_api_wrapper.cpp)
+  message(STATUS "BUILD_SHMEM_DEVICE_WRAPPER enabled for mori_shmem")
 else()
   add_library(mori_shmem init.cpp memory.cpp)
 endif()

src/shmem/shmem_device_api_wrapper.cpp

@@ -57,6 +57,14 @@ __device__ __attribute__((visibility("default"))) void mori_shmem_fence_thread_p
   mori::shmem::ShmemFenceThread(pe, qpId);
 }
 
+__device__ __attribute__((visibility("default"))) void mori_shmem_barrier_all_thread() {
+  mori::shmem::ShmemBarrierAllThread();
+}
+
+__device__ __attribute__((visibility("default"))) void mori_shmem_barrier_all_block() {
+  mori::shmem::ShmemBarrierAllBlock();
+}
+
 // ============================================================================
 // PutNbi APIs - Address-based only
 // ============================================================================