From 3c171550f602ea304d9c81acfa64b5742d654ea5 Mon Sep 17 00:00:00 2001 From: Aleksander Dudek Date: Tue, 29 Oct 2024 15:58:02 +0000 Subject: [PATCH] Batched gemm - messy validation check --- .../run_batched_gemm_example.inc | 21 ++++++-- .../ck_tile/host/reference/reference_gemm.hpp | 52 +++++++++++++------ .../ops/gemm/kernel/batched_gemm_kernel.hpp | 21 +++++--- 3 files changed, 66 insertions(+), 28 deletions(-) diff --git a/example/ck_tile/05_batched_gemm/run_batched_gemm_example.inc b/example/ck_tile/05_batched_gemm/run_batched_gemm_example.inc index 1270785693..2074895838 100644 --- a/example/ck_tile/05_batched_gemm/run_batched_gemm_example.inc +++ b/example/ck_tile/05_batched_gemm/run_batched_gemm_example.inc @@ -96,11 +96,13 @@ int run_batched_gemm_example(int argc, char* argv[]) [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { if constexpr(std::is_same_v) { - return ck_tile::HostTensorDescriptor({row, col}, {stride, 1_uz}); + return ck_tile::HostTensorDescriptor({static_cast(16), row, col}, + {row * col, stride, 1_uz}); } else { - return ck_tile::HostTensorDescriptor({row, col}, {1_uz, stride}); + return ck_tile::HostTensorDescriptor({static_cast(16), row, col}, + {row * col, 1_uz, stride}); } }; @@ -194,8 +196,19 @@ int run_batched_gemm_example(int argc, char* argv[]) CDataType, ALayout, BLayout, - CLayout>( - a_m_k_dev_buf, b_k_n_dev_buf, c_m_n_gpu_buf_ref, M, N, K, stride_A, stride_B, stride_C); + CLayout>(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_gpu_buf_ref, + M, + N, + K, + stride_A, + stride_B, + stride_C, + batch_stride_A, + batch_stride_B, + batch_stride_C, + batch_count); c_m_n_gpu_buf_ref.FromDevice(c_m_n_gpu_ref.data()); pass = ck_tile::check_err(c_m_n_dev_result, c_m_n_gpu_ref); diff --git a/include/ck_tile/host/reference/reference_gemm.hpp b/include/ck_tile/host/reference/reference_gemm.hpp index dbdef0e9c7..19ca1960ef 100644 --- a/include/ck_tile/host/reference/reference_gemm.hpp +++ b/include/ck_tile/host/reference/reference_gemm.hpp @@ -29,22 +29,22 @@ CK_TILE_HOST void reference_gemm(const HostTensor& a_m_k, const std::size_t N = b_k_n.get_length(1); const std::size_t K = a_m_k.get_length(1); - auto f_mn = [&](auto m, auto n) { + auto f_mn = [&](auto m, auto n, auto b) { AccDataType v_acc = 0; for(std::size_t k = 0; k < K; ++k) { - ADataType v_a = a_element_op(a_m_k(m, k)); - BDataType v_b = b_element_op(b_k_n(k, n)); + ADataType v_a = a_element_op(a_m_k(b, m, k)); + BDataType v_b = b_element_op(b_k_n(b, k, n)); v_acc += ck_tile::type_convert(v_a) * ck_tile::type_convert(v_b); } - c_m_n(m, n) = ck_tile::type_convert(acc_element_op(v_acc)); + c_m_n(b, m, n) = ck_tile::type_convert(acc_element_op(v_acc)); }; - make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); + make_ParallelTensorFunctor(f_mn, M, N, 16)(std::thread::hardware_concurrency()); } template - <<>>(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); + for(int i = 0; i < batch_count; ++i) + { + ADataType* d_ATemp = d_A + i * batch_stride_A; + BDataType* d_BTemp = d_B + i * batch_stride_B; + CDataType* d_CTemp = d_C + i * batch_stride_C; + naive_gemm_kernel + <<>>( + d_ATemp, d_BTemp, d_CTemp, M, N, K, stride_a, stride_b, stride_c); + } + + errC = hipMemcpy(c_device.GetDeviceBuffer(), + d_C, + batch_count * M * N * sizeof(CDataType), + hipMemcpyDeviceToHost); if(errC != hipSuccess) { std::cerr << "Error copying C to device: " << hipGetErrorString(errC) << std::endl; diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp index 15ed2a55d8..30f40979a6 100644 --- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp @@ -89,13 +89,20 @@ struct BatchedGemmKernel CK_TILE_DEVICE void operator()(BatchedGemmCommonKargs kargs) const { const auto [i_m, i_n] = TilePartitioner{}(); - // const auto i_k = blockIdx.z; + const auto i_k = blockIdx.z; // options - const ADataType* a_start = static_cast( - kargs.a_ptr); //+ __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_A); - const BDataType* b_start = static_cast( - kargs.b_ptr); //+ __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_B); + const ADataType* a_start = static_cast(kargs.a_ptr) + + __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_A); + const BDataType* b_start = static_cast(kargs.b_ptr) + + __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_B); // Convert pointers to tensor views + // if(threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0) + // { + // printf("__builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_A): %d\n", + // __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_A)); + // printf("__builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_B): %d\n", + // __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_B)); + // } auto a_tensor_view = [&]() { if constexpr(std::is_same_v) { @@ -172,8 +179,8 @@ struct BatchedGemmKernel auto c_block_tile = GemmPipeline{}.template operator()(a_block_window, b_block_window, num_loop, smem_ptr); - CDataType* c_start = static_cast( - kargs.c_ptr); //; + __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_C); + CDataType* c_start = static_cast(kargs.c_ptr) + + __builtin_amdgcn_readfirstlane(i_k * kargs.batch_stride_C); auto c_tensor_view = [&]() { if constexpr(std::is_same_v) {