From c423b38a2465d2de54c49bf69248299ec53bb30c Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Fri, 10 Jan 2025 15:32:01 +0100
Subject: [PATCH] Proper API locations

---
 arbor/include/arbor/gpu/cuda_api.hpp | 5 ++++-
 arbor/include/arbor/gpu/hip_api.hpp  | 7 +++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/arbor/include/arbor/gpu/cuda_api.hpp b/arbor/include/arbor/gpu/cuda_api.hpp
index 532b9ec8f9..18ee8a237e 100644
--- 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
index 019f235a8e..5c33ae9799 100644
--- 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);
 }