mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-03-25 09:37:42 +00:00
278 lines
7.6 KiB
C++
278 lines
7.6 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#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
|
|
/// @cond INTERNAL
|
|
namespace ck {
|
|
namespace wrapper {
|
|
/// @endcond
|
|
|
|
/**
|
|
* \brief Memory type, allowed members:
|
|
* - Generic,
|
|
* - Global,
|
|
* - Lds,
|
|
* - Sgpr,
|
|
* - Vgpr,
|
|
*/
|
|
using MemoryTypeEnum = AddressSpaceEnum;
|
|
|
|
// Disable from doxygen docs generation
|
|
/// @cond INTERNAL
|
|
// 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) {}
|
|
|
|
/**
|
|
* \brief Calculate slice range.
|
|
*
|
|
* \param dim Dimension size.
|
|
* \return Slice range.
|
|
*/
|
|
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());
|
|
/// @endcond
|
|
|
|
/**
|
|
* \brief Make tensor function.
|
|
*
|
|
* \tparam MemoryType Type of memory.
|
|
* \param pointer Pointer to the memory.
|
|
* \param layout Tensor layout.
|
|
* \return Constructed tensor.
|
|
*/
|
|
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);
|
|
}
|
|
|
|
/**
|
|
* \brief Make SGPR or VGPR tensor function.
|
|
*
|
|
* \tparam MemoryType Type of memory.
|
|
* \tparam ElementType Memory data type.
|
|
* \return Constructed tensor.
|
|
*/
|
|
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);
|
|
}
|
|
|
|
/**
|
|
* \brief Clear tensor. (Only for Vpgr/Sgpr)
|
|
*
|
|
* \param tensor Tensor to be cleared.
|
|
*/
|
|
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();
|
|
}
|
|
|
|
/**
|
|
* \brief Get Tensor Layout.
|
|
*
|
|
* \param tensor Tensor to get layout of.
|
|
* \return Requsted layout.
|
|
*/
|
|
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();
|
|
}
|
|
|
|
/**
|
|
* \brief Product of tensor shape dims.
|
|
*
|
|
* \tparam Idxs Indexes to access specific shape dim (optional).
|
|
* \param tensor Tensor to get Shape of.
|
|
* \return Requsted size.
|
|
*/
|
|
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());
|
|
}
|
|
|
|
/**
|
|
* \brief Rank of Shape tuple.
|
|
*
|
|
* \tparam Idxs Indexes to access specific shape dim (optional).
|
|
* \param tensor Tensor to get rank of.
|
|
* \return Requsted rank.
|
|
*/
|
|
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());
|
|
}
|
|
|
|
/**
|
|
* \brief Depth of Shape tuple.
|
|
*
|
|
* \tparam Idxs Indexes to access specific shape dim (optional).
|
|
* \param tensor Tensor to get depth of.
|
|
* \return Requsted depth.
|
|
*/
|
|
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());
|
|
}
|
|
|
|
/**
|
|
* \brief Get Tensor shape.
|
|
*
|
|
* \param tensor Tensor to get shape from.
|
|
* \return Requsted shape.
|
|
*/
|
|
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());
|
|
}
|
|
|
|
/**
|
|
* \brief Get dim slice.
|
|
*
|
|
* \param from Beginning of the interval.
|
|
* \param to End of the interval. (could be also negative to index from the end)
|
|
* \return Requested slice. Could be used to create sliced tensor from other tensor.
|
|
*/
|
|
template <typename FromType, typename ToType>
|
|
constexpr auto slice(const FromType from, const ToType to)
|
|
{
|
|
return Slice<FromType, ToType>(from, to);
|
|
}
|
|
|
|
/**
|
|
* \brief Get dim slice. (Assumed that from is equal to 1)
|
|
*
|
|
* \param to End of the interval. (could be also negative to index from the end)
|
|
* \return Requested slice. Could be used to create sliced tensor from other tensor.
|
|
*/
|
|
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);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* \brief Get whole dim slice (from = 0, to = -1).
|
|
*
|
|
* \return Requested slice. Could be used to create sliced tensor from other tensor.
|
|
*/
|
|
constexpr auto slice() { return Slice<Number<0>, Number<-1>>(Number<0>{}, Number<-1>{}); }
|
|
|
|
} // namespace wrapper
|
|
} // namespace ck
|