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

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

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/wrapper/utils/tensor_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/data_type.hpp"
 #include "ck/utility/number.hpp"
 #include "ck/utility/tuple.hpp"
 #include "ck/utility/tuple_helper.hpp"
 #include "ck/utility/dynamic_buffer.hpp"
 #include "ck/utility/amd_address_space.hpp"
 #include "ck/utility/multi_index.hpp"
  
 // Disable from doxygen docs generation
 namespace ck {
 namespace wrapper {
  
 using MemoryTypeEnum = AddressSpaceEnum;
  
 // Disable from doxygen docs generation
 // forward declarations
 template <typename Shape, typename UnrolledDescriptorType>
 struct Layout;
 template <MemoryTypeEnum BufferAddressSpace,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 struct Tensor;
  
 template <typename FromType, typename ToType>
 struct Slice
 {
     __host__ __device__ constexpr Slice() : from_(), to_() {}
     __host__ __device__ constexpr Slice(FromType from, ToType to) : from_(from), to_(to) {}
  
     template <typename T>
     __host__ __device__ constexpr auto range(const T& dim) const
     {
         if constexpr(is_same_v<FromType, index_t> || is_same_v<ToType, index_t> ||
                      is_same_v<std::remove_const_t<T>, index_t>)
         {
             if(to_ < 0)
             {
                 return dim - from_ + to_ + 1;
             }
             else
             {
                 // workaround if one end of the interval is index_t and the second one is Number
                 return static_cast<index_t>(to_) - static_cast<index_t>(from_);
             }
         }
         else
         {
             static_assert(T{} >= ToType{} && FromType{} >= Number<0>{} &&
                               (ToType{} < 0 || ToType{} > FromType{}),
                           "Invalid range");
             if constexpr(ToType{} < 0)
             {
                 return dim - from_ + to_ + Number<1>{};
             }
             else
             {
                 return to_ - from_;
             }
         }
     }
  
     __host__ __device__ static constexpr bool IsSlice() { return true; }
  
     const FromType from_;
     const ToType to_;
 };
  
 template <typename T>
 using is_slice = decltype(std::declval<T&>().IsSlice());
  
 template <typename T>
 using is_tuple = decltype(std::declval<T&>().IsTuple());
  
 template <MemoryTypeEnum MemoryType,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 constexpr auto make_tensor(ElementType* pointer,
                            const Layout<Shape, UnrolledDescriptorType>& layout)
 {
     return Tensor<MemoryType, ElementType, Shape, UnrolledDescriptorType>(pointer, layout);
 }
  
 template <MemoryTypeEnum MemoryType,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 constexpr auto make_register_tensor(const Layout<Shape, UnrolledDescriptorType>& layout)
 {
     return Tensor<MemoryType, ElementType, Shape, UnrolledDescriptorType>(layout);
 }
  
 template <MemoryTypeEnum BufferAddressSpace,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 __host__ __device__ void
 clear(Tensor<BufferAddressSpace, ElementType, Shape, UnrolledDescriptorType>& tensor)
 {
     static_assert(
         !Tensor<BufferAddressSpace, ElementType, Shape, UnrolledDescriptorType>::IsDynamicBuffer);
     return tensor.GetBuffer().Clear();
 }
  
 template <MemoryTypeEnum BufferAddressSpace,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 __host__ __device__ constexpr const auto&
 layout(const Tensor<BufferAddressSpace, ElementType, Shape, UnrolledDescriptorType>& tensor)
 {
     return tensor.GetLayout();
 }
  
 template <index_t... Idxs,
           MemoryTypeEnum BufferAddressSpace,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 __host__ __device__ constexpr auto
 size(const Tensor<BufferAddressSpace, ElementType, Shape, UnrolledDescriptorType>& tensor)
 {
     return size<Idxs...>(tensor.GetLayout());
 }
  
 template <index_t... Idxs,
           MemoryTypeEnum BufferAddressSpace,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 __host__ __device__ constexpr auto
 rank(const Tensor<BufferAddressSpace, ElementType, Shape, UnrolledDescriptorType>& tensor)
 {
     return rank<Idxs...>(tensor.GetLayout());
 }
  
 template <index_t... Idxs,
           MemoryTypeEnum BufferAddressSpace,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 __host__ __device__ constexpr auto
 depth(const Tensor<BufferAddressSpace, ElementType, Shape, UnrolledDescriptorType>& tensor)
 {
     return depth<Idxs...>(tensor.GetLayout());
 }
  
 template <MemoryTypeEnum BufferAddressSpace,
           typename ElementType,
           typename Shape,
           typename UnrolledDescriptorType>
 __host__ __device__ constexpr const auto&
 shape(const Tensor<BufferAddressSpace, ElementType, Shape, UnrolledDescriptorType>& tensor)
 {
     return shape(tensor.GetLayout());
 }
  
 template <typename FromType, typename ToType>
 constexpr auto slice(const FromType from, const ToType to)
 {
     return Slice<FromType, ToType>(from, to);
 }
  
 template <typename ToType>
 constexpr auto slice(const ToType to)
 {
     if constexpr(is_same_v<ToType, index_t>)
     {
         return Slice<index_t, ToType>(0, to);
     }
     else
     {
         return Slice<Number<0>, ToType>(Number<0>{}, to);
     }
 }
  
 constexpr auto slice() { return Slice<Number<0>, Number<-1>>(Number<0>{}, Number<-1>{}); }
  
 } // namespace wrapper
 } // namespace ck