From f09dc1f341d1be54a0f70551ed3ebd84c16e24ae Mon Sep 17 00:00:00 2001 From: carlushuang Date: Thu, 7 Nov 2024 00:24:00 +0800 Subject: [PATCH] compiler ok --- example/ck_tile/15_fused_moe/CMakeLists.txt | 2 + include/ck_tile/core.hpp | 1 + .../core/arch/amd_buffer_addressing.hpp | 5 + .../ck_tile/core/utility/static_counter.hpp | 116 +++++++++++ .../host/reference/reference_fused_moe.hpp | 4 +- .../fused_moegemm_pipeline_flatmm.hpp | 186 +++++++++--------- .../fused_moegemm_pipeline_flatmm_policy.hpp | 36 +++- .../pipeline/fused_moegemm_traits.hpp | 11 ++ .../gemm/warp/warp_gemm_attribute_mfma.hpp | 14 ++ .../warp/warp_gemm_attribute_mfma_impl.hpp | 2 +- .../ck_tile/ops/gemm/warp/warp_gemm_impl.hpp | 5 + 11 files changed, 280 insertions(+), 102 deletions(-) create mode 100644 include/ck_tile/core/utility/static_counter.hpp diff --git a/example/ck_tile/15_fused_moe/CMakeLists.txt b/example/ck_tile/15_fused_moe/CMakeLists.txt index 24ba06bc7f..8275ba0ce0 100644 --- a/example/ck_tile/15_fused_moe/CMakeLists.txt +++ b/example/ck_tile/15_fused_moe/CMakeLists.txt @@ -11,5 +11,7 @@ set(TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS) # NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal) +list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS -DCK_TILE_BUFFER_LOAD_AGPR=1) # TODO: enable load to a +list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker) target_compile_options(${TILE_EXAPMLE_FUSED_MOE} PRIVATE ${TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS}) diff --git a/include/ck_tile/core.hpp b/include/ck_tile/core.hpp index 3b198502d0..e05509d44f 100644 --- a/include/ck_tile/core.hpp +++ b/include/ck_tile/core.hpp @@ -62,6 +62,7 @@ #include "ck_tile/core/utility/philox_rand.hpp" #include "ck_tile/core/utility/random.hpp" #include "ck_tile/core/utility/reduce_operator.hpp" +#include "ck_tile/core/utility/static_counter.hpp" #include "ck_tile/core/utility/to_sequence.hpp" #include "ck_tile/core/utility/transpose_vectors.hpp" #include "ck_tile/core/utility/type_traits.hpp" diff --git a/include/ck_tile/core/arch/amd_buffer_addressing.hpp b/include/ck_tile/core/arch/amd_buffer_addressing.hpp index fed8da77d5..f55073fb2e 100644 --- a/include/ck_tile/core/arch/amd_buffer_addressing.hpp +++ b/include/ck_tile/core/arch/amd_buffer_addressing.hpp @@ -888,6 +888,11 @@ CK_TILE_DEVICE void buffer_store_fence(index_t cnt = 0) asm volatile("s_waitcnt vmcnt(%0)" : : "n"(cnt) : "memory"); } +CK_TILE_DEVICE auto async_load_fence_raw(index_t cnt = 0) +{ + asm volatile("s_waitcnt vmcnt(%0)" : : "n"(cnt) : "memory"); +} + // buffer load i8 CK_TILE_DEVICE_EXTERN int8_t llvm_amdgcn_raw_buffer_load_i8(int32x4_t srsrc, diff --git a/include/ck_tile/core/utility/static_counter.hpp b/include/ck_tile/core/utility/static_counter.hpp new file mode 100644 index 0000000000..84af3dd52f --- /dev/null +++ b/include/ck_tile/core/utility/static_counter.hpp @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core/config.hpp" + +namespace ck_tile { + +template +struct static_counter +{ + public: + template + static constexpr index_t next() + { + return next(0) * Step + Start; + } + + template + static constexpr index_t next() + { + struct Unique + { + }; + return next(0) * Step + Start; + } + + template + static constexpr index_t current() + { + return current(0) * Step + Start; + } + + template + static constexpr index_t current() + { + struct Unique + { + }; + return current(0) * Step + Start; + } + + private: + template + struct slot + { + _Pragma("GCC diagnostic push"); + _Pragma("GCC diagnostic ignored \"-Wundefined-internal\""); + friend constexpr bool slot_allocated(slot); + _Pragma("GCC diagnostic pop"); + }; + + template + struct allocate_slot + { + friend constexpr bool slot_allocated(slot) { return true; } + enum + { + value = I + }; + }; + + // If slot_allocated(slot) has NOT been defined, then SFINAE will keep this function out of + // the overload set... + template ())> + static constexpr index_t next(index_t) + { + return next(0); + } + + // ...And this function will be used, instead, which will define slot_allocated(slot) via + // allocate_slot. + template + static constexpr index_t next(double) + { + return allocate_slot::value; + } + + // If slot_allocated(slot) has NOT been defined, then SFINAE will keep this function out of + // the overload set... + template ())> + static constexpr index_t current(index_t) + { + return current(0); + } + + // ...And this function will be used, instead, which will return the current counter, or assert + // in case next() hasn't been called yet. + template + static constexpr index_t current(double) + { + static_assert(I != 0, "You must invoke next() first"); + + return I - 1; + } +}; + +namespace impl { +template +struct static_counter_uniq_; +} + +#define MAKE_SC() \ + ck_tile::static_counter> {} +#define MAKE_SC_WITH(start_, step_) \ + ck_tile::static_counter, start_, step_> {} +#define NEXT_SC(c_) c_.next<__COUNTER__>() +#define NEXT_SCI(c_, static_i_) c_.next<__COUNTER__ + static_i_>() + +// Usage: +// constexpr auto c = MAKE_SC() +// NEXT_SC(c) // -> constexpr 0 +// NEXT_SC(c) // -> constexpr 1 +// NEXT_SC(c) // -> constexpr 2 +} // namespace ck_tile diff --git a/include/ck_tile/host/reference/reference_fused_moe.hpp b/include/ck_tile/host/reference/reference_fused_moe.hpp index 3374e7d6aa..e4a1559edd 100644 --- a/include/ck_tile/host/reference/reference_fused_moe.hpp +++ b/include/ck_tile/host/reference/reference_fused_moe.hpp @@ -97,7 +97,7 @@ void reference_fused_moe( int max_num_tokens_padded = topk * tokens + experts * (block_m - 1); // assert(); auto f = [&](auto i_flatten) { - ck_tile::index_t i_tile = i_flatten / block_m; + ck_tile::index_t i_tile = i_flatten / block_m; if(i_tile >= num_sorted_tiles) return; ck_tile::index_t i_expert = sorted_expert_ids_host.mData[i_tile]; @@ -136,7 +136,7 @@ void reference_fused_moe( { AccDataType tmp; Activation{}(tmp, acc_0(0, i_n)); - y(0, i_n) = tmp * acc_0(0, i_n + hidden_size); // TODO: elementwise mul + y(0, i_n) = tmp * acc_0(0, i_n + hidden_size); // TODO: elementwise mul } } diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm.hpp index 186e27d26b..1ec6b19527 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm.hpp @@ -156,10 +156,10 @@ struct FusedMoeGemmPipeline_Flatmm using g_thread_type = decltype(load_tile(g_win)); using d_thread_type = decltype(load_tile(d_win)); - // using WarpGemm0 = Policy::template GetWarpGemm0(); - // using WarpGemm1 = Policy::template GetWarpGemm1(); - // auto warp_gemm_0 = WarpGemm0{}; - // auto warp_gemm_1 = WarpGemm1{}; + using WarpGemm0 = decltype(Policy::template GetWarpGemm0()); + using WarpGemm1 = decltype(Policy::template GetWarpGemm1()); + auto warp_gemm_0 = WarpGemm0{}; + auto warp_gemm_1 = WarpGemm1{}; // issues_warps_lanes auto a_sst_win0 = @@ -175,7 +175,7 @@ struct FusedMoeGemmPipeline_Flatmm {0, 0, 0}); // m*k auto a_sld_win0 = [&]() { - using WG = decltype(Policy::template GetWarpGemm0()); + using WG = WarpGemm0; constexpr auto a_outer_dstr_enc = tile_distribution_encoding< sequence<>, tuple, @@ -196,7 +196,7 @@ struct FusedMoeGemmPipeline_Flatmm // m*k auto a_sld_win1 = [&]() { - using WG = decltype(Policy::template GetWarpGemm0()); + using WG = WarpGemm0; constexpr auto a_outer_dstr_enc = tile_distribution_encoding< sequence<>, tuple, @@ -242,10 +242,12 @@ struct FusedMoeGemmPipeline_Flatmm constexpr auto issues_d = number{}; constexpr auto issues_o = number{}; constexpr auto issues_gemm0 = - number{}; + number{}; constexpr auto issues_gemm1 = - number{}; - constexpr auto issues_sld_a = number{}; + number{}; + // constexpr auto issues_sld_a = number{}; const index_t num_blocks_k0 = (hidden_size + BlockShape::Block_K0 - 1) / BlockShape::Block_K0; @@ -284,11 +286,9 @@ struct FusedMoeGemmPipeline_Flatmm } load_tile_raw(g_, g_win, i_access, FALSE, PreNop{}); }; - auto move_g = - [&]() { - move_tile_window(g_win, - {number<0>{}, number{}, number<0>{}}); - }; + auto move_g = [&]() { + move_tile_window(g_win, {number<0>{}, number{}, number<0>{}}); + }; statically_indexed_array ds; auto gld_d = [&]>( @@ -314,16 +314,17 @@ struct FusedMoeGemmPipeline_Flatmm // clang-format off auto gemm_0 = [&]> (auto& t_c, auto& t_a, auto& t_b, auto i_access, PostNop = {}) { - auto warp_gemm = Policy::template GetWarpGemm0(); - using WarpGemm = remove_cvref_t; + using WarpGemm = remove_cvref_t; + constexpr auto repeat_sub = WarpGemm::get_num_of_access(); constexpr auto repeat_m = BlockShape::Repeat_M0; // constexpr auto repeat_n = BlockShape::Repeat_N0; constexpr auto repeat_k = BlockShape::Repeat_K0; // loop order n->m->k - constexpr auto i_k = i_access % repeat_k; - constexpr auto i_m = (i_access / repeat_k) % repeat_m; - constexpr auto i_n = (i_access / repeat_k) / repeat_m; + constexpr auto i_sub = i_access % repeat_sub; + constexpr auto i_k = (i_access / repeat_sub) % repeat_k; + constexpr auto i_m = (i_access / (repeat_sub * repeat_k )) % repeat_m; + constexpr auto i_n = (i_access / (repeat_sub * repeat_k )) / repeat_m; using AWarpTensor = typename WarpGemm::AWarpTensor; using BWarpTensor = typename WarpGemm::BWarpTensor; @@ -355,7 +356,7 @@ struct FusedMoeGemmPipeline_Flatmm merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - WarpGemm{}(w_c, w_a, w_b, PostNop{}); + warp_gemm_0(w_c, w_a, w_b, number{}, PostNop{}); t_c.set_y_sliced_thread_data( merge_sequences(sequence{}, c_warp_y_index_zeros), @@ -367,16 +368,17 @@ struct FusedMoeGemmPipeline_Flatmm // clang-format off auto gemm_1 = [&]> (auto& t_c, auto& t_a, auto& t_b, auto i_access, PostNop = {}) { - auto warp_gemm = Policy::template GetWarpGemm1(); - using WarpGemm = remove_cvref_t; + using WarpGemm = remove_cvref_t; - constexpr auto repeat_m = BlockShape::Repeat_M1; - // constexpr auto repeat_n = BlockShape::Repeat_N1; - constexpr auto repeat_k = BlockShape::Repeat_K1; + constexpr auto repeat_sub = WarpGemm::get_num_of_access(); + constexpr auto repeat_m = BlockShape::Repeat_M0; + // constexpr auto repeat_n = BlockShape::Repeat_N0; + constexpr auto repeat_k = BlockShape::Repeat_K0; // loop order n->m->k - constexpr auto i_k = i_access % repeat_k; - constexpr auto i_m = (i_access / repeat_k) % repeat_m; - constexpr auto i_n = (i_access / repeat_k) / repeat_m; + constexpr auto i_sub = i_access % repeat_sub; + constexpr auto i_k = (i_access / repeat_sub) % repeat_k; + constexpr auto i_m = (i_access / (repeat_sub * repeat_k )) % repeat_m; + constexpr auto i_n = (i_access / (repeat_sub * repeat_k )) / repeat_m; using AWarpTensor = typename WarpGemm::AWarpTensor; using BWarpTensor = typename WarpGemm::BWarpTensor; @@ -408,7 +410,7 @@ struct FusedMoeGemmPipeline_Flatmm merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - WarpGemm{}(w_c, w_a, w_b, PostNop{}); + warp_gemm_1(w_c, w_a, w_b, number{}, PostNop{}); t_c.set_y_sliced_thread_data( merge_sequences(sequence{}, c_warp_y_index_zeros), @@ -416,84 +418,72 @@ struct FusedMoeGemmPipeline_Flatmm w_c.get_thread_buffer()); }; // clang-format on - _Pragma("clang diagnostic pop") + _Pragma("clang diagnostic pop"); - // this gemm pipeline is designed with assumption that issues of buffer-load/ds_read can - // be hide under mfma. In other words, issues of mfma is >= memory this is true if we - // pre-shuffle B matrix, and A matrix is relatively small we prefer use multiple mfma - // paired with 1 buffer-load B matrix, to get max throughput of buffer_load. and by - // preshuffle, we always pack to dwordx4 load, and this will already extend to multiple - // mfma but that is already consumed inside warpgemm-impl. So indeed how many extra - // mfma(that can reuse the B matrix) only affected by M repeat. - auto pipeline_gemm0 = [&]() { - constexpr index_t total_loops = issues_gemm0; - constexpr index_t mfma_per_ld = total_loops / (issues_g + issues_a + issues_sld_a); + // this gemm pipeline is designed with assumption that issues of buffer-load/ds_read can + // be hide under mfma. In other words, issues of mfma is >= memory this is true if we + // pre-shuffle B matrix, and A matrix is relatively small we prefer use multiple mfma + // paired with 1 buffer-load B matrix, to get max throughput of buffer_load. and by + // preshuffle, we always pack to dwordx4 load, and this will already extend to multiple + // mfma but that is already consumed inside warpgemm-impl. So indeed how many extra + // mfma(that can reuse the B matrix) only affected by M repeat. + auto pipeline_gemm0 = [&]() { + constexpr index_t total_loops = issues_gemm0; + constexpr auto sr = Policy::template GetSequencer_0(); + static_assert(sr.size() == total_loops); + constexpr index_t SLD_A = + static_cast(FusedMoeGemmPipelineSequencerEnum::SLD_A); + constexpr index_t GLD_A = + static_cast(FusedMoeGemmPipelineSequencerEnum::GLD_A); + constexpr index_t GLD_B = + static_cast(FusedMoeGemmPipelineSequencerEnum::GLD_B); + constexpr auto c_sld_a_0 = MAKE_SC(); + constexpr auto c_gld_a_0 = MAKE_SC(); + constexpr auto c_gld_b_0 = MAKE_SC(); + // compute buffer 1 + static_for<0, total_loops, 1>{}([&](auto i_issue) { + gemm_0(acc_0, as[I0], gs[I0], i_issue); + constexpr index_t slot = sr.at(i_issue); - // compute buffer 0 - static_for<0, total_loops, 1>{}([&](auto i_issue) { - gemm_0(acc_0, as[I0], gs[I0], i_issue); + if constexpr(slot & SLD_A) + sld_a(as[I1], a_sld_win1, number{}); + if constexpr(slot & GLD_A) + gld_a(a_sst_win0, number{}); + if constexpr(slot & GLD_B) + gld_g(gs[I0], number{}); + }); + move_g(); + move_a(); + block_sync_load_raw(issues_a + issues_g); + lds_load_fence(); - if constexpr(i_issue % mfma_per_ld == 0) - { - constexpr index_t ld_id = 0; + constexpr auto c_sld_a_1 = MAKE_SC(); + constexpr auto c_gld_a_1 = MAKE_SC(); + constexpr auto c_gld_b_1 = MAKE_SC(); - if constexpr(ld_id < issues_g) - { - gld_g(gs[I0], number{}); - } - if constexpr(ld_id - issues_g < +issues_a) - { - gld_a(a_sst_win0, number{}); - } - if constexpr(ld_id - issues_g - issues_a < issues_sld_a) - { - sld_a(as[I1], a_sld_win1, number{}); - } + // compute buffer 1 + static_for<0, total_loops, 1>{}([&](auto i_issue) { + gemm_0(acc_0, as[I1], gs[I1], i_issue); + constexpr index_t slot = sr.at(i_issue); - ld_id++; - } - - }); - move_g(); - move_a(); - block_sync_load_raw(issues_a + issues_g); - lds_load_fence(); - - // compute buffer 1 - static_for<0, total_loops, 1>{}([&](auto i_issue) { - gemm_0(acc_0, as[I1], gs[I1], i_issue); - - if constexpr(i_issue % mfma_per_ld == 0) - { - constexpr index_t ld_id = 0; - - if constexpr(ld_id < issues_g) - { - gld_g(gs[I1], number{}); - } - if constexpr(ld_id - issues_g < +issues_a) - { - gld_a(a_sst_win1, number{}); - } - if constexpr(ld_id - issues_g - issues_a < issues_sld_a) - { - sld_a(as[I0], a_sld_win0, number{}); - } - - ld_id++; - } - }); - move_g(); - move_a(); - block_sync_load_raw(issues_a + issues_g); - lds_load_fence(); - }; + if constexpr(slot & SLD_A) + sld_a(as[I0], a_sld_win0, number{}); + if constexpr(slot & GLD_A) + gld_a(a_sst_win1, number{}); + if constexpr(slot & GLD_B) + gld_g(gs[I1], number{}); + }); + move_g(); + move_a(); + block_sync_load_raw(issues_a + issues_g); + lds_load_fence(); + }; auto pipeline_gemm0_tail = [&]() { constexpr index_t total_loops = issues_gemm0; constexpr index_t mfma_per_gld_g = total_loops / issues_g; // BlockShape::Repeat_M0; // constexpr index_t mfma_per_gld_a = total_loops / issues_a; - constexpr index_t mfma_per_sld_a = total_loops / issues_sld_a; + // constexpr index_t mfma_per_sld_a = total_loops / issues_sld_a; // compute buffer 0 static_for<0, total_loops, 1>{}([&](auto i_issue) { @@ -515,7 +505,7 @@ struct FusedMoeGemmPipeline_Flatmm }); // if cycle_mfma>gld_a sync here block_sync_load_raw(issues_g); - sld_a(as[I1], a_sld_win1, NEG1{}); + sld_a(as[I1], a_sld_win1, NEG1); // compute buffer 1 static_for<0, total_loops, 1>{}([&](auto i_issue) { diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp index 3885ad0ffd..8673e03692 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp @@ -609,11 +609,45 @@ struct FusedMoeGemmPipelineFlatmmPolicy } } + template + CK_TILE_HOST_DEVICE static constexpr auto GetSequencer_0() + { + // this function return seq<...> used to identify gld/sld/valu... inside mfma sequence + // the purpose is to hide thoes instructions under mfma + // every value inside seq<...> is a mask, indicating a specific operation + using S_ = typename Problem::BlockShape; + constexpr index_t SLD_A = static_cast(FusedMoeGemmPipelineSequencerEnum::SLD_A); + constexpr index_t GLD_A = static_cast(FusedMoeGemmPipelineSequencerEnum::GLD_A); + constexpr index_t GLD_B = static_cast(FusedMoeGemmPipelineSequencerEnum::GLD_B); + if constexpr(std::is_same_v && + std::is_same_v && + S_::Warp_M0 == 32 && S_::Warp_N0 == 32 && S_::Warp_K0 == 16 && + S_::Block_M0 == 32 && S_::Block_N0 == 512 && S_::Block_K0 == 128 && + S_::Block_N1 == 128) + { + // Total 64 instructions, 32 buffer-load-dwordx4 gld_b, 8x buffer-load-dwordx1-async + // gld_a 8x ds_read_b128 sld_a total 64 slot :) + // clang-format off + constexpr auto seq_all = + // 0 1 2 3 4 5 6 7 + sequence{}; // 7 + return seq_all; + // clang-format on + } + } + template CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemm1() { using S_ = typename Problem::BlockShape; - constexpr auto wg_ctrl = WGAttrCtlEnum::Raw_vva; + constexpr auto wg_ctrl = WGAttrCtlEnum::Raw_vav; // TODO: ugly if constexpr(std::is_same_v && std::is_same_v && diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp index 6bfc0ec854..b9eaffd19a 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp @@ -33,4 +33,15 @@ struct FusedMoeGemmTraits static constexpr bool PadHiddenSize = PadHiddenSize_; static constexpr bool PadIntermediateSize = PadIntermediateSize_; }; + +// Note: this need to be a bit mask +enum class FusedMoeGemmPipelineSequencerEnum +{ + SLD_A = 1 << 0, // shared load a + SLD_B = 1 << 1, + GLD_A = 1 << 2, // global load a + GLD_B = 1 << 3, + SST_A = 1 << 4, // shared store a + SST_B = 1 << 5, +}; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp index 3c687c1f14..0a8d2dfbe3 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp @@ -25,6 +25,8 @@ struct WarpGemmAtrributeMfma static constexpr index_t kN = Impl::kN; static constexpr index_t kK = Impl::kK; + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; } + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -88,6 +90,8 @@ struct WarpGemmAtrributeMfmaIterateK static constexpr index_t kN = Impl::kN; static constexpr index_t kK = Impl::kK * kKIter; + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -197,6 +201,8 @@ struct WarpGemmAtrributeMfmaTransposedCDistribution static constexpr index_t kN = Impl::kM; static constexpr index_t kK = Impl::kK; + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; } + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -258,6 +264,8 @@ struct WarpGemmAtrributeMfmaTransposedCDistribution_SwizzleB static constexpr index_t kN = Impl::kM; static constexpr index_t kK = Impl::kK; + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; } + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -326,6 +334,8 @@ struct WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution static constexpr index_t kN = Impl::kM; static constexpr index_t kK = Impl::kK * kKIter; + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -439,6 +449,8 @@ struct WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB static constexpr index_t kK = Impl::kK * kKIter; static constexpr index_t SFactor = SFactor_; // group how many CM1 together + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple, sequence>, @@ -576,6 +588,8 @@ struct WarpGemmAtrributeMfmaIterateK_SwizzleA static constexpr index_t kK = Impl::kK * kKIter; static constexpr index_t SFactor = SFactor_; // group how many CM1 together + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; } + using AWarpDstrEncoding = tile_distribution_encoding< sequence<>, tuple; using CWarpTensor = static_distributed_tensor; + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() + { + return WarpGemmAttribute_::get_num_of_access(); + } + template CK_TILE_DEVICE void operator()(CTensor& c, const ATensor& a, const BTensor& b, bool_constant = {}) const