/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/wrapper/utils/tensor_partition.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/wrapper/utils/tensor_partition.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/wrapper/utils/tensor_partition.hpp Source File
Go to the documentation of this file.
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2023-2024, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "tensor_utils.hpp"
 #include "layout_utils.hpp"
  
 #include "ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp"
 #include "ck/tensor_description/cluster_descriptor.hpp"
  
 // Disable from doxygen docs generation
 namespace ck {
 namespace wrapper {
  
 // Disable from doxygen docs generation
 namespace {
  
 namespace detail {
  
 template <typename... Ts, typename... Ls>
 __host__ __device__ constexpr auto CalculateLocalPartitionShape(const Tuple<Ts...>& shape,
                                                                 const Tuple<Ls...>& thread_lengths)
 {
     static_assert(Tuple<Ts...>::Size() == Tuple<Ls...>::Size(), "Wrong thread_lengths shape.");
     return generate_tuple(
         [&](auto i) {
             constexpr auto num_i = Number<i>{};
             const auto slice_len =
                 ck::math::integer_divide_ceil(size<num_i>(shape), thread_lengths.At(num_i));
             return slice_len;
         },
         Number<Tuple<Ls...>::Size()>{});
 }
  
 template <typename MultiIndex, typename ProjectionTuple>
 __host__ __device__ constexpr auto
 ApplyProjection([[maybe_unused]] const MultiIndex& base_tuple,
                 [[maybe_unused]] const ProjectionTuple& projection)
 {
     if constexpr(is_same_v<ProjectionTuple, Tuple<>>)
     {
         return Tuple<>{};
     }
     else
     {
         auto base_tuple_after_projection = generate_tuple(
             [&](auto i) {
                 const auto i_num = Number<i.value>{};
                 static_assert(
                     is_detected<is_slice, tuple_element_t<i_num, ProjectionTuple>>::value ||
                     is_same_v<tuple_element_t<i_num, ProjectionTuple>, Number<1>>);
                 if constexpr(is_detected<is_slice, tuple_element_t<i_num, ProjectionTuple>>::value)
                 {
                     // When slice (to remove), then insert empty tuple (will be removed in next
                     // step).
                     return Tuple<>{};
                 }
                 else
                 {
                     return make_tuple(base_tuple.At(i_num));
                 }
             },
             Number<MultiIndex::Size()>{});
         // Remove empty tuples
         return UnrollNestedTuple<0, 1>(base_tuple_after_projection);
     }
 }
  
 template <typename... Ts, typename... Ps>
 __host__ __device__ constexpr auto CalculateShapeWithProjection(const Tuple<Ts...>& shape,
                                                                 const Tuple<Ps...>& projection)
 {
     return generate_tuple(
         [&](auto i) {
             if constexpr(is_detected<is_slice, tuple_element_t<i, Tuple<Ps...>>>::value)
             {
                 return size<i>(projection).to_;
             }
             else
             {
                 // number of shape element in actual fragment of shape and projection (method to
                 // calculate shape idx)
                 constexpr index_t shape_i =
                     detail::ApplyProjection(TupleSlice<0, i>(Tuple<Ts...>{}),
                                             TupleSlice<0, i>(Tuple<Ps...>{}))
                         .Size();
                 return size<shape_i>(shape);
             }
         },
         Number<Tuple<Ps...>::Size()>{});
 }
  
 template <typename... Ts, typename... Ls, typename... Ps>
 __host__ __device__ constexpr auto CalculateGridSize(const Tuple<Ts...>& shape,
                                                      const Tuple<Ls...>& tile_shape)
 {
     return generate_tuple(
         [&](auto i) { return ck::math::integer_divide_ceil(size<i>(shape), size<i>(tile_shape)); },
         Number<Tuple<Ls...>::Size()>{});
 }
  
 template <typename ThreadIdxs, typename PartitionLengthsSeq, typename OldOffsetIdxs>
 __host__ __device__ constexpr auto
 CalculateOffsetMultiIdxs(const ThreadIdxs& thread_idxs,
                          const PartitionLengthsSeq& partition_lengths_seq,
                          const OldOffsetIdxs& old_offset_idxs)
 {
     return thread_idxs * partition_lengths_seq + old_offset_idxs;
 }
  
 template <typename BlockIdxs>
 __host__ __device__ constexpr auto GetDimsToPartition([[maybe_unused]] const BlockIdxs& block_idxs)
 {
     const auto dims_to_partition = generate_tuple(
         [&](auto i) {
             if constexpr(!is_detected<is_slice, tuple_element_t<i, BlockIdxs>>::value)
             {
                 return Number<i>{};
             }
             else
             {
                 return Tuple<>{};
             }
         },
         Number<BlockIdxs::Size()>{});
     // Remove empty tuples
     return UnrollNestedTuple<0, 1>(dims_to_partition);
 }
  
 template <typename BlockIdxs>
 __host__ __device__ constexpr auto ReplaceSlicesWithZeros(const BlockIdxs& block_idxs)
 {
     return generate_tuple(
         [&](auto i) {
             if constexpr(!is_detected<is_slice, tuple_element_t<i, BlockIdxs>>::value)
             {
                 return block_idxs.At(i);
             }
             else
             {
                 return Number<0>{};
             }
         },
         Number<BlockIdxs::Size()>{});
 }
  
 template <typename TileShape>
 __host__ __device__ constexpr auto
 GenerateDefaultProjection([[maybe_unused]] const TileShape tile_shape)
 {
     return generate_tuple([&](auto) { return Number<1>{}; }, Number<TileShape::Size()>{});
 }
  
 template <typename ThreadShape, typename ThreadUnrolledDesc>
 __host__ __device__ constexpr auto CalculateThreadMultiIdx(
     [[maybe_unused]] const Layout<ThreadShape, ThreadUnrolledDesc>& thread_layout,
     const index_t thread_id)
 {
     static_assert(ThreadUnrolledDesc::GetNumOfTransform() == 1,
                   "Thread layout should not be transformed.");
     constexpr auto embed_transform = ThreadUnrolledDesc{}.GetTransforms().At(Number<0>{});
     constexpr auto shape           = ThreadShape{};
     constexpr auto strides         = embed_transform.coefficients_;
  
     return generate_tuple(
         [&](auto i) {
             constexpr auto num_i = Number<i>{};
             return (thread_id / strides.At(num_i)) % shape.At(num_i);
         },
         Number<ThreadShape::Size()>{});
 }
 } // namespace detail
 } // namespace
  
 template <typename TensorType,
           typename ThreadShape,
           typename ThreadUnrolledDesc,
           typename ProjectionTuple>
 __host__ __device__ constexpr auto
 make_local_partition(TensorType& tensor,
                      [[maybe_unused]] const Layout<ThreadShape, ThreadUnrolledDesc>& thread_layout,
                      const index_t thread_id,
                      const ProjectionTuple& projection)
 {
     static_assert(!IsNestedTuple(ThreadShape{}));
     // Calculate new partition shape
     const auto& tensor_shape = shape(tensor);
     // Calculate projected thread lengths
     constexpr auto projected_thread_lengths =
         detail::ApplyProjection(ThreadShape{}, ProjectionTuple{});
     constexpr auto partition_shape =
         detail::CalculateLocalPartitionShape(decltype(tensor_shape){}, projected_thread_lengths);
     constexpr auto partition_shape_seq =
         generate_sequence_v2([&](auto I) { return size<I>(partition_shape); },
                              Number<decltype(partition_shape)::Size()>{});
     // Calculate thread idxs and offsets
     const auto thread_idxs = detail::CalculateThreadMultiIdx(thread_layout, thread_id);
     // Apply projection on thread idxs to remove not needed idxs
     const auto projected_thread_idxs = detail::ApplyProjection(thread_idxs, projection);
     const auto offset_multi_idxs     = detail::CalculateOffsetMultiIdxs(
         projected_thread_idxs, partition_shape_seq, tensor.GetMultiIdxOffsets());
     // Create new layout and tensor
     auto& unrolled_desc = layout(tensor).GetUnrolledDescriptor();
     // Slice descriptor
     const auto transforms = generate_tuple(
         [&](auto i) {
             return make_slice_transform(partition_shape.At(i),
                                         offset_multi_idxs.At(i),
                                         partition_shape.At(i) + offset_multi_idxs.At(i));
         },
         Number<remove_reference_t<decltype(tensor_shape)>::Size()>{});
     const auto lower_upper_dims =
         generate_tuple([&](auto i) { return Sequence<i.value>{}; },
                        Number<remove_reference_t<decltype(tensor_shape)>::Size()>{});
     auto sliced_desc =
         transform_tensor_descriptor(unrolled_desc, transforms, lower_upper_dims, lower_upper_dims);
     // Create layout
     const auto partition_layout =
         Layout<remove_reference_t<decltype(partition_shape)>, decltype(sliced_desc)>(
             partition_shape, sliced_desc);
     auto partition_tensor =
         make_tensor<TensorType::TensorBufferAddressSpace>(tensor.GetPointer(), partition_layout);
     // Apply offsets
     return partition_tensor;
 }
  
 template <typename TensorType, typename ThreadShape, typename ThreadUnrolledDesc>
 __host__ __device__ constexpr auto
 make_local_partition(TensorType& tensor,
                      const Layout<ThreadShape, ThreadUnrolledDesc>& thread_lengths,
                      const index_t thread_id)
 {
     const auto projection = detail::GenerateDefaultProjection(ThreadShape{});
     return make_local_partition(tensor, thread_lengths, thread_id, projection);
 }
  
 template <typename TensorType,
           typename BlockShapeTuple,
           typename BlockIdxs,
           typename ProjectionTuple>
 __host__ __device__ constexpr auto make_local_tile(const TensorType& tensor,
                                                    const BlockShapeTuple& tile_shape,
                                                    const BlockIdxs& block_idxs,
                                                    const ProjectionTuple& projection)
 {
     static_assert(!IsNestedTuple(BlockShapeTuple{}));
     static_assert(!IsNestedTuple(BlockIdxs{}));
  
     constexpr auto I0 = Number<0>{};
     constexpr auto I1 = Number<1>{};
     constexpr auto I2 = Number<2>{};
  
     auto& aligned_desc = layout(tensor).GetMergedNestingDescriptor();
  
     constexpr auto projected_tile_shape =
         detail::ApplyProjection(BlockShapeTuple{}, ProjectionTuple{});
     // Number of dims which are partitioned
     constexpr auto dims_to_partition = detail::GetDimsToPartition(BlockIdxs{});
     const auto parsed_block_idxs     = detail::ReplaceSlicesWithZeros(block_idxs);
     if constexpr(decltype(dims_to_partition)::Size() == I2)
     {
         const auto shape_with_projection_dims =
             detail::CalculateShapeWithProjection(shape(tensor), projection);
         // Set Value for M, N partition
         const auto M             = shape_with_projection_dims.At(dims_to_partition.At(I0));
         const auto N             = shape_with_projection_dims.At(dims_to_partition.At(I1));
         constexpr auto MPerBlock = BlockShapeTuple{}.At(dims_to_partition.At(I0));
         constexpr auto NPerBlock = BlockShapeTuple{}.At(dims_to_partition.At(I1));
         auto m_n_desc            = make_naive_tensor_descriptor_packed(make_tuple(M, N));
         // Get 1D block id
         const auto grid_size = detail::CalculateGridSize(shape_with_projection_dims, tile_shape);
         const auto block_lengths_desc = make_naive_tensor_descriptor_packed(grid_size);
         const index_t block_id_1d     = block_lengths_desc.CalculateOffset(parsed_block_idxs);
         // Optimized version for 2d tile shape [MxN]
         const auto block_2_tile_map =
             BlockToCTileMap_M00_N0_M01Adapt<MPerBlock,
                                             NPerBlock,
                                             remove_cvref_t<decltype(m_n_desc)>>(m_n_desc);
         const auto block_work_idx =
             block_2_tile_map.CalculateBottomIndex(make_multi_index(block_id_1d));
         const index_t m_block_data_idx_on_grid =
             __builtin_amdgcn_readfirstlane(block_work_idx[I0] * MPerBlock);
         const index_t n_block_data_idx_on_grid =
             __builtin_amdgcn_readfirstlane(block_work_idx[I1] * NPerBlock);
         // Apply 0 for non partitioned dims
         const auto offset_multi_idxs = generate_tuple(
             [&](auto i) {
                 if constexpr(i == dims_to_partition.At(I0))
                 {
                     return m_block_data_idx_on_grid;
                 }
                 else if constexpr(i == dims_to_partition.At(I1))
                 {
                     return n_block_data_idx_on_grid;
                 }
                 else
                 {
                     return Number<0>{};
                 }
             },
             Number<BlockShapeTuple::Size()>{});
         const auto projected_offset_multi_idxs =
             detail::ApplyProjection(offset_multi_idxs, projection);
         // Create new layout and tensor
         const auto tile_layout =
             Layout<remove_reference_t<decltype(projected_tile_shape)>, decltype(aligned_desc)>(
                 projected_tile_shape, aligned_desc);
         auto tile_tensor =
             make_tensor<TensorType::TensorBufferAddressSpace>(tensor.GetPointer(), tile_layout);
         // Apply offsets
         tile_tensor.SetMultiIdxOffset(to_multi_index(projected_offset_multi_idxs));
         return tile_tensor;
     }
     else
     {
         // Calculate offsets
         // Sequence with data to process per block
         using ProjectedTileShapeTuple = decltype(projected_tile_shape);
         constexpr auto projected_tile_shape_seq =
             generate_sequence_v2([](auto I) { return ProjectedTileShapeTuple{}.At(I); },
                                  Number<ProjectedTileShapeTuple::Size()>{});
         // Tuple with number of blocks
         const auto projected_block_idxs =
             to_multi_index(detail::ApplyProjection(parsed_block_idxs, projection));
         const auto offset_multi_idxs = detail::CalculateOffsetMultiIdxs(
             projected_block_idxs, projected_tile_shape_seq, tensor.GetMultiIdxOffsets());
         // Create new layout and tensor
         const auto tile_layout =
             Layout<remove_reference_t<ProjectedTileShapeTuple>, decltype(aligned_desc)>(
                 projected_tile_shape, aligned_desc);
         auto tile_tensor =
             make_tensor<TensorType::TensorBufferAddressSpace>(tensor.GetPointer(), tile_layout);
         // Apply offsets
         tile_tensor.SetMultiIdxOffset(to_multi_index(offset_multi_idxs));
         return tile_tensor;
     }
 }
  
 template <typename TensorType, typename BlockShapeTuple, typename BlockIdxs>
 __host__ __device__ constexpr auto make_local_tile(const TensorType& tensor,
                                                    const BlockShapeTuple& tile_shape,
                                                    const BlockIdxs& block_idxs)
 {
     const auto projection = detail::GenerateDefaultProjection(BlockShapeTuple{});
     return make_local_tile(tensor, tile_shape, block_idxs, projection);
 }
  
 } // namespace wrapper
 } // namespace ck