Merge branch 'develop' into ck_migraphx_integration

example/ck_tile/03_gemm/gemm_basic.cpp

@@ -41,18 +41,39 @@ template <typename LayoutA,
 float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s)
 {
     // The kPadA, kPadB, kPadC & kBlockPerCu should also come from the Codegen part.
-    constexpr bool kPadA = true;
-    constexpr bool kPadB = true;
+    constexpr bool kPadA        = true;
+    constexpr bool kPadB        = true;
+    constexpr bool kTilePermute = false;
 
     constexpr int kBlockPerCu = 1;
 
     using TilePartitioner = ck_tile::GemmTilePartitioner<GemmShape>;
-    using GemmEpilogue    = ck_tile::Default2DEpilogue<
-        ck_tile::Default2DEpilogueProblem<AccDataType, CDataType, kPadA, kPadB>>;
+
+    // The rank and permutation will also be generate out by the CodeGen part.
+    constexpr ck_tile::index_t kOutputRank = 2;
+
+    // Whether doing the CShuffle (transpose before the global memory), depending on the output
+    // layout.
+    constexpr bool CShuffleEpilogue =
+        std::is_same_v<LayoutC, ck_tile::tensor_layout::gemm::ColumnMajor>;
+
+    using GemmEpilogue = std::conditional_t<
+        CShuffleEpilogue,
+        ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<AccDataType,
+                                                                   CDataType,
+                                                                   kPadA,
+                                                                   kPadB,
+                                                                   kTilePermute,
+                                                                   kOutputRank,
+                                                                   1,
+                                                                   0,
+                                                                   TilePartitioner::kM,
+                                                                   TilePartitioner::kN>>,
+        ck_tile::Default2DEpilogue<
+            ck_tile::Default2DEpilogueProblem<AccDataType, CDataType, kPadA, kPadB>>>;
     // ToDo: Will add the codegen part to test different pipeline policies in GEMM.
     // Now we only use the BlockGemmASmemBSmemCRegV1DefaultPolicy.
-    using Kernel =
-        ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue, LayoutA, LayoutB, LayoutC>;
+    using Kernel = ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
 
     auto kargs = Kernel::MakeKargs(args.p_a,
                                    args.p_b,
@@ -255,15 +276,13 @@ int main(int argc, char* argv[])
                                ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
                                ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
 
-    using CodegenPipelineProblem = ck_tile::BlockGemmPipelineProblem<ADataType,
-                                                                     BDataType,
-                                                                     AccDataType,
-                                                                     CodegenGemmShape,
-                                                                     kPadA,
-                                                                     kPadB,
-                                                                     kPadC>;
+    using CodegenGemmTraits = ck_tile::
+        TileGemmTraits<kPadA, kPadB, kPadC, matrix_a_layout, matrix_b_layout, matrix_c_layout>;
+
+    using CodegenPipelineProblem = ck_tile::
+        GemmPipelineProblem<ADataType, BDataType, AccDataType, CodegenGemmShape, CodegenGemmTraits>;
 
-    using CodegenGemmPipeline = ck_tile::BlockGemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
+    using CodegenGemmPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
 
     invoke_gemm<ck_tile::half_t,
                 matrix_a_layout,
@@ -341,7 +360,13 @@ int main(int argc, char* argv[])
         ck_tile::HostTensor<CDataType> c_host_gpu_ref(c_dimensions);
         ck_tile::DeviceMem c_gpu_buf(c_host_gpu_ref.get_element_space_size_in_bytes());
 
-        ck_tile::reference_gemm_gpu<ADataType, BDataType, AccDataType, CDataType>(
+        ck_tile::reference_gemm_gpu<ADataType,
+                                    BDataType,
+                                    AccDataType,
+                                    CDataType,
+                                    matrix_a_layout,
+                                    matrix_b_layout,
+                                    matrix_c_layout>(
             a_buf, b_buf, c_gpu_buf, M, N, K, stride_a, stride_b, stride_c);
 
         c_buf.FromDevice(c_host_gpu_ref.data());

example/ck_tile/04_img2col/README.md

@@ -6,7 +6,8 @@ This folder contains example for Image to Column using ck_tile tile-programming
 ```
 # in the root of ck_tile
 mkdir build && cd build
-sh ../script/cmake-ck-dev.sh  ../ <arch>  # you can replace this <arch> to gfx90a, gfx942...
+# you can replace <arch> with the appropriate architecture (for example gfx90a or gfx942) or leave it blank
+sh ../script/cmake-ck-dev.sh  ../ <arch>
 make tile_example_img2col -j
 ```
 This will result in an executable `build/bin/tile_example_img2col`

include/ck/config.h.in

@@ -97,13 +97,6 @@
 #cmakedefine CK_ENABLE_DL_KERNELS @CK_ENABLE_DL_KERNELS@
 #endif
 
-//
-// Instances supports in the current CK build
-//
-#ifndef CK_ENABLE_INSTANCES_ONLY
-#cmakedefine CK_ENABLE_INSTANCES_ONLY @CK_ENABLE_INSTANCES_ONLY@
-#endif
-
 //
 // CK kernels which support XDL (MI series)
 //

include/ck/host_utility/kernel_launch.hpp

@@ -66,6 +66,9 @@ float launch_and_time_kernel(const StreamConfig& stream_config,
 
         hip_check_error(hipEventElapsedTime(&total_time, start, stop));
 
+        hip_check_error(hipEventDestroy(start));
+        hip_check_error(hipEventDestroy(stop));
+
         return total_time / nrepeat;
     }
     else
@@ -143,6 +146,9 @@ float launch_and_time_kernel_with_preprocess(const StreamConfig& stream_config,
 
         hip_check_error(hipEventElapsedTime(&total_time, start, stop));
 
+        hip_check_error(hipEventDestroy(start));
+        hip_check_error(hipEventDestroy(stop));
+
         return total_time / nrepeat;
     }
     else

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

@@ -308,7 +308,7 @@ struct BlockwiseGemmXdlops_pipeline_v1_ab_scale<BlockGemmPipelineScheduler::Intr
                                 typename vector_type<ComputeDataType,
                                                      xdlops_gemm.K1PerXdlops>::type;
 
-                            xdlops_gemm.template Run(
+                            xdlops_gemm.template Run<>(
                                 a_thread_vec.template AsType<mfma_input_type>(),
                                 b_thread_vec.template AsType<mfma_input_type>(),
                                 c_thread_buf_per_scale.GetVectorTypeReference(I0));
@@ -390,9 +390,10 @@ struct BlockwiseGemmXdlops_pipeline_v1_ab_scale<BlockGemmPipelineScheduler::Intr
                         using mfma_input_type =
                             typename vector_type<ComputeDataType, xdlops_gemm.K1PerXdlops>::type;
 
-                        xdlops_gemm.template Run(a_thread_vec.template AsType<mfma_input_type>(),
-                                                 b_thread_vec.template AsType<mfma_input_type>(),
-                                                 c_thread_buf_per_scale.GetVectorTypeReference(I0));
+                        xdlops_gemm.template Run<>(
+                            a_thread_vec.template AsType<mfma_input_type>(),
+                            b_thread_vec.template AsType<mfma_input_type>(),
+                            c_thread_buf_per_scale.GetVectorTypeReference(I0));
                     });
                     static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) {
                         constexpr index_t c_offset =

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

@@ -350,7 +350,7 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale<BlockGemmPipelineScheduler::Intr
                                     typename vector_type<ComputeDataType,
                                                          xdlops_gemm.K1PerXdlops>::type;
 
-                                xdlops_gemm.template Run(
+                                xdlops_gemm.template Run<>(
                                     a_thread_vec.template AsType<mfma_input_type>(),
                                     b_thread_vec.template AsType<mfma_input_type>(),
                                     c_thread_buf_per_scale.GetVectorTypeReference(I0));
@@ -443,7 +443,7 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale<BlockGemmPipelineScheduler::Intr
                                 typename vector_type<ComputeDataType,
                                                      xdlops_gemm.K1PerXdlops>::type;
 
-                            xdlops_gemm.template Run(
+                            xdlops_gemm.template Run<>(
                                 a_thread_vec.template AsType<mfma_input_type>(),
                                 b_thread_vec.template AsType<mfma_input_type>(),
                                 c_thread_buf_per_scale.GetVectorTypeReference(I0));
@@ -518,9 +518,10 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale<BlockGemmPipelineScheduler::Intr
                         using mfma_input_type =
                             typename vector_type<ComputeDataType, xdlops_gemm.K1PerXdlops>::type;
 
-                        xdlops_gemm.template Run(a_thread_vec.template AsType<mfma_input_type>(),
-                                                 b_thread_vec.template AsType<mfma_input_type>(),
-                                                 c_thread_buf_per_scale.GetVectorTypeReference(I0));
+                        xdlops_gemm.template Run<>(
+                            a_thread_vec.template AsType<mfma_input_type>(),
+                            b_thread_vec.template AsType<mfma_input_type>(),
+                            c_thread_buf_per_scale.GetVectorTypeReference(I0));
                     });
                     static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) {
                         constexpr index_t c_offset =
@@ -575,9 +576,10 @@ struct BlockwiseGemmXdlops_pipeline_v2_ab_scale<BlockGemmPipelineScheduler::Intr
                         using mfma_input_type =
                             typename vector_type<ComputeDataType, xdlops_gemm.K1PerXdlops>::type;
 
-                        xdlops_gemm.template Run(a_thread_vec.template AsType<mfma_input_type>(),
-                                                 b_thread_vec.template AsType<mfma_input_type>(),
-                                                 c_thread_buf_per_scale.GetVectorTypeReference(I0));
+                        xdlops_gemm.template Run<>(
+                            a_thread_vec.template AsType<mfma_input_type>(),
+                            b_thread_vec.template AsType<mfma_input_type>(),
+                            c_thread_buf_per_scale.GetVectorTypeReference(I0));
                     });
                     static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) {
                         constexpr index_t c_offset =

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

@@ -427,7 +427,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale<BlockGemmPipelineScheduler::Intr
                                 typename vector_type<ComputeDataType,
                                                      xdlops_gemm.K1PerXdlops>::type;
 
-                            xdlops_gemm.template Run(
+                            xdlops_gemm.template Run<>(
                                 a_thread_vec.template AsType<mfma_input_type>(),
                                 b_thread_vec.template AsType<mfma_input_type>(),
                                 c_thread_buf_per_scale.GetVectorTypeReference(I0));
@@ -504,9 +504,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale<BlockGemmPipelineScheduler::Intr
                         using mfma_input_type =
                             typename vector_type<ComputeDataType, xdlops_gemm.K1PerXdlops>::type;
 
-                        xdlops_gemm.template Run(a_thread_vec.template AsType<mfma_input_type>(),
-                                                 b_thread_vec.template AsType<mfma_input_type>(),
-                                                 c_thread_buf_per_scale.GetVectorTypeReference(I0));
+                        xdlops_gemm.template Run<>(
+                            a_thread_vec.template AsType<mfma_input_type>(),
+                            b_thread_vec.template AsType<mfma_input_type>(),
+                            c_thread_buf_per_scale.GetVectorTypeReference(I0));
                     });
                     static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) {
                         constexpr index_t c_offset =

include/ck/tensor_operation/gpu/device/device_cgemm.hpp

 // 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.
 
 #pragma once
 #include "device_base.hpp"
@@ -31,13 +31,13 @@ struct DeviceCGemm : public BaseOperator
                                                               CElementwiseOperation c_element_op,
                                                               ck::index_t KBatch = 1) = 0;
 
-    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer()   = 0;
     virtual std::size_t GetWorkspaceSize(index_t MRaw,
                                          index_t NRaw,
                                          index_t KRaw,
                                          index_t StrideA,
                                          index_t StrideB,
-                                         index_t StrideC)     = 0;
+                                         index_t StrideC) const = 0;
 };
 
 template <typename AElementwiseOperation,

include/ck/tensor_operation/gpu/device/impl/device_cgemm_4gemm_xdl_cshuffle.hpp

@@ -598,10 +598,26 @@ struct DeviceCGemm_4Gemm_Xdl_CShuffle
                                  [[maybe_unused]] index_t K,
                                  [[maybe_unused]] index_t StrideA,
                                  [[maybe_unused]] index_t StrideB,
-                                 index_t StrideC) override
+                                 index_t StrideC) const override
     {
         return 2 * sizeof(CDataType) * GetCElementSpaceSize(M, N, StrideC);
     }
+
+    std::size_t GetWorkSpaceSize(const BaseArgument* base_arg) const override
+    {
+        const auto* parg = dynamic_cast<const Argument*>(base_arg);
+
+        if(!parg)
+        {
+            std::ostringstream err;
+            err << "Provided argument pointer is not of an Argument class!"
+                << " In " << __FILE__ << ":" << __LINE__ << ", in function: " << __func__;
+            throw std::runtime_error(err.str());
+        }
+
+        return GetWorkspaceSize(
+            parg->M, parg->N, parg->K, parg->StrideA, parg->StrideB, parg->StrideC);
+    }
 };
 
 } // namespace device

include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multi_abd_xdl_fixed_nk.hpp

@@ -64,7 +64,7 @@ __global__ void
     const index_t N = gemm_desc_ptr[group_id].N;
     const index_t K = gemm_desc_ptr[group_id].K;
 
-    if(M * N * K == 0)
+    if(M == 0 || N == 0 || K == 0)
         return;
 
     const auto StrideAs = gemm_desc_ptr[group_id].StrideAs;

include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_splitk_xdl_cshuffle_two_stage.hpp

@@ -345,7 +345,7 @@ struct DeviceGroupedGemmMultipleDSplitKXdlCShuffleTwoStage
                 const index_t N = gemm_descs[i].N_;
                 const index_t K = gemm_descs[i].K_;
 
-                if(M * N * K == 0)
+                if(M == 0 || N == 0 || K == 0)
                 {
                     skipped_group_count_++;
                     continue;

include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp

@@ -109,7 +109,7 @@ __global__ void
             N = gemm_desc_ptr[group_id].N;
             K = gemm_desc_ptr[group_id].K;
 
-            if(M * N * K == 0)
+            if(M == 0 || N == 0 || K == 0)
             {
                 grid_size_grp = 0;
                 continue;

include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_fixed_nk.hpp

@@ -68,7 +68,7 @@ __global__ void
     const index_t N = gemm_desc_ptr[group_id].N;
     const index_t K = gemm_desc_ptr[group_id].K;
 
-    if(M * N * K == 0)
+    if(M == 0 || N == 0 || K == 0)
         return;
 
     const auto StrideA  = gemm_desc_ptr[group_id].StrideA;

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

@@ -419,6 +419,12 @@ struct UnaryAbs
 
         y = ck::math::abs(x);
     };
+
+    template <>
+    __host__ __device__ void operator()(f8_t& y, const f8_t& x) const
+    {
+        y = ck::type_convert<f8_t>(ck::math::abs(ck::type_convert<float>(x)));
+    };
 };
 
 struct UnarySqrt

include/ck/utility/math_v2.hpp

@@ -81,6 +81,8 @@ static inline __host__ bool isnan(half_t x)
     return (xx & 0x7FFF) > 0x7C00;
 };
 
+static inline __host__ bool isnan(f8_t x) { return (x & 0x80); };
+
 #ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
 static inline __host__ bool isnan(int4_t x)
 {
@@ -531,6 +533,8 @@ static inline __device__ bool isnan(half_t x)
     return (xx & 0x7FFF) > 0x7C00;
 };
 
+static inline __device__ bool isnan(f8_t x) { return (x & 0x80); };
+
 static inline __device__ half_t sqrt(half_t x)
 {
     return static_cast<half_t>(__builtin_amdgcn_sqrtf(static_cast<float>(x)));

include/ck_tile/core/container/thread_buffer.hpp

@@ -58,7 +58,7 @@ struct thread_buffer {
     template <index_t I> CK_TILE_HOST_DEVICE constexpr const auto& at() const           { return get(I); }
     template <index_t I> CK_TILE_HOST_DEVICE constexpr auto& at(number<I>)              { return get(I); }
     template <index_t I> CK_TILE_HOST_DEVICE constexpr const auto& at(number<I>) const  { return get(I); }
-
+    
     template <typename X_,
               typename std::enable_if<has_same_scalar_type<value_type, X_>::value, bool>::type = false>
     CK_TILE_HOST_DEVICE constexpr auto _get_as() const

include/ck_tile/host/arg_parser.hpp

@@ -50,12 +50,22 @@ class ArgParser
         }
         return *this;
     }
-    void print()
+    void print() const
     {
+        // find max key length
+        std::string::size_type max_key_length = 11;
+        for(auto& key : keys)
+        {
+            if(max_key_length < key.length())
+            {
+                max_key_length = key.length();
+            }
+        }
+
         printf("args:\n");
         for(auto& key : keys)
         {
-            auto value = input_map[key];
+            auto value = input_map.at(key);
             std::vector<std::string> help_text_lines;
             size_t pos = 0;
             for(size_t next_pos = value.help_text.find('\n', pos); next_pos != std::string::npos;)
@@ -69,8 +79,7 @@ class ArgParser
                 std::string(value.help_text.begin() + pos, value.help_text.end()));
 
             std::string default_value = std::string("(default:") + value.value + std::string(")");
-
-            std::cout << std::setw(2) << std::setw(12 - value.name.length()) << "-" << key
+            std::cout << std::setw(1 + max_key_length - value.name.length()) << "-" << key
                       << std::setw(4) << " " << help_text_lines[0] << " " << default_value
                       << std::endl;
 
@@ -78,7 +87,8 @@ class ArgParser
                 help_next_line != help_text_lines.end();
                 ++help_next_line)
             {
-                std::cout << std::setw(17) << " " << *help_next_line << std::endl;
+                std::cout << std::setw(1 + max_key_length + 4) << " " << *help_next_line
+                          << std::endl;
             }
         }
     }

include/ck_tile/host/convolution_parameter.hpp

@@ -13,7 +13,6 @@ namespace conv {
 
 struct ConvParam
 {
-    ConvParam();
     ConvParam(ck_tile::index_t n_dim,
               ck_tile::index_t group_count,
               ck_tile::index_t n_batch,
@@ -199,11 +198,6 @@ struct ConvParam
     }
 };
 
-ConvParam::ConvParam()
-    : ConvParam::ConvParam(2, 1, 128, 256, 192, {3, 3}, {71, 71}, {2, 2}, {1, 1}, {1, 1}, {1, 1})
-{
-}
-
 CK_TILE_HOST std::string get_conv_param_parser_helper_msg()
 {
     std::string msg;

include/ck_tile/host/reference/reference_gemm.hpp

@@ -27,7 +27,9 @@ CK_TILE_HOST void reference_gemm(const HostTensor<ADataType>& a_m_k,
                                  const BElementOp& b_element_op     = {},
                                  const ACCElementOp& acc_element_op = {})
 {
-    const int N = b_n_k.mDesc.get_lengths()[0];
+    const int N = (std::is_same_v<LayoutB, tensor_layout::gemm::ColumnMajor>)
+                      ? b_n_k.mDesc.get_lengths()[0]
+                      : b_n_k.mDesc.get_lengths()[1];
     const int K = (std::is_same_v<LayoutA, tensor_layout::gemm::RowMajor>)
                       ? a_m_k.mDesc.get_lengths()[1]
                       : a_m_k.mDesc.get_lengths()[0];
@@ -45,20 +47,31 @@ CK_TILE_HOST void reference_gemm(const HostTensor<ADataType>& a_m_k,
                 ADataType v_a = (std::is_same_v<LayoutA, tensor_layout::gemm::RowMajor>)
                                     ? a_element_op(a_m_k(m, k))
                                     : a_element_op(a_m_k(k, m));
-                BDataType v_b = b_element_op(b_n_k(n, k));
+                BDataType v_b = (std::is_same_v<LayoutB, tensor_layout::gemm::ColumnMajor>)
+                                    ? b_element_op(b_n_k(n, k))
+                                    : b_element_op(b_n_k(k, n));
 
                 v_acc += ck_tile::type_convert<AccDataType>(v_a) *
                          ck_tile::type_convert<AccDataType>(v_b);
             }
 
-            c_m_n(m, n) = ck_tile::type_convert<CDataType>(acc_element_op(v_acc));
+            CDataType& c_ref = (std::is_same_v<LayoutC, tensor_layout::gemm::RowMajor>)
+                                   ? c_m_n(m, n)
+                                   : c_m_n(n, m);
+            c_ref            = ck_tile::type_convert<CDataType>(acc_element_op(v_acc));
         }
     };
 
     make_ParallelTensorFunctor(f, M)(std::thread::hardware_concurrency());
 }
 
-template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
+template <typename ADataType,
+          typename BDataType,
+          typename AccDataType,
+          typename CDataType,
+          typename LayoutA,
+          typename LayoutB,
+          typename LayoutC>
 __global__ void naive_gemm_kernel(ADataType* A,
                                   BDataType* B,
                                   CDataType* C,
@@ -76,18 +89,32 @@ __global__ void naive_gemm_kernel(ADataType* A,
     if(row < M && col < N)
     {
         AccDataType acc = 0.0;
-
         for(int k = 0; k < K; ++k)
         {
-            acc += static_cast<AccDataType>(A[row * strideA + k]) *
-                   static_cast<AccDataType>(B[col * strideB + k]);
+            // Adjust indexing based on matrix layout
+            int a_index = (std::is_same_v<LayoutA, tensor_layout::gemm::RowMajor>)
+                              ? row * strideA + k
+                              : k * strideA + row;
+            int b_index = (std::is_same_v<LayoutB, tensor_layout::gemm::ColumnMajor>)
+                              ? col * strideB + k
+                              : k * strideB + col;
+            acc += static_cast<AccDataType>(A[a_index]) * static_cast<AccDataType>(B[b_index]);
         }
 
-        C[row * strideC + col] = acc; // Store as AccDataType
+        int c_index = (std::is_same_v<LayoutC, tensor_layout::gemm::RowMajor>)
+                          ? row * strideC + col
+                          : col * strideC + row;
+        C[c_index]  = acc;
     }
 }
 
-template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
+template <typename ADataType,
+          typename BDataType,
+          typename AccDataType,
+          typename CDataType,
+          typename LayoutA,
+          typename LayoutB,
+          typename LayoutC>
 void reference_gemm_gpu(DeviceMem& a_device,
                         DeviceMem& b_device,
                         DeviceMem& c_device,
@@ -145,7 +172,7 @@ void reference_gemm_gpu(DeviceMem& a_device,
     int numThreadsPerBlock = 256; // Common choice for threads per block
     int numBlocks          = (totalElements + numThreadsPerBlock - 1) / numThreadsPerBlock;
 
-    naive_gemm_kernel<ADataType, BDataType, AccDataType, CDataType>
+    naive_gemm_kernel<ADataType, BDataType, AccDataType, CDataType, LayoutA, LayoutB, LayoutC>
         <<<numBlocks, numThreadsPerBlock>>>(d_A, d_B, d_C, M, N, K, stride_a, stride_b, stride_c);
     errC = hipMemcpy(
         c_device.GetDeviceBuffer(), d_C, M * N * sizeof(CDataType), hipMemcpyDeviceToHost);