Use __builtin_amdgcn_readfirstlane for buffer resource in fused_moe (#2893)

* Use __builtin_amdgcn_readfirstlane for buffer resource in fused_moe

* also do the same for amd_buffer_addressing_builtins.hpp

* merge with develop

* fix clang format

---------

Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com>
Co-authored-by: ThomasNing <thomas.ning@amd.com>
Co-authored-by: illsilin_amdeng <Illia.Silin@amd.com>

[ROCm/composable_kernel commit: ef43078788]
This commit is contained in:
Sami Remes
2025-10-01 01:12:30 +03:00
committed by GitHub
parent 7dcf623fcb
commit 93ba707be4
3 changed files with 128 additions and 112 deletions

View File

@@ -210,7 +210,8 @@ struct FusedMoeGemmPipeline_FlatmmUk
auto a_res =
make_wave_buffer_resource(reinterpret_cast<const ADataType*>(kargs.a_ptr),
kargs.num_tokens * kargs.stride_token * sizeof(ADataType));
kargs.num_tokens * kargs.stride_token * sizeof(ADataType),
std::true_type{});
auto make_gu_win = [&](const auto* ptr_) {
auto view_ = make_naive_tensor_view<address_space_enum::global>(
@@ -322,7 +323,8 @@ struct FusedMoeGemmPipeline_FlatmmUk
auto o_res =
make_wave_buffer_resource(reinterpret_cast<const ODataType*>(kargs.o_ptr),
kargs.num_tokens * kargs.stride_token * sizeof(ODataType));
kargs.num_tokens * kargs.stride_token * sizeof(ODataType),
std::true_type{});
auto row_coords_o = GetRowCoords_O(sorted_tile_id * BlockShape::Block_M0);
auto w_scale = GetWeightScale(
row_coords_o, reinterpret_cast<const TopkWeightDataType*>(kargs.sorted_weight_ptr));