From 9f66d9fbcabb18e9d9bdc20a400d0d25ce321832 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ville=20Pietil=C3=A4?= Date: Wed, 27 Aug 2025 14:32:00 +0000 Subject: [PATCH] Packed cast improvement. --- .../tensor_operation/gpu/grid/packed_cast.hpp | 5 +- .../threadwise_tensor_slice_transfer.hpp | 75 ++++++++++++------- 2 files changed, 51 insertions(+), 29 deletions(-) diff --git a/include/ck/tensor_operation/gpu/grid/packed_cast.hpp b/include/ck/tensor_operation/gpu/grid/packed_cast.hpp index 50ab9ea90a..6a7190aa05 100644 --- a/include/ck/tensor_operation/gpu/grid/packed_cast.hpp +++ b/include/ck/tensor_operation/gpu/grid/packed_cast.hpp @@ -62,9 +62,11 @@ namespace ck { constexpr auto idx_1d_1 = I2 * i_pair + I1; constexpr auto idx_md_0 = SpaceFillingCurve::GetIndex(idx_1d_0); constexpr auto idx_md_1 = SpaceFillingCurve::GetIndex(idx_1d_1); + constexpr auto idx_md_pair = SpaceFillingCurve::GetIndex(i_pair); constexpr index_t src_offset_0 = src_desc.CalculateOffset(src_slice_origin_idx + idx_md_0); constexpr index_t src_offset_1 = src_desc.CalculateOffset(src_slice_origin_idx + idx_md_1); + constexpr index_t pair_offset = src_desc.CalculateOffset(src_slice_origin_idx + idx_md_pair); if constexpr (src_offset_1 - src_offset_0 == 1) { @@ -78,14 +80,13 @@ namespace ck { } // Store the packed bfloat2 value back to the src buffer - // Note: The value is stored to the space of the first float value. const ck::bhalf2_t packed_value= bf16x2_convert_rne(float2_buffer[0], float2_buffer[1]); union { ck::bhalf2_t bhalf2; float fp32; } converter; converter.bhalf2 = packed_value; - src_buf(Number{}) = converter.fp32; + src_buf(Number{}) = converter.fp32; }); }; }; diff --git a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp index 4932809ec6..bd440c3f45 100644 --- a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp +++ b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp @@ -104,8 +104,6 @@ struct ThreadwiseTensorSliceTransfer_v1r3_pass_through // TODO: Use SpaceFillingCurve::ScalarsPerAccess instread of DstScalarPerVector? static_assert(DstScalarPerVector == SpaceFillingCurve::ScalarPerVector, "wrong!DstScalarPerVector != SpaceFillingCurve::ScalarPerVector"); - typename vector_type_maker::type dst_vector; - using dst_vector_t = typename vector_type_maker::type::type; constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); @@ -115,17 +113,15 @@ struct ThreadwiseTensorSliceTransfer_v1r3_pass_through static_assert(1 == SpaceFillingCurve::ScalarPerVector, "wrong!1 != SpaceFillingCurve::ScalarPerVector"); static_assert(1 == DstScalarPerVector, "wrong!1 != DstScalarPerVector"); - static_for<0, num_access, 1>{}([&](auto idx_1d) + constexpr index_t num_pairs = num_access / 2; + constexpr bool has_odd_element = (num_access % 2 == 1); + + // TODO: Enable also odd number of elements. + static_assert(!has_odd_element, "wrong!Slice should have even number of elements."); + + static_for<0, num_pairs, 1>{}([&](auto i_pair) { - // We need map the odd indices to the even indices, since - // the even indices contain a packed bf16x2 value, where - // the first value contains the bf16 value for the corresponding even index - // and the second value contains the bf16 value for the odd index following the even index. - // The odd indices are not used, so we can just ignore them. - constexpr auto pair_index = idx_1d % I2; - constexpr auto idx_src_1d = idx_1d - pair_index; - - constexpr auto idx_md = SpaceFillingCurve::GetIndex(idx_src_1d); + constexpr auto idx_md = SpaceFillingCurve::GetIndex(i_pair); constexpr index_t src_offset = src_desc.CalculateOffset(src_slice_origin_idx + idx_md); union @@ -135,26 +131,51 @@ struct ThreadwiseTensorSliceTransfer_v1r3_pass_through } packed_value; packed_value.src_float = src_buf[Number{}]; - - DstData v; - element_op_(v, packed_value.src_bf16x2[pair_index.value]); - dst_vector.template AsType()(I0) = v; - + const bool is_dst_valid = coordinate_has_valid_offset_assuming_visible_index_is_valid(dst_desc, dst_coord_); - // copy data from dst_vector into dst_buf - dst_buf.template Update( - dst_coord_.GetOffset(), - is_dst_valid, - dst_vector.template AsType()[Number<0>{}]); - - if constexpr(idx_1d.value != num_access - 1) - { - constexpr auto forward_step = SpaceFillingCurve::GetForwardStep(idx_1d); + constexpr auto idx_1d_0 = I2 * i_pair; + constexpr auto idx_1d_1 = I2 * i_pair + I1; + + const auto dst_offset_0 = dst_coord_.GetOffset(); + // Move to the next dst coordinate + constexpr auto forward_step_0 = SpaceFillingCurve::GetForwardStep(idx_1d_0); move_tensor_coordinate( - dst_desc, dst_coord_, make_tensor_coordinate_step(dst_desc, forward_step)); + dst_desc, dst_coord_, make_tensor_coordinate_step(dst_desc, forward_step_0)); + + const auto dst_offset_1 = dst_coord_.GetOffset(); + + if (dst_offset_1 - dst_offset_0 == 1) + { + // Store the packed value directly since we have consequtive locations is the dst buffer. + dst_buf.template Update( + dst_coord_.GetOffset(), + is_dst_valid, + packed_value.src_bf16x2); + } + else + { + // Store the packed value into the dst buffer one by one. + dst_buf.template Update( + dst_offset_0, + is_dst_valid, + packed_value.src_bf16x2[0]); + + dst_buf.template Update( + dst_offset_1, + is_dst_valid, + packed_value.src_bf16x2[1]); + } + + // Move to next dst coordinate, unless this was the last pair + if constexpr(i_pair.value != num_pairs - 1) + { + + constexpr auto forward_step_1 = SpaceFillingCurve::GetForwardStep(idx_1d_1); + move_tensor_coordinate( + dst_desc, dst_coord_, make_tensor_coordinate_step(dst_desc, forward_step_1)); } });