From 90b90dff086a41ff1253533fa126ea9118785467 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ville=20Pietil=C3=A4?= Date: Wed, 27 Aug 2025 13:29:38 +0000 Subject: [PATCH] Vectorized packed cast. --- .../threadwise_tensor_slice_transfer.hpp | 233 +++++++++++++++++- 1 file changed, 228 insertions(+), 5 deletions(-) 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 b3e45af89c..4932809ec6 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 @@ -239,7 +239,7 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast static constexpr auto I0 = Number<0>{}; static constexpr auto I1 = Number<1>{}; static constexpr auto I2 = Number<2>{}; - static constexpr index_t BufScalarPerVector = 2; + static constexpr bool SerpentineAccessPattern = false; static constexpr index_t nDim = SliceLengths::Size(); @@ -276,6 +276,227 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast dst_coord_ = make_tensor_coordinate(dst_desc, dst_slice_origin_idx); } + template + __device__ void Run(const SrcDesc&, + const SrcSliceOriginIdx&, + const SrcBuffer& src_buf, + const DstDesc& dst_desc, + DstBuffer& dst_buf) + { + static_assert(SrcDesc::IsKnownAtCompileTime(), + "wrong! SrcDesc need to known at compile-time"); + + static_assert(is_known_at_compile_time>::value, + "wrong! SrcSliceOrigin need to known at compile-time"); + + static_assert(SrcBuffer::IsStaticBuffer(), "wrong! SrcBuffer need to be StaticBuffer"); + + // SrcDesc and src_slice_origin_idx are known at compile-time + constexpr auto src_desc = remove_cvref_t{}; + constexpr auto src_slice_origin_idx = to_multi_index(SrcSliceOriginIdx{}); + + constexpr auto dst_scalar_per_access = generate_sequence( + detail::lambda_scalar_per_access{}, Number{}); + + using SpaceFillingCurve = SpaceFillingCurve, + SerpentineAccessPattern>; + + static_assert(1 == SpaceFillingCurve::ScalarPerVector, "wrong!1 != SpaceFillingCurve::ScalarPerVector"); + + constexpr index_t num_access = SpaceFillingCurve::GetNumOfAccess(); + 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."); + + ck::float2_t float2_buffer; + static_for<0, num_pairs, 1>{}([&](auto i_pair) + { + constexpr auto idx_1d_0 = I2 * i_pair; + 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 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); + + if constexpr (src_offset_1 - src_offset_0 == 1) + { + // Load two consecutive float values from the src buffer + float2_buffer = src_buf.template GetAsType(Number{}); + } + else + { + // Load the two float values one by one + float2_buffer= {src_buf[Number{}], src_buf[Number{}]}; + } + + const ck::bhalf2_t packed_value= bf16x2_convert_rne(float2_buffer[0], float2_buffer[1]); + + const bool is_dst_valid = + coordinate_has_valid_offset_assuming_visible_index_is_valid(dst_desc, dst_coord_); + + 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_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); + } + else + { + // Store the packed value into the dst buffer one by one. + dst_buf.template Update( + dst_offset_0, + is_dst_valid, + packed_value[0]); + + dst_buf.template Update( + dst_offset_1, + is_dst_valid, + packed_value[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)); + } + }); + + // move dst coordinate back to slice origin (or not) + if constexpr(DstResetCoordinateAfterRun) + { + const auto dst_reset_step = + make_tensor_coordinate_step(dst_desc, GetDstCoordinateResetStep()); + + move_tensor_coordinate(dst_desc, dst_coord_, dst_reset_step); + } + } + + + __device__ static constexpr auto GetDstCoordinateResetStep() + { + constexpr auto dst_scalar_per_access = generate_sequence( + detail::lambda_scalar_per_access{}, Number{}); + + using SpaceFillingCurve = SpaceFillingCurve, + SerpentineAccessPattern>; + + constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); + if constexpr(num_access == 0) + { + return typename SpaceFillingCurve::Index{}; + } + else + { + constexpr auto reset_step = + SpaceFillingCurve::GetStepBetween(Number{}, Number<0>{}); + + return reset_step; + } + } + + // dst_slice_origin_step_idx need to be known at compile-time, for performance reason + __device__ void MoveDstSliceWindow(const DstDesc& dst_desc, + const Index& dst_slice_origin_step_idx) + { + // if dst coord was not reset by Run(), then need to adjust the step here + const auto adjusted_step_idx = + DstResetCoordinateAfterRun ? dst_slice_origin_step_idx + : dst_slice_origin_step_idx + GetDstCoordinateResetStep(); + + // is it OK to construct a new step every time? + const auto adjusted_step = make_tensor_coordinate_step(dst_desc, adjusted_step_idx); + + move_tensor_coordinate(dst_desc, dst_coord_, adjusted_step); + } +private: + DstCoord dst_coord_; +}; + +// Assume: +// 1. src: +// 1. SrcDesc is known at compile-time +// 2. SrcBuffer is StaticBuffer +// 3. SrcSliceOrginIdx is known at compile-time +// 2. dst: +// 1. DstDesc is not known at compile-time +// 2. DstBuffer is DynamicBuffer +// 3. DstSliceOrginIdx is not known at compile time +template ::type = false> +struct ThreadwiseTensorSliceTransfer_v1r3_buffered_packed_cast +{ + static constexpr auto I0 = Number<0>{}; + static constexpr auto I1 = Number<1>{}; + static constexpr auto I2 = Number<2>{}; + static constexpr bool SerpentineAccessPattern = false; + + static constexpr index_t nDim = SliceLengths::Size(); + + using Index = MultiIndex; + + using DstCoord = decltype(make_tensor_coordinate(DstDesc{}, Index{})); + + using DstCoordStep = decltype(make_tensor_coordinate_step(DstDesc{}, Index{})); + + __device__ constexpr ThreadwiseTensorSliceTransfer_v1r3_buffered_packed_cast(const DstDesc& dst_desc, + const Index& dst_slice_origin_idx, + const ElementwiseOperation&) + : dst_coord_(make_tensor_coordinate(dst_desc, dst_slice_origin_idx)) + { + static_assert(SrcDesc::IsKnownAtCompileTime(), + "wrong! SrcDesc need to known at compile-time"); + static_assert(SliceLengths::At(Number{}) % DstScalarPerVector == 0, + "wrong! Not divisible"); + + // Assert that elementwise op is pass through. + static_assert( + std::is_same_v, ck::tensor_operation::element_wise::PassThrough>, + "wrong! ElementwiseOperation must be PassThrough"); + + // For now, SrcData must be float and DstData must be ck::bhalf_t + static_assert(std::is_same_v, + "wrong! SrcData must be float"); + static_assert(std::is_same_v, + "wrong! DstData must be bhalf_t"); + } + + __device__ void SetDstSliceOrigin(const DstDesc& dst_desc, const Index& dst_slice_origin_idx) + { + dst_coord_ = make_tensor_coordinate(dst_desc, dst_slice_origin_idx); + } + template __device__ void RunRead(const SrcBuffer& src_buf) { @@ -297,7 +518,8 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast using SpaceFillingCurve = SpaceFillingCurve>; + remove_cv_t, + SerpentineAccessPattern>; static_assert(1 == SpaceFillingCurve::ScalarPerVector, "wrong!1 != SpaceFillingCurve::ScalarPerVector"); @@ -337,14 +559,14 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast using SpaceFillingCurve = SpaceFillingCurve>; + remove_cv_t, + SerpentineAccessPattern>; static_assert(1 == SpaceFillingCurve::ScalarPerVector, "wrong!1 != SpaceFillingCurve::ScalarPerVector"); constexpr index_t num_access = SpaceFillingCurve::GetNumOfAccess(); static_assert(num_access == buffer_size_, "wrong!num_access != buffer_size_"); - static_for<0, num_access, 1>{}([&](auto idx_1d) { const ck::bhalf_t val = thread_scratch_.template GetAsType(Number{}); @@ -394,7 +616,8 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast using SpaceFillingCurve = SpaceFillingCurve>; + remove_cv_t, + SerpentineAccessPattern>; constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); if constexpr(num_access == 0)