ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-14 02:02:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use __builtin_amdgcn_readfirstlane for buffer resource in fused_moe (#2893)

Browse Source 
* 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 e17b20625e
commit 18f4a0728b
3 changed files with 128 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
include/ck_tile/core/arch/amd_buffer_addressing.hpp
View File

@@ -28,6 +28,60 @@ using as3_uint32_ptr = uint32_t __attribute__((address_space(3)))*;
namespace ck_tile {
// amd_wave_read_first_lane is the SGPR function from AMD GPU device to load 1 or a series of the
// memory to the SGPR registers.
__device__ inline uint32_t amd_wave_read_first_lane(uint16_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint8_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
__device__ inline int32_t amd_wave_read_first_lane(int32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
template <typename Object, std::enable_if_t<std::is_trivially_copyable_v<Object>, int> = 0>
__device__ inline auto amd_wave_read_first_lane(const Object& obj)
{
constexpr size_t ObjectSize = sizeof(Object);
constexpr size_t SGPR_size = 4;
constexpr size_t NumFull = ObjectSize / SGPR_size;
constexpr size_t Tail = ObjectSize % SGPR_size;
const unsigned char* src = reinterpret_cast<const unsigned char*>(&obj);
alignas(Object) unsigned char dst[ObjectSize];
static_for<0, NumFull, 1>{}([&](auto Ic) {
constexpr size_t offset = Ic * SGPR_size;
uint32_t read_src;
__builtin_memcpy(&read_src, src + offset, SGPR_size);
read_src = __builtin_amdgcn_readfirstlane(read_src);
__builtin_memcpy(dst + offset, &read_src, SGPR_size);
});
if constexpr(Tail != 0)
{
constexpr size_t offset = NumFull * SGPR_size;
uint32_t tail_loc = 0;
__builtin_memcpy(&tail_loc, src + offset, Tail);
tail_loc = __builtin_amdgcn_readfirstlane(tail_loc);
__builtin_memcpy(dst + offset, &tail_loc, Tail);
}
Object out;
__builtin_memcpy(&out, dst, ObjectSize);
return out;
}
// 128 bit SGPRs to supply buffer resource in buffer instructions
// https://rocm-documentation.readthedocs.io/en/latest/GCN_ISA_Manuals/testdocbook.html#vector-memory-buffer-instructions
struct __attribute__((packed)) buffer_resource
@@ -37,10 +91,17 @@ struct __attribute__((packed)) buffer_resource
uint32_t config;
};
CK_TILE_DEVICE int32x4_t make_wave_buffer_resource(const void* ptr, uint32_t size = 0xffffffff)
template <typename ForceSGPR = std::false_type>
CK_TILE_DEVICE int32x4_t make_wave_buffer_resource(const void* ptr,
uint32_t size = 0xffffffff,
ForceSGPR = {})
{
buffer_resource res{ptr, size, CK_TILE_BUFFER_RESOURCE_3RD_DWORD};
int32x4_t r = __builtin_bit_cast(int32x4_t, res);
if constexpr(std::is_same_v<ForceSGPR, std::true_type>)
{
r = amd_wave_read_first_lane(r);
}
return r;
}
@@ -2829,60 +2890,6 @@ __device__ auto amd_transpose_load_to_vgpr(const T* __restrict__ in_ptr)
}
#endif
// amd_wave_read_first_lane is the SGPR function from AMD GPU device to load 1 or a series of the
// memory to the SGPR registers.
__device__ inline uint32_t amd_wave_read_first_lane(uint16_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint8_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
__device__ inline int32_t amd_wave_read_first_lane(int32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
template <typename Object, std::enable_if_t<std::is_trivially_copyable_v<Object>, int> = 0>
__device__ inline auto amd_wave_read_first_lane(const Object& obj)
{
constexpr size_t ObjectSize = sizeof(Object);
constexpr size_t SGPR_size = 4;
constexpr size_t NumFull = ObjectSize / SGPR_size;
constexpr size_t Tail = ObjectSize % SGPR_size;
const unsigned char* src = reinterpret_cast<const unsigned char*>(&obj);
alignas(Object) unsigned char dst[ObjectSize];
static_for<0, NumFull, 1>{}([&](auto Ic) {
constexpr size_t offset = Ic * SGPR_size;
uint32_t read_src;
__builtin_memcpy(&read_src, src + offset, SGPR_size);
read_src = __builtin_amdgcn_readfirstlane(read_src);
__builtin_memcpy(dst + offset, &read_src, SGPR_size);
});
if constexpr(Tail != 0)
{
constexpr size_t offset = NumFull * SGPR_size;
uint32_t tail_loc = 0;
__builtin_memcpy(&tail_loc, src + offset, Tail);
tail_loc = __builtin_amdgcn_readfirstlane(tail_loc);
__builtin_memcpy(dst + offset, &tail_loc, Tail);
}
Object out;
__builtin_memcpy(&out, dst, ObjectSize);
return out;
}
} // namespace ck_tile
#endif // !CK_TILE_USE_BUFFER_ADDRESSING_BUILTIN

117
include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp
View File

@@ -19,6 +19,60 @@ using as3_uint32_ptr = uint32_t __attribute__((address_space(3)))*;
namespace ck_tile {
// amd_wave_read_first_lane is the SGPR function from AMD GPU device to load 1 or a series of the
// memory to the SGPR registers.
__device__ inline uint32_t amd_wave_read_first_lane(uint16_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint8_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
__device__ inline int32_t amd_wave_read_first_lane(int32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
template <typename Object, std::enable_if_t<std::is_trivially_copyable_v<Object>, int> = 0>
__device__ inline auto amd_wave_read_first_lane(const Object& obj)
{
constexpr size_t ObjectSize = sizeof(Object);
constexpr size_t SGPR_size = 4;
constexpr size_t NumFull = ObjectSize / SGPR_size;
constexpr size_t Tail = ObjectSize % SGPR_size;
const unsigned char* src = reinterpret_cast<const unsigned char*>(&obj);
alignas(Object) unsigned char dst[ObjectSize];
static_for<0, NumFull, 1>{}([&](auto Ic) {
constexpr size_t offset = Ic * SGPR_size;
uint32_t read_src;
__builtin_memcpy(&read_src, src + offset, SGPR_size);
read_src = __builtin_amdgcn_readfirstlane(read_src);
__builtin_memcpy(dst + offset, &read_src, SGPR_size);
});
if constexpr(Tail != 0)
{
constexpr size_t offset = NumFull * SGPR_size;
uint32_t tail_loc = 0;
__builtin_memcpy(&tail_loc, src + offset, Tail);
tail_loc = __builtin_amdgcn_readfirstlane(tail_loc);
__builtin_memcpy(dst + offset, &tail_loc, Tail);
}
Object out;
__builtin_memcpy(&out, dst, ObjectSize);
return out;
}
// 128 bit SGPRs to supply buffer resource in buffer instructions
// https://rocm-documentation.readthedocs.io/en/latest/GCN_ISA_Manuals/testdocbook.html#vector-memory-buffer-instructions
struct __attribute__((packed)) buffer_resource
@@ -28,10 +82,17 @@ struct __attribute__((packed)) buffer_resource
uint32_t config;
};
CK_TILE_DEVICE int32x4_t make_wave_buffer_resource(const void* ptr, uint32_t size = 0xffffffff)
template <typename ForceSGPR = std::false_type>
CK_TILE_DEVICE int32x4_t make_wave_buffer_resource(const void* ptr,
uint32_t size = 0xffffffff,
ForceSGPR = {})
{
buffer_resource res{ptr, size, CK_TILE_BUFFER_RESOURCE_3RD_DWORD};
int32x4_t r = __builtin_bit_cast(int32x4_t, res);
if constexpr(std::is_same_v<ForceSGPR, std::true_type>)
{
r = amd_wave_read_first_lane(r);
}
return r;
}
@@ -2570,60 +2631,6 @@ CK_TILE_DEVICE void amd_buffer_atomic_max(const thread_buffer<T, N>& src_thread_
#endif
}
// amd_wave_read_first_lane is the SGPR function from AMD GPU device to load 1 or a series of the
// memory to the SGPR registers.
__device__ inline uint32_t amd_wave_read_first_lane(uint16_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint8_t v)
{
return __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(v));
}
__device__ inline uint32_t amd_wave_read_first_lane(uint32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
__device__ inline int32_t amd_wave_read_first_lane(int32_t value)
{
return __builtin_amdgcn_readfirstlane(value);
}
template <typename Object, std::enable_if_t<std::is_trivially_copyable_v<Object>, int> = 0>
__device__ inline auto amd_wave_read_first_lane(const Object& obj)
{
constexpr size_t ObjectSize = sizeof(Object);
constexpr size_t SGPR_size = 4;
constexpr size_t NumFull = ObjectSize / SGPR_size;
constexpr size_t Tail = ObjectSize % SGPR_size;
const unsigned char* src = reinterpret_cast<const unsigned char*>(&obj);
alignas(Object) unsigned char dst[ObjectSize];
static_for<0, NumFull, 1>{}([&](auto Ic) {
constexpr size_t offset = Ic * SGPR_size;
uint32_t read_src;
__builtin_memcpy(&read_src, src + offset, SGPR_size);
read_src = __builtin_amdgcn_readfirstlane(read_src);
__builtin_memcpy(dst + offset, &read_src, SGPR_size);
});
if constexpr(Tail != 0)
{
constexpr size_t offset = NumFull * SGPR_size;
uint32_t tail_loc = 0;
__builtin_memcpy(&tail_loc, src + offset, Tail);
tail_loc = __builtin_amdgcn_readfirstlane(tail_loc);
__builtin_memcpy(dst + offset, &tail_loc, Tail);
}
Object out;
__builtin_memcpy(&out, dst, ObjectSize);
return out;
}
template <typename T, index_t NumElemsPerThread>
CK_TILE_DEVICE void amd_direct_load_global_to_lds(const T* global_base_ptr,
const index_t global_offset,

6
include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_uk.hpp
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));