Proper API locations

arbor-sim · Jan 10, 2025 · c423b38 · c423b38
1 parent 0312e2e
commit c423b38
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/arbor/include/arbor/gpu/cuda_api.hpp b/arbor/include/arbor/gpu/cuda_api.hpp
@@ -138,11 +138,14 @@ inline float gpu_atomic_sub(float* address, float val) {
 }
 
 /// Warp-Level Primitives
-
 __device__ __inline__ unsigned ballot(unsigned mask, unsigned is_root) {
     return __ballot_sync(mask, is_root);
 }
 
+__device__ __inline__ unsigned active_mask() {
+    return __active_mask();
+}
+
 __device__ __inline__ unsigned any(unsigned mask, unsigned width) {
     return __any_sync(mask, width);
 }

diff --git a/arbor/include/arbor/gpu/hip_api.hpp b/arbor/include/arbor/gpu/hip_api.hpp
@@ -122,8 +122,7 @@ inline float gpu_atomic_sub(float* address, float val) {
 template<typename T>
 __device__ __inline__
 std::enable_if_t< !std::is_same_v<std::decay_t<T>, double>, std::decay_t<T>>
-shfl(T x, int lane)
-{
+shfl(T x, int lane) {
     return __shfl(x, lane);
 }
 
@@ -142,6 +141,10 @@ __device__ __inline__ unsigned ballot(unsigned mask, unsigned is_root) {
     return __ballot(is_root);
 }
 
+__device__ __inline__ unsigned active_mask() {
+    return __active_mask();
+}
+
 __device__ __inline__ unsigned any(unsigned mask, unsigned width) {
     return __any(width);
 }