merge

docs/sphinx/requirements.txt

@@ -42,7 +42,7 @@ fastjsonschema==2.18.0
     # via rocm-docs-core
 gitdb==4.0.10
     # via gitpython
-gitpython==3.1.31
+gitpython==3.1.35
     # via rocm-docs-core
 idna==3.4
     # via requests
@@ -103,7 +103,7 @@ requests==2.28.2
     # via
     #   pygithub
     #   sphinx
-rocm-docs-core>=0.20.0
+rocm-docs-core==0.24.0
     # via -r requirements.in
 six==1.16.0
     # via

example/01_gemm/CMakeLists.txt

@@ -66,21 +66,17 @@ endif()
 add_example_executable(example_gemm_xdl_streamk gemm_xdl_streamk.cpp)
 
 
-if(GPU_TARGETS MATCHES "gfx940" OR GPU_TARGETS MATCHES "gfx941" OR GPU_TARGETS MATCHES "gfx942")
-  add_example_executable(example_gemm_xdl_fp8 gemm_xdl_fp8.cpp)
-  if(result EQUAL 0)
+add_example_executable(example_gemm_xdl_fp8 gemm_xdl_fp8.cpp)
+if(result EQUAL 0)
     add_dependencies(example_gemm_xdl example_gemm_xdl_fp8)
-  endif()
 endif()
 
-if(GPU_TARGETS MATCHES "gfx940" OR GPU_TARGETS MATCHES "gfx941" OR GPU_TARGETS MATCHES "gfx942")
-  add_example_executable(example_gemm_xdl_fp8_bf8 gemm_xdl_fp8_bf8.cpp)
-  if(result EQUAL 0)
+add_example_executable(example_gemm_xdl_fp8_bf8 gemm_xdl_fp8_bf8.cpp)
+if(result EQUAL 0)
     add_dependencies(example_gemm_xdl example_gemm_xdl_fp8_bf8)
-  endif()
 endif()
 
 add_example_executable(example_gemm_xdl_fp16_fp8 gemm_xdl_fp16_fp8.cpp)
 if(result EQUAL 0)
-  add_dependencies(example_gemm_xdl example_gemm_xdl_fp16_fp8)
+    add_dependencies(example_gemm_xdl example_gemm_xdl_fp16_fp8)
 endif()

example/01_gemm/gemm_xdl_fp8.cpp

@@ -7,9 +7,9 @@
 
 using ADataType        = ck::f8_t;
 using BDataType        = ck::f8_t;
-using CDataType        = ck::f8_t;
+using CDataType        = ck::half_t;
 using AccDataType      = float;
-using CShuffleDataType = ck::f8_t;
+using CShuffleDataType = float;
 
 using ALayout = Row;
 using BLayout = Col;
@@ -27,7 +27,7 @@ using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemm_Xdl_CShuffle
 // ######|        |        |        |      Type|      Type|      Type|        Type|         DataType| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|         _MBlock_MWaveMPerXdl| ScalarPerVector|
 // ######|        |        |        |          |          |          |            |                 |   Operation|   Operation|   Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|         _NBlock_NWaveNPerXdl|   _NWaveNPerXdl|
 // ######|        |        |        |          |          |          |            |                 |            |            |            |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |
-         < ALayout, BLayout, CLayout, ADataType, BDataType, CDataType, AccDataType, CShuffleDataType,  AElementOp,  BElementOp,  CElementOp,    GemmDefault,        1,   256,   256,   128,    64,  16,  16,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,             16,             16,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 64, 1, 4>,              16>;
+         < ALayout, BLayout, CLayout, ADataType, BDataType, CDataType, AccDataType, CShuffleDataType,  AElementOp,  BElementOp,  CElementOp,    GemmDefault,        1,   256,   256,   128,    64,  16,  16,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,             16,             16,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 64, 1, 4>,               8>;
 // clang-format on
 
 using ReferenceGemmInstance = ck::tensor_operation::host::

example/01_gemm/gemm_xdl_fp8_bf8.cpp

@@ -7,9 +7,9 @@
 
 using ADataType        = ck::f8_t;
 using BDataType        = ck::bf8_t;
-using CDataType        = ck::f8_t;
+using CDataType        = ck::half_t;
 using AccDataType      = float;
-using CShuffleDataType = ck::f8_t;
+using CShuffleDataType = float;
 
 using ALayout = Row;
 using BLayout = Col;
@@ -31,7 +31,7 @@ using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemm_Xdl_CShuffle
 // ######|        |        |        |      Type|      Type|      Type|        Type|         DataType| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|         _MBlock_MWaveMPerXdl| ScalarPerVector|
 // ######|        |        |        |          |          |          |            |                 |   Operation|   Operation|   Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|         _NBlock_NWaveNPerXdl|   _NWaveNPerXdl|
 // ######|        |        |        |          |          |          |            |                 |            |            |            |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |
-         < ALayout, BLayout, CLayout, ADataType, BDataType, CDataType, AccDataType, CShuffleDataType,  AElementOp,  BElementOp,  CElementOp,    GemmDefault,        1,   256,   256,   128,    64,  16,  16,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,             16,             16,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 64, 1, 4>,              16,  LoopSched, PipelineVer, ComputeTypeA, ComputeTypeB>;
+         < ALayout, BLayout, CLayout, ADataType, BDataType, CDataType, AccDataType, CShuffleDataType,  AElementOp,  BElementOp,  CElementOp,    GemmDefault,        1,   256,   256,   128,    64,  16,  16,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,             16,             16,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 64, 1, 4>,               8,  LoopSched, PipelineVer, ComputeTypeA, ComputeTypeB>;
 // clang-format on
 
 using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,

example/60_gemm_multi_ABD/CMakeLists.txt

+if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
+    list(APPEND gpu_list2 gfx908 gfx90a gfx940 gfx941 gfx942)
+    set(target 0)
+    foreach(gpu IN LISTS GPU_TARGETS)
+        if(gpu IN_LIST gpu_list2 AND target EQUAL 0)
+            add_example_executable(example_gemm_multi_ABD_xdl_fp16 gemm_multi_ABD_xdl_fp16.cpp)
+            set(target 1)
+        endif()
+    endforeach()
+endif()

example/60_gemm_multi_ABD/gemm_multi_ABD_xdl_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_abd_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/utility/literals.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/library/utility/check_err.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using F16 = ck::half_t;
+using F32 = float;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+using ADataType        = F16;
+using BDataType        = F16;
+using AccDataType      = F32;
+using CShuffleDataType = F32;
+using DDataType        = F16;
+using EDataType        = F16;
+
+using ALayout = Row;
+using BLayout = Col;
+using DLayout = Row;
+using ELayout = Row;
+
+struct AddScale
+{
+    static constexpr auto I0 = ck::Number<0>{};
+    static constexpr auto I1 = ck::Number<1>{};
+    static constexpr auto I2 = ck::Number<2>{};
+    static constexpr auto I3 = ck::Number<3>{};
+
+    __host__ __device__ constexpr void
+    operator()(ck::half4_t& a, const ck::half4_t& a0, const ck::half4_t& a1) const
+    {
+        const auto a0_v_t = ck::vector_type<ck::half_t, 4>{a0};
+        const auto a1_v_t = ck::vector_type<ck::half_t, 4>{a1};
+
+        auto r_v_t = ck::vector_type<ck::half_t, 4>{};
+
+        r_v_t.AsType<ck::half_t>()(I0) =
+            scale * (a0_v_t.AsType<ck::half_t>()[I0] + a1_v_t.AsType<ck::half_t>()[I0]);
+        r_v_t.AsType<ck::half_t>()(I1) =
+            scale * (a0_v_t.AsType<ck::half_t>()[I1] + a1_v_t.AsType<ck::half_t>()[I1]);
+        r_v_t.AsType<ck::half_t>()(I2) =
+            scale * (a0_v_t.AsType<ck::half_t>()[I2] + a1_v_t.AsType<ck::half_t>()[I2]);
+        r_v_t.AsType<ck::half_t>()(I3) =
+            scale * (a0_v_t.AsType<ck::half_t>()[I3] + a1_v_t.AsType<ck::half_t>()[I3]);
+
+        a = r_v_t.AsType<ck::half4_t>()[I0];
+    }
+
+    __host__ __device__ constexpr void
+    operator()(ck::half_t& a, const ck::half_t& a0, const ck::half_t& a1) const
+    {
+        a = scale * (a0 + a1);
+    }
+
+    // this attribute controls the copy_function applying element_wise_op with
+    // pack4_data
+    constexpr const static bool is_pack4_invocable = true;
+
+    float scale = 1.0;
+};
+
+struct AlphaBetaAdd
+{
+    AlphaBetaAdd(float alpha, float beta) : alpha_(alpha), beta_(beta){};
+
+    template <typename E, typename C, typename D>
+    __host__ __device__ constexpr void operator()(E& e, const C& c, const D& d) const;
+
+    template <>
+    __host__ __device__ constexpr void operator()<ck::half_t, float, ck::half_t>(
+        ck::half_t& e, const float& c, const ck::half_t& d) const
+    {
+        e = ck::type_convert<ck::half_t>(alpha_ * c + beta_ * ck::type_convert<float>(d));
+    };
+
+    float alpha_;
+    float beta_;
+};
+
+using AElementOp   = AddScale;
+using BElementOp   = PassThrough;
+using CDEElementOp = AlphaBetaAdd;
+
+static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKPadding;
+
+using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Xdl_CShuffle<
+    ck::Tuple<ALayout, ALayout>,
+    ck::Tuple<BLayout>,
+    ck::Tuple<DLayout>,
+    ELayout,
+    ck::Tuple<ADataType, ADataType>,
+    ck::Tuple<BDataType>,
+    AccDataType,
+    CShuffleDataType,
+    ck::Tuple<DDataType>,
+    EDataType,
+    AElementOp,
+    BElementOp,
+    CDEElementOp,
+    GemmSpec,
+    1,
+    256,
+    256,
+    128,
+    32,
+    8,
+    8,
+    32,
+    32,
+    4,
+    2,
+    S<4, 64, 1>,
+    S<1, 0, 2>,
+    S<1, 0, 2>,
+    2,
+    8,
+    8,
+    1,
+    S<4, 64, 1>,
+    S<1, 0, 2>,
+    S<1, 0, 2>,
+    2,
+    8,
+    8,
+    1,
+    1,
+    1,
+    S<1, 32, 1, 8>,
+    8>;
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = false;
+
+    // GEMM shape
+    ck::index_t M = 3840;
+    ck::index_t N = 4096;
+    ck::index_t K = 4096;
+
+    ck::index_t StrideA = 4096;
+    ck::index_t StrideB = 4096;
+    ck::index_t StrideD = 4096;
+    ck::index_t StrideE = 4096;
+
+    float alpha = 1.0f;
+    float beta  = 1.0f;
+
+    if(argc == 1)
+    {
+        // use default case
+    }
+    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 == 6)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+
+        alpha = std::stof(argv[4]);
+        beta  = std::stof(argv[5]);
+    }
+    else if(argc == 13)
+    {
+        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]);
+        StrideD = std::stoi(argv[9]);
+        StrideE = std::stoi(argv[10]);
+
+        alpha = std::stof(argv[11]);
+        beta  = std::stof(argv[12]);
+    }
+    else
+    {
+        printf("arg1: verification (0=no, 1=yes)\n");
+        printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
+        printf("arg3: time kernel (0=no, 1=yes)\n");
+        printf("arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideD, StrideE, alpha, "
+               "beta\n");
+        exit(0);
+    }
+
+    auto f_host_tensor_descriptor =
+        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+            using namespace ck::literals;
+
+            if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+            {
+                return HostTensorDescriptor({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    Tensor<ADataType> a0_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
+    Tensor<ADataType> a1_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
+    Tensor<BDataType> b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
+    Tensor<DDataType> d_m_n(f_host_tensor_descriptor(M, N, StrideD, DLayout{}));
+    Tensor<EDataType> e_m_n_host_result(f_host_tensor_descriptor(M, N, StrideE, ELayout{}));
+    Tensor<EDataType> e_m_n_device_result(f_host_tensor_descriptor(M, N, StrideE, ELayout{}));
+
+    std::cout << "a0_m_k: " << a0_m_k.mDesc << std::endl;
+    std::cout << "a1_m_k: " << a1_m_k.mDesc << std::endl;
+    std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
+    std::cout << "d_m_n: " << d_m_n.mDesc << std::endl;
+    std::cout << "e_m_n: " << e_m_n_host_result.mDesc << std::endl;
+
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        a0_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
+        a1_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
+        b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
+        d_m_n.GenerateTensorValue(GeneratorTensor_2<DDataType>{-5, 5});
+        break;
+    default:
+        a0_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
+        a1_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
+        b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
+        d_m_n.GenerateTensorValue(GeneratorTensor_3<DDataType>{-0.5, 0.5});
+    }
+
+    DeviceMem a0_device_buf(sizeof(ADataType) * a0_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem a1_device_buf(sizeof(ADataType) * a1_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
+    DeviceMem d_device_buf(sizeof(DDataType) * d_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n_device_result.mDesc.GetElementSpaceSize());
+
+    a0_device_buf.ToDevice(a0_m_k.mData.data());
+    a1_device_buf.ToDevice(a1_m_k.mData.data());
+    b_device_buf.ToDevice(b_k_n.mData.data());
+    d_device_buf.ToDevice(d_m_n.mData.data());
+    e_device_buf.ToDevice(e_m_n_device_result.mData.data());
+
+    auto a_element_op   = AElementOp{0.2};
+    auto b_element_op   = BElementOp{};
+    auto cde_element_op = CDEElementOp{alpha, beta};
+
+    // do GEMM
+    auto device_op = DeviceOpInstance{};
+    auto invoker   = device_op.MakeInvoker();
+    auto argument =
+        device_op.MakeArgument(std::array<const void*, 2>{a0_device_buf.GetDeviceBuffer(),
+                                                          a1_device_buf.GetDeviceBuffer()},
+                               std::array<const void*, 1>{b_device_buf.GetDeviceBuffer()},
+                               std::array<const void*, 1>{d_device_buf.GetDeviceBuffer()},
+                               e_device_buf.GetDeviceBuffer(),
+                               M,
+                               N,
+                               K,
+                               std::array<ck::index_t, 2>{StrideA, StrideA},
+                               std::array<ck::index_t, 1>{StrideB},
+                               std::array<ck::index_t, 1>{StrideD},
+                               StrideE,
+                               a_element_op,
+                               b_element_op,
+                               cde_element_op);
+
+    if(!device_op.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error(
+            "wrong! device_gemm with the specified compilation parameters does "
+            "not support this GEMM problem");
+    }
+
+    float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+
+    std::size_t flop = std::size_t(2) * M * N * K;
+    std::size_t num_btype =
+        sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(EDataType) * M * N;
+
+    float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
+
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
+              << std::endl;
+
+    e_device_buf.FromDevice(e_m_n_device_result.mData.data());
+
+    if(do_verification)
+    {
+        Tensor<CShuffleDataType> c_m_n({M, N});
+
+        Tensor<ADataType> a_m_k({M, K});
+
+        for(int m = 0; m < M; ++m)
+        {
+            for(int k = 0; k < K; ++k)
+            {
+                a_element_op(a_m_k(m, k), a0_m_k(m, k), a1_m_k(m, k));
+            }
+        }
+
+        using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                                BDataType,
+                                                                                CShuffleDataType,
+                                                                                AccDataType,
+                                                                                PassThrough,
+                                                                                BElementOp,
+                                                                                PassThrough>;
+        auto ref_gemm               = ReferenceGemmInstance{};
+        auto ref_invoker            = ref_gemm.MakeInvoker();
+
+        auto ref_argument =
+            ref_gemm.MakeArgument(a_m_k, b_k_n, c_m_n, PassThrough{}, b_element_op, PassThrough{});
+
+        ref_invoker.Run(ref_argument);
+
+        for(int m = 0; m < M; ++m)
+        {
+            for(int n = 0; n < N; ++n)
+            {
+                cde_element_op(e_m_n_host_result(m, n), c_m_n(m, n), d_m_n(m, n));
+            }
+        }
+
+        e_device_buf.FromDevice(e_m_n_device_result.mData.data());
+
+        return ck::utils::check_err(e_m_n_device_result, e_m_n_host_result) ? 0 : 1;
+    }
+
+    return 0;
+}

include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp

@@ -173,8 +173,7 @@ struct PassThrough
     template <>
     __host__ __device__ void operator()<bf8_t, half_t>(bf8_t& y, const half_t& x) const
     {
-        // to-do: fix half_t to bf8_t convert
-        y = ck::type_convert<bf8_t>(ck::type_convert<float>(x));
+        y = ck::type_convert<bf8_t>(x);
     }
 #endif
 };

include/ck/tensor_operation/gpu/grid/gridwise_gemm_multiple_abd_xdl_cshuffle.hpp

@@ -658,8 +658,8 @@ struct GridwiseGemmMultipleABD_xdl_cshuffle
 
         auto blockwise_gemm = BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_Selector<
             BlockSize,
-            ComputeDataType,
-            ComputeDataType,
+            ComputeDataType, // ComputeDataType for A
+            ComputeDataType, // ComputeDataType for B
             AccDataType,
             decltype(a_block_desc_ak0_m_ak1),
             decltype(b_block_desc_bk0_n_bk1),

include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp

@@ -945,7 +945,8 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3_ext
             }
         }();
 
-        if constexpr(GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding)
+        if constexpr(GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||
+                     GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding)
         {
             return transform_tensor_descriptor(c_grid_desc_m_n,
                                                make_tuple(make_right_pad_transform(M, MPad - M),

include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp

@@ -9,6 +9,7 @@
 #include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
 #include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp"
 #include "ck/tensor/static_tensor.hpp"
+#include "ck/utility/is_detected.hpp"
 
 namespace ck {
 
@@ -211,10 +212,44 @@ struct ThreadwiseTensorSliceTransfer_v3r1
             auto src_vector_container = src_vector_type{
                 src_buf.template Get<src_vector_t>(src_coord_.GetOffset(), is_src_valid)};
 
+            using dst_vector_type = vector_type_maker_t<DstData, SrcScalarPerVector>;
+            using dst_vector_t    = typename dst_vector_type::type;
+            dst_vector_type op_r_v;
+
+            constexpr auto get_elem_op_vec_len = []() {
+                if constexpr(is_detected<is_pack8_invocable_t, decltype(src_element_op_)>::value)
+                {
+                    if constexpr(decltype(src_element_op_)::is_pack8_invocable)
+                        return math::min(8, SrcScalarPerVector);
+                }
+                if constexpr(is_detected<is_pack4_invocable_t, decltype(src_element_op_)>::value)
+                {
+                    if constexpr(decltype(src_element_op_)::is_pack4_invocable)
+                        return math::min(4, SrcScalarPerVector);
+                }
+                if constexpr(is_detected<is_pack2_invocable_t, decltype(src_element_op_)>::value)
+                {
+                    if constexpr(decltype(src_element_op_)::is_pack2_invocable)
+                        return math::min(2, SrcScalarPerVector);
+                }
+                return 1;
+            };
+
+            constexpr index_t elem_op_vec_len = get_elem_op_vec_len();
+
+            using src_elem_op_vec_t = typename vector_type<SrcData, elem_op_vec_len>::type;
+            using dst_elem_op_vec_t = typename vector_type<DstData, elem_op_vec_len>::type;
+
+            static_for<0, SrcScalarPerVector / elem_op_vec_len, 1>{}([&](auto idx) {
+                // apply the src elementwise op and convert to DstData under the hood if needed
+                src_element_op_(op_r_v.template AsType<dst_elem_op_vec_t>()(idx),
+                                src_vector_container.template AsType<src_elem_op_vec_t>()[idx]);
+            });
+
             // copy data from src_vector_container into src_thread_scratch_
             src_thread_scratch_tuple_(thread_scratch_id)
-                .template SetAsType<src_vector_t>(
-                    src_data_idx_seq, src_vector_container.template AsType<src_vector_t>()[I0]);
+                .template SetAsType<dst_vector_t>(src_data_idx_seq,
+                                                  op_r_v.template AsType<dst_vector_t>()[I0]);
 
             constexpr auto move_on_dim = [&]() constexpr
             {
@@ -267,19 +302,15 @@ struct ThreadwiseTensorSliceTransfer_v3r1
     {
 #if !CK_EXPERIMENTAL_USE_IN_REGISTER_SUB_DWORD_TRANSPOSE
         static_ford<SliceLengths>{}([&](auto idx) {
-            // convert from SrcData to DstData here
-            dst_thread_scratch_(idx) =
-                type_convert<DstData>(src_thread_scratch_tuple_[thread_scratch_id][idx]);
+            dst_thread_scratch_(idx) = src_thread_scratch_tuple_[thread_scratch_id][idx];
         });
 #else
         // sub-dword transpose between src_thread_scratch_ and dst_thread_scratch_
         // TODO make this logic more generic for more sub-dword datatype
         if constexpr(SrcVectorDim != DstVectorDim &&
-                     ((is_same<half_t, remove_cvref_t<SrcData>>::value &&
-                       is_same<half_t, remove_cvref_t<DstData>>::value &&
+                     ((is_same<half_t, remove_cvref_t<DstData>>::value &&
                        SrcScalarPerVector % 2 == 0 && DstScalarPerVector % 2 == 0) ||
-                      (is_same<int8_t, remove_cvref_t<SrcData>>::value &&
-                       is_same<int8_t, remove_cvref_t<DstData>>::value &&
+                      (is_same<int8_t, remove_cvref_t<DstData>>::value &&
                        SrcScalarPerVector % 4 == 0 && DstScalarPerVector % 4 == 0)))
         {
             // each transpose does
@@ -313,7 +344,7 @@ struct ThreadwiseTensorSliceTransfer_v3r1
                 constexpr auto data_idx_seq = generate_sequence_v2(
                     [&](auto i) { return Number<data_idx[i]>{}; }, Number<nDim>{});
 
-                using src_vector_t = vector_type_maker_t<SrcData, SrcScalarPerVector>;
+                using src_vector_t = vector_type_maker_t<DstData, SrcScalarPerVector>;
                 using dst_vector_t = vector_type_maker_t<DstData, DstScalarPerVector>;
 
                 // get DstScalarPerVector # of read-only references to src vectors from
@@ -336,17 +367,16 @@ struct ThreadwiseTensorSliceTransfer_v3r1
                     Number<num_dst_vector>{});
 
                 // do data transpose
-                transpose_vectors<SrcData, DstScalarPerVector, SrcScalarPerVector>{}(
+                transpose_vectors<DstData, DstScalarPerVector, SrcScalarPerVector>{}(
                     src_vector_refs, dst_vector_refs);
             });
         }
-
-        static_ford<SliceLengths>{}([&](auto idx) {
-            // apply the src elementwise op and convert to DstData under the hood if needed
-            DstData dst_v;
-            src_element_op_(dst_v, src_thread_scratch_tuple_[thread_scratch_id][idx]);
-            dst_thread_scratch_(idx) = dst_v;
-        });
+        else
+        {
+            static_ford<SliceLengths>{}([&](auto idx) {
+                dst_thread_scratch_(idx) = src_thread_scratch_tuple_[thread_scratch_id][idx];
+            });
+        }
 #endif
     }
 
@@ -761,11 +791,12 @@ struct ThreadwiseTensorSliceTransfer_v3r1
     static constexpr auto src_thread_scratch_desc_ = decltype(GetSrcThreadScratchDescriptor()){};
     static constexpr auto dst_thread_scratch_desc_ = decltype(GetDstThreadScratchDescriptor()){};
 
-    using SrcThreadScratch = StaticTensorTupleOfVectorBuffer<AddressSpaceEnum::Vgpr,
-                                                             SrcData,
-                                                             SrcScalarPerVector,
-                                                             decltype(src_thread_scratch_desc_),
-                                                             true>;
+    using SrcThreadScratch =
+        StaticTensorTupleOfVectorBuffer<AddressSpaceEnum::Vgpr,
+                                        DstData, // apply data_convert with SrcThreadScratch
+                                        SrcScalarPerVector,
+                                        decltype(src_thread_scratch_desc_),
+                                        true>;
 
     using DstThreadScratch = StaticTensorTupleOfVectorBuffer<AddressSpaceEnum::Vgpr,
                                                              DstData,

include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v7r2.hpp

@@ -132,9 +132,6 @@ struct ThreadwiseTensorSliceTransfer_v7r2
             Number<num>{});
     }
 
-    template <typename T>
-    using has_vec_len = decltype(std::declval<T&>().vec_len);
-
     // SrcDescs: Tuple<const SrcDesc0&, const SrcDesc1&, ...>
     // SrcBuffers: Tuple<const SrcBuffer0&, const SrcBuffer1&, ...>
     template <typename SrcBuffers,
@@ -159,94 +156,63 @@ struct ThreadwiseTensorSliceTransfer_v7r2
                                                            is_src_valid);
             });
 
-            if constexpr(is_detected<has_vec_len, decltype(element_op_)>::value)
-            {
-                constexpr auto elem_op_vec_len = decltype(element_op_)::vec_len;
-
-                static_assert(is_same<remove_cvref_t<decltype(elem_op_vec_len)>, index_t>::value,
-                              "vec_len in element_op_ type is not index_t");
+            constexpr auto get_elem_op_vec_len = []() {
+                if constexpr(is_detected<is_pack8_invocable_t, decltype(element_op_)>::value)
+                {
+                    if constexpr(decltype(element_op_)::is_pack8_invocable)
+                        return math::min(8, SrcScalarPerVector);
+                }
+                if constexpr(is_detected<is_pack4_invocable_t, decltype(element_op_)>::value)
+                {
+                    if constexpr(decltype(element_op_)::is_pack4_invocable)
+                        return math::min(4, SrcScalarPerVector);
+                }
+                if constexpr(is_detected<is_pack2_invocable_t, decltype(element_op_)>::value)
+                {
+                    if constexpr(decltype(element_op_)::is_pack2_invocable)
+                        return math::min(2, SrcScalarPerVector);
+                }
+                return 1;
+            };
+
+            constexpr index_t elem_op_vec_len = get_elem_op_vec_len();
+
+            // apply pointwise function
+            static_for<0, SrcScalarPerVector / elem_op_vec_len, 1>{}([&](auto i) {
+                // get reference to src data
+                const auto src_data_refs = generate_tie(
+                    // return type should be lvalue
+                    [&](auto iSrc) -> const auto& {
+                        using SrcData = remove_cvref_t<tuple_element_t<iSrc.value, SrcDatas>>;
+
+                        using elem_op_vec_t = typename vector_type<SrcData, elem_op_vec_len>::type;
+
+                        return src_vectors[iSrc].template AsType<elem_op_vec_t>()[i];
+                    },
+                    Number<nSrc>{});
+
+                // get reference to dst data
+                auto dst_data_refs = generate_tie(
+                    // return type should be lvalue
+                    [&](auto iDst) -> auto& {
+                        using DstData = remove_cvref_t<tuple_element_t<iDst.value, DstDatas>>;
+
+                        using elem_op_vec_t = typename vector_type<DstData, elem_op_vec_len>::type;
+
+                        return dst_vectors(iDst).template AsType<elem_op_vec_t>()(i);
+                    },
+                    Number<nDst>{});
 
-                static_assert(elem_op_vec_len == 1 || elem_op_vec_len == 2 ||
-                                  elem_op_vec_len == 4 || elem_op_vec_len == 8,
-                              "vec_len in element_op_ must be 1, 2, 4, 8");
-
-                static_assert(SrcScalarPerVector % elem_op_vec_len == 0,
-                              "vec_len in element_op_ cannot be divided by SrcScalarPerVector!");
-
-                // apply pointwise function
-                static_for<0, SrcScalarPerVector / elem_op_vec_len, 1>{}([&](auto i) {
-                    // get reference to src data
-                    const auto src_data_refs = generate_tie(
-                        // return type should be lvalue
-                        [&](auto iSrc) -> const auto& {
-                            using SrcData = remove_cvref_t<tuple_element_t<iSrc.value, SrcDatas>>;
-
-                            using elem_op_vec_t =
-                                typename vector_type<SrcData, elem_op_vec_len>::type;
-
-                            return src_vectors[iSrc].template AsType<elem_op_vec_t>()[i];
-                        },
-                        Number<nSrc>{});
-
-                    // get reference to dst data
-                    auto dst_data_refs = generate_tie(
-                        // return type should be lvalue
-                        [&](auto iDst) -> auto& {
-                            using DstData = remove_cvref_t<tuple_element_t<iDst.value, DstDatas>>;
-
-                            using elem_op_vec_t =
-                                typename vector_type<DstData, elem_op_vec_len>::type;
-
-                            return dst_vectors(iDst).template AsType<elem_op_vec_t>()(i);
-                        },
-                        Number<nDst>{});
-
-                    // apply pointwise function
-                    // pointwise function signature:
-                    // element_op_(dst_data_refs[I0],
-                    //             dst_data_refs[I1],
-                    //             ...,
-                    //             src_data_refs[I0],
-                    //             src_data_refs[I1],
-                    //             ...)
-                    unpack2(element_op_, dst_data_refs, src_data_refs);
-                });
-            }
-            else
-            {
                 // apply pointwise function
-                static_for<0, SrcScalarPerVector, 1>{}([&](auto i) {
-                    // get reference to src data
-                    const auto src_data_refs = generate_tie(
-                        // return type should be lvalue
-                        [&](auto iSrc) -> const auto& {
-                            using SrcData = remove_cvref_t<tuple_element_t<iSrc.value, SrcDatas>>;
-
-                            return src_vectors[iSrc].template AsType<SrcData>()[i];
-                        },
-                        Number<nSrc>{});
-
-                    // get reference to dst data
-                    auto dst_data_refs = generate_tie(
-                        // return type should be lvalue
-                        [&](auto iDst) -> auto& {
-                            using DstData = remove_cvref_t<tuple_element_t<iDst.value, DstDatas>>;
-
-                            return dst_vectors(iDst).template AsType<DstData>()(i);
-                        },
-                        Number<nDst>{});
-
-                    // apply pointwise function
-                    // pointwise function signature:
-                    // element_op_(dst_data_refs[I0],
-                    //             dst_data_refs[I1],
-                    //             ...,
-                    //             src_data_refs[I0],
-                    //             src_data_refs[I1],
-                    //             ...)
-                    unpack2(element_op_, dst_data_refs, src_data_refs);
-                });
-            }
+                // pointwise function signature:
+                // element_op_(dst_data_refs[I0],
+                //             dst_data_refs[I1],
+                //             ...,
+                //             src_data_refs[I0],
+                //             src_data_refs[I1],
+                //             ...)
+                unpack2(element_op_, dst_data_refs, src_data_refs);
+            });
 
             dst_vectors_tuple_(iAccess) = dst_vectors;
 

include/ck/utility/amd_buffer_addressing.hpp

@@ -299,584 +299,255 @@ enum struct AmdBufferCoherenceEnum
     GLC_SLC          = 3,
 };
 
-template <typename T,
-          index_t N,
-          AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence>
-__device__ typename vector_type<T, N>::type amd_buffer_load_impl(int32x4_t src_wave_buffer_resource,
-                                                                 index_t src_thread_addr_offset,
-                                                                 index_t src_wave_addr_offset)
+template <index_t N, AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence>
+__device__ typename vector_type<int8_t, N>::type
+amd_buffer_load_impl_raw(int32x4_t src_wave_buffer_resource,
+                         index_t src_thread_addr_offset,
+                         index_t src_wave_addr_offset)
 {
-    static_assert(
-        (is_same<T, double>::value && (N == 1 || N == 2 || N == 4)) ||
-            (is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
-            (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
-            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
-            (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
-            (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
-        "wrong! not implemented");
+    static_assert(N == 1 || N == 2 || N == 4 || N == 8 || N == 16 || N == 32 || N == 64,
+                  "wrong! not implemented");
 
-    if constexpr(is_same<T, double>::value)
+    if constexpr(N == 1)
     {
-        // use fp32 load to mimic fp64 load
-        if constexpr(N == 1)
-        {
-            const float2_t tmp =
-                llvm_amdgcn_raw_buffer_load_fp32x2(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset,
-                                                   static_cast<index_t>(coherence));
-
-            return bit_cast<double>(tmp);
-        }
-        else if constexpr(N == 2)
-        {
-            const float4_t tmp =
-                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset,
-                                                   static_cast<index_t>(coherence));
-
-            return bit_cast<double2_t>(tmp);
-        }
-        else if constexpr(N == 4)
-        {
-            const float4_t f32_0 =
-                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset,
-                                                   static_cast<index_t>(coherence));
-
-            const float4_t f32_1 =
-                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset + 4 * sizeof(float),
-                                                   static_cast<index_t>(coherence));
-            vector_type<double, 4> tmp;
-
-            tmp.AsType<double2_t>()(Number<0>{}) = bit_cast<double2_t>(f32_0);
-            tmp.AsType<double2_t>()(Number<1>{}) = bit_cast<double2_t>(f32_1);
-
-            return tmp.AsType<double4_t>()(Number<0>{});
-        }
+        return llvm_amdgcn_raw_buffer_load_i8(src_wave_buffer_resource,
+                                              src_thread_addr_offset,
+                                              src_wave_addr_offset,
+                                              static_cast<index_t>(coherence));
     }
-    else if constexpr(is_same<T, float>::value)
+    else if constexpr(N == 2)
     {
-        if constexpr(N == 1)
-        {
-            return llvm_amdgcn_raw_buffer_load_fp32(src_wave_buffer_resource,
-                                                    src_thread_addr_offset,
-                                                    src_wave_addr_offset,
-                                                    static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            return llvm_amdgcn_raw_buffer_load_fp32x2(src_wave_buffer_resource,
-                                                      src_thread_addr_offset,
-                                                      src_wave_addr_offset,
-                                                      static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            return llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
+
+        int16_t tmp = llvm_amdgcn_raw_buffer_load_i16(src_wave_buffer_resource,
                                                       src_thread_addr_offset,
                                                       src_wave_addr_offset,
                                                       static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 8)
-        {
-            vector_type<float, 8> tmp;
-
-            tmp.AsType<float4_t>()(Number<0>{}) =
-                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset,
-                                                   static_cast<index_t>(coherence));
 
-            tmp.AsType<float4_t>()(Number<1>{}) =
-                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset + 4 * sizeof(float),
-                                                   static_cast<index_t>(coherence));
-
-            return tmp.AsType<float8_t>()(Number<0>{});
-        }
+        return bit_cast<int8x2_t>(tmp);
     }
-    else if constexpr(is_same<T, half_t>::value)
+    else if constexpr(N == 4)
     {
-        if constexpr(N == 1)
-        {
-            return llvm_amdgcn_raw_buffer_load_fp16(src_wave_buffer_resource,
-                                                    src_thread_addr_offset,
-                                                    src_wave_addr_offset,
-                                                    static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            return llvm_amdgcn_raw_buffer_load_fp16x2(src_wave_buffer_resource,
-                                                      src_thread_addr_offset,
-                                                      src_wave_addr_offset,
-                                                      static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            return llvm_amdgcn_raw_buffer_load_fp16x4(src_wave_buffer_resource,
+        int32_t tmp = llvm_amdgcn_raw_buffer_load_i32(src_wave_buffer_resource,
                                                       src_thread_addr_offset,
                                                       src_wave_addr_offset,
                                                       static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 8)
-        {
-            // use fp32 load to mimic fp16 load
-            float4_t tmp = llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
-                                                              src_thread_addr_offset,
-                                                              src_wave_addr_offset,
-                                                              static_cast<index_t>(coherence));
-
-            return bit_cast<half8_t>(tmp);
-        }
-    }
-    else if constexpr(is_same<T, bhalf_t>::value)
-    {
-        if constexpr(N == 1)
-        {
-            return llvm_amdgcn_raw_buffer_load_i16(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset,
-                                                   static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            return llvm_amdgcn_raw_buffer_load_i16x2(src_wave_buffer_resource,
-                                                     src_thread_addr_offset,
-                                                     src_wave_addr_offset,
-                                                     static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            return llvm_amdgcn_raw_buffer_load_i16x4(src_wave_buffer_resource,
-                                                     src_thread_addr_offset,
-                                                     src_wave_addr_offset,
-                                                     static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 8)
-        {
-            int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
-                                                              src_thread_addr_offset,
-                                                              src_wave_addr_offset,
-                                                              static_cast<index_t>(coherence));
 
-            return bit_cast<bhalf8_t>(tmp);
-        }
+        return bit_cast<int8x4_t>(tmp);
     }
-    else if constexpr(is_same<T, int32_t>::value)
+    else if constexpr(N == 8)
     {
-        if constexpr(N == 1)
-        {
-            return llvm_amdgcn_raw_buffer_load_i32(src_wave_buffer_resource,
-                                                   src_thread_addr_offset,
-                                                   src_wave_addr_offset,
-                                                   static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            return llvm_amdgcn_raw_buffer_load_i32x2(src_wave_buffer_resource,
-                                                     src_thread_addr_offset,
-                                                     src_wave_addr_offset,
-                                                     static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            return llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
-                                                     src_thread_addr_offset,
-                                                     src_wave_addr_offset,
-                                                     static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 8)
-        {
-            vector_type<int32_t, 8> tmp;
-
-            tmp.AsType<int32x4_t>()(Number<0>{}) =
-                llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
-                                                  src_thread_addr_offset,
-                                                  src_wave_addr_offset,
-                                                  static_cast<index_t>(coherence));
-
-            tmp.AsType<int32x4_t>()(Number<1>{}) =
-                llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
-                                                  src_thread_addr_offset,
-                                                  src_wave_addr_offset + 4 * sizeof(int32_t),
-                                                  static_cast<index_t>(coherence));
-            return tmp.AsType<int32x8_t>()(Number<0>{});
-        }
-    }
-    else if constexpr(is_same<T, int8_t>::value)
-    {
-        if constexpr(N == 1)
-        {
-            return llvm_amdgcn_raw_buffer_load_i8(src_wave_buffer_resource,
-                                                  src_thread_addr_offset,
-                                                  src_wave_addr_offset,
-                                                  static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
-            return llvm_amdgcn_raw_buffer_load_i8x2(src_wave_buffer_resource,
-                                                    src_thread_addr_offset,
-                                                    src_wave_addr_offset,
-                                                    static_cast<index_t>(coherence));
-#else
-            int16_t tmp = llvm_amdgcn_raw_buffer_load_i16(src_wave_buffer_resource,
+        int32x2_t tmp = llvm_amdgcn_raw_buffer_load_i32x2(src_wave_buffer_resource,
                                                           src_thread_addr_offset,
                                                           src_wave_addr_offset,
                                                           static_cast<index_t>(coherence));
 
-            return bit_cast<int8x2_t>(tmp);
-#endif
-        }
-        else if constexpr(N == 4)
-        {
-#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
-            return llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
-                                                    src_thread_addr_offset,
-                                                    src_wave_addr_offset,
-                                                    static_cast<index_t>(coherence));
-#else
-            int32_t tmp = llvm_amdgcn_raw_buffer_load_i32(src_wave_buffer_resource,
+        return bit_cast<int8x8_t>(tmp);
+    }
+    else if constexpr(N == 16)
+    {
+        int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
                                                           src_thread_addr_offset,
                                                           src_wave_addr_offset,
                                                           static_cast<index_t>(coherence));
+        return bit_cast<int8x16_t>(tmp);
+    }
+    else if constexpr(N == 32)
+    {
+        int32x4_t tmp0 = llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
+                                                           src_thread_addr_offset,
+                                                           src_wave_addr_offset,
+                                                           static_cast<index_t>(coherence));
+        int32x4_t tmp1 =
+            llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
+                                              src_thread_addr_offset,
+                                              src_wave_addr_offset + 4 * sizeof(int32_t),
+                                              static_cast<index_t>(coherence));
+        vector_type<int32_t, 8> tmp;
 
-            return bit_cast<int8x4_t>(tmp);
-#endif
-        }
-        else if constexpr(N == 8)
-        {
-#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
-            vector_type<int8_t, 8> tmp;
-
-            tmp.AsType<int8x4_t>()(Number<0>{}) =
-                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
-                                                 src_thread_addr_offset,
-                                                 src_wave_addr_offset,
-                                                 static_cast<index_t>(coherence));
-
-            tmp.AsType<int8x4_t>()(Number<1>{}) =
-                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
-                                                 src_thread_addr_offset,
-                                                 src_wave_addr_offset + 4 * sizeof(int8_t),
-                                                 static_cast<index_t>(coherence));
-
-            return tmp.AsType<int8x8_t>()(Number<0>{});
-#else
-            int32x2_t tmp = llvm_amdgcn_raw_buffer_load_i32x2(src_wave_buffer_resource,
-                                                              src_thread_addr_offset,
-                                                              src_wave_addr_offset,
-                                                              static_cast<index_t>(coherence));
+        tmp.AsType<int32x4_t>()(Number<0>{}) = tmp0;
+        tmp.AsType<int32x4_t>()(Number<1>{}) = tmp1;
 
-            return bit_cast<int8x8_t>(tmp);
-#endif
-        }
-        else if constexpr(N == 16)
-        {
-#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
-            vector_type<int8_t, 16> tmp;
-
-            tmp.AsType<int8x4_t>()(Number<0>{}) =
-                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
-                                                 src_thread_addr_offset,
-                                                 src_wave_addr_offset,
-                                                 static_cast<index_t>(coherence));
-
-            tmp.AsType<int8x4_t>()(Number<1>{}) =
-                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
-                                                 src_thread_addr_offset,
-                                                 src_wave_addr_offset + 4 * sizeof(int8_t),
-                                                 static_cast<index_t>(coherence));
-
-            tmp.AsType<int8x4_t>()(Number<2>{}) =
-                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
-                                                 src_thread_addr_offset,
-                                                 src_wave_addr_offset + 8 * sizeof(int8_t),
-                                                 static_cast<index_t>(coherence));
-
-            tmp.AsType<int8x4_t>()(Number<3>{}) =
-                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
-                                                 src_thread_addr_offset,
-                                                 src_wave_addr_offset + 12 * sizeof(int8_t),
-                                                 static_cast<index_t>(coherence));
-
-            return tmp.AsType<int8x16_t>()(Number<0>{});
-#else
-            int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
-                                                              src_thread_addr_offset,
-                                                              src_wave_addr_offset,
-                                                              static_cast<index_t>(coherence));
+        return bit_cast<int8x32_t>(tmp);
+    }
+    else if constexpr(N == 64)
+    {
+        int32x4_t tmp0 = llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
+                                                           src_thread_addr_offset,
+                                                           src_wave_addr_offset,
+                                                           static_cast<index_t>(coherence));
+        int32x4_t tmp1 =
+            llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
+                                              src_thread_addr_offset,
+                                              src_wave_addr_offset + 4 * sizeof(int32_t),
+                                              static_cast<index_t>(coherence));
+        int32x4_t tmp2 =
+            llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
+                                              src_thread_addr_offset,
+                                              src_wave_addr_offset + 8 * sizeof(int32_t),
+                                              static_cast<index_t>(coherence));
+        int32x4_t tmp3 =
+            llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
+                                              src_thread_addr_offset,
+                                              src_wave_addr_offset + 12 * sizeof(int32_t),
+                                              static_cast<index_t>(coherence));
 
-            return bit_cast<int8x16_t>(tmp);
-#endif
-        }
+        vector_type<int32_t, 16> tmp;
+
+        tmp.AsType<int32x4_t>()(Number<0>{}) = tmp0;
+        tmp.AsType<int32x4_t>()(Number<1>{}) = tmp1;
+        tmp.AsType<int32x4_t>()(Number<2>{}) = tmp2;
+        tmp.AsType<int32x4_t>()(Number<3>{}) = tmp3;
+
+        return bit_cast<int8x64_t>(tmp);
     }
 }
 
 template <typename T,
           index_t N,
           AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence>
-__device__ void amd_buffer_store_impl(const typename vector_type<T, N>::type src_thread_data,
-                                      int32x4_t dst_wave_buffer_resource,
-                                      index_t dst_thread_addr_offset,
-                                      index_t dst_wave_addr_offset)
+__device__ typename vector_type<T, N>::type amd_buffer_load_impl(int32x4_t src_wave_buffer_resource,
+                                                                 index_t src_thread_addr_offset,
+                                                                 index_t src_wave_addr_offset)
 {
     static_assert(
-        (is_same<T, double>::value && (N == 1 || N == 2)) ||
-            (is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
-            (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
-            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
-            (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4)) ||
+        (is_same<T, double>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
+            (is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, f8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, bf8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
             (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
         "wrong! not implemented");
 
-    if constexpr(is_same<T, double>::value)
+    using r_t     = typename vector_type<T, N>::type;
+    auto raw_data = amd_buffer_load_impl_raw<sizeof(T) * N, coherence>(
+        src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset);
+    return bit_cast<r_t>(raw_data);
+}
+
+template <index_t N, AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence>
+__device__ void
+amd_buffer_store_impl_raw(const typename vector_type<int8_t, N>::type src_thread_data,
+                          int32x4_t dst_wave_buffer_resource,
+                          index_t dst_thread_addr_offset,
+                          index_t dst_wave_addr_offset)
+{
+    static_assert(N == 1 || N == 2 || N == 4 || N == 8 || N == 16 || N == 32 || N == 64,
+                  "wrong! not implemented");
+
+    if constexpr(N == 1)
     {
-        // use fp32 store to mimic fp64 store
-        if constexpr(N == 1)
-        {
-            llvm_amdgcn_raw_buffer_store_fp32x2(bit_cast<float2_t>(src_thread_data),
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            llvm_amdgcn_raw_buffer_store_fp32x4(bit_cast<float4_t>(src_thread_data),
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-        }
+        llvm_amdgcn_raw_buffer_store_i8(src_thread_data,
+                                        dst_wave_buffer_resource,
+                                        dst_thread_addr_offset,
+                                        dst_wave_addr_offset,
+                                        static_cast<index_t>(coherence));
     }
-    else if constexpr(is_same<T, float>::value)
+    else if constexpr(N == 2)
     {
-        if constexpr(N == 1)
-        {
-            llvm_amdgcn_raw_buffer_store_fp32(src_thread_data,
-                                              dst_wave_buffer_resource,
-                                              dst_thread_addr_offset,
-                                              dst_wave_addr_offset,
-                                              static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            llvm_amdgcn_raw_buffer_store_fp32x2(src_thread_data,
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            llvm_amdgcn_raw_buffer_store_fp32x4(src_thread_data,
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 8)
-        {
-            vector_type<float, 8> tmp{src_thread_data};
-            llvm_amdgcn_raw_buffer_store_fp32x4(tmp.AsType<float4_t>()[Number<0>{}],
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-            llvm_amdgcn_raw_buffer_store_fp32x4(tmp.AsType<float4_t>()[Number<1>{}],
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset + 4 * sizeof(float),
-                                                static_cast<index_t>(coherence));
-        }
+
+        llvm_amdgcn_raw_buffer_store_i16(bit_cast<int16_t>(src_thread_data),
+                                         dst_wave_buffer_resource,
+                                         dst_thread_addr_offset,
+                                         dst_wave_addr_offset,
+                                         static_cast<index_t>(coherence));
     }
-    else if constexpr(is_same<T, half_t>::value)
+    else if constexpr(N == 4)
     {
-        if constexpr(N == 1)
-        {
-            llvm_amdgcn_raw_buffer_store_fp16(src_thread_data,
-                                              dst_wave_buffer_resource,
-                                              dst_thread_addr_offset,
-                                              dst_wave_addr_offset,
-                                              static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            llvm_amdgcn_raw_buffer_store_fp16x2(src_thread_data,
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            llvm_amdgcn_raw_buffer_store_fp16x4(src_thread_data,
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 8)
-        {
-#if 0
-            vector_type<half_t, 8> tmp{src_thread_data};
-
-            llvm_amdgcn_raw_buffer_store_fp16x4(tmp.AsType<half4_t>()[Number<0>{}],
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-
-            llvm_amdgcn_raw_buffer_store_fp16x4(tmp.AsType<half4_t>()[Number<1>{}],
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset + 4 * sizeof(half_t),
-                                                static_cast<index_t>(coherence));
-#else
-            llvm_amdgcn_raw_buffer_store_fp32x4(bit_cast<float4_t>(src_thread_data),
-                                                dst_wave_buffer_resource,
-                                                dst_thread_addr_offset,
-                                                dst_wave_addr_offset,
-                                                static_cast<index_t>(coherence));
-#endif
-        }
+        llvm_amdgcn_raw_buffer_store_i32(bit_cast<int32_t>(src_thread_data),
+                                         dst_wave_buffer_resource,
+                                         dst_thread_addr_offset,
+                                         dst_wave_addr_offset,
+                                         static_cast<index_t>(coherence));
     }
-    else if constexpr(is_same<T, bhalf_t>::value)
+    else if constexpr(N == 8)
     {
-        if constexpr(N == 1)
-        {
-            llvm_amdgcn_raw_buffer_store_i16(src_thread_data,
-                                             dst_wave_buffer_resource,
-                                             dst_thread_addr_offset,
-                                             dst_wave_addr_offset,
-                                             static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            llvm_amdgcn_raw_buffer_store_i16x2(src_thread_data,
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset,
-                                               static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            llvm_amdgcn_raw_buffer_store_i16x4(src_thread_data,
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset,
-                                               static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 8)
-        {
-            vector_type<bhalf_t, 8> tmp{src_thread_data};
-
-            llvm_amdgcn_raw_buffer_store_i16x4(tmp.AsType<bhalf4_t>()[Number<0>{}],
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset,
-                                               static_cast<index_t>(coherence));
-
-            llvm_amdgcn_raw_buffer_store_i16x4(tmp.AsType<bhalf4_t>()[Number<1>{}],
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset + 4 * sizeof(bhalf_t),
-                                               static_cast<index_t>(coherence));
-        }
+        llvm_amdgcn_raw_buffer_store_i32x2(bit_cast<int32x2_t>(src_thread_data),
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset,
+                                           static_cast<index_t>(coherence));
     }
-    else if constexpr(is_same<T, int32_t>::value)
+    else if constexpr(N == 16)
     {
-        if constexpr(N == 1)
-        {
-            llvm_amdgcn_raw_buffer_store_i32(src_thread_data,
-                                             dst_wave_buffer_resource,
-                                             dst_thread_addr_offset,
-                                             dst_wave_addr_offset,
-                                             static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-            llvm_amdgcn_raw_buffer_store_i32x2(src_thread_data,
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset,
-                                               static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 4)
-        {
-            llvm_amdgcn_raw_buffer_store_i32x4(src_thread_data,
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset,
-                                               static_cast<index_t>(coherence));
-        }
+        llvm_amdgcn_raw_buffer_store_i32x4(bit_cast<int32x4_t>(src_thread_data),
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset,
+                                           static_cast<index_t>(coherence));
     }
-    else if constexpr(is_same<T, int8_t>::value)
+    else if constexpr(N == 32)
     {
-        if constexpr(N == 1)
-        {
-            llvm_amdgcn_raw_buffer_store_i8(src_thread_data,
-                                            dst_wave_buffer_resource,
-                                            dst_thread_addr_offset,
-                                            dst_wave_addr_offset,
-                                            static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 2)
-        {
-#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
-            llvm_amdgcn_raw_buffer_store_i8x2(src_thread_data,
-                                              dst_wave_buffer_resource,
-                                              dst_thread_addr_offset,
-                                              dst_wave_addr_offset,
-                                              static_cast<index_t>(coherence));
-#else
-            llvm_amdgcn_raw_buffer_store_i16(bit_cast<int16_t>(src_thread_data),
-                                             dst_wave_buffer_resource,
-                                             dst_thread_addr_offset,
-                                             dst_wave_addr_offset,
-                                             static_cast<index_t>(coherence));
-#endif
-        }
-        else if constexpr(N == 4)
-        {
-#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
-            llvm_amdgcn_raw_buffer_store_i8x4(src_thread_data,
-                                              dst_wave_buffer_resource,
-                                              dst_thread_addr_offset,
-                                              dst_wave_addr_offset,
-                                              static_cast<index_t>(coherence));
-#else
-            llvm_amdgcn_raw_buffer_store_i32(bit_cast<int32_t>(src_thread_data),
-                                             dst_wave_buffer_resource,
-                                             dst_thread_addr_offset,
-                                             dst_wave_addr_offset,
-                                             static_cast<index_t>(coherence));
-#endif
-        }
-        else if constexpr(N == 8)
-        {
-            llvm_amdgcn_raw_buffer_store_i32x2(bit_cast<int32x2_t>(src_thread_data),
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset,
-                                               static_cast<index_t>(coherence));
-        }
-        else if constexpr(N == 16)
-        {
-            llvm_amdgcn_raw_buffer_store_i32x4(bit_cast<int32x4_t>(src_thread_data),
-                                               dst_wave_buffer_resource,
-                                               dst_thread_addr_offset,
-                                               dst_wave_addr_offset,
-                                               static_cast<index_t>(coherence));
-        }
+        vector_type<int32_t, 8> tmp{bit_cast<int32x8_t>(src_thread_data)};
+
+        llvm_amdgcn_raw_buffer_store_i32x4(tmp.template AsType<int32x4_t>()[Number<0>{}],
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset,
+                                           static_cast<index_t>(coherence));
+
+        llvm_amdgcn_raw_buffer_store_i32x4(tmp.template AsType<int32x4_t>()[Number<1>{}],
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset + sizeof(int32_t) * 4,
+                                           static_cast<index_t>(coherence));
     }
+    else if constexpr(N == 64)
+    {
+        vector_type<int32_t, 16> tmp{bit_cast<int32x16_t>(src_thread_data)};
+
+        llvm_amdgcn_raw_buffer_store_i32x4(tmp.template AsType<int32x4_t>()[Number<0>{}],
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset,
+                                           static_cast<index_t>(coherence));
+
+        llvm_amdgcn_raw_buffer_store_i32x4(tmp.template AsType<int32x4_t>()[Number<1>{}],
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset + sizeof(int32_t) * 4,
+                                           static_cast<index_t>(coherence));
+
+        llvm_amdgcn_raw_buffer_store_i32x4(tmp.template AsType<int32x4_t>()[Number<2>{}],
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset + sizeof(int32_t) * 8,
+                                           static_cast<index_t>(coherence));
+
+        llvm_amdgcn_raw_buffer_store_i32x4(tmp.template AsType<int32x4_t>()[Number<3>{}],
+                                           dst_wave_buffer_resource,
+                                           dst_thread_addr_offset,
+                                           dst_wave_addr_offset + sizeof(int32_t) * 12,
+                                           static_cast<index_t>(coherence));
+    }
+}
+
+template <typename T,
+          index_t N,
+          AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence>
+__device__ void amd_buffer_store_impl(const typename vector_type<T, N>::type src_thread_data,
+                                      int32x4_t dst_wave_buffer_resource,
+                                      index_t dst_thread_addr_offset,
+                                      index_t dst_wave_addr_offset)
+{
+    static_assert(
+        (is_same<T, double>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
+            (is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, f8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, bf8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
+        "wrong! not implemented");
+
+    using r_t = typename vector_type<int8_t, sizeof(T) * N>::type;
+
+    amd_buffer_store_impl_raw<sizeof(T) * N, coherence>(bit_cast<r_t>(src_thread_data),
+                                                        dst_wave_buffer_resource,
+                                                        dst_thread_addr_offset,
+                                                        dst_wave_addr_offset);
 }
 
 template <typename T, index_t N>
@@ -1127,54 +798,14 @@ amd_buffer_load_invalid_element_return_zero(const T* p_src_wave,
 
 #if CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK
     uint32_t src_addr_shift = src_thread_element_valid ? 0 : 0x80000000;
-#if defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-    if constexpr(is_same<scalar_t, f8_t>::value || is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 && !defined CK_ENABLE_BF8
-        if constexpr(is_same<scalar_t, f8_t>::value)
-#endif
-#if !defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-            if constexpr(is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-            {
-                auto tmp = amd_buffer_load_impl<int8_t, vector_size, coherence>(
-                    src_wave_buffer_resource, src_addr_shift + src_thread_addr_offset, 0);
-                return bit_cast<vector_t>(tmp);
-            }
-            else
-            {
-#endif
-                return amd_buffer_load_impl<scalar_t, vector_size, coherence>(
-                    src_wave_buffer_resource, src_addr_shift + src_thread_addr_offset, 0);
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-            }
-#endif
+    return amd_buffer_load_impl<scalar_t, vector_size, coherence>(
+        src_wave_buffer_resource, src_addr_shift + src_thread_addr_offset, 0);
+
 #else
-#if defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-    if constexpr(is_same<scalar_t, f8_t>::value || is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 && !defined CK_ENABLE_BF8
-        if constexpr(is_same<scalar_t, f8_t>::value)
-#endif
-#if !defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-            if constexpr(is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-            {
-                auto tmp = amd_buffer_load_impl<int8_t, vector_size, coherence>(
-                    src_wave_buffer_resource, src_thread_addr_offset, 0);
-                return src_thread_element_valid ? bit_cast<vector_t>(tmp) : vector_t(0);
-            }
-            else
-            {
-#endif
-                vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size, coherence>(
-                    src_wave_buffer_resource, src_thread_addr_offset, 0);
-                return src_thread_element_valid ? tmp : vector_t(0);
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-            }
-#endif
+
+    vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size, coherence>(
+        src_wave_buffer_resource, src_thread_addr_offset, 0);
+    return src_thread_element_valid ? tmp : vector_t(0);
 #endif
 }
 
@@ -1232,62 +863,13 @@ __device__ void amd_buffer_store(const typename vector_type_maker<T, N>::type::t
 
 #if CK_EXPERIMENTAL_USE_BUFFER_STORE_OOB_CHECK_OFFSET_TRICK
     uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000;
-
-#if defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-    if constexpr(is_same<scalar_t, f8_t>::value || is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 && !defined CK_ENABLE_BF8
-        if constexpr(is_same<scalar_t, f8_t>::value)
-#endif
-#if !defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-            if constexpr(is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-            {
-                auto tmp = bit_cast<typename vector_type_maker<int8_t, vector_size>::type::type>(
-                    src_thread_data);
-                amd_buffer_store_impl<int8_t, vector_size, coherence>(
-                    tmp, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
-            }
-            else
-            {
-#endif
-                amd_buffer_store_impl<scalar_t, vector_size, coherence>(src_thread_data,
-                                                                        dst_wave_buffer_resource,
-                                                                        dst_addr_shift +
-                                                                            dst_thread_addr_offset,
-                                                                        0);
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-            }
-#endif
+    amd_buffer_store_impl<scalar_t, vector_size, coherence>(
+        src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
 #else
     if(dst_thread_element_valid)
     {
-#if defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-        if constexpr(is_same<scalar_t, f8_t>::value || is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 && !defined CK_ENABLE_BF8
-            if constexpr(is_same<scalar_t, f8_t>::value)
-#endif
-#if !defined CK_ENABLE_FP8 && defined CK_ENABLE_BF8
-                if constexpr(is_same<scalar_t, bf8_t>::value)
-#endif
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-                {
-                    auto tmp =
-                        bit_cast<typename vector_type_maker<int8_t, vector_size>::type::type>(
-                            src_thread_data);
-                    amd_buffer_store_impl<int8_t, vector_size, coherence>(
-                        tmp, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
-                }
-                else
-                {
-#endif
-                    amd_buffer_store_impl<scalar_t, vector_size, coherence>(
-                        src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
-#if defined CK_ENABLE_FP8 || defined CK_ENABLE_BF8
-                }
-#endif
+        amd_buffer_store_impl<scalar_t, vector_size, coherence>(
+            src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
     }
 #endif
 }

include/ck/utility/is_detected.hpp

@@ -31,4 +31,13 @@ struct nonesuch
 template <template <class...> class Op, class... Args>
 using is_detected = typename detail::detector<nonesuch, void, Op, Args...>::value_t;
 
+template <typename T>
+using is_pack2_invocable_t = decltype(std::declval<T&>().is_pack2_invocable);
+
+template <typename T>
+using is_pack4_invocable_t = decltype(std::declval<T&>().is_pack4_invocable);
+
+template <typename T>
+using is_pack8_invocable_t = decltype(std::declval<T&>().is_pack8_invocable);
+
 } // namespace ck

include/ck/utility/type_convert.hpp

@@ -344,7 +344,7 @@ inline __host__ __device__ f8_t f8_convert_sr<f8_t, half_t>(half_t x)
 #if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
     // convert to float and use native converion
     return f8_convert_sr<f8_t>(type_convert<float>(x));
-#else
+#elif 0
     constexpr bool negative_zero_nan = true;
     constexpr bool clip              = true;
     constexpr f8_rounding_mode rm    = f8_rounding_mode::stochastic;
@@ -353,6 +353,8 @@ inline __host__ __device__ f8_t f8_convert_sr<f8_t, half_t>(half_t x)
     return utils::
         cast_to_f8<half_t, f8_t, negative_zero_nan, clip, (rm == f8_rounding_mode::stochastic)>(
             x, rng);
+#else
+    return type_convert<f8_t>(type_convert<float>(x));
 #endif
 }
 #endif
@@ -393,7 +395,7 @@ inline __host__ __device__ bf8_t f8_convert_sr<bf8_t, half_t>(half_t x)
 #if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
     // convert to float and use native converion
     return f8_convert_sr<f8_t>(type_convert<float>(x));
-#else
+#elif 0
     constexpr bool negative_zero_nan = true;
     constexpr bool clip              = true;
     constexpr f8_rounding_mode rm    = f8_rounding_mode::stochastic;
@@ -403,6 +405,8 @@ inline __host__ __device__ bf8_t f8_convert_sr<bf8_t, half_t>(half_t x)
     return utils::
         cast_to_f8<half_t, bf8_t, negative_zero_nan, clip, (rm == f8_rounding_mode::stochastic)>(
             x, rng);
+#else
+    return type_convert<bf8_t>(type_convert<float>(x));
 #endif
 }
 #endif

test/batched_gemm/CMakeLists.txt

@@ -2,22 +2,8 @@ list(APPEND gpu_list gfx908 gfx90a gfx940 gfx941 gfx942)
 set(target 0)
 foreach(gpu IN LISTS GPU_TARGETS)
  if(gpu IN_LIST gpu_list AND target EQUAL 0)
-   add_test_executable(test_batched_gemm_fp16 batched_gemm_fp16.cpp)
-   if(result EQUAL 0)
-      target_link_libraries(test_batched_gemm_fp16 PRIVATE utility device_batched_gemm_instance)
-   endif()
-   add_test_executable(test_batched_gemm_fp32 batched_gemm_fp32.cpp)
-   if(result EQUAL 0)
-      target_link_libraries(test_batched_gemm_fp32 PRIVATE utility device_batched_gemm_instance)
-   endif()
-   add_test_executable(test_batched_gemm_bf16 batched_gemm_bf16.cpp)
-   if(result EQUAL 0)
-      target_link_libraries(test_batched_gemm_bf16 PRIVATE utility device_batched_gemm_instance)
-   endif()
-   add_test_executable(test_batched_gemm_int8 batched_gemm_int8.cpp)
-   if(result EQUAL 0)
-      target_link_libraries(test_batched_gemm_int8 PRIVATE utility device_batched_gemm_instance)
-   endif()
+   add_gtest_executable(test_batched_gemm test_batched_gemm.cpp)
+   target_link_libraries(test_batched_gemm PRIVATE utility device_batched_gemm_instance)
    set(target 1)
  endif()
 endforeach()
\ No newline at end of file

test/batched_gemm/batched_gemm_bf16.cpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
-
-#include <iostream>
-
-#include "profiler/profile_batched_gemm_impl.hpp"
-
-#include "ck/library/tensor_operation_instance/gpu/batched_gemm.hpp"
-
-namespace {
-using ADataType = ck::bhalf_t;
-using BDataType = ck::bhalf_t;
-using CDataType = ck::bhalf_t;
-
-using Row = ck::tensor_layout::gemm::RowMajor;
-using Col = ck::tensor_layout::gemm::ColumnMajor;
-
-using PassThrough = ck::tensor_operation::element_wise::PassThrough;
-} // namespace
-
-int main()
-{
-    int M          = 256;
-    int N          = 256;
-    int K          = 128;
-    int BatchCount = 3;
-
-    bool pass = true;
-
-    using namespace ck::tensor_operation::device;
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, K, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, K, N, M * K, K * N, M * N, BatchCount);
-
-    std::cout << "test BatchedGEMM bf16: " << (pass ? "Pass" : "Fail") << std::endl;
-    return pass ? 0 : 1;
-}

test/batched_gemm/batched_gemm_fp16.cpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
-
-#include <iostream>
-
-#include "profiler/profile_batched_gemm_impl.hpp"
-
-#include "ck/library/tensor_operation_instance/gpu/batched_gemm.hpp"
-
-namespace {
-using ADataType = ck::half_t;
-using BDataType = ck::half_t;
-using CDataType = ck::half_t;
-
-using Row = ck::tensor_layout::gemm::RowMajor;
-using Col = ck::tensor_layout::gemm::ColumnMajor;
-
-using PassThrough = ck::tensor_operation::element_wise::PassThrough;
-} // namespace
-
-int main()
-{
-    int M          = 512;
-    int N          = 256;
-    int K          = 128;
-    int BatchCount = 3;
-
-    bool pass = true;
-
-    using namespace ck::tensor_operation::device;
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, K, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, K, N, M * K, K * N, M * N, BatchCount);
-
-    std::cout << "test BatchedGEMM fp16: " << (pass ? "Pass" : "Fail") << std::endl;
-    return pass ? 0 : 1;
-}

test/batched_gemm/batched_gemm_fp32.cpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
-
-#include <iostream>
-
-#include "profiler/profile_batched_gemm_impl.hpp"
-
-#include "ck/library/tensor_operation_instance/gpu/batched_gemm.hpp"
-
-namespace {
-using ADataType = float;
-using BDataType = float;
-using CDataType = float;
-
-using Row = ck::tensor_layout::gemm::RowMajor;
-using Col = ck::tensor_layout::gemm::ColumnMajor;
-
-using PassThrough = ck::tensor_operation::element_wise::PassThrough;
-} // namespace
-
-int main()
-{
-    int M          = 256;
-    int N          = 256;
-    int K          = 128;
-    int BatchCount = 3;
-
-    bool pass = true;
-
-    using namespace ck::tensor_operation::device;
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, K, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, K, N, M * K, K * N, M * N, BatchCount);
-
-    std::cout << "test BatchedGEMM fp32: " << (pass ? "Pass" : "Fail") << std::endl;
-    return pass ? 0 : 1;
-}

test/batched_gemm/batched_gemm_int8.cpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
-
-#include <iostream>
-
-#include "profiler/profile_batched_gemm_impl.hpp"
-
-#include "ck/library/tensor_operation_instance/gpu/batched_gemm.hpp"
-
-namespace {
-using ADataType = int8_t;
-using BDataType = int8_t;
-using CDataType = int8_t;
-
-using Row = ck::tensor_layout::gemm::RowMajor;
-using Col = ck::tensor_layout::gemm::ColumnMajor;
-
-using PassThrough = ck::tensor_operation::element_wise::PassThrough;
-} // namespace
-
-int main()
-{
-    int M          = 256;
-    int N          = 256;
-    int K          = 128;
-    int BatchCount = 3;
-
-    bool pass = true;
-
-    using namespace ck::tensor_operation::device;
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Row,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Row,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, K, K, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Row,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Row,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, N, N, M * K, K * N, M * N, BatchCount);
-
-    pass = pass && ck::profiler::profile_batched_gemm_impl<ADataType,
-                                                           BDataType,
-                                                           CDataType,
-                                                           Col,
-                                                           Col,
-                                                           Row,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           PassThrough,
-                                                           DeviceBatchedGemm<Col,
-                                                                             Col,
-                                                                             Row,
-                                                                             ADataType,
-                                                                             BDataType,
-                                                                             CDataType,
-                                                                             PassThrough,
-                                                                             PassThrough,
-                                                                             PassThrough>>(
-                       true, 1, false, 1, M, N, K, M, K, N, M * K, K * N, M * N, BatchCount);
-
-    std::cout << "test BatchedGEMM int8: " << (pass ? "Pass" : "Fail") << std::endl;
-    return pass ? 0 : 1;
-}

test/batched_gemm/test_batched_gemm.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <cstdlib>
+#include <iostream>
+#include <initializer_list>
+#include <tuple>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include "profiler/profile_batched_gemm_impl.hpp"
+
+#include "ck/library/tensor_operation_instance/gpu/batched_gemm.hpp"
+
+struct GemmParams
+{
+    ck::index_t M;
+    ck::index_t N;
+    ck::index_t K;
+    ck::index_t BatchCount;
+};
+
+class TestBatchedGemm : public ::testing::Test
+{
+    protected:
+    using Row = ck::tensor_layout::gemm::RowMajor;
+    using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+    using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+    std::vector<GemmParams> params;
+
+    template <typename DataType>
+    void Run()
+    {
+        using namespace ck::tensor_operation::device;
+
+        bool pass = true;
+        for(auto& param : params)
+        {
+            const auto M          = param.M;
+            const auto N          = param.N;
+            const auto K          = param.K;
+            const auto BatchCount = param.BatchCount;
+
+            pass =
+                pass && ck::profiler::profile_batched_gemm_impl<DataType,
+                                                                DataType,
+                                                                DataType,
+                                                                Row,
+                                                                Row,
+                                                                Row,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                DeviceBatchedGemm<Row,
+                                                                                  Row,
+                                                                                  Row,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  PassThrough,
+                                                                                  PassThrough,
+                                                                                  PassThrough>>(
+                            true, 1, false, 1, M, N, K, K, N, N, M * K, K * N, M * N, BatchCount);
+
+            pass =
+                pass && ck::profiler::profile_batched_gemm_impl<DataType,
+                                                                DataType,
+                                                                DataType,
+                                                                Row,
+                                                                Col,
+                                                                Row,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                DeviceBatchedGemm<Row,
+                                                                                  Col,
+                                                                                  Row,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  PassThrough,
+                                                                                  PassThrough,
+                                                                                  PassThrough>>(
+                            true, 1, false, 1, M, N, K, K, K, N, M * K, K * N, M * N, BatchCount);
+
+            pass =
+                pass && ck::profiler::profile_batched_gemm_impl<DataType,
+                                                                DataType,
+                                                                DataType,
+                                                                Col,
+                                                                Row,
+                                                                Row,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                DeviceBatchedGemm<Col,
+                                                                                  Row,
+                                                                                  Row,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  PassThrough,
+                                                                                  PassThrough,
+                                                                                  PassThrough>>(
+                            true, 1, false, 1, M, N, K, M, N, N, M * K, K * N, M * N, BatchCount);
+
+            pass =
+                pass && ck::profiler::profile_batched_gemm_impl<DataType,
+                                                                DataType,
+                                                                DataType,
+                                                                Col,
+                                                                Col,
+                                                                Row,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                DeviceBatchedGemm<Col,
+                                                                                  Col,
+                                                                                  Row,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  DataType,
+                                                                                  PassThrough,
+                                                                                  PassThrough,
+                                                                                  PassThrough>>(
+                            true, 1, false, 1, M, N, K, M, K, N, M * K, K * N, M * N, BatchCount);
+        }
+        EXPECT_TRUE(pass);
+    }
+};
+
+#ifdef CK_ENABLE_INT8
+TEST_F(TestBatchedGemm, i8)
+{
+    this->params.push_back({64, 64, 64, 2});
+    this->params.push_back({64, 64, 64, 1});
+    this->params.push_back({60, 60, 60, 2});
+    this->params.push_back({68, 68, 68, 2});
+    this->params.push_back({40, 40, 40, 2});
+    this->params.push_back({256, 256, 128, 3});
+    this->template Run<int8_t>();
+}
+#endif
+
+#ifdef CK_ENABLE_BF16
+TEST_F(TestBatchedGemm, bf16)
+{
+    this->params.push_back({64, 64, 64, 2});
+    this->params.push_back({64, 64, 64, 1});
+    this->params.push_back({60, 60, 60, 2});
+    this->params.push_back({68, 68, 68, 2});
+    this->params.push_back({40, 40, 40, 2});
+    this->params.push_back({256, 256, 128, 3});
+    this->template Run<ck::bhalf_t>();
+}
+#endif
+
+#ifdef CK_ENABLE_FP16
+TEST_F(TestBatchedGemm, fp16)
+{
+    this->params.push_back({64, 64, 64, 2});
+    this->params.push_back({64, 64, 64, 1});
+    this->params.push_back({60, 60, 60, 2});
+    this->params.push_back({68, 68, 68, 2});
+    this->params.push_back({40, 40, 40, 2});
+    this->params.push_back({256, 256, 128, 3});
+    this->template Run<ck::half_t>();
+}
+#endif
+
+#ifdef CK_ENABLE_FP32
+TEST_F(TestBatchedGemm, fp32)
+{
+    this->params.push_back({64, 64, 64, 2});
+    this->params.push_back({64, 64, 64, 1});
+    this->params.push_back({60, 60, 60, 2});
+    this->params.push_back({68, 68, 68, 2});
+    this->params.push_back({40, 40, 40, 2});
+    this->params.push_back({256, 256, 128, 3});
+    this->template Run<float>();
+}
+#endif