tensor_utils.hpp Source File

tensor_utils.hpp Source File#

Composable Kernel: tensor_utils.hpp Source 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