From 62c66a7d9c53ad8437f3b05f8a78edb99088de82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ville=20Pietil=C3=A4?= Date: Thu, 14 Aug 2025 12:39:18 +0000 Subject: [PATCH] WIP: packed bf16 cast v3. --- .../threadwise_tensor_slice_transfer.hpp | 168 +++++++++++------- .../test_tensor_slice_transfer.cpp | 47 ++--- test/utils/test_sequencies.cpp | 1 - 3 files changed, 132 insertions(+), 84 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 a26f8d79bd..ed0667301e 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 @@ -14,38 +14,79 @@ namespace ck { - -template + bool DstResetCoordinateAfterRun, + typename enable_if::type = false> 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 SrcScalarPerVector = 1; - static constexpr index_t DstScalarPerVector = 2 + static constexpr index_t BufScalarPerVector = 2; + + static constexpr auto FastestChangingDim = DimAccessOrder::At(Number{}); + + 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{})); - using BaseTransfer = ThreadwiseTensorSliceTransfer_v1r3< - float, ck::bhalf_t, SrcDesc, DstDesc, ElementwiseOperation, SliceLengths, - DimAccessOrder, DstVectorDim, SrcScalarPerVector, DstInMemOp, - DstScalarStrideInVector, DstResetCoordinateAfterRun, PackedInput>; - __device__ constexpr ThreadwiseTensorSliceTransfer_v1r3_packed_cast(const DstDesc& dst_desc, const Index& dst_slice_origin_idx, - const ElementwiseOperation& element_op) : base_transfer_(dst_desc, dst_slice_origin_idx, element_op) + 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(FastestChangingDim) % 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"); + + // Assert that DstScalarPerVector is divisible by 2, since we are using bf16x2_convert_rne. + static_assert(DstScalarPerVector % 2 == 0, "wrong! DstScalarPerVector must be divisible by 2"); + } + + __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 Run(const SrcDesc& src_desc, - const SrcSliceOriginIdx& src_slice_origin_idx, + __device__ void Run(const SrcDesc&, + const SrcSliceOriginIdx&, const SrcBuffer& src_buf, const DstDesc& dst_desc, DstBuffer& dst_buf) @@ -56,24 +97,21 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast static_assert(is_known_at_compile_time>::value, "wrong! SrcSliceOrigin need to known at compile-time"); - static_assert(base_transfer.SrcBuffer::IsStaticBuffer(), "wrong! SrcBuffer need to be StaticBuffer"); + 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{}); + constexpr auto src_scalar_per_access = generate_sequence( + detail::lambda_scalar_per_access{}, Number{}); using SpaceFillingCurve = SpaceFillingCurve>; + remove_cv_t>; typename vector_type_maker::type dst_vector; using dst_vector_t = typename vector_type_maker::type::type; - - static_assert(std::is_same_v>, - "wrong! DimAccessOrder must be the identity sequence <0, 1, 2, 3, 4, 5, 6, 7>"); static_assert(1 == SpaceFillingCurve::ScalarPerVector, "wrong!1 != SpaceFillingCurve::ScalarPerVector"); @@ -81,6 +119,9 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast 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) { constexpr auto idx_1d_0 = I2 * i_pair; @@ -91,37 +132,10 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast 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); - float& val_0 = src_buf(Number{}); - float& val_1 = src_buf(Number{}); + const float val_0 = src_buf[Number{}]; + const float val_1 = src_buf[Number{}]; - //static_cast_float_to_bhalf_packed_v2(val_0, val_1); - // Fill the packed value into the dst_vector - }); - - static_for<0, num_access, 1>{}([&](auto idx_1d) - { - // 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 index_t src_offset = src_desc.CalculateOffset(src_slice_origin_idx + idx_md); - - union - { - float src_float; - bhalf16_t src_bf16x2; - } packed_value; - - packed_value.src_float = src_buf[Number{}]; - - DstData v; - base_transfer.get_element_op()(v, packed_value.src_bf16x2[pair_index.value]); - dst_vector.template AsType()(I0) = v; + dst_vector.template AsType()(I0) = bf16x2_convert_rne(val_0, val_1); const bool is_dst_valid = coordinate_has_valid_offset_assuming_visible_index_is_valid(dst_desc, dst_coord_); @@ -132,12 +146,12 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast is_dst_valid, dst_vector.template AsType()[Number<0>{}]); - if constexpr(idx_1d.value != num_access - 1) + if constexpr(i_pair.value != num_pairs - 1 && !has_odd_element) { - constexpr auto forward_step = SpaceFillingCurve::GetForwardStep(idx_1d); + constexpr auto forward_step = SpaceFillingCurve::GetForwardStep(i_pair); move_tensor_coordinate( - dst_desc, base_transfer_.get_dst_coord(), make_tensor_coordinate_step(dst_desc, forward_step)); + dst_desc, dst_coord_, make_tensor_coordinate_step(dst_desc, forward_step)); } }); @@ -145,13 +159,52 @@ struct ThreadwiseTensorSliceTransfer_v1r3_packed_cast if constexpr(DstResetCoordinateAfterRun) { const auto dst_reset_step = - make_tensor_coordinate_step(dst_desc, base_transfer.GetDstCoordinateResetStep()); + make_tensor_coordinate_step(dst_desc, GetDstCoordinateResetStep()); - move_tensor_coordinate(dst_desc, base_transfer_.get_dst_coord(), dst_reset_step); + 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>; + + 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: - BaseTransfer base_transfer_; + DstCoord dst_coord_; }; // Assume: @@ -327,11 +380,6 @@ struct ThreadwiseTensorSliceTransfer_v1r3 move_tensor_coordinate(dst_desc, dst_coord_, adjusted_step); } - - DestCoord& get_dest_coord() { return dst_coord_; } - - const ElementwiseOperation& get_element_op() const { return element_op_; } - private: DstCoord dst_coord_; const ElementwiseOperation element_op_; diff --git a/test/tensor_slice_transfer/test_tensor_slice_transfer.cpp b/test/tensor_slice_transfer/test_tensor_slice_transfer.cpp index 3399261233..5fd84f80d5 100644 --- a/test/tensor_slice_transfer/test_tensor_slice_transfer.cpp +++ b/test/tensor_slice_transfer/test_tensor_slice_transfer.cpp @@ -10,7 +10,6 @@ #include "ck/ck.hpp" #include "ck/utility/common_header.hpp" #include "ck/tensor_description/tensor_space_filling_curve.hpp" -#include "ck/tensor_operation/gpu/grid/packed_cast.hpp" #include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" #include "ck/utility/data_type.hpp" @@ -95,17 +94,6 @@ __global__ void testVGPRToLDSTransfer_kernel( auto lds_buf = make_dynamic_buffer( &lds_data[thread_id * buffer_size], buffer_size); - // Test with packed cast -#if defined(__gfx950__) - if constexpr (UsePackedCast) { - auto packed_cast = PackedCastV2{}; - packed_cast.Run(c_thread_desc, - src_slice_origin_index, - src_thread_buf); - } -#endif // Create threadwise transfer using ElementOp = tensor_operation::element_wise::PassThrough; @@ -113,7 +101,25 @@ __global__ void testVGPRToLDSTransfer_kernel( const auto dst_slice_origin_index = make_multi_index(0, 0, 0, 0, 0, 0, 0, 0); - auto thread_transfer = ThreadwiseTensorSliceTransfer_v1r3< + using TransferType = std::conditional_t< + UsePackedCast, + ThreadwiseTensorSliceTransfer_v1r3_packed_cast< + SrcData, + DstData, + decltype(c_thread_desc), + decltype(lds_desc), + ElementOp, + Sequence, + Sequence<0, 1, 2, 3, 4, 5, 7, 6>, // Note: 7, 6 are swapped to enable vectorized transfer. + 7, // DstVectorDim + 2, // DstScalarPerVector + InMemoryDataOperationEnum::Set, + 1, // DstScalarStrideInVector + true // DstResetCoordinateAfterRun + >, + ThreadwiseTensorSliceTransfer_v1r3< SrcData, DstData, decltype(c_thread_desc), @@ -127,9 +133,11 @@ __global__ void testVGPRToLDSTransfer_kernel( 1, // DstScalarPerVector InMemoryDataOperationEnum::Set, 1, // DstScalarStrideInVector - true, // DstResetCoordinateAfterRun - UsePackedCast // PackedInput - >{lds_desc, dst_slice_origin_index, element_op}; + true // DstResetCoordinateAfterRun + > + >; + + auto thread_transfer = TransferType{lds_desc, dst_slice_origin_index, element_op}; // Perform the transfer if (thread_id < num_elements) { @@ -330,7 +338,6 @@ TEST_F(VGPRToLDSTransferTest, FloatToBhalf_device_NoPack) run(); } -#if defined(__gfx950__) TEST_F(VGPRToLDSTransferTest, FloatToBhalf_device_PackedCast) { constexpr bool UsePackedCast = true; @@ -358,9 +365,3 @@ TEST_F(VGPRToLDSTransferTest, FloatToBhalf_device_test_peformance) GTEST_FAIL() << "Failed to get average execution time for one or both runs."; } } -#endif // defined(__gfx950__) - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/test/utils/test_sequencies.cpp b/test/utils/test_sequencies.cpp index b5e8071122..cc88d042f9 100644 --- a/test/utils/test_sequencies.cpp +++ b/test/utils/test_sequencies.cpp @@ -19,7 +19,6 @@ TEST(Sequence, TestCreatingAccessSequence) constexpr index_t nDim = 8; constexpr auto SliceLengths = Sequence<4, 8, 1, 1, 4, 1, 2, 1>{}; - //constexpr auto DimAccessOrder = Sequence<0, 1, 2, 3, 4, 5, 7, 6>{}; constexpr auto ScalarPerAccess = generate_sequence( detail::lambda_scalar_per_access{}, Number{});