Merge remote-tracking branch 'origin/develop' into aosewski/ggemm_multi_d2

library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_multiply_bias_fastgelu_bf16_i8_bf16_mk_kn_mn_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "device_grouped_gemm_xdl_tile_loop_multiply_bf16_i8_bf16_mk_kn_mn.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+void add_device_grouped_gemm_xdl_tile_loop_multiply_bias_fastgelu_bf16_i8_bf16_mk_kn_mn_instances(
+    std::vector<std::unique_ptr<DeviceGroupedGemmTileLoop<Row,
+                                                          Row,
+                                                          ck::Tuple<Row, Row>,
+                                                          Row,
+                                                          BF16,
+                                                          I8,
+                                                          ck::Tuple<BF16, BF16>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyAddFastGelu>>>& instances)
+{
+    add_device_operation_instances(
+        instances,
+        device_grouped_gemm_xdl_tile_loop_bf16_i8_bf16_mk_kn_mn_comp_instances<
+            ck::Tuple<Row, Row>,
+            ck::Tuple<BF16, BF16>,
+            MultiplyAddFastGelu>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_gemm_xdl_tile_loop_bf16_i8_bf16_mk_kn_mn_mem_instances<
+            ck::Tuple<Row, Row>,
+            ck::Tuple<BF16, BF16>,
+            MultiplyAddFastGelu>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_multiply_fastgelu_bf16_i8_bf16_mk_kn_mn_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "device_grouped_gemm_xdl_tile_loop_multiply_bf16_i8_bf16_mk_kn_mn.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+void add_device_grouped_gemm_xdl_tile_loop_multiply_fastgelu_bf16_i8_bf16_mk_kn_mn_instances(
+    std::vector<std::unique_ptr<DeviceGroupedGemmTileLoop<Row,
+                                                          Row,
+                                                          ck::Tuple<Row>,
+                                                          Row,
+                                                          BF16,
+                                                          I8,
+                                                          ck::Tuple<BF16>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyFastGelu>>>& instances)
+{
+    add_device_operation_instances(
+        instances,
+        device_grouped_gemm_xdl_tile_loop_bf16_i8_bf16_mk_kn_mn_comp_instances<ck::Tuple<Row>,
+                                                                               ck::Tuple<BF16>,
+                                                                               MultiplyFastGelu>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_gemm_xdl_tile_loop_bf16_i8_bf16_mk_kn_mn_mem_instances<ck::Tuple<Row>,
+                                                                              ck::Tuple<BF16>,
+                                                                              MultiplyFastGelu>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

profiler/include/profiler/profile_grouped_gemm_multiply_tile_loop_impl.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iomanip>
+
+#include "ck/ck.hpp"
+#include "ck/host_utility/hip_check_error.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/gpu/grouped_gemm_tile_loop_multiply.hpp"
+
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/literals.hpp"
+#include "ck/library/utility/fill.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+
+namespace ck {
+namespace profiler {
+
+template <typename ADataType,
+          typename BDataType,
+          typename DDataType,
+          typename EDataType,
+          typename AccDataType,
+          typename ALayout,
+          typename BLayout,
+          typename DLayout,
+          typename ELayout>
+bool profile_grouped_gemm_multiply_tile_loop_impl(int do_verification,
+                                                  int init_method,
+                                                  bool do_log,
+                                                  bool time_kernel,
+                                                  const std::vector<int>& Ms,
+                                                  const std::vector<int>& Ns,
+                                                  const std::vector<int>& Ks,
+                                                  const std::vector<int>& StrideAs,
+                                                  const std::vector<int>& StrideBs,
+                                                  const std::vector<int>& StrideDs,
+                                                  const std::vector<int>& StrideEs,
+                                                  int n_warmup = 10,
+                                                  int n_iter   = 50)
+{
+    using CDataType = EDataType;
+    bool pass       = true;
+
+    auto f_host_tensor_descriptor =
+        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+            using namespace ck::literals;
+
+            if(is_same<decltype(layout), tensor_layout::gemm::RowMajor>::value)
+            {
+                return HostTensorDescriptor({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    std::size_t group_count = Ms.size();
+
+    if(!(group_count == Ns.size() && group_count == Ks.size() && group_count == StrideAs.size() &&
+         group_count == StrideBs.size() && group_count == StrideEs.size()))
+    {
+        throw std::runtime_error("wrong! inconsistent M/N/Ks, StrideA/B/Cs size\n");
+    }
+
+    std::vector<Tensor<ADataType>> a_m_k;
+    std::vector<Tensor<BDataType>> b_k_n;
+    std::vector<Tensor<DDataType>> d_m_n;
+    std::vector<Tensor<CDataType>> e_m_n_host_results;
+    std::vector<Tensor<CDataType>> e_m_n_device_results;
+
+    for(std::size_t i = 0; i < group_count; i++)
+    {
+        a_m_k.push_back(
+            Tensor<ADataType>(f_host_tensor_descriptor(Ms[i], Ks[i], StrideAs[i], ALayout{})));
+        b_k_n.push_back(
+            Tensor<BDataType>(f_host_tensor_descriptor(Ks[i], Ns[i], StrideBs[i], BLayout{})));
+        d_m_n.push_back(
+            Tensor<DDataType>(f_host_tensor_descriptor(Ms[i], Ns[i], StrideDs[i], DLayout{})));
+        e_m_n_device_results.push_back(
+            Tensor<CDataType>(f_host_tensor_descriptor(Ms[i], Ns[i], StrideEs[i], ELayout{})));
+        e_m_n_host_results.push_back(
+            Tensor<CDataType>(f_host_tensor_descriptor(Ms[i], Ns[i], StrideEs[i], ELayout{})));
+        if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))
+        {
+            std::cout << "group: " << i << " a_m_k[" << i << "]:" << a_m_k[i].mDesc << ", b_k_n["
+                      << i << "]:" << b_k_n[i].mDesc << ", e_m_n_device_results[" << i
+                      << "]:" << e_m_n_device_results[i].mDesc << std::endl;
+        }
+        switch(init_method)
+        {
+        case 0: break;
+        case 1:
+            ck::utils::FillUniformDistributionIntegerValue<ADataType>{-5, 5}(a_m_k[i]);
+            ck::utils::FillUniformDistributionIntegerValue<BDataType>{-5, 5}(b_k_n[i]);
+            ck::utils::FillUniformDistributionIntegerValue<DDataType>{-5, 5}(d_m_n[i]);
+            break;
+        case 2:
+            ck::utils::FillUniformDistribution<ADataType>{.0, 1.}(a_m_k[i]);
+            ck::utils::FillUniformDistribution<BDataType>{-0.5, 0.5}(b_k_n[i]);
+            ck::utils::FillUniformDistribution<DDataType>{-0.5, 0.5}(d_m_n[i]);
+            break;
+        default:
+            ck::utils::FillConstant<ADataType>{1}(a_m_k[i]);
+            ck::utils::FillConstant<BDataType>{1}(b_k_n[i]);
+            ck::utils::FillConstant<DDataType>{1}(d_m_n[i]);
+        }
+    }
+
+    using AElementOp   = ck::tensor_operation::element_wise::PassThrough;
+    using BElementOp   = ck::tensor_operation::element_wise::PassThrough;
+    using CElementOp   = ck::tensor_operation::element_wise::PassThrough;
+    using CDEElementOp = ck::tensor_operation::element_wise::Multiply;
+
+    const auto a_element_op   = AElementOp{};
+    const auto b_element_op   = BElementOp{};
+    const auto c_element_op   = CElementOp{};
+    const auto cde_element_op = CDEElementOp{};
+
+    using DeviceMemPtr = std::unique_ptr<DeviceMem>;
+    std::vector<DeviceMemPtr> a_device_buf, b_device_buf, d_device_buf, e_device_buf;
+
+    a_device_buf.reserve(group_count);
+    b_device_buf.reserve(group_count);
+    d_device_buf.reserve(group_count);
+    e_device_buf.reserve(group_count);
+
+    std::vector<const void*> p_a, p_b, p_d;
+    constexpr ck::index_t NumDTensor = 1;
+    auto p_ds                        = std::vector<std::array<const void*, NumDTensor>>{};
+    std::vector<void*> p_e;
+
+    p_a.reserve(group_count);
+    p_b.reserve(group_count);
+    p_ds.reserve(group_count);
+    p_e.reserve(group_count);
+
+    using KernelArguments =
+        ck::tensor_operation::device::GroupedGemmTileLoopKernelArguments<NumDTensor>;
+
+    std::vector<ck::tensor_operation::device::GemmDesc> gemm_descs;
+    std::vector<KernelArguments> gemm_kargs;
+
+    gemm_descs.reserve(group_count);
+    gemm_kargs.reserve(group_count);
+
+    for(std::size_t i = 0; i < group_count; i++)
+    {
+        a_device_buf.emplace_back(
+            std::make_unique<DeviceMem>(sizeof(ADataType) * a_m_k[i].mDesc.GetElementSpaceSize()));
+        b_device_buf.emplace_back(
+            std::make_unique<DeviceMem>(sizeof(BDataType) * b_k_n[i].mDesc.GetElementSpaceSize()));
+        d_device_buf.emplace_back(
+            std::make_unique<DeviceMem>(sizeof(DDataType) * d_m_n[i].mDesc.GetElementSpaceSize()));
+        e_device_buf.emplace_back(std::make_unique<DeviceMem>(
+            sizeof(CDataType) * e_m_n_device_results[i].mDesc.GetElementSpaceSize()));
+
+        a_device_buf[i]->ToDevice(a_m_k[i].mData.data());
+        b_device_buf[i]->ToDevice(b_k_n[i].mData.data());
+        d_device_buf[i]->ToDevice(d_m_n[i].mData.data());
+        e_device_buf[i]->SetZero();
+
+        p_a.push_back(a_device_buf[i]->GetDeviceBuffer());
+        p_b.push_back(b_device_buf[i]->GetDeviceBuffer());
+        p_ds.push_back({d_device_buf[i]->GetDeviceBuffer()});
+        p_e.push_back(e_device_buf[i]->GetDeviceBuffer());
+
+        gemm_descs.push_back(
+            {0, Ns[i], Ks[i], StrideAs[i], StrideBs[i], StrideEs[i], {StrideDs[i]}});
+        gemm_kargs.push_back({a_device_buf[i]->GetDeviceBuffer(),
+                              b_device_buf[i]->GetDeviceBuffer(),
+                              {d_device_buf[i]->GetDeviceBuffer()},
+                              e_device_buf[i]->GetDeviceBuffer(),
+                              Ms[i],
+                              Ns[i],
+                              Ks[i],
+                              StrideAs[i],
+                              StrideBs[i],
+                              {StrideDs[i]},
+                              StrideEs[i]});
+    }
+
+    using DeviceOp = ck::tensor_operation::device::DeviceGroupedGemmTileLoop<ALayout,
+                                                                             BLayout,
+                                                                             ck::Tuple<DLayout>,
+                                                                             ELayout,
+                                                                             ADataType,
+                                                                             BDataType,
+                                                                             ck::Tuple<DDataType>,
+                                                                             EDataType,
+                                                                             AElementOp,
+                                                                             BElementOp,
+                                                                             CDEElementOp>;
+
+    const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+        DeviceOp>::GetInstances();
+
+    if(op_ptrs.size() <= 0)
+    {
+        throw std::runtime_error("wrong! no device GEMM instance found");
+    }
+
+    std::string best_gemm_name;
+    float best_ave_time   = 0;
+    float best_tflops     = 0;
+    float best_gb_per_sec = 0;
+
+    if(do_verification)
+    {
+        for(std::size_t i = 0; i < gemm_descs.size(); i++)
+        {
+            Tensor<CDataType> c_m_n({Ms[i], Ns[i]});
+
+            using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                                                    BDataType,
+                                                                                    CDataType,
+                                                                                    AccDataType,
+                                                                                    AElementOp,
+                                                                                    BElementOp,
+                                                                                    CElementOp>;
+
+            auto ref_gemm     = ReferenceGemmInstance{};
+            auto ref_invoker  = ref_gemm.MakeInvoker();
+            auto ref_argument = ref_gemm.MakeArgument(
+                a_m_k[i], b_k_n[i], c_m_n, a_element_op, b_element_op, c_element_op);
+            ref_invoker.Run(ref_argument);
+
+            for(int m = 0; m < Ms[i]; ++m)
+            {
+                for(int n = 0; n < Ns[i]; ++n)
+                {
+                    cde_element_op(e_m_n_host_results[i](m, n), c_m_n(m, n), d_m_n[i](m, n));
+                }
+            }
+        }
+    }
+
+    // profile device GEMM instances
+    for(auto& gemm_ptr : op_ptrs)
+    {
+        auto argument_ptr =
+            gemm_ptr->MakeArgumentPointer(p_a,
+                                          p_b,
+                                          p_ds,
+                                          p_e,
+                                          gemm_descs,
+                                          ck::tensor_operation::element_wise::PassThrough{},
+                                          ck::tensor_operation::element_wise::PassThrough{},
+                                          cde_element_op);
+        auto invoker_ptr      = gemm_ptr->MakeInvokerPointer();
+        std::string gemm_name = gemm_ptr->GetTypeString();
+
+        DeviceMem gemm_arg_dev_mem(gemm_ptr->GetDeviceKernelArgSize(argument_ptr.get()));
+        hip_check_error(hipMemcpy(gemm_arg_dev_mem.GetDeviceBuffer(),
+                                  gemm_kargs.data(),
+                                  gemm_ptr->GetDeviceKernelArgSize(argument_ptr.get()),
+                                  hipMemcpyHostToDevice));
+        gemm_ptr->SetDeviceKernelArgs(argument_ptr.get(), gemm_arg_dev_mem.GetDeviceBuffer());
+
+        if(gemm_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false, 0, n_warmup, n_iter});
+            if(do_verification)
+            {
+                bool instance_pass = true;
+                for(std::size_t i = 0; i < gemm_descs.size(); i++)
+                {
+                    e_device_buf[i]->FromDevice(e_m_n_device_results[i].mData.data());
+                    instance_pass = instance_pass && ck::utils::check_err(e_m_n_device_results[i],
+                                                                          e_m_n_host_results[i]);
+
+                    if(do_log)
+                    {
+                        LogRangeAsType<float>(std::cout << "a : ", a_m_k[i].mData, ",")
+                            << std::endl;
+                        LogRangeAsType<float>(std::cout << "b: ", b_k_n[i].mData, ",") << std::endl;
+                        LogRangeAsType<float>(
+                            std::cout << "e_device: ", e_m_n_device_results[i].mData, ",")
+                            << std::endl;
+                        LogRangeAsType<float>(
+                            std::cout << "e_host  : ", e_m_n_host_results[i].mData, ",")
+                            << std::endl;
+                    }
+                }
+
+                std::cout << "Instance: " << gemm_name << " verification "
+                          << (instance_pass ? "SUCCEED" : "FAILED") << std::endl;
+
+                pass = pass && instance_pass;
+            }
+
+            if(time_kernel)
+            {
+                float ave_time = invoker_ptr->Run(
+                    argument_ptr.get(), StreamConfig{nullptr, time_kernel, 0, n_warmup, n_iter});
+
+                std::size_t flop = 0, num_btype = 0;
+                for(std::size_t i = 0; i < gemm_descs.size(); i++)
+                {
+                    flop += std::size_t(2) * Ms[i] * Ns[i] * Ks[i];
+
+                    num_btype += sizeof(ADataType) * Ms[i] * Ks[i] +
+                                 sizeof(BDataType) * Ks[i] * Ns[i] +
+                                 sizeof(EDataType) * Ms[i] * Ns[i] + // D matrix
+                                 sizeof(EDataType) * Ms[i] * Ns[i];
+                }
+
+                float tflops     = static_cast<float>(flop) / 1.E9 / ave_time;
+                float gb_per_sec = num_btype / 1.E6 / ave_time;
+                std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops
+                          << " TFlops, " << gb_per_sec << " GB/s, " << gemm_name << std::endl;
+
+                if(tflops > best_tflops)
+                {
+                    best_gemm_name  = gemm_name;
+                    best_tflops     = tflops;
+                    best_ave_time   = ave_time;
+                    best_gb_per_sec = gb_per_sec;
+                }
+            }
+        }
+        else
+        {
+            std::cout << "Instance: " << gemm_name << ", does not support this GEMM problem"
+                      << std::endl;
+        }
+    }
+
+    if(time_kernel)
+    {
+        std::cout << "Best Perf: " << best_ave_time << " ms, " << best_tflops << " TFlops, "
+                  << best_gb_per_sec << " GB/s, " << best_gemm_name << std::endl;
+    }
+
+    return pass;
+}
+
+} // namespace profiler
+} // namespace ck

profiler/src/CMakeLists.txt

@@ -44,6 +44,7 @@ if(GPU_TARGETS MATCHES "gfx9")
     list(APPEND PROFILER_SOURCES profile_grouped_gemm_fastgelu.cpp)
     list(APPEND PROFILER_SOURCES profile_grouped_gemm_multiple_d_splitk.cpp)
     list(APPEND PROFILER_SOURCES profile_grouped_gemm_tile_loop.cpp)
+    list(APPEND PROFILER_SOURCES profile_grouped_gemm_multiply_tile_loop.cpp)
   endif()
   list(APPEND PROFILER_SOURCES profile_gemm_multiply_add.cpp)
   list(APPEND PROFILER_SOURCES profile_batched_gemm.cpp)
@@ -59,7 +60,7 @@ if(GPU_TARGETS MATCHES "gfx9")
 
 endif()
 
-if(GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx9")
+if(GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12" OR GPU_TARGETS MATCHES "gfx9")
   if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
     list(APPEND PROFILER_SOURCES profile_gemm_bilinear.cpp)
   endif()
@@ -135,7 +136,7 @@ if(GPU_TARGETS MATCHES "gfx9")
   target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_grouped_conv2d_bwd_weight_instance)
 endif()
 
-if(GPU_TARGETS MATCHES "gfx9" OR GPU_TARGETS MATCHES "gfx11")
+if(GPU_TARGETS MATCHES "gfx9" OR GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12")
   if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
     target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_bilinear_instance)
   endif()

profiler/src/profile_grouped_gemm.cpp

@@ -98,8 +98,8 @@ int profile_grouped_gemm(int argc, char* argv[])
     int n_iter   = 10;
     if(argc == 17)
     {
-        n_warmup = std::stoi(argv[16]);
-        n_iter   = std::stoi(argv[17]);
+        n_warmup = std::stoi(argv[15]);
+        n_iter   = std::stoi(argv[16]);
     }
 
 #ifdef CK_ENABLE_FP16

profiler/src/profile_grouped_gemm_multiply_tile_loop.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <cstdlib>
+#include <iostream>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "profiler/profile_grouped_gemm_multiply_tile_loop_impl.hpp"
+#include "profiler_operation_registry.hpp"
+
+enum struct GemmMatrixLayout
+{
+    MK_KN_MN, // 0
+};
+
+enum struct GemmDataType
+{
+    BF16_INT8_BF16_BF16, // 0
+};
+
+#define OP_NAME "grouped_gemm_multiply_tile_loop"
+#define OP_DESC "Grouped GEMM Multiply Multiple D Tile Loop"
+
+namespace {
+
+std::vector<int> argToIntArray(char* input)
+{
+    std::vector<int> out;
+    std::istringstream in(input);
+    std::string item;
+
+    while(std::getline(in, item, ','))
+    {
+        out.push_back(std::stoi(item));
+    }
+    return out;
+}
+
+int profile_grouped_gemm_tile_loop(int argc, char* argv[])
+{
+    if(argc < 14)
+    {
+        std::cout
+            << "arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n"
+            << "arg2: data type (0: bf16@int8)\n"
+            << "arg3: matrix layout (0: A[m, k] * B[k, n] = C[m, n]);\n"
+            << "arg4: verification (0: no; 1: yes)\n"
+            << "arg5: initialization (0: no init; 1: integer value; 2: decimal value)\n"
+            << "arg6: print tensor value (0: no; 1: yes)\n"
+            << "arg7: time kernel (0=n0, 1=yes)\n"
+            << "arg8 to 13: Ms, Ns, Ks, StrideAs, StrideBs, StrideCs (e.g., 256,256 128,128 64,64 "
+               "64,64 64,64 128,128)\n"
+            << "optional:\n"
+            << "arg14: number of warm-up cycles (default 1)\n"
+            << "arg15: number of iterations (default 10)\n"
+            << std::endl;
+
+        exit(1);
+    }
+
+    const auto data_type       = static_cast<GemmDataType>(std::stoi(argv[2]));
+    const auto layout          = static_cast<GemmMatrixLayout>(std::stoi(argv[3]));
+    const bool do_verification = std::stoi(argv[4]);
+    const int init_method      = std::stoi(argv[5]);
+    const bool do_log          = std::stoi(argv[6]);
+    const bool time_kernel     = std::stoi(argv[7]);
+
+    const auto Ms = argToIntArray(argv[8]);
+    const auto Ns = argToIntArray(argv[9]);
+    const auto Ks = argToIntArray(argv[10]);
+
+    auto StrideAs = argToIntArray(argv[11]);
+    auto StrideBs = argToIntArray(argv[12]);
+    auto StrideCs = argToIntArray(argv[13]);
+
+    const int DefaultStrideA = Ks[0];
+    const int DefaultStrideB = Ns[0];
+    const int DefaultStrideC = Ns[0];
+
+    for(size_t i = 0; i < Ms.size(); ++i)
+    {
+        StrideAs[i] = StrideAs[i] == -1 ? DefaultStrideA : StrideAs[i];
+        StrideBs[i] = StrideBs[i] == -1 ? DefaultStrideB : StrideBs[i];
+        StrideCs[i] = StrideCs[i] == -1 ? DefaultStrideC : StrideCs[i];
+    }
+
+    std::vector<int> StrideDs(StrideCs);
+
+    int n_warmup = 10;
+    int n_iter   = 50;
+    if(argc == 16)
+    {
+        n_warmup = std::stoi(argv[14]);
+        n_iter   = std::stoi(argv[15]);
+    }
+
+    if(data_type == GemmDataType::BF16_INT8_BF16_BF16 && layout == GemmMatrixLayout::MK_KN_MN)
+    {
+        ck::profiler::profile_grouped_gemm_multiply_tile_loop_impl<
+            ck::bhalf_t,
+            int8_t,
+            ck::bhalf_t,
+            ck::bhalf_t,
+            float,
+            ck::tensor_layout::gemm::RowMajor,
+            ck::tensor_layout::gemm::RowMajor,
+            ck::tensor_layout::gemm::RowMajor,
+            ck::tensor_layout::gemm::RowMajor>(do_verification,
+                                               init_method,
+                                               do_log,
+                                               time_kernel,
+                                               Ms,
+                                               Ns,
+                                               Ks,
+                                               StrideAs,
+                                               StrideBs,
+                                               StrideDs,
+                                               StrideCs,
+                                               n_warmup,
+                                               n_iter);
+    }
+    else
+    {
+        throw std::runtime_error("wrong! this GEMM data_type & layout is not implemented");
+    }
+    return 0;
+}
+
+} // anonymous namespace
+
+REGISTER_PROFILER_OPERATION(OP_NAME, OP_DESC, profile_grouped_gemm_tile_loop);

test/CMakeLists.txt

@@ -60,7 +60,7 @@ function(add_test_executable TEST_NAME)
         endif()
     endforeach()
     foreach(source IN LISTS ARGN)
-        if(NOT TEST_TARGETS MATCHES "gfx11" AND source MATCHES "wmma")
+	if(NOT GPU_TARGETS MATCHES "gfx11" AND NOT GPU_TARGETS MATCHES "gfx12" AND source MATCHES "wmma")
             message("removing wmma test ${source} ")
             list(REMOVE_ITEM ARGN "${source}")
         endif()
@@ -139,7 +139,7 @@ function(add_gtest_executable TEST_NAME)
         endif()
     endforeach()
     foreach(source IN LISTS ARGN)
-        if(NOT TEST_TARGETS MATCHES "gfx11" AND source MATCHES "wmma")
+	if(NOT GPU_TARGETS MATCHES "gfx11" AND NOT GPU_TARGETS MATCHES "gfx12" AND source MATCHES "wmma")
             message("removing wmma test ${source} ")
             list(REMOVE_ITEM ARGN "${source}")
         endif()
@@ -210,4 +210,7 @@ add_subdirectory(wrapper)
 if(GPU_TARGETS MATCHES "gfx11")
     add_subdirectory(wmma_op)
 endif()
+if(GPU_TARGETS MATCHES "gfx942")
+    add_subdirectory(smfmac_op)
+endif()
 add_subdirectory(position_embedding)

test/contraction/test_contraction_xdl.cpp

@@ -212,4 +212,10 @@ TYPED_TEST(TestContractionScaleMixedPrecision, scale)
     this->template Run<6>({{1, 1, 1, 3, 2, 3}, {1, 1, 1, 3, 2, 3}, {1, 1, 1, 2, 2, 4}});
     this->template Run<2>({{16, 8}, {16, 8}, {16, 8}});
     this->template Run<2>({{8, 16}, {16, 8}, {8, 16}});
+
+    // special cases
+    this->template Run<2>({{1, 1}, {16, 8}, {8, 16}});
+    this->template Run<2>({{8, 16}, {16, 8}, {1, 1}});
+    this->template Run<2>({{8, 16}, {1, 1}, {8, 16}});
+    this->template Run<2>({{1, 1}, {1, 1}, {1, 1}});
 }

test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight.cpp

@@ -44,7 +44,7 @@ class TestGroupedConvndBwdWeight : public ::testing::Test
             }
         }
 
-        if(ck::is_gfx11_supported())
+        if(ck::is_gfx11_supported() || ck::is_gfx12_supported())
         {
             // on gfx11 only support for 3d is implemented
             if constexpr(NDimSpatial{} != 3)

test/grouped_convnd_fwd/CMakeLists.txt

-add_gtest_executable(test_grouped_convnd_fwd test_grouped_convnd_fwd_xdl_wmma.cpp)
-if(result EQUAL 0)
-    target_link_libraries(test_grouped_convnd_fwd PRIVATE utility device_grouped_conv1d_fwd_instance device_grouped_conv2d_fwd_instance device_grouped_conv3d_fwd_instance)
+if(GPU_TARGETS MATCHES "gfx9" OR GPU_TARGETS MATCHES "gfx11")
+    add_gtest_executable(test_grouped_convnd_fwd test_grouped_convnd_fwd.cpp)
+    if(GPU_TARGETS MATCHES "gfx11")
+        target_link_libraries(test_grouped_convnd_fwd PRIVATE utility device_grouped_conv2d_fwd_instance device_grouped_conv3d_fwd_instance)
+    else()
+        target_link_libraries(test_grouped_convnd_fwd PRIVATE utility device_grouped_conv1d_fwd_instance device_grouped_conv2d_fwd_instance device_grouped_conv3d_fwd_instance)
+    endif()
 endif()
 
 add_gtest_executable(test_grouped_convnd_fwd_multi_ab_interface test_grouped_convnd_fwd_multi_ab_interface.cpp)

test/grouped_convnd_fwd/test_grouped_convnd_fwd_xdl_wmma.cpp → test/grouped_convnd_fwd/test_grouped_convnd_fwd.cpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <cstdlib>
 #include <iostream>
@@ -69,6 +69,8 @@ using KernelTypes3d = ::testing::Types<std::tuple<float, GNDHWC, GKZYXC, GNDHWK>
                                        std::tuple<ck::bhalf_t, NDHWGC, GKZYXC, NDHWGK>,
                                        std::tuple<int8_t, NDHWGC, GKZYXC, NDHWGK>>;
 
+using KernelTypes2dLargeCases = ::testing::Types<std::tuple<float, NHWGC, GKYXC, NHWGK>>;
+
 template <typename Tuple>
 class TestGroupedConvndFwd1d : public TestGroupedConvndFwd<Tuple>
 {
@@ -84,9 +86,15 @@ class TestGroupedConvndFwd3d : public TestGroupedConvndFwd<Tuple>
 {
 };
 
+template <typename Tuple>
+class TestGroupedConvndFwd2dLargeCases : public TestGroupedConvndFwd<Tuple>
+{
+};
+
 TYPED_TEST_SUITE(TestGroupedConvndFwd1d, KernelTypes1d);
 TYPED_TEST_SUITE(TestGroupedConvndFwd2d, KernelTypes2d);
 TYPED_TEST_SUITE(TestGroupedConvndFwd3d, KernelTypes3d);
+TYPED_TEST_SUITE(TestGroupedConvndFwd2dLargeCases, KernelTypes2dLargeCases);
 
 TYPED_TEST(TestGroupedConvndFwd1d, Test1D)
 {
@@ -131,3 +139,11 @@ TYPED_TEST(TestGroupedConvndFwd3d, Test3D)
         {3, 1, 1, 1, 1, {3, 3, 3}, {32, 32, 32}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}});
     this->template Run<3>();
 }
+
+TYPED_TEST(TestGroupedConvndFwd2dLargeCases, Test2DLargeCases)
+{
+    // Case larger than 2GB
+    this->conv_params.push_back(
+        {2, 1, 64, 4, 192, {2, 2}, {224, 224}, {224, 224}, {0, 0}, {0, 0}, {0, 0}});
+    this->template Run<2>();
+}

test/grouped_convnd_fwd/test_grouped_convnd_fwd_multi_ab_interface.cpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2023-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <cstdlib>
 #include <iostream>
@@ -207,7 +207,7 @@ TEST_F(TestGroupedConvndFwdMultiAInterface, MultiA)
     std::array<const void*, NumAs> as{nullptr, nullptr};
     const void* b = nullptr;
 
-    EXPECT_TRUE(this->template Run(as, b));
+    EXPECT_TRUE(this->Run(as, b));
 }
 
 TEST_F(TestGroupedConvndFwdMultiBInterface, MultiB)
@@ -215,7 +215,7 @@ TEST_F(TestGroupedConvndFwdMultiBInterface, MultiB)
     const void* a = nullptr;
     std::array<const void*, NumBs> bs{nullptr, nullptr};
 
-    EXPECT_TRUE(this->template Run(a, bs));
+    EXPECT_TRUE(this->Run(a, bs));
 }
 
 TEST_F(TestGroupedConvndFwdMultiABInterface, MultiAB)
@@ -223,7 +223,7 @@ TEST_F(TestGroupedConvndFwdMultiABInterface, MultiAB)
     std::array<const void*, NumAs> as{nullptr, nullptr};
     std::array<const void*, NumBs> bs{nullptr, nullptr};
 
-    EXPECT_TRUE(this->template Run(as, bs));
+    EXPECT_TRUE(this->Run(as, bs));
 }
 
 TEST_F(TestGroupedConvndFwdInterface, SingleAB)
@@ -231,5 +231,5 @@ TEST_F(TestGroupedConvndFwdInterface, SingleAB)
     const void* a = nullptr;
     const void* b = nullptr;
 
-    EXPECT_TRUE(this->template Run(a, b));
+    EXPECT_TRUE(this->Run(a, b));
 }

test/smfmac_op/CMakeLists.txt

+add_gtest_executable(test_smfmac_op smfmac_op_xdl.cpp)
+target_link_libraries(test_smfmac_op PRIVATE utility)

test/smfmac_op/smfmac_op.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <algorithm>
+#include <cstdlib>
+#include <iostream>
+#include <numeric>
+#include <tuple>
+#include <vector>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "test/smfmac_op/smfmac_op_util.hpp"
+
+template <typename Src1Type,
+          ck::index_t Src1VecSize,
+          typename Src2Type,
+          ck::index_t Src2VecSize,
+          typename DstType,
+          ck::index_t AccVecSize,
+          typename GPUAccType,
+          typename CPUAccType,
+          ck::index_t M,
+          ck::index_t N,
+          ck::index_t K>
+bool run_test()
+{
+    using Row         = ck::tensor_layout::gemm::RowMajor;
+    using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+    bool pass         = true;
+
+    const auto matmul_default = ck::smfmac_op_util::matmul<Src1Type,
+                                                           Src1VecSize,
+                                                           Src2Type,
+                                                           Src2VecSize,
+                                                           GPUAccType,
+                                                           AccVecSize,
+                                                           DstType,
+                                                           M,
+                                                           N,
+                                                           K>;
+
+    const auto smfmac_kernel_container = std::make_tuple(matmul_default);
+
+    ck::static_for<0, 1, 1>{}([&](auto i) {
+        pass &=
+            ck::smfmac_op_util::TestSmfmac<decltype(std::get<ck::Number<i>{}>(
+                                               smfmac_kernel_container)),
+                                           Src1Type,
+                                           Src2Type,
+                                           DstType,
+                                           GPUAccType,
+                                           CPUAccType,
+                                           decltype(Row{}),
+                                           decltype(Row{}),
+                                           decltype(Row{}),
+                                           PassThrough,
+                                           PassThrough,
+                                           PassThrough,
+                                           AccVecSize,
+                                           M,
+                                           N,
+                                           K>{}(std::get<ck::Number<i>{}>(smfmac_kernel_container));
+    });
+
+    return pass;
+}
+int main(int, char*[])
+{
+    bool pass = true;
+    // clang-format off
+    //              |   Src1Type| Src1VecSize|    Src2Type| Src2VecSize| DstType| DstVecSize|  GPUAccType| CPUAccType| M| N| K|
+    pass &= run_test< ck::half_t,           4,  ck::half_t,           8,   float,          4,       float,      float,16,16,32>();
+    pass &= run_test<ck::bhalf_t,           4, ck::bhalf_t,           8,   float,          4,       float,      float,16,16,32>();
+    pass &= run_test< ck::half_t,           4,  ck::half_t,           8,   float,         16,       float,      float,32,32,16>();
+    pass &= run_test<ck::bhalf_t,           4, ck::bhalf_t,           8,   float,         16,       float,      float,32,32,16>();
+    // clang-format on
+
+    std::cout << "TestGemm ..... " << (pass ? "SUCCESS" : "FAILURE") << std::endl;
+    return pass;
+}

test/smfmac_op/smfmac_op_util.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/utility/amd_smfmac.hpp"
+#include "ck/library/utility/fill.hpp"
+
+namespace ck {
+namespace smfmac_op_util {
+
+template <typename src_vec1, typename src_vec2, typename acc_vec>
+__device__ void
+builtin_smfmac_naive_selector(const src_vec1&, const src_vec2&, const int32_t&, acc_vec&)
+{
+}
+
+template <>
+__device__ void
+builtin_smfmac_naive_selector<half4_t,
+                              half8_t,
+                              StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 4, true>>(
+    const half4_t& reg_a,
+    const half8_t& reg_b,
+    const int32_t& reg_idx,
+    StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 4, true>& reg_c)
+{
+    intrin_smfmac_f32_16x16x32f16<16, 16>::Run(
+        reg_a, reg_b, reg_idx, reg_c.GetVectorTypeReference(Number<0>{}));
+}
+
+template <>
+__device__ void
+builtin_smfmac_naive_selector<bhalf4_t,
+                              bhalf8_t,
+                              StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 4, true>>(
+    const bhalf4_t& reg_a,
+    const bhalf8_t& reg_b,
+    const int32_t& reg_idx,
+    StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 4, true>& reg_c)
+{
+    intrin_smfmac_f32_16x16x32bf16<16, 16>::Run(
+        reg_a, reg_b, reg_idx, reg_c.GetVectorTypeReference(Number<0>{}));
+}
+
+template <>
+__device__ void builtin_smfmac_naive_selector<
+    half4_t,
+    half8_t,
+    StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 16, true>>(
+    const half4_t& reg_a,
+    const half8_t& reg_b,
+    const int32_t& reg_idx,
+    StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 16, true>& reg_c)
+{
+    intrin_smfmac_f32_32x32x16f16<32, 32>::Run(
+        reg_a, reg_b, reg_idx, reg_c.GetVectorTypeReference(Number<0>{}));
+}
+
+template <>
+__device__ void builtin_smfmac_naive_selector<
+    bhalf4_t,
+    bhalf8_t,
+    StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 16, true>>(
+    const bhalf4_t& reg_a,
+    const bhalf8_t& reg_b,
+    const int32_t& reg_idx,
+    StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, float, 1, 16, true>& reg_c)
+{
+    intrin_smfmac_f32_32x32x16bf16<32, 32>::Run(
+        reg_a, reg_b, reg_idx, reg_c.GetVectorTypeReference(Number<0>{}));
+}
+
+// Smfmac instructions are using 4:2 structural sparsity, that means that in every contignuous
+// subgroup of 4 elements, atleast 2 must be equal to zero and the position of non-zero elements is
+// stored in idx register to allow selection of corresponding B matrix elements for multiplication.
+// Currently smfmac instructions support only A matrix as sparse
+template <typename src1_t,
+          index_t src1_vec_size,
+          typename src2_t,
+          index_t src2_vec_size,
+          typename acc_t,
+          index_t acc_vec_size,
+          typename dst_t,
+          int32_t M,
+          int32_t N,
+          int32_t K>
+__global__ void matmul(const src1_t* a, const src2_t* b, dst_t* c)
+{
+    __shared__ src1_t a_shared[M * K];
+    __shared__ src2_t b_shared[K * N];
+    const int lane = threadIdx.x;
+    // smfmac's A part is storing only non-zero elements in 2VGPRs
+    // smfmac's B part is storing all elements in 4VGPRs
+    using src1_vec      = typename vector_type<src1_t, src1_vec_size>::type;
+    using src1_full_vec = typename vector_type<src1_t, src1_vec_size * 2>::type;
+    using src2_vec      = typename vector_type<src2_t, src2_vec_size>::type;
+    src1_vec a_frag     = {};
+    src2_vec b_frag     = {};
+
+    src1_full_vec a_temp = {};
+    src2_vec b_temp      = {};
+    // initialize c fragment to 0
+    using acc_vec = StaticBufferTupleOfVector<AddressSpaceEnum::Vgpr, acc_t, 1, acc_vec_size, true>;
+    acc_vec c_thread_buf_;
+
+    for(int i = 0; i < 8; ++i)
+    {
+        a_temp[i] = a[(lane % M) * K + (lane / M) * 8 + i]; // M K
+    }
+
+    for(int i = 0; i < 8; ++i)
+    {
+        b_temp[i] = b[(8 * (lane / N) + i) * N + (lane % N)]; // K N
+    }
+
+    __syncthreads();
+
+    for(int i = 0; i < 8; ++i)
+    {
+        a_shared[(lane % M) * K + (lane / M) * 8 + i] = a_temp[i];
+    }
+    for(int i = 0; i < 8; ++i)
+    {
+        b_shared[(8 * (lane / N) + i) * N + (lane % N)] = b_temp[i];
+    }
+
+    __syncthreads();
+
+    // Idx must be a 32-bit register and it is storing 4 2-bit indexes of A's non zero elements.
+    // It starts with last two elements of every 4 elements subgroup set as non-zero
+    int32_t idx = 0b11101110;
+    // Bit masks are for zeroing 0-3rd position of idx
+    static constexpr int32_t bit_clear_masks[4] = {0b11, 0b1100, 0b110000, 0b11000000};
+
+    src1_t curr_val;
+    int32_t a_pos = 0;
+    for(int j = 0; j < 2; ++j)
+    {
+        a_pos = j * 2;
+        for(int i = 0; i < 4; ++i)
+        {
+            curr_val = a_shared[(lane % M) * K + (lane / M) * 8 + 4 * j + i];
+            if(curr_val != 0.0f)
+            {
+                idx &= ~bit_clear_masks[a_pos];
+                idx |= (i % 4) << 2 * a_pos;
+                a_frag[a_pos] = curr_val;
+                a_pos++;
+            }
+        }
+    }
+
+    for(int i = 0; i < 8; ++i)
+    {
+        b_frag[i] = b_shared[(8 * (lane / N) + i) * N + (lane % N)];
+    }
+
+    builtin_smfmac_naive_selector<src1_vec, src2_vec, acc_vec>(a_frag, b_frag, idx, c_thread_buf_);
+    __syncthreads();
+
+    // store results from unpacked c_thread_buf_ output
+    if constexpr(K == 32)
+    {
+        static_for<0, acc_vec_size, 1>{}([&](auto i) {
+            c[(4 * (lane / 16) + i) * N + lane % 16] =
+                ck::type_convert<dst_t>(c_thread_buf_[Number<i>{}]);
+        });
+    }
+    else
+    {
+        static_for<0, acc_vec_size, 1>{}([&](auto i) {
+            c[((8 * (i / 4)) % 32 + 4 * (lane / 32) + i % 4) * N + lane % 32] =
+                ck::type_convert<dst_t>(c_thread_buf_[Number<i>{}]);
+        });
+    }
+}
+
+struct GemmParams
+{
+    GemmParams() : M(16), N(16), K(32), StrideA(32), StrideB(16), StrideC(16), alpha(1), beta(0) {}
+
+    ck::index_t M;
+    ck::index_t N;
+    ck::index_t K;
+
+    ck::index_t StrideA;
+    ck::index_t StrideB;
+    ck::index_t StrideC;
+
+    float alpha;
+    float beta;
+};
+
+template <typename GemmInstance,
+          typename ADataType,
+          typename BDataType,
+          typename CDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CElementwiseOperation>
+void RunHostGEMM(const Tensor<ADataType>& A,
+                 const Tensor<BDataType>& B,
+                 Tensor<CDataType>& C,
+                 AElementwiseOperation a_element_op,
+                 BElementwiseOperation b_element_op,
+                 CElementwiseOperation c_element_op)
+{
+    auto ref_gemm     = GemmInstance{};
+    auto ref_invoker  = ref_gemm.MakeInvoker();
+    auto ref_argument = ref_gemm.MakeArgument(A, B, C, a_element_op, b_element_op, c_element_op);
+
+    ref_invoker.Run(ref_argument);
+}
+
+template <typename KernelType, typename ADataType, typename BDataType, typename CDataType>
+bool RunDeviceGEMM(KernelType kernel,
+                   const Tensor<ADataType>& A,
+                   const Tensor<BDataType>& B,
+                   Tensor<CDataType>& C)
+{
+    DeviceMem a_m_k_device_buf(sizeof(ADataType) * A.mDesc.GetElementSpaceSize());
+    DeviceMem b_n_k_device_buf(sizeof(BDataType) * B.mDesc.GetElementSpaceSize());
+    DeviceMem c_m_n_device_buf(sizeof(CDataType) * C.mDesc.GetElementSpaceSize());
+
+    a_m_k_device_buf.ToDevice(A.mData.data());
+    b_n_k_device_buf.ToDevice(B.mData.data());
+    kernel<<<1, 64>>>(static_cast<const ADataType*>(a_m_k_device_buf.GetDeviceBuffer()),
+                      static_cast<const BDataType*>(b_n_k_device_buf.GetDeviceBuffer()),
+                      static_cast<CDataType*>(c_m_n_device_buf.GetDeviceBuffer()));
+    c_m_n_device_buf.FromDevice(C.mData.data());
+
+    return true;
+}
+
+template <typename DeviceSmfmac,
+          typename ADataType,
+          typename BDataType,
+          typename CDataType,
+          typename GPUAccDataType,
+          typename CPUAccDataType,
+          typename ALayout,
+          typename BLayout,
+          typename CLayout,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CElementwiseOperation,
+          index_t CAccNum,
+          index_t M,
+          index_t N,
+          index_t K>
+struct TestSmfmac
+{
+    auto PrepareGemmTensor(const ck::smfmac_op_util::GemmParams& params)
+    {
+        auto f_host_tensor_descriptor =
+            [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+                if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+                {
+                    return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                                std::vector<std::size_t>({stride, 1}));
+                }
+                else
+                {
+                    return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
+                                                std::vector<std::size_t>({1, stride}));
+                }
+            };
+
+        Tensor<ADataType> a_m_k(
+            f_host_tensor_descriptor(params.M, params.K, params.StrideA, ALayout{}));
+        Tensor<BDataType> b_n_k(
+            f_host_tensor_descriptor(params.K, params.N, params.StrideB, BLayout{}));
+        Tensor<CDataType> c_m_n_host_result(
+            f_host_tensor_descriptor(params.M, params.N, params.StrideC, CLayout{}));
+        Tensor<CDataType> c_m_n_device_result(
+            f_host_tensor_descriptor(params.M, params.N, params.StrideC, CLayout{}));
+
+        auto f_generate_tensor_value = [](auto& tensor, auto type) {
+            using dataType = decltype(type);
+            tensor.GenerateTensorValue(GeneratorTensor_2<dataType>{-5, 5});
+        };
+
+        f_generate_tensor_value(a_m_k, ADataType{});
+        f_generate_tensor_value(b_n_k, BDataType{});
+        ck::utils::TransformIntoStructuralSparsity<ADataType>{}(a_m_k);
+
+        return std::make_tuple(a_m_k, b_n_k, c_m_n_host_result, c_m_n_device_result);
+    }
+
+    auto operator()(const DeviceSmfmac& smfmac_kernel)
+    {
+        std::cout << "ALayout = " << ALayout{}.name << ", BLayout = " << BLayout{}.name
+                  << ", CLayout = " << CLayout{}.name << std::endl;
+
+        // Arrange
+        ck::smfmac_op_util::GemmParams params;
+        params.M       = M;
+        params.N       = N;
+        params.K       = K;
+        params.StrideA = K; // M K
+        params.StrideB = N; // K N
+        params.StrideC = N; // M N
+
+        auto host_tensors = PrepareGemmTensor(params);
+
+        const Tensor<ADataType>& a  = std::get<0>(host_tensors);
+        const Tensor<BDataType>& b  = std::get<1>(host_tensors);
+        Tensor<CDataType>& c_host   = std::get<2>(host_tensors);
+        Tensor<CDataType>& c_device = std::get<3>(host_tensors);
+
+        auto a_element_op = AElementwiseOperation{};
+        auto b_element_op = BElementwiseOperation{};
+        auto c_element_op = CElementwiseOperation{};
+
+        using ReferenceGemmInstance =
+            ck::tensor_operation::host::ReferenceGemm<ADataType,
+                                                      BDataType,
+                                                      CDataType,
+                                                      CPUAccDataType,
+                                                      AElementwiseOperation,
+                                                      BElementwiseOperation,
+                                                      CElementwiseOperation>;
+        ck::smfmac_op_util::RunHostGEMM<ReferenceGemmInstance>(
+            a, b, c_host, a_element_op, b_element_op, c_element_op);
+
+        // Act
+        bool is_supported = ck::smfmac_op_util::RunDeviceGEMM(smfmac_kernel, a, b, c_device);
+
+        if(is_supported)
+        {
+            // Assert
+            bool res = false;
+            if(std::is_same<CDataType, float>::value)
+            {
+                res = ck::utils::check_err(c_device.mData, c_host.mData);
+                std::cout << (res ? "SUCCESS" : "FAILURE") << std::endl;
+            }
+            else
+            {
+                std::cout << "UNSUPPORTED CDataType" << std::endl;
+            }
+
+            return res;
+        }
+        else
+        {
+            return true;
+        }
+    }
+};
+
+} // namespace smfmac_op_util
+} // namespace ck

test/smfmac_op/smfmac_op_xdl.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <algorithm>
+#include <cstdlib>
+#include <iostream>
+#include <numeric>
+#include <tuple>
+#include <vector>
+
+#include "ck/ck.hpp"
+#include "gtest/gtest.h"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "test/smfmac_op/smfmac_op_util.hpp"
+
+using BF16        = ck::bhalf_t;
+using F16         = ck::half_t;
+using F32         = float;
+using Row         = ck::tensor_layout::gemm::RowMajor;
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+template <typename Tuple>
+class TestSmfmac : public ::testing::Test
+{
+    protected:
+    using Src1Type                           = std::tuple_element_t<0, Tuple>;
+    static constexpr ck::index_t Src1VecSize = std::tuple_element_t<1, Tuple>{}.value;
+    using Src2Type                           = std::tuple_element_t<2, Tuple>;
+    static constexpr ck::index_t Src2VecSize = std::tuple_element_t<3, Tuple>{}.value;
+    using DstType                            = std::tuple_element_t<4, Tuple>;
+    static constexpr ck::index_t AccVecSize  = std::tuple_element_t<5, Tuple>{}.value;
+    using GPUAccType                         = std::tuple_element_t<6, Tuple>;
+    using CPUAccType                         = std::tuple_element_t<7, Tuple>;
+    static constexpr ck::index_t M           = std::tuple_element_t<8, Tuple>{}.value;
+    static constexpr ck::index_t N           = std::tuple_element_t<9, Tuple>{}.value;
+    static constexpr ck::index_t K           = std::tuple_element_t<10, Tuple>{}.value;
+
+    void Run()
+    {
+        bool pass                     = true;
+        constexpr auto matmul_default = ck::smfmac_op_util::matmul<Src1Type,
+                                                                   Src1VecSize,
+                                                                   Src2Type,
+                                                                   Src2VecSize,
+                                                                   GPUAccType,
+                                                                   AccVecSize,
+                                                                   DstType,
+                                                                   M,
+                                                                   N,
+                                                                   K>;
+
+        constexpr auto smfmac_kernel_container = std::make_tuple(matmul_default);
+
+        ck::static_for<0, std::tuple_size_v<decltype(smfmac_kernel_container)>, 1>{}([&](auto i) {
+            pass &= ck::smfmac_op_util::TestSmfmac<
+                std::tuple_element_t<i.value, decltype(smfmac_kernel_container)>,
+                Src1Type,
+                Src2Type,
+                DstType,
+                GPUAccType,
+                CPUAccType,
+                decltype(Row{}),
+                decltype(Row{}),
+                decltype(Row{}),
+                PassThrough,
+                PassThrough,
+                PassThrough,
+                AccVecSize,
+                M,
+                N,
+                K>{}(std::get<ck::Number<i>{}>(smfmac_kernel_container));
+        });
+
+        EXPECT_TRUE(pass);
+    }
+};
+
+template <ck::index_t N>
+using I = ck::Number<N>;
+
+using KernelTypes =
+    ::testing::Types<std::tuple<F16, I<4>, F16, I<8>, F32, I<4>, F32, F32, I<16>, I<16>, I<32>>,
+                     std::tuple<BF16, I<4>, BF16, I<8>, F32, I<4>, F32, F32, I<16>, I<16>, I<32>>,
+                     std::tuple<F16, I<4>, F16, I<8>, F32, I<16>, F32, F32, I<32>, I<32>, I<16>>,
+                     std::tuple<BF16, I<4>, BF16, I<8>, F32, I<16>, F32, F32, I<32>, I<32>, I<16>>>;
+
+TYPED_TEST_SUITE(TestSmfmac, KernelTypes);
+TYPED_TEST(TestSmfmac, TestSmfmacFP16BF16) { this->Run(); }

test/wmma_op/wmma_op_util.hpp

@@ -140,10 +140,18 @@ __global__ void matmul(const src_t* a, const src_t* b, dst_t* c)
         p_shared[8 * 16 * lane_hi + 8 * lane_lo + ele + 16 * 16] = b_temp[ele];
     }
 
+#ifdef __gfx12__
+    asm volatile("\
+    s_wait_dscnt 0x0 \n \
+    s_barrier_signal -1 \n \
+    s_barrier_wait -1 \
+    " ::);
+#else
     asm volatile("\
     s_waitcnt lgkmcnt(0) \n \
     s_barrier \
     " ::);
+#endif
 
     for(int ele = 0; ele < 16; ++ele)
     {
@@ -155,10 +163,18 @@ __global__ void matmul(const src_t* a, const src_t* b, dst_t* c)
         a_frag[ele] = p_shared[(ele / 8) * 16 * 8 + 8 * lane + ele % 8];
     }
 
+#ifdef __gfx12__
+    asm volatile("\
+    s_wait_dscnt 0x0 \n \
+    s_barrier_signal -1 \n \
+    s_barrier_wait -1 \
+    " ::);
+#else
     asm volatile("\
     s_waitcnt lgkmcnt(0) \n \
     s_barrier \
     " ::);
+#endif
 
     // sync threads, similar to mma_sync
     // __syncthreads();