diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp index 77f2611317..6331379110 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp @@ -243,7 +243,6 @@ struct BlockUniversalGemmArBsCr : public BlockUniversalGemmBase{}([&](auto kIter) { LocalPrefetchB(b_block_window); 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 a1ddd59838..1adee65ff0 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 @@ -412,16 +412,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func); } } - else - { - if constexpr(is_a_col_major && !is_a_load_tr_v()) - { - auto a_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{})); - a_block_tiles.get(I0{}) = a_shuffle_tmp; - } - } + if constexpr(SkipBLds == false) { if constexpr(is_b_row_major && !is_b_load_tr_v()) @@ -436,16 +427,6 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); } } - else - { - if constexpr(is_b_row_major && !is_b_load_tr_v()) - { - auto b_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{})); - b_block_tiles.get(I0{}) = b_shuffle_tmp; - } - } // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { @@ -500,20 +481,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem a_element_func); } } - else - { - if constexpr(is_a_col_major && !is_a_load_tr_v()) - { - auto a_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d( - a_shuffle_tmp, - a_block_tiles.get( - number<(prefetch_idx + 1) % PrefetchStages>{})); - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) = - a_shuffle_tmp; - } - } + if constexpr(SkipBLds == false) { if constexpr(is_b_row_major && !is_b_load_tr_v()) @@ -536,20 +504,6 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem b_element_func); } } - else - { - if constexpr(is_b_row_major && !is_b_load_tr_v()) - { - auto b_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d( - b_shuffle_tmp, - b_block_tiles.get( - number<(prefetch_idx + 1) % PrefetchStages>{})); - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) = - b_shuffle_tmp; - } - } Base::GlobalPrefetch(a_block_tiles.get(number{}), a_copy_dram_window, @@ -595,19 +549,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem a_element_func); } } - else - { - if constexpr(is_a_col_major && !is_a_load_tr_v()) - { - auto a_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d( - a_shuffle_tmp, - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) = - a_shuffle_tmp; - } - } + if constexpr(SkipBLds == false) { if constexpr(is_b_row_major && !is_b_load_tr_v()) @@ -625,19 +567,6 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem b_element_func); } } - else - { - if constexpr(is_b_row_major && !is_b_load_tr_v()) - { - auto b_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d( - b_shuffle_tmp, - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) = - b_shuffle_tmp; - } - } }); block_sync_lds(); @@ -848,16 +777,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func); } } - else - { - if constexpr(is_a_col_major && !is_a_load_tr_v()) - { - auto a_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{})); - a_block_tiles.get(I0{}) = a_shuffle_tmp; - } - } + if constexpr(SkipBLds == false) { if constexpr(is_b_row_major && !is_b_load_tr_v()) @@ -872,16 +792,11 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); } } - else - { - if constexpr(is_b_row_major && !is_b_load_tr_v()) - { - auto b_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{})); - b_block_tiles.get(I0{}) = b_shuffle_tmp; - } - } + + // TODO add encoding and support for BRowMajor for SkipBLds, currently: + // MakeShuffledBRegTileDistribution takes into account shuffling encoding which is used + // for [Global -> Vgpr -> Lds -> Vgpr] reading, but for skipping lds we need to have + // different shuffled layout for [Global -> Vgpr] reads, similar for AColMajor // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { @@ -934,20 +849,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem a_element_func); } } - else - { - if constexpr(is_a_col_major && !is_a_load_tr_v()) - { - auto a_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d( - a_shuffle_tmp, - a_block_tiles.get( - number<(prefetch_idx + 1) % PrefetchStages>{})); - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) = - a_shuffle_tmp; - } - } + if constexpr(SkipBLds == false) { if constexpr(is_b_row_major && !is_b_load_tr_v()) diff --git a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp index e59a08b5cd..933e8ccdcb 100644 --- a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp +++ b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp @@ -68,10 +68,10 @@ struct TileGemmUniversalTraits static constexpr bool TransposeC = TransposeC_; /// @brief SkipALds_ Parameterize the data flow for matrix A: if true, global memory -> - /// registers; if false, global memory -> LDS -> registers. + /// VGPR; if false, global memory -> VGPR -> LDS -> VGPR. static constexpr bool SkipALds = SkipALds_; /// @brief SkipBLds_ Parameterize the data flow for matrix B: if true, global memory -> - /// registers; if false, global memory -> LDS -> registers. + /// VGPR; if false, global memory -> VGPR -> LDS -> VGPR. static constexpr bool SkipBLds = SkipBLds_; static constexpr bool UseStructuredSparsity = UseStructuredSparsity_; static constexpr bool UsePersistentKernel = UsePersistentKernel_; diff --git a/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp b/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp index ab59f9cd3b..2296922e9b 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp @@ -61,6 +61,7 @@ using KernelTypesMem = ::testing::Types< std::tuple< Col, Col, Row, F8, F8, F32, F16, I256, I256, I64, I32, I32, I16, Intrawave, Mem, UseLDS, UseLDS>, std::tuple< Col, Col, Row, F8, F8, F32, F16, I256, I256, I64, I32, I32, I16, Interwave, Mem, UseLDS, UseLDS>, // SkipALds/SkipBLds + // TODO: Add A column-wise/B row-wise layout implementation and tests for the SkipLDS feature std::tuple< Row, Col, Row, F16, F16, F32, F16, I256, I256, I64, I32, I32, I16, Interwave, Mem, UseLDS, SkipLDS>, std::tuple< Row, Col, Row, F16, F16, F32, F16, I256, I256, I64, I32, I32, I16, Interwave, Mem, SkipLDS, UseLDS>, std::tuple< Row, Col, Row, F16, F16, F32, F16, I256, I256, I64, I32, I32, I16, Interwave, Mem, SkipLDS, SkipLDS>,