From 008f64ad6d4bc42351a230cc998f706d7c82798a Mon Sep 17 00:00:00 2001 From: OscarXu Date: Tue, 22 Jul 2025 10:02:02 +0800 Subject: [PATCH] Merge branch 'develop' into moe_bs_fp8_no_asm_buf2lds --- .../blockwise_gemm_pipeline_xdlops_base.hpp | 5 - ...line_xdlops_blockscale_b_preshuffle_v1.hpp | 1 + ...line_xdlops_blockscale_b_preshuffle_v3.hpp | 23 +- .../threadwise_tensor_slice_transfer copy.hpp | 2101 ----------------- 4 files changed, 13 insertions(+), 2117 deletions(-) delete mode 100644 include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer copy.hpp diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp index d9a8e7bab6..9296b8136f 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp @@ -58,11 +58,6 @@ struct BlockwiseGemmXdlops_pipeline_base static constexpr index_t KPerThread = KPerBlock / xdlops_gemm.K0PerXdlops; static constexpr index_t KRepeat = KPerThread / KPack; static constexpr index_t KPerInnerLoop = KPack; - static constexpr index_t KGroup = - ((MPerXDL == 16 && MPerXDL == 16 && xdlops_gemm.KPerXdlops == 128) || - (MPerXDL == 32 && MPerXDL == 32 && xdlops_gemm.KPerXdlops == 64)) - ? 2 - : 1; static constexpr index_t KGroup = []() { if constexpr(is_same_v, f8_t>) diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v1.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v1.hpp index ce56720e03..8e2922e2ce 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v1.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v1.hpp @@ -814,6 +814,7 @@ struct BlockwiseGemmXdlops_pipeline_blockscale_bpreshuffle_v1(), c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); }); + constexpr index_t c_offset = c_thread_desc_.CalculateOffset(make_tuple(m0, n0, 0)); diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v3.hpp index 4b25fb93b2..cc4c5a2c36 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_blockscale_b_preshuffle_v3.hpp @@ -593,6 +593,15 @@ struct BlockwiseGemmXdlops_pipeline_blockscale_bpreshuffle_v3{})); } - b_scale_thread_copy.Run(b_scale_grid_desc, - b_scale_grid_buf, - b_scale_thread_desc, - make_tuple(I0, I0), - b_scale_thread_bufs(local_read_buf)); - - b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, - b_scale_thread_copy_step); - static_for<0, MRepeat, 1>{}([&](auto m0) { vector_type c_scale_thread_vec; c_scale_thread_vec.template AsType()(Number<0>{}) = @@ -744,8 +744,9 @@ struct BlockwiseGemmXdlops_pipeline_blockscale_bpreshuffle_v3::type = false> -struct ThreadwiseTensorSliceTransfer_v1r3 -{ - 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(const DstDesc& dst_desc, - const Index& dst_slice_origin_idx, - const ElementwiseOperation& element_op) - : dst_coord_(make_tensor_coordinate(dst_desc, dst_slice_origin_idx)), - element_op_{element_op} - { - static_assert(SrcDesc::IsKnownAtCompileTime(), - "wrong! SrcDesc need to known at compile-time"); - static_assert(SliceLengths::At(Number{}) % DstScalarPerVector == 0, - "wrong! Not divisible"); - } - - __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&, - 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{}); - - // scalar per access on each dim - // TODO: don't use lambda_scalar_per_access - constexpr auto dst_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto dst_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - using SpaceFillingCurve = SpaceFillingCurve>; - - // 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(); - - static_for<0, num_access, 1>{}([&](auto idx_1d) { - constexpr auto idx_md = SpaceFillingCurve::GetIndex(idx_1d); - - // copy data from src_buf into dst_vector - // TODO: It's a hack here to use \p dst_scalar_step_in_vector. Use SpaceFillingCurve? - static_for<0, DstScalarPerVector, 1>{}([&](auto i) { - constexpr index_t src_offset = src_desc.CalculateOffset( - src_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - DstData v; - - // apply element-wise operation - element_op_(v, src_buf[Number{}]); - - dst_vector.template AsType()(i) = 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); - - move_tensor_coordinate( - dst_desc, dst_coord_, make_tensor_coordinate_step(dst_desc, forward_step)); - } - }); - - // 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>; - - 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_; - const ElementwiseOperation element_op_; -}; // namespace ThreadwiseTensorSliceTransfer_v1r3 - -/** - * @brief Helper structure that facilitates transfer of source (grid) data to destination threads. - * - * @details The following assumptions are made: - * - For Source (Grid) Data: - * 1. The source tensor descriptor SrcDesc is not known at compile-time. - * 2. The source buffer is a dynamic buffer. - * 3. The source slice origin index src_slice_origin_idx is not known at compile-time. - * - For Destination (Thread) Data: - * 1. The destination tensor descriptor DstDesc is known at compile-time. - * 2. The destination buffer dst_buf is a static buffer. - * 3. The destination slice origin index dst_slice_origin_idx is known at compile-time. - * - * @tparam SrcData The data type of the source tensor. - * @tparam DstData The data type of the destination tensor. - * @tparam SrcDesc The descriptor type of the source tensor. - * @tparam DstDesc The descriptor type of the destination tensor. - * @tparam SliceLengths The lengths of the slice to be transferred. - * @tparam DimAccessOrder The order of dimension access for the space-filling curve. - * @tparam SrcVectorDim The dimension along which vectorized access is performed in the source - * tensor. - * @tparam SrcScalarPerVector The number of scalar elements per vector in the source tensor. - * @tparam SrcScalarStrideInVector The stride of scalar elements within a vector in the source - * tensor. - * @tparam SrcResetCoordinateAfterRun controls whether source coordinate is restored after each Run - * or rolled back one step in MoveSrcSliceWindow - * @tparam InvalidElementAsNaN Whether to fill invalid elements with NaN (only applicable for - * floating-point types). - * - */ -template ::type = false> -struct ThreadwiseTensorSliceTransfer_v2 -{ - static_assert((InvalidElementAsNaN && !ck::is_integral::value) || - (!InvalidElementAsNaN), - "Filling invalid element as NaN is only for floating point types"); - - static constexpr index_t nDim = SliceLengths::Size(); - - using Index = MultiIndex; - - using SrcCoord = decltype(make_tensor_coordinate(SrcDesc{}, Index{})); - - using SrcCoordStep = decltype(make_tensor_coordinate_step(SrcDesc{}, Index{})); - - static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t>) - return 2; - else - return 1; - }(); - - __device__ constexpr ThreadwiseTensorSliceTransfer_v2(const SrcDesc& src_desc, - const Index& src_slice_origin_idx) - : src_coord_(make_tensor_coordinate(src_desc, src_slice_origin_idx)) - { - static_assert(DstDesc::IsKnownAtCompileTime(), - "wrong! SrcDesc need to known at compile-time"); - static_assert(SliceLengths::At(Number{}) % SrcScalarPerVector == 0, - "wrong! Not divisible"); - - if constexpr(is_same_v, pk_i4_t>) - { - static_assert(SrcScalarPerVector % PackedSize == 0, "pk data N cannot be 1"); - } - } - - __device__ void SetSrcSliceOrigin(const SrcDesc& src_desc, const Index& src_slice_origin_idx) - { - src_coord_ = make_tensor_coordinate(src_desc, src_slice_origin_idx); - } - - template - __device__ void Run(const SrcDesc& src_desc, - const SrcBuffer& src_buf, - const DstDesc&, - const DstSliceOriginIdx&, - DstBuffer& dst_buf) - { - static_assert(DstDesc::IsKnownAtCompileTime(), - "wrong! DstDesc need to known at compile-time"); - - static_assert(is_known_at_compile_time>::value, - "wrong! DstSliceOrigin need to known at compile-time"); - - static_assert( - is_same, remove_cvref_t>::value && - "wrong! inconsistent type"); - - // DstDesc and dst_slice_origin_idx are known at compile-time - constexpr auto dst_desc = remove_cvref_t{}; - constexpr auto dst_slice_origin_idx = DstSliceOriginIdx{}; - - // scalar per access on each dim - // TODO: don't use lambda_scalar_per_access - constexpr auto src_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto src_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - using SpaceFillingCurve = SpaceFillingCurve>; - - // loop over tensor and copy - constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); - - static_for<0, num_access, 1>{}([&](auto idx_1d) { - typename vector_type_maker::type src_vector; - - using src_vector_t = - typename vector_type_maker::type::type; - constexpr auto src_data_idx = SpaceFillingCurve::GetIndex(idx_1d); - - const bool is_src_valid = - coordinate_has_valid_offset_assuming_visible_index_is_valid(src_desc, src_coord_); - - // copy data from src_buf into src_vector - src_vector.template AsType()(Number<0>{}) = - src_buf.template Get(src_coord_.GetOffset() / PackedSize, - is_src_valid); - - // copy data from src_vector into dst_buf - static_for<0, SrcScalarPerVector / PackedSize, 1>{}([&](auto i) { - constexpr index_t dst_offset = - dst_desc.CalculateOffset(to_multi_index(dst_slice_origin_idx) + src_data_idx + - i * src_scalar_step_in_vector); - - if constexpr(InvalidElementAsNaN) - { - dst_buf(Number{}) = - is_src_valid - ? type_convert(src_vector.template AsType()[i]) - : NumericLimits::QuietNaN(); - } - else - { - dst_buf(Number{}) = - type_convert(src_vector.template AsType()[i]); - } - }); - - if constexpr(idx_1d.value != num_access - 1) - { - constexpr auto forward_step = SpaceFillingCurve::GetForwardStep(idx_1d); - - move_tensor_coordinate( - src_desc, src_coord_, make_tensor_coordinate_step(src_desc, forward_step)); - } - }); - - // move src coordinate back to slice origin (or not) - if constexpr(SrcResetCoordinateAfterRun) - { - const auto src_reset_step = - make_tensor_coordinate_step(src_desc, GetSrcCoordinateResetStep()); - - move_tensor_coordinate(src_desc, src_coord_, src_reset_step); - } - } - - __device__ static constexpr auto GetSrcCoordinateResetStep() - { - constexpr auto src_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 MoveSrcSliceWindow(const SrcDesc& src_desc, - const Index& src_slice_origin_step_idx) - { - // if src coord was not reset by Run(), then need to adjust the step here - const auto adjusted_step_idx = - SrcResetCoordinateAfterRun ? src_slice_origin_step_idx - : src_slice_origin_step_idx + GetSrcCoordinateResetStep(); - - // is it OK to construct a new step every time? - const auto adjusted_step = make_tensor_coordinate_step(src_desc, adjusted_step_idx); - - move_tensor_coordinate(src_desc, src_coord_, adjusted_step); - } - - // src_slice_origin_step_idx need to be known at compile-time, for performance reason - template - __device__ void - MoveSrcSliceWindow(const SrcDesc& src_desc, - const Index& src_slice_origin_step_idx, - const SrcMoveSliceWindowStepHack& src_move_slice_window_step_hack) - { - // if src coord was not reset by RunRead(), then need to adjust the step here - const auto adjusted_step_idx = - SrcResetCoordinateAfterRun ? src_slice_origin_step_idx - : src_slice_origin_step_idx + GetSrcCoordinateResetStep(); - - // is it OK to construct a new step every time? - const auto adjusted_step = make_tensor_coordinate_step( - src_desc, adjusted_step_idx, src_move_slice_window_step_hack); - - move_tensor_coordinate(src_desc, src_coord_, adjusted_step); - } - - private: - SrcCoord src_coord_; -}; // namespace ck - - -template ::type = false> -struct ThreadwiseTensorSliceTransfer_v2_gather -{ - static_assert((InvalidElementAsNaN && !ck::is_integral::value) || - (!InvalidElementAsNaN), - "Filling invalid element as NaN is only for floating point types"); - - static constexpr index_t nDim = SliceLengths::Size(); - - using Index = MultiIndex; - - using SrcCoord = decltype(make_tensor_coordinate(SrcDesc{}, Index{})); - - using SrcCoordStep = decltype(make_tensor_coordinate_step(SrcDesc{}, Index{})); - - static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t>) - return 2; - else - return 1; - }(); - - __device__ constexpr ThreadwiseTensorSliceTransfer_v2_gather(const SrcDesc& src_desc, - const Index& src_slice_origin_idx, - const StaticallyIndexedArray& scale_gather_offsets) - : src_coord_(make_tensor_coordinate(src_desc, src_slice_origin_idx)), - scale_gather_offsets_(scale_gather_offsets) - { - static_assert(DstDesc::IsKnownAtCompileTime(), - "wrong! SrcDesc need to known at compile-time"); - static_assert(SliceLengths::At(Number{}) % SrcScalarPerVector == 0, - "wrong! Not divisible"); - - if constexpr(is_same_v, pk_i4_t>) - { - static_assert(SrcScalarPerVector % PackedSize == 0, "pk data N cannot be 1"); - } - } - - __device__ void SetSrcSliceOrigin(const SrcDesc& src_desc, const Index& src_slice_origin_idx) - { - auto adjusted_origin_idx = [&]() { - Index idx; - - static_for<0, nDim, 1>{}([&](auto i) { - idx(i) = i.value == 0 ? 0 : src_slice_origin_idx[Number{}]; - }); - - return idx; - }(); - - src_coord_ = make_tensor_coordinate(src_desc, adjusted_origin_idx); - } - - template - __device__ void Run(const SrcDesc& src_desc, - const SrcBuffer& src_buf, - const DstDesc&, - const DstSliceOriginIdx&, - DstBuffer& dst_buf) - { - static_assert(DstDesc::IsKnownAtCompileTime(), - "wrong! DstDesc need to known at compile-time"); - - static_assert(is_known_at_compile_time>::value, - "wrong! DstSliceOrigin need to known at compile-time"); - - static_assert( - is_same, remove_cvref_t>::value && - "wrong! inconsistent type"); - - // DstDesc and dst_slice_origin_idx are known at compile-time - constexpr auto dst_desc = remove_cvref_t{}; - constexpr auto dst_slice_origin_idx = DstSliceOriginIdx{}; - - // scalar per access on each dim - // TODO: don't use lambda_scalar_per_access - constexpr auto src_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto src_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - using SpaceFillingCurve = SpaceFillingCurve>; - - // loop over tensor and copy - constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); - - static_for<0, scale_gather_num, 1>{}([&](auto gather_idx) { - auto current_dst_origin = to_multi_index(dst_slice_origin_idx) + make_multi_index(gather_idx, 0); - - static_for<0, num_access, 1>{}([&](auto idx_1d) { - typename vector_type_maker::type src_vector; - - using src_vector_t = - typename vector_type_maker::type::type; - constexpr auto src_data_idx = SpaceFillingCurve::GetIndex(idx_1d); - - const bool is_src_valid = - coordinate_has_valid_offset_assuming_visible_index_is_valid(src_desc, src_coord_); - - // copy data from src_buf into src_vector - src_vector.template AsType()(Number<0>{}) = - src_buf.template Get(src_coord_.GetOffset() / PackedSize + scale_gather_offsets_(gather_idx), - is_src_valid); - - // copy data from src_vector into dst_buf - static_for<0, SrcScalarPerVector / PackedSize, 1>{}([&](auto i) { - constexpr index_t dst_offset = - dst_desc.CalculateOffset(to_multi_index(dst_slice_origin_idx) + src_data_idx + - i * src_scalar_step_in_vector); - auto full_dst_offset = dst_desc.CalculateOffset(current_dst_origin) + dst_offset; - - if constexpr(InvalidElementAsNaN) - { - dst_buf(full_dst_offset) = - is_src_valid - ? type_convert(src_vector.template AsType()[i]) - : NumericLimits::QuietNaN(); - } - else - { - dst_buf(full_dst_offset) = - type_convert(src_vector.template AsType()[i]); - } - }); - - if constexpr(idx_1d.value != num_access - 1) - { - constexpr auto forward_step = SpaceFillingCurve::GetForwardStep(idx_1d); - - move_tensor_coordinate( - src_desc, src_coord_, make_tensor_coordinate_step(src_desc, forward_step)); - } - }); - }); - - // move src coordinate back to slice origin (or not) - if constexpr(SrcResetCoordinateAfterRun) - { - const auto src_reset_step = - make_tensor_coordinate_step(src_desc, GetSrcCoordinateResetStep()); - - move_tensor_coordinate(src_desc, src_coord_, src_reset_step); - } - } - - __device__ static constexpr auto GetSrcCoordinateResetStep() - { - constexpr auto src_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 MoveSrcSliceWindow(const SrcDesc& src_desc, - const Index& src_slice_origin_step_idx) - { - // if src coord was not reset by Run(), then need to adjust the step here - const auto adjusted_step_idx = - SrcResetCoordinateAfterRun ? src_slice_origin_step_idx - : src_slice_origin_step_idx + GetSrcCoordinateResetStep(); - - // is it OK to construct a new step every time? - const auto adjusted_step = make_tensor_coordinate_step(src_desc, adjusted_step_idx); - - move_tensor_coordinate(src_desc, src_coord_, adjusted_step); - } - - // src_slice_origin_step_idx need to be known at compile-time, for performance reason - template - __device__ void - MoveSrcSliceWindow(const SrcDesc& src_desc, - const Index& src_slice_origin_step_idx, - const SrcMoveSliceWindowStepHack& src_move_slice_window_step_hack) - { - // if src coord was not reset by RunRead(), then need to adjust the step here - const auto adjusted_step_idx = - SrcResetCoordinateAfterRun ? src_slice_origin_step_idx - : src_slice_origin_step_idx + GetSrcCoordinateResetStep(); - - // is it OK to construct a new step every time? - const auto adjusted_step = make_tensor_coordinate_step( - src_desc, adjusted_step_idx, src_move_slice_window_step_hack); - - move_tensor_coordinate(src_desc, src_coord_, adjusted_step); - } - - private: - SrcCoord src_coord_; - StaticallyIndexedArray scale_gather_offsets_; -}; // namespace ck - - - -// Assume: -// 1. src_desc and dst_desc are not known at compile-time -// 2. SrcBuffer and DstBuffer are DynamicBuffer -// 3. src_slice_origin and dst_slice_origin are not known at compile-time, -// 4. Use thread buffer -template // control whether to move back dst coordinate after each - // RunWrite(), will be fused with MoveDstSliceWindow to - // save addr computation -struct ThreadwiseTensorSliceTransfer_v3 -{ - static constexpr index_t nDim = SliceLengths::Size(); - using Index = MultiIndex; - - using SrcCoord = decltype(make_tensor_coordinate(SrcDesc{}, Index{})); - using DstCoord = decltype(make_tensor_coordinate(DstDesc{}, Index{})); - - using SrcCoordStep = decltype(make_tensor_coordinate_step(SrcDesc{}, Index{})); - using DstCoordStep = decltype(make_tensor_coordinate_step(DstDesc{}, Index{})); - - __device__ constexpr ThreadwiseTensorSliceTransfer_v3(const SrcDesc& src_desc, - const Index& src_slice_origin, - const DstDesc& dst_desc, - const Index& dst_slice_origin) - : src_coord_(make_tensor_coordinate(src_desc, src_slice_origin)), - dst_coord_(make_tensor_coordinate(dst_desc, dst_slice_origin)) - { - static_assert(SliceLengths::At(Number{}) % SrcScalarPerVector == 0, - "wrong! Not divisible"); - static_assert(SliceLengths::At(Number{}) % DstScalarPerVector == 0, - "wrong! Not divisible"); - } - - __device__ void SetSrcSliceOrigin(const SrcDesc& src_desc, const Index& src_slice_origin_idx) - { - src_coord_ = make_tensor_coordinate(src_desc, src_slice_origin_idx); - } - - __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 SrcDesc& src_desc, const SrcBuffer& src_buf, const SrcStepHacks& src_step_hacks) - { - static_assert(SrcBuffer::GetAddressSpace() == AddressSpaceEnum::Global or - SrcBuffer::GetAddressSpace() == AddressSpaceEnum::Lds, - "wrong!"); - - static_assert( - is_same, remove_cvref_t>::value, - "wrong! SrcBuffer and SrcData data type are inconsistent"); - - constexpr auto I0 = Number<0>{}; - constexpr auto I1 = Number<1>{}; - - // scalar per access on each dim - // TODO: don't use lambda_scalar_per_access - constexpr auto src_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto src_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - constexpr auto src_access_lengths = SliceLengths{} / src_scalar_per_access; - - constexpr auto src_dim_access_order = SrcDimAccessOrder{}; - - constexpr auto ordered_src_access_lengths = - container_reorder_given_new2old(src_access_lengths, src_dim_access_order); - - // make forward steps - const auto src_forward_steps = generate_tuple( - [&](auto i) { - Index forward_step_idx; - - static_for<0, nDim, 1>{}([&](auto j) { - forward_step_idx(j) = (i.value == j.value) ? src_scalar_per_access[i] : 0; - }); - - return make_tensor_coordinate_step( - src_desc, forward_step_idx, src_step_hacks[I0][i]); - }, - Number{}); - - // make backward steps - const auto src_backward_steps = generate_tuple( - [&](auto i) { - Index backward_step_idx; - - static_for<0, nDim, 1>{}([&](auto j) { - backward_step_idx(j) = (i.value == j.value) ? -src_scalar_per_access[i] : 0; - }); - - return make_tensor_coordinate_step( - src_desc, backward_step_idx, src_step_hacks[I1][i]); - }, - Number{}); - - // loop over tensor and copy - static_ford{}([&](auto ordered_src_access_idx) { - // judge move forward or move backward - constexpr auto forward_sweep = [&]() { - StaticallyIndexedArray forward_sweep_; - - forward_sweep_(I0) = true; - - static_for<1, nDim, 1>{}([&](auto i) { - index_t tmp = ordered_src_access_idx[I0]; - - static_for<1, i, 1>{}([&](auto j) { - tmp = tmp * ordered_src_access_lengths[j] + ordered_src_access_idx[j]; - }); - - forward_sweep_(i) = tmp % 2 == 0; - }); - - return forward_sweep_; - }(); - - // calculate src data index - constexpr auto src_data_idx = [&]() { - Index ordered_idx; - - static_for<0, nDim, 1>{}([&](auto i) { - ordered_idx(i) = forward_sweep[i] ? ordered_src_access_idx[i] - : ordered_src_access_lengths[i] - 1 - - ordered_src_access_idx[i]; - }); - - return container_reorder_given_old2new(ordered_idx, src_dim_access_order) * - src_scalar_per_access; - }(); - - vector_type_maker_t src_tmp_vector; - - using src_vector_t = typename decltype(src_tmp_vector)::type; - - const bool is_src_valid = - coordinate_has_valid_offset_assuming_visible_index_is_valid(src_desc, src_coord_); - - // copy data from src_buf to src_tmp_vector - src_tmp_vector.template AsType()(Number<0>{}) = - src_buf.template Get(src_coord_.GetOffset(), is_src_valid); - - // copy data from src_tmp_vector to buffer_ - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t buffer_offset = - buffer_desc_.CalculateOffset(src_data_idx + i * src_scalar_step_in_vector); - - buffer_(Number{}) = src_tmp_vector.template AsType()[i]; - }); - - constexpr auto move_on_dim = [&]() constexpr - { - StaticallyIndexedArray move_on_dim_; - - static_for<0, nDim, 1>{}([&](auto i) { - move_on_dim_(i) = ordered_src_access_idx[i] < ordered_src_access_lengths[i] - 1; - - static_for{}([&](auto j) { - move_on_dim_(i) &= - ordered_src_access_idx[j] == ordered_src_access_lengths[j] - 1; - }); - }); - - return move_on_dim_; - } - (); - - // move - static_for<0, nDim, 1>{}([&](auto i) { - if constexpr(move_on_dim[i]) - { - if constexpr(forward_sweep[i]) - { - move_tensor_coordinate( - src_desc, src_coord_, src_forward_steps[src_dim_access_order[i]]); - } - else - { - move_tensor_coordinate( - src_desc, src_coord_, src_backward_steps[src_dim_access_order[i]]); - } - } - }); - }); - - // move src coordinate back to slice origin (or not) - if constexpr(SrcResetCoordinateAfterRun) - { - const auto src_reset_step = - make_tensor_coordinate_step(src_desc, GetSrcCoordinateResetStep()); - - move_tensor_coordinate(src_desc, src_coord_, src_reset_step); - } - } - - template - __device__ void - RunWrite(const DstDesc& dst_desc, DstBuffer& dst_buf, const DstStepHacks& dst_step_hacks) - { - static_assert(DstBuffer::GetAddressSpace() == AddressSpaceEnum::Global or - DstBuffer::GetAddressSpace() == AddressSpaceEnum::Lds, - "wrong!"); - - static_assert( - is_same, remove_cvref_t>::value, - "wrong! SrcBuffer or DstBuffer data type is wrong"); - - constexpr auto I0 = Number<0>{}; - constexpr auto I1 = Number<1>{}; - - // src scalar per access on each dim - // TODO: don't use this - constexpr auto dst_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto dst_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - constexpr auto dst_access_lengths = SliceLengths{} / dst_scalar_per_access; - - constexpr auto dst_dim_access_order = DstDimAccessOrder{}; - - constexpr auto ordered_dst_access_lengths = - container_reorder_given_new2old(dst_access_lengths, dst_dim_access_order); - - // make forward steps - const auto dst_forward_steps = generate_tuple( - [&](auto i) { - Index forward_step_idx; - - static_for<0, nDim, 1>{}([&](auto j) { - forward_step_idx(j) = (i.value == j.value) ? dst_scalar_per_access[i] : 0; - }); - - return make_tensor_coordinate_step( - dst_desc, forward_step_idx, dst_step_hacks[I0][i]); - }, - Number{}); - - // make backward steps - const auto dst_backward_steps = generate_tuple( - [&](auto i) { - Index backward_step_idx; - - static_for<0, nDim, 1>{}([&](auto j) { - backward_step_idx(j) = (i.value == j.value) ? -dst_scalar_per_access[i] : 0; - }); - - return make_tensor_coordinate_step( - dst_desc, backward_step_idx, dst_step_hacks[I1][i]); - }, - Number{}); - - // loop over tensor and copy - static_ford{}([&](auto ordered_dst_access_idx) { - // judge move forward or move backward - constexpr auto forward_sweep = [&]() { - StaticallyIndexedArray forward_sweep_; - - forward_sweep_(I0) = true; - - static_for<1, nDim, 1>{}([&](auto i) { - index_t tmp = ordered_dst_access_idx[I0]; - - static_for<1, i, 1>{}([&](auto j) { - tmp = tmp * ordered_dst_access_lengths[j] + ordered_dst_access_idx[j]; - }); - - forward_sweep_(i) = tmp % 2 == 0; - }); - - return forward_sweep_; - }(); - - // calculate dst data index - constexpr auto dst_data_idx = [&]() { - Index ordered_idx; - - static_for<0, nDim, 1>{}([&](auto i) { - ordered_idx(i) = forward_sweep[i] ? ordered_dst_access_idx[i] - : ordered_dst_access_lengths[i] - 1 - - ordered_dst_access_idx[i]; - }); - - return container_reorder_given_old2new(ordered_idx, dst_dim_access_order) * - dst_scalar_per_access; - }(); - - vector_type_maker_t dst_tmp_vector; - - // copy data from buffer_ to dst_tmp_vector - static_for<0, DstScalarPerVector, 1>{}([&](auto i) { - constexpr index_t buffer_offset = - buffer_desc_.CalculateOffset(dst_data_idx + i * dst_scalar_step_in_vector); - - dst_tmp_vector.template AsType()(i) = - type_convert(buffer_[Number{}]); - }); - - using dst_vector_t = typename decltype(dst_tmp_vector)::type; - - // copy data from dst_tmp_vector to dst_buf - const bool is_dst_valid = - coordinate_has_valid_offset_assuming_visible_index_is_valid(dst_desc, dst_coord_); - - dst_buf.template Set( - dst_coord_.GetOffset(), - is_dst_valid, - dst_tmp_vector.template AsType()[Number<0>{}]); - - constexpr auto move_on_dim = [&]() constexpr - { - StaticallyIndexedArray move_on_dim_; - - static_for<0, nDim, 1>{}([&](auto i) { - move_on_dim_(i) = ordered_dst_access_idx[i] < ordered_dst_access_lengths[i] - 1; - - static_for{}([&](auto j) { - move_on_dim_(i) &= - ordered_dst_access_idx[j] == ordered_dst_access_lengths[j] - 1; - }); - }); - - return move_on_dim_; - } - (); - - // move - static_for<0, nDim, 1>{}([&](auto i) { - if constexpr(move_on_dim[i]) - { - if constexpr(forward_sweep[i]) - { - move_tensor_coordinate( - dst_desc, dst_coord_, dst_forward_steps[dst_dim_access_order[i]]); - } - else - { - move_tensor_coordinate( - dst_desc, dst_coord_, dst_backward_steps[dst_dim_access_order[i]]); - } - } - }); - }); - - // 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); - } - } - - template - __device__ void RunRead(const SrcDesc& src_desc, const SrcBuffer& src_buf) - { - constexpr index_t ntransform_src = SrcDesc::GetNumOfTransform(); - - constexpr auto zeros = typename uniform_sequence_gen::type{}; - - constexpr auto src_step_hacks = - make_tuple(generate_tuple([&](auto) { return zeros; }, Number{}), - generate_tuple([&](auto) { return zeros; }, Number{})); - - RunRead(src_desc, src_buf, src_step_hacks); - } - - template - __device__ void RunWrite(const DstDesc& dst_desc, DstBuffer& dst_buf) - { - constexpr index_t ntransform_dst = DstDesc::GetNumOfTransform(); - - constexpr auto zeros = typename uniform_sequence_gen::type{}; - - constexpr auto dst_step_hacks = - make_tuple(generate_tuple([&](auto) { return zeros; }, Number{}), - generate_tuple([&](auto) { return zeros; }, Number{})); - - RunWrite(dst_desc, dst_buf, dst_step_hacks); - } - - __device__ static constexpr auto GetSrcCoordinateResetStep() - { - constexpr auto I0 = Number<0>{}; - - // scalar per access on each dim - // TODO: don't use lambda_scalar_per_access - constexpr auto src_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto src_access_lengths = SliceLengths{} / src_scalar_per_access; - - constexpr auto src_dim_access_order = SrcDimAccessOrder{}; - - constexpr auto ordered_src_access_lengths = - container_reorder_given_new2old(src_access_lengths, src_dim_access_order); - - // judge move forward or move backward during the last iteration - constexpr auto forward_sweep = [&]() { - StaticallyIndexedArray forward_sweep_; - - forward_sweep_(I0) = true; - - static_for<1, nDim, 1>{}([&](auto i) { - index_t tmp = ordered_src_access_lengths[I0] - 1; - - static_for<1, i, 1>{}([&](auto j) { - tmp = tmp * ordered_src_access_lengths[j] + ordered_src_access_lengths[j] - 1; - }); - - forward_sweep_(i) = tmp % 2 == 0; - }); - - return forward_sweep_; - }(); - - // calculate src data index after last iteration in RunRead(), if it has not being reset by - // RunRead() - constexpr auto src_data_idx = [&]() { - Index ordered_idx; - - static_for<0, nDim, 1>{}([&](auto i) { - ordered_idx(i) = forward_sweep[i] ? ordered_src_access_lengths[i] - 1 : 0; - }); - - return container_reorder_given_old2new(ordered_idx, src_dim_access_order) * - src_scalar_per_access; - }(); - - // - constexpr auto reset_src_data_step = [&]() { - Index reset_src_data_step_; - - static_for<0, nDim, 1>{}([&](auto i) { reset_src_data_step_(i) = -src_data_idx[i]; }); - - return reset_src_data_step_; - }(); - - return reset_src_data_step; - } - - __device__ static constexpr auto GetDstCoordinateResetStep() - { - constexpr auto I0 = Number<0>{}; - - // scalar per access on each dim - // TODO: don't use lambda_scalar_per_access - constexpr auto dst_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto dst_access_lengths = SliceLengths{} / dst_scalar_per_access; - - constexpr auto dst_dim_access_order = DstDimAccessOrder{}; - - constexpr auto ordered_dst_access_lengths = - container_reorder_given_new2old(dst_access_lengths, dst_dim_access_order); - - // judge move forward or move backward during the last iteration - constexpr auto forward_sweep = [&]() { - StaticallyIndexedArray forward_sweep_; - - forward_sweep_(I0) = true; - - static_for<1, nDim, 1>{}([&](auto i) { - index_t tmp = ordered_dst_access_lengths[I0] - 1; - - static_for<1, i, 1>{}([&](auto j) { - tmp = tmp * ordered_dst_access_lengths[j] + ordered_dst_access_lengths[j] - 1; - }); - - forward_sweep_(i) = tmp % 2 == 0; - }); - - return forward_sweep_; - }(); - - // calculate dst data index after last iteration in RunWrite(), if it has not being reset by - // RunWrite() - constexpr auto dst_data_idx = [&]() { - Index ordered_idx; - - static_for<0, nDim, 1>{}([&](auto i) { - ordered_idx(i) = forward_sweep[i] ? ordered_dst_access_lengths[i] - 1 : 0; - }); - - return container_reorder_given_old2new(ordered_idx, dst_dim_access_order) * - dst_scalar_per_access; - }(); - - // - constexpr auto reset_dst_data_step = [&]() { - Index reset_dst_data_step_; - - static_for<0, nDim, 1>{}([&](auto i) { reset_dst_data_step_(i) = -dst_data_idx[i]; }); - - return reset_dst_data_step_; - }(); - - return reset_dst_data_step; - } - - // src_slice_origin_step_idx need to be known at compile-time, for performance reason - __device__ void MoveSrcSliceWindow(const SrcDesc& src_desc, - const Index& src_slice_origin_step_idx) - { - // if src coord was not reset by RunRead(), then need to adjust the step here - const auto adjusted_step_idx = - SrcResetCoordinateAfterRun ? src_slice_origin_step_idx - : src_slice_origin_step_idx + GetSrcCoordinateResetStep(); - - // is it OK to construct a new step every time? - const auto adjusted_step = make_tensor_coordinate_step(src_desc, adjusted_step_idx); - - move_tensor_coordinate(src_desc, src_coord_, adjusted_step); - } - - // src_slice_origin_step_idx need to be known at compile-time, for performance reason - template - __device__ void - MoveSrcSliceWindow(const SrcDesc& src_desc, - const Index& src_slice_origin_step_idx, - const SrcMoveSliceWindowStepHack& src_move_slice_window_step_hack) - { - // if src coord was not reset by RunRead(), then need to adjust the step here - const auto adjusted_step_idx = - SrcResetCoordinateAfterRun ? src_slice_origin_step_idx - : src_slice_origin_step_idx + GetSrcCoordinateResetStep(); - - // is it OK to construct a new step every time? - const auto adjusted_step = make_tensor_coordinate_step( - src_desc, adjusted_step_idx, src_move_slice_window_step_hack); - - move_tensor_coordinate(src_desc, src_coord_, adjusted_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 RunWrite(), 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: - static constexpr auto buffer_desc_ = - make_naive_tensor_descriptor_packed(sequence_to_tuple_of_number(SliceLengths{})); - - static constexpr auto buffer_size_ = buffer_desc_.GetElementSpaceSize(); - - StaticBuffer buffer_; - - SrcCoord src_coord_; - DstCoord dst_coord_; -}; - -// Assume: -// 1. src: -// 1. SrcDesc is known at compile-time -// 2. SrcBuffer is DynamicBuffer -// 3. src_ref_idx is known at run-time -// 4. SrcRefToOriginDisplacement is known at compile-time -// 5. use #-step -// 2. dst: -// 1. DstDesc is known at compile-time -// 2. DstBuffer is StaticBuffer -// 3. DstOriginIdx is known at compile-time -// 4. use direct address calculation -// 3. vector access on src -template ::type = false> -struct ThreadwiseTensorSliceTransfer_v4 -{ - static constexpr index_t nDim = SliceLengths::Size(); - - using Index = MultiIndex; - - using SrcCoord = decltype(make_tensor_coordinate(SrcDesc{}, Index{})); - - using SrcCoordStep = decltype(make_tensor_coordinate_step(SrcDesc{}, Index{})); - - static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t>) - return 2; - else - return 1; - }(); - - __device__ constexpr ThreadwiseTensorSliceTransfer_v4(const Index& src_ref_idx) - : src_ref_coord_(make_tensor_coordinate(SrcDesc{}, src_ref_idx)) - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! SrcDesc and DstDesc need to known at compile-time"); - - static_assert(SliceLengths::At(Number{}) % SrcScalarPerVector == 0, - "wrong! Not divisible"); - - if constexpr(is_same_v, pk_i4_t>) - { - static_assert(SrcScalarPerVector % PackedSize == 0, "pk data N cannot be 1"); - } - } - - template - __device__ void Run(const SrcDesc&, - const SrcRefToOriginDisplacement&, - const SrcBuffer& src_buf, - const DstDesc&, - const DstOriginIdx&, - DstBuffer& dst_buf) const - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! SrcDesc and DstDesc need to known at compile-time"); - - static_assert( - is_same, remove_cvref_t>::value && - is_same, remove_cvref_t>::value, - "wrong! SrcBuffer or DstBuffer data type is wrong"); - - static_assert(DstBuffer::IsStaticBuffer(), "wrong! DstBuffer need to be StaticBuffer"); - - static_assert(is_known_at_compile_time>::value && - is_known_at_compile_time>::value, - "wrong! SrcOriginToRefDistance and DstOriginToRefDistance need to be known " - "at compile-time"); - - // SrcDesc and DstDesc are known at compile-time - constexpr auto src_desc = remove_cvref_t{}; - constexpr auto dst_desc = remove_cvref_t{}; - - // SrcOriginToRefDisttance and DstOriginToRefDistance are known at compile-time - constexpr auto src_ref_to_origin_disp_idx = to_multi_index(SrcRefToOriginDisplacement{}); - constexpr auto dst_origin_idx = to_multi_index(DstOriginIdx{}); - - // scalar per access of each dim - constexpr auto src_scalar_per_access = generate_sequence_v2( - [&](auto i) constexpr { - if constexpr(i == SrcVectorDim) - { - return Number{}; - } - else - { - return Number<1>{}; - } - }, - Number{}); - - // scalar step (if steping on SrcVectorDim) of each dim - constexpr auto src_scalar_step_in_vector = generate_sequence_v2( - [&](auto i) constexpr { - if constexpr(i == SrcVectorDim) - { - return Number<1>{}; - } - else - { - return Number<0>{}; - } - }, - Number{}); - - constexpr auto access_lengths = SliceLengths{} / src_scalar_per_access; - - constexpr auto dim_access_order = DimAccessOrder{}; - - constexpr auto ordered_access_lengths = - container_reorder_given_new2old(access_lengths, dim_access_order); - - static_ford{}([&](auto ordered_access_idx) { -#if 0 - // TODO: unable to compile - // position in slice window - constexpr auto data_to_origin_disp_idx = - container_reorder_given_old2new(ordered_access_idx, dim_access_order) * - src_scalar_per_access; -#else - // position in slice window - constexpr auto data_to_origin_disp_idx = - ordered_access_idx.ReorderGivenOld2New(dim_access_order) * src_scalar_per_access; -#endif - // src coordinate - constexpr auto src_ref_to_data_disp_idx = - src_ref_to_origin_disp_idx + data_to_origin_disp_idx; - - constexpr auto src_ref_to_data_disp_coord_step = - make_tensor_coordinate_step(src_desc, src_ref_to_data_disp_idx); - - auto src_data_coord = src_ref_coord_; - - move_tensor_coordinate(src_desc, src_data_coord, src_ref_to_data_disp_coord_step); - - vector_type_maker_t src_tmp_vector; - - using src_vector_t = typename decltype(src_tmp_vector)::type; - - const bool is_src_valid = coordinate_has_valid_offset_assuming_visible_index_is_valid( - src_desc, src_data_coord); - - // copy data from src_buf into src_tmp_vector - if constexpr(SrcBuffer::IsDynamicBuffer()) - { - src_tmp_vector.template AsType()(Number<0>{}) = - src_buf.template Get(src_data_coord.GetOffset() / PackedSize, - is_src_valid); - } - else if constexpr(SrcBuffer::IsStaticBuffer()) - { - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t src_offset = src_desc.CalculateOffset( - src_ref_to_origin_disp_idx + data_to_origin_disp_idx + - i * src_scalar_step_in_vector); - - src_tmp_vector.template AsType()(i) = src_buf[Number{}]; - }); - } - - if constexpr(is_same, pk_i4_t>::value) - { - // copy data from src_tmp_vector to dst_tmp_vector (data cast data from SrcData to - // DstData) - vector_type_maker_t dst_tmp_vector; - - constexpr index_t pack_size = 8; - - static_assert(SrcScalarPerVector % pack_size == 0, ""); - - using src_v_t = typename vector_type_maker_t::type; - using dst_v_t = typename vector_type_maker_t::type; - - static_for<0, SrcScalarPerVector / pack_size, 1>{}([&](auto i) { - ck::tensor_operation::element_wise::PassThroughPack8{}( - dst_tmp_vector.template AsType()(i), - src_tmp_vector.template AsType()[i]); - }); - - // copy data from dst_tmp_vector into dst_buf - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); - - dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; - }); - } - else if constexpr(is_same, f8_t>::value && - is_same, half_t>::value && - SrcScalarPerVector % 2 == 0) - { - // copy data from src_tmp_vector to dst_tmp_vector (data cast data from SrcData to - // DstData) - vector_type_maker_t dst_tmp_vector; - - constexpr index_t pack_size = 2; - - using dst_v_t = typename vector_type_maker_t::type; - using src_v_t = typename vector_type_maker_t::type; - static_for<0, SrcScalarPerVector / pack_size, 1>{}([&](auto i) { - ck::tensor_operation::element_wise::PassThroughPack2{}( - dst_tmp_vector.template AsType()(i), - src_tmp_vector.template AsType()[i]); - }); - - // copy data from dst_tmp_vector into dst_buf - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); - - dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; - }); - } - else - { - // copy data from src_tmp_vector to dst_tmp_vector (data cast data from SrcData to - // DstData) - vector_type_maker_t dst_tmp_vector; - - // TODO: if SrcData and DstData are vetor type, then static_cast may not compile - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - dst_tmp_vector.template AsType()(i) = - type_convert(src_tmp_vector.template AsType()[i]); - }); - - // copy data from dst_tmp_vector into dst_buf - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); - - dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; - }); - } - }); - } - - // Fuse scale - template - __device__ void Run(const SrcDesc&, - const SrcRefToOriginDisplacement&, - const SrcBuffer& src_buf, - const DstData& scale, - const DstDesc&, - const DstOriginIdx&, - DstBuffer& dst_buf) const - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! SrcDesc and DstDesc need to known at compile-time"); - - static_assert( - is_same, remove_cvref_t>::value && - is_same, remove_cvref_t>::value, - "wrong! SrcBuffer or DstBuffer data type is wrong"); - - static_assert(DstBuffer::IsStaticBuffer(), "wrong! DstBuffer need to be StaticBuffer"); - - static_assert(is_known_at_compile_time>::value && - is_known_at_compile_time>::value, - "wrong! SrcOriginToRefDistance and DstOriginToRefDistance need to be known " - "at compile-time"); - - // SrcDesc and DstDesc are known at compile-time - constexpr auto src_desc = remove_cvref_t{}; - constexpr auto dst_desc = remove_cvref_t{}; - - // SrcOriginToRefDisttance and DstOriginToRefDistance are known at compile-time - constexpr auto src_ref_to_origin_disp_idx = to_multi_index(SrcRefToOriginDisplacement{}); - constexpr auto dst_origin_idx = to_multi_index(DstOriginIdx{}); - - // scalar per access of each dim - constexpr auto src_scalar_per_access = generate_sequence_v2( - [&](auto i) constexpr { - if constexpr(i == SrcVectorDim) - { - return Number{}; - } - else - { - return Number<1>{}; - } - }, - Number{}); - - // scalar step (if steping on SrcVectorDim) of each dim - constexpr auto src_scalar_step_in_vector = generate_sequence_v2( - [&](auto i) constexpr { - if constexpr(i == SrcVectorDim) - { - return Number<1>{}; - } - else - { - return Number<0>{}; - } - }, - Number{}); - - constexpr auto access_lengths = SliceLengths{} / src_scalar_per_access; - - constexpr auto dim_access_order = DimAccessOrder{}; - - constexpr auto ordered_access_lengths = - container_reorder_given_new2old(access_lengths, dim_access_order); - - static_ford{}([&](auto ordered_access_idx) { -#if 0 - // TODO: unable to compile - // position in slice window - constexpr auto data_to_origin_disp_idx = - container_reorder_given_old2new(ordered_access_idx, dim_access_order) * - src_scalar_per_access; -#else - // position in slice window - constexpr auto data_to_origin_disp_idx = - ordered_access_idx.ReorderGivenOld2New(dim_access_order) * src_scalar_per_access; -#endif - // src coordinate - constexpr auto src_ref_to_data_disp_idx = - src_ref_to_origin_disp_idx + data_to_origin_disp_idx; - - constexpr auto src_ref_to_data_disp_coord_step = - make_tensor_coordinate_step(src_desc, src_ref_to_data_disp_idx); - - auto src_data_coord = src_ref_coord_; - - move_tensor_coordinate(src_desc, src_data_coord, src_ref_to_data_disp_coord_step); - - vector_type_maker_t src_tmp_vector; - - using src_vector_t = typename decltype(src_tmp_vector)::type; - - const bool is_src_valid = coordinate_has_valid_offset_assuming_visible_index_is_valid( - src_desc, src_data_coord); - - // copy data from src_buf into src_tmp_vector - if constexpr(SrcBuffer::IsDynamicBuffer()) - { - src_tmp_vector.template AsType()(Number<0>{}) = - src_buf.template Get(src_data_coord.GetOffset() / PackedSize, - is_src_valid); - } - else if constexpr(SrcBuffer::IsStaticBuffer()) - { - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t src_offset = src_desc.CalculateOffset( - src_ref_to_origin_disp_idx + data_to_origin_disp_idx + - i * src_scalar_step_in_vector); - - src_tmp_vector.template AsType()(i) = src_buf[Number{}]; - }); - } - - if constexpr(is_same, pk_i4_t>::value) - { - // copy data from src_tmp_vector to dst_tmp_vector (data cast data from SrcData to - // DstData) - vector_type_maker_t dst_tmp_vector; - vector_type scale_vector; - scale_vector.template AsType()(Number<0>{}) = scale; - scale_vector.template AsType()(Number<1>{}) = scale; - - constexpr index_t pack_size = 8; - - static_assert(SrcScalarPerVector % pack_size == 0, ""); - - using src_v_t = typename vector_type_maker_t::type; - using dst_v_t = typename vector_type_maker_t::type; - using scale_v_t = typename vector_type_maker_t::type; - - static_for<0, SrcScalarPerVector / pack_size, 1>{}([&](auto i) { - ck::tensor_operation::element_wise::DequantPack8{}( - dst_tmp_vector.template AsType()(i), - src_tmp_vector.template AsType()[i], - scale_vector.template AsType()[Number<0>{}]); - }); - - // copy data from dst_tmp_vector into dst_buf - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); - - dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; - }); - } - else if constexpr(is_same, f8_t>::value && - is_same, half_t>::value && - SrcScalarPerVector % 2 == 0) - { - // copy data from src_tmp_vector to dst_tmp_vector (data cast data from SrcData to - // DstData) - vector_type_maker_t dst_tmp_vector; - - constexpr index_t pack_size = 2; - - using dst_v_t = typename vector_type_maker_t::type; - using src_v_t = typename vector_type_maker_t::type; - static_for<0, SrcScalarPerVector / pack_size, 1>{}([&](auto i) { - ck::tensor_operation::element_wise::PassThroughPack2{}( - dst_tmp_vector.template AsType()(i), - src_tmp_vector.template AsType()[i]); - }); - - // copy data from dst_tmp_vector into dst_buf - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); - - dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; - }); - } - else - { - // copy data from src_tmp_vector to dst_tmp_vector (data cast data from SrcData to - // DstData) - vector_type_maker_t dst_tmp_vector; - - // TODO: if SrcData and DstData are vetor type, then static_cast may not compile - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - dst_tmp_vector.template AsType()(i) = - type_convert(src_tmp_vector.template AsType()[i]); - }); - - // copy data from dst_tmp_vector into dst_buf - static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); - - dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; - }); - } - }); - } - - template - __device__ void MoveSrcSliceWindow(const SrcDesc&, - const SrcSliceMoveStepIdx& src_slice_move_step_idx) - { - constexpr auto src_desc = SrcDesc{}; - - const auto src_slice_move_step_iter = - make_tensor_coordinate_step(src_desc, to_multi_index(src_slice_move_step_idx)); - - move_tensor_coordinate(SrcDesc{}, src_ref_coord_, src_slice_move_step_iter); - } - __device__ void SetSrcCoord(const Index& src_ref_idx) - { - src_ref_coord_ = make_tensor_coordinate(SrcDesc{}, src_ref_idx); - } - - private: - SrcCoord src_ref_coord_; -}; - -/** - * @brief Threadwise data transfer - * - * Do NOT involve any tensor coordinates with StaticBuffer - * - */ -template ::type = false> -struct ThreadwiseTensorSliceTransfer_StaticToStatic -{ - static constexpr index_t nDim = SliceLengths::Size(); - - using Index = MultiIndex; - - static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t>) - return 2; - else - return 1; - }(); - - __device__ constexpr ThreadwiseTensorSliceTransfer_StaticToStatic( - const ElementwiseOperation& element_op) - : element_op_{element_op} - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! Desc need to known at compile-time"); - - static_assert(SliceLengths::At(Number{}) % DstScalarPerVector == 0, - "wrong! Not divisible"); - } - - template - __device__ void Run(const SrcDesc&, - const SrcSliceOriginIdx&, - const SrcBuffer& src_buf, - const DstDesc&, - const DstSliceOriginIdx&, - DstBuffer& dst_buf) const - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! Desc need to known at compile-time"); - - static_assert(is_known_at_compile_time>::value && - is_known_at_compile_time>::value, - "wrong! SliceOrigin need to known at compile-time"); - - static_assert(SrcBuffer::IsStaticBuffer() && DstBuffer::IsStaticBuffer(), - "wrong! Buffer need to be StaticBuffer"); - - // SrcDesc and src_slice_origin_idx are known at compile-time - constexpr auto src_desc = remove_cvref_t{}; - constexpr auto dst_desc = remove_cvref_t{}; - constexpr auto src_slice_origin_idx = to_multi_index(SrcSliceOriginIdx{}); - constexpr auto dst_slice_origin_idx = to_multi_index(DstSliceOriginIdx{}); - - // scalar per access on each dim - constexpr auto dst_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto dst_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - using SpaceFillingCurve = SpaceFillingCurve>; - - static_assert(DstScalarPerVector == SpaceFillingCurve::ScalarPerVector, - "wrong!DstScalarPerVector != SpaceFillingCurve::ScalarPerVector"); - - constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); - - if constexpr(is_same, pk_i4_t>::value) - { - static_for<0, num_access, 1>{}([&](auto idx_1d) { - typename vector_type_maker::type - src_tmp_vector; - - constexpr auto idx_md = SpaceFillingCurve::GetIndex(idx_1d); - - // copy data from src_buf into dst_vector - static_for<0, DstScalarPerVector / PackedSize, 1>{}([&](auto i) { - constexpr index_t src_offset = src_desc.CalculateOffset( - src_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - src_tmp_vector.template AsType()(i) = src_buf[Number{}]; - }); - - // copy data from src_tmp_vector to dst_tmp_vector (data cast data from SrcData to - // DstData) - vector_type_maker_t dst_tmp_vector; - - constexpr index_t pack_size = 8; - - static_assert(DstScalarPerVector % pack_size == 0, ""); - - using src_v_t = typename vector_type_maker_t::type; - using dst_v_t = typename vector_type_maker_t::type; - - static_for<0, DstScalarPerVector / pack_size, 1>{}([&](auto i) { - ck::tensor_operation::element_wise::PassThroughPack8{}( - dst_tmp_vector.template AsType()(i), - src_tmp_vector.template AsType()[i]); - }); - - // copy data from dst_tmp_vector into dst_buf - static_for<0, DstScalarPerVector, 1>{}([&](auto i) { - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; - }); - }); - } - else - { - static_for<0, num_access, 1>{}([&](auto idx_1d) { - constexpr auto idx_md = SpaceFillingCurve::GetIndex(idx_1d); - - // copy data from src_buf into dst_vector - static_for<0, DstScalarPerVector, 1>{}([&](auto i) { - constexpr index_t src_offset = src_desc.CalculateOffset( - src_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - DstData v; - - // apply element-wise operation - element_op_(v, src_buf[Number{}]); - - // apply type convert - dst_buf(Number{}) = v; - }); - }); - } - } - - ElementwiseOperation element_op_; -}; - -// Specialized for gfx11 -// A single Wave32 is composed by double row -// Data exchange allowed between these two rows -// This RowLane Dst buf will be filled from two Src buf -// SrcA: From specific thread buffer hold by This RowLane on This Row -// SrcB: From specific thread buffer hold by This RowLane on The other Row -template ::type = false> -struct ThreadwiseTensorSliceTransfer_StaticToStatic_InterRow -{ - static constexpr index_t nDim = SliceLengths::Size(); - - using Index = MultiIndex; - - __device__ constexpr ThreadwiseTensorSliceTransfer_StaticToStatic_InterRow(const Index& src_idx) - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! Desc need to known at compile-time"); - - static_assert(SliceLengths::At(Number{}) % DstScalarPerVector == 0, - "wrong! Not divisible"); - ignore = src_idx; - } - - template - __device__ void Run(const SrcDesc&, - const SrcSliceOriginIdx&, - const SrcBuffer& src_buf, - const DstDesc&, - const DstSliceOriginIdx&, - DstBuffer& dst_buf) const - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! Desc need to known at compile-time"); - - static_assert(is_known_at_compile_time>::value && - is_known_at_compile_time>::value, - "wrong! SliceOrigin need to known at compile-time"); - - static_assert(SrcBuffer::IsStaticBuffer() && DstBuffer::IsStaticBuffer(), - "wrong! Buffer need to be StaticBuffer"); - - // SrcDesc and src_slice_origin_idx are known at compile-time - constexpr auto src_desc = remove_cvref_t{}; - constexpr auto dst_desc = remove_cvref_t{}; - constexpr auto src_slice_origin_idx = to_multi_index(SrcSliceOriginIdx{}); - constexpr auto dst_slice_origin_idx = to_multi_index(DstSliceOriginIdx{}); - - // scalar per access on each dim - constexpr auto dst_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto dst_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - using SpaceFillingCurve = SpaceFillingCurve>; - - static_assert(DstScalarPerVector == SpaceFillingCurve::ScalarPerVector, - "wrong!DstScalarPerVector != SpaceFillingCurve::ScalarPerVector"); - - constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); - - static_for<0, num_access, 1>{}([&](auto idx_1d) { - constexpr auto idx_md = SpaceFillingCurve::GetIndex(idx_1d); - - // copy data from src_buf into dst_vector - static_for<0, DstScalarPerVector, 1>{}([&](auto i) { - // src_desc error, non constexpr, caused by merge transform - constexpr index_t src_offset = src_desc.CalculateOffset( - src_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - SrcData v_this_row, v_theother_row; - // int type temp value due to intrinsic requirement - int temp = 0; - - // apply element-wise operation - element_op_(v_this_row, src_buf[Number{}]); - - // apply intra-row permute. - if constexpr(IntraRowSwizzlePerm) - { - temp = __builtin_amdgcn_permlane16( - temp, type_convert_sp(v_this_row), 0xb3a29180, 0xf7e6d5c4, 1, 0); - v_this_row = type_convert_sp(temp); - } - - // apply inter-row permute. - temp = __builtin_amdgcn_permlanex16(temp, - type_convert_sp(v_this_row), - LowEightRowlaneIdx, - HighEightRowLaneIdx, - 1, - 0); - v_theother_row = type_convert_sp(temp); - - if(get_thread_local_1d_id() % 32 < 16) - { - // apply type convert - dst_buf(Number{}) = type_convert_sp(v_this_row); - dst_buf(Number{}) = - type_convert_sp(v_theother_row); - } - else - { - // apply type convert - dst_buf(Number{}) = - type_convert_sp(v_this_row); - dst_buf(Number{}) = type_convert_sp(v_theother_row); - } - }); - }); - } - ElementwiseOperation element_op_{}; -}; - -// Specialized for gfx12 -template ::type = false> -struct ThreadwiseTensorSliceTransfer_StaticToStatic_IntraRow -{ - static constexpr index_t nDim = SliceLengths::Size(); - - using Index = MultiIndex; - - __device__ constexpr ThreadwiseTensorSliceTransfer_StaticToStatic_IntraRow(const Index& src_idx) - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! Desc need to known at compile-time"); - - static_assert(SliceLengths::At(Number{}) % DstScalarPerVector == 0, - "wrong! Not divisible"); - ignore = src_idx; - } - - template - __device__ void Run(const SrcDesc&, - const SrcSliceOriginIdx&, - const SrcBuffer& src_buf, - const DstDesc&, - const DstSliceOriginIdx&, - DstBuffer& dst_buf) const - { - static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), - "wrong! Desc need to known at compile-time"); - - static_assert(is_known_at_compile_time>::value && - is_known_at_compile_time>::value, - "wrong! SliceOrigin need to known at compile-time"); - - static_assert(SrcBuffer::IsStaticBuffer() && DstBuffer::IsStaticBuffer(), - "wrong! Buffer need to be StaticBuffer"); - - // SrcDesc and src_slice_origin_idx are known at compile-time - constexpr auto src_desc = remove_cvref_t{}; - constexpr auto dst_desc = remove_cvref_t{}; - constexpr auto src_slice_origin_idx = to_multi_index(SrcSliceOriginIdx{}); - constexpr auto dst_slice_origin_idx = to_multi_index(DstSliceOriginIdx{}); - - // scalar per access on each dim - constexpr auto dst_scalar_per_access = generate_sequence( - detail::lambda_scalar_per_access{}, Number{}); - - constexpr auto dst_scalar_step_in_vector = - generate_sequence(detail::lambda_scalar_step_in_vector{}, Number{}); - - using SpaceFillingCurve = SpaceFillingCurve>; - - static_assert(DstScalarPerVector == SpaceFillingCurve::ScalarPerVector, - "wrong!DstScalarPerVector != SpaceFillingCurve::ScalarPerVector"); - - constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); - - static_for<0, num_access, 1>{}([&](auto idx_1d) { - constexpr auto idx_md = SpaceFillingCurve::GetIndex(idx_1d); - - // copy data from src_buf into dst_vector - static_for<0, DstScalarPerVector, 1>{}([&](auto i) { - // src_desc error, non constexpr, caused by merge transform - constexpr index_t src_offset = src_desc.CalculateOffset( - src_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - constexpr index_t dst_offset = dst_desc.CalculateOffset( - dst_slice_origin_idx + idx_md + i * dst_scalar_step_in_vector); - - SrcData v_this_row; - // int type temp value due to intrinsic requirement - int temp = 0; - - // apply element-wise operation - element_op_(v_this_row, src_buf[Number{}]); - - // apply intra-row permute. - if constexpr(IntraRowSwizzlePerm) - { - temp = __builtin_amdgcn_permlane16( - temp, type_convert_sp(v_this_row), 0xb3a29180, 0xf7e6d5c4, 1, 0); - v_this_row = type_convert_sp(temp); - } - - // apply type convert - dst_buf(Number{}) = type_convert_sp(v_this_row); - }); - }); - } - ElementwiseOperation element_op_{}; -}; - -} // namespace ck