From 76201e9af5f748239a961ac9de2d181f277faba8 Mon Sep 17 00:00:00 2001 From: lalala-sh Date: Mon, 13 Oct 2025 02:52:46 +0000 Subject: [PATCH] enable atomic_add_bf16 in gfx950 --- example/ck_tile/03_gemm/run_gemm_example.inc | 10 +- .../core/arch/amd_buffer_addressing.hpp | 48 ++++- .../arch/amd_buffer_addressing_builtins.hpp | 49 ++++- include/ck_tile/core/numeric/vector_type.hpp | 24 ++- include/ck_tile/core/tensor/buffer_view.hpp | 14 +- .../core/tensor/tile_scatter_gather.hpp | 176 ++++++++++++++++++ 6 files changed, 310 insertions(+), 11 deletions(-) diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index 2abc66bc4a..754136c49a 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -195,6 +195,11 @@ auto shuffle_b(const ck_tile::HostTensor& t) } else { + assert(is_wave32() == false); + constexpr int MaxVecSize = 16 / sizeof(T); + constexpr int KLane = ck_tile::get_warp_size() / GemmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = std::min(MaxVecSize, GemmConfig::K_Warp_Tile / KLane); + if(ck_tile::is_gfx11_supported()) { int divisor = 1; @@ -208,11 +213,6 @@ auto shuffle_b(const ck_tile::HostTensor& t) } else { - assert(is_wave32() == false); - constexpr int MaxVecSize = 16 / sizeof(T); - constexpr int KLane = ck_tile::get_warp_size() / GemmConfig::N_Warp_Tile; - constexpr int ItemsPerAccess = std::min(MaxVecSize, GemmConfig::K_Warp_Tile / KLane); - ck_tile::HostTensor t_view({n_ / GemmConfig::N_Warp_Tile, GemmConfig::N_Warp_Tile, k_ / ItemsPerAccess, diff --git a/include/ck_tile/core/arch/amd_buffer_addressing.hpp b/include/ck_tile/core/arch/amd_buffer_addressing.hpp index 7a9c017eb2..a7b9790510 100644 --- a/include/ck_tile/core/arch/amd_buffer_addressing.hpp +++ b/include/ck_tile/core/arch/amd_buffer_addressing.hpp @@ -1242,6 +1242,15 @@ CK_TILE_DEVICE_EXTERN fp16x2_t llvm_amdgcn_raw_buffer_atomic_add_fp16x2( index_t soffset, index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fadd.v2f16"); +// buffer atomic-add bf16 +// TODO: Replace with bf16x2_t, but llvm builins only accept cktile_llvm_bf16x2_t now. +CK_TILE_DEVICE_EXTERN llvm_bf16x2_t llvm_amdgcn_raw_buffer_atomic_add_bf16x2( + llvm_bf16x2_t vdata, + int32x4_t rsrc, + index_t voffset, + index_t soffset, + index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fadd.v2bf16"); + // buffer atomic-add i32 CK_TILE_DEVICE_EXTERN int32_t llvm_amdgcn_raw_buffer_atomic_add_i32( int32_t vdata, @@ -1476,8 +1485,11 @@ CK_TILE_DEVICE thread_buffer amd_buffer_load_impl(int32x4_t src_wave_buffe (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || + (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (std::is_same::value && - (N == 1 || N == 2 || N == 4 || N == 8 || N == 16 || N == 32)), + (N == 1 || N == 2 || N == 4 || N == 8 || N == 16 || N == 32) || + (std::is_same::value && + (N == 1 || N == 2 || N == 4 || N == 8 || N == 16))), "wrong! not implemented"); using rtn_type = thread_buffer; @@ -2201,6 +2213,7 @@ CK_TILE_DEVICE void amd_buffer_atomic_add_impl(const thread_buffer& src_th { static_assert((std::is_same::value && (N == 1 || N == 2 || N == 4)) || (std::is_same::value && (N == 2 || N == 4 || N == 8)) || + (std::is_same::value && (N == 2 || N == 4 || N == 8)) || (std::is_same::value && (N == 1 || N == 2 || N == 4)), "wrong! not implemented"); @@ -2294,6 +2307,39 @@ CK_TILE_DEVICE void amd_buffer_atomic_add_impl(const thread_buffer& src_th }); } } + else if constexpr(std::is_same::value) + { + if constexpr(N == 2) + { + llvm_amdgcn_raw_buffer_atomic_add_bf16x2(bit_cast(src_thread_data), + dst_wave_buffer_resource, + dst_thread_addr_offset, + dst_wave_addr_offset, + 0); + } + else if constexpr(N == 4) + { + static_for<0, 2, 1>{}([&](auto i) { + llvm_amdgcn_raw_buffer_atomic_add_bf16x2( + src_thread_data.template get_as()[i], + dst_wave_buffer_resource, + dst_thread_addr_offset, + dst_wave_addr_offset + i * sizeof(bf16x2_t), + 0); + }); + } + else if constexpr(N == 8) + { + static_for<0, 4, 1>{}([&](auto i) { + llvm_amdgcn_raw_buffer_atomic_add_bf16x2( + src_thread_data.template get_as()[i], + dst_wave_buffer_resource, + dst_thread_addr_offset, + dst_wave_addr_offset + i * sizeof(bf16x2_t), + 0); + }); + } + } else if constexpr(std::is_same::value) { if constexpr(N == 1) diff --git a/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp b/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp index 4e0a86119a..8acfaebdbc 100644 --- a/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp +++ b/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp @@ -1110,6 +1110,15 @@ CK_TILE_DEVICE_EXTERN fp16x2_t llvm_amdgcn_raw_buffer_atomic_add_fp16x2( index_t soffset, index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fadd.v2f16"); +// buffer atomic-add bf16 +// TODO: Replace with bf16x2_t, but llvm builins only accept cktile_llvm_bf16x2_t now. +CK_TILE_DEVICE_EXTERN llvm_bf16x2_t llvm_amdgcn_raw_buffer_atomic_add_bf16x2( + llvm_bf16x2_t vdata, + int32x4_t rsrc, + index_t voffset, + index_t soffset, + index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fadd.v2bf16"); + // buffer atomic-add i32 CK_TILE_DEVICE_EXTERN int32_t llvm_amdgcn_raw_buffer_atomic_add_i32( int32_t vdata, @@ -1344,8 +1353,12 @@ CK_TILE_DEVICE thread_buffer amd_buffer_load_impl(int32x4_t src_wave_buffe (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (std::is_same::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || + (std::is_same::value && + (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) || (std::is_same::value && - (N == 1 || N == 2 || N == 4 || N == 8 || N == 16 || N == 32)), + (N == 1 || N == 2 || N == 4 || N == 8 || N == 16 || N == 32) || + (std::is_same::value && + (N == 1 || N == 2 || N == 4 || N == 8 || N == 16))), "wrong! not implemented"); using rtn_type = thread_buffer; @@ -1984,6 +1997,7 @@ CK_TILE_DEVICE void amd_buffer_atomic_add_impl(const thread_buffer& src_th { static_assert((std::is_same::value && (N == 1 || N == 2 || N == 4)) || (std::is_same::value && (N == 2 || N == 4 || N == 8)) || + (std::is_same::value && (N == 2 || N == 4 || N == 8)) || (std::is_same::value && (N == 1 || N == 2 || N == 4)), "wrong! not implemented"); @@ -2077,6 +2091,39 @@ CK_TILE_DEVICE void amd_buffer_atomic_add_impl(const thread_buffer& src_th }); } } + else if constexpr(std::is_same::value) + { + if constexpr(N == 2) + { + llvm_amdgcn_raw_buffer_atomic_add_bf16x2(bit_cast(src_thread_data), + dst_wave_buffer_resource, + dst_thread_addr_offset, + dst_wave_addr_offset, + 0); + } + else if constexpr(N == 4) + { + static_for<0, 2, 1>{}([&](auto i) { + llvm_amdgcn_raw_buffer_atomic_add_bf16x2( + src_thread_data.template get_as()[i], + dst_wave_buffer_resource, + dst_thread_addr_offset, + dst_wave_addr_offset + i * sizeof(bf16x2_t), + 0); + }); + } + else if constexpr(N == 8) + { + static_for<0, 4, 1>{}([&](auto i) { + llvm_amdgcn_raw_buffer_atomic_add_bf16x2( + src_thread_data.template get_as()[i], + dst_wave_buffer_resource, + dst_thread_addr_offset, + dst_wave_addr_offset + i * sizeof(bf16x2_t), + 0); + }); + } + } else if constexpr(std::is_same::value) { if constexpr(N == 1) diff --git a/include/ck_tile/core/numeric/vector_type.hpp b/include/ck_tile/core/numeric/vector_type.hpp index 5d8b109901..f594106a9f 100644 --- a/include/ck_tile/core/numeric/vector_type.hpp +++ b/include/ck_tile/core/numeric/vector_type.hpp @@ -11,6 +11,7 @@ #include "ck_tile/core/numeric/half.hpp" #include "ck_tile/core/numeric/bfloat16.hpp" #include "ck_tile/core/numeric/pk_int4.hpp" +#include "ck_tile/core/numeric/e8m0.hpp" #include "ck_tile/core/utility/type_traits.hpp" namespace ck_tile { @@ -88,7 +89,12 @@ template struct vector_traits { using scalar_type = - std::conditional_t, pk_int4_t>, int8_t, remove_cvref_t>; + std::conditional_t, pk_int4_t>, + int8_t, + std::conditional_t, pk_fp4_t> || + std::is_same_v, e8m0_t>, + uint8_t, + remove_cvref_t>>; static constexpr index_t vector_size = 1; }; @@ -96,7 +102,12 @@ struct vector_traits template struct vector_traits { - using scalar_type = std::conditional_t, int8_t, T>; + using scalar_type = std::conditional_t< + std::is_same_v, + int8_t, + std::conditional_t || std::is_same_v, e8m0_t>, + uint8_t, + T>>; static constexpr index_t vector_size = N; }; @@ -138,6 +149,9 @@ using bf16x16_t = bfloat16_t __attribute__((ext_vector_type(16))); using bf16x32_t = bfloat16_t __attribute__((ext_vector_type(32))); using bf16x64_t = bfloat16_t __attribute__((ext_vector_type(64))); +using llvm_bf16x2_t = __bf16 __attribute__((ext_vector_type(2))); +using llvm_bf16x4_t = __bf16 __attribute__((ext_vector_type(4))); + // i32 // using int32_t = ... using int32x2_t = int32_t __attribute__((ext_vector_type(2))); @@ -237,4 +251,10 @@ using pk_int4x4_t = int8_t __attribute__((ext_vector_type(4))); using pk_int4x8_t = int8_t __attribute__((ext_vector_type(8))); using pk_int4x16_t = int8_t __attribute__((ext_vector_type(16))); using pk_int4x32_t = int8_t __attribute__((ext_vector_type(32))); + +using pk_fp4x2_t = uint8_t __attribute((ext_vector_type(2))); +using pk_fp4x4_t = uint8_t __attribute((ext_vector_type(4))); +using pk_fp4x8_t = uint8_t __attribute((ext_vector_type(8))); +using pk_fp4x16_t = uint8_t __attribute((ext_vector_type(16))); +using pk_fp4x32_t = uint8_t __attribute((ext_vector_type(32))); } // namespace ck_tile diff --git a/include/ck_tile/core/tensor/buffer_view.hpp b/include/ck_tile/core/tensor/buffer_view.hpp index d1e770ef42..99637656ce 100644 --- a/include/ck_tile/core/tensor/buffer_view.hpp +++ b/include/ck_tile/core/tensor/buffer_view.hpp @@ -631,14 +631,24 @@ struct buffer_view, int32_t> || std::is_same_v, float> || - (std::is_same_v, half_t> && scalar_per_x_vector % 2 == 0); + (std::is_same_v, half_t> && scalar_per_x_vector % 2 == 0) +#if defined(__gfx950__) // only gfx950 support atomic_pk_add_bf16 + || + (std::is_same_v, bfloat16_t> && scalar_per_x_vector % 2 == 0) +#endif + ; #elif CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER && (!CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT) bool constexpr use_amd_buffer_addressing = std::is_same_v, int32_t>; #elif(!CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER) && CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT bool constexpr use_amd_buffer_addressing = std::is_same_v, float> || - (std::is_same_v, half_t> && scalar_per_x_vector % 2 == 0); + (std::is_same_v, half_t> && scalar_per_x_vector % 2 == 0) +#if defined(__gfx950__) // only gfx950 support atomic_pk_add_bf16 + || + (std::is_same_v, bfloat16_t> && scalar_per_x_vector % 2 == 0) +#endif + ; #else bool constexpr use_amd_buffer_addressing = false; #endif diff --git a/include/ck_tile/core/tensor/tile_scatter_gather.hpp b/include/ck_tile/core/tensor/tile_scatter_gather.hpp index c7811133d6..3b508aa66c 100644 --- a/include/ck_tile/core/tensor/tile_scatter_gather.hpp +++ b/include/ck_tile/core/tensor/tile_scatter_gather.hpp @@ -404,6 +404,100 @@ struct tile_scatter_gather }); } + template + CK_TILE_DEVICE auto async_load(LdsTileWindow_&& lds_tile, + number = {}, + bool_constant = {}) const + { + using LdsTileWindow = remove_cvref_t; + using LdsDataType = typename LdsTileWindow::DataType; + using Traits = load_store_traits; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = TileDstr{}; + + // Precompute invariant values outside loops + const auto window_origin = lds_tile.get_window_origin(); + const auto& bottom_tensor_view = lds_tile.get_bottom_tensor_view(); + const auto& tensor_descriptor = bottom_tensor_view.get_tensor_descriptor(); + auto smem_base_ptr = bottom_tensor_view.get_buffer_view().p_data_; + + // loop over thread tensor space [y0, y1, ...] + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + /// TODO: use structure binding (to be captured later) if compiled in C++20 + auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + auto lds_window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto lds_bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + + // Use precomputed window origin + auto lds_bottom_tensor_thread_idx = + window_origin + lds_window_adaptor_thread_coord.get_bottom_index(); + // Use precomputed tensor descriptor + const auto lds_coord = + make_tensor_coordinate(tensor_descriptor, lds_bottom_tensor_thread_idx); + // Calculate SMEM address using base pointer + CK_TILE_LDS_ADDR LdsDataType* smem = smem_base_ptr + lds_coord.get_offset(); + + // data index [y0, y1, ...] + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + constexpr auto idx_gather = idx_ys_start[number{}]; + const auto page_offset = page_idx_[idx_gather]; + + // merge page_offset into bottom_coord + auto mixed_bottom_thread_coord = bottom_tensor_thread_coord; + mixed_bottom_thread_coord.get_hidden_index()[number<0>{}] += page_offset; + + // read from bottom tensor + if constexpr(std::is_same_v) + this->get_bottom_tensor_view().template async_get_vectorized_elements( + smem, + mixed_bottom_thread_coord, + number<0>{}, + bool_constant{}); + else + this->get_bottom_tensor_view().template async_get_vectorized_elements( + smem, + mixed_bottom_thread_coord, + number<0>{}, + valids_[idx_gather], + bool_constant{}); + + // move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto forward_step_scatter = generate_tuple( + [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + forward_step_scatter); + // lds_diff doesn't need to mask the difference of the gather-dim. + constexpr auto lds_idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + idx_diff_ys); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + move_window_adaptor_and_bottom_tensor_thread_coordinate( + lds_window_adaptor_thread_coord, + lds_bottom_tensor_thread_coord, + lds_idx_diff_ps_ys); + } + }); + }); + } + // TODO: currently async load only implemented in inline asm template + CK_TILE_DEVICE void update(const static_distributed_tensor& dstr_tensor, + number = {}, + bool_constant = {}) const + { + using Traits = load_store_traits; + + // using vector_type_t = typename Traits::vector_type_t; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = TileDstr{}; + + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + + // data index [y0, y1, ...] + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + constexpr auto idx_gather = idx_ys_start[number<0>{}]; + const auto page_offset = page_idx_[idx_gather]; + + // read from distributed tensor + vector_t vec_value; + + static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) { + constexpr auto idx_ys = generate_tuple( + [&](auto jj) { + return jj == Traits::VectorDimY ? (idx_ys_start[jj] + j) + : idx_ys_start[jj]; + }, + number{}); + + constexpr index_t d = + tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys) / + Traits::PackedSize; + + vec_value.template get_as()(j / Traits::PackedSize) = + dstr_tensor.get_thread_buffer().template at(); + }); + + // write into bottom tensor + if constexpr(std::is_same_v) + { + get_bottom_tensor_view().template update_vectorized_elements( + bottom_tensor_thread_coord, + page_offset, + vec_value, + bool_constant{}); + } + else + { + get_bottom_tensor_view().template update_vectorized_elements( + bottom_tensor_thread_coord, + page_offset, + valids_[idx_gather], + vec_value, + bool_constant{}); + } + + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto forward_step_scatter = generate_tuple( + [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + forward_step_scatter); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + } + }); + }); + } + template CK_TILE_DEVICE void store(const static_distributed_tensor& dstr_tensor, number = {},