Merge remote-tracking branch 'origin/develop' into gelu

include/ck/utility/common_header.hpp

@@ -32,7 +32,7 @@
 #include "debug.hpp"
 
 #include "amd_buffer_addressing.hpp"
-#include "generic_memory_space_atomic_add.hpp"
+#include "generic_memory_space_atomic.hpp"
 #include "get_id.hpp"
 #include "synchronization.hpp"
 #include "amd_address_space.hpp"

include/ck/utility/dynamic_buffer.hpp

@@ -3,7 +3,7 @@
 #include "enable_if.hpp"
 #include "c_style_pointer_cast.hpp"
 #include "amd_buffer_addressing.hpp"
-#include "generic_memory_space_atomic_add.hpp"
+#include "generic_memory_space_atomic.hpp"
 
 namespace ck {
 
@@ -125,6 +125,10 @@ struct DynamicBuffer
         {
             this->template AtomicAdd<X>(i, is_valid_element, x);
         }
+        else if constexpr(Op == InMemoryDataOperationEnum::AtomicMax)
+        {
+            this->template AtomicMax<X>(i, is_valid_element, x);
+        }
         else if constexpr(Op == InMemoryDataOperationEnum::Add)
         {
             auto tmp = this->template Get<X>(i, is_valid_element);
@@ -326,6 +330,42 @@ struct DynamicBuffer
         }
     }
 
+    template <typename X,
+              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
+                                         typename scalar_type<remove_cvref_t<T>>::type>::value,
+                                 bool>::type = false>
+    __host__ __device__ void AtomicMax(index_t i, bool is_valid_element, const X& x)
+    {
+        // X contains multiple T
+        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
+
+        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
+
+        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
+                      "wrong! X should contain multiple T");
+
+        static_assert(GetAddressSpace() == AddressSpaceEnum::Global, "only support global mem");
+
+#if CK_USE_AMD_BUFFER_ATOMIC_MAX_FLOAT64
+        using scalar_t                           = typename scalar_type<remove_cvref_t<T>>::type;
+        bool constexpr use_amd_buffer_addressing = is_same_v<remove_cvref_t<scalar_t>, double>;
+#else
+        bool constexpr use_amd_buffer_addressing = false;
+#endif
+
+        if constexpr(use_amd_buffer_addressing)
+        {
+            constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;
+
+            amd_buffer_atomic_max<remove_cvref_t<T>, t_per_x>(
+                x, p_data_, i, is_valid_element, element_space_size_);
+        }
+        else if(is_valid_element)
+        {
+            atomic_max<X>(c_style_pointer_cast<X*>(&p_data_[i]), x);
+        }
+    }
+
     __host__ __device__ static constexpr bool IsStaticBuffer() { return false; }
 
     __host__ __device__ static constexpr bool IsDynamicBuffer() { return true; }

include/ck/utility/get_id.hpp

@@ -11,10 +11,14 @@ __host__ __device__ constexpr index_t get_warp_size()
 
 __device__ index_t get_thread_local_1d_id() { return threadIdx.x; }
 
+__device__ index_t get_thread_global_1d_id() { return blockIdx.x * blockDim.x + threadIdx.x; }
+
 __device__ index_t get_warp_local_1d_id() { return threadIdx.x / get_warp_size(); }
 
 __device__ index_t get_block_1d_id() { return blockIdx.x; }
 
 __device__ index_t get_grid_size() { return gridDim.x; }
 
+__device__ index_t get_block_size() { return blockDim.x; }
+
 } // namespace ck

include/ck/utility/type.hpp

@@ -29,6 +29,9 @@ using remove_cv_t = typename std::remove_cv<T>::type;
 template <typename T>
 using remove_cvref_t = remove_cv_t<std::remove_reference_t<T>>;
 
+template <typename T>
+using remove_pointer_t = typename std::remove_pointer<T>::type;
+
 template <typename T>
 inline constexpr bool is_pointer_v = std::is_pointer<T>::value;