added tuple

[ROCm/composable_kernel commit: 625838def0]

composable_kernel/include/utility/tuple.hpp

@@ -2,66 +2,99 @@
 #define CK_TUPLE_HPP
 
 #include "integral_constant.hpp"
+#include "Sequence.hpp"
 
 namespace ck {
 
-template <class... Ts>
-struct tuple : public std::tuple<Ts...>
-{
-    using type = tuple;
+namespace detail {
 
-    __host__ __device__ static constexpr index_t GetSize() { return std::tuple_size(tuple{}); }
+template <index_t>
+struct TupleElementKey
+{
+};
+
+template <typename Key, typename Data>
+struct TupleElement
+{
+    template <typename T>
+    __host__ __device__ explicit constexpr TupleElement(T&& v) : mData(static_cast<T&&>(v))
+    {
+    }
+
+    Data mData;
+};
+
+template <typename Key, typename Data>
+__host__ __device__ constexpr const Data& get_tuple_element(const TupleElement<Key, Data>& x)
+{
+    return x.mData;
+}
+
+template <typename Key, typename Data>
+__host__ __device__ constexpr Data& get_tuple_element(TupleElement<Key, Data>& x)
+{
+    return x.mData;
+}
+
+template <typename Key, typename Data>
+__host__ __device__ constexpr Data&& get_tuple_element(TupleElement<Key, Data>&& x)
+{
+    return static_cast<Data&&>(x.mData);
+}
+
+template <typename Indices, typename... Xs>
+struct TupleImpl;
+
+template <index_t... Is, typename... Xs>
+struct TupleImpl<Sequence<Is...>, Xs...> : TupleElement<TupleElementKey<Is>, Xs>...
+{
+    template <typename... Ys>
+    __host__ __device__ explicit constexpr TupleImpl(Ys&&... ys)
+        : TupleElement<TupleElementKey<Is>, Xs>(static_cast<Ys&&>(ys))...
+    {
+    }
+
+    __host__ __device__ static constexpr index_t Size() { return sizeof...(Xs); }
 
     template <index_t I>
-    __host__ __device__ constexpr auto Get(Number<I>) const
+    __host__ __device__ constexpr const auto& GetElementByKey(TupleElementKey<I>) const
     {
-        return std::get<I>(*this);
+        return get_tuple_element<TupleElementKey<I>>(*this);
     }
 
     template <index_t I>
-    __host__ __device__ constexpr auto operator[](Number<I>) const
+    __host__ __device__ constexpr auto& GetElementByKey(TupleElementKey<I>)
     {
-        return Get(Number<I>{}) :
+        return get_tuple_element<TupleElementKey<I>>(*this);
     }
 };
 
-// merge tuple
-template <class... Tuples>
-__host__ __device__ constexpr auto merge_tuple(Tuples&&... xs)
-{
-    return std::tuple_cat(xs...);
-};
+} // namespace detail
 
-// generate sequence
-template <index_t IBegin, index_t NRemain, class F>
-struct tuple_gen_impl
+template <typename... Xs>
+struct Tuple : detail::TupleImpl<typename arithmetic_sequence_gen<0, sizeof...(Xs), 1>::type, Xs...>
 {
-    static constexpr index_t NRemainLeft  = NRemain / 2;
-    static constexpr index_t NRemainRight = NRemain - NRemainLeft;
-    static constexpr index_t IMiddle      = IBegin + NRemainLeft;
+    using base =
+        detail::TupleImpl<typename arithmetic_sequence_gen<0, sizeof...(Xs), 1>::type, Xs...>;
 
-    using type =
-        typename tuple_merge<typename tuple_gen_impl<IBegin, NRemainLeft, F>::type,
-                             typename tuple_gen_impl<IMiddle, NRemainRight, F>::type>::type;
-};
+    template <typename... Ys>
+    __host__ __device__ explicit constexpr Tuple(Ys&&... ys) : base(static_cast<Ys&&>(ys)...)
+    {
+    }
 
-template <index_t I, class F>
-struct tuple_gen_impl<I, 1, F>
-{
-    static constexpr auto x = F{}(Number<I>{});
-    using type              = tuple<Is>;
-};
+    template <index_t I>
+    __host__ __device__ constexpr const auto& At(Number<I>) const
+    {
+        static_assert(I < base::Size(), "wrong! out of range");
+        return GetElementByKey(detail::TupleElementKey<I>{});
+    }
 
-template <index_t I, class F>
-struct sequence_gen_impl<I, 0, F>
-{
-    using type = Sequence<>;
-};
-
-template <index_t NSize, class F>
-struct sequence_gen
-{
-    using type = typename sequence_gen_impl<0, NSize, F>::type;
+    template <index_t I>
+    __host__ __device__ constexpr auto& At(Number<I>)
+    {
+        static_assert(I < base::Size(), "wrong! out of range");
+        return GetElementByKey(detail::TupleElementKey<I>{});
+    }
 };
 
 } // namespace ck