Refactor unit-test. Fallback to non-universal gemm.

Need to use GemmPipelineAGmemBGmemCRegV1 for now,
since GemmKernel is now supporting also non-K major vector reads.

example/ck_tile/16_batched_gemm/batched_gemm.cpp

@@ -72,9 +72,7 @@ float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stre
         ck_tile::TileGemmTraits<kPadM, kPadN, kPadK, ALayout, BLayout, CLayout>;
     using CodegenPipelineProblem = ck_tile::
         GemmPipelineProblem<ADataType, BDataType, AccDataType, CodegenGemmShape, CodegenGemmTraits>;
-    using CodegenGemmPolicy = ck_tile::UniversalGemmPipelineAgBgCrPolicy;
-    using CodegenGemmPipeline =
-        ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem, CodegenGemmPolicy>;
+    using CodegenGemmPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
     // ToDo: Will add the codegen part to test different pipeline policies in GEMM.
     // Now we only use the BlockGemmASmemBSmemCRegV1DefaultPolicy.
     using Kernel = ck_tile::BatchedGemmKernel<TilePartitioner, CodegenGemmPipeline, GemmEpilogue>;

example/ck_tile/16_batched_gemm/batched_gemm.hpp

@@ -39,7 +39,7 @@ auto create_args(int argc, char* argv[])
         .insert("stride_b", "0", "Tensor B stride")
         .insert("stride_c", "0", "Tensor C stride")
         .insert("a_layout", "R", "A tensor data layout - Row by default")
-        .insert("b_layout", "R", "B tensor data layout - Row by default")
+        .insert("b_layout", "C", "B tensor data layout - Row by default")
         .insert("c_layout", "R", "C tensor data layout - Row by default")
         .insert("batch_stride_a", "32768", "Batch A stride")
         .insert("batch_stride_b", "16384", "Batch B stride")

example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc

@@ -3,13 +3,6 @@
 
 #pragma once
 
-template <typename Layout>
-static constexpr inline auto is_row_major(Layout layout_)
-{
-    return ck_tile::bool_constant<std::is_same_v<ck_tile::remove_cvref_t<decltype(layout_)>,
-                                                 ck_tile::tensor_layout::gemm::RowMajor>>{};
-}
-
 auto calculate_rtol_atol(const ck_tile::index_t K,
                          const ck_tile::index_t kbatch,
                          const float max_accumulated_value)
@@ -113,16 +106,56 @@ int run_batched_gemm_example_with_layouts(int argc,
     int n_warmup = arg_parser.get_int("warmup");
     int n_repeat = arg_parser.get_int("repeat");
 
-    stride_A = ck_tile::get_default_stride(M, K, stride_A, is_row_major(a_layout));
-    stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout));
-    stride_C = ck_tile::get_default_stride(M, N, stride_C, is_row_major(c_layout));
+    using namespace ck_tile::literals;
 
-    ck_tile::HostTensor<ADataType> a_m_k(ck_tile::host_tensor_descriptor(
-        batch_count, M, K, stride_A, batch_stride_A, is_row_major(a_layout)));
-    ck_tile::HostTensor<BDataType> b_k_n(ck_tile::host_tensor_descriptor(
-        batch_count, K, N, stride_B, batch_stride_B, is_row_major(b_layout)));
-    ck_tile::HostTensor<CDataType> c_m_n_dev_result(ck_tile::host_tensor_descriptor(
-        batch_count, M, N, stride_C, batch_stride_C, is_row_major(c_layout)));
+    auto f_host_tensor_descriptor = [](std::size_t batch_count_,
+                                       std::size_t row,
+                                       std::size_t col,
+                                       std::size_t stride,
+                                       std::size_t batch_stride,
+                                       auto layout) {
+        if constexpr(std::is_same_v<decltype(layout), ck_tile::tensor_layout::gemm::RowMajor>)
+        {
+            return ck_tile::HostTensorDescriptor({batch_count_, row, col},
+                                                 {batch_stride, stride, 1_uz});
+        }
+        else
+        {
+            return ck_tile::HostTensorDescriptor({batch_count_, row, col},
+                                                 {batch_stride, 1_uz, stride});
+        }
+    };
+
+    auto f_get_default_stride = [](std::size_t row,
+                                   std::size_t col,
+                                   std::size_t stride,
+                                   auto layout) {
+        if(stride == 0)
+        {
+            // give a chance if stride is zero, return a default packed stride
+            if constexpr(std::is_same_v<decltype(layout), ck_tile::tensor_layout::gemm::RowMajor>)
+            {
+                return col;
+            }
+            else
+            {
+                return row;
+            }
+        }
+        else
+            return stride;
+    };
+
+    stride_A = f_get_default_stride(M, K, stride_A, a_layout);
+    stride_B = f_get_default_stride(K, N, stride_B, b_layout);
+    stride_C = f_get_default_stride(M, N, stride_C, c_layout);
+
+    ck_tile::HostTensor<ADataType> a_m_k(
+        f_host_tensor_descriptor(batch_count, M, K, stride_A, batch_stride_A, a_layout));
+    ck_tile::HostTensor<BDataType> b_k_n(
+        f_host_tensor_descriptor(batch_count, K, N, stride_B, batch_stride_B, b_layout));
+    ck_tile::HostTensor<CDataType> c_m_n_dev_result(
+        f_host_tensor_descriptor(batch_count, M, N, stride_C, batch_stride_C, c_layout));
 
     ck_tile::FillUniformDistribution<ADataType>{-5.f, 5.f}(a_m_k);
     ck_tile::FillUniformDistribution<BDataType>{-5.f, 5.f}(b_k_n);
@@ -158,8 +191,8 @@ int run_batched_gemm_example_with_layouts(int argc,
 
     if(arg_parser.get_int("v") == 1)
     {
-        ck_tile::HostTensor<CDataType> c_m_n_host_ref(ck_tile::host_tensor_descriptor(
-            batch_count, M, N, stride_C, batch_stride_C, is_row_major(CLayout){}));
+        ck_tile::HostTensor<CDataType> c_m_n_host_ref(
+            f_host_tensor_descriptor(batch_count, M, N, stride_C, batch_stride_C, CLayout{}));
         c_m_n_host_ref.SetZero();
 
         const auto b_n_k = b_k_n.transpose({0, 2, 1});
@@ -183,8 +216,8 @@ int run_batched_gemm_example_with_layouts(int argc,
     }
     else if(arg_parser.get_int("v") == 2)
     {
-        ck_tile::HostTensor<CDataType> c_m_n_gpu_ref(ck_tile::host_tensor_descriptor(
-            batch_count, M, N, stride_C, batch_stride_C, is_row_major(CLayout){}));
+        ck_tile::HostTensor<CDataType> c_m_n_gpu_ref(
+            f_host_tensor_descriptor(batch_count, M, N, stride_C, batch_stride_C, CLayout{}));
         ck_tile::DeviceMem c_m_n_gpu_buf_ref(c_m_n_gpu_ref.get_element_space_size_in_bytes());
         c_m_n_gpu_ref.SetZero();
         c_m_n_gpu_buf_ref.SetZero();
@@ -268,11 +301,11 @@ int run_batched_gemm_example(int argc, char* argv[])
     std::string a_layout = arg_parser.get_str("a_layout");
     std::string b_layout = arg_parser.get_str("b_layout");
 
-    if(a_layout == "R" && b_layout == "R")
-    {
-        return run_batched_gemm_example_with_layouts(argc, argv, Row{}, Row{}, Row{});
-    }
-    else if(a_layout == "R" && b_layout == "C")
+    // if(a_layout == "R" && b_layout == "R")
+    // {
+    //     return run_batched_gemm_example_with_layouts(argc, argv, Row{}, Row{}, Row{});
+    // }
+    if(a_layout == "R" && b_layout == "C")
     {
         return run_batched_gemm_example_with_layouts(argc, argv, Row{}, Col{}, Row{});
     }

example/ck_tile/17_grouped_gemm/grouped_gemm.cpp

@@ -88,12 +88,9 @@ using CodegenPipelineProblem =
                                  CodegenGemmShape,
                                  CodegenGemmTraits<ALayout, BLayout, CLayout>>;
 
-using CodegenGemmPolicy = ck_tile::UniversalGemmPipelineAgBgCrPolicy;
-
 template <typename ALayout, typename BLayout, typename CLayout>
 using CodegenGemmPipeline =
-    ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem<ALayout, BLayout, CLayout>,
-                                          CodegenGemmPolicy>;
+    ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem<ALayout, BLayout, CLayout>>;
 
 template <typename ALayout, typename BLayout, typename CLayout>
 using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner,

example/ck_tile/17_grouped_gemm/grouped_gemm.hpp

@@ -41,7 +41,7 @@ auto create_args(int argc, char* argv[])
         .insert("stride_Bs", "", "Tensor B strides - it is empty by default.")
         .insert("stride_Cs", "", "Tensor C strides - it is empty by default.")
         .insert("a_layout", "R", "A tensor data layout - Row by default.")
-        .insert("b_layout", "R", "B tensor data layout - Row by default.")
+        .insert("b_layout", "C", "B tensor data layout - Row by default.")
         .insert("c_layout", "R", "C tensor data layout - Row by default.")
         .insert("validate", "1", "0. No validation, 1. Validation on CPU.")
         .insert("warmup", "10", "number of iterations before benchmark the kernel.")

example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc

@@ -135,12 +135,9 @@ int run_grouped_gemm_example_with_layouts(int argc,
         const ck_tile::index_t N = Ns[i];
         const ck_tile::index_t K = Ks[i];
 
-        stride_As[i] =
-            ck_tile::get_default_stride(M, N, stride_As[i], is_row_major(a_layout));
-        stride_Bs[i] =
-            ck_tile::get_default_stride(K, N, stride_Bs[i], is_row_major(b_layout));
-        stride_Cs[i] =
-            ck_tile::get_default_stride(M, N, stride_Cs[i], is_row_major(CLayout{}));
+        stride_As[i] = ck_tile::get_default_stride(M, N, stride_As[i], is_row_major(a_layout));
+        stride_Bs[i] = ck_tile::get_default_stride(K, N, stride_Bs[i], is_row_major(b_layout));
+        stride_Cs[i] = ck_tile::get_default_stride(M, N, stride_Cs[i], is_row_major(CLayout{}));
 
         a_m_k_tensors.push_back(ck_tile::HostTensor<ADataType>(
             ck_tile::host_tensor_descriptor(M, K, stride_As[i], is_row_major(a_layout))));
@@ -229,10 +226,10 @@ int run_grouped_gemm_example(int argc, char* argv[])
     {
         return run_grouped_gemm_example_with_layouts(argc, argv, Row{}, Col{}, Row{});
     }
-    else if(a_layout == "R" && b_layout == "R")
-    {
-        return run_grouped_gemm_example_with_layouts(argc, argv, Row{}, Row{}, Row{});
-    }
+    // else if(a_layout == "R" && b_layout == "R")
+    // {
+    //     return run_grouped_gemm_example_with_layouts(argc, argv, Row{}, Row{}, Row{});
+    // }
     else
     {
         throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!");

include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp

@@ -50,7 +50,6 @@ struct GroupedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
     using GemmKernelArgs          = typename Base::GemmKernelArgs;
 
     static constexpr index_t KernelBlockSize = GemmPipeline::BlockSize;
-    static constexpr index_t KBatch          = 1;
 
     struct GemmTransKernelArg
     {
@@ -124,7 +123,7 @@ struct GroupedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
                                        stride_a,
                                        stride_b,
                                        stride_c,
-                                       KBatch};
+                                       gemm_descs[i].k_batch};
 
             gemm_kernel_args_.emplace_back(std::move(karg), block_start, block_end);
         }

test/ck_tile/batched_gemm/test_batched_gemm.cpp

@@ -17,7 +17,7 @@ using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
 // clang-format off
 using KernelTypes = ::testing::Types<
     //         ALayout, BLayout, CLayout, ADataType, BDataType, AccDataType, CDataType
-    std::tuple<    Row,     Row,     Row,       F16,       F16,         F32,      F16>,
+    // std::tuple<    Row,     Row,     Row,       F16,       F16,         F32,      F16>,
     //std::tuple<    Col,     Row,     Row,       F16,       F16,         F32,      F16>,
     std::tuple<    Row,     Col,     Row,       F16,       F16,         F32,      F16>//,
     //std::tuple<    Col,     Col,     Row,       F16,       F16,         F32,      F16>

test/ck_tile/gemm/test_gemm_pipeline_util.hpp

@@ -16,6 +16,7 @@ enum struct GemmPipelineType
     Mem,
     Comp
 };
+
 template <typename Tuple>
 class TestCkTileGemmPipeline : public ::testing::Test
 {

test/ck_tile/grouped_gemm/test_grouped_gemm.cpp

@@ -17,7 +17,7 @@ using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
 // clang-format off
 using KernelTypes = ::testing::Types<
     //         ALayout, BLayout, CLayout, ADataType, BDataType, AccDataType, CDataType
-    std::tuple<    Row,     Row,     Row,       F16,       F16,         F32,      F16>,
+    // std::tuple<    Row,     Row,     Row,       F16,       F16,         F32,      F16>,
     //std::tuple<    Col,     Row,     Row,       F16,       F16,         F32,      F16>,
     std::tuple<    Row,     Col,     Row,       F16,       F16,         F32,      F16>//,
     //std::tuple<    Col,     Col,     Row,       F16,       F16,         F32,      F16>

test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp

@@ -96,12 +96,9 @@ class TestCkTileGroupedGemm : public ::testing::Test
                                      CodegenGemmShape,
                                      CodegenGemmTraits<ALayout, BLayout, CLayout>>;
 
-    using CodegenGemmPolicy = ck_tile::UniversalGemmPipelineAgBgCrPolicy;
-
     template <typename ALayout, typename BLayout, typename CLayout>
     using CodegenGemmPipeline =
-        ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem<ALayout, BLayout, CLayout>,
-                                              CodegenGemmPolicy>;
+        ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem<ALayout, BLayout, CLayout>>;
 
     template <typename ALayout, typename BLayout, typename CLayout>
     using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner,