Merge branch 'develop' of...

Merge branch 'develop' of https://github.com/ROCmSoftwarePlatform/composable_kernel into add_fp16_wmma_conv_instance

client_example/19_pool_fwd/max_pool2d_fwd.cpp → client_example/19_pool/max_pool2d_fwd.cpp

@@ -10,14 +10,18 @@
 #include "ck/tensor_operation/gpu/device/device_pool_fwd.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
-#include "ck/library/tensor_operation_instance/gpu/pool2d_fwd.hpp"
+#include "ck/library/tensor_operation_instance/gpu/pool3d_fwd.hpp"
 
 using InDataType    = ck::half_t;
 using OutDataType   = ck::half_t;
 using IndexDataType = int32_t;
 
-constexpr ck::index_t InOutRank  = 4;
-constexpr ck::index_t WindowRank = 2;
+// We use pool3d to implement pool2d in this example
+using InLayout  = ck::tensor_layout::convolution::NDHWC;
+using OutLayout = ck::tensor_layout::convolution::NDHWC;
+
+constexpr ck::index_t InOutRank  = 5;
+constexpr ck::index_t WindowRank = 3;
 #if 1
 constexpr auto ReduceOpId  = ck::ReduceTensorOp::MAX;
 constexpr bool OutputIndex = true;
@@ -42,6 +46,35 @@ struct SimpleDeviceMem
     void* p_mem_;
 };
 
+void TransformPool2dparamToPool3d(std::vector<ck::index_t>& input_lengths,
+                                  std::vector<ck::index_t>& window_lengths,
+                                  std::vector<ck::index_t>& output_lengths,
+                                  std::vector<ck::index_t>& input_stride,
+                                  std::vector<ck::index_t>& output_stride,
+                                  std::vector<ck::index_t>& indices_stride,
+                                  std::vector<ck::index_t>& window_strides,
+                                  std::vector<ck::index_t>& window_dilations,
+                                  std::vector<ck::index_t>& input_left_pads,
+                                  std::vector<ck::index_t>& input_right_pads,
+                                  std::vector<ck::index_t>& pooling_dims)
+{
+    // NCHW to NCDHW
+    input_lengths.insert(input_lengths.begin() + 2, 1);
+    output_lengths.insert(output_lengths.begin() + 2, 1);
+    input_stride.insert(input_stride.begin() + 2, 0);
+    output_stride.insert(output_stride.begin() + 2, 0);
+    indices_stride.insert(indices_stride.begin() + 2, 0);
+
+    // YX to ZYX
+    window_lengths.insert(window_lengths.begin(), 1);
+    window_strides.insert(window_strides.begin(), 0);
+    window_dilations.insert(window_dilations.begin(), 0);
+    input_left_pads.insert(input_left_pads.begin(), 0);
+    input_right_pads.insert(input_right_pads.begin(), 0);
+
+    pooling_dims = {2, 3, 4};
+}
+
 int main(int argc, char* argv[])
 {
     ck::index_t N                 = 2;
@@ -52,21 +85,27 @@ int main(int argc, char* argv[])
     ck::index_t Wi                = 30;
     ck::index_t window_stride_h   = 2;
     ck::index_t window_stride_w   = 2;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_h     = 1;
     ck::index_t in_left_pad_w     = 1;
     ck::index_t in_right_pad_h    = 1;
     ck::index_t in_right_pad_w    = 1;
 
-    ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Y) / window_stride_h + 1;
-    ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - X) / window_stride_w + 1;
+    const ck::index_t Ys = (Y - 1) * window_dilation_h + 1;
+    const ck::index_t Xs = (X - 1) * window_dilation_w + 1;
+    ck::index_t Ho       = (Hi + in_left_pad_h + in_right_pad_h - Ys) / window_stride_h + 1;
+    ck::index_t Wo       = (Wi + in_left_pad_w + in_right_pad_w - Xs) / window_stride_w + 1;
 
     // Pool API only support the order of NCHW
     std::vector<ck::index_t> in_length              = {N, C, Hi, Wi};
     std::vector<ck::index_t> out_length             = {N, C, Ho, Wo};
     std::vector<ck::index_t> window_spatial_lengths = {Y, X};
     std::vector<ck::index_t> window_strides         = {window_stride_h, window_stride_w};
+    std::vector<ck::index_t> window_dilations       = {window_dilation_h, window_dilation_w};
     std::vector<ck::index_t> input_left_pads        = {in_left_pad_h, in_left_pad_w};
     std::vector<ck::index_t> input_right_pads       = {in_right_pad_h, in_right_pad_w};
+    std::vector<ck::index_t> pooling_dims           = {2, 3};
 
     std::size_t in_tensor_size  = N * C * Hi * Wi;
     std::size_t out_tensor_size = N * C * Ho * Wo;
@@ -75,6 +114,18 @@ int main(int argc, char* argv[])
     std::vector<ck::index_t> in_tensor_stride  = {C * Hi * Wi, 1, Wi * C, C};
     std::vector<ck::index_t> out_tensor_stride = {C * Ho * Wo, 1, Wo * C, C};
 
+    TransformPool2dparamToPool3d(in_length,
+                                 window_spatial_lengths,
+                                 out_length,
+                                 in_tensor_stride,
+                                 out_tensor_stride,
+                                 out_tensor_stride,
+                                 window_strides,
+                                 window_dilations,
+                                 input_left_pads,
+                                 input_right_pads,
+                                 pooling_dims);
+
     SimpleDeviceMem in_device_buf(sizeof(InDataType) * in_tensor_size);
     SimpleDeviceMem out_device_buf(sizeof(OutDataType) * out_tensor_size);
     SimpleDeviceMem out_indices_device_buf(sizeof(IndexDataType) * out_tensor_size);
@@ -84,6 +135,8 @@ int main(int argc, char* argv[])
                                                                  InDataType,
                                                                  OutDataType,
                                                                  IndexDataType,
+                                                                 InLayout,
+                                                                 OutLayout,
                                                                  ReduceOpId,
                                                                  OutputIndex>;
 
@@ -116,9 +169,10 @@ int main(int argc, char* argv[])
             out_tensor_stride,
             out_tensor_stride,
             window_strides,
+            window_dilations,
             input_left_pads,
             input_right_pads,
-            {2, 3});
+            pooling_dims);
 
         auto invoker_ptr = op_ptr->MakeInvokerPointer();
 
@@ -175,9 +229,10 @@ int main(int argc, char* argv[])
             out_tensor_stride,
             out_tensor_stride,
             window_strides,
+            window_dilations,
             input_left_pads,
             input_right_pads,
-            {2, 3});
+            pooling_dims);
 
         auto invoker_ptr = op_ptr->MakeInvokerPointer();
 

client_example/20_image_to_column/CMakeLists.txt

+add_executable(client_image_to_column image_to_column.cpp)
+target_link_libraries(client_image_to_column PRIVATE composable_kernel::device_operations)

client_example/20_image_to_column/image_to_column.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <cstdlib>
+#include <iomanip>
+#include <iostream>
+#include <iterator>
+#include <numeric>
+#include <vector>
+
+#include "ck/ck.hpp"
+#include "ck/library/tensor_operation_instance/gpu/image_to_column.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+
+using InDataType  = ck::half_t;
+using OutDataType = ck::half_t;
+
+using InLayout = ck::tensor_layout::convolution::GNHWC;
+
+static constexpr ck::index_t NumDimSpatial = 2;
+static constexpr ck::index_t G             = 1;
+static constexpr ck::index_t N             = 32; // batch size
+static constexpr ck::index_t C             = 32; // input channel (per group)
+static constexpr ck::index_t Y             = 3;  // filter H
+static constexpr ck::index_t X             = 3;  // filter W
+static constexpr ck::index_t Hi            = 28; // input H
+static constexpr ck::index_t Wi            = 28; // input W
+static constexpr ck::index_t Ho            = 28; // output H
+static constexpr ck::index_t Wo            = 28; // output W
+
+struct SimpleDeviceMem
+{
+    SimpleDeviceMem() = delete;
+
+    SimpleDeviceMem(std::size_t mem_size) : p_mem_{}
+    {
+        (void)hipMalloc(static_cast<void**>(&p_mem_), mem_size);
+    }
+
+    void* GetDeviceBuffer() { return p_mem_; }
+
+    ~SimpleDeviceMem() { (void)hipFree(p_mem_); }
+
+    void* p_mem_;
+};
+
+int main()
+{
+
+    std::array<ck::index_t, 2> in_spatial_lengths{Hi, Wi};
+    std::array<ck::index_t, 2> wei_spatial_lengths{Y, X};
+    std::array<ck::index_t, 2> out_spatial_lengths{Ho, Wo};
+
+    // We have NHWGC in memory space (G is dummy)
+    // However, CK's API only accept length and stride with order of GNCHW
+    // Hence, we need to adjust the order of stride
+    std::array<ck::index_t, 5> in_strides{C, Hi * Wi * G * C, 1, Wi * G * C, G * C};
+    std::array<ck::index_t, 2> out_strides{Y * X * C, 1};
+
+    std::array<ck::index_t, NumDimSpatial> filter_strides{1, 1};
+    std::array<ck::index_t, NumDimSpatial> filter_dilations{1, 1};
+    std::array<ck::index_t, NumDimSpatial> input_left_pads{1, 1};
+    std::array<ck::index_t, NumDimSpatial> input_right_pads{1, 1};
+
+    SimpleDeviceMem in(sizeof(InDataType) * N * Hi * Wi * G * C);
+    SimpleDeviceMem out(sizeof(OutDataType) * N * Ho * Wo * Y * X * C);
+
+    using DeviceOp = ck::tensor_operation::device::
+        DeviceImageToColumn<NumDimSpatial, InLayout, InDataType, OutDataType>;
+
+    // get device op instances
+    const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+        DeviceOp>::GetInstances();
+
+    std::cout << "found " << op_ptrs.size() << " instances" << std::endl;
+
+    std::string best_op_name;
+    int best_op_id        = -1;
+    float best_avg_time   = std::numeric_limits<float>::max();
+    float best_gb_per_sec = 0;
+
+    // profile device operation instances
+    std::cout << "Run all instances and do timing" << std::endl;
+
+    for(int i = 0; i < op_ptrs.size(); ++i)
+    {
+        auto& op_ptr        = op_ptrs[i];
+        auto argument_ptr   = op_ptr->MakeArgumentPointer(in.GetDeviceBuffer(),
+                                                        out.GetDeviceBuffer(),
+                                                        N,
+                                                        C,
+                                                        in_spatial_lengths,
+                                                        out_spatial_lengths,
+                                                        wei_spatial_lengths,
+                                                        in_strides,
+                                                        out_strides,
+                                                        filter_strides,
+                                                        filter_dilations,
+                                                        input_left_pads,
+                                                        input_right_pads);
+        auto invoker_ptr    = op_ptr->MakeInvokerPointer();
+        std::string op_name = op_ptr->GetTypeString();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true});
+
+            std::size_t num_bytes = sizeof(InDataType) * N * Hi * Wi * G * C +
+                                    sizeof(OutDataType) * N * Ho * Wo * Y * X * C;
+
+            float gb_per_sec = num_bytes / 1.E6 / avg_time;
+
+            std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << gb_per_sec << " GB/s, "
+                      << op_name << std::endl;
+
+            if(avg_time < best_avg_time)
+            {
+                best_op_id      = i;
+                best_op_name    = op_name;
+                best_avg_time   = avg_time;
+                best_gb_per_sec = gb_per_sec;
+            }
+        }
+        else
+        {
+            std::cerr << op_name << " does not support this problem" << std::endl;
+        }
+    }
+
+    if(best_op_id < 0)
+    {
+        std::cerr << "no suitable instance" << std::endl;
+        return EXIT_FAILURE;
+    }
+
+    std::cout << "Best Perf: " << std::setw(10) << best_avg_time << " ms, " << best_gb_per_sec
+              << " GB/s, " << best_op_name << std::endl;
+
+    // run the best intance
+    {
+        auto& op_ptr = op_ptrs[best_op_id];
+        std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString()
+                  << std::endl;
+        auto argument_ptr = op_ptr->MakeArgumentPointer(in.GetDeviceBuffer(),
+                                                        out.GetDeviceBuffer(),
+                                                        N,
+                                                        C,
+                                                        in_spatial_lengths,
+                                                        out_spatial_lengths,
+                                                        wei_spatial_lengths,
+                                                        in_strides,
+                                                        out_strides,
+                                                        filter_strides,
+                                                        filter_dilations,
+                                                        input_left_pads,
+                                                        input_right_pads);
+
+        auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
+        }
+
+        std::cout << "Done" << std::endl;
+    }
+}

client_example/20_splitk_gemm/CMakeLists.txt

+add_executable(client_splitK_gemm splitK_gemm_fp16_f8.cpp)
+target_link_libraries(client_splitK_gemm PRIVATE composable_kernel::device_operations)

client_example/20_splitk_gemm/splitK_gemm_fp16_f8.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iomanip>
+#include <vector>
+#include <iostream>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_gemm_splitk.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/gpu/gemm_splitk.hpp"
+
+using F8  = ck::f8_t;
+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 AElementOp = PassThrough;
+using BElementOp = PassThrough;
+using CElementOp = PassThrough;
+
+using ADataType = F8;
+using BDataType = F16;
+using CDataType = F16;
+
+using ALayout = Row;
+using BLayout = Col;
+using CLayout = Row;
+
+struct SimpleDeviceMem
+{
+    SimpleDeviceMem() = delete;
+
+    SimpleDeviceMem(std::size_t mem_size) : p_mem_{}
+    {
+        (void)hipMalloc(static_cast<void**>(&p_mem_), mem_size);
+    }
+
+    void* GetDeviceBuffer() { return p_mem_; }
+
+    ~SimpleDeviceMem() { (void)hipFree(p_mem_); }
+
+    void* p_mem_;
+};
+
+int main(int argc, char* argv[])
+{
+    // 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 StrideC = 4096;
+
+    ck::index_t KBatch = 1;
+
+    if(argc == 1)
+    {
+        // use default case
+    }
+    else if(argc == 8)
+    {
+        M = std::stoi(argv[1]);
+        N = std::stoi(argv[2]);
+        K = std::stoi(argv[3]);
+
+        StrideA = std::stoi(argv[4]);
+        StrideB = std::stoi(argv[5]);
+        StrideC = std::stoi(argv[6]);
+
+        KBatch = std::stoi(argv[7]);
+    }
+    else
+    {
+        printf("arg1 to 7: M, N, K, StrideA, StrideB, StrideC, KBatch\n");
+        exit(0);
+    }
+
+    auto f_matrix_space_size =
+        [](std::size_t nRow, std::size_t nCol, std::size_t stride, auto layout) {
+            using Layout = decltype(layout);
+
+            if constexpr(std::is_same<Layout, ck::tensor_layout::gemm::RowMajor>::value)
+            {
+                return (nRow - 1) * stride + nCol;
+            }
+            else
+            {
+                return (nCol - 1) * stride + nRow;
+            }
+        };
+
+    SimpleDeviceMem a_device_buf(sizeof(ADataType) * f_matrix_space_size(M, K, StrideA, ALayout{}));
+    SimpleDeviceMem b_device_buf(sizeof(BDataType) * f_matrix_space_size(K, N, StrideB, BLayout{}));
+    SimpleDeviceMem c_device_buf(sizeof(CDataType) * f_matrix_space_size(M, N, StrideC, CLayout{}));
+
+    using DeviceOp = ck::tensor_operation::device::DeviceGemmSplitK<
+        ALayout,
+        BLayout,
+        CLayout,
+        ADataType,
+        BDataType,
+        CDataType,
+        ck::tensor_operation::element_wise::PassThrough,
+        ck::tensor_operation::element_wise::PassThrough,
+        ck::tensor_operation::element_wise::PassThrough>;
+
+    // get device op instances
+    const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+        DeviceOp>::GetInstances();
+
+    std::cout << "found " << op_ptrs.size() << " instances" << std::endl;
+
+    const auto a_element_op = AElementOp{};
+    const auto b_element_op = BElementOp{};
+    const auto c_element_op = CElementOp{};
+
+    std::string best_op_name;
+    bool found            = false;
+    int best_op_id        = -1;
+    float best_ave_time   = 0;
+    float best_tflops     = 0;
+    float best_gb_per_sec = 0;
+
+    // profile device operation instances
+    std::cout << "Run all instances and do timing" << std::endl;
+
+    for(int i = 0; i < op_ptrs.size(); ++i)
+    {
+        auto& op_ptr = op_ptrs[i];
+
+        auto argument_ptr = op_ptr->MakeArgumentPointer(a_device_buf.GetDeviceBuffer(),
+                                                        b_device_buf.GetDeviceBuffer(),
+                                                        c_device_buf.GetDeviceBuffer(),
+                                                        M,
+                                                        N,
+                                                        K,
+                                                        StrideA,
+                                                        StrideB,
+                                                        StrideC,
+                                                        a_element_op,
+                                                        b_element_op,
+                                                        c_element_op,
+                                                        KBatch);
+
+        auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+        std::string op_name = op_ptr->GetTypeString();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true});
+
+            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(CDataType) * 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: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, "
+                      << gb_per_sec << " GB/s, " << op_name << std::endl;
+
+            if(tflops > best_tflops)
+            {
+                found           = true;
+                best_op_id      = i;
+                best_op_name    = op_name;
+                best_tflops     = tflops;
+                best_ave_time   = ave_time;
+                best_gb_per_sec = gb_per_sec;
+            }
+        }
+        else
+        {
+            std::cout << op_name << " does not support this problem" << std::endl;
+        }
+    }
+
+    std::cout << "Best Perf: " << best_ave_time << " ms, " << best_tflops << " TFlops, "
+              << best_gb_per_sec << " GB/s, " << best_op_name << std::endl;
+
+    // run the best intance
+    {
+        auto& op_ptr = op_ptrs[best_op_id];
+
+        std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString()
+                  << std::endl;
+
+        auto argument_ptr = op_ptr->MakeArgumentPointer(a_device_buf.GetDeviceBuffer(),
+                                                        b_device_buf.GetDeviceBuffer(),
+                                                        c_device_buf.GetDeviceBuffer(),
+                                                        M,
+                                                        N,
+                                                        K,
+                                                        StrideA,
+                                                        StrideB,
+                                                        StrideC,
+                                                        a_element_op,
+                                                        b_element_op,
+                                                        c_element_op,
+                                                        KBatch);
+
+        auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
+        }
+
+        std::cout << "Done" << std::endl;
+    }
+
+    return 0;
+}

client_example/21_grouped_gemm_bias/CMakeLists.txt

+add_executable(client_grouped_gemm_fixed_nk_bias_fp16 grouped_gemm_fixed_nk_bias_fp16.cpp)
+target_link_libraries(client_grouped_gemm_fixed_nk_bias_fp16 PRIVATE composable_kernel::device_operations)

client_example/21_grouped_gemm_bias/grouped_gemm_fixed_nk_bias_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iomanip>
+#include <iostream>
+#include <vector>
+#include <random>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_grouped_gemm_fixed_nk.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/gpu/grouped_gemm_bias.hpp"
+
+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 Add         = ck::tensor_operation::element_wise::Add;
+
+using ADataType  = F16;
+using BDataType  = F16;
+using D0DataType = F32;
+using DsDataType = ck::Tuple<D0DataType>;
+using EDataType  = F32;
+
+using ALayout  = Row;
+using BLayout  = Row;
+using D0Layout = Row;
+using DsLayout = ck::Tuple<D0Layout>;
+using ELayout  = Row;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = Add;
+
+struct SimpleDeviceMem
+{
+    SimpleDeviceMem() = delete;
+
+    SimpleDeviceMem(std::size_t mem_size) : p_mem_{}
+    {
+        (void)hipMalloc(static_cast<void**>(&p_mem_), mem_size);
+    }
+
+    void* GetDeviceBuffer() { return p_mem_; }
+
+    ~SimpleDeviceMem() { (void)hipFree(p_mem_); }
+
+    void* p_mem_;
+};
+
+int main()
+{
+    std::vector<int> Ms, Ns, Ks, StrideAs, StrideBs, StrideEs;
+
+    int sum_of_m = 0;
+
+    Ms = {167, 183, 177, 181, 153, 139, 156, 173, 163, 150, 204, 184, 168, 156, 168, 148};
+
+    int group_count = Ms.size();
+
+    for(int i = 0; i < group_count; ++i)
+    {
+        Ns.push_back(768);
+        Ks.push_back(4608);
+
+        StrideAs.push_back(std::is_same<Row, ALayout>::value ? Ks[i] : Ms[i]);
+        StrideBs.push_back(std::is_same<Row, BLayout>::value ? Ns[i] : Ks[i]);
+        StrideEs.push_back(std::is_same<Row, ELayout>::value ? Ns[i] : Ms[i]);
+
+        sum_of_m += Ms[i];
+    }
+
+    auto f_matrix_space_size =
+        [](std::size_t nRow, std::size_t nCol, std::size_t stride, auto layout) {
+            using Layout = decltype(layout);
+
+            if constexpr(std::is_same<Layout, ck::tensor_layout::gemm::RowMajor>::value)
+            {
+                return (nRow - 1) * stride + nCol;
+            }
+            else
+            {
+                return (nCol - 1) * stride + nRow;
+            }
+        };
+
+    std::vector<SimpleDeviceMem> a_dev_bufs, b_dev_bufs, d0_dev_bufs, e_dev_bufs;
+
+    a_dev_bufs.reserve(group_count);
+    b_dev_bufs.reserve(group_count);
+    d0_dev_bufs.reserve(group_count);
+    e_dev_bufs.reserve(group_count);
+
+    std::vector<void*> p_e;
+
+    p_e.reserve(group_count);
+
+    std::vector<ck::tensor_operation::device::GemmDesc> gemm_descs;
+
+    gemm_descs.reserve(group_count);
+
+    std::vector<ck::tensor_operation::device::GroupedGemmKernelArgument<1>>
+        grouped_gemm_kernel_args_;
+    grouped_gemm_kernel_args_.reserve(group_count);
+
+    for(int i = 0; i < group_count; ++i)
+    {
+        a_dev_bufs.emplace_back(sizeof(ADataType) *
+                                f_matrix_space_size(Ms[i], Ks[i], StrideAs[i], ALayout{}));
+        b_dev_bufs.emplace_back(sizeof(BDataType) *
+                                f_matrix_space_size(Ks[i], Ns[i], StrideBs[i], BLayout{}));
+        d0_dev_bufs.emplace_back(sizeof(D0DataType) *
+                                 f_matrix_space_size(Ms[i], Ns[i], 0, D0Layout{}));
+        e_dev_bufs.emplace_back(sizeof(EDataType) *
+                                f_matrix_space_size(Ms[i], Ns[i], StrideEs[i], ELayout{}));
+
+        gemm_descs.push_back({sum_of_m, Ns[i], Ks[i], 1, StrideBs[i], 1, {0}});
+
+        p_e.push_back(e_dev_bufs[i].GetDeviceBuffer());
+
+        grouped_gemm_kernel_args_.push_back(
+            {a_dev_bufs[i].GetDeviceBuffer(),
+             b_dev_bufs[i].GetDeviceBuffer(),
+             std::array<const void*, 1>{d0_dev_bufs[i].GetDeviceBuffer()},
+             e_dev_bufs[i].GetDeviceBuffer(),
+             Ms[i],
+             Ns[i],
+             Ks[i],
+             StrideAs[i],
+             StrideBs[i],
+             std::array<ck::index_t, 1>{0},
+             StrideEs[i]});
+    }
+
+    using DeviceOp = ck::tensor_operation::device::DeviceGroupedGemmFixedNK<ALayout,
+                                                                            BLayout,
+                                                                            DsLayout,
+                                                                            ELayout,
+                                                                            ADataType,
+                                                                            BDataType,
+                                                                            DsDataType,
+                                                                            EDataType,
+                                                                            AElementOp,
+                                                                            BElementOp,
+                                                                            CDEElementOp>;
+
+    // get device op instances
+    const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+        DeviceOp>::GetInstances();
+
+    std::cout << "found " << op_ptrs.size() << " instances" << std::endl;
+
+    const auto a_element_op   = AElementOp{};
+    const auto b_element_op   = BElementOp{};
+    const auto cde_element_op = CDEElementOp{};
+
+    std::string best_op_name;
+    bool found            = false;
+    int best_op_id        = -1;
+    float best_ave_time   = 0;
+    float best_tflops     = 0;
+    float best_gb_per_sec = 0;
+
+    // profile device operation instances
+    std::cout << "Run all instances and do timing" << std::endl;
+
+    std::vector<const void*> p_a = {}, p_b = {};
+    std::vector<std::array<const void*, 1>> p_ds = {};
+
+    for(int i = 0; i < op_ptrs.size(); ++i)
+    {
+        auto& op_ptr = op_ptrs[i];
+
+        auto argument_ptr = op_ptr->MakeArgumentPointer(
+            p_a, p_b, p_ds, p_e, gemm_descs, a_element_op, b_element_op, cde_element_op);
+
+        auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+        SimpleDeviceMem grouped_gemm_kernel_args_dev(
+            op_ptr->GetDeviceKernelArgSize(argument_ptr.get()));
+
+        SimpleDeviceMem grouped_gemm_workspace_dev(op_ptr->GetWorkSpaceSize(argument_ptr.get()));
+
+        std::string op_name = op_ptr->GetTypeString();
+
+        hipGetErrorString(hipMemcpy(grouped_gemm_kernel_args_dev.GetDeviceBuffer(),
+                                    grouped_gemm_kernel_args_.data(),
+                                    op_ptr->GetDeviceKernelArgSize(argument_ptr.get()),
+                                    hipMemcpyHostToDevice));
+
+        op_ptr->SetWorkSpacePointer(argument_ptr.get(),
+                                    grouped_gemm_workspace_dev.GetDeviceBuffer());
+
+        op_ptr->SetDeviceKernelArgs(argument_ptr.get(),
+                                    grouped_gemm_kernel_args_dev.GetDeviceBuffer());
+
+        op_ptr->SetKBatch(argument_ptr.get(), 2);
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true});
+
+            std::size_t flop = 0, num_btype = 0;
+            for(std::size_t j = 0; j < gemm_descs.size(); ++j)
+            {
+                flop += std::size_t(2) * Ms[j] * Ns[j] * Ks[j];
+
+                num_btype += sizeof(ADataType) * Ms[j] * Ks[j] + sizeof(BDataType) * Ks[j] * Ns[j] +
+                             sizeof(EDataType) * Ms[j] * Ns[j];
+            }
+
+            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, " << op_name << std::endl;
+
+            if(tflops > best_tflops)
+            {
+                found           = true;
+                best_op_id      = i;
+                best_op_name    = op_name;
+                best_tflops     = tflops;
+                best_ave_time   = ave_time;
+                best_gb_per_sec = gb_per_sec;
+            }
+        }
+        else
+        {
+            std::cout << op_name << " does not support this problem" << std::endl;
+        }
+    }
+
+    std::cout << "Best Perf: " << best_ave_time << " ms, " << best_tflops << " TFlops, "
+              << best_gb_per_sec << " GB/s, " << best_op_name << std::endl;
+
+    return 0;
+}

client_example/CMakeLists.txt

@@ -5,29 +5,50 @@ add_compile_options(-std=c++17)
 if (DTYPES)
     add_definitions(-DDTYPES)
     if (DTYPES MATCHES "int8")
-                add_definitions(-D__int8__)
+        add_definitions(-DCK_ENABLE_INT8)
+        if(NOT DEFINED ${CK_ENABLE_INT8})
+            set(CK_ENABLE_INT8 "ON")
+        endif()
     endif()
     if (DTYPES MATCHES "fp8")
-                add_definitions(-D__fp8__)
+        add_definitions(-DCK_ENABLE_FP8)
+        if(NOT DEFINED ${CK_ENABLE_FP8})
+            set(CK_ENABLE_FP8 "ON")
+        endif()
     endif()
     if (DTYPES MATCHES "fp16")
-                add_definitions(-D__fp16__)
+        add_definitions(-DCK_ENABLE_FP16)
+        if(NOT DEFINED ${CK_ENABLE_FP16})
+            set(CK_ENABLE_FP16 "ON")
+        endif()
     endif()
     if (DTYPES MATCHES "fp32")
-                add_definitions(-D__fp32__)
+        add_definitions(-DCK_ENABLE_FP32)
+        if(NOT DEFINED ${CK_ENABLE_FP32})
+            set(CK_ENABLE_FP32 "ON")
+        endif()
     endif()
     if (DTYPES MATCHES "fp64")
-                add_definitions(-D__fp64__)
+        add_definitions(-DCK_ENABLE_FP64)
+        if(NOT DEFINED ${CK_ENABLE_FP64})
+            set(CK_ENABLE_FP64 "ON")
+        endif()
     endif()
     if (DTYPES MATCHES "bf16")
-                add_definitions(-D__bf16__)
+        add_definitions(-DCK_ENABLE_BF16)
+        if(NOT DEFINED ${CK_ENABLE_BF16})
+            set(CK_ENABLE_BF16 "ON")
+        endif()
     endif()
     message("DTYPES macro set to ${DTYPES}")
 else()
-        add_definitions(-D__int8__ -D__fp8__ -D__fp16__ -D__fp32__ -D__fp64__ -D__bf16__)
+    add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16)
+    if(NOT DEFINED ${CK_ENABLE_ALL_DTYPES})
+        set(CK_ENABLE_ALL_DTYPES "ON")
+    endif()
 endif()
 
-find_package(composable_kernel 1.0.0 COMPONENTS device_operations)
+find_package(composable_kernel COMPONENTS device_operations)
 find_package(hip REQUIRED PATHS /opt/rocm)
 message(STATUS "Build with HIP ${hip_VERSION}")
 

cmake/DoxygenDoc.cmake

@@ -309,6 +309,8 @@ XML_OUTPUT
 XML_PROGRAMLISTING
 )
 
+set(WARN_AS_ERROR YES)
+
 set(DOXYGEN_CONFIG_FILE "${CMAKE_CURRENT_BINARY_DIR}/doxygen/doxygen.conf" CACHE PATH "Path to generated doxygen configuration file")
 
 function(add_doxygen_doc)

docs/API_Reference_Guide.rst

@@ -7,8 +7,8 @@ API Reference Guide
 Introduction
 =================
 
-This document contains details of the APIs for the Composable Kernel (CK) library and introduces some of the key design
-principles that are used to write new classes that extend CK functionality.
+This document contains details of the APIs for the Composable Kernel (CK) library and introduces
+some of the key design principles that are used to write new classes that extend CK functionality.
 
 =================
 Using CK API
@@ -30,8 +30,8 @@ DeviceMem
 Kernels For Flashattention
 ---------------------------
 
-The Flashattention algorithm is defined in :cite:t:`dao2022flashattention`.  This sections lists the classes that are
-used in the CK GPU implementation of Flashattention.
+The Flashattention algorithm is defined in :cite:t:`dao2022flashattention`. This sections lists
+the classes that are used in the CK GPU implementation of Flashattention.
 
 **Gridwise classes**
 

docs/Contributors_Guide.rst

@@ -2,7 +2,101 @@
 Contributor's Guide
 ===================
 
-Pull-request guidelines
-=======================
+This chapter explains how to get started contributing to the Composable Kernel project and what are
+the contributing rules.
 
-[TODO]
+Getting started
+===============
+
+#. **Documentation:** Before contributing to the library, familiarize yourself with the
+   `Composable Kernel User Guide <https://rocm.docs.amd.com/projects/composable_kernel/en/latest/>`_.
+   It provides insight into the core concepts, environment configuration, and steps to obtain or
+   build the library. You can also find some of this information in the
+   `README file <https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/develop/README.md>`_
+   on the project's GitHub page.
+#. **Additional reading:** We also recommend reading a `blog post
+   <https://community.amd.com/t5/instinct-accelerators/amd-composable-kernel-library-efficient-fused-kernels-for-ai/ba-p/553224>`_
+   from the AMD Community portal. It offers a deeper understanding of the library's objectives and
+   showcases its performance capabilities.
+#. **General information:** For broader information about AMD products, consider exploring the
+   `AMD Developer Central portal <https://www.amd.com/en/developer.html>`_.
+
+How do I contribute
+===================
+
+We deeply value contributions from our users. You can make an impact by reporting issues or
+proposing code enhancements through pull requests.
+
+Reporting issues
+----------------
+
+We use `Github issues <https://github.com/ROCmSoftwarePlatform/composable_kernel/issues>`_
+to track public bugs and enhancement requests.
+
+If you encounter an issue with the library, please check if the problem has already been
+reported by searching existing issues on GitHub. If your issue seems unique, please submit a new
+issue. All reported issues must include:
+
+* A comprehensive description of the problem, including:
+
+  * What did you observe?
+  * Why do you think it is a bug (if it seems like one)?
+  * What did you expect to happen? What would indicate the resolution of the problem?
+  * Are there any known workarounds?
+
+* Your configuration details, including:
+
+  * Which GPU are you using?
+  * Which OS version are you on?
+  * Which ROCm version are you using?
+  * Are you using a Docker image? If so, which one?
+
+* Steps to reproduce the issue, including:
+
+  * What actions trigger the issue? What are the reproduction steps?
+
+    * If you build the library from scratch, what CMake command did you use?
+
+  * How frequently does this issue happen? Does it reproduce every time? Or is it a sporadic issue?
+
+Before sumbitting any issue, ensure you have addressed all relevant questions from the checklist.
+
+Creating Pull Requests
+----------------------
+
+You can submit `Pull Requests (PR) on GitHub
+<https://github.com/ROCmSoftwarePlatform/composable_kernel/pulls>`_.
+
+All contributors are required to develop their changes on a separate branch and then create a
+pull requrest to merge their changes into the `develop` branch, which is the default
+development branch in the Composable Kernel project. All external contributors must use their own
+forks of the project to develop their changes.
+
+When submitting a Pull Request you should:
+
+* Describe the change providing information about the motivation for the change and a general
+  description of all code modifications.
+
+* Verify and test the change:
+
+  * Run any relevant existing tests.
+  * Write new tests if added functionality is not covered by current tests.
+
+* Ensure your changes align with the coding style defined in the ``.clang-format`` file located in
+  the project's root directory. We leverage `pre-commit` to run `clang-format` automatically. We
+  highly recommend contributors utilize this method to maintain consistent code formatting.
+  Instructions on setting up `pre-commit` can be found in the project's
+  `README file <https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/develop/README.md>`_
+
+* Link your PR to any related issues:
+
+  * If there is an issue that is resolved by your change, please provide a link to the issue in
+    the description of your pull request.
+
+* For larger contributions, structure your change into a sequence of smaller, focused commits, each
+  addressing a particular aspect or fix.
+
+Following the above guidelines ensures a seamless review process and faster assistance from our
+end.
+
+Thank you for your commitment to enhancing the Composable Kernel project! We look forward to collaborating with you.

docs/Supported_Primitives_Guide.rst

@@ -2,15 +2,16 @@
 Supported Primitives Guide
 ==========================
 
-This document contains details of supported primitives in Composable Kernel (CK). In contrast to the API Reference
-Guide, the Supported Primitives Guide is an introduction to the math which underpins the algorithms implemented in CK.
+This document contains details of supported primitives in Composable Kernel (CK). In contrast to the
+API Reference Guide, the Supported Primitives Guide is an introduction to the math which underpins
+the algorithms implemented in CK.
 
 ------------
 Softmax
 ------------
 
-For vectors :math:`x^{(1)}, x^{(2)}, \ldots, x^{(T)}` of size :math:`B` we can decompose the softmax of concatenated
-:math:`x = [ x^{(1)}\ | \ \ldots \ | \ x^{(T)} ]` as,
+For vectors :math:`x^{(1)}, x^{(2)}, \ldots, x^{(T)}` of size :math:`B` we can decompose the
+softmax of concatenated :math:`x = [ x^{(1)}\ | \ \ldots \ | \ x^{(T)} ]` as,
 
 .. math::
    :nowrap:
@@ -25,8 +26,8 @@ For vectors :math:`x^{(1)}, x^{(2)}, \ldots, x^{(T)}` of size :math:`B` we can d
 where :math:`f(x^{(j)}) = \exp( x^{(j)} - m(x^{(j)}) )` is of size :math:`B` and
 :math:`z(x^{(j)}) = f(x_1^{(j)})+ \ldots+ f(x_B^{(j)})` is a scalar.
 
-For a matrix :math:`X` composed of :math:`T_r \times T_c` tiles, :math:`X_{ij}`, of size :math:`B_r \times B_c` we can
-compute the row-wise softmax as follows.
+For a matrix :math:`X` composed of :math:`T_r \times T_c` tiles, :math:`X_{ij}`, of size
+:math:`B_r \times B_c` we can compute the row-wise softmax as follows.
 
 For :math:`j` from :math:`1` to :math:`T_c`, and :math:`i` from :math:`1` to :math:`T_r` calculate,
 

docs/dockerhub.rst

 ===================
-CK docker hub
+CK Docker Hub
 ===================
 
-`Docker hub <https://hub.docker.com/r/rocm/composable_kernel>`_
-
 -------------------------------------
 Why do I need this?
 -------------------------------------
 
-To make our lives easier and bring Composable Kernel dependencies together, we recommend using docker images.
+To make our lives easier and bring Composable Kernel dependencies together, we recommend using
+docker images that can be found on `Docker Hub <https://hub.docker.com/r/rocm/composable_kernel>`_.
 
 -------------------------------------
 So what is Composable Kernel?
 -------------------------------------
 
-Composable Kernel (CK) library aims to provide a programming model for writing performance critical kernels for machine learning workloads across multiple architectures including GPUs, CPUs, etc, through general purpose kernel languages, like HIP C++.
+Composable Kernel (CK) library aims to provide a programming model for writing performance critical
+kernels for machine learning workloads across multiple architectures including GPUs, CPUs, etc,
+through general purpose kernel languages, like HIP C++.
 
 To get the CK library::
 
     git clone https://github.com/ROCmSoftwarePlatform/composable_kernel.git
 
 
-
 run a docker container::
 
     docker run                                                            \
@@ -30,7 +30,7 @@ run a docker container::
     --group-add sudo                                                      \
     -w /root/workspace                                                    \
     -v ${PATH_TO_LOCAL_WORKSPACE}:/root/workspace                         \
-    rocm/composable_kernel:ck_ub20.04_rocm5.3_release                     \
+    rocm/composable_kernel:ck_ub20.04_rocm5.6                             \
     /bin/bash
 
 and build the CK::
@@ -58,7 +58,9 @@ We can also run specific examples or tests like::
     ./bin/example_gemm_xdl_fp16
     ./bin/test_gemm_fp16
 
-For more details visit `CK github repo <https://github.com/ROCmSoftwarePlatform/composable_kernel>`_, `CK examples <https://github.com/ROCmSoftwarePlatform/composable_kernel/tree/develop/example)>`_, `even more CK examples <https://github.com/ROCmSoftwarePlatform/composable_kernel/tree/develop/client_example>`_.
+For more details visit `CK github repository <https://github.com/ROCmSoftwarePlatform/composable_kernel>`_,
+`CK examples <https://github.com/ROCmSoftwarePlatform/composable_kernel/tree/develop/example)>`_,
+`even more CK examples <https://github.com/ROCmSoftwarePlatform/composable_kernel/tree/develop/client_example>`_.
 
 -------------------------------------
 And what is inside?
@@ -74,12 +76,11 @@ The docker images have everything you need for running CK including:
 Which image is right for me?
 -------------------------------------
 
-Let's take a look at the image naming, for example "ck_ub20.04_rocm5.4_release". The image specs are:
+Let's take a look at the image naming, for example ``ck_ub20.04_rocm5.6``. The image specs are:
 
-* "ck" - made for running Composable Kernel
-* "ub20.04" - based on Ubuntu 20.04
-* "rocm5.4" - ROCm platform version 5.4
-* "release" - compiler version is release
+* ``ck`` - made for running Composable Kernel;
+* ``ub20.04`` - based on Ubuntu 20.04;
+* ``rocm5.6`` - ROCm platform version 5.6.
 
 So just pick the right image for your project dependencies and you're all set.
 
@@ -87,7 +88,9 @@ So just pick the right image for your project dependencies and you're all set.
 DIY starts here
 -------------------------------------
 
-If you need to customize a docker image or just can't stop tinkering, feel free to adjust the `Dockerfile <https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/develop/Dockerfile>`_ for your needs.
+If you need to customize a docker image or just can't stop tinkering, feel free to adjust the
+`Dockerfile <https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/develop/Dockerfile>`_
+for your needs.
 
 -------------------------------------
 License

docs/index.rst

@@ -12,12 +12,15 @@ This document contains instructions for installing, using, and contributing to C
 Methodology
 -----------
 
-Composable Kernel (CK) library aims to provide a programming model for writing performance critical kernels for machine learning workloads across multiple architectures including GPUs, CPUs, etc, through general purpose kernel languages, like HIP C++.
+Composable Kernel (CK) library aims to provide a programming model for writing performance critical
+kernels for machine learning workloads across multiple architectures including GPUs, CPUs, etc,
+through general purpose kernel languages, like HIP C++.
 
 CK utilizes two concepts to achieve performance portability and code maintainability:
 
 * A tile-based programming model
-* Algorithm complexity reduction for complex ML operators, using innovative technique we call "Tensor Coordinate Transformation".
+* Algorithm complexity reduction for complex ML operators, using innovative technique we call
+  "Tensor Coordinate Transformation".
 
 .. image:: data/ck_component.png
    :alt: CK Components

docs/sphinx/requirements.in

-rocm-docs-core==0.10.3
+rocm-docs-core>=0.20.0
 sphinxcontrib-bibtex==2.5.0

docs/sphinx/requirements.txt

@@ -38,6 +38,8 @@ docutils==0.16
     #   pydata-sphinx-theme
     #   sphinx
     #   sphinxcontrib-bibtex
+fastjsonschema==2.18.0
+    # via rocm-docs-core
 gitdb==4.0.10
     # via gitpython
 gitpython==3.1.31
@@ -46,20 +48,12 @@ idna==3.4
     # via requests
 imagesize==1.4.1
     # via sphinx
-importlib-metadata==6.0.0
-    # via
-    #   sphinx
-    #   sphinxcontrib-bibtex
-importlib-resources==5.12.0
-    # via rocm-docs-core
 jinja2==3.1.2
     # via
     #   myst-parser
     #   sphinx
 latexcodec==2.0.1
     # via pybtex
-linkify-it-py==1.0.3
-    # via myst-parser
 markdown-it-py==2.2.0
     # via
     #   mdit-py-plugins
@@ -70,7 +64,7 @@ mdit-py-plugins==0.3.5
     # via myst-parser
 mdurl==0.1.2
     # via markdown-it-py
-myst-parser[linkify]==1.0.0
+myst-parser==1.0.0
     # via rocm-docs-core
 packaging==23.0
     # via
@@ -99,18 +93,17 @@ pyjwt[crypto]==2.6.0
     # via pygithub
 pynacl==1.5.0
     # via pygithub
-pytz==2023.3
-    # via babel
 pyyaml==6.0
     # via
     #   myst-parser
     #   pybtex
+    #   rocm-docs-core
     #   sphinx-external-toc
 requests==2.28.2
     # via
     #   pygithub
     #   sphinx
-rocm-docs-core==0.10.3
+rocm-docs-core>=0.20.0
     # via -r requirements.in
 six==1.16.0
     # via
@@ -160,13 +153,7 @@ sphinxcontrib-serializinghtml==1.1.5
     # via sphinx
 typing-extensions==4.5.0
     # via pydata-sphinx-theme
-uc-micro-py==1.0.1
-    # via linkify-it-py
 urllib3==1.26.15
     # via requests
 wrapt==1.15.0
     # via deprecated
-zipp==3.15.0
-    # via
-    #   importlib-metadata
-    #   importlib-resources

docs/tutorial_hello_world.rst

@@ -6,15 +6,26 @@ CK Hello world
 Motivation
 -------------------------------------
 
-This tutorial is aimed at engineers dealing with artificial intelligence and machine learning who would like to optimize their pipelines and squeeze every performance drop by adding Composable Kernel (CK) library to their projects. We would like to make the CK library approachable so the tutorial is not based on the latest release and doesn't have all the bleeding edge features, but it will be reproducible now and forever.
+This tutorial is aimed at engineers dealing with artificial intelligence and machine learning who
+would like to optimize their pipelines and squeeze every performance drop by adding Composable
+Kernel (CK) library to their projects. We would like to make the CK library approachable so
+the tutorial is not based on the latest release and doesn't have all the bleeding edge features,
+but it will be reproducible now and forever.
 
-During this tutorial we will have an introduction to the CK library, we will build it and run some examples and tests, so to say we will run a "Hello world" example. In future tutorials we will go in depth and breadth and get familiar with other tools and ways to integrate CK into your project.
+During this tutorial we will have an introduction to the CK library, we will build it and run some
+examples and tests, so to say we will run a "Hello world" example. In future tutorials we will go
+in depth and breadth and get familiar with other tools and ways to integrate CK into your project.
 
 -------------------------------------
 Description
 -------------------------------------
 
-Modern AI technology solves more and more problems in all imaginable fields, but crafting fast and efficient workflows is still challenging. CK is one of the tools to make AI heavy lifting as fast and efficient as possible. CK is a collection of optimized AI operator kernels and tools to create new ones. The library has components required for majority of modern neural networks architectures including matrix multiplication, convolution, contraction, reduction, attention modules, variety of activation functions, fused operators and many more.
+Modern AI technology solves more and more problems in all imaginable fields, but crafting fast and
+efficient workflows is still challenging. CK is one of the tools to make AI heavy lifting as fast
+and efficient as possible. CK is a collection of optimized AI operator kernels and tools to create
+new ones. The library has components required for majority of modern neural networks architectures
+including matrix multiplication, convolution, contraction, reduction, attention modules, variety of
+activation functions, fused operators and many more.
 
 So how do we (almost) reach the speed of light? CK acceleration abilities are based on:
 
@@ -24,15 +35,18 @@ So how do we (almost) reach the speed of light? CK acceleration abilities are ba
 * Hardware acceleration use.
 * Support of low precision data types including fp16, bf16, int8 and int4.
 
-If you are excited and need more technical details and benchmarking results - read this awesome `blog post <https://community.amd.com/t5/instinct-accelerators/amd-composable-kernel-library-efficient-fused-kernels-for-ai/ba-p/553224>`_.
+If you are excited and need more technical details and benchmarking results - read this awesome
+`blog post <https://community.amd.com/t5/instinct-accelerators/amd-composable-kernel-library-efficient-fused-kernels-for-ai/ba-p/553224>`_.
 
-For more details visit our `github repo <https://github.com/ROCmSoftwarePlatform/composable_kernel>`_.
+For more details visit our `github repository <https://github.com/ROCmSoftwarePlatform/composable_kernel>`_.
 
 -------------------------------------
 Hardware targets
 -------------------------------------
 
-CK library fully supports "gfx908" and "gfx90a" GPU architectures and only some operators are supported for "gfx1030". Let's check the hardware you have at hand and decide on the target GPU architecture
+CK library fully supports `gfx908` and `gfx90a` GPU architectures and only some operators are
+supported for `gfx1030`. Let's check the hardware you have at hand and decide on the target
+GPU architecture.
 
 ==========     =========
 GPU Target     AMD GPU
@@ -42,7 +56,8 @@ gfx90a 	       Radeon Instinct MI210, MI250, MI250X
 gfx1030        Radeon PRO V620, W6800, W6800X, W6800X Duo, W6900X, RX 6800, RX 6800 XT, RX 6900 XT, RX 6900 XTX, RX 6950 XT
 ==========     =========
 
-There are also `cloud options <https://aws.amazon.com/ec2/instance-types/g4/>`_ you can find if you don't have an AMD GPU at hand.
+There are also `cloud options <https://aws.amazon.com/ec2/instance-types/g4/>`_ you can find if
+you don't have an AMD GPU at hand.
 
 -------------------------------------
 Build the library
@@ -54,9 +69,13 @@ First let's clone the library and rebase to the tested version::
     cd composable_kernel/
     git checkout tutorial_hello_world
 
-To make our lives easier we prepared `docker images <https://hub.docker.com/r/rocm/composable_kernel>`_ with all the necessary dependencies. Pick the right image and create a container. In this tutorial we use "rocm/composable_kernel:ck_ub20.04_rocm5.3_release" image, it is based on Ubuntu 20.04, ROCm v5.3, compiler release version.
+To make our lives easier we prepared
+`docker images <https://hub.docker.com/r/rocm/composable_kernel>`_ with all the necessary
+dependencies. Pick the right image and create a container. In this tutorial we use
+``rocm/composable_kernel:ck_ub20.04_rocm5.6`` image, it is based on Ubuntu 20.04 and
+ROCm v5.6.
 
-If your current folder is ${HOME}, start the docker container with::
+If your current folder is ``${HOME}``, start the docker container with::
 
     docker run  \
     -it  \
@@ -64,20 +83,23 @@ If your current folder is ${HOME}, start the docker container with::
     --group-add sudo  \
     -w /root/workspace  \
     -v ${HOME}:/root/workspace  \
-    rocm/composable_kernel:ck_ub20.04_rocm5.3_release  \
+    rocm/composable_kernel:ck_ub20.04_rocm5.6  \
     /bin/bash
 
-If your current folder is different from ${HOME}, adjust the line `-v ${HOME}:/root/workspace` to fit your folder structure.
+If your current folder is different from ``${HOME}``, adjust the line ``-v ${HOME}:/root/workspace``
+to fit your folder structure.
 
-Inside the docker container current folder is "~/workspace", library path is "~/workspace/composable_kernel", navigate to the library::
+Inside the docker container current folder is ``~/workspace``, library path is
+``~/workspace/composable_kernel``, navigate to the library::
 
     cd composable_kernel/
 
-Create and go to the "build" directory::
+Create and go to the ``build`` directory::
 
     mkdir build && cd build
 
-In the previous section we talked about target GPU architecture. Once you decide which one is right for you, run cmake using the right GPU_TARGETS flag::
+In the previous section we talked about target GPU architecture. Once you decide which one is right
+for you, run CMake using the right ``GPU_TARGETS`` flag::
 
     cmake  \
     -D CMAKE_PREFIX_PATH=/opt/rocm  \
@@ -87,7 +109,7 @@ In the previous section we talked about target GPU architecture. Once you decide
     -D BUILD_DEV=OFF  \
     -D GPU_TARGETS="gfx908;gfx90a;gfx1030" ..
 
-If everything went well the cmake run will end up with::
+If everything went well the CMake run will end up with::
 
     -- Configuring done
     -- Generating done
@@ -118,9 +140,12 @@ We can also run them separately, here is a separate example execution::
 
     ./bin/example_gemm_xdl_fp16 1 1 1
 
-The arguments "1 1 1" mean that we want to run this example in the mode: verify results with CPU, initialize matrices with integers and benchmark the kernel execution. You can play around with these parameters and see how output and execution results change.
+The arguments ``1 1 1`` mean that we want to run this example in the mode: verify results with CPU,
+initialize matrices with integers and benchmark the kernel execution. You can play around with
+these parameters and see how output and execution results change.
 
-If everything goes well and you have a device based on gfx908 or gfx90a architecture you should see something like::
+If everything goes well and you have a device based on `gfx908` or `gfx90a` architecture you should see
+something like::
 
     a_m_k: dim 2, lengths {3840, 4096}, strides {4096, 1}
     b_k_n: dim 2, lengths {4096, 4096}, strides {1, 4096}
@@ -130,14 +155,15 @@ If everything goes well and you have a device based on gfx908 or gfx90a architec
     Start running 10 times...
     Perf: 1.10017 ms, 117.117 TFlops, 87.6854 GB/s, DeviceGemmXdl<256, 256, 128, 4, 8, 32, 32, 4, 2> NumPrefetch: 1, LoopScheduler: Default, PipelineVersion: v1
 
-Meanwhile, running it on a gfx1030 device should result in::
+Meanwhile, running it on a `gfx1030` device should result in::
 
     a_m_k: dim 2, lengths {3840, 4096}, strides {4096, 1}
     b_k_n: dim 2, lengths {4096, 4096}, strides {1, 4096}
     c_m_n: dim 2, lengths {3840, 4096}, strides {4096, 1}
     DeviceGemmXdl<256, 256, 128, 4, 8, 32, 32, 4, 2> NumPrefetch: 1, LoopScheduler: Default, PipelineVersion: v1 does not support this problem
 
-But don't panic, some of the operators are supported on gfx1030 architecture, so you can run a separate example like::
+But don't panic, some of the operators are supported on `gfx1030` architecture, so you can run a
+separate example like::
 
     ./bin/example_gemm_dl_fp16 1 1 1
 
@@ -154,7 +180,14 @@ and it should result in something nice similar to::
     Start running 10 times...
     Perf: 3.65695 ms, 35.234 TFlops, 26.3797 GB/s, DeviceGemmDl<256, 128, 128, 16, 2, 4, 4, 1>
 
-Or we can run a separate test::
+.. note::
+
+    There was a new CMake flag ``DL_KERNELS`` added in the latest versions of CK. If you use one of
+    the newest versions of the library and do not see the above results when running
+    ``example_gemm_dl_fp16``, it might be necessary to add ``-D DL_KERNELS=ON`` to your CMake command
+    in order to build the operators supported on the `gfx1030` architecture.
+
+We can also run a separate test::
 
     ctest -R test_gemm_fp16
 
@@ -169,6 +202,9 @@ If everything goes well you should see something like::
 Summary
 -----------
 
-In this tutorial we took the first look at the Composable Kernel library, built it on your system and ran some examples and tests. Stay tuned, in the next tutorial we will run kernels with different configs to find out the best one for your hardware and task.
+In this tutorial we took the first look at the Composable Kernel library, built it on your system
+and ran some examples and tests. Stay tuned, in the next tutorial we will run kernels with different
+configs to find out the best one for your hardware and task.
 
-P.S.: Don't forget to switch out the cloud instance if you have launched one, you can find better ways to spend your money for sure!
+P.S.: Don't forget to switch off the cloud instance if you have launched one, you can find better
+ways to spend your money for sure!

example/01_gemm/CMakeLists.txt

-add_custom_target(example_gemm_dl)
+if(DL_KERNELS)
+  add_custom_target(example_gemm_dl)
 
-add_example_executable(example_gemm_dl_fp32 gemm_dl_fp32.cpp)
-add_example_executable(example_gemm_dl_fp16 gemm_dl_fp16.cpp)
-
-add_dependencies(example_gemm_dl example_gemm_dl_fp32)
-add_dependencies(example_gemm_dl example_gemm_dl_fp16)
-
-if(DTYPES MATCHES "int8" OR NOT DEFINED DTYPES)
+  add_example_executable(example_gemm_dl_fp32 gemm_dl_fp32.cpp)
+  add_dependencies(example_gemm_dl example_gemm_dl_fp32)
+  if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
+    add_example_executable(example_gemm_dl_fp16 gemm_dl_fp16.cpp)
+    add_dependencies(example_gemm_dl example_gemm_dl_fp16)
+    add_example_executable(example_gemm_dpp_fp16 gemm_dpp_fp16.cpp)
+  endif()
+  if(DTYPES MATCHES "int8" OR NOT DEFINED DTYPES)
     add_example_executable(example_gemm_dl_int8 gemm_dl_int8.cpp)
     add_dependencies(example_gemm_dl example_gemm_dl_int8)
-endif()
+  endif()
 
-if(USE_BITINT_EXTENSION_INT4)
+  if(USE_BITINT_EXTENSION_INT4)
     add_example_executable(example_gemm_dl_int4 gemm_dl_int4.cpp)
     add_dependencies(example_gemm_dl example_gemm_dl_int4)
-endif(USE_BITINT_EXTENSION_INT4)
-
+  endif(USE_BITINT_EXTENSION_INT4)
+endif()
 
 add_custom_target(example_gemm_xdl)
+if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
+  add_example_executable(example_gemm_xdl_fp16 gemm_xdl_fp16.cpp)
+  add_example_executable(example_gemm_xdl_wavelet_fp16 gemm_xdl_wavelet_fp16.cpp)
+  add_dependencies(example_gemm_xdl example_gemm_xdl_fp16)
+  add_dependencies(example_gemm_xdl example_gemm_xdl_wavelet_fp16)
+  add_example_executable(example_gemm_xdl_skip_b_lds_fp16 gemm_xdl_skip_b_lds_fp16.cpp)
+  add_dependencies(example_gemm_xdl example_gemm_xdl_skip_b_lds_fp16)
+
+  if(GPU_TARGETS MATCHES "gfx1100" OR GPU_TARGETS MATCHES "gfx1101" OR GPU_TARGETS MATCHES "gfx1102")
+    add_custom_target(example_gemm_wmma)
+    add_example_executable(example_gemm_wmma_fp16 gemm_wmma_fp16.cpp)
+    add_dependencies(example_gemm_wmma example_gemm_wmma_fp16)
+  endif()
+
+endif()
 
-add_example_executable(example_gemm_xdl_fp16 gemm_xdl_fp16.cpp)
-add_example_executable(example_gemm_xdl_wavelet_fp16 gemm_xdl_wavelet_fp16.cpp)
-add_example_executable(example_gemm_xdl_bf16 gemm_xdl_bf16.cpp)
+if(DTYPES MATCHES "bf16" OR NOT DEFINED DTYPES)
+  add_example_executable(example_gemm_xdl_bf16 gemm_xdl_bf16.cpp)
+  add_dependencies(example_gemm_xdl example_gemm_xdl_bf16)
 
-add_dependencies(example_gemm_xdl example_gemm_xdl_fp16)
-add_dependencies(example_gemm_xdl example_gemm_xdl_bf16)
-add_dependencies(example_gemm_xdl example_gemm_xdl_wavelet_fp16)
+  add_example_executable(example_gemm_xdl_bf16_rtn gemm_xdl_bf16_rtn.cpp)
+  add_dependencies(example_gemm_xdl example_gemm_xdl_bf16_rtn)
+endif()
 
 if(DTYPES MATCHES "int8" OR NOT DEFINED DTYPES)
   add_example_executable(example_gemm_xdl_int8 gemm_xdl_int8.cpp)
@@ -37,20 +54,20 @@ if(USE_BITINT_EXTENSION_INT4)
   add_dependencies(example_gemm_xdl example_gemm_xdl_int4)
 endif(USE_BITINT_EXTENSION_INT4)
 
-add_example_executable(example_gemm_xdl_skip_b_lds_fp16 gemm_xdl_skip_b_lds_fp16.cpp)
-# FIXME: re-enable this exampe as test when SWDEV-335738 is fixed
-add_example_executable_no_testing(example_gemm_xdl_fp64 gemm_xdl_fp64.cpp)
-
-add_dependencies(example_gemm_xdl example_gemm_xdl_skip_b_lds_fp16)
-add_dependencies(example_gemm_xdl example_gemm_xdl_fp64)
-
-if(GPU_TARGETS MATCHES "gfx1100" OR GPU_TARGETS MATCHES "gfx1101" OR GPU_TARGETS MATCHES "gfx1102")
-  add_custom_target(example_gemm_wmma)
-  add_example_executable(example_gemm_wmma_fp16 gemm_wmma_fp16.cpp)
-  add_dependencies(example_gemm_wmma example_gemm_wmma_fp16)
+if(DTYPES MATCHES "fp64" OR NOT DEFINED DTYPES)
+  # FIXME: re-enable this exampe as test when SWDEV-335738 is fixed
+  add_example_executable_no_testing(example_gemm_xdl_fp64 gemm_xdl_fp64.cpp)
+  add_dependencies(example_gemm_xdl example_gemm_xdl_fp64)
 endif()
 
-if(GPU_TARGETS MATCHES "gfx940" OR GPU_TARGETS MATCHES "gfx941" OR GPU_TARGETS MATCHES "gfx942")
+add_example_executable(example_gemm_xdl_streamk gemm_xdl_streamk.cpp)
+
+if(DTYPES MATCHES "fp8" OR NOT DEFINED DTYPES)
+  if(GPU_TARGETS MATCHES "gfx940" OR GPU_TARGETS MATCHES "gfx941" OR GPU_TARGETS MATCHES "gfx942")
     add_example_executable(example_gemm_xdl_f8 gemm_xdl_f8.cpp)
     add_dependencies(example_gemm_xdl example_gemm_xdl_f8)
+  endif()
 endif()
+
+add_example_executable(example_gemm_xdl_fp16_f8 gemm_xdl_fp16_f8.cpp)
+add_dependencies(example_gemm_xdl example_gemm_xdl_fp16_f8)

example/01_gemm/common.hpp

@@ -33,6 +33,19 @@ struct ProblemSize final
     ck::index_t StrideC = 4096;
 };
 
+struct ProblemSizeStreamK final
+{
+    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 StrideC = 4096;
+
+    ck::index_t NumSKBlocks = -1;
+};
+
 struct ExecutionConfig final
 {
     bool do_verification = true;
@@ -48,8 +61,17 @@ using Col = ck::tensor_layout::gemm::ColumnMajor;
 
 using PassThrough = ck::tensor_operation::element_wise::PassThrough;
 
-inline bool
-parse_cmd_args(int argc, char* argv[], ProblemSize& problem_size, ExecutionConfig& config)
+template <typename ProblemType>
+bool parse_cmd_args(int, char*[], ProblemType&, ExecutionConfig&)
+{
+    return false;
+}
+
+template <>
+bool parse_cmd_args<ProblemSize>(int argc,
+                                 char* argv[],
+                                 ProblemSize& problem_size,
+                                 ExecutionConfig& config)
 {
     if(argc == 1)
     {
@@ -87,3 +109,52 @@ parse_cmd_args(int argc, char* argv[], ProblemSize& problem_size, ExecutionConfi
 
     return true;
 }
+
+template <>
+bool parse_cmd_args<ProblemSizeStreamK>(int argc,
+                                        char* argv[],
+                                        ProblemSizeStreamK& problem_size,
+                                        ExecutionConfig& config)
+{
+    if(argc == 1)
+    {
+        // use default case
+    }
+    else if(argc == 4)
+    {
+        config.do_verification = std::stoi(argv[1]);
+        config.init_method     = std::stoi(argv[2]);
+        config.time_kernel     = std::stoi(argv[3]);
+    }
+    else if(argc >= 10)
+    {
+        config.do_verification = std::stoi(argv[1]);
+        config.init_method     = std::stoi(argv[2]);
+        config.time_kernel     = std::stoi(argv[3]);
+
+        problem_size.M = std::stoi(argv[4]);
+        problem_size.N = std::stoi(argv[5]);
+        problem_size.K = std::stoi(argv[6]);
+
+        problem_size.StrideA = std::stoi(argv[7]);
+        problem_size.StrideB = std::stoi(argv[8]);
+        problem_size.StrideC = std::stoi(argv[9]);
+
+        if(argc >= 11)
+        {
+            problem_size.NumSKBlocks = std::stoi(argv[10]);
+        }
+    }
+    else
+    {
+        std::cerr << "arg1: verification (0=no, 1=yes)" << std::endl
+                  << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)"
+                  << std::endl
+                  << "arg3: time kernel (0=no, 1=yes)" << std::endl
+                  << "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC" << std::endl
+                  << "arg10: NumSKBlocks(optional)" << std::endl;
+        return false;
+    }
+
+    return true;
+}

example/01_gemm/gemm_dpp_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "common.hpp"
+
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_dpp.hpp"
+
+using ADataType   = ck::half_t;
+using BDataType   = ck::half_t;
+using AccDataType = float;
+using CDataType   = ck::half_t;
+
+using F16 = ck::half_t;
+
+using ALayout = Row;
+using BLayout = Col;
+using CLayout = Row;
+
+using AElementOp = PassThrough;
+using BElementOp = PassThrough;
+using CElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::MNPadding;
+
+// clang-format off
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmDpp
+// ######|     AData|     BData|     CData|     AccData| ALayout| BLayout| CLayout|           A|           B|           C|          GEMM| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer|    MDpp|    NDpp|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CThreadTransfer| CThreadTransfer|
+// ######|      Type|      Type|      Type|        Type|        |        |        | Elementwise| Elementwise| Elementwise|Spacialization|  Size| Block| Block| Block|    |    |  Dpp|  Dpp| PerWave| PerWave|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| SrcDstVectorDim|       DstScalar|
+// ######|          |          |          |            |        |        |        |   Operation|   Operation|   Operation|              |      |      |      |      |    |    |     |     |        |        | Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |                |       PerVector|
+// ######|          |          |          |            |        |        |        |            |            |            |              |      |      |      |      |    |    |     |     |        |        |                |               |               |               |               |               |          |                |               |               |              |               |               |          |                |                |
+         < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout,  AElementOp,  BElementOp,  CElementOp,   GemmDefault,   128,    64,    64,    64,   8,   2,   32,    8,       2,       2,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              2,              2,      true,               5,               1>;
+// // clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::
+    ReferenceGemm<ADataType, BDataType, CDataType, AccDataType, AElementOp, BElementOp, CElementOp>;
+
+#include "run_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_gemm_example(argc, argv); }