From d184438d8413c10894c0f77ddcf7a28fbade377f Mon Sep 17 00:00:00 2001 From: coderfeli Date: Tue, 22 Jul 2025 14:55:41 +0000 Subject: [PATCH] fix a dist --- example/ck_tile/18_flatmm/flatmm_basic.cpp | 66 ++++++++-------- .../ck_tile/18_flatmm/run_flatmm_example.inc | 9 ++- include/ck_tile/core/tensor/tile_window.hpp | 22 ++++++ .../flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 75 ++++++++----------- 4 files changed, 93 insertions(+), 79 deletions(-) diff --git a/example/ck_tile/18_flatmm/flatmm_basic.cpp b/example/ck_tile/18_flatmm/flatmm_basic.cpp index 858d275bc9..8929c170e3 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.cpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.cpp @@ -138,43 +138,43 @@ float flatmm_calc(const ck_tile::FlatmmHostArgs<>& args, const ck_tile::stream_c << std::endl; } - // if(s.flush_cache_) - // { - // std::cout << "Flushing cache..." << std::endl; - // static constexpr ck_tile::index_t APackedSize = - // std::is_same_v ? 2 : 1; - // static constexpr ck_tile::index_t BPackedSize = - // std::is_same_v ? 2 : 1; + if(s.flush_cache_) + { + std::cout << "Flushing cache..." << std::endl; + static constexpr ck_tile::index_t APackedSize = + std::is_same_v ? 2 : 1; + static constexpr ck_tile::index_t BPackedSize = + std::is_same_v ? 2 : 1; - // ck_tile::HostTensor a_m(ck_tile::host_tensor_descriptor( - // args.M, args.K, args.stride_A, is_row_major(ALayout{}))); - // ck_tile::HostTensor b_n(ck_tile::host_tensor_descriptor( - // args.K, args.N, args.stride_B, is_row_major(BLayout{}))); + ck_tile::HostTensor a_m(ck_tile::host_tensor_descriptor( + args.M, args.K, args.stride_A, is_row_major(ALayout{}))); + ck_tile::HostTensor b_n(ck_tile::host_tensor_descriptor( + args.K, args.N, args.stride_B, is_row_major(BLayout{}))); - // auto size_a_buffer = a_m.get_element_space_size_in_bytes() / APackedSize; - // auto size_b_buffer = b_n.get_element_space_size_in_bytes() / BPackedSize; + auto size_a_buffer = a_m.get_element_space_size_in_bytes() / APackedSize; + auto size_b_buffer = b_n.get_element_space_size_in_bytes() / BPackedSize; - // ck_tile::RotatingMemWrapper rotating_mem( - // kargs.a_ptr, kargs.b_ptr, s.rotating_count_, size_a_buffer, size_b_buffer); - // rotating_mem.Print(); + ck_tile::RotatingMemWrapper rotating_mem( + kargs.a_ptr, kargs.b_ptr, s.rotating_count_, size_a_buffer, size_b_buffer); + rotating_mem.Print(); - // auto run_flush_cache = [&]() { - // // flush icache - // ck_tile::flush_icache(); - // // rotating mem - // rotating_mem.Next(); - // // clear c mem - // if(args.k_batch > 1) - // hipGetErrorString(hipMemsetAsync( - // args.e_ptr, 0, args.M * args.N * sizeof(CDataType), s.stream_id_)); - // }; - // ave_time = ck_tile::launch_kernel_preprocess( - // s, - // run_flush_cache, - // ck_tile::make_kernel( - // Kernel{}, grids, blocks, 0, kargs)); - // } - // else + auto run_flush_cache = [&]() { + // flush icache + ck_tile::flush_icache(); + // rotating mem + rotating_mem.Next(); + // clear c mem + if(args.k_batch > 1) + hipGetErrorString(hipMemsetAsync( + args.e_ptr, 0, args.M * args.N * sizeof(CDataType), s.stream_id_)); + }; + ave_time = ck_tile::launch_kernel_preprocess( + s, + run_flush_cache, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + } + else { ave_time = ck_tile::launch_kernel(s, diff --git a/example/ck_tile/18_flatmm/run_flatmm_example.inc b/example/ck_tile/18_flatmm/run_flatmm_example.inc index 65b54b1ed3..a3ced34241 100644 --- a/example/ck_tile/18_flatmm/run_flatmm_example.inc +++ b/example/ck_tile/18_flatmm/run_flatmm_example.inc @@ -215,15 +215,16 @@ int run_flatmm_example_with_layouts(int argc, } else if(init_method == 3) { - ck_tile::FillUniformDistribution{1.f, 1.f}(a_host); - ck_tile::FillUniformDistribution{0.f, .5f}(b_origin_host); + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_host); + ck_tile::FillUniformDistribution{1.f, 1.f}(b_origin_host); + // ck_tile::FillUniformDistribution{-.5f, .5f}(b_origin_host); } else { a_host.SetZero(); b_origin_host.SetZero(); } - + b_origin_host.savetxt("b_origin_host.txt"); ck_tile::DeviceMem a_dev_buf(a_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem c_dev_buf(c_rslt_host.get_element_space_size_in_bytes()); @@ -278,6 +279,8 @@ int run_flatmm_example_with_layouts(int argc, *std::max_element(c_ref_host.mData.begin(), c_ref_host.mData.end()); const auto rtol_atol = calculate_rtol_atol( K, kbatch, max_accumulated_value); + c_rslt_host.savetxt("out.txt"); + c_ref_host.savetxt("ref.txt"); pass = ck_tile::check_err(c_rslt_host, c_ref_host, "Error: Incorrect results!", diff --git a/include/ck_tile/core/tensor/tile_window.hpp b/include/ck_tile/core/tensor/tile_window.hpp index ad5902f16e..bc155614dd 100644 --- a/include/ck_tile/core/tensor/tile_window.hpp +++ b/include/ck_tile/core/tensor/tile_window.hpp @@ -816,6 +816,28 @@ make_tile_window(const TensorView_& tensor_view, tensor_view, window_lengths, origin, tile_distribution}; } + + +// TODO: use strategy +// template +// CK_TILE_DEVICE constexpr auto +// make_tile_window_debug(const TensorView_& tensor_view, +// const WindowLengths_& window_lengths, +// const multi_index& origin, +// const StaticTileDistribution_& tile_distribution, +// number = {}) +// { +// return tile_window_with_static_distribution, +// remove_cvref_t, +// remove_cvref_t, +// NumCoord, +// 1>{ +// tensor_view, window_lengths, origin, tile_distribution}; +// } + // this version can't be called in a constexpr context template , + tile_distribution_encoding, tuple, sequence>, - tuple>, - tuple>, + tuple>, + tuple>, sequence<2>, sequence<0>>{}; return detail::make_embed_tile_distribution_encoding( @@ -619,13 +619,13 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV // ping-pong window for A LDS auto a_warp_window_ping_tmp = make_tile_window( a_lds_block_ping, - make_tuple(number{}, number{}), + make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, make_static_tile_distribution(MakeAWarpKIterDistributionEncode())); auto a_warp_window_pong_tmp = make_tile_window( a_lds_block_pong, - make_tuple(number{}, number{}), + make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, make_static_tile_distribution(MakeAWarpKIterDistributionEncode())); @@ -637,7 +637,7 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV move_tile_window(a_warp_windows_ping(mIter), {mIter * MPerBlockPerIter, 0}); a_warp_windows_pong(mIter) = a_warp_window_pong_tmp; - // move_tile_window(a_warp_windows_pong(mIter), {mIter * MPerBlockPerIter, 0}); + move_tile_window(a_warp_windows_pong(mIter), {mIter * MPerBlockPerIter, 0}); }); // Block GEMM auto block_flatmm = BlockFlatmm(); @@ -774,19 +774,18 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV CWarpTensor c_warp_tensor; AWarpTensor a_warp_tensor; - set_tile(a_warp_tensor, 1.0f); - // a_warp_tensor.get_thread_buffer() = a_warp_tensor_ping(AwarpIter).get_y_sliced_thread_data( - // merge_sequences(sequence{}, a_warp_y_index_zeros), - // merge_sequences(sequence<1>{}, a_warp_y_lengths)); + a_warp_tensor.get_thread_buffer() = a_warp_tensor_ping(AwarpIter).get_y_sliced_thread_data( + merge_sequences(sequence{}, a_warp_y_index_zeros), + merge_sequences(sequence<1>{}, a_warp_y_lengths)); c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - // if(threadIdx.x % 8 == 0){ + // if constexpr(mIter==0 && nIter ==0) + // if(threadIdx.x % 16== 0 && threadIdx.x<64){ // for(int i=0;i(a_warp_tensor.thread_buf_(i)), type_convert(b_warp_tensor_ping(nIter)(kIter).thread_buf_(i))); + // printf("tid=%u, i0 %d bval=%f\n", threadIdx.x, i, type_convert(b_warp_tensor_ping(nIter)(kIter).thread_buf_(i))); // } - // } // warp GEMM WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor_ping(nIter)(kIter)); @@ -796,7 +795,6 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), c_warp_tensor.get_thread_buffer()); - __builtin_amdgcn_sched_barrier(0x7F6); }); }); @@ -804,13 +802,6 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV if constexpr(mIter == MIterPerWarp - 2) { block_sync_lds(); - // if(threadIdx.x == 0) { - // for (int i = 0 ; i < a_lds_block_desc.get_element_space_size(); i++) { - // printf("%f,", type_convert(p_a_lds_pong[i])); - // } - // printf("\n111111111111111111\n"); - // } - // block_sync_lds(); a_warp_tensor_ping(AwarpIter) = load_tile(a_warp_windows_pong(I0)); } else if constexpr (mIter == MIterPerWarp - 1) { a_warp_tensor_ping(AwarpIter) = load_tile(a_warp_windows_pong(I1)); @@ -862,10 +853,9 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV AWarpTensor a_warp_tensor; - set_tile(a_warp_tensor, 1.0f); - // a_warp_tensor.get_thread_buffer() = a_warp_tensor_ping(AwarpIter).get_y_sliced_thread_data( - // merge_sequences(sequence{}, a_warp_y_index_zeros), - // merge_sequences(sequence<1>{}, a_warp_y_lengths)); + a_warp_tensor.get_thread_buffer() = a_warp_tensor_ping(AwarpIter).get_y_sliced_thread_data( + merge_sequences(sequence{}, a_warp_y_index_zeros), + merge_sequences(sequence<1>{}, a_warp_y_lengths)); // warp GEMM // if(threadIdx.x % 16 == 0 && threadIdx.x>=192){ // for(int i=0;i(a_warp_tensor_ping(AwarpIter).thread_buf_[i])); // } // } - WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor_ping(nIter)(kIter)); + // if constexpr(mIter==0 && nIter ==0) + // if(threadIdx.x % 16== 0 && threadIdx.x<64){ + // for(int i=0;i(b_warp_tensor_pong(nIter)(kIter).thread_buf_(i))); + + // } + // } + WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor_pong(nIter)(kIter)); // write C warp tensor into C block tensor c_block_tile.set_y_sliced_thread_data( merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), c_warp_tensor.get_thread_buffer()); - - __builtin_amdgcn_sched_barrier(0x7F6); }); }); @@ -1026,24 +1021,14 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV AWarpTensor a_warp_tensor; - // a_warp_tensor.get_thread_buffer() = a_warp_tensor_ping(AwarpIter).get_y_sliced_thread_data( - // merge_sequences(sequence{}, a_warp_y_index_zeros), - // merge_sequences(sequence<1>{}, a_warp_y_lengths)); + a_warp_tensor.get_thread_buffer() = a_warp_tensor_ping(AwarpIter).get_y_sliced_thread_data( + merge_sequences(sequence{}, a_warp_y_index_zeros), + merge_sequences(sequence<1>{}, a_warp_y_lengths)); - set_tile(a_warp_tensor, 1.0f); + // set_tile(a_warp_tensor, 1.0f); c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - - // if(threadIdx.x >= 128 && threadIdx.x % 1 == 0){ - // for(int i=0;i(a_warp_tensor.thread_buf_(i)), type_convert(b_warp_tensor_ping(nIter)(kIter).thread_buf_(i))); - // } - // for(int i=0;i(a_warp_tensor_ping(AwarpIter).thread_buf_[i])); - // } - // } - // warp GEMM WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor_ping(nIter)(kIter)); // write C warp tensor into C block tensor @@ -1051,9 +1036,13 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), c_warp_tensor.get_thread_buffer()); - __builtin_amdgcn_sched_barrier(0x7F6); + // __builtin_amdgcn_sched_barrier(0x7F6); }); }); + if constexpr(mIter < MIterPerWarp - 2) + { + a_warp_tensor_ping(AwarpIter) = load_tile(a_warp_windows_ping(number{})); + } }); } // }