Merge branch 'develop' into wavelet_model

example/16_gemm_multi_d_multi_reduces/gemm_max_xdl_fp32.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, 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/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-2022, 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/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-2022, 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/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-2022, 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/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-2022, Advanced Micro Devices, Inc. All rights reserved.
 
-#include <iostream>
-#include <numeric>
-#include <initializer_list>
-#include <cstdlib>
+#include "gemm_reduce_xdl_common.hpp"
 
-#include "ck/ck.hpp"
-#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
-#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
-#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
-#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
-
-#include "ck/library/utility/device_memory.hpp"
-#include "ck/library/utility/host_tensor.hpp"
-#include "ck/library/utility/host_tensor_generator.hpp"
 #include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
-#include "ck/library/utility/check_err.hpp"
-
-template <ck::index_t... Is>
-using S = ck::Sequence<Is...>;
-
-using F16 = ck::half_t;
-using F32 = float;
-
-using Row = ck::tensor_layout::gemm::RowMajor;
-using Col = ck::tensor_layout::gemm::ColumnMajor;
+#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 
 // DataType
 using ADataType         = F16;
@@ -45,7 +25,6 @@ using BLayout = Col;
 using ELayout = Row;
 
 // Elementwise op
-using PassThrough  = ck::tensor_operation::element_wise::PassThrough;
 using Square       = ck::tensor_operation::element_wise::UnarySquare;
 using Div          = ck::tensor_operation::element_wise::UnaryDivide;
 using AElementOp   = PassThrough;
@@ -67,61 +46,71 @@ static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecializa
 
 // 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>;
+        <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,
-                                                                        EDataType,
+                                                                        ReduceAccDataType,
                                                                         GemmAccDataType,
                                                                         AElementOp,
                                                                         BElementOp,
                                                                         CDEElementOp>;
 
-template <typename ADataType,
-          typename BDataType,
-          typename EDataType,
-          typename R0DataType,
-          typename R1DataType>
-void DumpPerf(float ave_time, int M, int N, int K)
+int main(int argc, char* argv[])
 {
-    std::size_t flop          = std::size_t(2) * M * N * K;
-    std::size_t gemm_num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
-                                sizeof(EDataType) * M * N + sizeof(R0DataType) * M +
-                                sizeof(R1DataType) * M;
-
-    float tflops          = static_cast<float>(flop) / 1.E9 / ave_time;
-    float gemm_gb_per_sec = gemm_num_byte / 1.E6 / ave_time;
-
-    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gemm_gb_per_sec
-              << " GB/s, " << std::endl;
-}
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
 
-auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
-    return HostTensorDescriptor(std::vector<std::size_t>({len}),
-                                std::vector<std::size_t>({stride}));
-};
-
-auto f_host_tensor_descriptor2d =
-    [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
-        if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
-        {
-            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
-                                        std::vector<std::size_t>({stride, 1}));
-        }
-        else
-        {
-            return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
-                                        std::vector<std::size_t>({1, stride}));
-        }
-    };
-
-int main()
-{
+    // GEMM shape
     ck::index_t M = 1024;
     ck::index_t N = 1024;
     ck::index_t K = 1024;
@@ -130,125 +119,56 @@ int main()
     ck::index_t StrideB = 1024;
     ck::index_t StrideE = 1024;
 
-    Tensor<ADataType> a_m_k(f_host_tensor_descriptor2d(M, K, StrideA, ALayout{}));
-    Tensor<BDataType> b_k_n(f_host_tensor_descriptor2d(K, N, StrideB, BLayout{}));
-    Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
-    Tensor<R0DataType> r0_m(f_host_tensor_descriptor1d(M, 1));
-    Tensor<R1DataType> r1_m(f_host_tensor_descriptor1d(M, 1));
-
-    a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{-1, 1});
-    b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-1, 1});
-
-    DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
-    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
-    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n.mDesc.GetElementSpaceSize());
-    DeviceMem r0_device_buf(sizeof(R0DataType) * r0_m.mDesc.GetElementSpaceSize());
-    DeviceMem r1_device_buf(sizeof(R1DataType) * r1_m.mDesc.GetElementSpaceSize());
-
-    a_device_buf.ToDevice(a_m_k.mData.data());
-    b_device_buf.ToDevice(b_k_n.mData.data());
-
-    auto a_element_op   = AElementOp{};
-    auto b_element_op   = BElementOp{};
-    auto cde_element_op = CDEElementOp{};
-    auto qs_element_op  = QsElementOp{};
-    auto rs_element_op  = RsElementOp{N, N};
-
-    // Prepare GEMM, mean, mean_square
-    auto device_op = DeviceOpInstance{};
-    auto invoker   = device_op.MakeInvoker();
-    auto argument =
-        device_op.MakeArgument(a_device_buf.GetDeviceBuffer(),
-                               b_device_buf.GetDeviceBuffer(),
-                               {},
-                               e_device_buf.GetDeviceBuffer(),
-                               {r0_device_buf.GetDeviceBuffer(), r1_device_buf.GetDeviceBuffer()},
-                               M,
-                               N,
-                               K,
-                               StrideA,
-                               StrideB,
-                               {},
-                               StrideE,
-                               a_element_op,
-                               b_element_op,
-                               cde_element_op,
-                               qs_element_op,
-                               rs_element_op);
-
-    if(!device_op.IsSupportedArgument(argument))
+    if(argc == 1)
     {
-        throw std::runtime_error("wrong! this device_op instance does not support this problem");
+        // do nothing
     }
-
-    // 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)
+    else if(argc == 4)
     {
-        auto I0 = ck::Number<0>{};
-        auto I1 = ck::Number<1>{};
-
-        Tensor<EDataType> e_m_n_host(e_m_n.mDesc);
-        Tensor<R0DataType> r0_m_host(r0_m.mDesc);
-        Tensor<R1DataType> r1_m_host(r1_m.mDesc);
-
-        auto ref_gemm    = ReferenceGemmInstance{};
-        auto ref_invoker = ref_gemm.MakeInvoker();
-
-        auto ref_argument = ref_gemm.MakeArgument(
-            a_m_k, b_k_n, e_m_n_host, a_element_op, b_element_op, PassThrough{});
-
-        ref_invoker.Run(ref_argument);
-
-        auto reduce0_op = R0ThreadReduceOp{};
-        auto reduce1_op = R1ThreadReduceOp{};
-
-        for(int m = 0; m < M; ++m)
-        {
-            auto reduce0_acc = reduce0_op.GetIdentityValue<ReduceAccDataType>();
-            auto reduce1_acc = reduce1_op.GetIdentityValue<ReduceAccDataType>();
-
-            for(int n = 0; n < N; ++n)
-            {
-                ReduceAccDataType square_e_val;
-                auto e_val = ck::type_convert<ReduceAccDataType>(e_m_n_host(m, n));
-                qs_element_op[I1](square_e_val, e_val);
-
-                reduce0_op(reduce0_acc, e_val);
-                reduce1_op(reduce1_acc, square_e_val);
-            }
-
-            rs_element_op[I0](reduce0_acc, reduce0_acc);
-            rs_element_op[I1](reduce1_acc, reduce1_acc);
-            r0_m_host(m) = ck::type_convert<R0DataType>(reduce0_acc);
-            r1_m_host(m) = ck::type_convert<R1DataType>(reduce1_acc);
-        }
-
-        e_device_buf.FromDevice(e_m_n.mData.data());
-        r0_device_buf.FromDevice(r0_m.mData.data());
-        r1_device_buf.FromDevice(r1_m.mData.data());
-
-        pass = ck::utils::check_err(
-            e_m_n.mData, e_m_n_host.mData, "Error: Incorrect results c", 1e-2, 1e-2);
-        pass &= ck::utils::check_err(
-            r0_m.mData, r0_m_host.mData, "Error: Incorrect results d0", 1e-2, 1e-2);
-        pass &= ck::utils::check_err(
-            r1_m.mData, r1_m_host.mData, "Error: Incorrect results d1", 1e-2, 1e-2);
+        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]);
 
-    bool time_kernel = true;
-    if(time_kernel)
+        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
     {
-        float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
-        DumpPerf<ADataType, BDataType, EDataType, R0DataType, R1DataType>(ave_time, M, N, K);
+        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 pass ? 0 : 1;
+    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-2022, 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/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-2022, 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.begin(),
+                                                                             a_m_k.end());
+        ck::utils::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_k_n.begin(),
+                                                                             b_k_n.end());
+        break;
+    default:
+        ck::utils::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k.begin(), a_m_k.end());
+        ck::utils::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n.begin(), b_k_n.end());
+        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.mData,
+                                        e_m_n_host_converted.mData,
+                                        "Error: Incorrect results c",
+                                        1e-2,
+                                        1e-2);
+        }
+        else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        {
+            pass = ck::utils::check_err(
+                e_m_n.mData, e_m_n_host_converted.mData, "Error: Incorrect results c", 1e-2, 1e-2);
+        }
+
+        r0_device_buf.FromDevice(r0_m.mData.data());
+        pass &= ck::utils::check_err(
+            r0_m.mData, r0_m_host.mData, "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.begin(),
+                                                                             a_m_k.end());
+        ck::utils::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_k_n.begin(),
+                                                                             b_k_n.end());
+        break;
+    default:
+        ck::utils::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k.begin(), a_m_k.end());
+        ck::utils::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n.begin(), b_k_n.end());
+        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.mData,
+                                        e_m_n_host_converted.mData,
+                                        "Error: Incorrect results c",
+                                        1e-2,
+                                        1e-2);
+        }
+        else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        {
+            pass = ck::utils::check_err(
+                e_m_n.mData, e_m_n_host_converted.mData, "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.mData, r0_m_host.mData, "Error: Incorrect results d0", 1e-2, 1e-2);
+        pass &= ck::utils::check_err(
+            r1_m.mData, r1_m_host.mData, "Error: Incorrect results d1", 1e-2, 1e-2);
+
+        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/19_binary_elementwise/broadcast_add_2d_amn_bn.cpp

@@ -6,7 +6,7 @@
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp"
-#include "ck/tensor_operation/gpu/device/device_binary_elementwise.hpp"
+#include "ck/tensor_operation/gpu/device/device_elementwise.hpp"
 
 #include "ck/library/utility/check_err.hpp"
 #include "ck/library/utility/device_memory.hpp"
@@ -16,28 +16,23 @@
 using F16 = ck::half_t;
 using F32 = float;
 
-using ABDataType             = F16;
-using CDataType              = F16;
-using EltwiseComputeDataType = F32;
+using ABDataType = F16;
+using CDataType  = F16;
 
 using Add = ck::tensor_operation::element_wise::Add;
 
 using DeviceElementwiseAddInstance =
-    ck::tensor_operation::device::DeviceBinaryElementwise<ABDataType,
-                                                          ABDataType,
-                                                          CDataType,
-                                                          EltwiseComputeDataType,
-                                                          Add,
-                                                          2,
-                                                          8,
-                                                          8,
-                                                          8,
-                                                          8>;
+    ck::tensor_operation::device::DeviceElementwise<ck::Tuple<ABDataType, ABDataType>,
+                                                    ck::Tuple<CDataType>,
+                                                    Add,
+                                                    2,
+                                                    8,
+                                                    ck::Sequence<8, 8>,
+                                                    ck::Sequence<8>>;
 
 template <typename HostTensorA,
           typename HostTensorB,
           typename HostTensorC,
-          typename ComputeDataType,
           typename Functor,
           int broadcastDim>
 void host_broadcast2D(
@@ -49,19 +44,19 @@ void host_broadcast2D(
     {
         for(int n = 0; n < N; ++n)
         {
-            ComputeDataType Amn = ck::type_convert<ComputeDataType>(A(m, n));
-            ComputeDataType Cmn = 0;
+            auto Amn  = A(m, n);
+            ctype Cmn = 0;
             if constexpr(broadcastDim == 0)
             {
-                ComputeDataType Bn = ck::type_convert<ComputeDataType>(B(n));
+                auto Bn = B(n);
                 functor(Cmn, Amn, Bn);
             }
             else
             {
-                ComputeDataType Bm = ck::type_convert<ComputeDataType>(B(m));
+                auto Bm = B(m);
                 functor(Cmn, Amn, Bm);
             }
-            C(m, n) = ck::type_convert<ctype>(Cmn);
+            C(m, n) = Cmn;
         }
     }
 }
@@ -103,18 +98,19 @@ int main()
                                         b_n_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {c_m_n_device_buf.GetDeviceBuffer()};
 
-    std::vector<ck::index_t> a_strides = {Stride, 1};
-    std::vector<ck::index_t> b_strides = {0, 1};
-    std::vector<ck::index_t> c_strides = {Stride, 1};
+    std::array<ck::index_t, 2> abc_lengths = {M, N};
+    std::array<ck::index_t, 2> a_strides   = {Stride, 1};
+    std::array<ck::index_t, 2> b_strides   = {0, 1};
+    std::array<ck::index_t, 2> c_strides   = {Stride, 1};
 
     auto broadcastAdd = DeviceElementwiseAddInstance{};
     auto argument     = broadcastAdd.MakeArgumentPointer(
-        input, output, {M, N}, {a_strides, b_strides}, {c_strides}, Add{});
+        abc_lengths, {a_strides, b_strides}, {c_strides}, input, output, Add{});
 
     if(!broadcastAdd.IsSupportedArgument(argument.get()))
     {
-        throw std::runtime_error("The runtime parameters seems not supported by the "
-                                 "DeviceBinaryElementwise instance, exiting!");
+        throw std::runtime_error(
+            "The runtime parameters seems not supported by the device instance, exiting!");
     };
 
     auto broadcastAdd_invoker_ptr = broadcastAdd.MakeInvokerPointer();
@@ -129,12 +125,8 @@ int main()
         c_m_n_device_buf.FromDevice(c_m_n.mData.data());
         Tensor<CDataType> host_c_m_n(f_host_tensor_descriptor2d(M, N, Stride));
 
-        host_broadcast2D<Tensor<ABDataType>,
-                         Tensor<ABDataType>,
-                         Tensor<CDataType>,
-                         EltwiseComputeDataType,
-                         Add,
-                         0>(host_c_m_n, a_m_n, b_n, M, N, Add{});
+        host_broadcast2D<Tensor<ABDataType>, Tensor<ABDataType>, Tensor<CDataType>, Add, 0>(
+            host_c_m_n, a_m_n, b_n, M, N, Add{});
 
         pass &= ck::utils::check_err(
             c_m_n.mData, host_c_m_n.mData, "Error: Incorrect results c", 1e-3, 1e-3);

example/19_binary_elementwise/broadcast_add_3d_am_bmnk.cpp

@@ -6,7 +6,7 @@
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp"
-#include "ck/tensor_operation/gpu/device/device_binary_elementwise.hpp"
+#include "ck/tensor_operation/gpu/device/device_elementwise.hpp"
 
 #include "ck/library/utility/check_err.hpp"
 #include "ck/library/utility/device_memory.hpp"
@@ -16,29 +16,21 @@
 using F16 = ck::half_t;
 using F32 = float;
 
-using ABDataType             = F16;
-using CDataType              = F16;
-using EltwiseComputeDataType = F32;
+using ABDataType = F16;
+using CDataType  = F16;
 
 using Add = ck::tensor_operation::element_wise::Add;
 
 using DeviceElementwiseAddInstance =
-    ck::tensor_operation::device::DeviceBinaryElementwise<ABDataType,
-                                                          ABDataType,
-                                                          CDataType,
-                                                          EltwiseComputeDataType,
-                                                          Add,
-                                                          3,
-                                                          8,
-                                                          1,
-                                                          8,
-                                                          8>;
-
-template <typename HostTensorA,
-          typename HostTensorB,
-          typename HostTensorC,
-          typename ComputeDataType,
-          typename Functor>
+    ck::tensor_operation::device::DeviceElementwise<ck::Tuple<ABDataType, ABDataType>,
+                                                    ck::Tuple<CDataType>,
+                                                    Add,
+                                                    3,
+                                                    8,
+                                                    ck::Sequence<1, 8>,
+                                                    ck::Sequence<8>>;
+
+template <typename HostTensorA, typename HostTensorB, typename HostTensorC, typename Functor>
 void host_broadcast3D_am_bmnk(HostTensorC& C,
                               const HostTensorA& A,
                               const HostTensorB& B,
@@ -51,11 +43,11 @@ void host_broadcast3D_am_bmnk(HostTensorC& C,
         for(std::size_t n = 0; n < shape[1]; ++n)
             for(std::size_t k = 0; k < shape[2]; ++k)
             {
-                ComputeDataType a_val = ck::type_convert<ComputeDataType>(A(m));
-                ComputeDataType b_val = ck::type_convert<ComputeDataType>(B(m, n, k));
-                ComputeDataType c_val = 0;
+                auto a_val  = A(m);
+                auto b_val  = B(m, n, k);
+                ctype c_val = 0;
                 functor(c_val, a_val, b_val);
-                C(m, n, k) = ck::type_convert<ctype>(c_val);
+                C(m, n, k) = c_val;
             }
 }
 
@@ -85,25 +77,25 @@ int main()
                                         b_m_n_k_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {c_m_n_k_device_buf.GetDeviceBuffer()};
 
-    std::vector<ck::index_t> a_strides = {1, 0, 0};
-    std::vector<ck::index_t> b_strides{b_m_n_k.mDesc.GetStrides().begin(),
-                                       b_m_n_k.mDesc.GetStrides().end()};
-    std::vector<ck::index_t> c_strides{c_m_n_k.mDesc.GetStrides().begin(),
-                                       c_m_n_k.mDesc.GetStrides().end()};
+    std::array<ck::index_t, 3> abc_lengths;
+    std::array<ck::index_t, 3> a_strides = {1, 0, 0};
+    std::array<ck::index_t, 3> b_strides;
+    std::array<ck::index_t, 3> c_strides;
+
+    std::copy(mnk.begin(), mnk.end(), abc_lengths.begin());
+    std::copy(
+        b_m_n_k.mDesc.GetStrides().begin(), b_m_n_k.mDesc.GetStrides().end(), b_strides.begin());
+    std::copy(
+        c_m_n_k.mDesc.GetStrides().begin(), c_m_n_k.mDesc.GetStrides().end(), c_strides.begin());
 
     auto broadcastAdd = DeviceElementwiseAddInstance{};
-    auto argument =
-        broadcastAdd.MakeArgumentPointer(input,
-                                         output,
-                                         std::vector<ck::index_t>{mnk.begin(), mnk.end()},
-                                         {a_strides, b_strides},
-                                         {c_strides},
-                                         Add{});
+    auto argument     = broadcastAdd.MakeArgumentPointer(
+        abc_lengths, {a_strides, b_strides}, {c_strides}, input, output, Add{});
 
     if(!broadcastAdd.IsSupportedArgument(argument.get()))
     {
-        throw std::runtime_error("The runtime parameters seems not supported by the "
-                                 "DeviceBinaryElementwise instance, exiting!");
+        throw std::runtime_error(
+            "The runtime parameters seems not supported by the device instance, exiting!");
     };
 
     auto broadcastAdd_invoker_ptr = broadcastAdd.MakeInvokerPointer();
@@ -118,11 +110,8 @@ int main()
         c_m_n_k_device_buf.FromDevice(c_m_n_k.mData.data());
         Tensor<CDataType> host_c_m_n_k(mnk);
 
-        host_broadcast3D_am_bmnk<Tensor<ABDataType>,
-                                 Tensor<ABDataType>,
-                                 Tensor<CDataType>,
-                                 EltwiseComputeDataType,
-                                 Add>(host_c_m_n_k, a_m, b_m_n_k, mnk, Add{});
+        host_broadcast3D_am_bmnk<Tensor<ABDataType>, Tensor<ABDataType>, Tensor<CDataType>, Add>(
+            host_c_m_n_k, a_m, b_m_n_k, mnk, Add{});
 
         pass &= ck::utils::check_err(
             c_m_n_k.mData, host_c_m_n_k.mData, "Error: Incorrect results c", 1e-3, 1e-3);

example/19_binary_elementwise/elementwise_add_1d.cpp

@@ -5,7 +5,7 @@
 #include <cstdlib>
 
 #include "ck/ck.hpp"
-#include "ck/tensor_operation/gpu/device/device_binary_elementwise.hpp"
+#include "ck/tensor_operation/gpu/device/device_elementwise.hpp"
 #include "ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp"
 #include "ck/library/utility/check_err.hpp"
 #include "ck/library/utility/device_memory.hpp"
@@ -15,29 +15,21 @@
 using F16 = ck::half_t;
 using F32 = float;
 
-using ABDataType             = F16;
-using CDataType              = F16;
-using EltwiseComputeDataType = F32;
+using ABDataType = F16;
+using CDataType  = F16;
 
 using Add = ck::tensor_operation::element_wise::Add;
 
 using DeviceElementwiseAddInstance =
-    ck::tensor_operation::device::DeviceBinaryElementwise<ABDataType,
-                                                          ABDataType,
-                                                          CDataType,
-                                                          EltwiseComputeDataType,
-                                                          Add,
-                                                          1,
-                                                          8,
-                                                          8,
-                                                          8,
-                                                          8>;
-
-template <typename HostTensorA,
-          typename HostTensorB,
-          typename HostTensorC,
-          typename ComputeDataType,
-          typename Functor>
+    ck::tensor_operation::device::DeviceElementwise<ck::Tuple<ABDataType, ABDataType>,
+                                                    ck::Tuple<CDataType>,
+                                                    Add,
+                                                    1,
+                                                    8,
+                                                    ck::Sequence<8, 8>,
+                                                    ck::Sequence<8>>;
+
+template <typename HostTensorA, typename HostTensorB, typename HostTensorC, typename Functor>
 void host_elementwise1D(
     HostTensorC& C, const HostTensorA& A, const HostTensorB& B, int M, Functor functor)
 {
@@ -45,11 +37,11 @@ void host_elementwise1D(
 
     for(int m = 0; m < M; ++m)
     {
-        ComputeDataType Am = ck::type_convert<ComputeDataType>(A(m));
-        ComputeDataType Bm = ck::type_convert<ComputeDataType>(B(m));
-        ComputeDataType Cm = 0;
+        auto Am  = A(m);
+        auto Bm  = B(m);
+        ctype Cm = 0;
         functor(Cm, Am, Bm);
-        C(m) = ck::type_convert<ctype>(Cm);
+        C(m) = Cm;
     }
 }
 
@@ -83,18 +75,19 @@ int main()
                                         b_m_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {c_m_device_buf.GetDeviceBuffer()};
 
-    std::vector<ck::index_t> a_strides = {1};
-    std::vector<ck::index_t> b_strides = {1};
-    std::vector<ck::index_t> c_strides = {1};
+    std::array<ck::index_t, 1> abc_lengths = {M};
+    std::array<ck::index_t, 1> a_strides   = {1};
+    std::array<ck::index_t, 1> b_strides   = {1};
+    std::array<ck::index_t, 1> c_strides   = {1};
 
     auto broadcastAdd = DeviceElementwiseAddInstance{};
     auto argument     = broadcastAdd.MakeArgumentPointer(
-        input, output, {M}, {{a_strides}, b_strides}, {c_strides}, Add{});
+        abc_lengths, {a_strides, b_strides}, {c_strides}, input, output, Add{});
 
     if(!broadcastAdd.IsSupportedArgument(argument.get()))
     {
-        throw std::runtime_error("The runtime parameters seems not supported by the "
-                                 "DeviceBinaryElementwise instance, exiting!");
+        throw std::runtime_error(
+            "The runtime parameters seems not supported by the device instance, exiting!");
     };
 
     auto broadcastAdd_invoker_ptr = broadcastAdd.MakeInvokerPointer();
@@ -109,11 +102,8 @@ int main()
         c_m_device_buf.FromDevice(c_m.mData.data());
         Tensor<CDataType> host_c_m(f_host_tensor_descriptor1d(M, 1));
 
-        host_elementwise1D<Tensor<ABDataType>,
-                           Tensor<ABDataType>,
-                           Tensor<CDataType>,
-                           EltwiseComputeDataType,
-                           Add>(host_c_m, a_m, b_m, M, Add{});
+        host_elementwise1D<Tensor<ABDataType>, Tensor<ABDataType>, Tensor<CDataType>, Add>(
+            host_c_m, a_m, b_m, M, Add{});
 
         pass &= ck::utils::check_err(
             c_m.mData, host_c_m.mData, "Error: Incorrect results c", 1e-3, 1e-3);

example/19_binary_elementwise/elementwise_add_4d.cpp

@@ -6,7 +6,7 @@
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp"
-#include "ck/tensor_operation/gpu/device/device_binary_elementwise.hpp"
+#include "ck/tensor_operation/gpu/device/device_elementwise.hpp"
 
 #include "ck/library/utility/check_err.hpp"
 #include "ck/library/utility/device_memory.hpp"
@@ -16,29 +16,21 @@
 using F16 = ck::half_t;
 using F32 = float;
 
-using ABDataType             = F16;
-using CDataType              = F16;
-using EltwiseComputeDataType = F32;
+using ABDataType = F16;
+using CDataType  = F16;
 
 using Add = ck::tensor_operation::element_wise::Add;
 
 using DeviceElementwiseAddInstance =
-    ck::tensor_operation::device::DeviceBinaryElementwise<ABDataType,
-                                                          ABDataType,
-                                                          CDataType,
-                                                          EltwiseComputeDataType,
-                                                          Add,
-                                                          4,
-                                                          8,
-                                                          8,
-                                                          8,
-                                                          8>;
-
-template <typename HostTensorA,
-          typename HostTensorB,
-          typename HostTensorC,
-          typename ComputeDataType,
-          typename Functor>
+    ck::tensor_operation::device::DeviceElementwise<ck::Tuple<ABDataType, ABDataType>,
+                                                    ck::Tuple<CDataType>,
+                                                    Add,
+                                                    4,
+                                                    8,
+                                                    ck::Sequence<8, 8>,
+                                                    ck::Sequence<8>>;
+
+template <typename HostTensorA, typename HostTensorB, typename HostTensorC, typename Functor>
 void host_elementwise4D(HostTensorC& C,
                         const HostTensorA& A,
                         const HostTensorB& B,
@@ -52,11 +44,11 @@ void host_elementwise4D(HostTensorC& C,
             for(std::size_t h = 0; h < shape[2]; ++h)
                 for(std::size_t w = 0; w < shape[3]; ++w)
                 {
-                    ComputeDataType a_val = ck::type_convert<ComputeDataType>(A(n, c, h, w));
-                    ComputeDataType b_val = ck::type_convert<ComputeDataType>(B(n, c, h, w));
-                    ComputeDataType c_val = 0;
+                    auto a_val  = A(n, c, h, w);
+                    auto b_val  = B(n, c, h, w);
+                    ctype c_val = 0;
                     functor(c_val, a_val, b_val);
-                    C(n, c, h, w) = ck::type_convert<ctype>(c_val);
+                    C(n, c, h, w) = c_val;
                 }
 }
 
@@ -85,23 +77,24 @@ int main()
                                         b_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {c_device_buf.GetDeviceBuffer()};
 
-    std::vector<ck::index_t> a_strides{a.mDesc.GetStrides().begin(), a.mDesc.GetStrides().end()};
-    std::vector<ck::index_t> b_strides{b.mDesc.GetStrides().begin(), b.mDesc.GetStrides().end()};
-    std::vector<ck::index_t> c_strides{c.mDesc.GetStrides().begin(), c.mDesc.GetStrides().end()};
+    std::array<ck::index_t, 4> abc_lengths;
+    std::array<ck::index_t, 4> a_strides;
+    std::array<ck::index_t, 4> b_strides;
+    std::array<ck::index_t, 4> c_strides;
+
+    std::copy(nchw.begin(), nchw.end(), abc_lengths.begin());
+    std::copy(a.mDesc.GetStrides().begin(), a.mDesc.GetStrides().end(), a_strides.begin());
+    std::copy(b.mDesc.GetStrides().begin(), b.mDesc.GetStrides().end(), b_strides.begin());
+    std::copy(c.mDesc.GetStrides().begin(), c.mDesc.GetStrides().end(), c_strides.begin());
 
     auto broadcastAdd = DeviceElementwiseAddInstance{};
-    auto argument =
-        broadcastAdd.MakeArgumentPointer(input,
-                                         output,
-                                         std::vector<ck::index_t>{nchw.begin(), nchw.end()},
-                                         {{a_strides}, b_strides},
-                                         {c_strides},
-                                         Add{});
+    auto argument     = broadcastAdd.MakeArgumentPointer(
+        abc_lengths, {a_strides, b_strides}, {c_strides}, input, output, Add{});
 
     if(!broadcastAdd.IsSupportedArgument(argument.get()))
     {
-        throw std::runtime_error("The runtime parameters seems not supported by the "
-                                 "DeviceBinaryElementwise instance, exiting!");
+        throw std::runtime_error(
+            "The runtime parameters seems not supported by the device instance, exiting!");
     };
 
     auto broadcastAdd_invoker_ptr = broadcastAdd.MakeInvokerPointer();
@@ -116,11 +109,8 @@ int main()
         c_device_buf.FromDevice(c.mData.data());
         Tensor<CDataType> host_c(nchw);
 
-        host_elementwise4D<Tensor<ABDataType>,
-                           Tensor<ABDataType>,
-                           Tensor<CDataType>,
-                           EltwiseComputeDataType,
-                           Add>(host_c, a, b, nchw, Add{});
+        host_elementwise4D<Tensor<ABDataType>, Tensor<ABDataType>, Tensor<CDataType>, Add>(
+            host_c, a, b, nchw, Add{});
 
         pass &=
             ck::utils::check_err(c.mData, host_c.mData, "Error: Incorrect results c", 1e-3, 1e-3);

example/21_gemm_layernorm/gemm_bias_relu_add_layernorm_xdl_fp16.cpp

@@ -10,7 +10,7 @@
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
-#include "ck/tensor_operation/gpu/device/device_5ary_elementwise.hpp"
+#include "ck/tensor_operation/gpu/device/device_elementwise.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
 #include "ck/library/utility/device_memory.hpp"
@@ -94,23 +94,18 @@ using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataTyp
 using NormalizeFunctor = ck::tensor_operation::element_wise::Normalize;
 
 // A:x, B:E[x], C:E[x^2], D:Gamma, E:Beta , F:y
-using DeviceNormalizeInstance =
-    ck::tensor_operation::device::Device5AryElementwise<EDataType,
-                                                        R0DataType,
-                                                        R1DataType,
-                                                        GammaDataType,
-                                                        BetaDataType,
-                                                        LayerNormOutDataType,
-                                                        NormalizeComputeDataType,
-                                                        NormalizeFunctor,
-                                                        2,
-                                                        8,
-                                                        8,  // scalarPerVector: gemm_out
-                                                        1,  // scalarPerVector: reduce_mean
-                                                        1,  // scalarPerVector: reduce_mean_square
-                                                        8,  // scalarPerVector: Gamma
-                                                        8,  // scalarPerVector: Beta
-                                                        8>; // scalarPerVector: LayerNorm_out
+using DeviceNormalizeInstance = ck::tensor_operation::device::DeviceElementwise<
+    ck::Tuple<EDataType,
+              R0DataType,
+              R1DataType,
+              GammaDataType,
+              BetaDataType>,         // x(gemm_out), mean, meansquare, gamma, beta
+    ck::Tuple<LayerNormOutDataType>, // y
+    NormalizeFunctor,
+    2,
+    8,                           // MPerthread
+    ck::Sequence<8, 1, 1, 8, 8>, // scalarPerVector: x(gemm_out), mean, meansquare, gamma, beta
+    ck::Sequence<8>>;            // scalarPerVector: y(layerNorm_out)
 
 auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
     return HostTensorDescriptor(std::vector<std::size_t>({len}),
@@ -197,14 +192,9 @@ void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
     {
         for(int n = 0; n < N; ++n)
         {
-            NormalizeComputeDataType out_acc = 0;
-            layerNormInst(out_acc,
-                          ck::type_convert<NormalizeComputeDataType>(e_m_n(m, n)),
-                          ck::type_convert<NormalizeComputeDataType>(mean_m(m)),
-                          ck::type_convert<NormalizeComputeDataType>(meanSquare_m(m)),
-                          ck::type_convert<NormalizeComputeDataType>(gamma_n(n)),
-                          ck::type_convert<NormalizeComputeDataType>(beta_n(n)));
-            out_m_n(m, n) = ck::type_convert<LayerNormOutDataType>(out_acc);
+            LayerNormOutDataType out_val = 0;
+            layerNormInst(out_val, e_m_n(m, n), mean_m(m), meanSquare_m(m), gamma_n(n), beta_n(n));
+            out_m_n(m, n) = out_val;
         }
     }
 }
@@ -339,28 +329,28 @@ int main()
                                         beta_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {layerNorm_device_buf.GetDeviceBuffer()};
 
-    auto normalize          = DeviceNormalizeInstance{};
-    auto normalize_invoker  = normalize.MakeInvoker();
-    auto normalize_argument = normalize.MakeArgument(input,
-                                                     output,
-                                                     {M, N},
-                                                     {StrideE, 1},
-                                                     {1, 0},
-                                                     {1, 0},
-                                                     {0, 1},
-                                                     {0, 1},
-                                                     {StrideE, 1},
-                                                     NormalizeFunctor{});
-
-    if(!normalize.IsSupportedArgument(normalize_argument))
+    std::array<ck::index_t, 2> xyLengths = {M, N};
+    std::array<ck::index_t, 2> xyStrides = {StrideE, 1};
+
+    auto normalize         = DeviceNormalizeInstance{};
+    auto normalize_invoker = normalize.MakeInvoker();
+    auto normalize_argument_ptr =
+        normalize.MakeArgumentPointer(xyLengths,
+                                      {xyStrides, {1, 0}, {1, 0}, {0, 1}, {0, 1}},
+                                      {xyStrides},
+                                      input,
+                                      output,
+                                      NormalizeFunctor{});
+
+    if(!normalize.IsSupportedArgument(normalize_argument_ptr.get()))
     {
-        throw std::runtime_error("The runtime parameters seems not supported by the "
-                                 "Device5AryElementwise instance, exiting!");
+        throw std::runtime_error(
+            "The runtime parameters seems not supported by the device, exiting!");
     }
 
     // run kernel
     gemmReduce_invoker.Run(gemmReduce_argument, StreamConfig{nullptr, false});
-    normalize_invoker.Run(normalize_argument, StreamConfig{nullptr, false});
+    normalize_invoker.Run(normalize_argument_ptr.get(), StreamConfig{nullptr, false});
 
     bool pass = true;
     {
@@ -396,7 +386,7 @@ int main()
         float gemm_reduce_mean_reduce_square_mean_ave_time =
             gemmReduce_invoker.Run(gemmReduce_argument, StreamConfig{nullptr, time_kernel});
         float normalize_ave_time =
-            normalize_invoker.Run(normalize_argument, StreamConfig{nullptr, time_kernel});
+            normalize_invoker.Run(normalize_argument_ptr.get(), StreamConfig{nullptr, time_kernel});
 
         if(time_kernel)
             DumpGemmLayerNormPerf<ADataType,

example/21_gemm_layernorm/gemm_layernorm_xdl_fp16.cpp

@@ -10,7 +10,7 @@
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_multiple_r_xdl_cshuffle.hpp"
-#include "ck/tensor_operation/gpu/device/device_5ary_elementwise.hpp"
+#include "ck/tensor_operation/gpu/device/device_elementwise.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
 #include "ck/library/utility/device_memory.hpp"
@@ -91,23 +91,20 @@ using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataTyp
 using NormalizeFunctor = ck::tensor_operation::element_wise::Normalize;
 
 // A:x, B:E[x], C:E[x^2], D:Gamma, E:Beta , F:y
-using DeviceNormalizeInstance =
-    ck::tensor_operation::device::Device5AryElementwise<EDataType,
-                                                        R0DataType,
-                                                        R1DataType,
-                                                        GammaDataType,
-                                                        BetaDataType,
-                                                        LayerNormOutDataType,
-                                                        NormalizeComputeDataType,
-                                                        NormalizeFunctor,
-                                                        2,
-                                                        8,
-                                                        8,  // scalarPerVector: gemm_out
-                                                        1,  // scalarPerVector: reduce_mean
-                                                        1,  // scalarPerVector: reduce_mean_square
-                                                        8,  // scalarPerVector: Gamma
-                                                        8,  // scalarPerVector: Beta
-                                                        8>; // scalarPerVector: LayerNorm_out
+using DeviceNormalizeInstance = ck::tensor_operation::device::DeviceElementwise<
+    ck::Tuple<EDataType,
+              R0DataType,
+              R1DataType,
+              GammaDataType,
+              BetaDataType>,         // x(gemm_out), mean,
+                                     // meansquare,
+                                     // gamma, beta
+    ck::Tuple<LayerNormOutDataType>, // y
+    NormalizeFunctor,
+    2,
+    8,                           // MPerthread
+    ck::Sequence<8, 1, 1, 8, 8>, // scalarPerVector: x(gemm_out), mean, meansquare, gamma, beta
+    ck::Sequence<8>>;            // scalarPerVector: y(layerNorm_out)
 
 auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
     return HostTensorDescriptor(std::vector<std::size_t>({len}),
@@ -139,7 +136,6 @@ void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
                          int M,
                          int N)
 {
-
     int StrideE = N;
     Tensor<EDataType> e_m_n(f_host_tensor_descriptor2d(M, N, StrideE, ELayout{}));
     Tensor<R0DataType> mean_m(f_host_tensor_descriptor1d(M, 1));
@@ -184,14 +180,9 @@ void host_gemm_layernorm(Tensor<LayerNormOutDataType>& out_m_n,
     {
         for(int n = 0; n < N; ++n)
         {
-            NormalizeComputeDataType out_acc = 0;
-            layerNormInst(out_acc,
-                          ck::type_convert<NormalizeComputeDataType>(e_m_n(m, n)),
-                          ck::type_convert<NormalizeComputeDataType>(mean_m(m)),
-                          ck::type_convert<NormalizeComputeDataType>(meanSquare_m(m)),
-                          ck::type_convert<NormalizeComputeDataType>(gamma_n(n)),
-                          ck::type_convert<NormalizeComputeDataType>(beta_n(n)));
-            out_m_n(m, n) = ck::type_convert<LayerNormOutDataType>(out_acc);
+            LayerNormOutDataType out_val = 0;
+            layerNormInst(out_val, e_m_n(m, n), mean_m(m), meanSquare_m(m), gamma_n(n), beta_n(n));
+            out_m_n(m, n) = out_val;
         }
     }
 }
@@ -314,28 +305,28 @@ int main()
                                         beta_device_buf.GetDeviceBuffer()};
     std::array<void*, 1> output      = {layerNorm_device_buf.GetDeviceBuffer()};
 
-    auto normalize          = DeviceNormalizeInstance{};
-    auto normalize_invoker  = normalize.MakeInvoker();
-    auto normalize_argument = normalize.MakeArgument(input,
-                                                     output,
-                                                     {M, N},
-                                                     {StrideE, 1},
-                                                     {1, 0},
-                                                     {1, 0},
-                                                     {0, 1},
-                                                     {0, 1},
-                                                     {StrideE, 1},
-                                                     NormalizeFunctor{});
-
-    if(!normalize.IsSupportedArgument(normalize_argument))
+    std::array<ck::index_t, 2> xyLengths = {M, N};
+    std::array<ck::index_t, 2> xyStrides = {StrideE, 1};
+
+    auto normalize         = DeviceNormalizeInstance{};
+    auto normalize_invoker = normalize.MakeInvoker();
+    auto normalize_argument_ptr =
+        normalize.MakeArgumentPointer(xyLengths,
+                                      {xyStrides, {1, 0}, {1, 0}, {0, 1}, {0, 1}},
+                                      {xyStrides},
+                                      input,
+                                      output,
+                                      NormalizeFunctor{});
+
+    if(!normalize.IsSupportedArgument(normalize_argument_ptr.get()))
     {
-        throw std::runtime_error("The runtime parameters seems not supported by the "
-                                 "Device5AryElementwise instance, exiting!");
+        throw std::runtime_error(
+            "The runtime parameters seems not supported by the device, exiting");
     }
 
     // run kernel
     gemmReduce_invoker.Run(gemmReduce_argument, StreamConfig{nullptr, false});
-    normalize_invoker.Run(normalize_argument, StreamConfig{nullptr, false});
+    normalize_invoker.Run(normalize_argument_ptr.get(), StreamConfig{nullptr, false});
 
     bool pass = true;
     {
@@ -369,7 +360,7 @@ int main()
         float gemm_reduce_mean_reduce_square_mean_ave_time =
             gemmReduce_invoker.Run(gemmReduce_argument, StreamConfig{nullptr, time_kernel});
         float normalize_ave_time =
-            normalize_invoker.Run(normalize_argument, StreamConfig{nullptr, time_kernel});
+            normalize_invoker.Run(normalize_argument_ptr.get(), StreamConfig{nullptr, time_kernel});
 
         if(time_kernel)
             DumpGemmLayerNormPerf<ADataType,

example/22_cgemm/CMakeLists.txt

@@ -5,7 +5,13 @@ add_example_executable(example_cgemm_xdl_fp16 cgemm_xdl_fp16.cpp)
 add_example_executable(example_cgemm_xdl_fp32 cgemm_xdl_fp32.cpp)
 add_example_executable(example_cgemm_xdl_int8 cgemm_xdl_int8.cpp)
 
-add_dependencies(example_cgemm_xdl example_cgemm_xdl_bf16)
-add_dependencies(example_cgemm_xdl example_cgemm_xdl_fp16)
-add_dependencies(example_cgemm_xdl example_cgemm_xdl_fp32)
-add_dependencies(example_cgemm_xdl example_cgemm_xdl_int8)
+add_dependencies(example_cgemm_xdl 
+                 example_cgemm_xdl_bf16
+                 example_cgemm_xdl_fp16
+                 example_cgemm_xdl_fp32
+                 example_cgemm_xdl_int8)
+
+if(USE_BITINT_EXTENSION_INT4)
+  add_example_executable(example_cgemm_xdl_int4 cgemm_xdl_int4.cpp)
+  add_dependencies(example_cgemm_xdl example_cgemm_xdl_int4)
+endif()

example/22_cgemm/cgemm_xdl_bf16.cpp

@@ -117,16 +117,16 @@ int main(int argc, char* argv[])
         exit(0);
     }
 
-    return run_cgemm_xdl<ADataType,
-                         BDataType,
-                         CDataType,
-                         ALayout,
-                         BLayout,
-                         CLayout,
-                         PassThrough,
-                         PassThrough,
-                         PassThrough,
-                         DeviceCGemmInstance,
-                         ReferenceCGemmInstance>(
+    return !run_cgemm_xdl<ADataType,
+                          BDataType,
+                          CDataType,
+                          ALayout,
+                          BLayout,
+                          CLayout,
+                          PassThrough,
+                          PassThrough,
+                          PassThrough,
+                          DeviceCGemmInstance,
+                          ReferenceCGemmInstance>(
         M, N, K, StrideA, StrideB, StrideC, do_verification, init_method, time_kernel);
 }

example/22_cgemm/cgemm_xdl_common.hpp

@@ -21,6 +21,9 @@ using F32   = float;
 using BF16  = ck::bhalf_t;
 using INT8  = std::int8_t;
 using INT32 = std::int32_t;
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+using INT4 = ck::int4_t;
+#endif
 
 template <typename ADataType,
           typename BDataType,
@@ -32,17 +35,31 @@ template <typename ADataType,
           typename BElementwiseOperation,
           typename CElementwiseOperation,
           typename DeviceCGemmInstance,
-          typename ReferenceCGemmInstance>
-int run_cgemm_xdl(ck::index_t M,
-                  ck::index_t N,
-                  ck::index_t K,
-                  ck::index_t StrideA,
-                  ck::index_t StrideB,
-                  ck::index_t StrideC,
-                  bool do_verification,
-                  int init_method,
-                  bool time_kernel)
+          typename ReferenceCGemmInstance,
+          typename KernelADataType = ADataType,
+          typename KernelBDataType = BDataType,
+          typename KernelCDataType = CDataType>
+bool run_cgemm_xdl(ck::index_t M,
+                   ck::index_t N,
+                   ck::index_t K,
+                   ck::index_t StrideA,
+                   ck::index_t StrideB,
+                   ck::index_t StrideC,
+                   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),
+                  "sizeof ck::int4_t and int8_t is different!");
+    static_assert(sizeof(ADataType) == sizeof(KernelADataType),
+                  "sizeof ADataType and KernelADataType is different!");
+    static_assert(sizeof(BDataType) == sizeof(KernelBDataType),
+                  "sizeof BDataType and KernelBDataType is different!");
+    static_assert(sizeof(CDataType) == sizeof(KernelCDataType),
+                  "sizeof CDataType and KernelCDataType is different!");
+#endif
+
     auto f_host_tensor_descriptor =
         [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
             if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
@@ -61,8 +78,10 @@ int run_cgemm_xdl(ck::index_t M,
     Tensor<ADataType> a_m_k_imag(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
     Tensor<BDataType> b_k_n_real(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
     Tensor<BDataType> b_k_n_imag(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
-    Tensor<CDataType> c_m_n_real_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
-    Tensor<CDataType> c_m_n_imag_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
+    Tensor<KernelCDataType> c_m_n_real_device_result(
+        f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
+    Tensor<KernelCDataType> c_m_n_imag_device_result(
+        f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
 
     std::cout << "a_m_k_real: " << a_m_k_real.mDesc << std::endl;
     std::cout << "a_m_k_imag: " << a_m_k_imag.mDesc << std::endl;
@@ -89,20 +108,41 @@ int run_cgemm_xdl(ck::index_t M,
 
     auto cgemm = DeviceCGemmInstance{};
 
-    DeviceMem a_m_k_real_device_buf(sizeof(ADataType) * a_m_k_real.mDesc.GetElementSpaceSize());
-    DeviceMem a_m_k_imag_device_buf(sizeof(ADataType) * a_m_k_imag.mDesc.GetElementSpaceSize());
-    DeviceMem b_k_n_real_device_buf(sizeof(BDataType) * b_k_n_real.mDesc.GetElementSpaceSize());
-    DeviceMem b_k_n_imag_device_buf(sizeof(BDataType) * b_k_n_imag.mDesc.GetElementSpaceSize());
-    DeviceMem c_m_n_real_device_buf(sizeof(CDataType) *
+    DeviceMem a_m_k_real_device_buf(sizeof(KernelADataType) *
+                                    a_m_k_real.mDesc.GetElementSpaceSize());
+    DeviceMem a_m_k_imag_device_buf(sizeof(KernelADataType) *
+                                    a_m_k_imag.mDesc.GetElementSpaceSize());
+    DeviceMem b_k_n_real_device_buf(sizeof(KernelBDataType) *
+                                    b_k_n_real.mDesc.GetElementSpaceSize());
+    DeviceMem b_k_n_imag_device_buf(sizeof(KernelBDataType) *
+                                    b_k_n_imag.mDesc.GetElementSpaceSize());
+    DeviceMem c_m_n_real_device_buf(sizeof(KernelCDataType) *
                                     c_m_n_real_device_result.mDesc.GetElementSpaceSize());
-    DeviceMem c_m_n_imag_device_buf(sizeof(CDataType) *
+    DeviceMem c_m_n_imag_device_buf(sizeof(KernelCDataType) *
                                     c_m_n_imag_device_result.mDesc.GetElementSpaceSize());
     DeviceMem workspace_device_buf(cgemm.GetWorkspaceSize(M, N, K, StrideA, StrideB, StrideC));
 
-    a_m_k_real_device_buf.ToDevice(a_m_k_real.mData.data());
-    a_m_k_imag_device_buf.ToDevice(a_m_k_imag.mData.data());
-    b_k_n_real_device_buf.ToDevice(b_k_n_real.mData.data());
-    b_k_n_imag_device_buf.ToDevice(b_k_n_imag.mData.data());
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    if constexpr(std::is_same_v<ADataType, ck::int4_t>)
+    {
+        Tensor<KernelADataType> a_m_k_real_converted(a_m_k_real);
+        Tensor<KernelADataType> a_m_k_imag_converted(a_m_k_imag);
+        Tensor<KernelBDataType> b_k_n_real_converted(b_k_n_real);
+        Tensor<KernelBDataType> b_k_n_imag_converted(b_k_n_imag);
+
+        a_m_k_real_device_buf.ToDevice(a_m_k_real_converted.mData.data());
+        a_m_k_imag_device_buf.ToDevice(a_m_k_imag_converted.mData.data());
+        b_k_n_real_device_buf.ToDevice(b_k_n_real_converted.mData.data());
+        b_k_n_imag_device_buf.ToDevice(b_k_n_imag_converted.mData.data());
+    }
+    else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+    {
+        a_m_k_real_device_buf.ToDevice(a_m_k_real.mData.data());
+        a_m_k_imag_device_buf.ToDevice(a_m_k_imag.mData.data());
+        b_k_n_real_device_buf.ToDevice(b_k_n_real.mData.data());
+        b_k_n_imag_device_buf.ToDevice(b_k_n_imag.mData.data());
+    }
 
     auto a_element_op = AElementwiseOperation{};
     auto b_element_op = BElementwiseOperation{};
@@ -111,13 +151,13 @@ int run_cgemm_xdl(ck::index_t M,
     // do GEMM
     auto invoker = cgemm.MakeInvoker();
     auto argument =
-        cgemm.MakeArgument(static_cast<ADataType*>(a_m_k_real_device_buf.GetDeviceBuffer()),
-                           static_cast<ADataType*>(a_m_k_imag_device_buf.GetDeviceBuffer()),
-                           static_cast<BDataType*>(b_k_n_real_device_buf.GetDeviceBuffer()),
-                           static_cast<BDataType*>(b_k_n_imag_device_buf.GetDeviceBuffer()),
-                           static_cast<CDataType*>(c_m_n_real_device_buf.GetDeviceBuffer()),
-                           static_cast<CDataType*>(c_m_n_imag_device_buf.GetDeviceBuffer()),
-                           static_cast<CDataType*>(workspace_device_buf.GetDeviceBuffer()),
+        cgemm.MakeArgument(static_cast<KernelADataType*>(a_m_k_real_device_buf.GetDeviceBuffer()),
+                           static_cast<KernelADataType*>(a_m_k_imag_device_buf.GetDeviceBuffer()),
+                           static_cast<KernelBDataType*>(b_k_n_real_device_buf.GetDeviceBuffer()),
+                           static_cast<KernelBDataType*>(b_k_n_imag_device_buf.GetDeviceBuffer()),
+                           static_cast<KernelCDataType*>(c_m_n_real_device_buf.GetDeviceBuffer()),
+                           static_cast<KernelCDataType*>(c_m_n_imag_device_buf.GetDeviceBuffer()),
+                           static_cast<KernelCDataType*>(workspace_device_buf.GetDeviceBuffer()),
                            M,
                            N,
                            K,
@@ -142,16 +182,12 @@ int run_cgemm_xdl(ck::index_t M,
         std::size_t(2) *
         (sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N);
 
-    float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
-
+    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, "
               << cgemm.GetTypeString() << std::endl;
 
-    c_m_n_real_device_buf.FromDevice(c_m_n_real_device_result.mData.data());
-    c_m_n_imag_device_buf.FromDevice(c_m_n_imag_device_result.mData.data());
-
     if(do_verification)
     {
         Tensor<CDataType> c_m_n_real_host_result(
@@ -159,9 +195,8 @@ int run_cgemm_xdl(ck::index_t M,
         Tensor<CDataType> c_m_n_imag_host_result(
             f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
 
-        auto ref_cgemm   = ReferenceCGemmInstance{};
-        auto ref_invoker = ref_cgemm.MakeInvoker();
-
+        auto ref_cgemm    = ReferenceCGemmInstance{};
+        auto ref_invoker  = ref_cgemm.MakeInvoker();
         auto ref_argument = ref_cgemm.MakeArgument(a_m_k_real,
                                                    a_m_k_imag,
                                                    b_k_n_real,
@@ -174,19 +209,45 @@ int run_cgemm_xdl(ck::index_t M,
 
         ref_invoker.Run(ref_argument);
 
+        c_m_n_real_device_buf.FromDevice(c_m_n_real_device_result.mData.data());
+        c_m_n_imag_device_buf.FromDevice(c_m_n_imag_device_result.mData.data());
+
         bool result = true;
-        result      = ck::utils::check_err(c_m_n_real_device_result.mData,
-                                      c_m_n_real_host_result.mData,
-                                      "Verification error: incorrect results in real part!",
-                                      1e-2f,
-                                      1e-1f);
-        result      = result &&
-                 ck::utils::check_err(c_m_n_imag_device_result.mData,
-                                      c_m_n_imag_host_result.mData,
-                                      "Verification error: incorrect results in imaginary part!",
-                                      1e-2f,
-                                      1e-1f);
-        return result ? 0 : 1;
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        if constexpr(std::is_same_v<ADataType, ck::int4_t>)
+        {
+            const Tensor<CDataType> c_m_n_real_device_result_converted(c_m_n_real_device_result);
+            const Tensor<CDataType> c_m_n_imag_device_result_converted(c_m_n_imag_device_result);
+
+            result = ck::utils::check_err(c_m_n_real_device_result_converted.mData,
+                                          c_m_n_real_host_result.mData,
+                                          "Verification error: incorrect results in real part!",
+                                          1e-2f,
+                                          1e-1f);
+            result = result && ck::utils::check_err(
+                                   c_m_n_imag_device_result_converted.mData,
+                                   c_m_n_imag_host_result.mData,
+                                   "Verification error: incorrect results in imaginary part!",
+                                   1e-2f,
+                                   1e-1f);
+        }
+        else
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+        {
+            result = ck::utils::check_err(c_m_n_real_device_result.mData,
+                                          c_m_n_real_host_result.mData,
+                                          "Verification error: incorrect results in real part!",
+                                          1e-2f,
+                                          1e-1f);
+            result = result && ck::utils::check_err(
+                                   c_m_n_imag_device_result.mData,
+                                   c_m_n_imag_host_result.mData,
+                                   "Verification error: incorrect results in imaginary part!",
+                                   1e-2f,
+                                   1e-1f);
+        }
+
+        return result;
     }
-    return 0;
+    return true;
 }

example/22_cgemm/cgemm_xdl_fp16.cpp

@@ -116,16 +116,16 @@ int main(int argc, char* argv[])
         exit(0);
     }
 
-    return run_cgemm_xdl<ADataType,
-                         BDataType,
-                         CDataType,
-                         ALayout,
-                         BLayout,
-                         CLayout,
-                         PassThrough,
-                         PassThrough,
-                         PassThrough,
-                         DeviceCGemmInstance,
-                         ReferenceCGemmInstance>(
+    return !run_cgemm_xdl<ADataType,
+                          BDataType,
+                          CDataType,
+                          ALayout,
+                          BLayout,
+                          CLayout,
+                          PassThrough,
+                          PassThrough,
+                          PassThrough,
+                          DeviceCGemmInstance,
+                          ReferenceCGemmInstance>(
         M, N, K, StrideA, StrideB, StrideC, do_verification, init_method, time_kernel);
 }

example/22_cgemm/cgemm_xdl_fp32.cpp

@@ -117,16 +117,16 @@ int main(int argc, char* argv[])
         exit(0);
     }
 
-    return run_cgemm_xdl<ADataType,
-                         BDataType,
-                         CDataType,
-                         ALayout,
-                         BLayout,
-                         CLayout,
-                         PassThrough,
-                         PassThrough,
-                         PassThrough,
-                         DeviceCGemmInstance,
-                         ReferenceCGemmInstance>(
+    return !run_cgemm_xdl<ADataType,
+                          BDataType,
+                          CDataType,
+                          ALayout,
+                          BLayout,
+                          CLayout,
+                          PassThrough,
+                          PassThrough,
+                          PassThrough,
+                          DeviceCGemmInstance,
+                          ReferenceCGemmInstance>(
         M, N, K, StrideA, StrideB, StrideC, do_verification, init_method, time_kernel);
 }

example/22_cgemm/cgemm_xdl_int4.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+
+#include "cgemm_xdl_common.hpp"
+
+#include "ck/library/reference_tensor_operation/cpu/reference_cgemm.hpp"
+
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/device/device_cgemm_4gemm_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+
+using ADataType        = INT4;
+using BDataType        = INT4;
+using CDataType        = INT4;
+using AccDataType      = INT32;
+using CShuffleDataType = INT32;
+
+using KernelADataType = INT8;
+using KernelBDataType = INT8;
+using KernelCDataType = INT8;
+
+using ALayout = ck::tensor_layout::gemm::RowMajor;
+using BLayout = ck::tensor_layout::gemm::ColumnMajor;
+using CLayout = ck::tensor_layout::gemm::RowMajor;
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+using ReferenceCGemmInstance = ck::tensor_operation::host::
+    ReferenceCGemm<ADataType, BDataType, CDataType, PassThrough, PassThrough, PassThrough>;
+
+// clang-format off
+using DeviceCGemmInstance = ck::tensor_operation::device::DeviceCGemm_4Gemm_Xdl_CShuffle
+    <ALayout,                    // typename ALayout
+     BLayout,                    // typename BLayout
+     CLayout,                    // typename CLayout
+     KernelADataType,            // typename ADataType
+     KernelBDataType,            // typename BDataType
+     KernelCDataType,            // typename CDataType
+     AccDataType,                // typename GemmAccDataType
+     CShuffleDataType,           // typename CShuffleDataType
+     PassThrough,                // typename AElementwiseOperation
+     PassThrough,                // typename BElementwiseOperation
+     PassThrough,                // typename CElementwiseOperation
+     GemmDefault,                // GemmSpecialization GemmSpec
+     1,                          // index_t NumGemmKPrefetchStage
+     256,                        // index_t BlockSize
+     256,                        // index_t MPerBlock
+     128,                        // index_t NPerBlock
+     64,                         // index_t KPerBlock
+     16,                         // index_t AK1
+     16,                         // index_t BK1
+     32,                         // index_t MPerXDL
+     32,                         // index_t NPerXDL
+     4,                          // index_t MXdlPerWave
+     2,                          // index_t NXdlPerWave
+     S<4, 64, 1>,                // typename ABlockTransferThreadClusterLengths_AK0_M_AK1
+     S<1, 0, 2>,                 // typename ABlockTransferThreadClusterArrangeOrder
+     S<1, 0, 2>,                 // typename ABlockTransferSrcAccessOrder
+     2,                          // index_t ABlockTransferSrcVectorDim
+     16,                         // index_t ABlockTransferSrcScalarPerVector
+     16,                         // index_t ABlockTransferDstScalarPerVector_AK1
+     1,                          // index_t ABlockLdsExtraM
+     S<4, 64, 1>,                // typename BBlockTransferThreadClusterLengths_BK0_N_BK1
+     S<1, 0, 2>,                 // typename BBlockTransferThreadClusterArrangeOrder
+     S<1, 0, 2>,                 // typename BBlockTransferSrcAccessOrder
+     2,                          // index_t BBlockTransferSrcVectorDim
+     8,                          // index_t BBlockTransferSrcScalarPerVector
+     8,                          // index_t BBlockTransferDstScalarPerVector_BK1
+     1,                          // index_t BBlockLdsExtraN
+     1,                          // index_t CShuffleMXdlPerWavePerShuffle
+     1,                          // index_t CShuffleNXdlPerWavePerShuffle
+     S<1, 64, 1, 4>,             // typename CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
+     16>;                        // index_t CShuffleBlockTransferScalarPerVector_NPerBlock
+// clang-format on
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = true;
+
+    // CGEMM shape
+    ck::index_t M = 1024;
+    ck::index_t N = 1152;
+    ck::index_t K = 512;
+
+    ck::index_t StrideA = K;
+    ck::index_t StrideB = K;
+    ck::index_t StrideC = N;
+
+    if(argc == 4)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+    }
+    else if(argc == 10)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+
+        M = std::stoi(argv[4]);
+        N = std::stoi(argv[5]);
+        K = std::stoi(argv[6]);
+
+        StrideA = std::stoi(argv[7]);
+        StrideB = std::stoi(argv[8]);
+        StrideC = std::stoi(argv[9]);
+    }
+    else
+    {
+        std::cout << "arg1: verification (0=no, 1=yes)\n"
+                  << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+                  << "arg3: time kernel (0=no, 1=yes)\n"
+                  << "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC\n"
+                  << std::endl;
+        exit(EXIT_SUCCESS);
+    }
+
+    return !run_cgemm_xdl<ADataType,
+                          BDataType,
+                          CDataType,
+                          ALayout,
+                          BLayout,
+                          CLayout,
+                          PassThrough,
+                          PassThrough,
+                          PassThrough,
+                          DeviceCGemmInstance,
+                          ReferenceCGemmInstance,
+                          KernelADataType,
+                          KernelBDataType,
+                          KernelCDataType>(
+        M, N, K, StrideA, StrideB, StrideC, do_verification, init_method, time_kernel);
+}

example/22_cgemm/cgemm_xdl_int8.cpp

@@ -117,16 +117,16 @@ int main(int argc, char* argv[])
         exit(0);
     }
 
-    return run_cgemm_xdl<ADataType,
-                         BDataType,
-                         CDataType,
-                         ALayout,
-                         BLayout,
-                         CLayout,
-                         PassThrough,
-                         PassThrough,
-                         PassThrough,
-                         DeviceCGemmInstance,
-                         ReferenceCGemmInstance>(
+    return !run_cgemm_xdl<ADataType,
+                          BDataType,
+                          CDataType,
+                          ALayout,
+                          BLayout,
+                          CLayout,
+                          PassThrough,
+                          PassThrough,
+                          PassThrough,
+                          DeviceCGemmInstance,
+                          ReferenceCGemmInstance>(
         M, N, K, StrideA, StrideB, StrideC, do_verification, init_method, time_kernel);
 }