From a5b164bb77aa49c481a7cbbb8805b4fd23658599 Mon Sep 17 00:00:00 2001
From: Nara <nara@streamhpc.com>
Date: Mon, 7 Jul 2025 19:10:41 +0200
Subject: [PATCH] refactor(rocrand): replace hip vector accessor macro and use
 underlying native vector (#446)

Original change by @NB4444, I'm porting this over from our internal
repository. This changes the original fix by @umfranzw s.t. it uses
get_native_vector() instead of the newer operator[] accessor.
get_native_vector() would result in better generated code since
operator[] may spill to memory (see
https://github.com/ROCm/clr/issues/167).

(Already reviewed at: https://github.com/ROCm/rocm-libraries/pull/203)

Co-authored-by: Nick Breed <nick@streamhpc.com>
---
 .../library/include/rocrand/rocrand_common.h  | 63 +++++++++++--------
 .../include/rocrand/rocrand_philox4x32_10.h   |  7 +--
 .../include/rocrand/rocrand_threefry2_impl.h  | 10 ++-
 .../include/rocrand/rocrand_threefry4_impl.h  |  7 +--
 4 files changed, 45 insertions(+), 42 deletions(-)

diff --git a/projects/rocrand/library/include/rocrand/rocrand_common.h b/projects/rocrand/library/include/rocrand/rocrand_common.h
index e4c3df65ebc4..40a6062988f0 100644
--- a/projects/rocrand/library/include/rocrand/rocrand_common.h
+++ b/projects/rocrand/library/include/rocrand/rocrand_common.h
@@ -87,32 +87,6 @@
     #define ROCRAND_DEPRECATED(msg)
 #endif
 
-// This is an accessor macro for HIP vector types (eg. int2, float4, etc.).
-// Prior to HIP 7.0, individual elements could be accessed through the
-// data member:
-//
-// int2 vec;
-// vec.data[0] = 1;
-//
-// Beginning with HIP 7.0, the data member is hidden, and individual
-// elements must be accessed like this:
-//
-// int2 vec;
-// vec[0] = 1;
-//
-// You can use the macro like this:
-//
-// int2 vec;
-// ROCRAND_HIPVEC_ACCESS(vec)[0];
-//
-#if defined(__HIP_PLATFORM_AMD__)
-    #if HIP_VERSION_MAJOR < 7
-        #define ROCRAND_HIPVEC_ACCESS(x) x.data
-    #else
-        #define ROCRAND_HIPVEC_ACCESS(x) x
-    #endif
-#endif
-
 namespace rocrand_device {
 namespace detail {
 
@@ -188,6 +162,43 @@ __forceinline__ __device__ __host__ void
     hi = 0;
 }
 
+template<typename T, typename = void>
+struct has_data : std::false_type
+{};
+
+template<typename T>
+struct has_data<T, decltype(std::declval<T>().data, void())> : std::true_type
+{};
+
+template<typename T>
+static constexpr bool has_data_v = has_data<T>::value;
+
+/// This is an accessor utility for HIP vector types (eg. int2, float4, etc.).
+/// Prior to HIP 7.0, individual elements could be accessed through the
+/// data member:
+///   int2 vec;
+///   vec.data[0] = 1;
+///
+/// Beginning with HIP 7.0, the data member is hidden, and individual
+/// elements can instead be accessed through the underlaying native vector.
+template<typename V, typename Index>
+__forceinline__ __device__ __host__
+auto get_element_at(V vector, Index i)
+{
+#if defined(__HIP_PLATFORM_AMD__)
+    if constexpr(has_data_v<V>)
+    {
+        return vector.data[i];
+    }
+    else
+    {
+        return get_native_vector(vector)[i];
+    }
+#else
+    return (&vector.x)[i];
+#endif
+}
+
 } // end namespace detail
 } // end namespace rocrand_device
 
diff --git a/projects/rocrand/library/include/rocrand/rocrand_philox4x32_10.h b/projects/rocrand/library/include/rocrand/rocrand_philox4x32_10.h
index 022e8a2f389d..d5f850291081 100644
--- a/projects/rocrand/library/include/rocrand/rocrand_philox4x32_10.h
+++ b/projects/rocrand/library/include/rocrand/rocrand_philox4x32_10.h
@@ -171,11 +171,8 @@ class philox4x32_10_engine
 
     __forceinline__ __device__ __host__ unsigned int next()
     {
-    #if defined(__HIP_PLATFORM_AMD__)
-        unsigned int ret = ROCRAND_HIPVEC_ACCESS(m_state.result)[m_state.substate];
-    #else
-        unsigned int ret = (&m_state.result.x)[m_state.substate];
-    #endif
+        unsigned int ret = detail::get_element_at(m_state.result, m_state.substate);
+
         m_state.substate++;
         if(m_state.substate == 4)
         {
diff --git a/projects/rocrand/library/include/rocrand/rocrand_threefry2_impl.h b/projects/rocrand/library/include/rocrand/rocrand_threefry2_impl.h
index 97ee5ffc2bc9..dc8282e5ee98 100644
--- a/projects/rocrand/library/include/rocrand/rocrand_threefry2_impl.h
+++ b/projects/rocrand/library/include/rocrand/rocrand_threefry2_impl.h
@@ -135,13 +135,11 @@ class threefry_engine2_base
         return this->next();
     }
 
-    __forceinline__ __device__ __host__ value next()
+    __forceinline__ __device__ __host__
+    value next()
     {
-#if defined(__HIP_PLATFORM_AMD__)
-        value ret = ROCRAND_HIPVEC_ACCESS(m_state.result)[m_state.substate];
-#else
-        value ret = (&m_state.result.x)[m_state.substate];
-#endif
+        value ret = detail::get_element_at(m_state.result, m_state.substate);
+
         m_state.substate++;
         if(m_state.substate == 2)
         {
diff --git a/projects/rocrand/library/include/rocrand/rocrand_threefry4_impl.h b/projects/rocrand/library/include/rocrand/rocrand_threefry4_impl.h
index 4c93216e98c8..aff969d740ba 100644
--- a/projects/rocrand/library/include/rocrand/rocrand_threefry4_impl.h
+++ b/projects/rocrand/library/include/rocrand/rocrand_threefry4_impl.h
@@ -152,11 +152,8 @@ class threefry_engine4_base
 
     __forceinline__ __device__ __host__ value next()
     {
-#if defined(__HIP_PLATFORM_AMD__)
-        value ret = ROCRAND_HIPVEC_ACCESS(m_state.result)[m_state.substate];
-#else
-        value ret = (&m_state.result.x)[m_state.substate];
-#endif
+        value ret = detail::get_element_at(m_state.result, m_state.substate);
+
         m_state.substate++;
         if(m_state.substate == 4)
         {