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

include/ck/utility/generic_memory_space_atomic.hpp

@@ -28,6 +28,12 @@ __device__ float atomic_add<float>(float* p_dst, const float& x)
     return atomicAdd(p_dst, x);
 }
 
+template <>
+__device__ double atomic_add<double>(double* p_dst, const double& x)
+{
+    return atomicAdd(p_dst, x);
+}
+
 template <>
 __device__ float2_t atomic_add<float2_t>(float2_t* p_dst, const float2_t& x)
 {
@@ -45,6 +51,23 @@ __device__ float2_t atomic_add<float2_t>(float2_t* p_dst, const float2_t& x)
     return vy.template AsType<float2_t>()[I0];
 }
 
+template <>
+__device__ double2_t atomic_add<double2_t>(double2_t* p_dst, const double2_t& x)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+
+    const vector_type<double, 2> vx{x};
+    vector_type<double, 2> vy{0};
+
+    vy.template AsType<double>()(I0) =
+        atomicAdd(c_style_pointer_cast<double*>(p_dst), vx.template AsType<double>()[I0]);
+    vy.template AsType<double>()(I1) =
+        atomicAdd(c_style_pointer_cast<double*>(p_dst) + 1, vx.template AsType<double>()[I1]);
+
+    return vy.template AsType<double2_t>()[I0];
+}
+
 // Caution: DO NOT REMOVE
 // intentionally have only declaration but no definition to cause compilation failure when trying to
 // instantiate this template. The purpose is to make the implementation of atomic_max explicit for

include/ck/utility/inner_product.hpp

-#ifndef CK_INNER_PRODUCT_HPP
-#define CK_INNER_PRODUCT_HPP
-
+#pragma once
 #include "data_type.hpp"
 
 namespace ck {
@@ -138,7 +136,7 @@ template <>
 __device__ void
 inner_product<int8x4_t, int8x4_t, int32_t>(const int8x4_t& a, const int8x4_t& b, int32_t& c)
 {
-#if defined(CK_USE_DOT4_I32_I8)
+#if defined(CK_USE_AMD_V_DOT4_I32_I8)
 #if CK_USE_AMD_INNER_PRODUCT_INLINE_ASM
     asm volatile("\n \
             v_dot4_i32_i8 %0, %1, %2, %0\n \
@@ -202,4 +200,3 @@ inner_product<int8x16_t, int8x16_t, int32_t>(const int8x16_t& a, const int8x16_t
 }
 
 } // namespace ck
-#endif

include/ck/utility/reduction_operator.hpp

@@ -26,7 +26,8 @@
 #ifndef CK_REDUCTION_OPERATOR_HPP
 #define CK_REDUCTION_OPERATOR_HPP
 
-#include "common_header.hpp"
+#include "config.hpp"
+#include "data_type.hpp"
 
 namespace ck {
 
@@ -41,12 +42,10 @@ namespace reduce {
 //                    when operated against them, and the concept is similar to zero vector in
 //                    vector space
 //                    (http://pages.cs.wisc.edu/~matthewb/pages/notes/pdf/linearalgebra/VectorSpaces.pdf).
-// 2) indexable -- boolean value indicating whether indices of the operated elements could be
-// recorded. Usually, Min/Max operator could
-//                 need to record the indices of elements. For operator like Add/Mul, no need to
-//                 record the indices.
-// 3) operator() -- the first argument of the operator must be both an input & output, and the
-// corresponding variable usually stores
+// 2) IsCompatibleInMemoryDataOperation() -- return true if the reduction task corresponding to this
+// operator can use the InMemoryDataOperation to finalize, or else it return false 3) operator() --
+// the first argument of the operator must be both an input & output, and the corresponding variable
+// usually stores
 //                  the accumulated result of many operator() calls; the second argument is only an
 //                  input. For indexable binary
 //                  operator, the second version of operator() has third argument (which is an
@@ -62,6 +61,13 @@ struct Add
 
     __host__ __device__ static constexpr T GetReductionZeroVal() { return static_cast<T>(0.0f); };
 
+    __device__ static constexpr bool
+    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
+    {
+        return operation == InMemoryDataOperationEnum::AtomicAdd ||
+               operation == InMemoryDataOperationEnum::Set;
+    };
+
     __host__ __device__ inline constexpr void operator()(T& a, T b) const { a = a + b; }
 };
 
@@ -72,6 +78,12 @@ struct Mul
 
     __host__ __device__ static constexpr T GetReductionZeroVal() { return static_cast<T>(1.0f); };
 
+    __device__ static constexpr bool
+    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
+    {
+        return operation == InMemoryDataOperationEnum::Set;
+    };
+
     __host__ __device__ inline constexpr void operator()(T& a, T b) const { a = a * b; }
 };
 
@@ -85,6 +97,13 @@ struct Max
         return NumericLimits<T>::Lowest();
     };
 
+    __device__ static constexpr bool
+    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
+    {
+        // ToChange: atomic_max to be added
+        return operation == InMemoryDataOperationEnum::Set;
+    };
+
     __host__ __device__ inline constexpr void operator()(T& a, T b) const
     {
         if(a < b)
@@ -111,6 +130,13 @@ struct Min
         return NumericLimits<T>::Max();
     };
 
+    __device__ static constexpr bool
+    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
+    {
+        // ToChange: atomic_min to be added
+        return operation == InMemoryDataOperationEnum::Set;
+    };
+
     __host__ __device__ inline constexpr void operator()(T& a, T b) const
     {
         if(a > b)
@@ -134,6 +160,13 @@ struct AMax
 
     __host__ __device__ static constexpr T GetReductionZeroVal() { return static_cast<T>(0.0f); };
 
+    __device__ static constexpr bool
+    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
+    {
+        // ToChange: atomic_max to be added
+        return operation == InMemoryDataOperationEnum::Set;
+    };
+
     __host__ __device__ inline constexpr void operator()(T& a, T b) const
     {
         if(a < b)
@@ -150,6 +183,17 @@ struct AMax
     }
 };
 
+template <typename T>
+T GetReductionZeroValueForInMemoryDataOperation(InMemoryDataOperationEnum operation)
+{
+    T result = ck::type_convert<T>(0.0f);
+
+    if(operation == InMemoryDataOperationEnum::AtomicMax)
+        result = ck::NumericLimits<T>::Lowest();
+
+    return (result);
+};
+
 }; // end of namespace reduce
 
 } // end of namespace ck

include/ck/utility/static_buffer.hpp

@@ -36,6 +36,11 @@ struct StaticBuffer : public StaticallyIndexedArray<T, N>
     {
         return base::operator()(i);
     }
+
+    __host__ __device__ void Clear()
+    {
+        static_for<0, N, 1>{}([&](auto i) { operator()(i) = T{0}; });
+    }
 };
 
 // static buffer for vector
@@ -146,9 +151,9 @@ struct StaticBufferTupleOfVector
 
     __host__ __device__ void Clear()
     {
-        const index_t numScalars = NumOfVector * ScalarPerVector;
+        constexpr index_t NumScalars = NumOfVector * ScalarPerVector;
 
-        static_for<0, Number<numScalars>{}, 1>{}([&](auto i) { SetAsType(i, S{0}); });
+        static_for<0, NumScalars, 1>{}([&](auto i) { SetAsType(i, S{0}); });
     }
 };
 

library/include/ck/library/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_b16_f32_b16.hpp

 #ifndef DEVICE_REDUCE_INSTANCE_BLOCKWISE_B16_F32_B16_HPP
 #define DEVICE_REDUCE_INSTANCE_BLOCKWISE_B16_F32_B16_HPP
 
-#include "reduction_enums.hpp"
-#include "reduction_operator_mapping.hpp"
+#include "data_type.hpp"
 #include "device_reduce_instance_blockwise.hpp"
 
 namespace ck {

library/include/ck/library/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f16_f16_f16.hpp

 #ifndef DEVICE_REDUCE_INSTANCE_BLOCKWISE_F16_F16_F16_HPP
 #define DEVICE_REDUCE_INSTANCE_BLOCKWISE_F16_F16_F16_HPP
 
-#include "reduction_enums.hpp"
-#include "reduction_operator_mapping.hpp"
+#include "data_type.hpp"
 #include "device_reduce_instance_blockwise.hpp"
 
 namespace ck {