Merge branch 'develop' into feature/simplify-karg-for-device-gemm-xdl-improved

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_bias_softmax_gemm_permute.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_batched_gemm_softmax_gemm_permute.hpp"

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_gemm.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_batched_gemm_gemm.hpp"

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_softmax_gemm_permute.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_batched_gemm_softmax_gemm_permute.hpp"

library/include/ck/library/tensor_operation_instance/gpu/contraction_bilinear.hpp

@@ -3,10 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
 #include <vector>
 #include <memory>
-
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_contraction_multiple_d.hpp"

library/include/ck/library/tensor_operation_instance/gpu/contraction_scale.hpp

@@ -3,10 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
 #include <vector>
 #include <memory>
-
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_contraction_multiple_d.hpp"

library/include/ck/library/tensor_operation_instance/gpu/convolution_backward_data.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_conv_bwd_data.hpp"

library/include/ck/library/tensor_operation_instance/gpu/convolution_forward.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_conv_fwd.hpp"

library/include/ck/library/tensor_operation_instance/gpu/device_elementwise_instance.hpp

@@ -3,8 +3,7 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/impl/device_elementwise_impl.hpp"

library/include/ck/library/tensor_operation_instance/gpu/device_gemm_mean_squaremean_instance.hpp

@@ -4,7 +4,7 @@
 #pragma once
 
 #include <cstdlib>
-
+#include <vector>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm_reduce.hpp"

library/include/ck/library/tensor_operation_instance/gpu/gemm.hpp

@@ -3,10 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
 #include <memory>
 #include <vector>
-
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm.hpp"

library/include/ck/library/tensor_operation_instance/gpu/gemm_add_add_fastgelu.hpp

@@ -3,10 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
 #include <vector>
 #include <memory>
-
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp"

library/include/ck/library/tensor_operation_instance/gpu/gemm_bilinear.hpp

@@ -3,10 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
 #include <vector>
 #include <memory>
-
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp"

library/include/ck/library/tensor_operation_instance/gpu/gemm_splitk.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm_splitk.hpp"

library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward.hpp

@@ -4,7 +4,7 @@
 #pragma once
 
 #include <vector>
-
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_d.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"

library/include/ck/library/tensor_operation_instance/gpu/grouped_gemm.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_grouped_gemm.hpp"

library/include/ck/library/tensor_operation_instance/gpu/normalization.hpp

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <cstdlib>
-
+#include <vector>
+#include <memory>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_normalization.hpp"

library/include/ck/library/tensor_operation_instance/gpu/pool2d_fwd.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_pool_fwd.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto InOutRank  = 4;
+static constexpr auto WindowRank = 2;
+
+static constexpr auto MaxOp = ck::ReduceTensorOp::MAX;
+static constexpr auto AvgOp = ck::ReduceTensorOp::AVG;
+
+// FP16
+void add_device_pool2d_fwd_nhwc_f16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, MaxOp, false>>>&);
+
+void add_device_pool2d_fwd_nhwc_f16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, AvgOp, false>>>&);
+
+// FP16 - return index
+void add_device_pool2d_fwd_nhwc_index_f16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, MaxOp, true>>>&);
+
+// FP32
+void add_device_pool2d_fwd_nhwc_f32_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, MaxOp, false>>>&);
+
+void add_device_pool2d_fwd_nhwc_f32_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, AvgOp, false>>>&);
+
+// FP32 - return index
+void add_device_pool2d_fwd_nhwc_index_f32_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, MaxOp, true>>>&);
+
+template <typename InDataType,
+          typename OutDataType,
+          typename IndexDataType,
+          ck::ReduceTensorOp ReduceOpId,
+          bool OutputIndex>
+struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DevicePoolFwd<InOutRank,
+                                                                                  WindowRank,
+                                                                                  InDataType,
+                                                                                  OutDataType,
+                                                                                  IndexDataType,
+                                                                                  ReduceOpId,
+                                                                                  OutputIndex>>
+{
+    using DeviceOp = DevicePoolFwd<InOutRank,
+                                   WindowRank,
+                                   InDataType,
+                                   OutDataType,
+                                   IndexDataType,
+                                   ReduceOpId,
+                                   OutputIndex>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+
+        if constexpr(is_same_v<InDataType, F16> && is_same_v<OutDataType, F16> &&
+                     is_same_v<IndexDataType, I32>)
+        {
+            if constexpr(OutputIndex && ReduceOpId == MaxOp)
+            {
+                add_device_pool2d_fwd_nhwc_index_f16_instances(op_ptrs);
+            }
+            else
+            {
+                add_device_pool2d_fwd_nhwc_f16_instances(op_ptrs);
+            }
+        }
+        else if constexpr(is_same_v<InDataType, F32> && is_same_v<OutDataType, F32> &&
+                          is_same_v<IndexDataType, I32>)
+        {
+            if constexpr(OutputIndex && ReduceOpId == MaxOp)
+            {
+                add_device_pool2d_fwd_nhwc_index_f32_instances(op_ptrs);
+            }
+            else
+            {
+                add_device_pool2d_fwd_nhwc_f32_instances(op_ptrs);
+            }
+        }
+
+        return op_ptrs;
+    }
+};
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/tensor_operation_instance/gpu/pool3d_fwd.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_pool_fwd.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto InOutRank  = 5;
+static constexpr auto WindowRank = 3;
+
+static constexpr auto MaxOp = ck::ReduceTensorOp::MAX;
+static constexpr auto AvgOp = ck::ReduceTensorOp::AVG;
+
+// FP16
+void add_device_pool3d_fwd_ndhwc_f16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, MaxOp, false>>>&);
+
+void add_device_pool3d_fwd_ndhwc_f16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, AvgOp, false>>>&);
+
+// FP16 - return index
+void add_device_pool3d_fwd_ndhwc_index_f16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, MaxOp, true>>>&);
+
+// FP32
+void add_device_pool3d_fwd_ndhwc_f32_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, MaxOp, false>>>&);
+
+void add_device_pool3d_fwd_ndhwc_f32_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, AvgOp, false>>>&);
+
+// FP32 - return index
+void add_device_pool3d_fwd_ndhwc_index_f32_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, MaxOp, true>>>&);
+
+template <typename InDataType,
+          typename OutDataType,
+          typename IndexDataType,
+          ck::ReduceTensorOp ReduceOpId,
+          bool OutputIndex>
+struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DevicePoolFwd<InOutRank,
+                                                                                  WindowRank,
+                                                                                  InDataType,
+                                                                                  OutDataType,
+                                                                                  IndexDataType,
+                                                                                  ReduceOpId,
+                                                                                  OutputIndex>>
+{
+    using DeviceOp = DevicePoolFwd<InOutRank,
+                                   WindowRank,
+                                   InDataType,
+                                   OutDataType,
+                                   IndexDataType,
+                                   ReduceOpId,
+                                   OutputIndex>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+
+        if constexpr(is_same_v<InDataType, F16> && is_same_v<OutDataType, F16> &&
+                     is_same_v<IndexDataType, I32>)
+        {
+            if constexpr(OutputIndex && ReduceOpId == MaxOp)
+            {
+                add_device_pool3d_fwd_ndhwc_index_f16_instances(op_ptrs);
+            }
+            else
+            {
+                add_device_pool3d_fwd_ndhwc_f16_instances(op_ptrs);
+            }
+        }
+        else if constexpr(is_same_v<InDataType, F32> && is_same_v<OutDataType, F32> &&
+                          is_same_v<IndexDataType, I32>)
+        {
+            if constexpr(OutputIndex && ReduceOpId == MaxOp)
+            {
+                add_device_pool3d_fwd_ndhwc_index_f32_instances(op_ptrs);
+            }
+            else
+            {
+                add_device_pool3d_fwd_ndhwc_f32_instances(op_ptrs);
+            }
+        }
+
+        return op_ptrs;
+    }
+};
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

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

@@ -90,6 +90,7 @@ void add_device_reduce_instance_threadwise(
                                                             AccElementwiseOp,
                                                             PropagateNan,
                                                             OutputIndex,
+                                                            false,
                                                             false, // HaveIndexInputIfOutputIndex
                                                             cfg1::BlockSize_,
                                                             cfg2::MThreadSliceSize_,

library/include/ck/library/utility/host_conv.hpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-#include "host_tensor.hpp"
-#include "conv_common.hpp"
-
-template <typename TIn,
-          typename TWei,
-          typename TOut,
-          typename ConvStrides,
-          typename ConvDilations,
-          typename InLeftPads,
-          typename InRightPads>
-void host_conv_nchw_kcyx_nkhw(const Tensor<TIn>& in,
-                              const Tensor<TWei>& wei,
-                              Tensor<TOut>& out,
-                              const ConvStrides& conv_strides,
-                              const ConvDilations& conv_dilations,
-                              const InLeftPads& in_left_pads,
-                              const InRightPads&)
-{
-    constexpr auto I0 = ck::Number<0>{};
-    constexpr auto I1 = ck::Number<1>{};
-
-    auto f_nchw = [&](auto n, auto k, auto ho, auto wo) {
-        float v = 0;
-        for(int c = 0; c < wei.mDesc.GetLengths()[1]; ++c)
-        {
-            for(int y = 0; y < wei.mDesc.GetLengths()[2]; ++y)
-            {
-                int hi = ho * conv_strides[I0] + y * conv_dilations[I0] - in_left_pads[I0];
-                for(int x = 0; x < wei.mDesc.GetLengths()[3]; ++x)
-                {
-                    int wi = wo * conv_strides[I1] + x * conv_dilations[I1] - in_left_pads[I1];
-                    if(hi >= 0 && hi < in.mDesc.GetLengths()[2] && wi >= 0 &&
-                       wi < in.mDesc.GetLengths()[3])
-                    {
-                        v += ck::type_convert<float>(in(n, c, hi, wi)) *
-                             ck::type_convert<float>(wei(k, c, y, x));
-                    }
-                }
-            }
-        }
-        out(n, k, ho, wo) = ck::type_convert<TOut>(v);
-    };
-
-    make_ParallelTensorFunctor(f_nchw,
-                               out.mDesc.GetLengths()[0],
-                               out.mDesc.GetLengths()[1],
-                               out.mDesc.GetLengths()[2],
-                               out.mDesc.GetLengths()[3])(std::thread::hardware_concurrency());
-}
-
-template <typename TIn,
-          typename TWei,
-          typename TOut,
-          typename ConvStrides,
-          typename ConvDilations,
-          typename InLeftPads,
-          typename InRightPads>
-void host_conv3d_ndhwc_kzyxc_ndhwk(const Tensor<TIn>& in,
-                                   const Tensor<TWei>& wei,
-                                   Tensor<TOut>& out,
-                                   const ConvStrides& conv_strides,
-                                   const ConvDilations& conv_dilations,
-                                   const InLeftPads& in_left_pads,
-                                   const InRightPads&)
-{
-    using namespace ck;
-
-    constexpr auto I0 = Number<0>{};
-    constexpr auto I1 = Number<1>{};
-    constexpr auto I2 = Number<2>{};
-    const auto Di     = in.mDesc.GetLengths()[1];
-    const auto Hi     = in.mDesc.GetLengths()[2];
-    const auto Wi     = in.mDesc.GetLengths()[3];
-    const auto Z      = wei.mDesc.GetLengths()[1];
-    const auto Y      = wei.mDesc.GetLengths()[2];
-    const auto X      = wei.mDesc.GetLengths()[3];
-    const auto C      = wei.mDesc.GetLengths()[4];
-
-    auto f_ndhwc = [&](auto n, auto do_tmp, auto ho_tmp, auto wo_tmp, auto k) {
-        // do__ must be converted to signed integer, otherwise zmin might be wrong in cases
-        // negative values.
-        const int do_ = static_cast<int>(do_tmp);
-        const int ho  = static_cast<int>(ho_tmp);
-        const int wo  = static_cast<int>(wo_tmp);
-        const int zmin =
-            std::max(0,
-                     (in_left_pads[I0] - do_ * conv_strides[I0] + conv_dilations[I0] - 1) /
-                         conv_dilations[I0]);
-        const int ymin =
-            std::max(0,
-                     (in_left_pads[I1] - ho * conv_strides[I1] + conv_dilations[I1] - 1) /
-                         conv_dilations[I1]);
-        const int xmin =
-            std::max(0,
-                     (in_left_pads[I2] - wo * conv_strides[I2] + conv_dilations[I2] - 1) /
-                         conv_dilations[I2]);
-        const int zmax =
-            std::min(Z, (in_left_pads[I0] - do_ * conv_strides[I0] + Di) / conv_dilations[I0]);
-        const int ymax =
-            std::min(Y, (in_left_pads[I1] - ho * conv_strides[I1] + Hi) / conv_dilations[I1]);
-        const int xmax =
-            std::min(X, (in_left_pads[I2] - wo * conv_strides[I2] + Wi) / conv_dilations[I2]);
-        const int di_min = do_ * conv_strides[I0] + zmin * conv_dilations[I0] - in_left_pads[I0];
-        const int hi_min = ho * conv_strides[I1] + ymin * conv_dilations[I1] - in_left_pads[I1];
-        const int wi_min = wo * conv_strides[I2] + xmin * conv_dilations[I2] - in_left_pads[I2];
-
-        double v = 0;
-
-        const TIn* in_n   = in.mData.data() + n * Di * Hi * Wi * C;
-        const TWei* wei_k = wei.mData.data() + k * Z * Y * X * C;
-
-        int di = di_min;
-        for(int z = zmin; z < zmax; ++z, di += conv_dilations[I0])
-        {
-            const TIn* in_n_di  = in_n + di * Hi * Wi * C;
-            const TWei* wei_k_z = wei_k + z * Y * X * C;
-            int hi              = hi_min;
-
-            for(int y = ymin; y < ymax; ++y, hi += conv_dilations[I1])
-            {
-                const TIn* in_n_di_hi = in_n_di + hi * Wi * C;
-                const TWei* wei_k_z_y = wei_k_z + y * X * C;
-                int wi                = wi_min;
-
-                for(int x = xmin; x < xmax; ++x, wi += conv_dilations[I2])
-                {
-                    const TIn* in_n_di_hi_wi = in_n_di_hi + wi * C;
-                    const TWei* wei_k_z_y_x  = wei_k_z_y + x * C;
-
-                    for(int c = 0; c < C; ++c)
-                    {
-                        v += static_cast<const double>(in_n_di_hi_wi[c]) *
-                             static_cast<const double>(wei_k_z_y_x[c]);
-                    }
-                }
-            }
-        }
-
-        out(n, do_, ho, wo, k) = v;
-    };
-
-    make_ParallelTensorFunctor(f_ndhwc,
-                               out.mDesc.GetLengths()[0],
-                               out.mDesc.GetLengths()[1],
-                               out.mDesc.GetLengths()[2],
-                               out.mDesc.GetLengths()[3],
-                               out.mDesc.GetLengths()[4])(std::thread::hardware_concurrency() - 4);
-}