mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-14 11:07:44 +00:00
60 lines
1.8 KiB
C++
60 lines
1.8 KiB
C++
// SPDX-License-Identifier: MIT
|
|
// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "layout.hpp"
|
|
#include "utils/tensor_utils.hpp"
|
|
|
|
namespace ck {
|
|
namespace wrapper {
|
|
|
|
template <AddressSpaceEnum BufferAddressSpace,
|
|
typename ElementType,
|
|
typename Shape,
|
|
typename Strides>
|
|
struct Tensor
|
|
{
|
|
using ElementSpaceSize = decltype(Layout<Shape, Strides>{
|
|
Shape{}, Strides{}}.GetElementSpaceSize());
|
|
|
|
__host__ __device__ Tensor() = delete;
|
|
__host__ __device__ Tensor(ElementType* pointer, const Layout<Shape, Strides>& layout)
|
|
: layout_(layout),
|
|
dynamic_buffer_(
|
|
make_dynamic_buffer<BufferAddressSpace>(pointer, layout.GetElementSpaceSize()))
|
|
{
|
|
}
|
|
|
|
template <typename... Ts>
|
|
__host__ __device__ index_t operator[](const Tuple<Ts...>& Idx) const
|
|
{
|
|
// Padding is not supported, so we can assume that read should be valid.
|
|
return dynamic_buffer_.template Get<ElementType>(layout_(Idx), true /*is_valid*/);
|
|
}
|
|
|
|
template <typename... Ts>
|
|
__host__ __device__ index_t operator()(const Tuple<Ts...>& Idx) const
|
|
{
|
|
return dynamic_buffer_.template Get<ElementType>(layout_(Idx), true /*is_valid*/);
|
|
}
|
|
|
|
template <typename... Idxs>
|
|
__host__ __device__ index_t operator()(Idxs... idxs) const
|
|
{
|
|
const auto idxs_tuple = make_tuple(idxs...);
|
|
return dynamic_buffer_.template Get<ElementType>(layout_(idxs_tuple), true /*is_valid*/);
|
|
}
|
|
|
|
private:
|
|
const Layout<Shape, Strides>& layout_;
|
|
DynamicBuffer<BufferAddressSpace,
|
|
ElementType,
|
|
ElementSpaceSize,
|
|
true /*InvalidElementUseNumericalZeroValue*/>
|
|
dynamic_buffer_;
|
|
};
|
|
|
|
} // namespace wrapper
|
|
} // namespace ck
|