tmp

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<decltype(MakeNaiveDescriptor(Shape{}, DeducedStrides{}))>;
 
+    constexpr auto GetElementSpaceSize() const { return descriptor_.GetElementSpaceSize(); }
+
     /**
      * \brief Layout constructor.
      *

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 <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

include/ck/wrapper/utils/layout_utils.hpp

@@ -22,7 +22,7 @@ namespace wrapper {
 // Disable from doxygen docs generation
 /// @cond
 // forward declaration
-template <typename Shape, typename Strides = Tuple<>>
+template <typename Shape, typename Strides>
 struct Layout;
 
 template <typename T>
@@ -52,9 +52,9 @@ __host__ __device__ constexpr Layout<Shape, Strides> make_layout(const Shape& sh
  * \return Constructed layout.
  */
 template <typename Shape>
-__host__ __device__ constexpr Layout<Shape> make_layout(const Shape& shape)
+__host__ __device__ constexpr Layout<Shape, Tuple<>> make_layout(const Shape& shape)
 {
-    return Layout<Shape>(shape);
+    return Layout<Shape, Tuple<>>(shape);
 }
 
 // Layout helpers

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 <typename Shape, typename Strides>
+struct Layout;
+template <AddressSpaceEnum BufferAddressSpace,
+          typename ElementType,
+          typename Shape,
+          typename Strides>
+struct Tensor;
+/// @endcond
+
+template <typename PointerElementType>
+struct MemoryPointerTag
+{
+    MemoryPointerTag(PointerElementType* pointer) : pointer_(pointer) {}
+
+    using ElementType = PointerElementType;
+
+    AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Generic;
+    ElementType* pointer_;
+};
+
+template <typename ElementType>
+struct GlobalMemoryPointerTag : public MemoryPointerTag<ElementType>
+{
+    AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Global;
+};
+
+template <typename ElementType>
+struct SharedMemoryPointerTag : public MemoryPointerTag<ElementType>
+{
+    AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Lds;
+};
+
+template <typename ElementType>
+struct SgprMemoryPointerTag : public MemoryPointerTag<ElementType>
+{
+    AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Sgpr;
+};
+
+template <typename ElementType>
+struct VgprMemoryPointerTag : public MemoryPointerTag<ElementType>
+{
+    AddressSpaceEnum buffer_adress_space_ = AddressSpaceEnum::Vgpr;
+};
+
+template <typename ElementType>
+constexpr auto make_gmem_ptr(const ElementType* pointer)
+{
+    return GlobalMemoryPointerTag<ElementType>(pointer);
+}
+
+template <typename ElementType>
+constexpr auto make_smem_ptr(const ElementType* pointer)
+{
+    return SharedMemoryPointerTag<ElementType>(pointer);
+}
+
+template <typename ElementType>
+constexpr auto make_sgprmem_ptr(const ElementType* pointer)
+{
+    return SgprMemoryPointerTag<ElementType>(pointer);
+}
+
+template <typename ElementType>
+constexpr auto make_vgprmem_ptr(const ElementType* pointer)
+{
+    return VgprMemoryPointerTag<ElementType>(pointer);
+}
+
+template <typename ElementType, typename Shape, typename Strides>
+constexpr auto make_tensor(ElementType* pointer, const Layout<Shape, Strides>& layout)
+{
+    return Tensor<AddressSpaceEnum::Generic, ElementType, Shape, Strides>(pointer, layout);
+}
+
+template <typename ElementType, typename Shape, typename Strides>
+constexpr auto make_tensor(MemoryPointerTag<ElementType>& mem_tag,
+                           const Layout<Shape, Strides>& layout)
+{
+    return Tensor<mem_tag.buffer_adress_space_, ElementType, Shape, Strides>(mem_tag.pointer_,
+                                                                             layout);
+}
+
+} // namespace wrapper
+} // namespace ck