Merge remote-tracking branch 'origin/develop' into fused-gemm

example/01_gemm/CMakeLists.txt

@@ -4,5 +4,6 @@ add_example_executable(example_gemm_dl_int8 gemm_dl_int8.cpp)
 add_example_executable(example_gemm_xdl_fp16 gemm_xdl_fp16.cpp)
 add_example_executable(example_gemm_xdl_bf16 gemm_xdl_bf16.cpp)
 add_example_executable(example_gemm_xdl_int8 gemm_xdl_int8.cpp)
+add_example_executable(example_gemm_xdl_skip_b_lds_fp16 gemm_xdl_skip_b_lds_fp16.cpp)
 # FIXME: re-enable this exampe as test when SWDEV-335738 is fixed
 add_example_executable_no_testing(example_gemm_xdl_fp64 gemm_xdl_fp64.cpp)

example/01_gemm/gemm_xdl_skip_b_lds_fp16.cpp

+#include <iostream>
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_xdl.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_xdl_skip_b_lds.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/host_tensor/device_memory.hpp"
+#include "ck/library/host_tensor/host_tensor.hpp"
+#include "ck/library/host_tensor/host_tensor_generator.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using F16 = ck::half_t;
+using F32 = float;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+using ALayout = ck::tensor_layout::gemm::RowMajor;
+using BLayout = ck::tensor_layout::gemm::ColumnMajor;
+using CLayout = ck::tensor_layout::gemm::RowMajor;
+
+using AElementOp = ck::tensor_operation::element_wise::PassThrough;
+using BElementOp = ck::tensor_operation::element_wise::PassThrough;
+using CElementOp = ck::tensor_operation::element_wise::PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+#define USING_SKIP_LDS 1
+
+// clang-format off
+#if USING_SKIP_LDS
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmXdlSkipBLds
+        //###########| AData| BData| CData| AccData| ALayout| BLayout| CLayout|           A|           B|           C|          GEMM| Block|  MPer|  NPer| K0Per| K1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BThreadTransfer|  BBlock| CThreadTransfer| CThreadTransfer|
+        //###########|  Type|  Type|  Type|    Type|        |        |        | Elementwise| Elementwise| Elementwise|Spacialization|  Size| Block| Block| Block|   |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|       SrcScalar|  buffer| SrcDstVectorDim|       DstScalar|
+        //###########|      |      |      |        |        |        |        |   Operation|   Operation|   Operation|              |      |      |      |      |   |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          |       PerVector|  size  |                |       PerVector|
+        //###########|      |      |      |        |        |        |        |            |            |            |              |      |      |      |      |   |     |     |     |     |                |               |               |               |               |               |          |                |        |                |                |
+#if 0       
+                    <   F16,   F16,   F16,     F32,     Row,     Col,     Row, PassThrough, PassThrough, PassThrough,   GemmDefault,   256,    16,   64,     4,  8,   16,   16,    1,    1,     S<16, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,            8,       8,        7,               1>;
+using ADataType   = ck::half_t;
+using BDataType   = ck::half_t;
+using CDataType   = ck::half_t;
+using AccDataType = float;
+#else  
+                    <   F32,   F32,   F32,     F32,     Row,     Col,     Row, PassThrough, PassThrough, PassThrough,   GemmDefault,   256,    16,   64,     4,  4,   16,   16,    1,    1,     S<16, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              4,              4,      true,             4,      4,        7,               1>;
+using ADataType   = float;
+using BDataType   = float;
+using CDataType   = float;
+using AccDataType = float;
+#endif
+
+#else
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmXdl
+        //###########| AData| BData| CData| AccData| ALayout| BLayout| CLayout|           A|           B|           C|          GEMM| Block|  MPer|  NPer| K0Per| K1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CThreadTransfer| CThreadTransfer|
+        //###########|  Type|  Type|  Type|    Type|        |        |        | Elementwise| Elementwise| Elementwise|Spacialization|  Size| Block| Block| Block|   |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| SrcDstVectorDim|       DstScalar|
+        //###########|      |      |      |        |        |        |        |   Operation|   Operation|   Operation|              |      |      |      |      |   |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |                |       PerVector|
+        //###########|      |      |      |        |        |        |        |            |            |            |              |      |      |      |      |   |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |                |                |
+                   <   F32,   F32,   F32,     F32,     Row,     Col,     Row, PassThrough, PassThrough, PassThrough,   GemmDefault,   256,    16,   64,     4,  4,   16,   16,    1,    1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              4,              4,      true,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              4,              4,      true,               7,              1,      2>;
+using ADataType   = float;
+using BDataType   = float;
+using CDataType   = float;
+using AccDataType = float;
+
+#endif
+    // clang-format on
+
+    using ReferenceGemmInstance = ck::tensor_operation::host::
+        ReferenceGemm<ADataType, BDataType, CDataType, float, AElementOp, BElementOp, CElementOp>;
+
+template <typename DataType>
+std::ostream& show_2d_matrix(std::ostream& os, Tensor<DataType>& matrix)
+{
+    os << "[" << std::endl;
+    for(size_t x = 0; x < matrix.mDesc.GetLengths()[0]; x++)
+    {
+        os << "[";
+        for(size_t y = 0; y < matrix.mDesc.GetLengths()[1]; y++)
+        {
+            os << std::setw(5) << static_cast<float>(matrix(x, y));
+        }
+        os << "]" << std::endl;
+    }
+    os << "]";
+    return os;
+}
+int main(int argc, char* argv[])
+{
+    bool do_verification = 0;
+    int init_method      = 0;
+    bool time_kernel     = false;
+
+    // GEMM shape
+#if 1
+    ck::index_t M = 16;
+    ck::index_t N = 64 * 120;
+    ck::index_t K = 4096;
+
+    ck::index_t StrideA = K;
+    ck::index_t StrideB = K;
+    ck::index_t StrideC = N;
+#else
+    ck::index_t M = 16;
+    ck::index_t N = 16;
+    ck::index_t K = 32;
+
+    ck::index_t StrideA = 8;
+    ck::index_t StrideB = 8;
+    ck::index_t StrideC = 16;
+#endif
+
+    if(argc == 4)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+    }
+    else if(argc == 10)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+
+        M = std::stoi(argv[4]);
+        N = std::stoi(argv[5]);
+        K = std::stoi(argv[6]);
+
+        StrideA = std::stoi(argv[7]);
+        StrideB = std::stoi(argv[8]);
+        StrideC = std::stoi(argv[9]);
+    }
+    else
+    {
+        printf("arg1: verification (0=no, 1=yes)\n");
+        printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
+        printf("arg3: time kernel (0=n0, 1=yes)\n");
+        printf("arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC\n");
+        exit(0);
+    }
+
+    auto f_host_tensor_descriptor =
+        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+            if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+            {
+                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                            std::vector<std::size_t>({stride, 1}));
+            }
+            else
+            {
+                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                            std::vector<std::size_t>({1, stride}));
+            }
+        };
+
+    Tensor<ADataType> a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
+    Tensor<BDataType> b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
+    Tensor<CDataType> c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
+    Tensor<CDataType> c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
+
+    std::cout << "a_m_k: " << a_m_k.mDesc << std::endl;
+    std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
+    std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl;
+
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
+        b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
+        break;
+    case 2:
+        a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
+        b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
+        break;
+    default:
+        // a_m_k.GenerateTensorValue(GeneratorTensor_1<ADataType>{1});
+        a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
+        b_k_n.GenerateTensorValue(GeneratorTensor_1<ADataType>{1});
+    }
+
+    DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpace());
+    DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpace());
+    DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpace());
+
+    a_m_k_device_buf.ToDevice(a_m_k.mData.data());
+    b_k_n_device_buf.ToDevice(b_k_n.mData.data());
+
+    auto a_element_op = AElementOp{};
+    auto b_element_op = BElementOp{};
+    auto c_element_op = CElementOp{};
+
+    // do GEMM
+    auto gemm     = DeviceGemmInstance{};
+    auto invoker  = gemm.MakeInvoker();
+    auto argument = gemm.MakeArgument(static_cast<ADataType*>(a_m_k_device_buf.GetDeviceBuffer()),
+                                      static_cast<BDataType*>(b_k_n_device_buf.GetDeviceBuffer()),
+                                      static_cast<CDataType*>(c_m_n_device_buf.GetDeviceBuffer()),
+                                      M,
+                                      N,
+                                      K,
+                                      StrideA,
+                                      StrideB,
+                                      StrideC,
+                                      a_element_op,
+                                      b_element_op,
+                                      c_element_op);
+
+    if(!gemm.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error(
+            "wrong! device_gemm with the specified compilation parameters does "
+            "not support this GEMM problem");
+    }
+
+    float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+
+    std::size_t flop = std::size_t(2) * M * N * K;
+    std::size_t num_btype =
+        sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
+
+    float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
+
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
+              << gemm.GetTypeString() << std::endl;
+
+    c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data());
+
+    if(do_verification)
+    {
+        auto ref_gemm    = ReferenceGemmInstance{};
+        auto ref_invoker = ref_gemm.MakeInvoker();
+
+        auto ref_argument = ref_gemm.MakeArgument(
+            a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op);
+
+        ref_invoker.Run(ref_argument);
+
+#if 0
+        {
+            show_2d_matrix(std::cout << "a : ", a_m_k) << std::endl;
+            show_2d_matrix(std::cout << "b: ", b_k_n) << std::endl;
+            show_2d_matrix(std::cout << "c_device: ", c_m_n_device_result) << std::endl;
+            show_2d_matrix(std::cout << "c_host  :", c_m_n_host_result) << std::endl;
+        }
+#endif
+        ck::utils::check_err(c_m_n_device_result.mData, c_m_n_host_result.mData);
+    }
+
+    return 0;
+}

example/12_reduce/CMakeLists.txt

 add_example_executable(example_reduce_blockwise reduce_blockwise.cpp)
+add_example_executable(example_reduce_multiblock_atomic_add reduce_multiblock_atomic_add.cpp)
 add_example_executable(example_reduce_blockwise_two_call reduce_blockwise_two_call.cpp)

example/12_reduce/README.md

@@ -2,20 +2,41 @@
 
 ## Run ```example_reduce_blockwise```
 ```bash
-# -D <xxx> : input 4-d tensor lengths
+# -D <xxx> : input 3d/4d/5d tensor lengths
+# -R <xxx> : reduce dimension ids
 # -v <x> :   verification (0=no, 1=yes)
-#arg1: initialization (0=no init, 1=single integer value, 2=scope integer value, 3=decimal value)
-#arg2: time kernel (0=no, 1=yes)
-./bin/example_reduce_blockwise -D 16,64,32,960 -v 1 1 1
+#arg1: data type (0: fp16, 1: fp32, 3: int8, 5: bp16, 6: fp64, 7: int4)
+#arg2: initialization (0=no init, 1=single integer value, 2=scope integer value, 3=decimal value)
+#arg3: time kernel (0=no, 1=yes)
+./bin/example_reduce_blockwise -D 16,64,32,960 -v 1 0 2 1
 ```
 
 Result
 ```
-./bin/example_reduce_blockwise -D 16,64,32,960 -v 1 1 1
-launch_and_time_kernel: grid_dim {240, 1, 1}, block_dim {256, 1, 1}
+./bin/example_reduce_blockwise -D 16,64,32,960 -v 1 0 2 1
+launch_and_time_kernel: grid_dim {240, 1, 1}, block_dim {256, 1, 1} 
 Warm up 1 time
 Start running 10 times...
-Perf: 0.282592 ms, 222.641 GB/s, DeviceReduceBlockWise<256,M_C4_S1,K_C64_S1,InSrcVectorDim_0_InSrcVectorSize_1_OutDstVectorSize_1>
+Perf: 0.238063 ms, 264.285 GB/s, DeviceReduceBlockWise<256,M_C4_S1,K_C64_S1,InSrcVectorDim_0_InSrcVectorSize_1_OutDstVectorSize_1>
+```
+
+## Run ```example_reduce_multiblock_atomic_add```
+```bash
+# -D <xxx> : input 3d/4d/5d tensor lengths
+# -R <xxx> : reduce dimension ids
+# -v <x> :   verification (0=no, 1=yes)
+#arg1: data type (0: fp32, 1: fp64)
+#arg2: initialization (0=no init, 1=single integer value, 2=scope integer value, 3=decimal value)
+#arg3: time kernel (0=no, 1=yes)
+./bin/example_reduce_multiblock_atomic_add -D 16,64,32,960 -v 1 0 2 0
+```
+
+Result
+```
+./bin/example_reduce_multiblock_atomic_add -D 16,64,32,960 -v 1 0 2 0
+Perf: 0 ms, inf GB/s, DeviceReduceMultiBlock<256,M_C4_S1,K_C64_S1,InSrcVectorDim_0_InSrcVectorSize_1_OutDstVectorSize_1>
+echo $?
+0
 ```
 
 # Instructions for ```example_reduce_blockwise_two_call```

example/12_reduce/reduce_blockwise.cpp

@@ -2,64 +2,17 @@
 // Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <iostream>
-#include <numeric>
 #include <initializer_list>
 #include <cstdlib>
 #include <getopt.h>
 
-#include "ck/ck.hpp"
 #include "ck/utility/reduction_enums.hpp"
-#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
-#include "ck/tensor_operation/gpu/device/device_reduce_multiblock.hpp"
-
-#include "ck/library/utility/check_err.hpp"
-#include "ck/library/utility/device_memory.hpp"
-#include "ck/library/utility/host_tensor.hpp"
-#include "ck/library/utility/host_tensor_generator.hpp"
-#include "ck/library/utility/host_common_util.hpp"
-#include "ck/library/utility/host_reduction.hpp"
+#include "reduce_blockwise_impl.hpp"
+#include "reduce_example_common.hpp"
 
 using namespace ck;
 using namespace ck::tensor_operation::device;
 
-using InDataType  = ck::half_t;
-using OutDataType = ck::half_t;
-using AccDataType = float;
-
-constexpr int Rank         = 4;
-constexpr int NumReduceDim = 3;
-
-constexpr ReduceTensorOp ReduceOpId = ReduceTensorOp::NORM2;
-constexpr bool PropagateNan         = true;
-constexpr bool OutputIndex          = false;
-
-using ReduceOperation = typename reduce_binary_operator<ReduceOpId>::opType;
-using InElementwiseOperation =
-    typename reduce_unary_operator<ReduceOpId, true, true>::InElementwiseOperation;
-using AccElementwiseOperation =
-    typename reduce_unary_operator<ReduceOpId, true, true>::AccElementwiseOperation;
-
-using DeviceReduceInstance = DeviceReduceMultiBlock<InDataType,
-                                                    AccDataType,
-                                                    OutDataType,
-                                                    Rank,
-                                                    NumReduceDim,
-                                                    ReduceOperation,
-                                                    InElementwiseOperation,
-                                                    AccElementwiseOperation,
-                                                    InMemoryDataOperationEnum::Set,
-                                                    PropagateNan,
-                                                    OutputIndex,
-                                                    false, // HaveIndexInputIfOutputIndex
-                                                    256,
-                                                    4,
-                                                    64,
-                                                    1,
-                                                    1,
-                                                    0,
-                                                    1,
-                                                    1>;
-
 static struct option long_options[] = {{"inLengths", required_argument, nullptr, 'D'},
                                        {"verify", required_argument, nullptr, 'v'},
                                        {"help", no_argument, nullptr, '?'},
@@ -72,10 +25,12 @@ class SimpleAppArgs
 
     public:
     std::vector<size_t> inLengths = {16, 64, 32, 960};
+    std::vector<int> reduceDims   = {0, 1, 2};
     std::vector<float> scales     = {1.0f, 0.0f};
 
     bool do_verification = true;
-    int init_method      = 1;
+    int data_type        = 1;
+    int init_method      = 2;
     bool time_kernel     = true;
 
     public:
@@ -84,13 +39,17 @@ class SimpleAppArgs
         std::cout << "Usage of " << cmd << std::endl;
         std::cout << "--inLengths or -D, comma separated list of input tensor dimension lengths"
                   << std::endl;
+        std::cout << "--reduceDims or -R, comma separated list of to-reduce dimensions"
+                  << std::endl;
         std::cout << "--verify or -v, 1/0 to indicate whether to verify the reduction result by "
                      "comparing with the host-based reduction"
                   << std::endl;
-        std::cout << "Arg1 -- init method (0=no init, 1=single integer value, 2=scope integer "
+        std::cout << "Arg1: data type (0: fp16, 1: fp32, 3: int8, 5: bp16, 6: fp64, 7: int4)"
+                  << std::endl;
+        std::cout << "Arg2 -- init method (0=no init, 1=single integer value, 2=scope integer "
                      "value, 3=decimal value)"
                   << std::endl;
-        std::cout << "Arg2 -- time kernel (0=no, 1=yes)" << std::endl;
+        std::cout << "Arg3 -- time kernel (0=no, 1=yes)" << std::endl;
     };
 
     int processArgs(int argc, char* argv[])
@@ -101,7 +60,7 @@ class SimpleAppArgs
 
         while(1)
         {
-            ch = getopt_long(argc, argv, "D:v:l:", long_options, &option_index);
+            ch = getopt_long(argc, argv, "D:R:v:l:", long_options, &option_index);
             if(ch == -1)
                 break;
             switch(ch)
@@ -112,6 +71,12 @@ class SimpleAppArgs
 
                 inLengths = getTypeValuesFromString<size_t>(optarg);
                 break;
+            case 'R':
+                if(!optarg)
+                    throw std::runtime_error("Invalid option format!");
+
+                reduceDims = getTypeValuesFromString<int>(optarg);
+                break;
             case 'v':
                 if(!optarg)
                     throw std::runtime_error("Invalid option format!");
@@ -129,9 +94,12 @@ class SimpleAppArgs
             };
         };
 
-        if(optind + 2 > argc)
+        if(optind + 3 > argc)
+        {
             throw std::runtime_error("Invalid cmd-line arguments, more argumetns are needed!");
+        };
 
+        data_type   = std::atoi(argv[optind++]);
         init_method = std::atoi(argv[optind++]);
         time_kernel = static_cast<bool>(std::atoi(argv[optind]));
 
@@ -145,198 +113,152 @@ class SimpleAppArgs
     };
 };
 
-int main(int argc, char* argv[])
+template <typename InOutDataType,
+          typename AccDataType,
+          ReduceTensorOp ReduceOpId,
+          index_t PropagateNan,
+          index_t OutputIndex>
+bool reduce_blockwise_test(bool do_verification,
+                           int init_method,
+                           bool time_kernel,
+                           const std::vector<size_t>& inLengths,
+                           const std::vector<int>& reduceDims,
+                           float alpha,
+                           float beta)
 {
-    const std::vector<int> reduceDims{0, 1, 2};
-    const std::vector<int> invariantDims{3};
+    bool matched = false;
+    int result   = 0;
 
-    SimpleAppArgs args;
+    const auto tuple_object = reduce_shape_instances{};
 
-    if(argc > 1)
-    {
-        if(args.processArgs(argc, argv) < 0)
-            return (-1);
-    };
-
-    constexpr bool op_support_indices =
-        (ReduceOpId == ReduceTensorOp::MIN || ReduceOpId == ReduceTensorOp::MAX ||
-         ReduceOpId == ReduceTensorOp::AMAX);
-
-    // if input is half type, no reason to use float for indiced reduction operation and must use
-    // float for non-indiced reduction operation for accuracy
-    constexpr bool invalid_reduce_1 =
-        std::is_same<InDataType, ck::half_t>::value &&
-        ((!op_support_indices && !std::is_same<AccDataType, float>::value) ||
-         (op_support_indices && !std::is_same<AccDataType, ck::half_t>::value));
-
-    // if input is float type, no reason to use double for indiced reduction operation
-    constexpr bool invalid_reduce_2 =
-        std::is_same<InDataType, float>::value &&
-        (op_support_indices && !std::is_same<AccDataType, float>::value);
-
-    // indices option can only be used when it is really needed
-    constexpr bool invalid_reduce_3 = (!op_support_indices && OutputIndex);
+    static_for<0, std::tuple_size<reduce_shape_instances>::value, 1>{}([&](auto i) {
+        if(matched)
+            return;
 
-    constexpr bool invalid_reduce = (invalid_reduce_1 || invalid_reduce_2 || invalid_reduce_3);
+        using ShapeType = remove_cvref_t<decltype(std::get<i>(tuple_object))>;
 
-    if constexpr(invalid_reduce)
-        std::cout << "Reduction setting is not supported, exiting!" << std::endl;
+        if(ShapeType::Rank_ != inLengths.size() || ShapeType::NumReduceDim_ != reduceDims.size())
+            return;
 
-    Tensor<InDataType> in(args.inLengths);
+        result = reduce_blockwise_impl<InOutDataType,
+                                       AccDataType,
+                                       ReduceOpId,
+                                       ShapeType::Rank_,
+                                       ShapeType::NumReduceDim_,
+                                       PropagateNan,
+                                       OutputIndex>(
+            do_verification, init_method, time_kernel, inLengths, reduceDims, alpha, beta);
 
-    std::vector<size_t> outLengths;
+        matched = true;
+    });
 
-    if(invariantDims.empty())
-        outLengths.push_back(1);
-    else
-        for(auto dim : invariantDims)
-            outLengths.push_back(args.inLengths[dim]);
-
-    Tensor<OutDataType> out_ref(outLengths);
-    Tensor<OutDataType> out(outLengths);
-    Tensor<int> out_indices_ref(outLengths);
-    Tensor<int> out_indices(outLengths);
+    return (result == 0) ? true : false;
+};
 
-    auto inStrides  = in.mDesc.GetStrides();
-    auto outStrides = out.mDesc.GetStrides();
+constexpr ReduceTensorOp ReduceOpId = ReduceTensorOp::AVG;
+constexpr bool PropagateNan         = true;
+constexpr bool OutputIndex          = false;
 
-    size_t invariant_total_length = out.mDesc.GetElementSize();
-    size_t reduce_total_length    = in.mDesc.GetElementSize() / invariant_total_length;
+int main(int argc, char* argv[])
+{
+    bool pass = true;
 
-    float alpha = args.scales[0];
-    float beta  = args.scales[1];
+    if(argc > 1)
+    {
+        SimpleAppArgs arg;
 
-    std::size_t num_thread = 1;
+        if(arg.processArgs(argc, argv) < 0)
+            return (-1);
 
-    if(args.do_verification)
-    {
-        switch(args.init_method)
+        if(arg.data_type == 0)
         {
-        case 0: break;
-        case 1:
-            in.GenerateTensorValue(GeneratorTensor_1<InDataType>{1}, num_thread);
-            if(beta != 0.0f)
-                out_ref.GenerateTensorValue(GeneratorTensor_1<InDataType>{1}, num_thread);
-            break;
-        case 2:
-            in.GenerateTensorValue(GeneratorTensor_2<InDataType>{-5, 5}, num_thread);
-            if(beta != 0.0f)
-                out_ref.GenerateTensorValue(GeneratorTensor_2<InDataType>{-5, 5}, num_thread);
-            break;
-        default:
-            in.GenerateTensorValue(GeneratorTensor_3<InDataType>{-5.0, 5.0}, num_thread);
-            if(beta != 0.0f)
-                out_ref.GenerateTensorValue(GeneratorTensor_3<InDataType>{-5.0, 5.0}, num_thread);
+            pass = reduce_blockwise_test<ck::half_t, float, ReduceOpId, PropagateNan, OutputIndex>(
+                arg.do_verification,
+                arg.init_method,
+                arg.time_kernel,
+                arg.inLengths,
+                arg.reduceDims,
+                arg.scales[0],
+                arg.scales[1]);
         }
-
-        if(beta != 0.0f)
-            for(size_t i = 0; i < out_ref.mDesc.GetElementSpaceSize(); i++)
-                out.mData[i] = out_ref.mData[i];
-    };
-
-    // these buffers are usually provided by the user application
-    DeviceMem in_dev(sizeof(InDataType) * in.mDesc.GetElementSpaceSize());
-    DeviceMem out_dev(sizeof(OutDataType) * out.mDesc.GetElementSpaceSize());
-
-    in_dev.ToDevice(in.mData.data());
-
-    if(beta != 0.0f)
-        out_dev.ToDevice(out.mData.data());
-
-    size_t indicesSizeInBytes = OutputIndex ? out.mDesc.GetElementSize() * sizeof(int32_t) : 0;
-
-    DeviceMem out_index_dev(indicesSizeInBytes);
-
-    InElementwiseOperation in_elementwise_op;
-    AccElementwiseOperation acc_elementwise_op;
-
-    std::tie(in_elementwise_op, acc_elementwise_op) =
-        reduce_unary_operator<ReduceOpId, true, true>::GetElementwiseOperator(
-            static_cast<int32_t>(reduce_total_length));
-
-    if(args.do_verification)
-    {
-        ReductionHost<InDataType,
-                      AccDataType,
-                      OutDataType,
-                      ReduceOperation,
-                      InElementwiseOperation,
-                      AccElementwiseOperation,
-                      Rank,
-                      NumReduceDim,
-                      PropagateNan,
-                      OutputIndex>
-            hostReduce(in.mDesc, out_ref.mDesc, invariantDims, reduceDims);
-
-        hostReduce.Run(alpha,
-                       in.mData.data(),
-                       beta,
-                       out_ref.mData.data(),
-                       out_indices_ref.mData.data(),
-                       in_elementwise_op,
-                       acc_elementwise_op);
-    };
-
-    std::vector<ck::index_t> i_inLengths;
-    std::vector<ck::index_t> i_inStrides;
-    std::vector<ck::index_t> i_outLengths;
-    std::vector<ck::index_t> i_outStrides;
-
-    i_inLengths.assign(args.inLengths.begin(), args.inLengths.end());
-    i_inStrides.assign(inStrides.begin(), inStrides.end());
-    i_outLengths.assign(outLengths.begin(), outLengths.end());
-    i_outStrides.assign(outStrides.begin(), outStrides.end());
-
-    auto reduce = DeviceReduceInstance{};
-
-    auto argument_ptr = reduce.MakeArgumentPointer(i_inLengths,
-                                                   i_inStrides,
-                                                   i_outLengths,
-                                                   i_outStrides,
-                                                   reduceDims,
-                                                   alpha,
-                                                   beta,
-                                                   in_dev.GetDeviceBuffer(),
-                                                   nullptr,
-                                                   out_dev.GetDeviceBuffer(),
-                                                   out_index_dev.GetDeviceBuffer(),
-                                                   in_elementwise_op,
-                                                   acc_elementwise_op);
-
-    if(!reduce.IsSupportedArgument(argument_ptr.get()))
-    {
-        std::cout
-            << "The runtime parameters seems not supported by the DeviceReduce instance, exiting!"
-            << std::endl;
-    };
-
-    std::string reduce_name = reduce.GetTypeString();
-
-    auto invoker_ptr = reduce.MakeInvokerPointer();
-
-    float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, args.time_kernel});
-
-    std::size_t num_bytes = invariant_total_length * reduce_total_length * sizeof(InDataType) +
-                            invariant_total_length * sizeof(OutDataType);
-
-    float gb_per_sec = num_bytes / 1.E6 / avg_time;
-
-    std::cout << "Perf: " << avg_time << " ms, " << gb_per_sec << " GB/s, " << reduce_name
-              << std::endl;
-
-    bool pass = true;
-
-    if(args.do_verification)
-    {
-        out_dev.FromDevice(out.mData.data());
-        pass = pass && ck::utils::check_err(out.mData, out_ref.mData);
-
-        if(OutputIndex)
+        else if(arg.data_type == 1)
         {
-            out_index_dev.FromDevice(out_indices.mData.data());
-            pass = pass && ck::utils::check_err(out_indices.mData, out_indices_ref.mData);
-        };
+            pass = reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
+                arg.do_verification,
+                arg.init_method,
+                arg.time_kernel,
+                arg.inLengths,
+                arg.reduceDims,
+                arg.scales[0],
+                arg.scales[1]);
+        }
+        else if(arg.data_type == 3)
+        {
+            pass = reduce_blockwise_test<int8_t, float, ReduceOpId, PropagateNan, OutputIndex>(
+                arg.do_verification,
+                arg.init_method,
+                arg.time_kernel,
+                arg.inLengths,
+                arg.reduceDims,
+                arg.scales[0],
+                arg.scales[1]);
+        }
+        else if(arg.data_type == 5)
+        {
+            pass = reduce_blockwise_test<ck::bhalf_t, float, ReduceOpId, PropagateNan, OutputIndex>(
+                arg.do_verification,
+                arg.init_method,
+                arg.time_kernel,
+                arg.inLengths,
+                arg.reduceDims,
+                arg.scales[0],
+                arg.scales[1]);
+        }
+        else if(arg.data_type == 6)
+        {
+            pass = reduce_blockwise_test<double, double, ReduceOpId, PropagateNan, OutputIndex>(
+                arg.do_verification,
+                arg.init_method,
+                arg.time_kernel,
+                arg.inLengths,
+                arg.reduceDims,
+                arg.scales[0],
+                arg.scales[1]);
+        }
+    }
+    else
+    {
+        // for testing half_t
+        pass =
+            pass && reduce_blockwise_test<ck::half_t, float, ReduceOpId, PropagateNan, OutputIndex>(
+                        true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
+
+        // for testing float
+        pass = pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
+                           true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
+
+        // for testing double
+        pass = pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
+                           true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
+
+        // for testing bhalf_t
+        pass = pass &&
+               reduce_blockwise_test<ck::bhalf_t, float, ReduceOpId, PropagateNan, OutputIndex>(
+                   true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
+
+        // for testing int8_t
+        pass =
+            pass && reduce_blockwise_test<int8_t, int32_t, ReduceOpId, PropagateNan, OutputIndex>(
+                        true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
+
+        // for testing 3D input
+        pass = pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
+                           true, 2, true, {16, 64, 960}, {0, 1}, 1.0f, 0.0f);
+
+        // for testing 5D input
+        pass = pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
+                           true, 2, true, {16, 64, 32, 2, 960}, {0, 1, 2, 3}, 1.0f, 0.0f);
     };
 
     return (pass ? 0 : 1);
-}
+};

example/12_reduce/reduce_blockwise_impl.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+
+#include "ck/ck.hpp"
+#include "ck/utility/reduction_enums.hpp"
+#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
+#include "ck/tensor_operation/gpu/device/device_reduce_multiblock.hpp"
+
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/utility/host_common_util.hpp"
+#include "ck/library/utility/host_reduction.hpp"
+
+#include "reduce_example_common.hpp"
+
+template <typename InOutDataType,
+          typename AccDataType,
+          ck::ReduceTensorOp ReduceOpId,
+          ck::index_t Rank,
+          ck::index_t NumReduceDim,
+          bool PropagateNan,
+          bool OutputIndex>
+int reduce_blockwise_impl(bool do_verification,
+                          int init_method,
+                          bool time_kernel,
+                          const std::vector<size_t>& inLengths,
+                          const std::vector<int>& reduceDims,
+                          float alpha,
+                          float beta)
+
+{
+    using namespace ck;
+    using namespace ck::tensor_operation::device;
+
+    constexpr bool op_support_indices =
+        (ReduceOpId == ReduceTensorOp::MIN || ReduceOpId == ReduceTensorOp::MAX ||
+         ReduceOpId == ReduceTensorOp::AMAX);
+
+    constexpr bool invalid_reduce_1 = OutputIndex && !op_support_indices;
+
+    // 1) If InOutDataType is half_t, must use half_t as AccDataType for indexable reduction
+    // operations 2) If InOutDataType is half_t, must use float as AccDataType for non-indexable
+    // reduction operations
+    constexpr bool invalid_reduce_2 =
+        std::is_same<InOutDataType, half_t>::value &&
+        ((!op_support_indices && !std::is_same<AccDataType, float>::value) ||
+         (op_support_indices && !std::is_same<AccDataType, half_t>::value));
+
+    // 1) If InOutDataType is float, must use float as AccDataType for indexable reduction
+    // operations
+    constexpr bool invalid_reduce_3 =
+        std::is_same<InOutDataType, float>::value &&
+        (op_support_indices && !std::is_same<AccDataType, float>::value);
+
+    // 1) If InOutDataType is int8_t, must use int8_t as AccDataType for indexable reduction
+    // operations 2) If InOutDataType is int8_t, must use int32_t as AccDataType for non-indexable
+    // reduction operations
+    constexpr bool invalid_reduce_4 =
+        std::is_same<InOutDataType, int8_t>::value &&
+        ((!op_support_indices && !std::is_same<AccDataType, int32_t>::value) ||
+         (op_support_indices && !std::is_same<AccDataType, int8_t>::value));
+
+    // 1) If InOutDataType is int8_t, the supported operation must be either indexable operations or
+    // ADD/AVG
+    constexpr bool invalid_reduce_5 = std::is_same<InOutDataType, int8_t>::value &&
+                                      (!op_support_indices && ReduceOpId != ReduceTensorOp::ADD &&
+                                       ReduceOpId != ReduceTensorOp::AVG);
+
+    // 1) If InOutDataType is bhalf_t, must use float as AccDataType for all reduction operations
+    constexpr bool invalid_reduce_6 =
+        std::is_same<InOutDataType, bhalf_t>::value && !std::is_same<AccDataType, float>::value;
+
+    constexpr bool invalid_reduce = (invalid_reduce_1 || invalid_reduce_2 || invalid_reduce_3 ||
+                                     invalid_reduce_4 || invalid_reduce_5 || invalid_reduce_6);
+
+    if(invalid_reduce)
+    {
+        std::cerr << "The reduction setting is invalid, exiting!" << std::endl;
+        return (-1);
+    };
+
+    using ReduceOperation = typename reduce_binary_operator<ReduceOpId>::opType;
+    using InElementwiseOperation =
+        typename reduce_unary_operator<ReduceOpId, true, true>::InElementwiseOperation;
+    using AccElementwiseOperation =
+        typename reduce_unary_operator<ReduceOpId, true, true>::AccElementwiseOperation;
+
+    using DeviceReduceInstance =
+        ck::tensor_operation::device::DeviceReduceMultiBlock<InOutDataType,
+                                                             AccDataType,
+                                                             InOutDataType,
+                                                             Rank,
+                                                             NumReduceDim,
+                                                             ReduceOperation,
+                                                             InElementwiseOperation,
+                                                             AccElementwiseOperation,
+                                                             InMemoryDataOperationEnum::Set,
+                                                             PropagateNan,
+                                                             OutputIndex,
+                                                             false, // HaveIndexInputIfOutputIndex
+                                                             256,   // BlockSize
+                                                             4,     // MThreadClusterSize
+                                                             64,    // KThreadClusterSize
+                                                             1,     // MThreadSliceSize
+                                                             1,     // KThreadSliceSize
+                                                             0,     // InSrcVectorDim
+                                                             1,     // InSrceVectorSize
+                                                             1>;    // OutDstVectorSize
+
+    Tensor<InOutDataType> in(inLengths);
+
+    std::vector<size_t> outLengths;
+
+    std::vector<int> invariantDims = get_invariant_dims<Rank, NumReduceDim>(reduceDims);
+
+    if(invariantDims.empty())
+        outLengths.push_back(1);
+    else
+        for(auto dim : invariantDims)
+            outLengths.push_back(inLengths[dim]);
+
+    Tensor<InOutDataType> out_ref(outLengths);
+    Tensor<InOutDataType> out(outLengths);
+    Tensor<int> out_indices_ref(outLengths);
+    Tensor<int> out_indices(outLengths);
+
+    auto inStrides  = in.mDesc.GetStrides();
+    auto outStrides = out.mDesc.GetStrides();
+
+    size_t invariant_total_length = out.mDesc.GetElementSize();
+    size_t reduce_total_length    = in.mDesc.GetElementSize() / invariant_total_length;
+
+    std::size_t num_thread = 1;
+
+    if(do_verification)
+    {
+        switch(init_method)
+        {
+        case 0: break;
+        case 1:
+            in.GenerateTensorValue(GeneratorTensor_1<InOutDataType>{1}, num_thread);
+            if(beta != 0.0f)
+                out_ref.GenerateTensorValue(GeneratorTensor_1<InOutDataType>{1}, num_thread);
+            break;
+        case 2:
+            in.GenerateTensorValue(GeneratorTensor_2<InOutDataType>{-5, 5}, num_thread);
+            if(beta != 0.0f)
+                out_ref.GenerateTensorValue(GeneratorTensor_2<InOutDataType>{-5, 5}, num_thread);
+            break;
+        default:
+            in.GenerateTensorValue(GeneratorTensor_3<InOutDataType>{-5.0, 5.0}, num_thread);
+            if(beta != 0.0f)
+                out_ref.GenerateTensorValue(GeneratorTensor_3<InOutDataType>{-5.0, 5.0},
+                                            num_thread);
+        }
+
+        if(beta != 0.0f)
+            for(size_t i = 0; i < out_ref.mDesc.GetElementSpaceSize(); i++)
+                out.mData[i] = out_ref.mData[i];
+    };
+
+    // these buffers are usually provided by the user application
+    DeviceMem in_dev(sizeof(InOutDataType) * in.mDesc.GetElementSpaceSize());
+    DeviceMem out_dev(sizeof(InOutDataType) * out.mDesc.GetElementSpaceSize());
+
+    in_dev.ToDevice(in.mData.data());
+
+    if(beta != 0.0f)
+        out_dev.ToDevice(out.mData.data());
+
+    size_t indicesSizeInBytes = OutputIndex ? out.mDesc.GetElementSize() * sizeof(int32_t) : 0;
+
+    DeviceMem out_index_dev(indicesSizeInBytes);
+
+    InElementwiseOperation in_elementwise_op;
+    AccElementwiseOperation acc_elementwise_op;
+
+    std::tie(in_elementwise_op, acc_elementwise_op) =
+        reduce_unary_operator<ReduceOpId, true, true>::GetElementwiseOperator(
+            static_cast<int32_t>(reduce_total_length));
+
+    if(do_verification)
+    {
+        ReductionHost<InOutDataType,
+                      AccDataType,
+                      InOutDataType,
+                      ReduceOperation,
+                      InElementwiseOperation,
+                      AccElementwiseOperation,
+                      Rank,
+                      NumReduceDim,
+                      PropagateNan,
+                      OutputIndex>
+            hostReduce(in.mDesc, out_ref.mDesc, invariantDims, reduceDims);
+
+        hostReduce.Run(alpha,
+                       in.mData.data(),
+                       beta,
+                       out_ref.mData.data(),
+                       out_indices_ref.mData.data(),
+                       in_elementwise_op,
+                       acc_elementwise_op);
+    };
+
+    std::vector<ck::index_t> i_inLengths;
+    std::vector<ck::index_t> i_inStrides;
+    std::vector<ck::index_t> i_outLengths;
+    std::vector<ck::index_t> i_outStrides;
+
+    i_inLengths.assign(inLengths.begin(), inLengths.end());
+    i_inStrides.assign(inStrides.begin(), inStrides.end());
+    i_outLengths.assign(outLengths.begin(), outLengths.end());
+    i_outStrides.assign(outStrides.begin(), outStrides.end());
+
+    auto reduce = DeviceReduceInstance{};
+
+    auto argument_ptr = reduce.MakeArgumentPointer(i_inLengths,
+                                                   i_inStrides,
+                                                   i_outLengths,
+                                                   i_outStrides,
+                                                   reduceDims,
+                                                   alpha,
+                                                   beta,
+                                                   in_dev.GetDeviceBuffer(),
+                                                   nullptr,
+                                                   out_dev.GetDeviceBuffer(),
+                                                   out_index_dev.GetDeviceBuffer(),
+                                                   in_elementwise_op,
+                                                   acc_elementwise_op);
+
+    if(!reduce.IsSupportedArgument(argument_ptr.get()))
+    {
+        std::cerr
+            << "The runtime parameters seems not supported by the DeviceReduce instance, exiting!"
+            << std::endl;
+
+        return (-2);
+    };
+
+    std::string reduce_name = reduce.GetTypeString();
+
+    auto invoker_ptr = reduce.MakeInvokerPointer();
+
+    float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
+
+    std::size_t num_bytes = invariant_total_length * reduce_total_length * sizeof(InOutDataType) +
+                            invariant_total_length * sizeof(InOutDataType);
+
+    float gb_per_sec = num_bytes / 1.E6 / avg_time;
+
+    std::cout << "Perf: " << avg_time << " ms, " << gb_per_sec << " GB/s, " << reduce_name
+              << std::endl;
+
+    bool pass = true;
+
+    if(do_verification)
+    {
+        out_dev.FromDevice(out.mData.data());
+        pass = pass && ck::utils::check_err(out.mData, out_ref.mData);
+
+        if(OutputIndex)
+        {
+            out_index_dev.FromDevice(out_indices.mData.data());
+            pass = pass && ck::utils::check_err(out_indices.mData, out_indices_ref.mData);
+        };
+    };
+
+    return (pass ? 0 : 1);
+}

example/12_reduce/reduce_example_common.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck/ck.hpp"
+
+template <ck::index_t Rank, ck::index_t NumReduceDim>
+std::vector<int> get_invariant_dims(const std::vector<int>& reduceDims)
+{
+    assert(NumReduceDim == reduceDims.size());
+
+    int reduceFlag = 0;
+
+    // flag the bits for the reduceDims
+    for(int i = 0; i < NumReduceDim; i++)
+    {
+        reduceFlag |= 1 << reduceDims[i];
+    };
+
+    std::vector<int> invariantDims;
+
+    // collect invariant dimensions
+    for(int i = 0; i < Rank; i++)
+        if((reduceFlag & (1 << i)) == 0)
+        {
+            invariantDims.push_back(i);
+        };
+
+    return invariantDims;
+};
+
+template <ck::index_t Rank, ck::index_t NumReduceDim>
+struct ReduceShape
+{
+    static constexpr ck::index_t Rank_         = Rank;
+    static constexpr ck::index_t NumReduceDim_ = NumReduceDim;
+};
+
+using reduce_shape_instances = std::tuple<ReduceShape<3, 1>,
+                                          ReduceShape<3, 2>,
+                                          ReduceShape<4, 1>,
+                                          ReduceShape<4, 2>,
+                                          ReduceShape<4, 3>,
+                                          ReduceShape<5, 1>,
+                                          ReduceShape<5, 2>,
+                                          ReduceShape<5, 3>,
+                                          ReduceShape<5, 4>>;

example/12_reduce/reduce_multiblock_atomic_add.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <initializer_list>
+#include <cstdlib>
+#include <getopt.h>
+
+#include "ck/utility/reduction_enums.hpp"
+#include "reduce_multiblock_atomic_add_impl.hpp"
+#include "reduce_example_common.hpp"
+
+using namespace ck;
+using namespace ck::tensor_operation::device;
+
+static struct option long_options[] = {{"inLengths", required_argument, nullptr, 'D'},
+                                       {"verify", required_argument, nullptr, 'v'},
+                                       {"help", no_argument, nullptr, '?'},
+                                       {nullptr, 0, nullptr, 0}};
+
+class SimpleAppArgs
+{
+    private:
+    int option_index = 0;
+
+    public:
+    std::vector<size_t> inLengths = {16, 64, 32, 960};
+    std::vector<int> reduceDims   = {0, 1, 2};
+    std::vector<float> scales     = {1.0f, 0.0f};
+
+    bool do_verification = true;
+    int data_type        = 1;
+    int init_method      = 2;
+    bool time_kernel     = true;
+
+    public:
+    void show_usage(const char* cmd)
+    {
+        std::cout << "Usage of " << cmd << std::endl;
+        std::cout << "--inLengths or -D, comma separated list of input tensor dimension lengths"
+                  << std::endl;
+        std::cout << "--reduceDims or -R, comma separated list of to-reduce dimensions"
+                  << std::endl;
+        std::cout << "--verify or -v, 1/0 to indicate whether to verify the reduction result by "
+                     "comparing with the host-based reduction"
+                  << std::endl;
+        std::cout << "Arg1: data type (0: fp32, 1: fp64)" << std::endl;
+        std::cout << "Arg2 -- init method (0=no init, 1=single integer value, 2=scope integer "
+                     "value, 3=decimal value)"
+                  << std::endl;
+        std::cout << "Arg3 -- time kernel (0=no, 1=yes)" << std::endl;
+    };
+
+    int processArgs(int argc, char* argv[])
+    {
+        using ck::host_common::getTypeValuesFromString;
+
+        int ch;
+
+        while(1)
+        {
+            ch = getopt_long(argc, argv, "D:R:v:l:", long_options, &option_index);
+            if(ch == -1)
+                break;
+            switch(ch)
+            {
+            case 'D':
+                if(!optarg)
+                    throw std::runtime_error("Invalid option format!");
+
+                inLengths = getTypeValuesFromString<size_t>(optarg);
+                break;
+            case 'R':
+                if(!optarg)
+                    throw std::runtime_error("Invalid option format!");
+
+                reduceDims = getTypeValuesFromString<int>(optarg);
+                break;
+            case 'v':
+                if(!optarg)
+                    throw std::runtime_error("Invalid option format!");
+
+                do_verification = static_cast<bool>(std::atoi(optarg));
+                break;
+            case '?':
+                if(std::string(long_options[option_index].name) == "help")
+                {
+                    show_usage(argv[0]);
+                    return (-1);
+                };
+                break;
+            default: show_usage(argv[0]); return (-1);
+            };
+        };
+
+        if(optind + 3 > argc)
+        {
+            throw std::runtime_error("Invalid cmd-line arguments, more argumetns are needed!");
+        };
+
+        data_type   = std::atoi(argv[optind++]);
+        init_method = std::atoi(argv[optind++]);
+        time_kernel = static_cast<bool>(std::atoi(argv[optind]));
+
+        if(scales.empty())
+        {
+            scales.push_back(1.0f);
+            scales.push_back(0.0f);
+        };
+
+        return (0);
+    };
+};
+
+template <typename InOutDataType,
+          typename AccDataType,
+          ReduceTensorOp ReduceOpId,
+          index_t PropagateNan>
+bool reduce_multiblock_atomic_add_test(bool do_verification,
+                                       int init_method,
+                                       bool time_kernel,
+                                       const std::vector<size_t>& inLengths,
+                                       const std::vector<int>& reduceDims,
+                                       float alpha,
+                                       float beta)
+{
+    bool matched = false;
+    int result   = 0;
+
+    const auto tuple_object = reduce_shape_instances{};
+
+    static_for<0, std::tuple_size<reduce_shape_instances>::value, 1>{}([&](auto i) {
+        if(matched)
+            return;
+
+        using ShapeType = remove_cvref_t<decltype(std::get<i>(tuple_object))>;
+
+        if(ShapeType::Rank_ != inLengths.size() || ShapeType::NumReduceDim_ != reduceDims.size())
+            return;
+
+        result = reduce_multiblock_atomic_add_impl<InOutDataType,
+                                                   AccDataType,
+                                                   ReduceOpId,
+                                                   ShapeType::Rank_,
+                                                   ShapeType::NumReduceDim_,
+                                                   PropagateNan>(
+            do_verification, init_method, time_kernel, inLengths, reduceDims, alpha, beta);
+
+        matched = true;
+    });
+
+    return (result == 0) ? true : false;
+};
+
+constexpr ReduceTensorOp ReduceOpId = ReduceTensorOp::AVG;
+constexpr bool PropagateNan         = true;
+
+int main(int argc, char* argv[])
+{
+    bool pass = true;
+
+    if(argc > 1)
+    {
+        SimpleAppArgs arg;
+
+        if(arg.processArgs(argc, argv) < 0)
+            return (-1);
+
+        if(arg.data_type == 0)
+        {
+            pass = reduce_multiblock_atomic_add_test<float, float, ReduceOpId, PropagateNan>(
+                arg.do_verification,
+                arg.init_method,
+                arg.time_kernel,
+                arg.inLengths,
+                arg.reduceDims,
+                arg.scales[0],
+                arg.scales[1]);
+        }
+        else if(arg.data_type == 1)
+        {
+            pass = reduce_multiblock_atomic_add_test<double, double, ReduceOpId, PropagateNan>(
+                arg.do_verification,
+                arg.init_method,
+                arg.time_kernel,
+                arg.inLengths,
+                arg.reduceDims,
+                arg.scales[0],
+                arg.scales[1]);
+        }
+    }
+    else
+    {
+        // for testing float
+        pass = pass && reduce_multiblock_atomic_add_test<float, float, ReduceOpId, PropagateNan>(
+                           true, 2, false, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
+
+        // for testing double
+        pass = pass && reduce_multiblock_atomic_add_test<double, double, ReduceOpId, PropagateNan>(
+                           true, 2, false, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
+
+        // for testing 3D input
+        pass = pass && reduce_multiblock_atomic_add_test<float, float, ReduceOpId, PropagateNan>(
+                           true, 2, false, {16, 64, 960}, {0, 1}, 1.0f, 0.0f);
+
+        // for testing 5D input
+        pass = pass && reduce_multiblock_atomic_add_test<float, float, ReduceOpId, PropagateNan>(
+                           true, 2, false, {16, 64, 32, 2, 960}, {0, 1, 2, 3}, 1.0f, 0.0f);
+    };
+
+    return (pass ? 0 : 1);
+};

example/12_reduce/reduce_multiblock_atomic_add_impl.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+
+#include "ck/ck.hpp"
+#include "ck/utility/reduction_enums.hpp"
+#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
+#include "ck/tensor_operation/gpu/device/device_reduce_multiblock.hpp"
+
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/utility/host_common_util.hpp"
+#include "ck/library/utility/host_reduction.hpp"
+
+#include "reduce_example_common.hpp"
+
+template <typename InOutDataType,
+          typename AccDataType,
+          ck::ReduceTensorOp ReduceOpId,
+          ck::index_t Rank,
+          ck::index_t NumReduceDim,
+          bool PropagateNan>
+int reduce_multiblock_atomic_add_impl(bool do_verification,
+                                      int init_method,
+                                      bool time_kernel,
+                                      const std::vector<size_t>& inLengths,
+                                      const std::vector<int>& reduceDims,
+                                      float alpha,
+                                      float beta)
+
+{
+    using namespace ck;
+    using namespace ck::tensor_operation::device;
+
+    constexpr bool op_support_atomic_add =
+        (ReduceOpId == ReduceTensorOp::ADD || ReduceOpId == ReduceTensorOp::AVG);
+
+    constexpr bool invalid_reduce_1 = !op_support_atomic_add;
+    constexpr bool invalid_reduce_2 =
+        !(std::is_same<InOutDataType, float>::value || std::is_same<InOutDataType, double>::value);
+
+    constexpr bool invalid_reduce = (invalid_reduce_1 || invalid_reduce_2);
+
+    if(invalid_reduce)
+    {
+        std::cerr << "The reduction setting is invalid, exiting!" << std::endl;
+        return (-1);
+    };
+
+    using ReduceOperation = typename reduce_binary_operator<ReduceOpId>::opType;
+    using InElementwiseOperation =
+        typename reduce_unary_operator<ReduceOpId, true, true>::InElementwiseOperation;
+    using AccElementwiseOperation =
+        typename reduce_unary_operator<ReduceOpId, true, true>::AccElementwiseOperation;
+
+    using DeviceReduceInstance =
+        ck::tensor_operation::device::DeviceReduceMultiBlock<InOutDataType,
+                                                             AccDataType,
+                                                             InOutDataType,
+                                                             Rank,
+                                                             NumReduceDim,
+                                                             ReduceOperation,
+                                                             InElementwiseOperation,
+                                                             AccElementwiseOperation,
+                                                             InMemoryDataOperationEnum::AtomicAdd,
+                                                             PropagateNan,
+                                                             false,
+                                                             false, // HaveIndexInputIfOutputIndex
+                                                             256,
+                                                             4,
+                                                             64,
+                                                             1,
+                                                             1,
+                                                             0,
+                                                             1,
+                                                             1>;
+
+    Tensor<InOutDataType> in(inLengths);
+
+    std::vector<size_t> outLengths;
+
+    std::vector<int> invariantDims = get_invariant_dims<Rank, NumReduceDim>(reduceDims);
+
+    if(invariantDims.empty())
+        outLengths.push_back(1);
+    else
+        for(auto dim : invariantDims)
+            outLengths.push_back(inLengths[dim]);
+
+    Tensor<InOutDataType> out_ref(outLengths);
+    Tensor<InOutDataType> out(outLengths);
+
+    auto inStrides  = in.mDesc.GetStrides();
+    auto outStrides = out.mDesc.GetStrides();
+
+    size_t invariant_total_length = out.mDesc.GetElementSize();
+    size_t reduce_total_length    = in.mDesc.GetElementSize() / invariant_total_length;
+
+    std::size_t num_thread = 1;
+
+    if(do_verification)
+    {
+        switch(init_method)
+        {
+        case 0: break;
+        case 1:
+            in.GenerateTensorValue(GeneratorTensor_1<InOutDataType>{1}, num_thread);
+            if(beta != 0.0f)
+                out_ref.GenerateTensorValue(GeneratorTensor_1<InOutDataType>{1}, num_thread);
+            break;
+        case 2:
+            in.GenerateTensorValue(GeneratorTensor_2<InOutDataType>{-5, 5}, num_thread);
+            if(beta != 0.0f)
+                out_ref.GenerateTensorValue(GeneratorTensor_2<InOutDataType>{-5, 5}, num_thread);
+            break;
+        default:
+            in.GenerateTensorValue(GeneratorTensor_3<InOutDataType>{-5.0, 5.0}, num_thread);
+            if(beta != 0.0f)
+                out_ref.GenerateTensorValue(GeneratorTensor_3<InOutDataType>{-5.0, 5.0},
+                                            num_thread);
+        }
+
+        if(beta != 0.0f)
+            for(size_t i = 0; i < out_ref.mDesc.GetElementSpaceSize(); i++)
+                out.mData[i] = out_ref.mData[i];
+    };
+
+    // these buffers are usually provided by the user application
+    DeviceMem in_dev(sizeof(InOutDataType) * in.mDesc.GetElementSpaceSize());
+    DeviceMem out_dev(sizeof(InOutDataType) * out.mDesc.GetElementSpaceSize());
+
+    in_dev.ToDevice(in.mData.data());
+
+    if(beta != 0.0f)
+        out_dev.ToDevice(out.mData.data());
+
+    InElementwiseOperation in_elementwise_op;
+    AccElementwiseOperation acc_elementwise_op;
+
+    std::tie(in_elementwise_op, acc_elementwise_op) =
+        reduce_unary_operator<ReduceOpId, true, true>::GetElementwiseOperator(
+            static_cast<int32_t>(reduce_total_length));
+
+    if(do_verification)
+    {
+        ReductionHost<InOutDataType,
+                      AccDataType,
+                      InOutDataType,
+                      ReduceOperation,
+                      InElementwiseOperation,
+                      AccElementwiseOperation,
+                      Rank,
+                      NumReduceDim,
+                      PropagateNan,
+                      false>
+            hostReduce(in.mDesc, out_ref.mDesc, invariantDims, reduceDims);
+
+        hostReduce.Run(alpha,
+                       in.mData.data(),
+                       beta,
+                       out_ref.mData.data(),
+                       nullptr,
+                       in_elementwise_op,
+                       acc_elementwise_op);
+    };
+
+    std::vector<ck::index_t> i_inLengths;
+    std::vector<ck::index_t> i_inStrides;
+    std::vector<ck::index_t> i_outLengths;
+    std::vector<ck::index_t> i_outStrides;
+
+    i_inLengths.assign(inLengths.begin(), inLengths.end());
+    i_inStrides.assign(inStrides.begin(), inStrides.end());
+    i_outLengths.assign(outLengths.begin(), outLengths.end());
+    i_outStrides.assign(outStrides.begin(), outStrides.end());
+
+    auto reduce = DeviceReduceInstance{};
+
+    auto argument_ptr = reduce.MakeArgumentPointer(i_inLengths,
+                                                   i_inStrides,
+                                                   i_outLengths,
+                                                   i_outStrides,
+                                                   reduceDims,
+                                                   alpha,
+                                                   beta,
+                                                   in_dev.GetDeviceBuffer(),
+                                                   nullptr,
+                                                   out_dev.GetDeviceBuffer(),
+                                                   nullptr,
+                                                   in_elementwise_op,
+                                                   acc_elementwise_op);
+
+    if(!reduce.IsSupportedArgument(argument_ptr.get()))
+    {
+        std::cerr
+            << "The runtime parameters seems not supported by the DeviceReduce instance, exiting!"
+            << std::endl;
+
+        return (-2);
+    };
+
+    std::string reduce_name = reduce.GetTypeString();
+
+    auto invoker_ptr = reduce.MakeInvokerPointer();
+
+    float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
+
+    std::size_t num_bytes = invariant_total_length * reduce_total_length * sizeof(InOutDataType) +
+                            invariant_total_length * sizeof(InOutDataType);
+
+    float gb_per_sec = num_bytes / 1.E6 / avg_time;
+
+    std::cout << "Perf: " << avg_time << " ms, " << gb_per_sec << " GB/s, " << reduce_name
+              << std::endl;
+
+    bool pass = true;
+
+    if(do_verification)
+    {
+        out_dev.FromDevice(out.mData.data());
+        pass = pass && ck::utils::check_err(out.mData, out_ref.mData);
+    };
+
+    return (pass ? 0 : 1);
+}

example/16_gemm_multi_d_multi_reduces/CMakeLists.txt

+add_example_executable(example_gemm_add_add_mean_meansquare_xdl_fp16 gemm_add_add_mean_meansquare_xdl_fp16.cpp)
+add_example_executable(example_gemm_max_xdl_fp16 gemm_max_xdl_fp16.cpp)
+add_example_executable(example_gemm_mean_meansquare_xdl_fp16 gemm_mean_meansquare_xdl_fp16.cpp)

example/16_gemm_multi_d_multi_reduces/gemm_add_add_mean_meansquare_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/library/utility/check_err.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using F16 = ck::half_t;
+using F32 = float;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+// DataType
+using ADataType         = F16;
+using BDataType         = F16;
+using GemmAccDataType   = F32;
+using CShuffleDataType  = F32;
+using D0DataType        = F16;
+using D1DataType        = F16;
+using DsDataType        = ck::Tuple<D0DataType, D1DataType>;
+using EDataType         = F16;
+using ReduceAccDataType = F32;
+using R0DataType        = F32;
+using R1DataType        = F32;
+using RsDataType        = ck::Tuple<R0DataType, R1DataType>;
+
+// Layout
+using ALayout  = Row;
+using BLayout  = Col;
+using D1Layout = Row;
+using ELayout  = D1Layout;
+
+// Elementwise op
+using PassThrough  = ck::tensor_operation::element_wise::PassThrough;
+using AddAdd       = ck::tensor_operation::element_wise::AddAdd;
+using Square       = ck::tensor_operation::element_wise::UnarySquare;
+using Div          = ck::tensor_operation::element_wise::UnaryDivide;
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = AddAdd;
+using QsElementOp  = ck::Tuple<PassThrough, Square>;
+using RsElementOp  = ck::Tuple<Div, Div>;
+
+// ReduceOp
+using R0ThreadReduceOp = ck::reduce::Add;
+using R1ThreadReduceOp = ck::reduce::Add;
+using RsThreadReduceOp = ck::Tuple<R0ThreadReduceOp, R1ThreadReduceOp>;
+
+static constexpr auto R0GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+static constexpr auto R1GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+using RsGlobalReduceOp = ck::InMemoryDataOperationEnumSequence<R0GlobalReduceOp, R1GlobalReduceOp>;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+// clang-format off
+using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleDMultipleR_Xdl_CShuffle
+//######| ALayout| BLayout| ELayout|     AData|     BData|     GemmAccData|         CShuffle|     DsData|     EData|     ReduceAccData|     RsData|           A|           B|          CDE|          Qs|          Rs|           Thread|           Global|           GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle|    CDRThreadTransfer|                  CDE|    RThreadTransfer|
+//######|        |        |        |      Type|      Type|            Type|         DataType|       Type|      Type|              Type|       Type| Elementwise| Elementwise|  Elementwise| Elementwise| Elementwise|           Reduce|           Reduce| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|       ClusterLengths| ReduceThreadTransfer| DstScalarPerVector|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |   Operation|   Operation|    Operation|   Operation|   Operation|        Operation|        Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle| _MPerBlock_NPerBlock|      ScalarPerVector|         _MPerBlock|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |            |            |             |            |            |                 |                 |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                     |           _NPerBlock|                   |
+        < ALayout, BLayout, ELayout, ADataType, BDataType, GemmAccDataType, CShuffleDataType, DsDataType, EDataType, ReduceAccDataType, RsDataType,  AElementOp,  BElementOp, CDEElementOp, QsElementOp, RsElementOp, RsThreadReduceOp, RsGlobalReduceOp,    GemmDefault,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,             S<64, 4>,                    4,                  1>;
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        EDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        PassThrough>;
+
+template <typename ADataType,
+          typename BDataType,
+          typename D0DataType,
+          typename D1DataType,
+          typename EDataType,
+          typename R0DataType,
+          typename R1DataType>
+void DumpPerf(float ave_time, int M, int N, int K)
+{
+    std::size_t flop          = std::size_t(2) * M * N * K + std::size_t(2) * M * N;
+    std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
+                                sizeof(D0DataType) * M * N + sizeof(D1DataType) * M * N +
+                                sizeof(EDataType) * M * N + sizeof(R0DataType) * M +
+                                sizeof(R1DataType) * M;
+
+    float tflops          = static_cast<float>(flop) / 1.E9 / ave_time;
+    float gemm_gb_per_sec = gemm_num_byte / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gemm_gb_per_sec
+              << " GB/s, " << std::endl;
+}
+
+auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
+    return HostTensorDescriptor(std::vector<std::size_t>({len}),
+                                std::vector<std::size_t>({stride}));
+};
+
+auto f_host_tensor_descriptor2d =
+    [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+        if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+        {
+            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                        std::vector<std::size_t>({stride, 1}));
+        }
+        else
+        {
+            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                        std::vector<std::size_t>({1, stride}));
+        }
+    };
+
+int main()
+{
+    ck::index_t M = 1024;
+    ck::index_t N = 1024;
+    ck::index_t K = 1024;
+
+    ck::index_t StrideA  = 1024;
+    ck::index_t StrideB  = 1024;
+    ck::index_t StrideD0 = 0;
+    ck::index_t StrideD1 = 1024;
+    ck::index_t StrideE  = 1024;
+
+    Tensor<ADataType> a_m_k(f_host_tensor_descriptor2d(M, K, StrideA, ALayout{}));
+    Tensor<BDataType> b_k_n(f_host_tensor_descriptor2d(K, N, StrideB, BLayout{}));
+    Tensor<D0DataType> d0_n(f_host_tensor_descriptor1d(N, 1));
+    Tensor<D1DataType> d1_m_n(f_host_tensor_descriptor2d(M, N, StrideD1, D1Layout{}));
+    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> r0_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<R1DataType> r1_m(f_host_tensor_descriptor1d(M, 1));
+
+    a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{-1, 1});
+    b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-1, 1});
+    d0_n.GenerateTensorValue(GeneratorTensor_3<D0DataType>{-1, 1});
+    d1_m_n.GenerateTensorValue(GeneratorTensor_3<D1DataType>{-1, 1});
+
+    DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
+    DeviceMem d0_device_buf(sizeof(D0DataType) * d0_n.mDesc.GetElementSpaceSize());
+    DeviceMem d1_device_buf(sizeof(D1DataType) * d1_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem r0_device_buf(sizeof(R0DataType) * r0_m.mDesc.GetElementSpaceSize());
+    DeviceMem r1_device_buf(sizeof(R1DataType) * r1_m.mDesc.GetElementSpaceSize());
+
+    a_device_buf.ToDevice(a_m_k.mData.data());
+    b_device_buf.ToDevice(b_k_n.mData.data());
+    d0_device_buf.ToDevice(d0_n.mData.data());
+    d1_device_buf.ToDevice(d1_m_n.mData.data());
+
+    auto a_element_op   = AElementOp{};
+    auto b_element_op   = BElementOp{};
+    auto cde_element_op = CDEElementOp{};
+    auto qs_element_op  = QsElementOp{};
+    auto rs_element_op  = RsElementOp{N, N};
+
+    // Prepare GEMM, mean, mean_square
+    auto device_op = DeviceOpInstance{};
+    auto invoker   = device_op.MakeInvoker();
+    auto argument =
+        device_op.MakeArgument(a_device_buf.GetDeviceBuffer(),
+                               b_device_buf.GetDeviceBuffer(),
+                               {d0_device_buf.GetDeviceBuffer(), d1_device_buf.GetDeviceBuffer()},
+                               e_device_buf.GetDeviceBuffer(),
+                               {r0_device_buf.GetDeviceBuffer(), r1_device_buf.GetDeviceBuffer()},
+                               M,
+                               N,
+                               K,
+                               StrideA,
+                               StrideB,
+                               {StrideD0, StrideD1},
+                               StrideE,
+                               a_element_op,
+                               b_element_op,
+                               cde_element_op,
+                               qs_element_op,
+                               rs_element_op);
+
+    if(!device_op.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error("wrong! this device_op instance does not support this problem");
+    }
+
+    // init reducetion buffer to 0
+    r0_device_buf.SetZero();
+    r1_device_buf.SetZero();
+
+    invoker.Run(argument, StreamConfig{nullptr, false});
+
+    bool do_verification = true;
+    bool pass            = true;
+
+    if(do_verification)
+    {
+        auto I0 = ck::Number<0>{};
+        auto I1 = ck::Number<1>{};
+
+        Tensor<EDataType> e_m_n_host(e_m_n.mDesc);
+        Tensor<R0DataType> r0_m_host(r0_m.mDesc);
+        Tensor<R1DataType> r1_m_host(r1_m.mDesc);
+
+        auto ref_gemm    = ReferenceGemmInstance{};
+        auto ref_invoker = ref_gemm.MakeInvoker();
+
+        auto ref_argument = ref_gemm.MakeArgument(
+            a_m_k, b_k_n, e_m_n_host, a_element_op, b_element_op, PassThrough{});
+
+        ref_invoker.Run(ref_argument);
+
+        auto reduce0_op = R0ThreadReduceOp{};
+        auto reduce1_op = R1ThreadReduceOp{};
+
+        for(int m = 0; m < M; ++m)
+        {
+            auto reduce0_acc = reduce0_op.GetIdentityValue<ReduceAccDataType>();
+            auto reduce1_acc = reduce1_op.GetIdentityValue<ReduceAccDataType>();
+
+            for(int n = 0; n < N; ++n)
+            {
+                ReduceAccDataType square_e_val;
+
+                auto e_val  = ck::type_convert<GemmAccDataType>(e_m_n_host(m, n));
+                auto d0_val = ck::type_convert<GemmAccDataType>(d0_n(n));
+                auto d1_val = ck::type_convert<GemmAccDataType>(d1_m_n(m, n));
+                cde_element_op(e_val, e_val, d0_val, d1_val);
+                e_m_n_host(m, n) = ck::type_convert<EDataType>(e_val);
+
+                auto e_val_reduce = ck::type_convert<ReduceAccDataType>(e_val);
+                qs_element_op[I1](square_e_val, e_val_reduce);
+
+                reduce0_op(reduce0_acc, e_val_reduce);
+                reduce1_op(reduce1_acc, square_e_val);
+            }
+
+            rs_element_op[I0](reduce0_acc, reduce0_acc);
+            rs_element_op[I1](reduce1_acc, reduce1_acc);
+            r0_m_host(m) = ck::type_convert<R0DataType>(reduce0_acc);
+            r1_m_host(m) = ck::type_convert<R1DataType>(reduce1_acc);
+        }
+
+        e_device_buf.FromDevice(e_m_n.mData.data());
+        r0_device_buf.FromDevice(r0_m.mData.data());
+        r1_device_buf.FromDevice(r1_m.mData.data());
+
+        pass = ck::utils::check_err(
+            e_m_n.mData, e_m_n_host.mData, "Error: Incorrect results c", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(
+            r0_m.mData, r0_m_host.mData, "Error: Incorrect results d0", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(
+            r1_m.mData, r1_m_host.mData, "Error: Incorrect results d1", 1e-2, 1e-2);
+    }
+
+    bool time_kernel = true;
+    if(time_kernel)
+    {
+        float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+        DumpPerf<ADataType, BDataType, D0DataType, D1DataType, EDataType, R0DataType, R1DataType>(
+            ave_time, M, N, K);
+    }
+
+    return pass ? 0 : 1;
+}

example/16_gemm_multi_d_multi_reduces/gemm_max_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/library/utility/check_err.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using F16 = ck::half_t;
+using F32 = float;
+using F64 = double;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+// DataType
+using ADataType         = F16;
+using BDataType         = F16;
+using GemmAccDataType   = F32;
+using CShuffleDataType  = F32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = F16;
+using ReduceAccDataType = F32;
+using R0DataType        = F32;
+using RsDataType        = ck::Tuple<R0DataType>;
+
+// Layout
+using ALayout = Row;
+using BLayout = Col;
+using ELayout = Row;
+
+// Elementwise op
+using PassThrough  = ck::tensor_operation::element_wise::PassThrough;
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+using QsElementOp  = ck::Tuple<PassThrough>;
+using RsElementOp  = ck::Tuple<PassThrough>;
+
+// ReduceOp
+using RsThreadReduceOp = ck::Tuple<ck::reduce::Max>;
+
+using RsGlobalReduceOp =
+    ck::InMemoryDataOperationEnumSequence<ck::InMemoryDataOperationEnum::AtomicMax>;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+// clang-format off
+using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleDMultipleR_Xdl_CShuffle
+//######| ALayout| BLayout| ELayout|     AData|     BData|     GemmAccData|         CShuffle|     DsData|     EData|     ReduceAccData|     RsData|           A|           B|          CDE|          Qs|          Rs|           Thread|           Global|           GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle|    CDRThreadTransfer|                  CDE|    RThreadTransfer|
+//######|        |        |        |      Type|      Type|            Type|         DataType|       Type|      Type|              Type|       Type| Elementwise| Elementwise|  Elementwise| Elementwise| Elementwise|           Reduce|           Reduce| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|       ClusterLengths| ReduceThreadTransfer| DstScalarPerVector|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |   Operation|   Operation|    Operation|   Operation|   Operation|        Operation|        Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle| _MPerBlock_NPerBlock|      ScalarPerVector|         _MPerBlock|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |            |            |             |            |            |                 |                 |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                     |           _NPerBlock|                   |
+        < ALayout, BLayout, ELayout, ADataType, BDataType, GemmAccDataType, CShuffleDataType, DsDataType, EDataType, ReduceAccDataType, RsDataType,  AElementOp,  BElementOp, CDEElementOp, QsElementOp, RsElementOp, RsThreadReduceOp, RsGlobalReduceOp,    GemmDefault,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,             S<64, 4>,                    4,                  1>;
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        EDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+template <typename ADataType, typename BDataType, typename EDataType, typename R0DataType>
+void DumpPerf(float ave_time, int M, int N, int K)
+{
+    std::size_t flop          = std::size_t(2) * M * N * K;
+    std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
+                                sizeof(EDataType) * M * N + sizeof(R0DataType) * M;
+
+    float tflops          = static_cast<float>(flop) / 1.E9 / ave_time;
+    float gemm_gb_per_sec = gemm_num_byte / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gemm_gb_per_sec
+              << " GB/s, " << std::endl;
+}
+
+auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
+    return HostTensorDescriptor(std::vector<std::size_t>({len}),
+                                std::vector<std::size_t>({stride}));
+};
+
+auto f_host_tensor_descriptor2d =
+    [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+        if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+        {
+            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                        std::vector<std::size_t>({stride, 1}));
+        }
+        else
+        {
+            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                        std::vector<std::size_t>({1, stride}));
+        }
+    };
+
+int main()
+{
+    ck::index_t M = 1024;
+    ck::index_t N = 1024;
+    ck::index_t K = 1024;
+
+    ck::index_t StrideA = 1024;
+    ck::index_t StrideB = 1024;
+    ck::index_t StrideE = 1024;
+
+    Tensor<ADataType> a_m_k(f_host_tensor_descriptor2d(M, K, StrideA, ALayout{}));
+    Tensor<BDataType> b_k_n(f_host_tensor_descriptor2d(K, N, StrideB, BLayout{}));
+    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> r0_m(f_host_tensor_descriptor1d(M, 1));
+
+    a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{-1, 1});
+    b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-1, 1});
+
+    DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem r0_device_buf(sizeof(R0DataType) * r0_m.mDesc.GetElementSpaceSize());
+
+    a_device_buf.ToDevice(a_m_k.mData.data());
+    b_device_buf.ToDevice(b_k_n.mData.data());
+
+    auto a_element_op   = AElementOp{};
+    auto b_element_op   = BElementOp{};
+    auto cde_element_op = CDEElementOp{};
+    auto qs_element_op  = QsElementOp{};
+    auto rs_element_op  = RsElementOp{};
+
+    // Prepare GEMM, max
+    auto device_op = DeviceOpInstance{};
+    auto invoker   = device_op.MakeInvoker();
+    auto argument  = device_op.MakeArgument(a_device_buf.GetDeviceBuffer(),
+                                           b_device_buf.GetDeviceBuffer(),
+                                           {},
+                                           e_device_buf.GetDeviceBuffer(),
+                                           {r0_device_buf.GetDeviceBuffer()},
+                                           M,
+                                           N,
+                                           K,
+                                           StrideA,
+                                           StrideB,
+                                           {},
+                                           StrideE,
+                                           a_element_op,
+                                           b_element_op,
+                                           cde_element_op,
+                                           qs_element_op,
+                                           rs_element_op);
+
+    if(!device_op.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error("wrong! this device_op instance does not support this problem");
+    }
+
+    // [CAUSION]: launch_and_time_kernel will not initialize D.
+    // If we evaluate kernel multiple time but without initialize D. Verification will fail
+    r0_device_buf.SetValue(ck::NumericLimits<R0DataType>::Lowest());
+
+    invoker.Run(argument, StreamConfig{nullptr, false});
+
+    bool do_verification = true;
+    bool pass            = true;
+
+    if(do_verification)
+    {
+        auto I0 = ck::Number<0>{};
+
+        Tensor<EDataType> e_m_n_host(e_m_n.mDesc);
+        Tensor<R0DataType> r0_m_host(r0_m.mDesc);
+
+        auto ref_gemm    = ReferenceGemmInstance{};
+        auto ref_invoker = ref_gemm.MakeInvoker();
+
+        auto ref_argument = ref_gemm.MakeArgument(
+            a_m_k, b_k_n, e_m_n_host, a_element_op, b_element_op, cde_element_op);
+
+        ref_invoker.Run(ref_argument);
+
+        auto reduce0_op = RsThreadReduceOp{}[I0];
+
+        for(int m = 0; m < M; ++m)
+        {
+            auto reduce0_acc = reduce0_op.GetIdentityValue<ReduceAccDataType>();
+
+            for(int n = 0; n < N; ++n)
+            {
+                auto e_val = ck::type_convert<ReduceAccDataType>(e_m_n_host(m, n));
+                reduce0_op(reduce0_acc, e_val);
+            };
+
+            r0_m_host(m) = ck::type_convert<R0DataType>(reduce0_acc);
+        }
+
+        e_device_buf.FromDevice(e_m_n.mData.data());
+        r0_device_buf.FromDevice(r0_m.mData.data());
+
+        pass = ck::utils::check_err(
+            e_m_n.mData, e_m_n_host.mData, "Error: Incorrect results c", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(
+            r0_m.mData, r0_m_host.mData, "Error: Incorrect results d0", 1e-2, 1e-2);
+    }
+
+    bool time_kernel = true;
+    if(time_kernel)
+    {
+        float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+        DumpPerf<ADataType, BDataType, EDataType, R0DataType>(ave_time, M, N, K);
+    }
+
+    return pass ? 0 : 1;
+}

example/16_gemm_multi_d_multi_reduces/gemm_mean_meansquare_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/library/utility/check_err.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using F16 = ck::half_t;
+using F32 = float;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+// DataType
+using ADataType         = F16;
+using BDataType         = F16;
+using GemmAccDataType   = F32;
+using CShuffleDataType  = F32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = F16;
+using ReduceAccDataType = F32;
+using R0DataType        = F32;
+using R1DataType        = F32;
+using RsDataType        = ck::Tuple<R0DataType, R1DataType>;
+
+// Layout
+using ALayout = Row;
+using BLayout = Col;
+using ELayout = Row;
+
+// Elementwise op
+using PassThrough  = ck::tensor_operation::element_wise::PassThrough;
+using Square       = ck::tensor_operation::element_wise::UnarySquare;
+using Div          = ck::tensor_operation::element_wise::UnaryDivide;
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+using QsElementOp  = ck::Tuple<PassThrough, Square>;
+using RsElementOp  = ck::Tuple<Div, Div>;
+
+// ReduceOp
+using R0ThreadReduceOp = ck::reduce::Add;
+using R1ThreadReduceOp = ck::reduce::Add;
+using RsThreadReduceOp = ck::Tuple<R0ThreadReduceOp, R1ThreadReduceOp>;
+
+static constexpr auto R0GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+static constexpr auto R1GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+using RsGlobalReduceOp = ck::InMemoryDataOperationEnumSequence<R0GlobalReduceOp, R1GlobalReduceOp>;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+// clang-format off
+using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleDMultipleR_Xdl_CShuffle
+//######| ALayout| BLayout| ELayout|     AData|     BData|     GemmAccData|         CShuffle|     DsData|     EData|     ReduceAccData|     RsData|           A|           B|          CDE|          Qs|          Rs|           Thread|           Global|           GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle|    CDRThreadTransfer|                  CDE|    RThreadTransfer|
+//######|        |        |        |      Type|      Type|            Type|         DataType|       Type|      Type|              Type|       Type| Elementwise| Elementwise|  Elementwise| Elementwise| Elementwise|           Reduce|           Reduce| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|       ClusterLengths| ReduceThreadTransfer| DstScalarPerVector|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |   Operation|   Operation|    Operation|   Operation|   Operation|        Operation|        Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle| _MPerBlock_NPerBlock|      ScalarPerVector|         _MPerBlock|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |            |            |             |            |            |                 |                 |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                     |           _NPerBlock|                   |
+        < ALayout, BLayout, ELayout, ADataType, BDataType, GemmAccDataType, CShuffleDataType, DsDataType, EDataType, ReduceAccDataType, RsDataType,  AElementOp,  BElementOp, CDEElementOp, QsElementOp, RsElementOp, RsThreadReduceOp, RsGlobalReduceOp,    GemmDefault,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,             S<64, 4>,                    4,                  1>;
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        EDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+template <typename ADataType,
+          typename BDataType,
+          typename EDataType,
+          typename R0DataType,
+          typename R1DataType>
+void DumpPerf(float ave_time, int M, int N, int K)
+{
+    std::size_t flop          = std::size_t(2) * M * N * K;
+    std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
+                                sizeof(EDataType) * M * N + sizeof(R0DataType) * M +
+                                sizeof(R1DataType) * M;
+
+    float tflops          = static_cast<float>(flop) / 1.E9 / ave_time;
+    float gemm_gb_per_sec = gemm_num_byte / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gemm_gb_per_sec
+              << " GB/s, " << std::endl;
+}
+
+auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
+    return HostTensorDescriptor(std::vector<std::size_t>({len}),
+                                std::vector<std::size_t>({stride}));
+};
+
+auto f_host_tensor_descriptor2d =
+    [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+        if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+        {
+            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                        std::vector<std::size_t>({stride, 1}));
+        }
+        else
+        {
+            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                        std::vector<std::size_t>({1, stride}));
+        }
+    };
+
+int main()
+{
+    ck::index_t M = 1024;
+    ck::index_t N = 1024;
+    ck::index_t K = 1024;
+
+    ck::index_t StrideA = 1024;
+    ck::index_t StrideB = 1024;
+    ck::index_t StrideE = 1024;
+
+    Tensor<ADataType> a_m_k(f_host_tensor_descriptor2d(M, K, StrideA, ALayout{}));
+    Tensor<BDataType> b_k_n(f_host_tensor_descriptor2d(K, N, StrideB, BLayout{}));
+    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> r0_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<R1DataType> r1_m(f_host_tensor_descriptor1d(M, 1));
+
+    a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{-1, 1});
+    b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-1, 1});
+
+    DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem r0_device_buf(sizeof(R0DataType) * r0_m.mDesc.GetElementSpaceSize());
+    DeviceMem r1_device_buf(sizeof(R1DataType) * r1_m.mDesc.GetElementSpaceSize());
+
+    a_device_buf.ToDevice(a_m_k.mData.data());
+    b_device_buf.ToDevice(b_k_n.mData.data());
+
+    auto a_element_op   = AElementOp{};
+    auto b_element_op   = BElementOp{};
+    auto cde_element_op = CDEElementOp{};
+    auto qs_element_op  = QsElementOp{};
+    auto rs_element_op  = RsElementOp{N, N};
+
+    // Prepare GEMM, mean, mean_square
+    auto device_op = DeviceOpInstance{};
+    auto invoker   = device_op.MakeInvoker();
+    auto argument =
+        device_op.MakeArgument(a_device_buf.GetDeviceBuffer(),
+                               b_device_buf.GetDeviceBuffer(),
+                               {},
+                               e_device_buf.GetDeviceBuffer(),
+                               {r0_device_buf.GetDeviceBuffer(), r1_device_buf.GetDeviceBuffer()},
+                               M,
+                               N,
+                               K,
+                               StrideA,
+                               StrideB,
+                               {},
+                               StrideE,
+                               a_element_op,
+                               b_element_op,
+                               cde_element_op,
+                               qs_element_op,
+                               rs_element_op);
+
+    if(!device_op.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error("wrong! this device_op instance does not support this problem");
+    }
+
+    // init reducetion buffer to 0
+    r0_device_buf.SetZero();
+    r1_device_buf.SetZero();
+
+    invoker.Run(argument, StreamConfig{nullptr, false});
+
+    bool do_verification = true;
+    bool pass            = true;
+
+    if(do_verification)
+    {
+        auto I0 = ck::Number<0>{};
+        auto I1 = ck::Number<1>{};
+
+        Tensor<EDataType> e_m_n_host(e_m_n.mDesc);
+        Tensor<R0DataType> r0_m_host(r0_m.mDesc);
+        Tensor<R1DataType> r1_m_host(r1_m.mDesc);
+
+        auto ref_gemm    = ReferenceGemmInstance{};
+        auto ref_invoker = ref_gemm.MakeInvoker();
+
+        auto ref_argument = ref_gemm.MakeArgument(
+            a_m_k, b_k_n, e_m_n_host, a_element_op, b_element_op, PassThrough{});
+
+        ref_invoker.Run(ref_argument);
+
+        auto reduce0_op = R0ThreadReduceOp{};
+        auto reduce1_op = R1ThreadReduceOp{};
+
+        for(int m = 0; m < M; ++m)
+        {
+            auto reduce0_acc = reduce0_op.GetIdentityValue<ReduceAccDataType>();
+            auto reduce1_acc = reduce1_op.GetIdentityValue<ReduceAccDataType>();
+
+            for(int n = 0; n < N; ++n)
+            {
+                ReduceAccDataType square_e_val;
+                auto e_val = ck::type_convert<ReduceAccDataType>(e_m_n_host(m, n));
+                qs_element_op[I1](square_e_val, e_val);
+
+                reduce0_op(reduce0_acc, e_val);
+                reduce1_op(reduce1_acc, square_e_val);
+            }
+
+            rs_element_op[I0](reduce0_acc, reduce0_acc);
+            rs_element_op[I1](reduce1_acc, reduce1_acc);
+            r0_m_host(m) = ck::type_convert<R0DataType>(reduce0_acc);
+            r1_m_host(m) = ck::type_convert<R1DataType>(reduce1_acc);
+        }
+
+        e_device_buf.FromDevice(e_m_n.mData.data());
+        r0_device_buf.FromDevice(r0_m.mData.data());
+        r1_device_buf.FromDevice(r1_m.mData.data());
+
+        pass = ck::utils::check_err(
+            e_m_n.mData, e_m_n_host.mData, "Error: Incorrect results c", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(
+            r0_m.mData, r0_m_host.mData, "Error: Incorrect results d0", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(
+            r1_m.mData, r1_m_host.mData, "Error: Incorrect results d1", 1e-2, 1e-2);
+    }
+
+    bool time_kernel = true;
+    if(time_kernel)
+    {
+        float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+        DumpPerf<ADataType, BDataType, EDataType, R0DataType, R1DataType>(ave_time, M, N, K);
+    }
+
+    return pass ? 0 : 1;
+}

example/16_gemm_reduce/CMakeLists.txt

-add_example_executable(example_gemm_reduce_xdl_max_fp16 gemm_reduce_xdl_max_fp16.cpp)
-add_example_executable(example_gemm_reduce_xdl_mean_squaremean_fp16 gemm_reduce_xdl_mean_squaremean_fp16.cpp)

example/16_gemm_reduce/gemm_reduce_xdl_max_fp16.cpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
-
-#include <iostream>
-#include <numeric>
-#include <initializer_list>
-#include <cstdlib>
-
-#include "ck/ck.hpp"
-#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
-#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
-#include "ck/tensor_operation/gpu/device/device_gemm_reduce_xdl_cshuffle.hpp"
-#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
-
-#include "ck/library/utility/device_memory.hpp"
-#include "ck/library/utility/host_tensor.hpp"
-#include "ck/library/utility/host_tensor_generator.hpp"
-#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
-#include "ck/library/utility/check_err.hpp"
-
-template <ck::index_t... Is>
-using S = ck::Sequence<Is...>;
-
-using F16 = ck::half_t;
-using F32 = float;
-using F64 = double;
-
-using Row = ck::tensor_layout::gemm::RowMajor;
-using Col = ck::tensor_layout::gemm::ColumnMajor;
-
-using ADataType         = F16;
-using BDataType         = F16;
-using CDataType         = F16;
-using GemmAccDataType   = F32;
-using ReduceAccDataType = F32;
-using ReduceDataType    = F64;
-using ReducePtrsGlobal  = ck::Tuple<ReduceDataType*>;
-
-using ALayout = ck::tensor_layout::gemm::RowMajor;
-using BLayout = ck::tensor_layout::gemm::ColumnMajor;
-using CLayout = ck::tensor_layout::gemm::RowMajor;
-
-using AElementOp       = ck::tensor_operation::element_wise::PassThrough;
-using BElementOp       = ck::tensor_operation::element_wise::PassThrough;
-using CElementOp       = ck::tensor_operation::element_wise::PassThrough;
-using ReduceOps        = ck::Tuple<ck::reduce::Max>;
-using ReduceElementOps = ck::Tuple<ck::tensor_operation::element_wise::PassThrough>;
-using ReduceGlobalMemOps =
-    ck::InMemoryDataOperationEnumSequence<ck::InMemoryDataOperationEnum::AtomicMax>;
-
-static constexpr auto GemmSpecialization =
-    ck::tensor_operation::device::GemmSpecialization::Default;
-
-// clang-format off
-using DeviceGemmReduceInstance = ck::tensor_operation::device::DeviceGemmReduce_Xdl_CShuffle
-//######| ALayout| BLayout| CLayout|AData| BData| CData|  GemmAcc| CShuffle|         ReduceAcc|       ReduceData|           A|           B|           C|      Reduce|    ReduceInEleOp|   ReduceAccEleOp|             Reduce|               GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer|              CReduce| CReduceThreadLds2VGprCopy| CReduceThreadVgpr2GlobalCopy|
-//######|        |        |        | Type|  Type|  Type| DataType| DataType|          DataType|       Type Tuple| Elementwise| Elementwise| Elementwise|   Operation|                 |                 |         MemoryData|     Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN| MXdlPerWave| NXdlPerWave|            _MBlock_MPerBlock| ScalarPerVector| ThreadClusterLengths|     SrcDstScalarPerVector|        SrcDstScalarPerVector|
-//######|        |        |        |     |      |      |         |         |                  |                 |   Operation|   Operation|   Operation|            |                 |                 |          Operation|                   |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|            _NBlock_NPerBlock|      _NPerBlock| _MPerBlock_NPerBlock|                _NPerBlock|                   _MPerBlock|
-//######|        |        |        |     |      |      |         |         |                  |                 |            |            |            |            |                 |                 |                   |                   |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |                     |                          |                             |
-        <     Row,     Col,     Row,  F16,   F16,   F16,      F32,      F32, ReduceAccDataType, ReducePtrsGlobal,  AElementOp,  BElementOp,  CElementOp,   ReduceOps, ReduceElementOps, ReduceElementOps, ReduceGlobalMemOps, GemmSpecialization,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,             S<64, 4>,                         4,                            1>;
-// clang-format on
-
-using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
-                                                                        BDataType,
-                                                                        CDataType,
-                                                                        GemmAccDataType,
-                                                                        AElementOp,
-                                                                        BElementOp,
-                                                                        CElementOp>;
-
-template <typename ADataType, typename BDataType, typename CDataType, typename ReduceDataType>
-void DumpGemmLayerNormPerf(float gemm_reduce_time, int M, int N, int K)
-{
-    std::size_t gemm_flop     = std::size_t(2) * M * N * K;
-    std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
-                                sizeof(CDataType) * M * N + sizeof(ReduceDataType) * M;
-
-    float tflops          = static_cast<float>(gemm_flop) / 1.E9 / gemm_reduce_time;
-    float gemm_gb_per_sec = gemm_num_byte / 1.E6 / gemm_reduce_time;
-
-    std::cout << "gemm + reduceMax Perf: " << gemm_reduce_time << " ms, " << tflops << " TFlops, "
-              << gemm_gb_per_sec << " GB/s, " << std::endl;
-}
-
-int main(int argc, char* argv[])
-{
-    bool do_verification = true;
-    int init_method      = 1;
-    bool time_kernel     = false;
-
-    // GEMM shape
-    ck::index_t M = 3840;
-    ck::index_t N = 4096;
-    ck::index_t K = 4096;
-
-    ck::index_t StrideA = 4096;
-    ck::index_t StrideB = 4096;
-    ck::index_t StrideC = 4096;
-
-    if(argc == 1)
-    {
-        // do nothing
-    }
-    else if(argc == 4)
-    {
-        do_verification = std::stoi(argv[1]);
-        init_method     = std::stoi(argv[2]);
-        time_kernel     = std::stoi(argv[3]);
-    }
-    else if(argc == 10)
-    {
-        do_verification = std::stoi(argv[1]);
-        init_method     = std::stoi(argv[2]);
-        time_kernel     = std::stoi(argv[3]);
-
-        M = std::stoi(argv[4]);
-        N = std::stoi(argv[5]);
-        K = std::stoi(argv[6]);
-
-        StrideA = std::stoi(argv[7]);
-        StrideB = std::stoi(argv[8]);
-        StrideC = std::stoi(argv[9]);
-    }
-    else
-    {
-        printf("arg1: verification (0=no, 1=yes)\n");
-        printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
-        printf("arg3: run kernel # of times (>1)\n");
-        printf("arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC\n");
-        exit(0);
-    }
-
-    auto f_host_tensor_descriptor =
-        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
-            if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
-            {
-                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
-                                            std::vector<std::size_t>({stride, 1}));
-            }
-            else
-            {
-                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
-                                            std::vector<std::size_t>({1, stride}));
-            }
-        };
-
-    Tensor<ADataType> a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
-    Tensor<BDataType> b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
-
-    Tensor<CDataType> c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> reduce_m_host_result(
-        HostTensorDescriptor(std::vector<std::size_t>({static_cast<std::size_t>(M)})));
-
-    Tensor<CDataType> c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> reduce_m_device_result(
-        HostTensorDescriptor(std::vector<std::size_t>({static_cast<std::size_t>(M)})));
-
-    std::cout << "a_m_k: " << a_m_k.mDesc << std::endl;
-    std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
-    std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl;
-    std::cout << "reduce_m: " << reduce_m_host_result.mDesc << std::endl;
-
-    switch(init_method)
-    {
-    case 0: break;
-    case 1:
-        a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
-        b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
-        break;
-    default:
-        a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
-        b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
-        break;
-    }
-
-    DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
-    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
-    DeviceMem c_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize());
-    DeviceMem reduce_device_buf(sizeof(ReduceDataType) *
-                                reduce_m_device_result.mDesc.GetElementSpaceSize());
-
-    a_device_buf.ToDevice(a_m_k.mData.data());
-    b_device_buf.ToDevice(b_k_n.mData.data());
-
-    auto a_element_op                       = AElementOp{};
-    auto b_element_op                       = BElementOp{};
-    auto c_element_op                       = CElementOp{};
-    auto reduce_element_op                  = ReduceElementOps{}[ck::Number<0>{}];
-    std::array<void*, 3> gemm_element_ops   = {&a_element_op, &b_element_op, &c_element_op};
-    std::array<void*, 1> reduce_element_ops = {&reduce_element_op};
-    std::array<void*, 1> p_reduces          = {reduce_device_buf.GetDeviceBuffer()};
-
-    // do GEMM
-    auto gemm     = DeviceGemmReduceInstance{};
-    auto invoker  = gemm.MakeInvoker();
-    auto argument = gemm.MakeArgument(a_device_buf.GetDeviceBuffer(),
-                                      b_device_buf.GetDeviceBuffer(),
-                                      nullptr,
-                                      {},
-                                      c_device_buf.GetDeviceBuffer(),
-                                      p_reduces,
-                                      M,
-                                      N,
-                                      K,
-                                      StrideA,
-                                      StrideB,
-                                      StrideC,
-                                      {},
-                                      gemm_element_ops,
-                                      {},
-                                      reduce_element_ops,
-                                      reduce_element_ops);
-
-    if(!gemm.IsSupportedArgument(argument))
-    {
-        throw std::runtime_error(
-            "wrong! device_gemm with the specified compilation parameters does "
-            "not support this GEMM problem");
-    }
-
-    // [CAUSION]: launch_and_time_kernel will not initialize D.
-    // If we evaluate kernel multiple time but without initialize D. Verification will fail
-    reduce_device_buf.SetValue(ck::NumericLimits<ReduceDataType>::Lowest());
-    invoker.Run(argument, StreamConfig{nullptr, false});
-
-    bool pass = true;
-
-    if(do_verification)
-    {
-        c_device_buf.FromDevice(c_m_n_device_result.mData.data());
-        reduce_device_buf.FromDevice(reduce_m_device_result.mData.data());
-
-        auto ref_gemm    = ReferenceGemmInstance{};
-        auto ref_invoker = ref_gemm.MakeInvoker();
-
-        auto ref_argument = ref_gemm.MakeArgument(
-            a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op);
-
-        ref_invoker.Run(ref_argument);
-
-        auto reduce_op = ReduceOps{}[ck::Number<0>{}];
-
-        for(int m = 0; m < M; ++m)
-        {
-            ReduceAccDataType reduce_acc = reduce_op.GetIdentityValue<ReduceAccDataType>();
-
-            for(int n = 0; n < N; ++n)
-            {
-                ReduceAccDataType curr_val =
-                    ck::type_convert<ReduceAccDataType>(c_m_n_host_result(m, n));
-                reduce_op(reduce_acc, curr_val);
-            };
-
-            reduce_m_host_result(m) = reduce_acc;
-        }
-
-        pass = ck::utils::check_err(c_m_n_device_result.mData,
-                                    c_m_n_host_result.mData,
-                                    "Error: Incorrect results c") &&
-               ck::utils::check_err(reduce_m_device_result.mData,
-                                    reduce_m_host_result.mData,
-                                    "Error: Incorrect results d",
-                                    1e-3,
-                                    1e-3);
-    }
-
-    if(time_kernel)
-    {
-        float gemm_reduceMax_ave_time = invoker.Run(argument, StreamConfig{nullptr, true});
-
-        DumpGemmLayerNormPerf<ADataType, BDataType, CDataType, ReduceDataType>(
-            gemm_reduceMax_ave_time, M, N, K);
-    }
-
-    return pass ? 0 : 1;
-}

example/16_gemm_reduce/gemm_reduce_xdl_mean_squaremean_fp16.cpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
-
-#include <iostream>
-#include <numeric>
-#include <initializer_list>
-#include <cstdlib>
-
-#include "ck/ck.hpp"
-#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
-#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
-#include "ck/tensor_operation/gpu/device/device_gemm_reduce_xdl_cshuffle.hpp"
-#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
-#include "ck/utility/reduction_operator.hpp"
-
-#include "ck/library/utility/check_err.hpp"
-#include "ck/library/utility/device_memory.hpp"
-#include "ck/library/utility/host_tensor.hpp"
-#include "ck/library/utility/host_tensor_generator.hpp"
-#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
-
-template <ck::index_t... Is>
-using S = ck::Sequence<Is...>;
-
-using F16 = ck::half_t;
-using F32 = float;
-
-using Row = ck::tensor_layout::gemm::RowMajor;
-using Col = ck::tensor_layout::gemm::ColumnMajor;
-
-using ADataType         = F16;
-using BDataType         = F16;
-using CDataType         = F16;
-using GemmAccDataType   = F32;
-using ReduceAccDataType = F32;
-using ReduceDataType    = F32;
-using ReducePtrsGlobal  = ck::Tuple<ReduceDataType*, ReduceDataType*>;
-
-using ALayout = ck::tensor_layout::gemm::RowMajor;
-using BLayout = ck::tensor_layout::gemm::ColumnMajor;
-using CLayout = ck::tensor_layout::gemm::RowMajor;
-
-using AElementOp = ck::tensor_operation::element_wise::PassThrough;
-using BElementOp = ck::tensor_operation::element_wise::PassThrough;
-using CElementOp = ck::tensor_operation::element_wise::PassThrough;
-using ReduceOp0  = ck::reduce::Add;
-using ReduceOp1  = ck::reduce::Add;
-using ReduceOps  = ck::Tuple<ReduceOp0, ReduceOp1>;
-
-using UnaryIdenticElementOp = ck::tensor_operation::element_wise::PassThrough;
-using UnaryDivElementOp     = ck::tensor_operation::element_wise::UnaryDivide;
-using UnarySquareElementOp  = ck::tensor_operation::element_wise::UnarySquare;
-using ReduceInElementOps    = ck::Tuple<UnaryIdenticElementOp, UnarySquareElementOp>;
-using ReduceOutElementOps   = ck::Tuple<UnaryDivElementOp, UnaryDivElementOp>;
-
-using ReduceGlobalMemOps =
-    ck::InMemoryDataOperationEnumSequence<ck::InMemoryDataOperationEnum::AtomicAdd,
-                                          ck::InMemoryDataOperationEnum::AtomicAdd>;
-
-static constexpr auto GemmSpecialization =
-    ck::tensor_operation::device::GemmSpecialization::Default;
-
-// clang-format off
-using DeviceGemmReduceInstance = ck::tensor_operation::device::DeviceGemmReduce_Xdl_CShuffle
-//######| ALayout| BLayout| CLayout|AData| BData| CData|  GemmAcc| CShuffle| ReduceAcc|      ReduceDData|           A|           B|           C|      Reduce|      ReduceInEleOp|      ReduceOutEleOp|              Reduce|               GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer|              CReduce| CReduceThreadLds2VGprCopy| CReduceThreadVgpr2GlobalCopy|
-//######|        |        |        | Type|  Type|  Type| DataType| DataType|  DataType|       Type Tuple| Elementwise| Elementwise| Elementwise|   Operation|                   |                    |          MemoryData|     Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN| MXdlPerWave| NXdlPerWave|            _MBlock_MPerBlock| ScalarPerVector| ThreadClusterLengths|     SrcDstScalarPerVector|        SrcDstScalarPerVector|
-//######|        |        |        |     |      |      |         |         |          |                 |   Operation|   Operation|   Operation|            |                   |                    |           Operation|                   |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|            _NBlock_NPerBlock|      _NPerBlock| _MPerBlock_NPerBlock|                _NPerBlock|                   _MPerBlock|
-//######|        |        |        |     |      |      |         |         |          |                 |            |            |            |            |                   |                    |                    |                   |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |                     |                          |                             |
-        <     Row,     Col,     Row,  F16,   F16,   F16,      F32,      F32,       F32, ReducePtrsGlobal,  AElementOp,  BElementOp,  CElementOp,   ReduceOps, ReduceInElementOps, ReduceOutElementOps,  ReduceGlobalMemOps, GemmSpecialization,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,             S<64, 4>,                         4,                            1>;
-// clang-format on
-
-using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
-                                                                        BDataType,
-                                                                        CDataType,
-                                                                        GemmAccDataType,
-                                                                        AElementOp,
-                                                                        BElementOp,
-                                                                        CElementOp>;
-
-template <typename ADataType, typename BDataType, typename CDataType, typename ReduceDataType>
-void DumpGemmLayerNormPerf(float gemm_reduce_time, int M, int N, int K)
-{
-    std::size_t gemm_flop     = std::size_t(2) * M * N * K;
-    std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
-                                sizeof(CDataType) * M * N + sizeof(ReduceDataType) * M +
-                                sizeof(ReduceDataType) * M;
-
-    float tflops          = static_cast<float>(gemm_flop) / 1.E9 / gemm_reduce_time;
-    float gemm_gb_per_sec = gemm_num_byte / 1.E6 / gemm_reduce_time;
-
-    std::cout << "gemm + reduce_mean + reduce_mean_square Perf: " << gemm_reduce_time << " ms, "
-              << tflops << " TFlops, " << gemm_gb_per_sec << " GB/s, " << std::endl;
-}
-
-int main(int argc, char* argv[])
-{
-    bool do_verification = true;
-    int init_method      = 1;
-    bool time_kernel     = false;
-
-    // GEMM shape
-    ck::index_t M = 3840;
-    ck::index_t N = 4096;
-    ck::index_t K = 4096;
-
-    ck::index_t StrideA = 4096;
-    ck::index_t StrideB = 4096;
-    ck::index_t StrideC = 4096;
-
-    if(argc == 1)
-    {
-        // do nothing
-    }
-    else if(argc == 4)
-    {
-        do_verification = std::stoi(argv[1]);
-        init_method     = std::stoi(argv[2]);
-        time_kernel     = std::stoi(argv[3]);
-    }
-    else if(argc == 10)
-    {
-        do_verification = std::stoi(argv[1]);
-        init_method     = std::stoi(argv[2]);
-        time_kernel     = std::stoi(argv[3]);
-
-        M = std::stoi(argv[4]);
-        N = std::stoi(argv[5]);
-        K = std::stoi(argv[6]);
-
-        StrideA = std::stoi(argv[7]);
-        StrideB = std::stoi(argv[8]);
-        StrideC = std::stoi(argv[9]);
-    }
-    else
-    {
-        printf("arg1: verification (0=no, 1=yes)\n");
-        printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
-        printf("arg3: time kernel (0=n0, 1=yes)\n");
-        printf("arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC\n");
-        exit(0);
-    }
-
-    auto f_host_tensor_descriptor =
-        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
-            if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
-            {
-                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
-                                            std::vector<std::size_t>({stride, 1}));
-            }
-            else
-            {
-                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
-                                            std::vector<std::size_t>({1, stride}));
-            }
-        };
-
-    Tensor<ADataType> a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
-    Tensor<BDataType> b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
-
-    Tensor<CDataType> c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> reduce0_m_host_result(
-        HostTensorDescriptor(std::vector<std::size_t>({static_cast<std::size_t>(M)})));
-    Tensor<ReduceDataType> reduce1_m_host_result(
-        HostTensorDescriptor(std::vector<std::size_t>({static_cast<std::size_t>(M)})));
-
-    Tensor<CDataType> c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> reduce0_m_device_result(
-        HostTensorDescriptor(std::vector<std::size_t>({static_cast<std::size_t>(M)})));
-    Tensor<ReduceDataType> reduce1_m_device_result(
-        HostTensorDescriptor(std::vector<std::size_t>({static_cast<std::size_t>(M)})));
-
-    std::cout << "a_m_k: " << a_m_k.mDesc << std::endl;
-    std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
-    std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl;
-    std::cout << "reduce0_m: " << reduce0_m_host_result.mDesc << std::endl;
-    std::cout << "reduce1_m: " << reduce1_m_host_result.mDesc << std::endl;
-
-    switch(init_method)
-    {
-    case 0: break;
-    case 1:
-        a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
-        b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
-        break;
-    default:
-        a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
-        b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
-        break;
-    }
-
-    DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
-    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
-    DeviceMem c_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize());
-    DeviceMem reduce0_device_buf(sizeof(ReduceDataType) *
-                                 reduce0_m_device_result.mDesc.GetElementSpaceSize());
-    DeviceMem reduce1_device_buf(sizeof(ReduceDataType) *
-                                 reduce1_m_device_result.mDesc.GetElementSpaceSize());
-
-    a_device_buf.ToDevice(a_m_k.mData.data());
-    b_device_buf.ToDevice(b_k_n.mData.data());
-
-    auto a_element_op                     = AElementOp{};
-    auto b_element_op                     = BElementOp{};
-    auto c_element_op                     = CElementOp{};
-    std::array<void*, 3> gemm_element_ops = {&a_element_op, &b_element_op, &c_element_op};
-
-    auto passthrough                            = UnaryIdenticElementOp{};
-    auto square                                 = UnarySquareElementOp{};
-    auto div                                    = UnaryDivElementOp{N};
-    std::array<void*, 2> reduce_in_element_ops  = {&passthrough, &square};
-    std::array<void*, 2> reduce_out_element_ops = {&div, &div};
-
-    std::array<void*, 2> p_reduces = {reduce0_device_buf.GetDeviceBuffer(),
-                                      reduce1_device_buf.GetDeviceBuffer()};
-
-    // do GEMM
-    auto gemm     = DeviceGemmReduceInstance{};
-    auto invoker  = gemm.MakeInvoker();
-    auto argument = gemm.MakeArgument(a_device_buf.GetDeviceBuffer(),
-                                      b_device_buf.GetDeviceBuffer(),
-                                      nullptr,
-                                      {},
-                                      c_device_buf.GetDeviceBuffer(),
-                                      p_reduces,
-                                      M,
-                                      N,
-                                      K,
-                                      StrideA,
-                                      StrideB,
-                                      StrideC,
-                                      {},
-                                      gemm_element_ops,
-                                      {},
-                                      reduce_in_element_ops,
-                                      reduce_out_element_ops);
-
-    if(!gemm.IsSupportedArgument(argument))
-    {
-        throw std::runtime_error(
-            "wrong! device_gemm with the specified compilation parameters does "
-            "not support this GEMM problem");
-    }
-
-    // init reducetion buffer to 0
-    reduce0_device_buf.SetZero();
-    reduce1_device_buf.SetZero();
-
-    // if time_kernel == true, kernel will run multiple times. This kernel use atomic-add so result
-    // will not be correct. need to set time_kernel = false for correctness test
-    invoker.Run(argument, StreamConfig{nullptr, false});
-    bool pass = true;
-
-    if(do_verification)
-    {
-        c_device_buf.FromDevice(c_m_n_device_result.mData.data());
-        reduce0_device_buf.FromDevice(reduce0_m_device_result.mData.data());
-        reduce1_device_buf.FromDevice(reduce1_m_device_result.mData.data());
-
-        auto ref_gemm    = ReferenceGemmInstance{};
-        auto ref_invoker = ref_gemm.MakeInvoker();
-
-        auto ref_argument = ref_gemm.MakeArgument(
-            a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op);
-
-        ref_invoker.Run(ref_argument);
-
-        auto reduce0_op = ReduceOp0{};
-        auto reduce1_op = ReduceOp1{};
-
-        for(int m = 0; m < M; ++m)
-        {
-            auto reduce0_acc = reduce0_op.GetIdentityValue<ReduceAccDataType>();
-            auto reduce1_acc = reduce1_op.GetIdentityValue<ReduceAccDataType>();
-
-            for(int n = 0; n < N; ++n)
-            {
-                auto c_val = ck::type_convert<ReduceAccDataType>(c_m_n_host_result(m, n));
-                ReduceAccDataType square_c_val;
-                square(square_c_val, c_val);
-
-                reduce0_op(reduce0_acc, c_val);
-                reduce1_op(reduce1_acc, square_c_val);
-            }
-
-            div(reduce0_acc, reduce0_acc);
-            div(reduce1_acc, reduce1_acc);
-            reduce0_m_host_result(m) = ck::type_convert<ReduceDataType>(reduce0_acc);
-            reduce1_m_host_result(m) = ck::type_convert<ReduceDataType>(reduce1_acc);
-        }
-
-        pass = ck::utils::check_err(c_m_n_device_result.mData,
-                                    c_m_n_host_result.mData,
-                                    "Error: Incorrect results c") &&
-               ck::utils::check_err(reduce0_m_device_result.mData,
-                                    reduce0_m_host_result.mData,
-                                    "Error: Incorrect results d0",
-                                    1e-4,
-                                    1e-5) &&
-               ck::utils::check_err(reduce1_m_device_result.mData,
-                                    reduce1_m_host_result.mData,
-                                    "Error: Incorrect results d1",
-                                    1e-3,
-                                    1e-5);
-    }
-
-    if(time_kernel)
-    {
-        float ave_time = invoker.Run(argument, StreamConfig{nullptr, true});
-
-        DumpGemmLayerNormPerf<ADataType, BDataType, CDataType, ReduceDataType>(ave_time, M, N, K);
-    }
-
-    return pass ? 0 : 1;
-}

example/21_gemm_layernorm/gemm_bias_relu_add_layernorm_xdl_fp16.cpp

@@ -9,7 +9,7 @@
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
-#include "ck/tensor_operation/gpu/device/device_gemm_bias_add_reduce_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
 #include "ck/tensor_operation/gpu/device/device_5ary_elementwise.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
@@ -28,57 +28,64 @@ using F32 = float;
 using Row = ck::tensor_layout::gemm::RowMajor;
 using Col = ck::tensor_layout::gemm::ColumnMajor;
 
+// DataType
 using ADataType                = F16;
 using BDataType                = F16;
-using CDataType                = F16;
-using BiasDataType             = F32;
-using D0DataType               = F16;
 using GemmAccDataType          = F32;
+using CShuffleDataType         = F32;
+using D0DataType               = F16;
+using D1DataType               = F16;
+using DsDataType               = ck::Tuple<D0DataType, D1DataType>;
+using EDataType                = F16;
 using ReduceAccDataType        = F32;
-using ReduceDataType           = F32;
-using ReducePtrsGlobal         = ck::Tuple<ReduceDataType*, ReduceDataType*>;
+using R0DataType               = F32;
+using R1DataType               = F32;
+using RsDataType               = ck::Tuple<R0DataType, R1DataType>;
 using GammaDataType            = F16;
 using BetaDataType             = F16;
 using LayerNormOutDataType     = F16;
 using NormalizeComputeDataType = F32;
 
-using ALayout = ck::tensor_layout::gemm::RowMajor;
-using BLayout = ck::tensor_layout::gemm::ColumnMajor;
-using CLayout = ck::tensor_layout::gemm::RowMajor;
-
-using PassThrough = ck::tensor_operation::element_wise::PassThrough;
-using AElementOp  = PassThrough;
-using BElementOp  = PassThrough;
-using CElementOp  = ck::tensor_operation::element_wise::Relu;
-using D0ElementOp = PassThrough;
-using ReduceSumOp = ck::reduce::Add;
-using ReduceOps   = ck::Tuple<ReduceSumOp, ReduceSumOp>;
-
-using UnaryIdenticElementOp = ck::tensor_operation::element_wise::PassThrough;
-using UnaryDivElementOp     = ck::tensor_operation::element_wise::UnaryDivide;
-using UnarySquareElementOp  = ck::tensor_operation::element_wise::UnarySquare;
-using ReduceInElementOps    = ck::Tuple<UnaryIdenticElementOp, UnarySquareElementOp>;
-using ReduceOutElementOps   = ck::Tuple<UnaryDivElementOp, UnaryDivElementOp>;
-
-using ReduceGlobalMemOps =
-    ck::InMemoryDataOperationEnumSequence<ck::InMemoryDataOperationEnum::AtomicAdd,
-                                          ck::InMemoryDataOperationEnum::AtomicAdd>;
-
-static constexpr auto GemmSpecialization =
-    ck::tensor_operation::device::GemmSpecialization::Default;
+// Layout
+using ALayout  = Row;
+using BLayout  = Col;
+using D1Layout = Row;
+using ELayout  = D1Layout;
+
+// Elementwise op
+using PassThrough  = ck::tensor_operation::element_wise::PassThrough;
+using AddReluAdd   = ck::tensor_operation::element_wise::AddReluAdd;
+using Square       = ck::tensor_operation::element_wise::UnarySquare;
+using Div          = ck::tensor_operation::element_wise::UnaryDivide;
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = AddReluAdd;
+using QsElementOp  = ck::Tuple<PassThrough, Square>;
+using RsElementOp  = ck::Tuple<Div, Div>;
+
+// ReduceOp
+using R0ThreadReduceOp = ck::reduce::Add;
+using R1ThreadReduceOp = ck::reduce::Add;
+using RsThreadReduceOp = ck::Tuple<R0ThreadReduceOp, R1ThreadReduceOp>;
+
+static constexpr auto R0GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+static constexpr auto R1GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+using RsGlobalReduceOp = ck::InMemoryDataOperationEnumSequence<R0GlobalReduceOp, R1GlobalReduceOp>;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
 
 // clang-format off
-using DeviceGemmBiasAddReduceInstance = ck::tensor_operation::device::DeviceGemmBiasAddReduce_Xdl_CShuffle
-//######| ALayout| BLayout| CLayout|AData| BData| CData|C0Data|C1Data|  GemmAcc| CShuffle| ReduceAcc|       ReduceData|           A|           B|           C|          C1|    Reduce|     ReduceInEleOp|      ReduceAccEleOp|              Reduce|               GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer|              CReduce| CReduceThreadLds2VGprCopy| CReduceThreadVgpr2GlobalCopy|
-//######|        |        |        | Type|  Type|  Type|  Type|  Type| DataType| DataType|  DataType|       Type Tuple| Elementwise| Elementwise| Elementwise| Elementwise| Operation|                  |                    |          MemoryData|     Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN| MXdlPerWave| NXdlPerWave|            _MBlock_MPerBlock| ScalarPerVector| ThreadClusterLengths|     SrcDstScalarPerVector|        SrcDstScalarPerVector|
-//######|        |        |        |     |      |      |      |      |         |         |          |                 |   Operation|   Operation|   Operation|   Operation|          |                  |                    |           Operation|                   |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|            _NBlock_NPerBlock|      _NPerBlock| _MPerBlock_NPerBlock|                _NPerBlock|                   _MPerBlock|
-//######|        |        |        |     |      |      |      |      |         |         |          |                 |            |            |            |            |          |                  |                    |                    |                   |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |                     |                          |                             |
-        <     Row,     Col,     Row,  F16,   F16,   F16,   F32,   F16,      F32,      F32,       F32, ReducePtrsGlobal,  AElementOp,  BElementOp,  CElementOp, D0ElementOp, ReduceOps,ReduceInElementOps, ReduceOutElementOps,  ReduceGlobalMemOps, GemmSpecialization,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,             S<64, 4>,                         4,                            1>;
+using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleDMultipleR_Xdl_CShuffle
+//######| ALayout| BLayout| ELayout|     AData|     BData|     GemmAccData|         CShuffle|     DsData|     EData|     ReduceAccData|     RsData|           A|           B|          CDE|          Qs|          Rs|           Thread|           Global|           GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle|    CDRThreadTransfer|                  CDE|    RThreadTransfer|
+//######|        |        |        |      Type|      Type|            Type|         DataType|       Type|      Type|              Type|       Type| Elementwise| Elementwise|  Elementwise| Elementwise| Elementwise|           Reduce|           Reduce| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|       ClusterLengths| ReduceThreadTransfer| DstScalarPerVector|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |   Operation|   Operation|    Operation|   Operation|   Operation|        Operation|        Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle| _MPerBlock_NPerBlock|      ScalarPerVector|         _MPerBlock|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |            |            |             |            |            |                 |                 |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                     |           _NPerBlock|                   |
+        < ALayout, BLayout, ELayout, ADataType, BDataType, GemmAccDataType, CShuffleDataType, DsDataType, EDataType, ReduceAccDataType, RsDataType,  AElementOp,  BElementOp, CDEElementOp, QsElementOp, RsElementOp, RsThreadReduceOp, RsGlobalReduceOp,    GemmDefault,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,             S<64, 4>,                    4,                  1>;
 // clang-format on
 
 using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
                                                                         BDataType,
-                                                                        CDataType,
+                                                                        EDataType,
                                                                         GemmAccDataType,
                                                                         AElementOp,
                                                                         BElementOp,
@@ -88,9 +95,9 @@ using NormalizeFunctor = ck::tensor_operation::element_wise::Normalize;
 
 // A:x, B:E[x], C:E[x^2], D:Gamma, E:Beta , F:y
 using DeviceNormalizeInstance =
-    ck::tensor_operation::device::Device5AryElementwise<CDataType,
-                                                        ReduceDataType,
-                                                        ReduceDataType,
+    ck::tensor_operation::device::Device5AryElementwise<EDataType,
+                                                        R0DataType,
+                                                        R1DataType,
                                                         GammaDataType,
                                                         BetaDataType,
                                                         LayerNormOutDataType,
@@ -124,41 +131,31 @@ auto f_host_tensor_descriptor2d =
         }
     };
 
-template <typename CDataType,
-          typename ReduceDataType,
-          typename AccDataType,
-          typename BiasDataType,
-          typename D0DataType,
-          typename A_functor,
-          typename B_functor,
-          typename C_functor,
-          typename C1_functor>
 void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
                          const Tensor<ADataType>& a_m_k,
-                         const Tensor<ADataType>& b_k_n,
-                         const Tensor<BiasDataType>& bias_n,
-                         const Tensor<D0DataType>& c1_m_n,
+                         const Tensor<BDataType>& b_k_n,
+                         const Tensor<D0DataType>& bias_n,
+                         const Tensor<D1DataType>& d1_m_n,
                          const Tensor<GammaDataType>& gamma_n,
-                         const Tensor<GammaDataType>& beta_n,
-                         A_functor a_element_op,
-                         B_functor b_element_op,
-                         C_functor c_element_op,
-                         C1_functor c1_element_op,
+                         const Tensor<BetaDataType>& beta_n,
+                         AElementOp a_element_op,
+                         BElementOp b_element_op,
+                         CDEElementOp cde_element_op,
                          int M,
                          int N)
 {
 
-    int StrideC = N;
-    Tensor<CDataType> c_m_n(f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> mean_m(f_host_tensor_descriptor1d(M, 1));
-    Tensor<ReduceDataType> meanSquare_m(f_host_tensor_descriptor1d(M, 1));
-    auto averageOpInst = UnaryDivElementOp{N};
+    int StrideE = N;
+    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> mean_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<R1DataType> meanSquare_m(f_host_tensor_descriptor1d(M, 1));
+    auto averageOpInst = Div{N};
 
     auto ref_gemm    = ReferenceGemmInstance{};
     auto ref_invoker = ref_gemm.MakeInvoker();
 
     auto ref_argument =
-        ref_gemm.MakeArgument(a_m_k, b_k_n, c_m_n, a_element_op, b_element_op, PassThrough{});
+        ref_gemm.MakeArgument(a_m_k, b_k_n, e_m_n, a_element_op, b_element_op, PassThrough{});
 
     ref_invoker.Run(ref_argument);
 
@@ -166,38 +163,32 @@ void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
     for(int m = 0; m < M; ++m)
         for(int n = 0; n < N; ++n)
         {
-            AccDataType acc = ck::type_convert<AccDataType>(c_m_n(m, n)) +
-                              ck::type_convert<AccDataType>(bias_n(n));
-
-            AccDataType c1 = ck::type_convert<AccDataType>(c1_m_n(m, n));
-
-            c_element_op(acc, acc);
-            c1_element_op(c1, c1);
-            acc += c1;
-            c_m_n(m, n) = ck::type_convert<CDataType>(acc);
+            auto acc = ck::type_convert<GemmAccDataType>(e_m_n(m, n));
+            cde_element_op(e_m_n(m, n), acc, bias_n(n), d1_m_n(m, n));
         }
 
     // reduce_mean and reduce_square_mean
-    auto reduceSumOpInst = ReduceSumOp{};
+    auto r0Op = R0ThreadReduceOp{};
+    auto r1Op = R1ThreadReduceOp{};
     for(int m = 0; m < M; ++m)
     {
-        auto mean_acc        = reduceSumOpInst.GetIdentityValue<AccDataType>();
-        auto square_mean_acc = reduceSumOpInst.GetIdentityValue<AccDataType>();
+        auto mean_acc        = r0Op.GetIdentityValue<ReduceAccDataType>();
+        auto mean_square_acc = r1Op.GetIdentityValue<ReduceAccDataType>();
 
         for(int n = 0; n < N; ++n)
         {
-            AccDataType c_val        = ck::type_convert<AccDataType>(c_m_n(m, n));
-            AccDataType square_c_val = 0;
-            UnarySquareElementOp{}(square_c_val, c_val);
+            auto e_val                     = ck::type_convert<ReduceAccDataType>(e_m_n(m, n));
+            ReduceAccDataType square_e_val = 0;
+            Square{}(square_e_val, e_val);
 
-            reduceSumOpInst(mean_acc, c_val);
-            reduceSumOpInst(square_mean_acc, square_c_val);
+            r0Op(mean_acc, e_val);
+            r1Op(mean_square_acc, square_e_val);
         }
 
         averageOpInst(mean_acc, mean_acc);
-        averageOpInst(square_mean_acc, square_mean_acc);
-        mean_m(m)       = ck::type_convert<ReduceDataType>(mean_acc);
-        meanSquare_m(m) = ck::type_convert<ReduceDataType>(square_mean_acc);
+        averageOpInst(mean_square_acc, mean_square_acc);
+        mean_m(m)       = ck::type_convert<R0DataType>(mean_acc);
+        meanSquare_m(m) = ck::type_convert<R1DataType>(mean_square_acc);
     }
 
     // LayerNorm
@@ -206,24 +197,25 @@ void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
     {
         for(int n = 0; n < N; ++n)
         {
-            AccDataType out_acc = 0;
+            NormalizeComputeDataType out_acc = 0;
             layerNormInst(out_acc,
-                          ck::type_convert<AccDataType>(c_m_n(m, n)),
-                          ck::type_convert<AccDataType>(mean_m(m)),
-                          ck::type_convert<AccDataType>(meanSquare_m(m)),
-                          ck::type_convert<AccDataType>(gamma_n(n)),
-                          ck::type_convert<AccDataType>(beta_n(n)));
-            out_m_n(m, n) = ck::type_convert<ReduceDataType>(out_acc);
+                          ck::type_convert<NormalizeComputeDataType>(e_m_n(m, n)),
+                          ck::type_convert<NormalizeComputeDataType>(mean_m(m)),
+                          ck::type_convert<NormalizeComputeDataType>(meanSquare_m(m)),
+                          ck::type_convert<NormalizeComputeDataType>(gamma_n(n)),
+                          ck::type_convert<NormalizeComputeDataType>(beta_n(n)));
+            out_m_n(m, n) = ck::type_convert<LayerNormOutDataType>(out_acc);
         }
     }
 }
 
 template <typename ADataType,
           typename BDataType,
-          typename CDataType,
-          typename BiasDataType,
+          typename EDataType,
           typename D0DataType,
-          typename ReduceDataType,
+          typename D1DataType,
+          typename R0DataType,
+          typename R1DataType,
           typename GammaDataType,
           typename BetaDataType,
           typename NormalizeDataType>
@@ -231,12 +223,12 @@ void DumpGemmLayerNormPerf(float gemm_reduce_time, float normalize_time, int M,
 {
     std::size_t gemm_flop     = std::size_t(2) * M * N * K + std::size_t(2) * M * N;
     std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
-                                sizeof(CDataType) * M * N + sizeof(BiasDataType) * M * N +
-                                sizeof(D0DataType) * M * N + sizeof(ReduceDataType) * M +
-                                sizeof(ReduceDataType) * M;
+                                sizeof(EDataType) * M * N + sizeof(D0DataType) * M * N +
+                                sizeof(D0DataType) * M * N + sizeof(R0DataType) * M +
+                                sizeof(R1DataType) * M;
 
-    std::size_t normalize_num_byte = sizeof(CDataType) * M * N + sizeof(ReduceDataType) * M +
-                                     sizeof(ReduceDataType) * M + sizeof(GammaDataType) * N +
+    std::size_t normalize_num_byte = sizeof(EDataType) * M * N + sizeof(R0DataType) * M +
+                                     sizeof(R1DataType) * M + sizeof(GammaDataType) * N +
                                      sizeof(BetaDataType) * N + sizeof(NormalizeDataType) * M * N;
 
     float tflops               = static_cast<float>(gemm_flop) / 1.E9 / gemm_reduce_time;
@@ -259,37 +251,37 @@ int main()
 
     ck::index_t StrideA  = 1024;
     ck::index_t StrideB  = 1024;
-    ck::index_t StrideC  = 1024;
-    ck::index_t StrideD0 = 1024;
+    ck::index_t StrideD0 = 0;
+    ck::index_t StrideD1 = 1024;
+    ck::index_t StrideE  = 1024;
 
     Tensor<ADataType> a_m_k(f_host_tensor_descriptor2d(M, K, StrideA, ALayout{}));
     Tensor<BDataType> b_k_n(f_host_tensor_descriptor2d(K, N, StrideB, BLayout{}));
-    Tensor<CDataType> c_m_n(f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
-    Tensor<BiasDataType> bias_n(f_host_tensor_descriptor1d(N, 1));
-    Tensor<D0DataType> c1_m_n(f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> reduceMean_m(f_host_tensor_descriptor1d(M, 1));
-    Tensor<ReduceDataType> reduceMeanSquare_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<D0DataType> bias_n(f_host_tensor_descriptor1d(N, 1));
+    Tensor<D1DataType> d1_m_n(f_host_tensor_descriptor2d(M, N, StrideD1, ELayout{}));
+    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> r0_Mean_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<R1DataType> r1_MeanSquare_m(f_host_tensor_descriptor1d(M, 1));
     Tensor<GammaDataType> gamma_n(f_host_tensor_descriptor1d(N, 1));
     Tensor<BetaDataType> beta_n(f_host_tensor_descriptor1d(N, 1));
     Tensor<LayerNormOutDataType> layerNorm_m_n(
-        f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
+        f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
 
     a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{-1, 1});
     b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-1, 1});
-    bias_n.GenerateTensorValue(GeneratorTensor_3<BiasDataType>{-1, 1});
-    c1_m_n.GenerateTensorValue(GeneratorTensor_3<D0DataType>{-5, 5});
+    bias_n.GenerateTensorValue(GeneratorTensor_3<D0DataType>{-1, 1});
+    d1_m_n.GenerateTensorValue(GeneratorTensor_3<D1DataType>{-5, 5});
     gamma_n.GenerateTensorValue(GeneratorTensor_3<GammaDataType>{-1, 1});
     beta_n.GenerateTensorValue(GeneratorTensor_3<BetaDataType>{-1, 1});
 
     DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
     DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
-    DeviceMem c_device_buf(sizeof(CDataType) * c_m_n.mDesc.GetElementSpaceSize());
-    DeviceMem bias_device_buf(sizeof(BiasDataType) * bias_n.mDesc.GetElementSpaceSize());
-    DeviceMem d0_device_buf(sizeof(D0DataType) * c1_m_n.mDesc.GetElementSpaceSize());
-    DeviceMem reduceMean_device_buf(sizeof(ReduceDataType) *
-                                    reduceMean_m.mDesc.GetElementSpaceSize());
-    DeviceMem reduceMeanSquare_device_buf(sizeof(ReduceDataType) *
-                                          reduceMeanSquare_m.mDesc.GetElementSpaceSize());
+    DeviceMem bias_device_buf(sizeof(D0DataType) * bias_n.mDesc.GetElementSpaceSize());
+    DeviceMem d1_device_buf(sizeof(D1DataType) * d1_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem r0_Mean_device_buf(sizeof(R0DataType) * r0_Mean_m.mDesc.GetElementSpaceSize());
+    DeviceMem r1_MeanSquare_device_buf(sizeof(R1DataType) *
+                                       r1_MeanSquare_m.mDesc.GetElementSpaceSize());
     DeviceMem gamma_device_buf(sizeof(GammaDataType) * gamma_n.mDesc.GetElementSpaceSize());
     DeviceMem beta_device_buf(sizeof(BetaDataType) * beta_n.mDesc.GetElementSpaceSize());
     DeviceMem layerNorm_device_buf(sizeof(LayerNormOutDataType) *
@@ -298,60 +290,51 @@ int main()
     a_device_buf.ToDevice(a_m_k.mData.data());
     b_device_buf.ToDevice(b_k_n.mData.data());
     bias_device_buf.ToDevice(bias_n.mData.data());
-    d0_device_buf.ToDevice(c1_m_n.mData.data());
+    d1_device_buf.ToDevice(d1_m_n.mData.data());
     gamma_device_buf.ToDevice(gamma_n.mData.data());
     beta_device_buf.ToDevice(beta_n.mData.data());
 
-    auto a_element_op                     = AElementOp{};
-    auto b_element_op                     = BElementOp{};
-    auto c_element_op                     = CElementOp{};
-    auto d_element_op                     = D0ElementOp{};
-    std::array<void*, 3> gemm_element_ops = {&a_element_op, &b_element_op, &c_element_op};
-
-    auto passthrough                            = UnaryIdenticElementOp{};
-    auto square                                 = UnarySquareElementOp{};
-    auto div                                    = UnaryDivElementOp{N};
-    std::array<void*, 2> reduce_in_element_ops  = {&passthrough, &square};
-    std::array<void*, 2> reduce_out_element_ops = {&div, &div};
-
-    std::array<void*, 2> p_reduces = {reduceMean_device_buf.GetDeviceBuffer(),
-                                      reduceMeanSquare_device_buf.GetDeviceBuffer()};
+    auto a_element_op   = AElementOp{};
+    auto b_element_op   = BElementOp{};
+    auto cde_element_op = CDEElementOp{};
+    auto qs_element_op  = QsElementOp{};
+    auto rs_element_op  = RsElementOp{N, N};
 
-    // Prepare GEMM, reduce_mean, reduce_mean_square
-    auto gemmReduce          = DeviceGemmBiasAddReduceInstance{};
+    // Prepare GEMM, mean, mean_square
+    auto gemmReduce          = DeviceOpInstance{};
     auto gemmReduce_invoker  = gemmReduce.MakeInvoker();
-    auto gemmReduce_argument = gemmReduce.MakeArgument(a_device_buf.GetDeviceBuffer(),
-                                                       b_device_buf.GetDeviceBuffer(),
-                                                       bias_device_buf.GetDeviceBuffer(),
-                                                       {d0_device_buf.GetDeviceBuffer()},
-                                                       c_device_buf.GetDeviceBuffer(),
-                                                       p_reduces,
-                                                       M,
-                                                       N,
-                                                       K,
-                                                       StrideA,
-                                                       StrideB,
-                                                       StrideC,
-                                                       {StrideD0},
-                                                       gemm_element_ops,
-                                                       {&d_element_op},
-                                                       reduce_in_element_ops,
-                                                       reduce_out_element_ops);
+    auto gemmReduce_argument = gemmReduce.MakeArgument(
+        a_device_buf.GetDeviceBuffer(),
+        b_device_buf.GetDeviceBuffer(),
+        {bias_device_buf.GetDeviceBuffer(), d1_device_buf.GetDeviceBuffer()},
+        e_device_buf.GetDeviceBuffer(),
+        {r0_Mean_device_buf.GetDeviceBuffer(), r1_MeanSquare_device_buf.GetDeviceBuffer()},
+        M,
+        N,
+        K,
+        StrideA,
+        StrideB,
+        {StrideD0, StrideD1},
+        StrideE,
+        a_element_op,
+        b_element_op,
+        cde_element_op,
+        qs_element_op,
+        rs_element_op);
 
     if(!gemmReduce.IsSupportedArgument(gemmReduce_argument))
     {
-        throw std::runtime_error(
-            "wrong! device_gemm with the specified compilation parameters does "
-            "not support this GEMM problem");
+        throw std::runtime_error("wrong! this device_op instance does not support this problem");
     }
 
-    reduceMean_device_buf.SetZero();
-    reduceMeanSquare_device_buf.SetZero();
+    // init reducetion buffer to 0
+    r0_Mean_device_buf.SetZero();
+    r1_MeanSquare_device_buf.SetZero();
 
     // Prepare LayerNorm
-    std::array<const void*, 5> input = {c_device_buf.GetDeviceBuffer(),
-                                        reduceMean_device_buf.GetDeviceBuffer(),
-                                        reduceMeanSquare_device_buf.GetDeviceBuffer(),
+    std::array<const void*, 5> input = {e_device_buf.GetDeviceBuffer(),
+                                        r0_Mean_device_buf.GetDeviceBuffer(),
+                                        r1_MeanSquare_device_buf.GetDeviceBuffer(),
                                         gamma_device_buf.GetDeviceBuffer(),
                                         beta_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {layerNorm_device_buf.GetDeviceBuffer()};
@@ -361,12 +344,12 @@ int main()
     auto normalize_argument = normalize.MakeArgument(input,
                                                      output,
                                                      {M, N},
-                                                     {StrideC, 1},
+                                                     {StrideE, 1},
                                                      {1, 0},
                                                      {1, 0},
                                                      {0, 1},
                                                      {0, 1},
-                                                     {StrideC, 1},
+                                                     {StrideE, 1},
                                                      NormalizeFunctor{});
 
     if(!normalize.IsSupportedArgument(normalize_argument))
@@ -383,21 +366,20 @@ int main()
     {
         // verification
         Tensor<LayerNormOutDataType> host_layerNorm_m_n(
-            f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
-
-        host_gemm_layernorm<CDataType, ReduceDataType, ReduceAccDataType>(host_layerNorm_m_n,
-                                                                          a_m_k,
-                                                                          b_k_n,
-                                                                          bias_n,
-                                                                          c1_m_n,
-                                                                          gamma_n,
-                                                                          beta_n,
-                                                                          a_element_op,
-                                                                          b_element_op,
-                                                                          c_element_op,
-                                                                          d_element_op,
-                                                                          M,
-                                                                          N);
+            f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+
+        host_gemm_layernorm(host_layerNorm_m_n,
+                            a_m_k,
+                            b_k_n,
+                            bias_n,
+                            d1_m_n,
+                            gamma_n,
+                            beta_n,
+                            a_element_op,
+                            b_element_op,
+                            cde_element_op,
+                            M,
+                            N);
 
         layerNorm_device_buf.FromDevice(layerNorm_m_n.mData.data());
         pass &= ck::utils::check_err(layerNorm_m_n.mData,
@@ -419,10 +401,11 @@ int main()
         if(time_kernel)
             DumpGemmLayerNormPerf<ADataType,
                                   BDataType,
-                                  CDataType,
-                                  BiasDataType,
+                                  EDataType,
                                   D0DataType,
-                                  ReduceDataType,
+                                  D1DataType,
+                                  R0DataType,
+                                  R1DataType,
                                   GammaDataType,
                                   BetaDataType,
                                   LayerNormOutDataType>(

example/21_gemm_layernorm/gemm_layernorm_xdl_fp16.cpp

@@ -9,7 +9,7 @@
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
-#include "ck/tensor_operation/gpu/device/device_gemm_reduce_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
 #include "ck/tensor_operation/gpu/device/device_5ary_elementwise.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
@@ -28,65 +28,73 @@ using F32 = float;
 using Row = ck::tensor_layout::gemm::RowMajor;
 using Col = ck::tensor_layout::gemm::ColumnMajor;
 
+// DataType
 using ADataType                = F16;
 using BDataType                = F16;
-using CDataType                = F16;
 using GemmAccDataType          = F32;
+using CShuffleDataType         = F32;
+using DsDataType               = ck::Tuple<>;
+using EDataType                = F16;
 using ReduceAccDataType        = F32;
-using ReduceDataType           = F32;
-using ReducePtrsGlobal         = ck::Tuple<ReduceDataType*, ReduceDataType*>;
+using R0DataType               = F32;
+using R1DataType               = F32;
+using RsDataType               = ck::Tuple<R0DataType, R1DataType>;
 using GammaDataType            = F16;
 using BetaDataType             = F16;
 using LayerNormOutDataType     = F16;
 using NormalizeComputeDataType = F32;
 
-using ALayout = ck::tensor_layout::gemm::RowMajor;
-using BLayout = ck::tensor_layout::gemm::ColumnMajor;
-using CLayout = ck::tensor_layout::gemm::RowMajor;
-
-using AElementOp  = ck::tensor_operation::element_wise::PassThrough;
-using BElementOp  = ck::tensor_operation::element_wise::PassThrough;
-using CElementOp  = ck::tensor_operation::element_wise::PassThrough;
-using ReduceSumOp = ck::reduce::Add;
-using ReduceOps   = ck::Tuple<ReduceSumOp, ReduceSumOp>;
-
-using UnaryIdenticElementOp = ck::tensor_operation::element_wise::PassThrough;
-using UnaryDivElementOp     = ck::tensor_operation::element_wise::UnaryDivide;
-using UnarySquareElementOp  = ck::tensor_operation::element_wise::UnarySquare;
-using ReduceInElementOps    = ck::Tuple<UnaryIdenticElementOp, UnarySquareElementOp>;
-using ReduceOutElementOps   = ck::Tuple<UnaryDivElementOp, UnaryDivElementOp>;
-
-using ReduceGlobalMemOps =
-    ck::InMemoryDataOperationEnumSequence<ck::InMemoryDataOperationEnum::AtomicAdd,
-                                          ck::InMemoryDataOperationEnum::AtomicAdd>;
-
-static constexpr auto GemmSpecialization =
-    ck::tensor_operation::device::GemmSpecialization::Default;
+// Layout
+using ALayout  = Row;
+using BLayout  = Col;
+using D1Layout = Row;
+using ELayout  = D1Layout;
+
+// Elementwise op
+using PassThrough  = ck::tensor_operation::element_wise::PassThrough;
+using Square       = ck::tensor_operation::element_wise::UnarySquare;
+using Div          = ck::tensor_operation::element_wise::UnaryDivide;
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+using QsElementOp  = ck::Tuple<PassThrough, Square>;
+using RsElementOp  = ck::Tuple<Div, Div>;
+
+// ReduceOp
+using R0ThreadReduceOp = ck::reduce::Add;
+using R1ThreadReduceOp = ck::reduce::Add;
+using RsThreadReduceOp = ck::Tuple<R0ThreadReduceOp, R1ThreadReduceOp>;
+
+static constexpr auto R0GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+static constexpr auto R1GlobalReduceOp = ck::InMemoryDataOperationEnum::AtomicAdd;
+using RsGlobalReduceOp = ck::InMemoryDataOperationEnumSequence<R0GlobalReduceOp, R1GlobalReduceOp>;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
 
 // clang-format off
-using DeviceGemmReduceInstance = ck::tensor_operation::device::DeviceGemmReduce_Xdl_CShuffle
-//######| ALayout| BLayout| CLayout|AData| BData| CData|  GemmAcc| CShuffle| ReduceAcc|        ReduceData|           A|           B|           C|    Reduce|     ReduceInEleOp|      ReduceAccEleOp|             Reduce|               GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer|              CReduce| CReduceThreadLds2VGprCopy| CReduceThreadVgpr2GlobalCopy|
-//######|        |        |        | Type|  Type|  Type| DataType| DataType|  DataType|        Type Tuple| Elementwise| Elementwise| Elementwise| Operation|                  |                    |         MemoryData|     Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN| MXdlPerWave| NXdlPerWave|            _MBlock_MPerBlock| ScalarPerVector| ThreadClusterLengths|     SrcDstScalarPerVector|        SrcDstScalarPerVector|
-//######|        |        |        |     |      |      |         |         |          |                  |   Operation|   Operation|   Operation|          |                  |                    |          Operation|                   |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|            _NBlock_NPerBlock|      _NPerBlock| _MPerBlock_NPerBlock|                _NPerBlock|                   _MPerBlock|
-//######|        |        |        |     |      |      |         |         |          |                  |            |            |            |          |                  |                    |                   |                   |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |                     |                          |                             |
-        <     Row,     Col,     Row,  F16,   F16,   F16,      F32,      F32,       F32,  ReducePtrsGlobal,  AElementOp,  BElementOp,  CElementOp, ReduceOps,ReduceInElementOps, ReduceOutElementOps, ReduceGlobalMemOps, GemmSpecialization,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,             S<64, 4>,                         4,                            1>;
+using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleDMultipleR_Xdl_CShuffle
+//######| ALayout| BLayout| ELayout|     AData|     BData|     GemmAccData|         CShuffle|     DsData|     EData|     ReduceAccData|     RsData|           A|           B|          CDE|          Qs|          Rs|           Thread|           Global|           GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle|    CDRThreadTransfer|                  CDE|    RThreadTransfer|
+//######|        |        |        |      Type|      Type|            Type|         DataType|       Type|      Type|              Type|       Type| Elementwise| Elementwise|  Elementwise| Elementwise| Elementwise|           Reduce|           Reduce| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|       ClusterLengths| ReduceThreadTransfer| DstScalarPerVector|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |   Operation|   Operation|    Operation|   Operation|   Operation|        Operation|        Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle| _MPerBlock_NPerBlock|      ScalarPerVector|         _MPerBlock|
+//######|        |        |        |          |          |                |                 |           |          |                  |           |            |            |             |            |            |                 |                 |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                     |           _NPerBlock|                   |
+        < ALayout, BLayout, ELayout, ADataType, BDataType, GemmAccDataType, CShuffleDataType, DsDataType, EDataType, ReduceAccDataType, RsDataType,  AElementOp,  BElementOp, CDEElementOp, QsElementOp, RsElementOp, RsThreadReduceOp, RsGlobalReduceOp,    GemmDefault,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,             S<64, 4>,                    4,                  1>;
 // clang-format on
 
 using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
                                                                         BDataType,
-                                                                        CDataType,
+                                                                        EDataType,
                                                                         GemmAccDataType,
                                                                         AElementOp,
                                                                         BElementOp,
-                                                                        CElementOp>;
+                                                                        PassThrough>;
 
 using NormalizeFunctor = ck::tensor_operation::element_wise::Normalize;
 
 // A:x, B:E[x], C:E[x^2], D:Gamma, E:Beta , F:y
 using DeviceNormalizeInstance =
-    ck::tensor_operation::device::Device5AryElementwise<CDataType,
-                                                        ReduceDataType,
-                                                        ReduceDataType,
+    ck::tensor_operation::device::Device5AryElementwise<EDataType,
+                                                        R0DataType,
+                                                        R1DataType,
                                                         GammaDataType,
                                                         BetaDataType,
                                                         LayerNormOutDataType,
@@ -120,60 +128,54 @@ auto f_host_tensor_descriptor2d =
         }
     };
 
-template <typename CDataType,
-          typename ReduceDataType,
-          typename A_functor,
-          typename B_functor,
-          typename C_functor>
 void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
                          const Tensor<ADataType>& a_m_k,
-                         const Tensor<ADataType>& b_k_n,
+                         const Tensor<BDataType>& b_k_n,
                          const Tensor<GammaDataType>& gamma_n,
                          const Tensor<BetaDataType>& beta_n,
-                         A_functor a_element_op,
-                         B_functor b_element_op,
-                         C_functor c_element_op,
+                         AElementOp a_element_op,
+                         BElementOp b_element_op,
+                         CDEElementOp c_element_op,
                          int M,
                          int N)
 {
-    using out_type = ck::remove_reference_t<decltype(out_m_n(0, 0))>;
 
-    int StrideC = N;
-    Tensor<CDataType> c_m_n(f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> mean_m(f_host_tensor_descriptor1d(M, 1));
-    Tensor<ReduceDataType> meanSquare_m(f_host_tensor_descriptor1d(M, 1));
-    auto averageOpInst = UnaryDivElementOp{N};
+    int StrideE = N;
+    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> mean_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<R1DataType> meanSquare_m(f_host_tensor_descriptor1d(M, 1));
+    auto averageOpInst = Div{N};
 
     auto ref_gemm    = ReferenceGemmInstance{};
     auto ref_invoker = ref_gemm.MakeInvoker();
 
     auto ref_argument =
-        ref_gemm.MakeArgument(a_m_k, b_k_n, c_m_n, a_element_op, b_element_op, c_element_op);
+        ref_gemm.MakeArgument(a_m_k, b_k_n, e_m_n, a_element_op, b_element_op, c_element_op);
 
     ref_invoker.Run(ref_argument);
 
     // reduce_mean and reduce_square_mean
-    auto reduceSumOpInst = ReduceSumOp{};
+    auto r0Op = R0ThreadReduceOp{};
+    auto r1Op = R1ThreadReduceOp{};
     for(int m = 0; m < M; ++m)
     {
-        auto mean_acc        = reduceSumOpInst.GetIdentityValue<ReduceAccDataType>();
-        auto square_mean_acc = reduceSumOpInst.GetIdentityValue<ReduceAccDataType>();
+        auto mean_acc        = r0Op.GetIdentityValue<ReduceAccDataType>();
+        auto mean_square_acc = r1Op.GetIdentityValue<ReduceAccDataType>();
 
         for(int n = 0; n < N; ++n)
         {
-            auto c_val        = ck::type_convert<ReduceAccDataType>(c_m_n(m, n));
-            auto square_c_val = reduceSumOpInst.GetIdentityValue<ReduceAccDataType>();
+            auto e_val                     = ck::type_convert<ReduceAccDataType>(e_m_n(m, n));
+            ReduceAccDataType square_e_val = 0;
+            Square{}(square_e_val, e_val);
 
-            UnarySquareElementOp{}(square_c_val, c_val);
-
-            reduceSumOpInst(mean_acc, c_val);
-            reduceSumOpInst(square_mean_acc, square_c_val);
+            r0Op(mean_acc, e_val);
+            r1Op(mean_square_acc, square_e_val);
         }
 
         averageOpInst(mean_acc, mean_acc);
-        averageOpInst(square_mean_acc, square_mean_acc);
-        mean_m(m)       = ck::type_convert<ReduceDataType>(mean_acc);
-        meanSquare_m(m) = ck::type_convert<ReduceDataType>(square_mean_acc);
+        averageOpInst(mean_square_acc, mean_square_acc);
+        mean_m(m)       = ck::type_convert<R0DataType>(mean_acc);
+        meanSquare_m(m) = ck::type_convert<R1DataType>(mean_square_acc);
     }
 
     // LayerNorm
@@ -182,22 +184,23 @@ void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
     {
         for(int n = 0; n < N; ++n)
         {
-            float out_f32 = 0;
-            layerNormInst(out_f32,
-                          static_cast<float>(c_m_n(m, n)),
-                          static_cast<float>(mean_m(m)),
-                          static_cast<float>(meanSquare_m(m)),
-                          static_cast<float>(gamma_n(n)),
-                          static_cast<float>(beta_n(n)));
-            out_m_n(m, n) = static_cast<out_type>(out_f32);
+            NormalizeComputeDataType out_acc = 0;
+            layerNormInst(out_acc,
+                          ck::type_convert<NormalizeComputeDataType>(e_m_n(m, n)),
+                          ck::type_convert<NormalizeComputeDataType>(mean_m(m)),
+                          ck::type_convert<NormalizeComputeDataType>(meanSquare_m(m)),
+                          ck::type_convert<NormalizeComputeDataType>(gamma_n(n)),
+                          ck::type_convert<NormalizeComputeDataType>(beta_n(n)));
+            out_m_n(m, n) = ck::type_convert<LayerNormOutDataType>(out_acc);
         }
     }
 }
 
 template <typename ADataType,
           typename BDataType,
-          typename CDataType,
-          typename ReduceDataType,
+          typename EDataType,
+          typename R0DataType,
+          typename R1DataType,
           typename GammaDataType,
           typename BetaDataType,
           typename NormalizeDataType>
@@ -205,11 +208,11 @@ void DumpGemmLayerNormPerf(float gemm_reduce_time, float normalize_time, int M,
 {
     std::size_t gemm_flop     = std::size_t(2) * M * N * K;
     std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
-                                sizeof(CDataType) * M * N + sizeof(ReduceDataType) * M +
-                                sizeof(ReduceDataType) * M;
+                                sizeof(EDataType) * M * N + sizeof(R0DataType) * M +
+                                sizeof(R1DataType) * M;
 
-    std::size_t normalize_num_btye = sizeof(CDataType) * M * N + sizeof(ReduceDataType) * M +
-                                     sizeof(ReduceDataType) * M + sizeof(GammaDataType) * N +
+    std::size_t normalize_num_btye = sizeof(EDataType) * M * N + sizeof(R0DataType) * M +
+                                     sizeof(R1DataType) * M + sizeof(GammaDataType) * N +
                                      sizeof(BetaDataType) * N + sizeof(NormalizeDataType) * M * N;
 
     float tflops               = static_cast<float>(gemm_flop) / 1.E9 / gemm_reduce_time;
@@ -232,17 +235,17 @@ int main()
 
     ck::index_t StrideA = 1024;
     ck::index_t StrideB = 1024;
-    ck::index_t StrideC = 1024;
+    ck::index_t StrideE = 1024;
 
     Tensor<ADataType> a_m_k(f_host_tensor_descriptor2d(M, K, StrideA, ALayout{}));
     Tensor<BDataType> b_k_n(f_host_tensor_descriptor2d(K, N, StrideB, BLayout{}));
-    Tensor<CDataType> c_m_n(f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
-    Tensor<ReduceDataType> reduceMean_m(f_host_tensor_descriptor1d(M, 1));
-    Tensor<ReduceDataType> reduceMeanSquare_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> r0_Mean_m(f_host_tensor_descriptor1d(M, 1));
+    Tensor<R1DataType> r1_MeanSquare_m(f_host_tensor_descriptor1d(M, 1));
     Tensor<GammaDataType> gamma_n(f_host_tensor_descriptor1d(N, 1));
     Tensor<BetaDataType> beta_n(f_host_tensor_descriptor1d(N, 1));
     Tensor<LayerNormOutDataType> layerNorm_m_n(
-        f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
+        f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
 
     a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{-1, 1});
     b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-1, 1});
@@ -251,11 +254,10 @@ int main()
 
     DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
     DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
-    DeviceMem c_device_buf(sizeof(CDataType) * c_m_n.mDesc.GetElementSpaceSize());
-    DeviceMem reduceMean_device_buf(sizeof(ReduceDataType) *
-                                    reduceMean_m.mDesc.GetElementSpaceSize());
-    DeviceMem reduceMeanSquare_device_buf(sizeof(ReduceDataType) *
-                                          reduceMeanSquare_m.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem r0_Mean_device_buf(sizeof(R0DataType) * r0_Mean_m.mDesc.GetElementSpaceSize());
+    DeviceMem r1_MeanSquare_device_buf(sizeof(R1DataType) *
+                                       r1_MeanSquare_m.mDesc.GetElementSpaceSize());
     DeviceMem gamma_device_buf(sizeof(GammaDataType) * gamma_n.mDesc.GetElementSpaceSize());
     DeviceMem beta_device_buf(sizeof(BetaDataType) * beta_n.mDesc.GetElementSpaceSize());
     DeviceMem layerNorm_device_buf(sizeof(LayerNormOutDataType) *
@@ -266,40 +268,33 @@ int main()
     gamma_device_buf.ToDevice(gamma_n.mData.data());
     beta_device_buf.ToDevice(beta_n.mData.data());
 
-    auto a_element_op                     = AElementOp{};
-    auto b_element_op                     = BElementOp{};
-    auto c_element_op                     = CElementOp{};
-    std::array<void*, 3> gemm_element_ops = {&a_element_op, &b_element_op, &c_element_op};
-
-    auto passthrough                            = UnaryIdenticElementOp{};
-    auto square                                 = UnarySquareElementOp{};
-    auto div                                    = UnaryDivElementOp{N};
-    std::array<void*, 2> reduce_in_element_ops  = {&passthrough, &square};
-    std::array<void*, 2> reduce_out_element_ops = {&div, &div};
-
-    std::array<void*, 2> p_reduces = {reduceMean_device_buf.GetDeviceBuffer(),
-                                      reduceMeanSquare_device_buf.GetDeviceBuffer()};
+    auto a_element_op   = AElementOp{};
+    auto b_element_op   = BElementOp{};
+    auto cde_element_op = CDEElementOp{};
+    auto qs_element_op  = QsElementOp{};
+    auto rs_element_op  = RsElementOp{N, N};
 
-    // Prepare GEMM, reduce_mean, reduce_mean_square
-    auto gemmReduce          = DeviceGemmReduceInstance{};
+    // Prepare GEMM, mean, mean_square
+    auto gemmReduce          = DeviceOpInstance{};
     auto gemmReduce_invoker  = gemmReduce.MakeInvoker();
-    auto gemmReduce_argument = gemmReduce.MakeArgument(a_device_buf.GetDeviceBuffer(),
-                                                       b_device_buf.GetDeviceBuffer(),
-                                                       nullptr,
-                                                       {},
-                                                       c_device_buf.GetDeviceBuffer(),
-                                                       p_reduces,
-                                                       M,
-                                                       N,
-                                                       K,
-                                                       StrideA,
-                                                       StrideB,
-                                                       StrideC,
-                                                       {},
-                                                       gemm_element_ops,
-                                                       {},
-                                                       reduce_in_element_ops,
-                                                       reduce_out_element_ops);
+    auto gemmReduce_argument = gemmReduce.MakeArgument(
+        a_device_buf.GetDeviceBuffer(),
+        b_device_buf.GetDeviceBuffer(),
+        {},
+        e_device_buf.GetDeviceBuffer(),
+        {r0_Mean_device_buf.GetDeviceBuffer(), r1_MeanSquare_device_buf.GetDeviceBuffer()},
+        M,
+        N,
+        K,
+        StrideA,
+        StrideB,
+        {},
+        StrideE,
+        a_element_op,
+        b_element_op,
+        cde_element_op,
+        qs_element_op,
+        rs_element_op);
 
     if(!gemmReduce.IsSupportedArgument(gemmReduce_argument))
     {
@@ -308,13 +303,13 @@ int main()
             "not support this GEMM problem");
     }
 
-    reduceMean_device_buf.SetZero();
-    reduceMeanSquare_device_buf.SetZero();
+    r0_Mean_device_buf.SetZero();
+    r1_MeanSquare_device_buf.SetZero();
 
     // Prepare LayerNorm
-    std::array<const void*, 5> input = {c_device_buf.GetDeviceBuffer(),
-                                        reduceMean_device_buf.GetDeviceBuffer(),
-                                        reduceMeanSquare_device_buf.GetDeviceBuffer(),
+    std::array<const void*, 5> input = {e_device_buf.GetDeviceBuffer(),
+                                        r0_Mean_device_buf.GetDeviceBuffer(),
+                                        r1_MeanSquare_device_buf.GetDeviceBuffer(),
                                         gamma_device_buf.GetDeviceBuffer(),
                                         beta_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {layerNorm_device_buf.GetDeviceBuffer()};
@@ -324,12 +319,12 @@ int main()
     auto normalize_argument = normalize.MakeArgument(input,
                                                      output,
                                                      {M, N},
-                                                     {StrideC, 1},
+                                                     {StrideE, 1},
                                                      {1, 0},
                                                      {1, 0},
                                                      {0, 1},
                                                      {0, 1},
-                                                     {StrideC, 1},
+                                                     {StrideE, 1},
                                                      NormalizeFunctor{});
 
     if(!normalize.IsSupportedArgument(normalize_argument))
@@ -346,18 +341,18 @@ int main()
     {
         // verification
         Tensor<LayerNormOutDataType> host_layerNorm_m_n(
-            f_host_tensor_descriptor2d(M, N, StrideC, CLayout{}));
-
-        host_gemm_layernorm<CDataType, ReduceDataType>(host_layerNorm_m_n,
-                                                       a_m_k,
-                                                       b_k_n,
-                                                       gamma_n,
-                                                       beta_n,
-                                                       a_element_op,
-                                                       b_element_op,
-                                                       c_element_op,
-                                                       M,
-                                                       N);
+            f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+
+        host_gemm_layernorm(host_layerNorm_m_n,
+                            a_m_k,
+                            b_k_n,
+                            gamma_n,
+                            beta_n,
+                            a_element_op,
+                            b_element_op,
+                            cde_element_op,
+                            M,
+                            N);
 
         layerNorm_device_buf.FromDevice(layerNorm_m_n.mData.data());
         pass &= ck::utils::check_err(layerNorm_m_n.mData,
@@ -379,8 +374,9 @@ int main()
         if(time_kernel)
             DumpGemmLayerNormPerf<ADataType,
                                   BDataType,
-                                  CDataType,
-                                  ReduceDataType,
+                                  EDataType,
+                                  R0DataType,
+                                  R1DataType,
                                   GammaDataType,
                                   BetaDataType,
                                   LayerNormOutDataType>(

example/CMakeLists.txt

@@ -30,7 +30,7 @@ add_subdirectory(12_reduce)
 add_subdirectory(13_pool2d_fwd)
 add_subdirectory(14_gemm_xdl_requant_relu_requant)
 add_subdirectory(15_grouped_gemm)
-add_subdirectory(16_gemm_reduce)
+add_subdirectory(16_gemm_multi_d_multi_reduces)
 add_subdirectory(17_convnd_bwd_data)
 add_subdirectory(18_batched_gemm_reduce)
 add_subdirectory(19_binary_elementwise)

include/ck/tensor_operation/gpu/block/blockwise_gemm_xdlops_skip_b_lds.hpp

+#ifndef CK_BLOCKWISE_GEMM_XDLOPS_B_REGISTER_HPP
+#define CK_BLOCKWISE_GEMM_XDLOPS_B_REGISTER_HPP
+
+#include "ck/utility/common_header.hpp"
+#include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp"
+#include "ck/tensor_operation/gpu/warp/xdlops_gemm.hpp"
+#include "ck/tensor_description/tensor_adaptor.hpp"
+
+namespace ck {
+
+template <index_t BlockSize,
+          typename FloatAB,
+          typename FloatAcc,
+          typename AK0MK1BlockDesc,
+          typename BK0K0BN0N1N2N3K1BlockDesc,
+          index_t MPerBlock,
+          index_t NPerBlock,
+          index_t K0PerBlock,
+          index_t MPerXDL,
+          index_t NPerXDL,
+          index_t MRepeat,
+          index_t NRepeat,
+          index_t KPack>
+struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1
+{
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+    static constexpr auto I2 = Number<2>{};
+    static constexpr auto I3 = Number<3>{};
+
+    static constexpr index_t WaveSize = 64;
+
+    static constexpr index_t KPerBlock = K0PerBlock * KPack;
+
+    static constexpr index_t A_K0 = AK0MK1BlockDesc{}.GetLength(I0);
+    static constexpr index_t A_K1 = AK0MK1BlockDesc{}.GetLength(I2);
+
+    static constexpr auto xdlops_gemm = XdlopsGemm<FloatAB, MPerXDL, NPerXDL, KPack>{};
+
+    static constexpr index_t KPerThread  = KPerBlock / xdlops_gemm.K0PerXdlops;
+    static constexpr index_t K0PerThread = K0PerBlock / xdlops_gemm.K0PerXdlops;
+
+    static constexpr index_t MWaves = MPerBlock / (MRepeat * MPerXDL);
+    static constexpr index_t NWaves = NPerBlock / (NRepeat * NPerXDL);
+
+    StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr,
+                              FloatAcc,
+                              MRepeat * NRepeat,
+                              xdlops_gemm.GetRegSizePerXdlops(),
+                              true>
+        c_thread_buf_;
+
+    __host__ __device__ constexpr auto& GetCThreadBuffer() { return c_thread_buf_; }
+
+    __device__ static auto GetWaveIdx()
+    {
+        const index_t thread_id = get_thread_local_1d_id();
+
+        constexpr auto threadid_to_wave_idx_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_merge_transform(make_tuple(MWaves, NWaves, WaveSize))),
+            make_tuple(Sequence<0, 1, 2>{}),
+            make_tuple(Sequence<0>{}));
+
+        return threadid_to_wave_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id));
+    }
+
+    __device__ static auto CalculateAThreadOriginDataIndex()
+    {
+        const auto wave_idx = GetWaveIdx();
+
+        const auto waveId_m = wave_idx[I0];
+
+        const auto xdlops_a_idx = xdlops_gemm.CalculateAThreadOriginDataIndex();
+
+        return make_tuple(0, waveId_m, xdlops_a_idx[I1], KPerThread * xdlops_a_idx[I0]);
+    }
+
+    __device__ static auto CalculateBThreadOriginDataIndex()
+    {
+        const auto wave_idx = GetWaveIdx();
+
+        const auto waveId_n = wave_idx[I1];
+
+        const auto xdlops_b_idx = xdlops_gemm.CalculateBThreadOriginDataIndex();
+
+        return make_tuple(0, waveId_n, xdlops_b_idx[I1], KPerThread * xdlops_b_idx[I0]);
+    }
+
+    template <index_t m0, index_t n0, index_t xdlops_i, index_t blk_i>
+    __device__ static auto
+        CalculateCThreadOriginDataIndex(Number<m0>, Number<n0>, Number<xdlops_i>, Number<blk_i>)
+    {
+        const auto wave_idx = GetWaveIdx();
+
+        const auto waveId_m = wave_idx[I0];
+        const auto waveId_n = wave_idx[I1];
+
+        const auto blk_idx = xdlops_gemm.GetBeginOfThreadBlk(xdlops_i, blk_i);
+
+        constexpr auto mrepeat_mwave_mperxdl_to_m_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_unmerge_transform(make_tuple(MRepeat, MWaves, MPerXDL))),
+            make_tuple(Sequence<0>{}),
+            make_tuple(Sequence<0, 1, 2>{}));
+
+        constexpr auto nrepeat_nwave_nperxdl_to_n_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_unmerge_transform(make_tuple(NRepeat, NWaves, NPerXDL))),
+            make_tuple(Sequence<0>{}),
+            make_tuple(Sequence<0, 1, 2>{}));
+
+        const index_t c_thread_m = mrepeat_mwave_mperxdl_to_m_adaptor.CalculateBottomIndex(
+            make_tuple(m0, waveId_m, blk_idx[I0]))[I0];
+        const index_t c_thread_n = nrepeat_nwave_nperxdl_to_n_adaptor.CalculateBottomIndex(
+            make_tuple(n0, waveId_n, blk_idx[I1]))[I0];
+
+        return make_tuple(c_thread_m, c_thread_n);
+    }
+
+    __host__ __device__ BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1()
+    {
+        static_assert(AK0MK1BlockDesc::IsKnownAtCompileTime() &&
+                          BK0K0BN0N1N2N3K1BlockDesc::IsKnownAtCompileTime(),
+                      "wrong! Desc should be known at compile-time");
+
+        static_assert(BlockSize == MWaves * NWaves * WaveSize,
+                      "BlockSize != MWaves * NWaves * WaveSize\n");
+
+        static_assert(MPerBlock % (MPerXDL * MRepeat) == 0 && NPerBlock % (NPerXDL * NRepeat) == 0,
+                      "wrong!");
+    }
+
+    __host__ __device__ static constexpr auto GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2()
+    {
+        constexpr auto c_m0_m1_m2_n_tblk_lens = xdlops_gemm.GetCM0M1M2NThreadBlkLengths();
+
+        constexpr auto M0 = c_m0_m1_m2_n_tblk_lens[I0];
+        constexpr auto M1 = c_m0_m1_m2_n_tblk_lens[I1];
+        constexpr auto M2 = c_m0_m1_m2_n_tblk_lens[I2];
+        constexpr auto N  = c_m0_m1_m2_n_tblk_lens[I3];
+
+        return make_naive_tensor_descriptor_packed(
+            make_tuple(Number<MRepeat>{}, Number<NRepeat>{}, I1, I1, M0, M1, M2, N));
+    }
+
+    __host__ __device__ static constexpr auto GetCThreadDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2()
+    {
+        constexpr auto c_m0_m1_m2_n_tblk_lens = xdlops_gemm.GetCM0M1M2NThreadBlkLengths();
+
+        constexpr auto M0 = c_m0_m1_m2_n_tblk_lens[I0];
+        constexpr auto M1 = c_m0_m1_m2_n_tblk_lens[I1];
+        constexpr auto M2 = c_m0_m1_m2_n_tblk_lens[I2];
+        constexpr auto N  = c_m0_m1_m2_n_tblk_lens[I3];
+
+        return make_naive_tensor_descriptor_packed(
+            make_tuple(I1, Number<MRepeat>{}, Number<NRepeat>{}, I1, I1, M0, M1, M2, N));
+    }
+
+    __host__ __device__ static constexpr auto GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2()
+    {
+        constexpr auto c_block_desc_m0_n0_m1_n1_m2_n2 =
+            make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{},
+                                                           Number<NRepeat>{},
+                                                           Number<MWaves>{},
+                                                           Number<NWaves>{},
+                                                           Number<MPerXDL>{},
+                                                           Number<NPerXDL>{}));
+
+        return xdlops_gemm.MakeCDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(c_block_desc_m0_n0_m1_n1_m2_n2);
+    }
+
+    __host__ __device__ static constexpr auto GetCBlockDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2()
+    {
+        constexpr auto c_block_desc_g_m0_n0_m1_n1_m2_n2 =
+            make_naive_tensor_descriptor_packed(make_tuple(I1,
+                                                           Number<MRepeat>{},
+                                                           Number<NRepeat>{},
+                                                           Number<MWaves>{},
+                                                           Number<NWaves>{},
+                                                           Number<MPerXDL>{},
+                                                           Number<NPerXDL>{}));
+
+        return xdlops_gemm.MakeCDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2(
+            c_block_desc_g_m0_n0_m1_n1_m2_n2);
+    }
+
+    template <typename CGridDesc_M_N>
+    __host__ __device__ static constexpr auto
+    MakeCGridDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(const CGridDesc_M_N& c_grid_desc_m_n)
+    {
+        const auto M = c_grid_desc_m_n.GetLength(I0);
+        const auto N = c_grid_desc_m_n.GetLength(I1);
+
+        const auto c_grid_desc_m0_n0_m1_n1_m2_n2 = transform_tensor_descriptor(
+            c_grid_desc_m_n,
+            make_tuple(make_unmerge_transform(make_tuple(M / (MWaves * MPerXDL), MWaves, MPerXDL)),
+                       make_unmerge_transform(make_tuple(N / (NWaves * NPerXDL), NWaves, NPerXDL))),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3, 5>{}));
+
+        return xdlops_gemm.MakeCDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(c_grid_desc_m0_n0_m1_n1_m2_n2);
+    }
+
+    template <typename CGridDesc_G_M_N>
+    __host__ __device__ static constexpr auto
+    MakeCGridDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2(const CGridDesc_G_M_N& c_grid_desc_g_m_n)
+    {
+        const auto G = c_grid_desc_g_m_n.GetLength(I0);
+        const auto M = c_grid_desc_g_m_n.GetLength(I1);
+        const auto N = c_grid_desc_g_m_n.GetLength(I2);
+
+        const auto c_grid_desc_g_m0_n0_m1_n1_m2_n2 = transform_tensor_descriptor(
+            c_grid_desc_g_m_n,
+            make_tuple(make_pass_through_transform(G),
+                       make_unmerge_transform(make_tuple(M / (MWaves * MPerXDL), MWaves, MPerXDL)),
+                       make_unmerge_transform(make_tuple(N / (NWaves * NPerXDL), NWaves, NPerXDL))),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+            make_tuple(Sequence<0>{}, Sequence<1, 3, 5>{}, Sequence<2, 4, 6>{}));
+
+        return xdlops_gemm.MakeCDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2(
+            c_grid_desc_g_m0_n0_m1_n1_m2_n2);
+    }
+
+    __host__ __device__ static constexpr auto MakeABlockDescriptor_M0_M1_M2_K()
+    {
+        return transform_tensor_descriptor(
+            AK0MK1BlockDesc{},
+            make_tuple(
+                make_merge_transform_v3_division_mod(make_tuple(Number<A_K0>{}, Number<A_K1>{})),
+                make_unmerge_transform(
+                    make_tuple(Number<MRepeat>{}, Number<MWaves>{}, Number<MPerXDL>{}))),
+            make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
+            make_tuple(Sequence<3>{}, Sequence<0, 1, 2>{}));
+    }
+
+    __device__ void MoveABlockSliceWindow()
+    {
+        a_thread_copy_.MoveSrcSliceWindow(a_block_desc_m0_m1_m2_k,
+                                          make_multi_index(0, 0, 0, K0PerBlock * KPack));
+    }
+    __device__ void ResetABlockStartWindow()
+    {
+        a_thread_copy_.SetSrcCoord(CalculateAThreadOriginDataIndex());
+    }
+
+    static constexpr auto a_block_desc_m0_m1_m2_k = MakeABlockDescriptor_M0_M1_M2_K();
+
+    template <typename ABlockBuffer, typename BBlockBuffer, typename CThreadBuffer>
+    __device__ void Run(const ABlockBuffer& a_block_buf,
+                        const BBlockBuffer& b_thread_buf,
+                        CThreadBuffer& c_thread_buf) const
+    {
+        auto a_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, FloatAB>(
+            a_thread_desc_.GetElementSpaceSize());
+
+        static_for<0, MRepeat, 1>{}([&](auto m0) {
+            // read A
+            a_thread_copy_.Run(a_block_desc_m0_m1_m2_k,
+                               make_tuple(m0, I0, I0, I0),
+                               a_block_buf,
+                               a_thread_desc_,
+                               make_tuple(I0, I0, I0, I0),
+                               a_thread_buf);
+
+            static_for<0, NRepeat, 1>{}([&](auto n0) {
+                // read B
+                static_for<0, KPerThread, KPack>{}([&](auto k) {
+                    vector_type<FloatAB, KPack> a_thread_vec;
+                    vector_type<FloatAB, KPack> b_thread_vec;
+                    constexpr index_t k0 = k / KPack;
+                    static_for<0, KPack, 1>{}([&](auto i) {
+                        a_thread_vec.template AsType<FloatAB>()(i) = a_thread_buf
+                            [Number<a_thread_desc_.CalculateOffset(make_tuple(0, 0, 0, k + i))>{}];
+                        b_thread_vec.template AsType<FloatAB>()(i) = b_thread_buf
+                            [Number<b_thread_desc_.CalculateOffset(make_tuple(k0, n0, i))>{}];
+                    });
+
+                    using mfma_input_type =
+                        typename vector_type<FloatAB, xdlops_gemm.K1PerXdlops>::type;
+
+                    constexpr index_t c_offset =
+                        c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0));
+
+                    xdlops_gemm.template Run(
+                        a_thread_vec.template AsType<mfma_input_type>(),
+                        b_thread_vec.template AsType<mfma_input_type>(),
+                        c_thread_buf.GetVectorTypeReference(Number<c_offset>{}));
+                });
+            });
+        });
+    }
+
+    private:
+    // A[M0, M1, M2, KPerThread]
+    static constexpr auto a_thread_desc_ =
+        make_naive_tensor_descriptor_packed(make_tuple(I1, I1, I1, Number<KPerThread>{}));
+
+    // B[N0, N1, N2, KPerThread]
+    static constexpr auto b_thread_desc_ =
+        make_naive_tensor_descriptor_packed(make_tuple(Number<K0PerThread>{}, // KPerThread
+                                                       Number<NRepeat>{},     // repeat
+                                                       Number<KPack>{}));
+
+    // C[M, N, NumRegXdlops]
+    static constexpr auto c_thread_desc_ = make_naive_tensor_descriptor_packed(
+        make_tuple(Number<MRepeat>{}, Number<NRepeat>{}, xdlops_gemm.GetRegSizePerXdlops()));
+
+    using AThreadCopy = ThreadwiseTensorSliceTransfer_v4<FloatAB,
+                                                         FloatAB,
+                                                         decltype(a_block_desc_m0_m1_m2_k),
+                                                         decltype(a_thread_desc_),
+                                                         Sequence<1, 1, 1, KPerThread>,
+                                                         Sequence<0, 1, 2, 3>,
+                                                         3,
+                                                         A_K1,
+                                                         A_K1>;
+
+    AThreadCopy a_thread_copy_{CalculateAThreadOriginDataIndex()};
+};
+
+} // namespace ck
+#endif