diff --git a/include/ck/wrapper/layout.hpp b/include/ck/wrapper/layout.hpp index 18ce8eb2a3..2ddd50b0e9 100644 --- a/include/ck/wrapper/layout.hpp +++ b/include/ck/wrapper/layout.hpp @@ -3,7 +3,7 @@ #pragma once -#include "ck/wrapper/layout_utils.hpp" +#include "ck/wrapper/utils/layout_utils.hpp" namespace ck { namespace wrapper { @@ -232,6 +232,8 @@ struct Layout using NaiveDescriptorType = remove_cvref_t; + constexpr auto GetElementSpaceSize() const { return descriptor_.GetElementSpaceSize(); } + /** * \brief Layout constructor. * diff --git a/include/ck/wrapper/tensor.hpp b/include/ck/wrapper/tensor.hpp new file mode 100644 index 0000000000..03fb5dd46a --- /dev/null +++ b/include/ck/wrapper/tensor.hpp @@ -0,0 +1,59 @@ +// 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 +struct Tensor +{ + using ElementSpaceSize = decltype(Layout{ + Shape{}, Strides{}}.GetElementSpaceSize()); + + __host__ __device__ Tensor() = delete; + __host__ __device__ Tensor(ElementType* pointer, const Layout& layout) + : layout_(layout), + dynamic_buffer_( + make_dynamic_buffer(pointer, layout.GetElementSpaceSize())) + { + } + + template + __host__ __device__ index_t operator[](const Tuple& Idx) const + { + // Padding is not supported, so we can assume that read should be valid. + return dynamic_buffer_.template Get(layout_(Idx), true /*is_valid*/); + } + + template + __host__ __device__ index_t operator()(const Tuple& Idx) const + { + return dynamic_buffer_.template Get(layout_(Idx), true /*is_valid*/); + } + + template + __host__ __device__ index_t operator()(Idxs... idxs) const + { + const auto idxs_tuple = make_tuple(idxs...); + return dynamic_buffer_.template Get(layout_(idxs_tuple), true /*is_valid*/); + } + + private: + const Layout& layout_; + DynamicBuffer + dynamic_buffer_; +}; + +} // namespace wrapper +} // namespace ck diff --git a/include/ck/wrapper/layout_utils.hpp b/include/ck/wrapper/utils/layout_utils.hpp similarity index 97% rename from include/ck/wrapper/layout_utils.hpp rename to include/ck/wrapper/utils/layout_utils.hpp index fac8f33854..214c143d55 100644 --- a/include/ck/wrapper/layout_utils.hpp +++ b/include/ck/wrapper/utils/layout_utils.hpp @@ -22,7 +22,7 @@ namespace wrapper { // Disable from doxygen docs generation /// @cond // forward declaration -template > +template struct Layout; template @@ -52,9 +52,9 @@ __host__ __device__ constexpr Layout make_layout(const Shape& sh * \return Constructed layout. */ template -__host__ __device__ constexpr Layout make_layout(const Shape& shape) +__host__ __device__ constexpr Layout> make_layout(const Shape& shape) { - return Layout(shape); + return Layout>(shape); } // Layout helpers diff --git a/include/ck/wrapper/utils/tensor_utils.hpp b/include/ck/wrapper/utils/tensor_utils.hpp new file mode 100644 index 0000000000..a93243267f --- /dev/null +++ b/include/ck/wrapper/utils/tensor_utils.hpp @@ -0,0 +1,107 @@ +#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/dynamic_buffer.hpp" +#include "ck/utility/amd_address_space.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" + +namespace ck { +namespace wrapper { + +// Disable from doxygen docs generation +/// @cond +// forward declarations +template +struct Layout; +template +struct Tensor; +/// @endcond + +template +struct MemoryPointerTag +{ + MemoryPointerTag(PointerElementType* pointer) : pointer_(pointer) {} + + using ElementType = PointerElementType; + + AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Generic; + ElementType* pointer_; +}; + +template +struct GlobalMemoryPointerTag : public MemoryPointerTag +{ + AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Global; +}; + +template +struct SharedMemoryPointerTag : public MemoryPointerTag +{ + AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Lds; +}; + +template +struct SgprMemoryPointerTag : public MemoryPointerTag +{ + AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Sgpr; +}; + +template +struct VgprMemoryPointerTag : public MemoryPointerTag +{ + AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Vgpr; +}; + +template +constexpr auto make_gmem_ptr(const ElementType* pointer) +{ + return GlobalMemoryPointerTag(pointer); +} + +template +constexpr auto make_smem_ptr(const ElementType* pointer) +{ + return SharedMemoryPointerTag(pointer); +} + +template +constexpr auto make_sgprmem_ptr(const ElementType* pointer) +{ + return SgprMemoryPointerTag(pointer); +} + +template +constexpr auto make_vgprmem_ptr(const ElementType* pointer) +{ + return VgprMemoryPointerTag(pointer); +} + +template +constexpr auto make_tensor(ElementType* pointer, const Layout& layout) +{ + return Tensor(pointer, layout); +} + +template +constexpr auto make_tensor(MemoryPointerTag& mem_tag, + const Layout& layout) +{ + return Tensor(mem_tag.pointer_, + layout); +} + +} // namespace wrapper +} // namespace ck diff --git a/test/wrapper/CMakeLists.txt b/test/wrapper/CMakeLists.txt index e25ef176dd..4a3c75cfba 100644 --- a/test/wrapper/CMakeLists.txt +++ b/test/wrapper/CMakeLists.txt @@ -1,2 +1,4 @@ add_gtest_executable(test_layout test_layout.cpp) target_link_libraries(test_layout PRIVATE utility) +add_gtest_executable(test_tensor_read_write test_tensor_read_write.cpp) +target_link_libraries(test_tensor_read_write PRIVATE utility) diff --git a/test/wrapper/test_tensor_read_write.cpp b/test/wrapper/test_tensor_read_write.cpp new file mode 100644 index 0000000000..10ec56fa9b --- /dev/null +++ b/test/wrapper/test_tensor_read_write.cpp @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved. + +#include +#include +#include +#include +#include +#include + +#include "ck/utility/common_header.hpp" + +#include "ck/wrapper/layout.hpp" +#include "ck/wrapper/tensor.hpp" + +// Compare data in tensor with offset with layout. +// Data and offset should match if physical memory has been initialized with +// sequentially increasing values. +template +__host__ __device__ bool TestTensorReadWriteCheckCustom(const TensorType& tensor, + const LayoutType& layout, + const std::vector idxs) +{ + for(size_t i = 0; i < idxs.size(); i++) + { + if(tensor(idxs[i]) != layout(idxs[i])) + { + return false; + } + } + return true; +} + +template +__host__ __device__ bool TestTensorReadWriteCheck1d(const TensorType& tensor, + const LayoutType& layout) +{ + for(ck::index_t w = 0; w < ck::wrapper::size<0>(layout); w++) + { + if(tensor(w) != layout(ck::make_tuple(w))) + { + return false; + } + } + return true; +} + +TEST(TestTensorReadWrite, HostMemory) +{ + constexpr ck::index_t nelems = 8; + + std::array data; + std::iota(data.begin(), data.end(), 0); + const auto layout = ck::wrapper::make_layout(ck::make_tuple(4, 2)); + const auto tensor = ck::wrapper::make_tensor(&data[0], layout); + + std::vector> idxs; + for(ck::index_t h = 0; h < ck::wrapper::size<0>(layout); h++) + { + for(ck::index_t w = 0; w < ck::wrapper::size<1>(layout); w++) + { + idxs.emplace_back(h, w); + } + } + + EXPECT_TRUE(TestTensorReadWriteCheck1d(tensor, layout)); + EXPECT_TRUE(TestTensorReadWriteCheckCustom(tensor, layout, idxs)); +} + +TEST(TestTensorReadWrite, HostMemoryNested) +{ + constexpr ck::index_t nelems = 8; + + std::array data; + std::iota(data.begin(), data.end(), 0); + const auto layout = ck::wrapper::make_layout(ck::make_tuple(ck::make_tuple(2, 2), 2)); + const auto tensor = ck::wrapper::make_tensor(&data[0], layout); + + std::vector, ck::index_t>> idxs; + for(ck::index_t d = 0; d < ck::wrapper::size<0>(ck::wrapper::get<0>(layout)); d++) + { + for(ck::index_t h = 0; h < ck::wrapper::size<1>(ck::wrapper::get<0>(layout)); h++) + { + for(ck::index_t w = 0; w < ck::wrapper::size<1>(layout); w++) + { + idxs.emplace_back(ck::make_tuple(d, h), w); + } + } + } + + EXPECT_TRUE(TestTensorReadWriteCheck1d(tensor, layout)); + EXPECT_TRUE(TestTensorReadWriteCheckCustom(tensor, layout, idxs)); +}