From b50c8ee3500e6ec753fa1b13072a3e0381928ac1 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Tue, 9 Sep 2025 08:44:20 +0000 Subject: [PATCH 1/9] enable async_load on gfx950 for a16w4 flatmm pipeline --- ...ec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 113 +++++++++++------- 1 file changed, 68 insertions(+), 45 deletions(-) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index acd75c4a55..90f2e01bfb 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -119,7 +119,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 static constexpr index_t KPerBlockPerIter = kKPerBlock / KIterPerWarp; static constexpr int MXFP4PackedSize = 2; - static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType); + static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType); static constexpr index_t BK1 = Problem::VectorLoadSize / sizeof(BDataType) * MXFP4PackedSize; static constexpr index_t m_preload = (MIterPerWarp * KIterPerWarp >= DsReadPreload) ? DsReadPreload @@ -171,8 +171,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 kNPerBlock * kKPerBlock / NWarp / 32 / ScaleBload_K1 / WaveSize; // BlockN * BlockK / NWarp / ScalePerK / ScaleB_K1 / wavesize static constexpr index_t KPerScaleLoad = KIterPerWarp / ScaleBload_num; - static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2; - static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter; + static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2; + static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter; static constexpr index_t mfma_perM_perK = NIterPerWarp * mfma_per_wg; static constexpr index_t dswrite_mIter = (DsWritePreIssue - 1) % MIterPerWarp; @@ -396,7 +396,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } if((kIter % KPerScaleLoad == 0) && (mIter == 0)) { - load_perM = load_perM + 1; + load_perM = load_perM + 1; } SchedulerPerM(dsread_perM, dswrite_perM, load_perM); } @@ -639,9 +639,34 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 ScaleNPerWarp> scale_b_warp_tensor_pong; + using ABlockTile = decltype(load_tile(a_copy_dram_window)); + ABlockTile a_block_tile; + + auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a) { +#if defined(__gfx950__) + // global -> lds + async_load_tile(lds_tile_a, dram_tile_a); +#else + // global -> vgpr + a_block_tile = load_tile(dram_tile_a); +#endif + }; + auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { +#if defined(__gfx950__) + // data has been stored in lds, no need more operation. + static_assert(std::is_same_v, + "buffer_load_lds don't support element func fot A before mfma"); +#else + // vgpr -> lds + auto a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); + store_tile(lds_tile_a, a_block_tile_transformed); +#endif + }; + // HEAD // Prefetch A0 - auto a_block_tile = load_tile(a_copy_dram_window); + prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window); + // move A window to next k move_tile_window(a_copy_dram_window, {0, kKPerBlock}); @@ -678,15 +703,12 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); - // A_Lds_TileDist may differ with ADramTileDistribution - - auto a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window_ping, a_block_tile_transformed); + prefill_lds_a_stage2(a_copy_lds_window_ping); __builtin_amdgcn_sched_barrier(0); // Prefetch A1 - a_block_tile = load_tile(a_copy_dram_window); + prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window); // move A window to next k move_tile_window(a_copy_dram_window, {0, kKPerBlock}); @@ -820,11 +842,10 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 }); // Prefill A(2i+1) - a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window_pong, a_block_tile_transformed); + prefill_lds_a_stage2(a_copy_lds_window_pong); // Prefetch A(2i+2) - a_block_tile = load_tile(a_copy_dram_window); + prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window); // move A window to next k move_tile_window(a_copy_dram_window, {0, kKPerBlock}); @@ -901,8 +922,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 scale_b_flat_dram_window; move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), - {scale_n_iter * NFlatPerBlockPerIter, - scale_k_iter * ScaleKFlatPerWarp}); + {scale_n_iter * NFlatPerBlockPerIter, + scale_k_iter * ScaleKFlatPerWarp}); scale_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) = load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); @@ -915,19 +936,18 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 move_tile_window( b_flat_dram_windows(nIter)(kIter), {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + - packed_n_rank, - kIter * KFlatPerBlockPerIter}); + packed_n_rank, + kIter * KFlatPerBlockPerIter}); ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); b_warp_tensor_ping(nIter)(kIter) = ub.u; }); }); // Prefill A(2i+2) - a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window_ping, a_block_tile_transformed); + prefill_lds_a_stage2(a_copy_lds_window_ping); // Prefetch A(2i+3) - a_block_tile = load_tile(a_copy_dram_window); + prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window); // move A window to next k move_tile_window(a_copy_dram_window, {0, kKPerBlock}); @@ -996,7 +1016,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 if constexpr(TailNum == TailNumber::Even) { // prefetch B(loopK) - const int b_k_off = b_flat_dram_window.get_tile_distribution().calculate_index()[I1] / ContinuousKPerThread / WG::kN * ContinuousKPerThread; + const int b_k_off = b_flat_dram_window.get_tile_distribution().calculate_index()[I1] / + ContinuousKPerThread / WG::kN * ContinuousKPerThread; static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) @@ -1008,8 +1029,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 scale_b_flat_dram_window; move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter), - {scale_n_iter * NFlatPerBlockPerIter, - scale_k_iter * ScaleKFlatPerWarp}); + {scale_n_iter * NFlatPerBlockPerIter, + scale_k_iter * ScaleKFlatPerWarp}); scale_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) = load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); @@ -1017,7 +1038,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 }); const int b_k_off_inter = kIter * kKPerBlock / MXFP4KPerWarp + b_k_off; - if( b_k_off_inter < kKPerBlock - k_padded_zeros) { + if(b_k_off_inter < kKPerBlock - k_padded_zeros) + { static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { auto packed_n_idx = nIter / number{}; auto packed_n_rank = nIter % number{}; @@ -1027,18 +1049,17 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 move_tile_window( b_flat_dram_windows(nIter)(kIter), {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + - packed_n_rank, - kIter * KFlatPerBlockPerIter}); + packed_n_rank, + kIter * KFlatPerBlockPerIter}); - ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); + ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); b_warp_tensor_pong(nIter)(kIter) = ub.u; }); } }); // Prefill A(loopK) - a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window_pong, a_block_tile_transformed); + prefill_lds_a_stage2(a_copy_lds_window_pong); // GEMM loopK-1 static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { @@ -1099,17 +1120,18 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // GEMM loopK static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - - if(kIter * WG::kK < kKPerBlock - k_padded_zeros) { + if(kIter * WG::kK < kKPerBlock - k_padded_zeros) + { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload; static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { // read C warp tensor from C block tensor CWarpTensor c_warp_tensor; - 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)); + 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 constexpr(mIter == 0) dequant_mxfp4( @@ -1120,7 +1142,9 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 kIter); // warp GEMM - WG{}(c_warp_tensor, a_warp_tensor(number{}), dequant_B_n[nIter]); + WG{}(c_warp_tensor, + a_warp_tensor(number{}), + dequant_B_n[nIter]); // write C warp tensor into C block tensor c_block_tile.set_y_sliced_thread_data( @@ -1129,7 +1153,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 c_warp_tensor.get_thread_buffer()); }); if constexpr((kIter * MIterPerWarp + mIter) < - (KIterPerWarp * MIterPerWarp - m_preload)) + (KIterPerWarp * MIterPerWarp - m_preload)) { constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); @@ -1210,15 +1234,14 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 void* p_smem_ping, void* p_smem_pong) const { - return operator()( - a_dram_block_window_tmp, - [](const ADataType & a) { return a; }, - b_flat_dram_block_window_tmp, - scale_b_flat_window, - num_loop, - k_padded_zeros, - p_smem_ping, - p_smem_pong); + return operator()(a_dram_block_window_tmp, + identity{}, + b_flat_dram_block_window_tmp, + scale_b_flat_window, + num_loop, + k_padded_zeros, + p_smem_ping, + p_smem_pong); } }; From e7c1c77120b5dbf0808b0da01aede8f10f21d9e4 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Wed, 10 Sep 2025 11:41:54 +0000 Subject: [PATCH 2/9] Enable async_load in the F16xF4_flatmm_pipelineo on gfx950 --- ...ec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 113 ++++++++++++------ ...mm_pipeline_agmem_bgmem_creg_v1_policy.hpp | 98 ++++++++++----- 2 files changed, 144 insertions(+), 67 deletions(-) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 90f2e01bfb..51019b8e47 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -184,6 +184,11 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 CK_TILE_HOST_DEVICE static constexpr auto SchedulerPerM(index_t dsread_perM, index_t dswrite_perM, index_t load_perM) { +#if defined(__gfx950__) + // GFX950 use BUFFER_LOAD_LDS to fill lds_buffer_A. + // There is no separate DS_WRITE instruction at all. + dswrite_perM = 0; +#endif // Init inst order index_t max_data_inst = dsread_perM > load_perM ? (dsread_perM > dswrite_perM ? dsread_perM : dswrite_perM) @@ -488,7 +493,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 typename AElementFunction, typename BFlatBlockWindowTmp, typename DequantBFlatWindow> - CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindowTmp a_copy_dram_window, + CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindowTmp a_copy_dram_window_, const AElementFunction& a_element_func, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, const DequantBFlatWindow& scale_b_flat_window, @@ -516,39 +521,50 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; - __builtin_amdgcn_sched_barrier(0); + auto a_dram_view = a_copy_dram_window_.get_bottom_tensor_view(); + auto a_copy_dram_window = make_tile_window( + PipelinePolicy::template TransformF16xF4_ATensorView(a_dram_view), + a_copy_dram_window_.get_window_lengths(), + a_copy_dram_window_.get_window_origin(), + a_copy_dram_window_.get_tile_distribution()); // A tile in LDS ADataType* p_a_lds_ping = static_cast(p_smem_ping); ADataType* p_a_lds_pong = static_cast(p_smem_pong); - constexpr auto a_lds_block_desc = - PipelinePolicy::template MakeF16xF4_ALdsBlockDescriptor(); + constexpr auto write_a_lds_block_desc = + PipelinePolicy::template MakeF16xF4_WriteALdsBlockDescriptor(); + constexpr auto read_a_lds_block_desc = + PipelinePolicy::template MakeF16xF4_ReadALdsBlockDescriptor(); - auto a_lds_block_ping = - make_tensor_view(p_a_lds_ping, a_lds_block_desc); - auto a_lds_block_pong = - make_tensor_view(p_a_lds_pong, a_lds_block_desc); + auto write_a_lds_block_ping = + make_tensor_view(p_a_lds_ping, write_a_lds_block_desc); + auto write_a_lds_block_pong = + make_tensor_view(p_a_lds_pong, write_a_lds_block_desc); + auto read_a_lds_block_ping = + make_tensor_view(p_a_lds_ping, read_a_lds_block_desc); + auto read_a_lds_block_pong = + make_tensor_view(p_a_lds_pong, read_a_lds_block_desc); auto a_copy_lds_window_ping = - make_tile_window(a_lds_block_ping, + make_tile_window(write_a_lds_block_ping, make_tuple(number{}, number{}), {0, 0}, PipelinePolicy::template MakeADramTileDistribution()); auto a_copy_lds_window_pong = - make_tile_window(a_lds_block_pong, + make_tile_window(write_a_lds_block_pong, make_tuple(number{}, number{}), {0, 0}, PipelinePolicy::template MakeADramTileDistribution()); // ping-pong window for A LDS auto a_warp_window_ping_tmp = - make_tile_window(a_lds_block_ping, + make_tile_window(read_a_lds_block_ping, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, PipelinePolicy::template MakeF16xF4_ALDS_TileDistribution()); auto a_warp_window_pong_tmp = - make_tile_window(a_lds_block_pong, + make_tile_window(read_a_lds_block_pong, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, PipelinePolicy::template MakeF16xF4_ALDS_TileDistribution()); @@ -642,26 +658,27 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 using ABlockTile = decltype(load_tile(a_copy_dram_window)); ABlockTile a_block_tile; - auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a) { #if defined(__gfx950__) + auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a) { // global -> lds async_load_tile(lds_tile_a, dram_tile_a); -#else - // global -> vgpr - a_block_tile = load_tile(dram_tile_a); -#endif }; auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { -#if defined(__gfx950__) // data has been stored in lds, no need more operation. static_assert(std::is_same_v, "buffer_load_lds don't support element func fot A before mfma"); + }; #else + auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a) { + // global -> vgpr + a_block_tile = load_tile(dram_tile_a); + }; + auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { // vgpr -> lds auto a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile); store_tile(lds_tile_a, a_block_tile_transformed); -#endif }; +#endif // HEAD // Prefetch A0 @@ -841,14 +858,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 }); }); - // Prefill A(2i+1) - prefill_lds_a_stage2(a_copy_lds_window_pong); - - // Prefetch A(2i+2) - prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window); - // move A window to next k - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - // GEMM 2i static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { @@ -900,6 +909,14 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); + // Prefill A(2i+1) + prefill_lds_a_stage2(a_copy_lds_window_pong); + + // Prefetch A(2i+2) + prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window); + // move A window to next k + move_tile_window(a_copy_dram_window, {0, kKPerBlock}); + static_for<0, m_preload, 1>{}([&](auto loadIter) { constexpr auto mIter = loadIter % MIterPerWarp; constexpr auto kIter = loadIter / MIterPerWarp; @@ -943,14 +960,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 b_warp_tensor_ping(nIter)(kIter) = ub.u; }); }); - // Prefill A(2i+2) - prefill_lds_a_stage2(a_copy_lds_window_ping); - - // Prefetch A(2i+3) - prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window); - // move A window to next k - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - // GEMM 2i+1 static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { @@ -997,6 +1006,14 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 }); }); + // Prefill A(2i+2) + prefill_lds_a_stage2(a_copy_lds_window_ping); + + // Prefetch A(2i+3) + prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window); + // move A window to next k + move_tile_window(a_copy_dram_window, {0, kKPerBlock}); + // move B window to next flat K move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); @@ -1058,9 +1075,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } }); - // Prefill A(loopK) - prefill_lds_a_stage2(a_copy_lds_window_pong); - // GEMM loopK-1 static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { @@ -1108,6 +1122,9 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 }); }); + // Prefill A(loopK) + prefill_lds_a_stage2(a_copy_lds_window_pong); + static_for<0, m_preload, 1>{}([&](auto loadIter) { constexpr auto mIter = loadIter % MIterPerWarp; constexpr auto kIter = loadIter / MIterPerWarp; @@ -1243,6 +1260,26 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 p_smem_ping, p_smem_pong); } + + template + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + const DequantBFlatWindow& scale_b_flat_window, + const index_t num_loop, + void* p_smem_ping, + void* p_smem_pong) const + { + return operator()(a_dram_block_window_tmp, + identity{}, + b_flat_dram_block_window_tmp, + scale_b_flat_window, + num_loop, + 0, + p_smem_ping, + p_smem_pong); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index f2b9672ad4..df60e0da00 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -17,16 +17,72 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy static constexpr index_t N_Pack = 2; // it's fixed for fp4 static constexpr index_t K_Pack = 2; // it's fixed for fp4 - template - CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_ALdsBlockDescriptor() + template + CK_TILE_HOST_DEVICE static constexpr auto + TransformF16xF4_ATensorView(const NativeADramTensorView& a_dram_view) { - using namespace ck_tile; +#if defined(__gfx950__) //|| defined(__gfx942__) + constexpr int DynamicTileOffsetFlag = 0; constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1); static_assert(MPerXdl == 16 && NPerXdl == 16); + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t KPack = GetSmemPackA(); + + constexpr int ContiguousThreadsCntInDS_READ_16B = 4; + + // implement swizzle pattern on global side + // because we can't adjust the ds_write pattern of BUFFER_LOAD_LDS. + auto swizzle_a_dram_view_1 = transform_tensor_view( + a_dram_view, + make_tuple( + // M-dim is not affected by swizzle pattern + make_unmerge_transform( + make_tuple(number{}, number{})), + // K-dim is the swizzle dimension + make_unmerge_transform(make_tuple(number{}, + number{}, + number{}))), + make_tuple(sequence<0>{}, sequence<1>{}), + make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{})); + + auto swizzle_a_dram_view_2 = transform_tensor_view( + swizzle_a_dram_view_1, + make_tuple(make_pass_through_transform(number{}), + make_xor_transform(make_tuple(number{}, + number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<0>{}, sequence<1, 3>{}, sequence<2>{}, sequence<4>{}), + make_tuple(sequence<0>{}, sequence<1, 3>{}, sequence<2>{}, sequence<4>{})); + + return transform_tensor_view( + swizzle_a_dram_view_2, + make_tuple( + make_merge_transform_v3_division_mod( + make_tuple(number{}, number{})), + make_merge_transform_v3_division_mod(make_tuple(number{}, + number{}, + number{}))), + make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{}), + make_tuple(sequence<0>{}, sequence<1>{})); +#else + return a_dram_view; +#endif + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_ReadALdsBlockDescriptor() + { + constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); + constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1); + + static_assert(MPerXdl == 16 && NPerXdl == 16); + /*reduce transform layers,compare with old ck*/ constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; @@ -60,35 +116,19 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy } template - CK_TILE_HOST_DEVICE static constexpr auto MakeFp16xF4_ADramTileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_WriteALdsBlockDescriptor() { - using ADataType = remove_cvref_t; - using ALayout = remove_cvref_t; - - constexpr index_t BlockSize = Problem::kBlockSize; - +#if defined(__gfx950__) //|| defined(__gfx942__) constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - - constexpr index_t K1 = Problem::VectorLoadSize / sizeof(ADataType); - constexpr index_t K0 = KPerBlock / K1; - constexpr index_t M2 = get_warp_size() / K0; - - constexpr index_t M1 = BlockSize / get_warp_size(); - static_assert(M2 != 0, "M2 is zero, which will lead to a division by zero error."); - static_assert(M1 != 0, "M1 is zero, which will lead to a division by zero error."); - constexpr index_t M0 = MPerBlock / (M2 * M1); - static_assert(M0 * M1 * M2 == MPerBlock, - "Incorrect M0, M2, M1 configuration! " - "M0, M1, M2 must cover whole MPerBlock!"); - - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, - sequence<1, 2>, - sequence<0, 1>>{}); + constexpr index_t KPack = GetSmemPackA(); + return make_naive_tensor_descriptor(make_tuple(number{}, number{}), + make_tuple(number{}, number<1>{}), + number{}, + number<1>{}); +#else + return MakeF16xF4_ReadALdsBlockDescriptor(); +#endif } template From f4fdaedf4c68fb9705677aa34327f4eb7ba51365 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Thu, 11 Sep 2025 05:16:29 +0000 Subject: [PATCH 3/9] Add macro option to enable BUFFER_LOAD_LDS --- ...ec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 65 ++++++++++++------- ...mm_pipeline_agmem_bgmem_creg_v1_policy.hpp | 6 +- 2 files changed, 44 insertions(+), 27 deletions(-) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 51019b8e47..2d80d6d620 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -184,7 +184,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 CK_TILE_HOST_DEVICE static constexpr auto SchedulerPerM(index_t dsread_perM, index_t dswrite_perM, index_t load_perM) { -#if defined(__gfx950__) +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) // GFX950 use BUFFER_LOAD_LDS to fill lds_buffer_A. // There is no separate DS_WRITE instruction at all. dswrite_perM = 0; @@ -658,10 +658,17 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 using ABlockTile = decltype(load_tile(a_copy_dram_window)); ABlockTile a_block_tile; -#if defined(__gfx950__) - auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a) { + enum + { + PrefillBeforeGemm = 1, + PrefillAfterGemm = 2, + PrefillAlways = PrefillBeforeGemm | PrefillAfterGemm, + }; +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) + auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a, auto prefill_location) { // global -> lds - async_load_tile(lds_tile_a, dram_tile_a); + if constexpr(prefill_location & PrefillAfterGemm) + async_load_tile(lds_tile_a, dram_tile_a); }; auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { // data has been stored in lds, no need more operation. @@ -669,9 +676,10 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 "buffer_load_lds don't support element func fot A before mfma"); }; #else - auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a) { + auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a, auto prefill_location) { // global -> vgpr - a_block_tile = load_tile(dram_tile_a); + if constexpr(prefill_location & PrefillBeforeGemm) + a_block_tile = load_tile(dram_tile_a); }; auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { // vgpr -> lds @@ -682,7 +690,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // HEAD // Prefetch A0 - prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window); + prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window, number{}); // move A window to next k move_tile_window(a_copy_dram_window, {0, kKPerBlock}); @@ -725,7 +733,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 __builtin_amdgcn_sched_barrier(0); // Prefetch A1 - prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window); + prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window, number{}); // move A window to next k move_tile_window(a_copy_dram_window, {0, kKPerBlock}); @@ -858,6 +866,12 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 }); }); + // Prefill A(2i+1) + prefill_lds_a_stage2(a_copy_lds_window_pong); + + // Prefetch A(2i+2) + prefill_lds_a_stage1( + a_copy_lds_window_ping, a_copy_dram_window, number{}); // GEMM 2i static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { @@ -904,19 +918,16 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } }); }); + prefill_lds_a_stage1( + a_copy_lds_window_ping, a_copy_dram_window, number{}); + + // move A window to next k + move_tile_window(a_copy_dram_window, {0, kKPerBlock}); // move B window to next flat K move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); - // Prefill A(2i+1) - prefill_lds_a_stage2(a_copy_lds_window_pong); - - // Prefetch A(2i+2) - prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window); - // move A window to next k - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - static_for<0, m_preload, 1>{}([&](auto loadIter) { constexpr auto mIter = loadIter % MIterPerWarp; constexpr auto kIter = loadIter / MIterPerWarp; @@ -960,6 +971,14 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 b_warp_tensor_ping(nIter)(kIter) = ub.u; }); }); + + // Prefill A(2i+2) + prefill_lds_a_stage2(a_copy_lds_window_ping); + + // Prefetch A(2i+3) + prefill_lds_a_stage1( + a_copy_lds_window_pong, a_copy_dram_window, number{}); + // GEMM 2i+1 static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { @@ -1005,15 +1024,11 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } }); }); + prefill_lds_a_stage1( + a_copy_lds_window_pong, a_copy_dram_window, number{}); - // Prefill A(2i+2) - prefill_lds_a_stage2(a_copy_lds_window_ping); - - // Prefetch A(2i+3) - prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window); // move A window to next k move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - // move B window to next flat K move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter}); move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp}); @@ -1075,6 +1090,9 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } }); + // Prefill A(loopK) + prefill_lds_a_stage2(a_copy_lds_window_pong); + // GEMM loopK-1 static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { @@ -1122,9 +1140,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 }); }); - // Prefill A(loopK) - prefill_lds_a_stage2(a_copy_lds_window_pong); - static_for<0, m_preload, 1>{}([&](auto loadIter) { constexpr auto mIter = loadIter % MIterPerWarp; constexpr auto kIter = loadIter / MIterPerWarp; diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index df60e0da00..c23bb98bd9 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -7,6 +7,8 @@ namespace ck_tile { +#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE 1 + struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy { static constexpr auto I0 = number<0>{}; @@ -21,7 +23,7 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy CK_TILE_HOST_DEVICE static constexpr auto TransformF16xF4_ATensorView(const NativeADramTensorView& a_dram_view) { -#if defined(__gfx950__) //|| defined(__gfx942__) +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) constexpr int DynamicTileOffsetFlag = 0; constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); @@ -118,7 +120,7 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_WriteALdsBlockDescriptor() { -#if defined(__gfx950__) //|| defined(__gfx942__) +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; constexpr index_t KPack = GetSmemPackA(); From 1c05d3be89d9c37d93d6656fd7b037e3ca97638f Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Thu, 11 Sep 2025 06:40:11 +0000 Subject: [PATCH 4/9] Disable hot schduler when enable buffer_load_lds --- ...ed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 12 ++++++++++-- ..._flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp | 14 ++++++++++++-- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 2d80d6d620..36724c25a6 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -184,7 +184,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 CK_TILE_HOST_DEVICE static constexpr auto SchedulerPerM(index_t dsread_perM, index_t dswrite_perM, index_t load_perM) { -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS // GFX950 use BUFFER_LOAD_LDS to fill lds_buffer_A. // There is no separate DS_WRITE instruction at all. dswrite_perM = 0; @@ -347,6 +347,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // 0 M7N2: 63 - - 8 - // 0 M7N3: 64 4 - - - +#if !CKTILE_FLATMM_USE_BUFFER_LOAD_LDS _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) { _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) @@ -410,10 +411,12 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 if(Aload_num_perK == 0) __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read __builtin_amdgcn_sched_barrier(0); +#endif } CK_TILE_HOST_DEVICE static constexpr auto Last2ndHotLoopScheduler() { +#if !CKTILE_FLATMM_USE_BUFFER_LOAD_LDS _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) { _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) @@ -462,10 +465,12 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } } __builtin_amdgcn_sched_barrier(0); +#endif } CK_TILE_HOST_DEVICE static constexpr auto LastHotLoopScheduler() { +#if !CKTILE_FLATMM_USE_BUFFER_LOAD_LDS _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) { _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) @@ -482,6 +487,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } } // __builtin_amdgcn_sched_barrier(0); +#endif } CK_TILE_HOST_DEVICE static constexpr auto GetADramTileDistribution() @@ -521,6 +527,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + __builtin_amdgcn_sched_barrier(0); + auto a_dram_view = a_copy_dram_window_.get_bottom_tensor_view(); auto a_copy_dram_window = make_tile_window( PipelinePolicy::template TransformF16xF4_ATensorView(a_dram_view), @@ -664,7 +672,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 PrefillAfterGemm = 2, PrefillAlways = PrefillBeforeGemm | PrefillAfterGemm, }; -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a, auto prefill_location) { // global -> lds if constexpr(prefill_location & PrefillAfterGemm) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index c23bb98bd9..fb2e4de0ce 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -9,6 +9,16 @@ namespace ck_tile { #define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE 1 +#if defined(__gfx950__) +#define CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4 1 +#else +#define CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4 0 +#endif + +#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS \ + (CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && \ + CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4) + struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy { static constexpr auto I0 = number<0>{}; @@ -23,7 +33,7 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy CK_TILE_HOST_DEVICE static constexpr auto TransformF16xF4_ATensorView(const NativeADramTensorView& a_dram_view) { -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS constexpr int DynamicTileOffsetFlag = 0; constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); @@ -120,7 +130,7 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_WriteALdsBlockDescriptor() { -#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && defined(__gfx950__) +#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; constexpr index_t KPack = GetSmemPackA(); From 2e36b3c3ff6fcc4fffc5ac73736ad41cd0ea082b Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Thu, 11 Sep 2025 07:06:52 +0000 Subject: [PATCH 5/9] Add buffer_load_lds synchronization --- .../pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 1 + 1 file changed, 1 insertion(+) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 36724c25a6..2908845d3c 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -679,6 +679,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 async_load_tile(lds_tile_a, dram_tile_a); }; auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { + async_load_fence(); // data has been stored in lds, no need more operation. static_assert(std::is_same_v, "buffer_load_lds don't support element func fot A before mfma"); From a44d35172db438a49df3611faa9d92379b19e9d3 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Thu, 11 Sep 2025 07:39:41 +0000 Subject: [PATCH 6/9] Add async_load for tile_scatter_gather --- .../core/tensor/tile_scatter_gather.hpp | 115 ++++++++++++++++++ include/ck_tile/core/tensor/tile_window.hpp | 27 ++++ ...ec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 10 +- 3 files changed, 146 insertions(+), 6 deletions(-) diff --git a/include/ck_tile/core/tensor/tile_scatter_gather.hpp b/include/ck_tile/core/tensor/tile_scatter_gather.hpp index 082207d1df..9d16b7cde2 100644 --- a/include/ck_tile/core/tensor/tile_scatter_gather.hpp +++ b/include/ck_tile/core/tensor/tile_scatter_gather.hpp @@ -475,6 +475,96 @@ struct tile_scatter_gather }); } + template + CK_TILE_DEVICE auto async_load(LdsTileWindow_&& lds_tile, + number = {}, + bool_constant = {}) const + { + using LdsTileWindow = remove_cvref_t; + using LdsDataType = typename LdsTileWindow::DataType; + using Traits = load_store_traits; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = TileDstr{}; + + // Precompute invariant values outside loops + const auto window_origin = lds_tile.get_window_origin(); + const auto& bottom_tensor_view = lds_tile.get_bottom_tensor_view(); + const auto& tensor_descriptor = bottom_tensor_view.get_tensor_descriptor(); + auto smem_base_ptr = bottom_tensor_view.get_buffer_view().p_data_; + + // loop over thread tensor space [y0, y1, ...] + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + /// TODO: use structure binding (to be captured later) if compiled in C++20 + auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + auto lds_window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto lds_bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + + // Use precomputed window origin + auto lds_bottom_tensor_thread_idx = + window_origin + lds_window_adaptor_thread_coord.get_bottom_index(); + // Use precomputed tensor descriptor + const auto lds_coord = + make_tensor_coordinate(tensor_descriptor, lds_bottom_tensor_thread_idx); + // Calculate SMEM address using base pointer + CK_TILE_LDS_ADDR LdsDataType* smem = smem_base_ptr + lds_coord.get_offset(); + + // data index [y0, y1, ...] + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + constexpr auto idx_gather = idx_ys_start[number{}]; + const auto page_offset = page_idx_[idx_gather]; + + // read from bottom tensor + if constexpr(std::is_same_v) + this->get_bottom_tensor_view().template async_get_vectorized_elements( + smem, + bottom_tensor_thread_coord, + page_offset, + bool_constant{}); + else + this->get_bottom_tensor_view().template async_get_vectorized_elements( + smem, + bottom_tensor_thread_coord, + page_offset, + valids_[idx_gather], + bool_constant{}); + + // move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto forward_step_scatter = generate_tuple( + [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + forward_step_scatter); + // lds_diff doesn't need to mask the difference of the gather-dim. + constexpr auto lds_idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + idx_diff_ys); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + move_window_adaptor_and_bottom_tensor_thread_coordinate( + lds_window_adaptor_thread_coord, + lds_bottom_tensor_thread_coord, + lds_idx_diff_ps_ys); + } + }); + }); + } + // TODO: currently async load only implemented in inline asm template +CK_TILE_DEVICE auto replace_bottom_tensor_view(const NewTensorView_& new_tensor_view, + const tile_scatter_gather& tile_window) +{ + return make_tile_scatter_gather(new_tensor_view, + tile_window.window_lengths_, + tile_window.window_origin_, + tile_window.tile_dstr_, + tile_window.page_idx_, + tile_window.valids_); +} + template +CK_TILE_DEVICE auto +replace_bottom_tensor_view(const NewTensorView_& new_tensor_view, + const tile_window_with_static_distribution& tile_window) +{ + return make_tile_window(new_tensor_view, + tile_window.get_window_lengths(), + tile_window.get_window_origin(), + tile_window.get_tile_distribution()); +} + template +CK_TILE_DEVICE auto replace_bottom_tensor_view( + const NewTensorView_& new_tensor_view, + const tile_window_with_static_lengths& tile_window) +{ + return make_tile_window( + new_tensor_view, tile_window.get_window_lengths(), tile_window.get_window_origin()); +} + template CK_TILE_DEVICE void move_tile_window( tile_window_with_static_lengths& window, diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 2908845d3c..69e5917949 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -529,12 +529,10 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 __builtin_amdgcn_sched_barrier(0); - auto a_dram_view = a_copy_dram_window_.get_bottom_tensor_view(); - auto a_copy_dram_window = make_tile_window( - PipelinePolicy::template TransformF16xF4_ATensorView(a_dram_view), - a_copy_dram_window_.get_window_lengths(), - a_copy_dram_window_.get_window_origin(), - a_copy_dram_window_.get_tile_distribution()); + auto a_copy_dram_window = replace_bottom_tensor_view( + PipelinePolicy::template TransformF16xF4_ATensorView( + a_copy_dram_window_.get_bottom_tensor_view()), + a_copy_dram_window_); // A tile in LDS ADataType* p_a_lds_ping = static_cast(p_smem_ping); From fdbb22d2b619d1fd8982642215471d2f82f97bab Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Fri, 12 Sep 2025 05:32:45 +0000 Subject: [PATCH 7/9] Add async_load for tile_scatter_gather --- include/ck_tile/core/tensor/tile_scatter_gather.hpp | 12 ++++++++---- ...ec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp | 2 +- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/include/ck_tile/core/tensor/tile_scatter_gather.hpp b/include/ck_tile/core/tensor/tile_scatter_gather.hpp index 9d16b7cde2..c8c7e5eb6d 100644 --- a/include/ck_tile/core/tensor/tile_scatter_gather.hpp +++ b/include/ck_tile/core/tensor/tile_scatter_gather.hpp @@ -522,18 +522,22 @@ struct tile_scatter_gather constexpr auto idx_gather = idx_ys_start[number{}]; const auto page_offset = page_idx_[idx_gather]; + // merge page_offset into bottom_coord + auto mixed_bottom_thread_coord = bottom_tensor_thread_coord; + mixed_bottom_thread_coord.get_hidden_index()[number<0>{}] += page_offset; + // read from bottom tensor if constexpr(std::is_same_v) this->get_bottom_tensor_view().template async_get_vectorized_elements( smem, - bottom_tensor_thread_coord, - page_offset, + mixed_bottom_thread_coord, + number<0>{}, bool_constant{}); else this->get_bottom_tensor_view().template async_get_vectorized_elements( smem, - bottom_tensor_thread_coord, - page_offset, + mixed_bottom_thread_coord, + number<0>{}, valids_[idx_gather], bool_constant{}); diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index fb2e4de0ce..01535d06dd 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -7,7 +7,7 @@ namespace ck_tile { -#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE 1 +#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE 0 #if defined(__gfx950__) #define CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4 1 From 186c69960ed23fd2dafa79a803f3b897092b6a77 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Mon, 15 Sep 2025 06:07:34 +0000 Subject: [PATCH 8/9] Change the preshuffle pattern of MFMA_16x16x128_F8 --- example/ck_tile/18_flatmm/flatmm_basic.cpp | 34 +++++---- example/ck_tile/18_flatmm/flatmm_basic.hpp | 4 +- example/ck_tile/21_moe_flatmm/moe_flatmm.cpp | 17 ++--- .../flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 12 +-- ...mm_pipeline_agmem_bgmem_creg_v1_policy.hpp | 76 ++++++++++++++----- .../moe_flatmm_pipeline_agmem_bgmem_creg.hpp | 4 +- 6 files changed, 92 insertions(+), 55 deletions(-) diff --git a/example/ck_tile/18_flatmm/flatmm_basic.cpp b/example/ck_tile/18_flatmm/flatmm_basic.cpp index 60dced003b..2fecdc1663 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.cpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.cpp @@ -50,35 +50,41 @@ template auto shuffle_b(const ck_tile::HostTensor& t) { assert(t.get_lengths().size() == 2); - int n_ = t.get_lengths()[1]; - int k_ = t.get_lengths()[0]; - constexpr int divisor = FlatmmConfig::N_Warp_Tile == 32 ? 2 : 4; + int n_ = t.get_lengths()[1]; + int k_ = t.get_lengths()[0]; + + constexpr int MaxVecSize = 16 / sizeof(T); + constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane); + ck_tile::HostTensor t_view({n_ / FlatmmConfig::N_Warp_Tile, FlatmmConfig::N_Warp_Tile, - k_ / FlatmmConfig::K_Warp_Tile, - divisor, - FlatmmConfig::K_Warp_Tile / divisor}); + k_ / ItemsPerAccess, + ItemsPerAccess}); std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + return ck_tile::reference_permute(t_view, {0, 2, 1, 3}); } template auto shuffle_b_v1(const ck_tile::HostTensor& t) { assert(t.get_lengths().size() == 2); - int n_ = t.get_lengths()[1]; - int k_ = t.get_lengths()[0]; - constexpr int divisor = FlatmmConfig::N_Warp_Tile == 32 ? 2 : 4; + int n_ = t.get_lengths()[1]; + int k_ = t.get_lengths()[0]; + + constexpr int MaxVecSize = 16 / sizeof(T); + constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane); constexpr int NRepeat = FlatmmConfig::N_Tile / FlatmmConfig::N_Warp_Tile / FlatmmConfig::N_Warp; + ck_tile::HostTensor t_view({n_ / FlatmmConfig::N_Tile, FlatmmConfig::N_Warp, FlatmmConfig::N_Warp_Tile, NRepeat, - k_ / FlatmmConfig::K_Warp_Tile, - divisor, - FlatmmConfig::K_Warp_Tile / divisor}); + k_ / ItemsPerAccess, + ItemsPerAccess}); std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 5, 2, 6}); + return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 2, 5}); } template diff --git a/example/ck_tile/18_flatmm/flatmm_basic.hpp b/example/ck_tile/18_flatmm/flatmm_basic.hpp index 063a981daf..20b9c6ee15 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.hpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.hpp @@ -81,7 +81,7 @@ struct FlatmmConfig16 static constexpr bool DoubleSmemBuffer = false; static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp; - static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; + static constexpr bool TiledMMAPermuteN = N_Repeat % 4 == 0; }; template @@ -94,7 +94,7 @@ struct FlatmmConfig16_950 : public FlatmmConfig16 static constexpr int N_Repeat = N_Tile / FlatmmConfig16::N_Warp_Tile / FlatmmConfig16::N_Warp; - static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; + static constexpr bool TiledMMAPermuteN = N_Repeat % 4 == 0; }; template diff --git a/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp b/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp index d2b9e666de..e002dcc57a 100644 --- a/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp +++ b/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp @@ -35,17 +35,16 @@ auto shuffle_b(const ck_tile::HostTensor& t) int n_ = t.get_lengths()[1]; int k_ = t.get_lengths()[0]; - constexpr int N_Warp_Tile = FlatmmConfig::N_Warp_Tile; - constexpr int N_Warp = FlatmmConfig::N_Warp; - constexpr int KPerLane = FlatmmConfig::K_Warp_Tile / (64 / N_Warp_Tile); + constexpr int MaxVecSize = 16 / sizeof(T); + constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane); - ck_tile::HostTensor t_view({n_ / N_Warp_Tile, - N_Warp_Tile, - k_ / (64 * KPerLane / N_Warp_Tile), - 64 / N_Warp_Tile, - KPerLane}); + ck_tile::HostTensor t_view({n_ / FlatmmConfig::N_Warp_Tile, + FlatmmConfig::N_Warp_Tile, + k_ / ItemsPerAccess, + ItemsPerAccess}); std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + return ck_tile::reference_permute(t_view, {0, 2, 1, 3}); } template diff --git a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp index ee88be1466..ceb6ef6734 100644 --- a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -565,13 +565,13 @@ defined(USING_MFMA_32x32x64) && defined(ENABLE_FP4) // mi350 fp4 32c 1*K1 make_tile_window(a_lds_block_ping, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); auto a_warp_window_pong_tmp = make_tile_window(a_lds_block_pong, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); statically_indexed_array< statically_indexed_array, @@ -586,16 +586,10 @@ defined(USING_MFMA_32x32x64) && defined(ENABLE_FP4) // mi350 fp4 32c 1*K1 static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { a_warp_windows_ping(mIter)(kIter) = a_warp_window_ping_tmp; + a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp; move_tile_window(a_warp_windows_ping(mIter)(kIter), {mIter * MPerBlockPerIter, kIter * KPerBlockPerIter}); - }); - }); - - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp; - move_tile_window(a_warp_windows_pong(mIter)(kIter), {mIter * MPerBlockPerIter, kIter * KPerBlockPerIter}); }); diff --git a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index f546cd7bac..cc882db8d6 100644 --- a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -250,6 +250,37 @@ struct UniversalFlatmmPipelineAgBgCrPolicy } } + template + CK_TILE_HOST_DEVICE static constexpr auto MakeALDS_WarpTileDistribution() + { + using TileShape = typename Problem::BlockGemmShape; + using ADataType = remove_cvref_t; + using ALayout = remove_cvref_t; + + static_assert(TileShape::BlockWarps::at(I0) == 1, "requires Wave_M == 1"); + + constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); + constexpr index_t KPerXdl = Problem::BlockGemmShape::WarpTile::at(I2); + + constexpr int Repeat = TileShape::BlockWarps::at(number<1>{}); + + constexpr int KLane = get_warp_size() / MPerXdl; + constexpr int KPerThread = KPerXdl / KLane; + + constexpr int MaxVecSize = 16 / sizeof(ADataType); + constexpr int KItemsPerLoad = min(MaxVecSize, KPerThread); + constexpr int KFragment = KPerThread / KItemsPerLoad; + + return make_static_tile_distribution( + tile_distribution_encoding< + sequence, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<1, 0>>, + sequence<2, 2>, + sequence<0, 2>>{}); + } + template CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution() { @@ -303,10 +334,10 @@ struct UniversalFlatmmPipelineAgBgCrPolicy { constexpr index_t K1 = Problem::VectorLoadSize / sizeof(ADataType); constexpr index_t K0 = KPerBlock / K1; - constexpr index_t M2 = get_warp_size() / K0; // coalesce reading for each blocks - if constexpr(get_warp_size() % (M2 * K0) == 0) + if constexpr(get_warp_size() % K0 == 0) { + constexpr index_t M2 = get_warp_size() / K0; constexpr index_t M1 = BlockSize / get_warp_size(); static_assert(M2 != 0, "M2 is zero, which will lead to a division by zero error."); static_assert(M1 != 0, "M1 is zero, which will lead to a division by zero error."); @@ -325,18 +356,18 @@ struct UniversalFlatmmPipelineAgBgCrPolicy } else { - constexpr index_t M0 = BlockSize / get_warp_size(); - constexpr index_t M1 = MPerBlock / (M2 * M0); - static_assert(M0 * M1 * M2 == MPerBlock, - "Incorrect M0, M1, M2 configuration! " - "M0, M1, M2 must cover whole MPerBlock!"); + constexpr index_t KWave = K0 / get_warp_size(); + constexpr index_t M0 = BlockSize / get_warp_size() / KWave; + constexpr index_t M1 = MPerBlock / M0; + return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, - sequence<1, 2>, - sequence<1, 1>>{}); + tile_distribution_encoding< + sequence<1>, + tuple, sequence>, + tuple, sequence<2>>, + tuple, sequence<1>>, + sequence<1, 2>, + sequence<1, 2>>{}); } } } @@ -381,11 +412,17 @@ struct UniversalFlatmmPipelineAgBgCrPolicy constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNum = BlockSize / WaveSize; - constexpr index_t KBPerLoad = GetKBPerLoad(); - constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim + constexpr index_t KBPerLoad = GetKBPerLoad(); + + constexpr index_t MaxVecSize = 16 / sizeof(typename Problem::BDataType); + constexpr index_t KItemsPerLoad = min(KBPerLoad, MaxVecSize); + constexpr index_t KFragment = KBPerLoad / KItemsPerLoad; + static_assert(KFragment * KItemsPerLoad == KBPerLoad); + + constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim./ constexpr index_t KWavePerBlk = 1; - constexpr index_t KRepeat = 1; static_assert(TileShape::flatKPerWarp == KThdPerWave * KBPerLoad, "wrong"); + static_assert(TileShape::BlockWarps::at(number<2>{}) == 1, "Requires K_Warp == 1"); constexpr index_t NBPerLoad = 1; constexpr index_t NThdPerWave = 1; @@ -396,9 +433,10 @@ struct UniversalFlatmmPipelineAgBgCrPolicy return make_static_tile_distribution( tile_distribution_encoding< - sequence, // ? - tuple, // second direction - sequence>, // first direction + sequence, // ? + tuple, // second direction + sequence>, // first + // direction // wave in blk, // thd in wave // // tuple, sequence<1, 2>>, // which direction diff --git a/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp b/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp index cd91f33a79..7a47bd7cd2 100644 --- a/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp +++ b/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp @@ -508,13 +508,13 @@ struct MoeFlatmmPipelineAGmemBGmemCRegV1 make_tile_window(a_lds_block_ping, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); auto a_warp_window_pong_tmp = make_tile_window(a_lds_block_pong, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); statically_indexed_array< statically_indexed_array, From 31ec8bd33a1bb5104b44a6e2fd746e8b816cd975 Mon Sep 17 00:00:00 2001 From: AMD-dteng Date: Wed, 17 Sep 2025 06:54:32 -0500 Subject: [PATCH 9/9] use builtin function to wait A load's data --- ...ec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 67 +++---------------- 1 file changed, 8 insertions(+), 59 deletions(-) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 69e5917949..969662d98e 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -170,6 +170,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 static constexpr index_t ScaleBload_num = kNPerBlock * kKPerBlock / NWarp / 32 / ScaleBload_K1 / WaveSize; // BlockN * BlockK / NWarp / ScalePerK / ScaleB_K1 / wavesize + static constexpr index_t Bload_total_num = Bload_num_perK * KIterPerWarp + ScaleBload_num + 0X3f0; static constexpr index_t KPerScaleLoad = KIterPerWarp / ScaleBload_num; static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2; static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter; @@ -347,7 +348,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // 0 M7N2: 63 - - 8 - // 0 M7N3: 64 4 - - - -#if !CKTILE_FLATMM_USE_BUFFER_LOAD_LDS _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) { _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) @@ -359,32 +359,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // Calculate ds_read number per M dsread_perM = dsread_per_wg; - // Calculate ds_write number per M - if(mIter == 0) - { - dswrite_perM = - (dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep) > 0 - ? dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep - : 0; - } - else if(mIter >= MIterPerWarp - DsWritePreIssue + 1) - { - dswrite_perM = 0; - } - else - { - dswrite_perM = (dswrite_num_perK - - (MIterPerWarp - DsWritePreIssue - mIter) * dswrite_rep) > 0 - ? dswrite_rep - : 0; - } - // Add ds write when ds write data > needed - if(dswrite_num_perK == 0 && kIter == (KIterPerWarp - 1 - dswrite_kIter)) - { - if(mIter == MIterPerWarp - 1 - dswrite_mIter) - dswrite_perM = 1; - } - // Calculate buffer_load number per M if(mIter < HalfMIter) { @@ -411,12 +385,10 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 if(Aload_num_perK == 0) __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read __builtin_amdgcn_sched_barrier(0); -#endif } CK_TILE_HOST_DEVICE static constexpr auto Last2ndHotLoopScheduler() { -#if !CKTILE_FLATMM_USE_BUFFER_LOAD_LDS _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) { _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) @@ -428,32 +400,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // Calculate ds_read number per M dsread_perM = dsread_per_wg; - // Calculate ds_write number per M - if(mIter == 0) - { - dswrite_perM = - (dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep) > 0 - ? dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep - : 0; - } - else if(mIter >= MIterPerWarp - DsWritePreIssue + 1) - { - dswrite_perM = 0; - } - else - { - dswrite_perM = (dswrite_num_perK - - (MIterPerWarp - DsWritePreIssue - mIter) * dswrite_rep) > 0 - ? dswrite_rep - : 0; - } - // Add ds write when ds write data > needed - if(dswrite_num_perK == 0 && kIter == (KIterPerWarp - 1 - dswrite_kIter)) - { - if(mIter == MIterPerWarp - 1 - dswrite_mIter) - dswrite_perM = 1; - } - // Calculate buffer_load number per M if(mIter < HalfMIter) { @@ -465,12 +411,10 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } } __builtin_amdgcn_sched_barrier(0); -#endif } CK_TILE_HOST_DEVICE static constexpr auto LastHotLoopScheduler() { -#if !CKTILE_FLATMM_USE_BUFFER_LOAD_LDS _Pragma("unroll") for(int kIter = 0; kIter < KIterPerWarp; kIter++) { _Pragma("unroll") for(int mIter = 0; mIter < MIterPerWarp; mIter++) @@ -487,7 +431,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 } } // __builtin_amdgcn_sched_barrier(0); -#endif } CK_TILE_HOST_DEVICE static constexpr auto GetADramTileDistribution() @@ -677,7 +620,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 async_load_tile(lds_tile_a, dram_tile_a); }; auto prefill_lds_a_stage2 = [&](auto lds_tile_a) { - async_load_fence(); + // async_load_fence(); + // __builtin_amdgcn_s_waitcnt(0x03fc); // data has been stored in lds, no need more operation. static_assert(std::is_same_v, "buffer_load_lds don't support element func fot A before mfma"); @@ -747,6 +691,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + __builtin_amdgcn_s_waitcnt(Bload_total_num); block_sync_lds(); // preload A00,A10... from lds @@ -921,6 +866,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // barrier if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) { + __builtin_amdgcn_s_waitcnt(Bload_total_num); block_sync_lds(); } }); @@ -1027,6 +973,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // barrier if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) { + __builtin_amdgcn_s_waitcnt(Bload_total_num); block_sync_lds(); } }); @@ -1142,6 +1089,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // barrier if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) { + __builtin_amdgcn_s_waitcnt(Bload_total_num); block_sync_lds(); } }); @@ -1252,6 +1200,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 // barrier if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) { + __builtin_amdgcn_s_waitcnt(Bload_total_num); block_sync_lds(); } });