diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp index 8a2198f85c..4a3f860276 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp @@ -147,12 +147,14 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem template <> struct PipelineImpl { - template - CK_TILE_DEVICE void GlobalPrefetch(BlockTile& block_tile, + template + CK_TILE_DEVICE void GlobalPrefetch(DstBlockTile& dst_block_tile, SrcTileWindow& dram_tile_window) const { - load_tile_raw(block_tile, dram_tile_window); + // TODO: we need to have an api of load_tile which takes as param output tile + load_tile_raw(dst_block_tile, dram_tile_window); move_tile_window(dram_tile_window, {0, KPerBlock}); + buffer_load_fence(); } template @@ -216,6 +218,8 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem make_tuple(number{}, number{}), a_dram_block_window_tmp.get_window_origin(), Policy::template MakeADramTileDistribution()); + a_copy_dram_window.init_raw(); + // A LDS tile window for store auto a_copy_lds_window = make_tile_window(a_lds_block, @@ -228,6 +232,8 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem make_tuple(number{}, number{}), b_dram_block_window_tmp.get_window_origin(), Policy::template MakeBDramTileDistribution()); + b_copy_dram_window.init_raw(); + // B LDS tile window for store auto b_copy_lds_window = make_tile_window(b_lds_block, @@ -283,7 +289,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func); LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); - // Global prefetch [2, PrefetchStages] + // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { GlobalPrefetch(a_block_tiles.get(number{}), a_copy_dram_window); GlobalPrefetch(b_block_tiles.get(number{}), b_copy_dram_window); @@ -295,7 +301,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem index_t i = 0; do { - static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { + static_for<0, PrefetchStages, 1>{}([&](auto prefetch_idx) { block_sync_lds(); // block_gemm.LocalPrefetch(); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); @@ -330,10 +336,10 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem block_sync_lds(); LocalPrefill(a_copy_lds_window, - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), + a_block_tiles.get(number{}), a_element_func); LocalPrefill(b_copy_lds_window, - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), + b_block_tiles.get(number{}), b_element_func); });