[CK_TILE] Refactor GemmKernel - update tests

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

@@ -95,7 +95,7 @@ struct GemmKernel
         index_t stride_C;
     };
 
-    CK_TILE_HOST static constexpr GemmKernelArgs MakeKernelArgs(GemmHostArgs& hostArgs)
+    CK_TILE_HOST static constexpr GemmKernelArgs MakeKernelArgs(const GemmHostArgs& hostArgs)
     {
         return GemmKernelArgs{hostArgs.a_ptr,
                               hostArgs.b_ptr,

test/ck_tile/batched_gemm/test_batched_gemm_util.hpp

@@ -91,19 +91,7 @@ class TestCkTileBatchedGemm : public ::testing::Test
         using Kernel =
             ck_tile::BatchedGemmKernel<TilePartitioner, CodegenGemmPipeline, GemmEpilogue>;
 
-        auto kargs = Kernel::MakeKernelArgs(args.a_ptr,
-                                            args.b_ptr,
-                                            args.c_ptr,
-                                            args.M,
-                                            args.N,
-                                            args.K,
-                                            args.stride_A,
-                                            args.stride_B,
-                                            args.stride_C,
-                                            args.batch_stride_A,
-                                            args.batch_stride_B,
-                                            args.batch_stride_C,
-                                            args.batch_count);
+        auto kargs = Kernel::MakeKernelArgs(args);
 
         const dim3 grids      = Kernel::GridSize(args.M, args.N, args.batch_count);
         constexpr dim3 blocks = Kernel::BlockSize();

test/ck_tile/gemm/test_gemm_pipeline_util.hpp

@@ -31,22 +31,8 @@ class TestCkTileGemmPipeline : public ::testing::Test
     static constexpr auto PipelineType = std::tuple_element_t<8, Tuple>::value;
     // TODO: expose tile size through test t-param ?
 
-    struct gemm_args
-    {
-        const void* p_a;
-        const void* p_b;
-        void* p_c;
-        ck_tile::index_t kbatch;
-        ck_tile::index_t M;
-        ck_tile::index_t N;
-        ck_tile::index_t K;
-        ck_tile::index_t stride_A;
-        ck_tile::index_t stride_B;
-        ck_tile::index_t stride_C;
-    };
-
     template <bool PadM, bool PadN, bool PadK>
-    void invoke_gemm(const gemm_args& args, const ck_tile::stream_config& s)
+    void invoke_gemm(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s)
     {
         // TODO: This should be parameterized in tests
         constexpr ck_tile::index_t M_Tile = 128;
@@ -117,17 +103,9 @@ class TestCkTileGemmPipeline : public ::testing::Test
                                                                              has_hot_loop_v,
                                                                              tail_number_v>>>;
             using Kernel = ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
-            auto kargs   = Kernel::MakeKernelArgs(args.p_a,
-                                                args.p_b,
-                                                args.p_c,
-                                                args.M,
-                                                args.N,
-                                                args.K,
-                                                args.stride_A,
-                                                args.stride_B,
-                                                args.stride_C);
-
-            const dim3 grids      = Kernel::GridSize(args.M, args.N, args.kbatch);
+            auto kargs   = Kernel::MakeKernelArgs(args);
+
+            const dim3 grids      = Kernel::GridSize(args.M, args.N, args.k_batch);
             constexpr dim3 blocks = Kernel::BlockSize();
 
             if(!Kernel::IsSupportedArgument(kargs))
@@ -319,11 +297,11 @@ class TestCkTileGemmPipeline : public ::testing::Test
         c_m_n_dev_buf.SetZero();
         c_m_n_dev_result.SetZero();
 
-        gemm_args args;
-        args.p_a      = a_m_k_dev_buf.GetDeviceBuffer();
-        args.p_b      = b_k_n_dev_buf.GetDeviceBuffer();
-        args.p_c      = c_m_n_dev_buf.GetDeviceBuffer();
-        args.kbatch   = kbatch;
+        ck_tile::GemmHostArgs args;
+        args.a_ptr    = a_m_k_dev_buf.GetDeviceBuffer();
+        args.b_ptr    = b_k_n_dev_buf.GetDeviceBuffer();
+        args.c_ptr    = c_m_n_dev_buf.GetDeviceBuffer();
+        args.k_batch  = kbatch;
         args.M        = M;
         args.N        = N;
         args.K        = K;