initial commit

example/15_grouped_gemm/grouped_gemm_xdl_fixed_nk_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, 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/impl/device_grouped_gemm_xdl_fixed_nk.hpp"
+#include "ck/tensor_operation/gpu/device/device_grouped_gemm.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.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/literals.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 ADataType        = F16;
+using BDataType        = F16;
+using AccDataType      = F32;
+using CShuffleDataType = F32;
+using DsDataType       = ck::Tuple<>;
+using EDataType        = F32;
+
+using ALayout  = Row;
+using BLayout  = Col;
+using DsLayout = ck::Tuple<>;
+using ELayout  = Row;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::MNPadding;
+
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGroupedGemm_Xdl_Fixed_NK
+    // clang-format off
+//######| ALayout| BLayout| DsLayout| ELayout|     AData|     BData|     AccData|         CShuffle|     DsData|     EData|           A|           B|          CDE|           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|
+//######|        |        |         |        |      Type|      Type|        Type|         DataType|       Type|      Type| Elementwise| Elementwise|  Elementwise| 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|         _MBlock_MWaveMPerXdl| ScalarPerVector|
+//######|        |        |         |        |          |          |            |                 |           |          |   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_NWaveNPerXdl|   _NWaveNPerXdl|
+//######|        |        |         |        |          |          |            |                 |           |          |            |            |             |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |
+        < ALayout, BLayout, DsLayout, ELayout, ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType,  AElementOp,  BElementOp, CDEElementOp,    GemmDefault,        1,   256,    64,   128,    32,   8,   8,   32,   32,    1,    2,  S<1, 4, 64, 1>,  S<0, 2, 1, 3>,  S<0, 2, 1, 3>,              3,              8,              8,         1,  S<1, 4, 64, 1>,  S<0, 2, 1, 3>,  S<0, 2, 1, 3>,             3,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,              4>;
+// clang-format on
+
+struct ProblemSize final
+{
+    std::vector<ck::index_t> Ms;
+    std::vector<ck::index_t> Ns;
+    std::vector<ck::index_t> Ks;
+
+    std::vector<ck::index_t> stride_As;
+    std::vector<ck::index_t> stride_Bs;
+    std::vector<ck::index_t> stride_Cs;
+
+    ck::index_t group_count;
+};
+
+struct ExecutionConfig final
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    int k_batch          = 1;
+    bool time_kernel     = false;
+};
+
+bool run_grouped_gemm(const ProblemSize& problem_size, const ExecutionConfig& config)
+{
+    auto group_count = problem_size.group_count;
+
+    // GEMM shape
+    std::vector<ck::tensor_operation::device::GemmDesc> gemm_descs;
+    std::vector<void*> p_Cs;
+
+    gemm_descs.reserve(group_count);
+
+    int sum_of_m = 0;
+
+    auto f_host_tensor_descriptor =
+        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+            using namespace ck::literals;
+
+            if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+            {
+                return HostTensorDescriptor({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    std::vector<Tensor<ADataType>> a_tensors;
+    std::vector<Tensor<BDataType>> b_tensors;
+    std::vector<Tensor<EDataType>> c_host_tensors;
+    std::vector<Tensor<EDataType>> c_device_tensors;
+
+    a_tensors.reserve(group_count);
+    b_tensors.reserve(group_count);
+    c_host_tensors.reserve(group_count);
+    c_device_tensors.reserve(group_count);
+
+    using DeviceMemPtr = std::unique_ptr<DeviceMem>;
+
+    std::vector<DeviceMemPtr> a_tensors_device, b_tensors_device, c_tensors_device;
+
+    a_tensors_device.reserve(group_count);
+    b_tensors_device.reserve(group_count);
+    c_tensors_device.reserve(group_count);
+
+    std::size_t flop = 0, num_btype = 0;
+
+    for(int i = 0; i < group_count; i++)
+    {
+        sum_of_m += problem_size.Ms[i];
+        a_tensors.push_back(Tensor<ADataType>(f_host_tensor_descriptor(
+            problem_size.Ms[i], problem_size.Ks[i], problem_size.stride_As[i], ALayout{})));
+        b_tensors.push_back(Tensor<BDataType>(f_host_tensor_descriptor(
+            problem_size.Ks[i], problem_size.Ns[i], problem_size.stride_Bs[i], BLayout{})));
+        c_host_tensors.push_back(Tensor<EDataType>(f_host_tensor_descriptor(
+            problem_size.Ms[i], problem_size.Ns[i], problem_size.stride_Cs[i], ELayout{})));
+        c_device_tensors.push_back(Tensor<EDataType>(f_host_tensor_descriptor(
+            problem_size.Ms[i], problem_size.Ns[i], problem_size.stride_Cs[i], ELayout{})));
+        std::cout << "gemm[" << i << "] a_m_k: " << a_tensors[i].mDesc
+                  << " b_k_n: " << b_tensors[i].mDesc << " c_m_n: " << c_device_tensors[i].mDesc
+                  << std::endl;
+
+        flop += std::size_t(2) * problem_size.Ms[i] * problem_size.Ks[i] * problem_size.Ns[i];
+        num_btype += sizeof(ADataType) * a_tensors[i].mDesc.GetElementSize() +
+                     sizeof(BDataType) * b_tensors[i].mDesc.GetElementSize() +
+                     sizeof(EDataType) * c_device_tensors[i].mDesc.GetElementSize();
+
+        switch(config.init_method)
+        {
+        case 0: break;
+        case 1:
+            a_tensors[i].GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
+            b_tensors[i].GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
+            break;
+        case 2:
+            a_tensors[i].GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
+            b_tensors[i].GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
+            break;
+        default:
+            a_tensors[i].GenerateTensorValue(GeneratorTensor_Sequential<0>{});
+            b_tensors[i].GenerateTensorValue(GeneratorTensor_Sequential<1>{});
+        }
+    }
+
+    using GroupedGemmKernelArgument = ck::tensor_operation::device::GroupedGemmKernelArgument<>;
+
+    std::vector<GroupedGemmKernelArgument> grouped_gemm_kernel_args_;
+    grouped_gemm_kernel_args_.reserve(group_count);
+
+    for(int i = 0; i < group_count; i++)
+    {
+        a_tensors_device.emplace_back(
+            std::make_unique<DeviceMem>(sizeof(ADataType) * sum_of_m * problem_size.Ks[i]));
+
+        b_tensors_device.emplace_back(std::make_unique<DeviceMem>(
+            sizeof(BDataType) * problem_size.Ns[i] * problem_size.Ks[i]));
+
+        c_tensors_device.emplace_back(
+            std::make_unique<DeviceMem>(sizeof(EDataType) * sum_of_m * problem_size.Ns[i]));
+
+        a_tensors_device[i]->ToDevice(a_tensors[i].mData.data(),
+                                      a_tensors[i].mDesc.GetElementSpaceSize() * sizeof(ADataType));
+        b_tensors_device[i]->ToDevice(b_tensors[i].mData.data(),
+                                      b_tensors[i].mDesc.GetElementSpaceSize() * sizeof(BDataType));
+        c_tensors_device[i]->SetZero();
+
+        p_Cs.push_back(c_tensors_device[i]->GetDeviceBuffer());
+
+        gemm_descs.push_back({sum_of_m,
+                              problem_size.Ns[i],
+                              problem_size.Ks[i],
+                              1,
+                              problem_size.stride_Bs[i],
+                              1,
+                              {}});
+
+        grouped_gemm_kernel_args_.push_back({a_tensors_device[i]->GetDeviceBuffer(),
+                                             b_tensors_device[i]->GetDeviceBuffer(),
+                                             {},
+                                             c_tensors_device[i]->GetDeviceBuffer(),
+                                             problem_size.Ms[i],
+                                             problem_size.Ns[i],
+                                             problem_size.Ks[i],
+                                             problem_size.stride_As[i],
+                                             problem_size.stride_Bs[i],
+                                             {},
+                                             problem_size.stride_Cs[i]});
+    }
+
+    auto a_element_op = AElementOp{};
+    auto b_element_op = BElementOp{};
+    auto c_element_op = CDEElementOp{};
+
+    auto gemm    = DeviceGemmInstance{};
+    auto invoker = gemm.MakeInvoker();
+
+    std::vector<const void*> p_As                = {};
+    std::vector<const void*> p_Bs                = {};
+    std::vector<std::array<const void*, 0>> p_Ds = {};
+
+    // do GEMM
+    auto argument = gemm.MakeArgument(
+        p_As, p_Bs, p_Ds, p_Cs, gemm_descs, a_element_op, b_element_op, c_element_op);
+
+    DeviceMem gemm_arg_dev_mem(gemm.GetDeviceKernelArgSize(&argument));
+    DeviceMem gemm_workspace_dev(gemm.GetWorkSpaceSize(&argument));
+
+    gemm.SetWorkSpacePointer(&argument, gemm_workspace_dev.GetDeviceBuffer());
+
+    hip_check_error(hipMemcpy(gemm_arg_dev_mem.GetDeviceBuffer(),
+                              grouped_gemm_kernel_args_.data(),
+                              gemm.GetDeviceKernelArgSize(&argument),
+                              hipMemcpyHostToDevice));
+
+    if(!gemm.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error(
+            "wrong! device_gemm with the specified compilation parameters does "
+            "not support this GEMM problem");
+    }
+
+    gemm.SetDeviceKernelArgs(argument, gemm_arg_dev_mem.GetDeviceBuffer());
+    gemm.SetKBatch(argument, config.k_batch);
+
+    invoker.Run(argument, StreamConfig{nullptr, false});
+
+    if(config.time_kernel)
+    {
+        float ave_time   = invoker.Run(argument, StreamConfig{nullptr, config.time_kernel});
+        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;
+    }
+
+    bool pass = true;
+    if(config.do_verification)
+    {
+        using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                                BDataType,
+                                                                                EDataType,
+                                                                                AccDataType,
+                                                                                AElementOp,
+                                                                                BElementOp,
+                                                                                CDEElementOp>;
+
+        for(std::size_t i = 0; i < gemm_descs.size(); i++)
+        {
+            c_tensors_device[i]->FromDevice(c_device_tensors[i].mData.data(),
+                                            c_device_tensors[i].mDesc.GetElementSize() *
+                                                sizeof(EDataType));
+            auto ref_gemm    = ReferenceGemmInstance{};
+            auto ref_invoker = ref_gemm.MakeInvoker();
+
+            auto ref_argument = ref_gemm.MakeArgument(a_tensors[i],
+                                                      b_tensors[i],
+                                                      c_host_tensors[i],
+                                                      a_element_op,
+                                                      b_element_op,
+                                                      c_element_op);
+
+            ref_invoker.Run(ref_argument);
+
+            pass &= ck::utils::check_err(c_device_tensors[i], c_host_tensors[i]);
+        }
+    }
+
+    return pass;
+}
+
+int main(int argc, char* argv[])
+{
+    ProblemSize problem_size;
+    ExecutionConfig config;
+
+    problem_size.group_count = 16;
+
+    problem_size.Ms = {
+        167, 183, 177, 181, 153, 139, 156, 173, 163, 150, 204, 184, 168, 156, 168, 148};
+
+    for(int i = 0; i < problem_size.group_count; i++)
+    {
+        problem_size.Ns.push_back(768);
+        problem_size.Ks.push_back(4608);
+
+        problem_size.stride_As.push_back(problem_size.Ks[i]);
+        problem_size.stride_Bs.push_back(problem_size.Ks[i]);
+        problem_size.stride_Cs.push_back(problem_size.Ns[i]);
+    }
+
+    if(argc == 5)
+    {
+        config.do_verification = std::stoi(argv[1]);
+        config.init_method     = std::stoi(argv[2]);
+        config.time_kernel     = std::stoi(argv[3]);
+        config.k_batch         = std::stoi(argv[4]);
+    }
+    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: k_batch (> 0)\n");
+        exit(0);
+    }
+
+    return !run_grouped_gemm(problem_size, config);
+}

example/15_grouped_gemm/grouped_gemm_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, 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/impl/device_grouped_gemm_xdl.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.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/literals.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 ADataType        = F16;
+using BDataType        = F16;
+using AccDataType      = F32;
+using CShuffleDataType = F16;
+using DsDataType       = ck::Tuple<>;
+using EDataType        = F16;
+
+using ALayout  = Row;
+using BLayout  = Col;
+using DsLayout = ck::Tuple<>;
+using ELayout  = Row;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGroupedGemm_Xdl
+    // clang-format off
+//######| ALayout| BLayout| DsLayout| ELayout|     AData|     BData|     AccData|         CShuffle|     DsData|     EData|           A|           B|          CDE|           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|
+//######|        |        |         |        |      Type|      Type|        Type|         DataType|       Type|      Type| Elementwise| Elementwise|  Elementwise| 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|         _MBlock_MWaveMPerXdl| ScalarPerVector|
+//######|        |        |         |        |          |          |            |                 |           |          |   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_NWaveNPerXdl|   _NWaveNPerXdl|
+//######|        |        |         |        |          |          |            |                 |           |          |            |            |             |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |
+        < ALayout, BLayout, DsLayout, ELayout, ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType,  AElementOp,  BElementOp, CDEElementOp,    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<1, 32, 1, 8>,               8>;
+// clang-format on
+
+#include "run_grouped_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_grouped_gemm_example(argc, argv); }

example/15_grouped_gemm/grouped_gemm_xdl_fp32.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, 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/impl/device_grouped_gemm_xdl.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.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/literals.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 ADataType        = F32;
+using BDataType        = F32;
+using AccDataType      = F32;
+using CShuffleDataType = F32;
+using DsDataType       = ck::Tuple<>;
+using EDataType        = F32;
+
+using ALayout  = Row;
+using BLayout  = Col;
+using DsLayout = ck::Tuple<>;
+using ELayout  = Row;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGroupedGemm_Xdl
+    // clang-format off
+//######| ALayout| BLayout| DsLayout| ELayout|     AData|     BData|     AccData|         CShuffle|     DsData|     EData|           A|           B|          CDE|           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|
+//######|        |        |         |        |      Type|      Type|        Type|         DataType|       Type|      Type| Elementwise| Elementwise|  Elementwise| 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|         _MBlock_MWaveMPerXdl| ScalarPerVector|
+//######|        |        |         |        |          |          |            |                 |           |          |   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_NWaveNPerXdl|   _NWaveNPerXdl|
+//######|        |        |         |        |          |          |            |                 |           |          |            |            |             |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |
+        < ALayout, BLayout, DsLayout, ELayout, ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType,  AElementOp,  BElementOp, CDEElementOp,    GemmDefault,        1,   256,   256,   128,    16,   4,   4,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              4,              4,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              4,              4,         1,           1,           1,               S<1, 32, 1, 8>,               4>;
+// clang-format on
+
+#include "run_grouped_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_grouped_gemm_example(argc, argv); }

example/15_grouped_gemm/grouped_gemm_xdl_int4.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, 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/impl/device_grouped_gemm_xdl.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.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/literals.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+using ADataType        = ck::int4_t;
+using BDataType        = ck::int4_t;
+using AccDataType      = int32_t;
+using CShuffleDataType = int32_t;
+using DsDataType       = ck::Tuple<>;
+using EDataType        = ck::int4_t;
+
+using KernelADataType = int8_t;
+using KernelBDataType = int8_t;
+using KernelEDataType = int8_t;
+
+using ALayout  = Row;
+using BLayout  = Col;
+using DsLayout = ck::Tuple<>;
+using ELayout  = Row;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGroupedGemm_Xdl
+    // clang-format off
+        < ALayout,              //ALayout
+          BLayout,              //BLayout
+          DsLayout,             //DsLayout
+          ELayout,              //ELayout
+          KernelADataType,      //ADataType    
+          KernelBDataType,      //BDataType   
+          AccDataType,          //AccDataType
+          CShuffleDataType,     //CShuffleDataType
+          DsDataType,           //DsDataType
+          KernelEDataType,      //EDataType
+          AElementOp,           //AElementwiseOperation
+          BElementOp,           //BElementwiseOperation
+          CDEElementOp,         //CDEElementwiseOperation
+          GemmDefault,          //GEMMSpecialization
+          1,                    // NumGemmKPrefetchStage
+          256,                  // BlockSize
+          256,                  // MPerBlock
+          128,                  // NPerBlock
+          64,                   // KPerBlock
+          16,                   // AK1
+          16,                   // BK1
+          32,                   // MPerXdl
+          32,                   // NPerXdl
+          4,                    // MXdlPerWave
+          2,                    // NXdlPerWave
+          S<4, 64, 1>,          // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+          S<1, 0, 2>,           // ABlockTransfer ThreadCluster ArrangeOrder
+          S<1, 0, 2>,           // ABlockTransfer SrcAccessOrder
+          2,                    // ABlockTransfer SrcVectorDim
+          16,                   // ABlockTransfer SrcScalarPerVector
+          16,                   // ABlockTransfer DstScalarPerVector_K1
+          1,                    // ABlockLdsExtraM
+          S<4, 64, 1>,          // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+          S<1, 0, 2>,           // BBlockTransfer ThreadCluster ArrangeOrder
+          S<1, 0, 2>,           // BBlockTransfer SrcAccessOrder
+          2,                    // BBlockTransfer SrcVectorDim
+          16,                   // BBlockTransfer SrcScalarPerVector
+          16,                   // BBlockTransfer DstScalarPerVector_K1
+          1,                    // BBlockLdsExtraN
+          1,                    // CShuffleMXdlPerWavePerShuffle
+          1,                    // CShuffleNXdlPerWavePerShuffle
+          S<1, 64, 1, 4>,       // CBlockTransferClusterLengths_MBlock_MWaveMPerXdl_NBlock_NWaveNPerXdl
+          16>;                  // CBlockTransferScalarPerVector_NWaveNPerXdl
+// clang-format on
+
+#define BUILD_INT4_EXAMPLE
+#include "run_grouped_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_grouped_gemm_example(argc, argv); }

example/15_grouped_gemm/grouped_gemm_xdl_int8.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, 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/impl/device_grouped_gemm_xdl.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.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/literals.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+using ADataType        = int8_t;
+using BDataType        = int8_t;
+using AccDataType      = int32_t;
+using CShuffleDataType = int8_t;
+using DsDataType       = ck::Tuple<>;
+using EDataType        = int8_t;
+
+using ALayout  = Row;
+using BLayout  = Col;
+using DsLayout = ck::Tuple<>;
+using ELayout  = Row;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGroupedGemm_Xdl
+    // clang-format off
+//######| ALayout| BLayout| DsLayout| ELayout|     AData|     BData|     AccData|         CShuffle|     DsData|     EData|           A|           B|          CDE|           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|
+//######|        |        |         |        |      Type|      Type|        Type|         DataType|       Type|      Type| Elementwise| Elementwise|  Elementwise| 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|         _MBlock_MWaveMPerXdl| ScalarPerVector|
+//######|        |        |         |        |          |          |            |                 |           |          |   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_NWaveNPerXdl|   _NWaveNPerXdl|
+//######|        |        |         |        |          |          |            |                 |           |          |            |            |             |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |
+        < ALayout, BLayout, DsLayout, ELayout, ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType,  AElementOp,  BElementOp, CDEElementOp,    GemmDefault,        1,   256,   256,   128,    64,  16,  16,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,             16,             16,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,             16,             16,         1,           1,           1,               S<1, 64, 1, 4>,             16>;
+// clang-format on
+
+#include "run_grouped_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_grouped_gemm_example(argc, argv); }

example/15_grouped_gemm/grouped_gemm_xdl_splitk_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, 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/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.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/literals.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 ADataType        = F16;
+using BDataType        = F16;
+using AccDataType      = F32;
+using CShuffleDataType = F16;
+using DsDataType       = ck::Tuple<>;
+using EDataType        = F16;
+
+using ALayout  = Row;
+using BLayout  = Col;
+using DsLayout = ck::Tuple<>;
+using ELayout  = Row;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::MNKPadding;
+
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGroupedGemmXdlSplitKCShuffle
+    // clang-format off
+//######| ALayout| BLayout| DsLayout| ELayout|     AData|     BData|     AccData|         CShuffle|     DsData|     EData|           A|           B|          CDE|           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|
+//######|        |        |         |        |      Type|      Type|        Type|         DataType|       Type|      Type| Elementwise| Elementwise|  Elementwise| 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|         _MBlock_MWaveMPerXdl| ScalarPerVector|
+//######|        |        |         |        |          |          |            |                 |           |          |   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_NWaveNPerXdl|   _NWaveNPerXdl|
+//######|        |        |         |        |          |          |            |                 |           |          |            |            |             |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |
+        < ALayout, BLayout, DsLayout, ELayout, ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType,  AElementOp,  BElementOp, CDEElementOp,    GemmDefault,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,  S<1, 4, 64, 1>,  S<0, 2, 1, 3>,  S<0, 2, 1, 3>,              3,              8,              8,         1,  S<1, 4, 64, 1>,  S<0, 2, 1, 3>,  S<0, 2, 1, 3>,             3,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8>;
+// clang-format on
+
+#include "run_grouped_gemm_example.inc"
+
+int main(int argc, char* argv[])
+{
+    ProblemSize problem_size;
+    ExecutionConfig config;
+
+    problem_size.group_count = 16;
+
+    problem_size.Ms = {
+        167, 183, 177, 181, 153, 139, 156, 173, 163, 150, 204, 184, 168, 156, 168, 148};
+
+    for(int i = 0; i < problem_size.group_count; i++)
+    {
+        problem_size.Ns.push_back(768);
+        problem_size.Ks.push_back(4608);
+
+        problem_size.stride_As.push_back(problem_size.Ks[i]);
+        problem_size.stride_Bs.push_back(problem_size.Ks[i]);
+        problem_size.stride_Cs.push_back(problem_size.Ns[i]);
+    }
+
+    if(argc == 4)
+    {
+        config.do_verification = std::stoi(argv[1]);
+        config.init_method     = std::stoi(argv[2]);
+        config.time_kernel     = std::stoi(argv[3]);
+    }
+    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");
+        exit(0);
+    }
+
+    return !run_grouped_gemm(problem_size, config);
+}

example/15_grouped_gemm/run_grouped_gemm_example.inc

+#pragma once
+
+struct ProblemSize final
+{
+    std::vector<ck::index_t> Ms;
+    std::vector<ck::index_t> Ns;
+    std::vector<ck::index_t> Ks;
+
+    std::vector<ck::index_t> stride_As;
+    std::vector<ck::index_t> stride_Bs;
+    std::vector<ck::index_t> stride_Cs;
+
+    ck::index_t group_count;
+};
+
+struct ExecutionConfig final
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = false;
+};
+
+bool run_grouped_gemm(const ProblemSize& problem_size, const ExecutionConfig& config)
+{
+#if defined(BUILD_INT4_EXAMPLE) && defined(CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4)
+    static_assert(sizeof(ck::int4_t) == sizeof(int8_t));
+    static_assert(sizeof(ADataType) == sizeof(KernelADataType));
+    static_assert(sizeof(BDataType) == sizeof(KernelBDataType));
+    static_assert(sizeof(EDataType) == sizeof(KernelEDataType));
+#endif
+    int group_count = problem_size.group_count;
+
+    // GEMM shape
+    std::vector<ck::tensor_operation::device::GemmDesc> gemm_descs;
+    std::vector<const void*> p_a, p_b;
+    std::vector<void*> p_c;
+
+    gemm_descs.reserve(group_count);
+
+    for(int i = 0; i < group_count; i++)
+    {
+        int M = problem_size.Ms[i];
+        int N = problem_size.Ns[i];
+        int K = problem_size.Ks[i];
+
+        int stride_A = problem_size.stride_As[i];
+        int stride_B = problem_size.stride_Bs[i];
+        int stride_C = problem_size.stride_Cs[i];
+
+        gemm_descs.push_back({M, N, K, stride_A, stride_B, stride_C, {}});
+    }
+
+    auto f_host_tensor_descriptor =
+        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+            using namespace ck::literals;
+
+            if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+            {
+                return HostTensorDescriptor({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    std::vector<Tensor<ADataType>> a_tensors;
+    std::vector<Tensor<BDataType>> b_tensors;
+    std::vector<Tensor<EDataType>> c_host_tensors;
+#ifdef BUILD_INT4_EXAMPLE
+    std::vector<Tensor<KernelEDataType>> c_device_tensors;
+#else
+    std::vector<Tensor<EDataType>> c_device_tensors;
+#endif
+
+    a_tensors.reserve(group_count);
+    b_tensors.reserve(group_count);
+    c_host_tensors.reserve(group_count);
+    c_device_tensors.reserve(group_count);
+
+    using DeviceMemPtr = std::unique_ptr<DeviceMem>;
+
+    std::vector<DeviceMemPtr> a_tensors_device, b_tensors_device, c_tensors_device;
+
+    a_tensors_device.reserve(group_count);
+    b_tensors_device.reserve(group_count);
+    c_tensors_device.reserve(group_count);
+
+    std::size_t flop = 0, num_btype = 0;
+
+    for(std::size_t i = 0; i < gemm_descs.size(); i++)
+    {
+        a_tensors.push_back(Tensor<ADataType>(f_host_tensor_descriptor(
+            gemm_descs[i].M_, gemm_descs[i].K_, gemm_descs[i].stride_A_, ALayout{})));
+        b_tensors.push_back(Tensor<BDataType>(f_host_tensor_descriptor(
+            gemm_descs[i].K_, gemm_descs[i].N_, gemm_descs[i].stride_B_, BLayout{})));
+        c_host_tensors.push_back(Tensor<EDataType>(f_host_tensor_descriptor(
+            gemm_descs[i].M_, gemm_descs[i].N_, gemm_descs[i].stride_C_, ELayout{})));
+#ifdef BUILD_INT4_EXAMPLE
+        c_device_tensors.push_back(Tensor<KernelEDataType>(f_host_tensor_descriptor(
+            gemm_descs[i].M_, gemm_descs[i].N_, gemm_descs[i].stride_C_, ELayout{})));
+#else
+        c_device_tensors.push_back(Tensor<EDataType>(f_host_tensor_descriptor(
+            gemm_descs[i].M_, gemm_descs[i].N_, gemm_descs[i].stride_C_, ELayout{})));
+#endif
+        std::cout << "gemm[" << i << "] a_m_k: " << a_tensors[i].mDesc
+                  << " b_k_n: " << b_tensors[i].mDesc << " c_m_n: " << c_device_tensors[i].mDesc
+                  << std::endl;
+
+        flop += std::size_t(2) * gemm_descs[i].M_ * gemm_descs[i].K_ * gemm_descs[i].N_;
+        num_btype += sizeof(ADataType) * a_tensors[i].mDesc.GetElementSize() +
+                     sizeof(BDataType) * b_tensors[i].mDesc.GetElementSize() +
+                     sizeof(EDataType) * c_device_tensors[i].mDesc.GetElementSize();
+
+        switch(config.init_method)
+        {
+        case 0: break;
+        case 1:
+            a_tensors[i].GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
+            b_tensors[i].GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
+            break;
+        case 2:
+            a_tensors[i].GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
+            b_tensors[i].GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
+            break;
+        default:
+            a_tensors[i].GenerateTensorValue(GeneratorTensor_Sequential<0>{});
+            b_tensors[i].GenerateTensorValue(GeneratorTensor_Sequential<1>{});
+        }
+    }
+
+    for(std::size_t i = 0; i < gemm_descs.size(); i++)
+    {
+        a_tensors_device.emplace_back(std::make_unique<DeviceMem>(
+            sizeof(ADataType) * a_tensors[i].mDesc.GetElementSpaceSize()));
+        b_tensors_device.emplace_back(std::make_unique<DeviceMem>(
+            sizeof(BDataType) * b_tensors[i].mDesc.GetElementSpaceSize()));
+        c_tensors_device.emplace_back(std::make_unique<DeviceMem>(
+            sizeof(EDataType) * c_device_tensors[i].mDesc.GetElementSpaceSize()));
+
+#ifdef BUILD_INT4_EXAMPLE
+        const Tensor<KernelADataType> a_converted(a_tensors[i]);
+        const Tensor<KernelBDataType> b_converted(b_tensors[i]);
+
+        a_tensors_device[i]->ToDevice(a_converted.mData.data());
+        b_tensors_device[i]->ToDevice(b_converted.mData.data());
+#else
+        a_tensors_device[i]->ToDevice(a_tensors[i].mData.data());
+        b_tensors_device[i]->ToDevice(b_tensors[i].mData.data());
+        c_tensors_device[i]->SetZero();
+#endif
+
+        p_a.push_back(a_tensors_device[i]->GetDeviceBuffer());
+        p_b.push_back(b_tensors_device[i]->GetDeviceBuffer());
+        p_c.push_back(c_tensors_device[i]->GetDeviceBuffer());
+    }
+
+    auto a_element_op = AElementOp{};
+    auto b_element_op = BElementOp{};
+    auto c_element_op = CDEElementOp{};
+
+    auto gemm    = DeviceGemmInstance{};
+    auto invoker = gemm.MakeInvoker();
+
+    std::vector<std::array<const void*, 0>> p_Ds = {};
+
+    // do GEMM
+    auto argument = gemm.MakeArgument(
+        p_a, p_b, p_Ds, p_c, gemm_descs, a_element_op, b_element_op, c_element_op);
+
+    DeviceMem gemm_desc_workspace(gemm.GetWorkSpaceSize(&argument));
+
+    gemm.SetWorkSpacePointer(&argument, gemm_desc_workspace.GetDeviceBuffer());
+
+    if(!gemm.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error(
+            "wrong! device_gemm with the specified compilation parameters does "
+            "not support this GEMM problem");
+    }
+
+    invoker.Run(argument, StreamConfig{nullptr, false});
+
+    bool pass = true;
+    if(config.do_verification)
+    {
+        using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                                BDataType,
+                                                                                EDataType,
+                                                                                AccDataType,
+                                                                                AElementOp,
+                                                                                BElementOp,
+                                                                                CDEElementOp>;
+
+        for(std::size_t i = 0; i < gemm_descs.size(); i++)
+        {
+            c_tensors_device[i]->FromDevice(c_device_tensors[i].mData.data());
+            auto ref_gemm    = ReferenceGemmInstance{};
+            auto ref_invoker = ref_gemm.MakeInvoker();
+
+            auto ref_argument = ref_gemm.MakeArgument(a_tensors[i],
+                                                      b_tensors[i],
+                                                      c_host_tensors[i],
+                                                      a_element_op,
+                                                      b_element_op,
+                                                      c_element_op);
+
+            ref_invoker.Run(ref_argument);
+
+#ifdef BUILD_INT4_EXAMPLE
+            const Tensor<EDataType> c_device_result_converted(c_device_tensors[i]);
+            pass &= ck::utils::check_err(c_device_result_converted, c_host_tensors[i]);
+
+#else
+            pass &= ck::utils::check_err(c_device_tensors[i], c_host_tensors[i]);
+#endif
+        }
+    }
+
+    if(config.time_kernel)
+    {
+        float ave_time   = invoker.Run(argument, StreamConfig{nullptr, config.time_kernel});
+        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;
+    }
+
+    return pass;
+}
+
+bool run_grouped_gemm_example(int argc, char* argv[])
+{
+    ProblemSize problem_size;
+    ExecutionConfig config;
+
+    problem_size.group_count = 16;
+
+    for(int i = 0; i < problem_size.group_count; i++)
+    {
+        problem_size.Ms.push_back(256 + 256 * i);
+        problem_size.Ns.push_back(128 + 128 * i);
+        problem_size.Ks.push_back(128 + 64 * i);
+
+        problem_size.stride_As.push_back(problem_size.Ks[i]);
+        problem_size.stride_Bs.push_back(problem_size.Ks[i]);
+        problem_size.stride_Cs.push_back(problem_size.Ns[i]);
+    }
+
+    if(argc == 4)
+    {
+        config.do_verification = std::stoi(argv[1]);
+        config.init_method     = std::stoi(argv[2]);
+        config.time_kernel     = std::stoi(argv[3]);
+    }
+    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");
+        exit(0);
+    }
+
+    return run_grouped_gemm(problem_size, config);
+}

example/16_gemm_multi_d_multi_reduces/CMakeLists.txt

+list(APPEND gpu_list gfx908 gfx90a gfx940 gfx941 gfx942)
+set(target 0)
+foreach(gpu IN LISTS GPU_TARGETS)
+ if(gpu IN_LIST gpu_list AND target EQUAL 0)
+   add_custom_target(example_gemm_reduce_xdl)
+   add_custom_target(example_gemm_reduce_xdl_max)
+   add_custom_target(example_gemm_reduce_xdl_mean_meansquare)
+   add_custom_target(example_gemm_add_add_mean_meansquare_xdl)
+   if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
+    add_example_executable(example_gemm_max_xdl_fp16 gemm_max_xdl_fp16.cpp)
+    add_example_executable(example_gemm_add_add_mean_meansquare_xdl_fp16 gemm_add_add_mean_meansquare_xdl_fp16.cpp)
+    add_example_executable(example_gemm_mean_meansquare_xdl_fp16 gemm_mean_meansquare_xdl_fp16.cpp)
+    add_dependencies(example_gemm_reduce_xdl_max example_gemm_max_xdl_fp16)
+    add_dependencies(example_gemm_add_add_mean_meansquare_xdl example_gemm_add_add_mean_meansquare_xdl_fp16)
+    add_dependencies(example_gemm_reduce_xdl_mean_meansquare example_gemm_mean_meansquare_xdl_fp16)
+   endif()
+   if(DTYPES MATCHES "int8" OR NOT DEFINED DTYPES)
+    add_example_executable(example_gemm_max_xdl_int8 gemm_max_xdl_int8.cpp)
+    add_example_executable(example_gemm_add_addsquare_xdl_int8 gemm_add_addsquare_xdl_int8.cpp)
+    add_dependencies(example_gemm_reduce_xdl_max example_gemm_max_xdl_int8)
+    add_dependencies(example_gemm_reduce_xdl_mean_meansquare example_gemm_add_addsquare_xdl_int8)
+   endif()
+   if(DTYPES MATCHES "fp32" OR NOT DEFINED DTYPES)
+    add_example_executable(example_gemm_max_xdl_fp32 gemm_max_xdl_fp32.cpp)
+    add_example_executable(example_gemm_mean_meansquare_xdl_fp32 gemm_mean_meansquare_xdl_fp32.cpp)
+    add_dependencies(example_gemm_reduce_xdl_max example_gemm_max_xdl_fp32)
+    add_dependencies(example_gemm_reduce_xdl_mean_meansquare example_gemm_mean_meansquare_xdl_fp32)
+   endif()
+   if(DTYPES MATCHES "bf16" OR NOT DEFINED DTYPES)
+    add_example_executable(example_gemm_max_xdl_bf16 gemm_max_xdl_bf16.cpp)
+    add_example_executable(example_gemm_mean_meansquare_xdl_bf16 gemm_mean_meansquare_xdl_bf16.cpp)
+    add_dependencies(example_gemm_reduce_xdl_max example_gemm_max_xdl_bf16)
+    add_dependencies(example_gemm_reduce_xdl_mean_meansquare example_gemm_mean_meansquare_xdl_bf16)
+   endif()
+   
+   add_dependencies(example_gemm_reduce_xdl
+                 example_gemm_reduce_xdl_mean_meansquare
+                 example_gemm_reduce_xdl_max
+                 example_gemm_add_add_mean_meansquare_xdl)
+
+   if(USE_BITINT_EXTENSION_INT4)
+      add_example_executable(example_gemm_max_xdl_int4 gemm_max_xdl_int4.cpp)
+      add_dependencies(example_gemm_reduce_xdl_max example_gemm_max_xdl_int4)
+   endif()
+   set(target 1)
+ endif()
+endforeach()

example/16_gemm_multi_d_multi_reduces/gemm_add_add_mean_meansquare_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, 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/impl/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/utility/literals.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({len}, {stride});
+};
+
+auto f_host_tensor_descriptor2d =
+    [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+        using namespace ck::literals;
+
+        if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+        {
+            return HostTensorDescriptor({row, col}, {stride, 1_uz});
+        }
+        else
+        {
+            return HostTensorDescriptor({row, col}, {1_uz, 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, e_m_n_host, "Error: Incorrect results c", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(r0_m, r0_m_host, "Error: Incorrect results d0", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(r1_m, r1_m_host, "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_add_addsquare_xdl_int8.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/library/utility/literals.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+// DataType
+using ADataType         = INT8;
+using BDataType         = INT8;
+using GemmAccDataType   = INT32;
+using CShuffleDataType  = INT32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = INT8;
+using ReduceAccDataType = INT32;
+using R0DataType        = INT32;
+using R1DataType        = INT32;
+using RsDataType        = ck::Tuple<R0DataType, R1DataType>;
+
+// Layout
+using ALayout = Row;
+using BLayout = Col;
+using ELayout = Row;
+
+// Elementwise op
+using Square       = ck::tensor_operation::element_wise::UnarySquare;
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = PassThrough;
+using QsElementOp  = ck::Tuple<PassThrough, Square>;
+using RsElementOp  = ck::Tuple<PassThrough, PassThrough>;
+
+// 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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         64,                        // KPerBlock
+         16,                        // AK1
+         16,                        // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         16,                        // ABlockTransfer SrcScalarPerVector
+         16,                        // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         16,                        // BBlockTransfer SrcScalarPerVector
+         16,                        // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+using namespace ck::literals;
+
+template <typename ADataType,
+          typename BDataType,
+          typename EDataType,
+          typename R0DataType,
+          typename R1DataType,
+          typename ALayout,
+          typename BLayout,
+          typename ELayout,
+          typename AElementOp,
+          typename BElementOp,
+          typename CDEElementOp,
+          typename QsElementOp,
+          typename RsElementOp,
+          typename RsThreadReduceOp,
+          typename ReduceAccDataType,
+          typename DeviceOpInstance,
+          typename ReferenceGemmInstance>
+bool run_gemm_reduce_add_addsquare_xdl(ck::index_t M,
+                                       ck::index_t N,
+                                       ck::index_t K,
+                                       ck::index_t StrideA,
+                                       ck::index_t StrideB,
+                                       ck::index_t StrideE,
+                                       bool do_verification,
+                                       int init_method,
+                                       bool time_kernel)
+{
+
+    auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
+        return HostTensorDescriptor({len}, {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({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    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));
+
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        ck::utils::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_m_k);
+        ck::utils::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_k_n);
+        break;
+    default:
+        ck::utils::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
+        ck::utils::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
+        break;
+    }
+
+    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{};
+
+    // Prepare GEMM, add, add_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 pass = true;
+
+    if(do_verification)
+    {
+        auto I0 = ck::Number<0>{};
+        auto I1 = ck::Number<1>{};
+
+        Tensor<ReduceAccDataType> 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 = RsThreadReduceOp{}[I0];
+        auto reduce1_op = RsThreadReduceOp{}[I1];
+
+        for(int m = 0; m < M; ++m)
+        {
+            auto reduce0_acc = reduce0_op.template GetIdentityValue<ReduceAccDataType>();
+            auto reduce1_acc = reduce1_op.template 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);
+            }
+
+            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());
+
+        Tensor<EDataType> e_m_n_host_converted(e_m_n_host);
+
+        pass = ck::utils::check_err(
+            e_m_n, e_m_n_host_converted, "Error: Incorrect results c", 1e-2, 1e-2);
+
+        r0_device_buf.FromDevice(r0_m.mData.data());
+        r1_device_buf.FromDevice(r1_m.mData.data());
+
+        pass &= ck::utils::check_err(r0_m, r0_m_host, "Error: Incorrect results d0", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(r1_m, r1_m_host, "Error: Incorrect results d1", 1e-2, 1e-2);
+
+        if(pass)
+        {
+            std::cout << "Success!" << std::endl;
+        }
+    }
+
+    if(time_kernel)
+    {
+        float ave_time            = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+        std::size_t flop          = 2_uz * M * N * K + 3_uz * M * N;
+        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;
+    }
+
+    return pass;
+}
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    ck::index_t M = 1024;
+    ck::index_t N = 1152;
+    ck::index_t K = 512;
+
+    ck::index_t StrideA = 512;
+    ck::index_t StrideB = 512;
+    ck::index_t StrideE = 1152;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return !run_gemm_reduce_add_addsquare_xdl<ADataType,
+                                              BDataType,
+                                              EDataType,
+                                              R0DataType,
+                                              R1DataType,
+                                              ALayout,
+                                              BLayout,
+                                              ELayout,
+                                              AElementOp,
+                                              BElementOp,
+                                              CDEElementOp,
+                                              QsElementOp,
+                                              RsElementOp,
+                                              RsThreadReduceOp,
+                                              ReduceAccDataType,
+                                              DeviceOpInstance,
+                                              ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_max_xdl_bf16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+// DataType
+using ADataType         = BF16;
+using BDataType         = BF16;
+using GemmAccDataType   = F32;
+using CShuffleDataType  = F32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = BF16;
+using ReduceAccDataType = F32;
+using R0DataType        = F32;
+using RsDataType        = ck::Tuple<R0DataType>;
+
+// Layout
+using ALayout = Row;
+using BLayout = Col;
+using ELayout = Row;
+
+// Elementwise op
+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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         32,                        // KPerBlock
+         8,                         // AK1
+         8,                         // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         8,                         // ABlockTransfer SrcScalarPerVector
+         8,                         // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         8,                         // BBlockTransfer SrcScalarPerVector
+         8,                         // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    ck::index_t M = 1024;
+    ck::index_t N = 1152;
+    ck::index_t K = 256;
+
+    ck::index_t StrideA = 256;
+    ck::index_t StrideB = 256;
+    ck::index_t StrideE = 1152;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return run_gemm_reduce_max_xdl<ADataType,
+                                   BDataType,
+                                   EDataType,
+                                   R0DataType,
+                                   ALayout,
+                                   BLayout,
+                                   ELayout,
+                                   AElementOp,
+                                   BElementOp,
+                                   CDEElementOp,
+                                   QsElementOp,
+                                   RsElementOp,
+                                   RsThreadReduceOp,
+                                   ReduceAccDataType,
+                                   DeviceOpInstance,
+                                   ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_max_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+// 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 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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         32,                        // KPerBlock
+         8,                         // AK1
+         8,                         // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         8,                         // ABlockTransfer SrcScalarPerVector
+         8,                         // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         8,                         // BBlockTransfer SrcScalarPerVector
+         8,                         // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    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;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return run_gemm_reduce_max_xdl<ADataType,
+                                   BDataType,
+                                   EDataType,
+                                   R0DataType,
+                                   ALayout,
+                                   BLayout,
+                                   ELayout,
+                                   AElementOp,
+                                   BElementOp,
+                                   CDEElementOp,
+                                   QsElementOp,
+                                   RsElementOp,
+                                   RsThreadReduceOp,
+                                   ReduceAccDataType,
+                                   DeviceOpInstance,
+                                   ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_max_xdl_fp32.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+// DataType
+using ADataType         = F32;
+using BDataType         = F32;
+using GemmAccDataType   = F32;
+using CShuffleDataType  = F32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = F32;
+using ReduceAccDataType = F32;
+using R0DataType        = F32;
+using RsDataType        = ck::Tuple<R0DataType>;
+
+// Layout
+using ALayout = Row;
+using BLayout = Col;
+using ELayout = Row;
+
+// Elementwise op
+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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         16,                        // KPerBlock
+         4,                         // AK1
+         4,                         // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         4,                         // ABlockTransfer SrcScalarPerVector
+         4,                         // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         4,                         // BBlockTransfer SrcScalarPerVector
+         4,                         // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    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;
+
+    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]);
+        StrideE = 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: Measure kernel execution time (1=ON, 0=Off)\n");
+        printf("arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n");
+        exit(0);
+    }
+
+    return run_gemm_reduce_max_xdl<ADataType,
+                                   BDataType,
+                                   EDataType,
+                                   R0DataType,
+                                   ALayout,
+                                   BLayout,
+                                   ELayout,
+                                   AElementOp,
+                                   BElementOp,
+                                   CDEElementOp,
+                                   QsElementOp,
+                                   RsElementOp,
+                                   RsThreadReduceOp,
+                                   ReduceAccDataType,
+                                   DeviceOpInstance,
+                                   ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_max_xdl_int4.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+using ADataType         = INT4;
+using ADataKernelType   = INT8;
+using BDataType         = INT4;
+using BDataKernelType   = INT8;
+using GemmAccDataType   = INT32;
+using CShuffleDataType  = INT32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = INT4;
+using EDataKernelType   = INT8;
+using ReduceAccDataType = INT32;
+using R0DataType        = INT32;
+using RsDataType        = ck::Tuple<R0DataType>;
+
+// Layout
+using ALayout = Row;
+using BLayout = Col;
+using ELayout = Row;
+
+// Elementwise op
+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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataKernelType,           // ADataType
+         BDataKernelType,           // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataKernelType,           // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         64,                        // KPerBlock
+         16,                        // AK1
+         16,                        // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         16,                        // ABlockTransfer SrcScalarPerVector
+         16,                        // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         16,                        // BBlockTransfer SrcScalarPerVector
+         16,                        // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    ck::index_t M = 1024;
+    ck::index_t N = 1152;
+    ck::index_t K = 256;
+
+    ck::index_t StrideA = 256;
+    ck::index_t StrideB = 256;
+    ck::index_t StrideE = 1152;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return run_gemm_reduce_max_xdl<ADataType,
+                                   BDataType,
+                                   EDataType,
+                                   R0DataType,
+                                   ALayout,
+                                   BLayout,
+                                   ELayout,
+                                   AElementOp,
+                                   BElementOp,
+                                   CDEElementOp,
+                                   QsElementOp,
+                                   RsElementOp,
+                                   RsThreadReduceOp,
+                                   ReduceAccDataType,
+                                   DeviceOpInstance,
+                                   ReferenceGemmInstance,
+                                   ADataKernelType,
+                                   BDataKernelType,
+                                   EDataKernelType>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_max_xdl_int8.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+using ADataType         = INT8;
+using BDataType         = INT8;
+using GemmAccDataType   = INT32;
+using CShuffleDataType  = INT32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = INT8;
+using ReduceAccDataType = INT32;
+using R0DataType        = INT32;
+using RsDataType        = ck::Tuple<R0DataType>;
+
+// Layout
+using ALayout = Row;
+using BLayout = Col;
+using ELayout = Row;
+
+// Elementwise op
+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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         64,                        // KPerBlock
+         16,                        // AK1
+         16,                        // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         16,                        // ABlockTransfer SrcScalarPerVector
+         16,                        // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         16,                        // BBlockTransfer SrcScalarPerVector
+         16,                        // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    ck::index_t M = 1024;
+    ck::index_t N = 1152;
+    ck::index_t K = 512;
+
+    ck::index_t StrideA = 512;
+    ck::index_t StrideB = 512;
+    ck::index_t StrideE = 1152;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return run_gemm_reduce_max_xdl<ADataType,
+                                   BDataType,
+                                   EDataType,
+                                   R0DataType,
+                                   ALayout,
+                                   BLayout,
+                                   ELayout,
+                                   AElementOp,
+                                   BElementOp,
+                                   CDEElementOp,
+                                   QsElementOp,
+                                   RsElementOp,
+                                   RsThreadReduceOp,
+                                   ReduceAccDataType,
+                                   DeviceOpInstance,
+                                   ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_mean_meansquare_xdl_bf16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+// DataType
+using ADataType         = BF16;
+using BDataType         = BF16;
+using GemmAccDataType   = F32;
+using CShuffleDataType  = F32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = BF16;
+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 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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         32,                        // KPerBlock
+         8,                         // AK1
+         8,                         // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         8,                         // ABlockTransfer SrcScalarPerVector
+         8,                         // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         8,                         // BBlockTransfer SrcScalarPerVector
+         8,                         // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    ck::index_t M = 1024;
+    ck::index_t N = 1152;
+    ck::index_t K = 192;
+
+    ck::index_t StrideA = 192;
+    ck::index_t StrideB = 192;
+    ck::index_t StrideE = 1152;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return !run_gemm_reduce_mean_meansquare_xdl<ADataType,
+                                                BDataType,
+                                                EDataType,
+                                                R0DataType,
+                                                R1DataType,
+                                                ALayout,
+                                                BLayout,
+                                                ELayout,
+                                                AElementOp,
+                                                BElementOp,
+                                                CDEElementOp,
+                                                QsElementOp,
+                                                RsElementOp,
+                                                RsThreadReduceOp,
+                                                ReduceAccDataType,
+                                                DeviceOpInstance,
+                                                ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_mean_meansquare_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+// 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 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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         32,                        // KPerBlock
+         8,                         // AK1
+         8,                         // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         8,                         // ABlockTransfer SrcScalarPerVector
+         8,                         // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         8,                         // BBlockTransfer SrcScalarPerVector
+         8,                         // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    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;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return !run_gemm_reduce_mean_meansquare_xdl<ADataType,
+                                                BDataType,
+                                                EDataType,
+                                                R0DataType,
+                                                R1DataType,
+                                                ALayout,
+                                                BLayout,
+                                                ELayout,
+                                                AElementOp,
+                                                BElementOp,
+                                                CDEElementOp,
+                                                QsElementOp,
+                                                RsElementOp,
+                                                RsThreadReduceOp,
+                                                ReduceAccDataType,
+                                                DeviceOpInstance,
+                                                ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_mean_meansquare_xdl_fp32.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gemm_reduce_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+// DataType
+using ADataType         = F32;
+using BDataType         = F32;
+using GemmAccDataType   = F32;
+using CShuffleDataType  = F32;
+using DsDataType        = ck::Tuple<>;
+using EDataType         = F32;
+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 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,                   // ALayout
+         BLayout,                   // BLayout
+         ELayout,                   // ELayout
+         ADataType,                 // ADataType
+         BDataType,                 // BDataType
+         GemmAccDataType,           // GemmAccDataType
+         CShuffleDataType,          // CShuffleDataType
+         DsDataType,                // DsDataType
+         EDataType,                 // EDataType
+         ReduceAccDataType,         // ReduceAccDataType
+         RsDataType,                // RsDataType
+         AElementOp,                // AElementwiseOperation
+         BElementOp,                // BElementwiseOperation
+         CDEElementOp,              // CDE ElementwiseOperation
+         QsElementOp,               // Qs Elementwise Operation
+         RsElementOp,               // Rs Elementwise Operation
+         RsThreadReduceOp,          // Thread Reduce Operation
+         RsGlobalReduceOp,          // Global Reduce Operation
+         GemmDefault,               // GEMM Specialization
+         1,                         // NumGemmKPrefetchStage
+         256,                       // BlockSize
+         256,                       // MPerBlock
+         128,                       // NPerBlock
+         16,                        // KPerBlock
+         4,                         // AK1
+         4,                         // BK1
+         32,                        // MPerXdl
+         32,                        // NPerXdl
+         4,                         // MXdlPerWave
+         2,                         // NXdlPerWave
+         S<4, 64, 1>,               // ABlockTransfer ThreadCluster Lengths_K0_M_K1
+         S<1, 0, 2>,                // ABlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // ABlockTransfer SrcAccessOrder
+         2,                         // ABlockTransfer SrcVectorDim
+         4,                         // ABlockTransfer SrcScalarPerVector
+         4,                         // ABlockTransfer DstScalarPerVector_K1
+         1,                         // ABlockLdsExtraM
+         S<4, 64, 1>,               // BBlockTransfer ThreadCluster Lengths_K0_N_K1
+         S<1, 0, 2>,                // BBlockTransfer ThreadCluster ArrangeOrder
+         S<1, 0, 2>,                // BBlockTransfer SrcAccessOrder
+         2,                         // BBlockTransfer SrcVectorDim
+         4,                         // BBlockTransfer SrcScalarPerVector
+         4,                         // BBlockTransfer DstScalarPerVector_K1
+         1,                         // BBlockLdsExtraN
+         1,                         // CShuffleMXdlPerWavePerShuffle
+         1,                         // CShuffleNXdlPerWavePerShuffle
+         S<64, 4>,                  // CD Reduce Thread Transfer ClusterLengths _MPerBlock_NPerBlock
+         4,                         // CDE ReduceThreadTransfer ScalarPerVector _NPerBlock
+         1>;                        // RThread DstScalarPerVector _MPerBlock
+// clang-format on
+using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                        BDataType,
+                                                                        ReduceAccDataType,
+                                                                        GemmAccDataType,
+                                                                        AElementOp,
+                                                                        BElementOp,
+                                                                        CDEElementOp>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // GEMM shape
+    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;
+
+    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]);
+        StrideE = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << " arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << " arg3: Measure kernel execution time (1=ON, 0=Off)\n"
+                  << " arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideE\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return !run_gemm_reduce_mean_meansquare_xdl<ADataType,
+                                                BDataType,
+                                                EDataType,
+                                                R0DataType,
+                                                R1DataType,
+                                                ALayout,
+                                                BLayout,
+                                                ELayout,
+                                                AElementOp,
+                                                BElementOp,
+                                                CDEElementOp,
+                                                QsElementOp,
+                                                RsElementOp,
+                                                RsThreadReduceOp,
+                                                ReduceAccDataType,
+                                                DeviceOpInstance,
+                                                ReferenceGemmInstance>(
+        M, N, K, StrideA, StrideB, StrideE, do_verification, init_method, time_kernel);
+}

example/16_gemm_multi_d_multi_reduces/gemm_reduce_xdl_common.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+
+#include <iostream>
+
+#include "ck/ck.hpp"
+#include "ck/host_utility/io.hpp"
+#include "ck/stream_config.hpp"
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/fill.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/literals.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+using F16         = ck::half_t;
+using BF16        = ck::bhalf_t;
+using F32         = float;
+using F64         = double;
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+using INT4 = ck::int4_t;
+#endif
+using INT8  = std::int8_t;
+using INT32 = std::int32_t;
+
+template <typename ADataType, typename BDataType, typename EDataType, typename R0DataType>
+void DumpGemmReduceMaxPerf(float ave_time, int M, int N, int K)
+{
+    using namespace ck::literals;
+
+    std::size_t flop          = 2_uz * 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;
+}
+
+template <typename ADataType,
+          typename BDataType,
+          typename EDataType,
+          typename R0DataType,
+          typename R1DataType>
+void DumpGemmReduceMeanSquareMeanPerf(float ave_time, int M, int N, int K)
+{
+    using namespace ck::literals;
+
+    std::size_t flop          = 2_uz * M * N * K + M * (3_uz * N + 2_uz);
+    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;
+}
+
+template <typename ADataType,
+          typename BDataType,
+          typename EDataType,
+          typename R0DataType,
+          typename ALayout,
+          typename BLayout,
+          typename ELayout,
+          typename AElementOp,
+          typename BElementOp,
+          typename CDEElementOp,
+          typename QsElementOp,
+          typename RsElementOp,
+          typename RsThreadReduceOp,
+          typename ReduceAccDataType,
+          typename DeviceOpInstance,
+          typename ReferenceGemmInstance,
+          typename ADataKernelType = ADataType,
+          typename BDataKernelType = BDataType,
+          typename EDataKernelType = EDataType>
+auto run_gemm_reduce_max_xdl(ck::index_t M,
+                             ck::index_t N,
+                             ck::index_t K,
+                             ck::index_t StrideA,
+                             ck::index_t StrideB,
+                             ck::index_t StrideE,
+                             bool do_verification,
+                             int init_method,
+                             bool time_kernel)
+{
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    static_assert(sizeof(ck::int4_t) == sizeof(int8_t));
+    static_assert(sizeof(ADataType) == sizeof(ADataKernelType));
+    static_assert(sizeof(BDataType) == sizeof(BDataKernelType));
+    static_assert(sizeof(EDataType) == sizeof(EDataKernelType));
+#endif
+    using namespace ck::literals;
+
+    auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
+        return HostTensorDescriptor({len}, {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({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    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<EDataKernelType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
+    Tensor<R0DataType> r0_m(f_host_tensor_descriptor1d(M, 1));
+
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        ck::utils::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_m_k);
+        ck::utils::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_k_n);
+        break;
+    default:
+        ck::utils::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
+        ck::utils::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
+        break;
+    }
+
+    DeviceMem a_device_buf(sizeof(ADataKernelType) * a_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem b_device_buf(sizeof(BDataKernelType) * b_k_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataKernelType) * e_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem r0_device_buf(sizeof(R0DataType) * r0_m.mDesc.GetElementSpaceSize());
+
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    if constexpr(std::is_same_v<ADataType, ck::int4_t>)
+    {
+        Tensor<ADataKernelType> a_m_k_converted = a_m_k.template CopyAsType<ADataKernelType>();
+        Tensor<BDataKernelType> b_k_n_converted = b_k_n.template CopyAsType<BDataKernelType>();
+
+        a_device_buf.ToDevice(a_m_k_converted.mData.data());
+        b_device_buf.ToDevice(b_k_n_converted.mData.data());
+    }
+    else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    {
+        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");
+    }
+
+    // [CAUTION]: 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 pass = true;
+
+    if(do_verification)
+    {
+        auto I0 = ck::Number<0>{};
+
+        Tensor<ReduceAccDataType> 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.template GetIdentityValue<ReduceAccDataType>();
+
+            for(int n = 0; n < N; ++n)
+            {
+                auto e_val = 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());
+        Tensor<EDataType> e_m_n_host_converted(e_m_n_host);
+
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        if constexpr(std::is_same_v<ADataType, ck::int4_t>)
+        {
+            Tensor<EDataType> e_m_n_device_converted(e_m_n);
+            pass = ck::utils::check_err(e_m_n_device_converted,
+                                        e_m_n_host_converted,
+                                        "Error: Incorrect results c",
+                                        1e-2,
+                                        1e-2);
+        }
+        else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        {
+            pass = ck::utils::check_err(
+                e_m_n, e_m_n_host_converted, "Error: Incorrect results c", 1e-2, 1e-2);
+        }
+
+        r0_device_buf.FromDevice(r0_m.mData.data());
+        pass &= ck::utils::check_err(r0_m, r0_m_host, "Error: Incorrect results d0", 1e-2, 1e-2);
+
+        if(pass)
+        {
+            std::cout << "Success!" << std::endl;
+        }
+    }
+
+    if(time_kernel)
+    {
+        float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+        DumpGemmReduceMaxPerf<ADataType, BDataType, EDataType, R0DataType>(ave_time, M, N, K);
+    }
+
+    return pass ? 0 : 1;
+}
+
+template <typename ADataType,
+          typename BDataType,
+          typename EDataType,
+          typename R0DataType,
+          typename R1DataType,
+          typename ALayout,
+          typename BLayout,
+          typename ELayout,
+          typename AElementOp,
+          typename BElementOp,
+          typename CDEElementOp,
+          typename QsElementOp,
+          typename RsElementOp,
+          typename RsThreadReduceOp,
+          typename ReduceAccDataType,
+          typename DeviceOpInstance,
+          typename ReferenceGemmInstance,
+          typename ADataKernelType = ADataType,
+          typename BDataKernelType = BDataType,
+          typename EDataKernelType = EDataType>
+bool run_gemm_reduce_mean_meansquare_xdl(ck::index_t M,
+                                         ck::index_t N,
+                                         ck::index_t K,
+                                         ck::index_t StrideA,
+                                         ck::index_t StrideB,
+                                         ck::index_t StrideE,
+                                         bool do_verification,
+                                         int init_method,
+                                         bool time_kernel)
+{
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    static_assert(sizeof(ck::int4_t) == sizeof(int8_t));
+    static_assert(sizeof(ADataType) == sizeof(ADataKernelType));
+    static_assert(sizeof(BDataType) == sizeof(BDataKernelType));
+    static_assert(sizeof(EDataType) == sizeof(EDataKernelType));
+#endif
+    using namespace ck::literals;
+
+    auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
+        return HostTensorDescriptor({len}, {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({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    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<EDataKernelType> 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));
+
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        ck::utils::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_m_k);
+        ck::utils::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_k_n);
+        break;
+    default:
+        ck::utils::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
+        ck::utils::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);
+        break;
+    }
+
+    DeviceMem a_device_buf(sizeof(ADataKernelType) * a_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem b_device_buf(sizeof(BDataKernelType) * b_k_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataKernelType) * 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());
+
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    if constexpr(std::is_same_v<ADataType, ck::int4_t>)
+    {
+        Tensor<ADataKernelType> a_m_k_converted = a_m_k.template CopyAsType<ADataKernelType>();
+        Tensor<BDataKernelType> b_k_n_converted = b_k_n.template CopyAsType<BDataKernelType>();
+
+        a_device_buf.ToDevice(a_m_k_converted.mData.data());
+        b_device_buf.ToDevice(b_k_n_converted.mData.data());
+    }
+    else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    {
+        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 pass = true;
+
+    if(do_verification)
+    {
+        auto I0 = ck::Number<0>{};
+        auto I1 = ck::Number<1>{};
+
+        Tensor<ReduceAccDataType> 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 = RsThreadReduceOp{}[I0];
+        auto reduce1_op = RsThreadReduceOp{}[I1];
+
+        for(int m = 0; m < M; ++m)
+        {
+            auto reduce0_acc = reduce0_op.template GetIdentityValue<ReduceAccDataType>();
+            auto reduce1_acc = reduce1_op.template 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());
+        Tensor<EDataType> e_m_n_host_converted(e_m_n_host);
+
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        if constexpr(std::is_same_v<ADataType, ck::int4_t>)
+        {
+            Tensor<EDataType> e_m_n_device_converted(e_m_n);
+            pass = ck::utils::check_err(e_m_n_device_converted,
+                                        e_m_n_host_converted,
+                                        "Error: Incorrect results c",
+                                        1e-2,
+                                        1e-2);
+        }
+        else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        {
+            pass = ck::utils::check_err(
+                e_m_n, e_m_n_host_converted, "Error: Incorrect results c", 1e-2, 1e-2);
+        }
+
+        r0_device_buf.FromDevice(r0_m.mData.data());
+        r1_device_buf.FromDevice(r1_m.mData.data());
+
+        pass &= ck::utils::check_err(r0_m, r0_m_host, "Error: Incorrect results d0", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(r1_m, r1_m_host, "Error: Incorrect results d1", 1e-2, 1e-2);
+
+        if(pass)
+        {
+            std::cout << "Success!" << std::endl;
+        }
+    }
+
+    if(time_kernel)
+    {
+        float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+        DumpGemmReduceMeanSquareMeanPerf<ADataType, BDataType, EDataType, R0DataType, R1DataType>(
+            ave_time, M, N, K);
+    }
+
+    return pass;
+}

example/17_convnd_bwd_data/CMakeLists.txt

+if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
+list(APPEND gpu_list gfx908 gfx90a gfx940 gfx941 gfx942)
+set(target 0)
+foreach(gpu IN LISTS GPU_TARGETS)
+ if(gpu IN_LIST gpu_list AND target EQUAL 0)
+   add_example_executable(example_convnd_bwd_data_xdl_fp16 convnd_bwd_data_xdl_fp16.cpp)
+   target_link_libraries(example_convnd_bwd_data_xdl_fp16 PRIVATE utility)
+   set(target 1)
+ endif()
+endforeach()
+  if(DL_KERNELS)
+    add_example_executable(example_convnd_bwd_data_dl_fp16 convnd_bwd_data_dl_fp16.cpp)
+    target_link_libraries(example_convnd_bwd_data_dl_fp16 PRIVATE utility)
+  endif()
+endif()