[CK TILE] Refactor sequence_reverse_inclusive_scan

include/ck_tile/core.hpp

@@ -41,6 +41,7 @@
 #include "ck_tile/core/container/span.hpp"
 #include "ck_tile/core/container/statically_indexed_array.hpp"
 #include "ck_tile/core/container/thread_buffer.hpp"
+#include "ck_tile/core/container/trivial_array.hpp"
 #include "ck_tile/core/container/tuple.hpp"
 #include "ck_tile/core/numeric/bfloat16.hpp"
 #include "ck_tile/core/numeric/e8m0.hpp"

include/ck_tile/core/container/sequence.hpp

@@ -4,6 +4,7 @@
 #pragma once
 
 #include "ck_tile/core/config.hpp"
+#include "ck_tile/core/container/trivial_array.hpp"
 #include "ck_tile/core/numeric/integer.hpp"
 #include "ck_tile/core/numeric/integral_constant.hpp"
 #include "ck_tile/core/numeric/math.hpp"
@@ -35,6 +36,7 @@ template <typename Seq>
 CK_TILE_HOST_DEVICE constexpr auto sequence_pop_back(Seq);
 
 namespace impl {
+
 // static_assert(__has_builtin(__type_pack_element), "can't find __type_pack_element");
 template <index_t I, typename... Ts>
 using at_index_t = __type_pack_element<I, Ts...>;
@@ -332,29 +334,46 @@ struct uniform_sequence_gen
 };
 
 // reverse inclusive scan (with init) sequence
-template <typename, typename, index_t>
-struct sequence_reverse_inclusive_scan;
-
-template <index_t I, index_t... Is, typename Reduce, index_t Init>
-struct sequence_reverse_inclusive_scan<sequence<I, Is...>, Reduce, Init>
-{
-    using old_scan = typename sequence_reverse_inclusive_scan<sequence<Is...>, Reduce, Init>::type;
-
-    static constexpr index_t new_reduce = Reduce{}(I, old_scan{}.front());
+namespace impl {
 
-    using type = typename sequence_merge<sequence<new_reduce>, old_scan>::type;
-};
+template <typename Seq, typename Reduce, index_t Init>
+struct sequence_reverse_inclusive_scan_impl;
 
-template <index_t I, typename Reduce, index_t Init>
-struct sequence_reverse_inclusive_scan<sequence<I>, Reduce, Init>
+template <index_t... Is, typename Reduce, index_t Init>
+struct sequence_reverse_inclusive_scan_impl<sequence<Is...>, Reduce, Init>
 {
-    using type = sequence<Reduce{}(I, Init)>;
+    template <index_t... Indices>
+    static constexpr auto compute(sequence<Indices...>)
+    {
+        constexpr index_t size = sizeof...(Is);
+        if constexpr(size == 0)
+        {
+            return sequence<>{};
+        }
+        else
+        {
+            constexpr auto arr = []() {
+                trivial_array<index_t, size> values = {Is...};
+                trivial_array<index_t, size> result = {0};
+                result[size - 1]                    = Reduce{}(values[size - 1], Init);
+                for(index_t i = size - 1; i > 0; --i)
+                {
+                    result[i - 1] = Reduce{}(values[i - 1], result[i]);
+                }
+                return result;
+            }();
+            return sequence<arr[Indices]...>{};
+        }
+    }
+
+    using type = decltype(compute(make_index_sequence<sizeof...(Is)>{}));
 };
+} // namespace impl
 
-template <typename Reduce, index_t Init>
-struct sequence_reverse_inclusive_scan<sequence<>, Reduce, Init>
+template <typename Seq, typename Reduce, index_t Init>
+struct sequence_reverse_inclusive_scan
 {
-    using type = sequence<>;
+    using type = typename impl::sequence_reverse_inclusive_scan_impl<Seq, Reduce, Init>::type;
 };
 
 // split sequence
@@ -1105,16 +1124,13 @@ struct sorted_sequence_histogram<h_idx, sequence<x>, sequence<r, rs...>>
 };
 } // namespace detail
 
-template <typename, index_t>
-struct array; // declare for later use (array->seq utility)
-
 // SeqSortedSamples: <0, 2, 3, 5, 7>, SeqRange: <0, 3, 6, 9> -> SeqHistogram : <2, 2, 1>
 template <typename SeqSortedSamples, index_t r, index_t... rs>
 CK_TILE_HOST_DEVICE constexpr auto histogram_sorted_sequence(SeqSortedSamples, sequence<r, rs...>)
 {
     constexpr auto bins      = sizeof...(rs); // or categories
     constexpr auto histogram = [&]() {
-        array<index_t, bins> h{0}; // make sure this can clear all element to zero
+        trivial_array<index_t, bins> h{0}; // make sure this can clear all element to zero
         detail::sorted_sequence_histogram<0, SeqSortedSamples, sequence<rs...>>{}(h);
         return h;
     }();
@@ -1127,7 +1143,7 @@ CK_TILE_HOST_DEVICE constexpr auto generate_array(F&& f, number<N>)
 {
     using T = remove_cvref_t<decltype(f(number<0>{}))>;
 
-    return unpack([&f](auto&&... is) { return array<T, N>{f(is)...}; },
+    return unpack([&f](auto&&... is) { return trivial_array<T, N>{f(is)...}; },
                   typename arithmetic_sequence_gen<0, N, 1>::type{});
 }
 

include/ck_tile/core/container/trivial_array.hpp

@@ -0,0 +1,21 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+#include "ck_tile/core/numeric/integer.hpp"
+
+namespace ck_tile {
+// trivial_array has no custom constructor, allowing it to be used as a constexpr variable type
+template <typename T, index_t N>
+struct trivial_array
+{
+    // Public aggregate initialization makes this a literal type
+    T data[N];
+
+    // Basic constexpr accessors
+    constexpr const T& operator[](index_t i) const { return data[i]; }
+    constexpr T& operator[](index_t i) { return data[i]; }
+
+    constexpr index_t size() const { return N; }
+};
+} // namespace ck_tile

test/ck_tile/core/CMakeLists.txt

@@ -2,3 +2,4 @@
 # SPDX-License-Identifier: MIT
 
 add_subdirectory(arch)
+add_subdirectory(container)

test/ck_tile/core/container/CMakeLists.txt

@@ -0,0 +1,7 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+add_gtest_executable(ck_tile_unit_sequence unit_sequence.cpp)
+if(result EQUAL 0)
+  target_link_libraries(ck_tile_unit_sequence PRIVATE utility)
+endif()