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

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

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/wrapper/utils/layout_utils.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 "ck/ck.hpp"
  
 #include "ck/utility/number.hpp"
 #include "ck/utility/tuple.hpp"
 #include "ck/utility/tuple_helper.hpp"
 #include "ck/utility/sequence.hpp"
 #include "ck/utility/sequence_helper.hpp"
 #include "ck/utility/is_detected.hpp"
  
 #include "ck/tensor_description/tensor_descriptor.hpp"
 #include "ck/tensor_description/tensor_descriptor_helper.hpp"
 #include "ck/tensor_description/multi_index_transform_helper.hpp"
 #include "ck/tensor_operation/gpu/device/matrix_padder.hpp"
  
 // Disable from doxygen docs generation
 namespace ck {
 namespace wrapper {
  
 // Disable from doxygen docs generation
 // forward declaration
 template <typename Shape, typename UnrolledDescriptorType>
 struct Layout;
  
 template <typename T>
 using is_tuple = decltype(std::declval<T&>().IsTuple());
  
 namespace {
 namespace detail {
 template <typename... Ts>
 __host__ __device__ constexpr static auto
 GenerateColumnMajorPackedStrides(const Tuple<Ts...>& shape)
 {
     const auto unrolled_shape = UnrollNestedTuple(shape);
     return generate_tuple(
         [&](auto i) {
             if constexpr(i.value == 0)
             {
                 return Number<1>{};
             }
             else
             {
                 return TupleReduce<Number<0>{}.value, i.value>([](auto x, auto y) { return x * y; },
                                                                unrolled_shape);
             }
         },
         Number<decltype(unrolled_shape)::Size()>{});
 }
  
 template <typename LayoutShape, typename LayoutStrides>
 __host__ __device__ constexpr auto MakeUnrolledDescriptor(const LayoutShape& shape,
                                                           const LayoutStrides& strides)
 {
     const auto unrolled_shape = UnrollNestedTuple(shape);
     if constexpr(is_same_v<LayoutStrides, Tuple<>>)
     {
         // if not passed, then generate
         const auto unrolled_strides = GenerateColumnMajorPackedStrides(unrolled_shape);
         static_assert(unrolled_shape.Size() == unrolled_strides.Size(),
                       "Size of strides and shape are not consistent.");
         return make_naive_tensor_descriptor(unrolled_shape, unrolled_strides);
     }
     else
     {
         const auto unrolled_strides = UnrollNestedTuple(strides);
         static_assert(unrolled_shape.Size() == unrolled_strides.Size(),
                       "Size of strides and shape are not consistent.");
         return make_naive_tensor_descriptor(unrolled_shape, unrolled_strides);
     }
 }
 } // namespace detail
 } // namespace
  
  
 // make_*
 template <typename Shape, typename Strides>
 __host__ __device__ constexpr auto make_layout(const Shape& shape, const Strides& strides)
 {
     using UnrolledDescriptorType = decltype(detail::MakeUnrolledDescriptor(Shape{}, Strides{}));
     return Layout<Shape, UnrolledDescriptorType>(shape,
                                                  detail::MakeUnrolledDescriptor(shape, strides));
 }
  
 template <typename Shape>
 __host__ __device__ constexpr auto make_layout(const Shape& shape)
 {
     using UnrolledDescriptorType = decltype(detail::MakeUnrolledDescriptor(Shape{}, Tuple<>{}));
     return Layout<Shape, UnrolledDescriptorType>(shape,
                                                  detail::MakeUnrolledDescriptor(shape, Tuple<>{}));
 }
 // Layout helpers
 // get
 template <typename T>
 __host__ __device__ T constexpr get(const T& dim)
 {
     return dim;
 }
  
 template <index_t idx, typename... Dims>
 __host__ __device__ constexpr auto get(const Tuple<Dims...>& tuple)
 {
     return tuple.At(Number<idx>{});
 }
  
 template <index_t idx, typename Shape, typename UnrolledDesc>
 __host__ __device__ constexpr auto get(const Layout<Shape, UnrolledDesc>& layout)
 {
     const auto& shape    = layout.GetShape();
     const auto new_shape = get<idx>(shape);
     static_assert(is_detected<is_tuple, decltype(new_shape)>::value,
                   "Shape of sub layout must be tuple");
  
     constexpr auto old_shape_dims = decltype(UnrollNestedTuple(shape))::Size();
     constexpr auto new_shape_dims = decltype(UnrollNestedTuple(new_shape))::Size();
     constexpr auto shape_offset   = decltype(UnrollNestedTuple(TupleSlice<0, idx>(shape)))::Size();
  
     const auto unrolled_shape = UnrollNestedTuple(shape);
     const auto transforms     = generate_tuple(
         [&](auto i) {
             // Compare Idx with shape
             if constexpr(i < shape_offset || i >= shape_offset + new_shape_dims)
             {
                 // Remove dimension
                 return make_freeze_transform(Number<0>{});
             }
             else
             {
                 return make_pass_through_transform(unrolled_shape.At(i));
             }
         },
         Number<old_shape_dims>{});
  
     const auto lower_dims =
         generate_tuple([&](auto i) { return Sequence<i.value>{}; }, Number<old_shape_dims>{});
     const auto upper_dims = generate_tuple(
         [&](auto i) {
             if constexpr(i < shape_offset || i >= shape_offset + new_shape_dims)
                 return Sequence<>{};
  
             else
             {
                 return Sequence<i.value - shape_offset>{};
             }
         },
         Number<old_shape_dims>{});
  
     const auto& flatten_desc = layout.GetUnrolledDescriptor();
     auto new_desc = transform_tensor_descriptor(flatten_desc, transforms, lower_dims, upper_dims);
     return Layout<decltype(new_shape), decltype(new_desc)>(new_shape, new_desc);
 }
  
 template <index_t Idx, index_t... Idxs, typename T>
 __host__ __device__ constexpr auto get(const T& elem)
 {
     return get<Idxs...>(get<Idx>(elem));
 }
  
 // size
 template <typename T>
 __host__ __device__ T constexpr size(const T& dim)
 {
     return dim;
 }
  
 template <index_t idx, typename Shape, typename UnrolledDescriptorType>
 __host__ __device__ constexpr auto size(const Layout<Shape, UnrolledDescriptorType>& layout)
 {
     return layout.template GetLength<idx>();
 }
  
 template <typename... ShapeDims>
 __host__ __device__ constexpr auto size(const Tuple<ShapeDims...>& shape)
 {
     const auto unrolled_shape = UnrollNestedTuple(shape);
     return TupleReduce<0, unrolled_shape.Size()>([](auto x, auto y) { return x * y; },
                                                  unrolled_shape);
 }
  
 template <typename Shape, typename UnrolledDescriptorType>
 __host__ __device__ constexpr auto size(const Layout<Shape, UnrolledDescriptorType>& layout)
 {
     return layout.GetLengths();
 }
  
 template <index_t idx, typename... Ts>
 __host__ __device__ constexpr auto size(const Tuple<Ts...>& tuple)
 {
     return size(tuple.At(Number<idx>{}));
 }
  
 template <index_t Idx, index_t... Idxs, typename T>
 __host__ __device__ constexpr auto size(const T& elem)
 {
     return size(get<Idx, Idxs...>(elem));
 }
  
 // rank
 template <typename Shape, typename UnrolledDescriptorType>
 __host__ __device__ constexpr auto
 rank([[maybe_unused]] const Layout<Shape, UnrolledDescriptorType>& layout)
 {
     return Shape::Size();
 }
  
 template <typename... Dims>
 __host__ __device__ constexpr auto rank([[maybe_unused]] const Tuple<Dims...>& tuple)
 {
     return Tuple<Dims...>::Size();
 }
  
 template <index_t IDim>
 __host__ __device__ constexpr index_t rank([[maybe_unused]] const Number<IDim>& dim)
 {
     return 1;
 }
  
 __host__ __device__ constexpr index_t rank([[maybe_unused]] const index_t& dim) { return 1; }
  
 template <index_t... Idxs, typename T>
 __host__ __device__ constexpr auto rank(const T& elem)
 {
     return rank(get<Idxs...>(elem));
 }
  
 // depth
 template <typename Shape, typename UnrolledDescriptorType>
 __host__ __device__ constexpr auto depth(const Layout<Shape, UnrolledDescriptorType>& layout)
 {
     const auto& shape = layout.GetShape();
     return TupleDepth(shape);
 }
  
 template <typename... Dims>
 __host__ __device__ constexpr auto depth(const Tuple<Dims...>& tuple)
 {
     return TupleDepth(tuple);
 }
  
 template <index_t IDim>
 __host__ __device__ constexpr index_t depth([[maybe_unused]] const Number<IDim>& dim)
 {
     return 0;
 }
  
 __host__ __device__ constexpr index_t depth([[maybe_unused]] const index_t& dim) { return 0; }
  
 template <index_t... Idxs, typename T>
 __host__ __device__ constexpr auto depth(const T& elem)
 {
     return depth(get<Idxs...>(elem));
 }
  
 template <typename LayoutType>
 __host__ __device__ constexpr const auto& shape(const LayoutType& layout)
 {
     return layout.GetShape();
 }
  
 // pad
 template <typename Shape, typename UnrolledDesc, typename TileLengths>
 __host__ __device__ constexpr auto pad(const Layout<Shape, UnrolledDesc>& layout,
                                        const TileLengths& tile_lengths)
 {
     auto& unrolled_desc = layout.GetUnrolledDescriptor();
     // Generate sequence with ones to mark that all dims will be padded
     constexpr auto do_pads_seq =
         generate_sequence_v2([](auto) { return Number<1>{}; }, Number<Shape::Size()>{});
     // Create descriptor with padding
     auto padded_desc =
         tensor_operation::device::PadTensorDescriptor(unrolled_desc, tile_lengths, do_pads_seq);
     // Generate padded shape
     const auto padded_shape = generate_tuple(
         [&](auto i) { return padded_desc.GetLength(Number<i>{}); }, Number<TileLengths::Size()>{});
     // Create layout
     return Layout<decltype(padded_shape), decltype(padded_desc)>(padded_shape, padded_desc);
 }
  
 // unmerge
 template <index_t Idx, typename Shape, typename UnrolledDesc, typename NewLengths, typename NewIdxs>
 __host__ __device__ constexpr auto unmerge(const Layout<Shape, UnrolledDesc>& layout,
                                            const NewLengths& new_lengths,
                                            [[maybe_unused]] const NewIdxs& new_indexes)
 {
     const auto& layout_shape = shape(layout);
     auto& unrolled_desc      = layout.GetUnrolledDescriptor();
     constexpr auto dims      = Shape::Size();
     // Generate transforms
     const auto transforms = generate_tuple(
         [&](auto i) {
             if constexpr(i == Idx)
             {
                 return make_unmerge_transform(new_lengths);
             }
             else
             {
                 return make_pass_through_transform(layout_shape.At(i));
             }
         },
         Number<dims>{});
  
     constexpr auto lower_dims =
         generate_tuple([&](auto i) { return Sequence<i.value>{}; }, Number<dims>{});
     constexpr auto upper_dims = generate_tuple(
         [&](auto i) {
             if constexpr(is_detected<is_tuple, tuple_element_t<i.value, NewIdxs>>::value)
             {
                 constexpr auto idxs_tuple = tuple_element_t<i.value, NewIdxs>{};
                 return to_sequence(idxs_tuple);
             }
             else
             {
                 constexpr index_t index = tuple_element_t<i.value, NewIdxs>{};
                 return Sequence<index>{};
             }
         },
         Number<dims>{});
  
     const auto unmerged_desc =
         transform_tensor_descriptor(unrolled_desc, transforms, lower_dims, upper_dims);
     const auto unmerged_shape =
         generate_tuple([&](auto i) { return unmerged_desc.GetLength(Number<i>{}); },
                        Number<decltype(unmerged_desc)::GetNumOfVisibleDimension()>{});
  
     // Create layout
     return Layout<decltype(unmerged_shape), decltype(unmerged_desc)>(unmerged_shape, unmerged_desc);
 }
  
 } // namespace wrapper
 } // namespace ck