Add container and tuple optimization helpers

- Replace lambdas with named functors in container_concat - Add make_uniform_tuple helper for repeated value patterns - Add container_product helper with O(1) depth fold expression - Add merge_sequences_functor and unpack_and_merge_sequences - Add 16 unit tests for container helpers Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-07 15:47:05 +00:00 · 2026-01-22 01:47:46 +00:00
parent fcc9372c00
commit 99872ec8b3
5 changed files with 275 additions and 4 deletions
--- a/include/ck/utility/container_helper.hpp
+++ b/include/ck/utility/container_helper.hpp
@@ -189,6 +189,14 @@ __host__ __device__ constexpr auto container_reduce(const Container& x,
 }
 #endif

+// O(1) template depth alternative to container_reduce for computing products.
+// Uses fold expression via unpack instead of O(N) linear recursion.
+template <typename Container>
+__host__ __device__ constexpr auto container_product(const Container& x)
+{
+    return unpack([](auto... xs) { return (xs * ...); }, x);
+}
+
 template <typename TData, index_t NSize, typename Reduce>
 __host__ __device__ constexpr auto
 container_reverse_inclusive_scan(const Array<TData, NSize>& x, Reduce f, TData init)
@@ -316,6 +324,26 @@ container_reverse_inclusive_scan(const Tuple<Xs...>& x, Reduce f, TData init)
    return y;
 }

+// Named functors for container_concat to reduce template instantiations
+// (lambdas create unique types per call site, functors are shared)
+struct make_tuple_functor
+{
+    template <typename... Ts>
+    __host__ __device__ constexpr auto operator()(Ts&&... xs) const
+    {
+        return make_tuple(ck::forward<Ts>(xs)...);
+    }
+};
+
+struct make_array_functor
+{
+    template <typename T, typename... Ts>
+    __host__ __device__ constexpr auto operator()(T&& x, Ts&&... xs) const
+    {
+        return make_array(ck::forward<T>(x), ck::forward<Ts>(xs)...);
+    }
+};
+
 template <typename X, typename... Ys>
 __host__ __device__ constexpr auto container_concat(const X& x, const Ys&... ys)
 {
@@ -325,15 +353,13 @@ __host__ __device__ constexpr auto container_concat(const X& x, const Ys&... ys)
 template <typename T, index_t NX, index_t NY>
 __host__ __device__ constexpr auto container_concat(const Array<T, NX>& ax, const Array<T, NY>& ay)
 {
-    return unpack2(
-        [&](auto&&... zs) { return make_array(ck::forward<decltype(zs)>(zs)...); }, ax, ay);
+    return unpack2(make_array_functor{}, ax, ay);
 }

 template <typename... X, typename... Y>
 __host__ __device__ constexpr auto container_concat(const Tuple<X...>& tx, const Tuple<Y...>& ty)
 {
-    return unpack2(
-        [&](auto&&... zs) { return make_tuple(ck::forward<decltype(zs)>(zs)...); }, tx, ty);
+    return unpack2(make_tuple_functor{}, tx, ty);
 }

 template <typename Container>
--- a/include/ck/utility/sequence_helper.hpp
+++ b/include/ck/utility/sequence_helper.hpp
@@ -34,4 +34,22 @@ __host__ __device__ constexpr auto to_sequence(Tuple<Number<Is>...>)
    return Sequence<Is...>{};
 }

+// Functor for merge_sequences to avoid lambda instantiation overhead
+struct merge_sequences_functor
+{
+    template <typename... Seqs>
+    __host__ __device__ constexpr auto operator()(Seqs... seqs) const
+    {
+        return merge_sequences(seqs...);
+    }
+};
+
+// Helper to unpack a tuple of sequences and merge them
+// Replaces: unpack([](auto... xs) { return merge_sequences(xs...); }, tuple_of_sequences)
+template <typename TupleOfSequences>
+__host__ __device__ constexpr auto unpack_and_merge_sequences(TupleOfSequences)
+{
+    return unpack(merge_sequences_functor{}, TupleOfSequences{});
+}
+
 } // namespace ck
--- a/include/ck/utility/tuple_helper.hpp
+++ b/include/ck/utility/tuple_helper.hpp
@@ -37,6 +37,50 @@ __host__ __device__ constexpr auto generate_tie(F&& f, Number<N>)
                  typename arithmetic_sequence_gen<0, N, 1>::type{});
 }

+// Optimized helper for common pattern: generate_tuple([](auto i) { return Sequence<i.value>{}; },
+// N) Creates Tuple<Sequence<0>, Sequence<1>, ..., Sequence<N-1>> without lambda instantiation
+namespace detail {
+template <index_t... Is>
+__host__ __device__ constexpr auto make_identity_sequences_impl(Sequence<Is...>)
+{
+    return make_tuple(Sequence<Is>{}...);
+}
+} // namespace detail
+
+template <index_t N>
+__host__ __device__ constexpr auto generate_identity_sequences()
+{
+    return detail::make_identity_sequences_impl(make_index_sequence<N>{});
+}
+
+template <index_t N>
+__host__ __device__ constexpr auto generate_identity_sequences(Number<N>)
+{
+    return generate_identity_sequences<N>();
+}
+
+// Optimized helper for common pattern: generate_tuple([&](auto) { return value; }, Number<N>{})
+// Creates Tuple<T, T, ..., T> (N copies) without lambda instantiation
+namespace detail {
+template <typename T, index_t... Is>
+__host__ __device__ constexpr auto make_uniform_tuple_impl(T&& value, Sequence<Is...>)
+{
+    return make_tuple(((void)Is, value)...);
+}
+} // namespace detail
+
+template <index_t N, typename T>
+__host__ __device__ constexpr auto make_uniform_tuple(T&& value)
+{
+    return detail::make_uniform_tuple_impl(static_cast<T&&>(value), make_index_sequence<N>{});
+}
+
+template <typename T, index_t N>
+__host__ __device__ constexpr auto make_uniform_tuple(T&& value, Number<N>)
+{
+    return make_uniform_tuple<N>(static_cast<T&&>(value));
+}
+
 // tx and ty are tuple of references, return type of will tuple of referennce (not rvalue)
 template <typename... X, typename... Y>
 __host__ __device__ constexpr auto concat_tuple_of_reference(const Tuple<X&...>& tx,
--- a/test/util/CMakeLists.txt
+++ b/test/util/CMakeLists.txt
@@ -5,3 +5,8 @@ add_gtest_executable(unit_sequence unit_sequence.cpp)
 if(result EQUAL 0)
  target_link_libraries(unit_sequence PRIVATE utility)
 endif()
+
+add_gtest_executable(unit_container_helper unit_container_helper.cpp)
+if(result EQUAL 0)
+  target_link_libraries(unit_container_helper PRIVATE utility)
+endif()
--- a/test/util/unit_container_helper.cpp
+++ b/test/util/unit_container_helper.cpp
@@ -0,0 +1,178 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <gtest/gtest.h>
+#include "ck/utility/container_helper.hpp"
+#include "ck/utility/tuple_helper.hpp"
+
+using namespace ck;
+
+// Test container_concat with tuples
+TEST(ContainerConcat, ConcatTwoTuples)
+{
+    constexpr auto t1     = make_tuple(Number<7>{}, Number<11>{});
+    constexpr auto t2     = make_tuple(Number<13>{}, Number<17>{});
+    constexpr auto result = container_concat(t1, t2);
+
+    EXPECT_EQ(result.Size(), 4);
+    EXPECT_EQ(result[Number<0>{}], 7);
+    EXPECT_EQ(result[Number<1>{}], 11);
+    EXPECT_EQ(result[Number<2>{}], 13);
+    EXPECT_EQ(result[Number<3>{}], 17);
+}
+
+TEST(ContainerConcat, ConcatThreeTuples)
+{
+    constexpr auto t1     = make_tuple(Number<19>{});
+    constexpr auto t2     = make_tuple(Number<23>{}, Number<29>{});
+    constexpr auto t3     = make_tuple(Number<31>{});
+    constexpr auto result = container_concat(t1, t2, t3);
+
+    EXPECT_EQ(result.Size(), 4);
+    EXPECT_EQ(result[Number<0>{}], 19);
+    EXPECT_EQ(result[Number<1>{}], 23);
+    EXPECT_EQ(result[Number<2>{}], 29);
+    EXPECT_EQ(result[Number<3>{}], 31);
+}
+
+TEST(ContainerConcat, ConcatWithEmptyTuple)
+{
+    constexpr auto t1     = make_tuple(Number<37>{}, Number<41>{});
+    constexpr auto empty  = make_tuple();
+    constexpr auto result = container_concat(t1, empty);
+
+    EXPECT_EQ(result.Size(), 2);
+    EXPECT_EQ(result[Number<0>{}], 37);
+    EXPECT_EQ(result[Number<1>{}], 41);
+}
+
+TEST(ContainerConcat, ConcatSingleTuple)
+{
+    constexpr auto t1     = make_tuple(Number<43>{}, Number<47>{}, Number<53>{});
+    constexpr auto result = container_concat(t1);
+
+    EXPECT_EQ(result.Size(), 3);
+    EXPECT_EQ(result[Number<0>{}], 43);
+    EXPECT_EQ(result[Number<1>{}], 47);
+    EXPECT_EQ(result[Number<2>{}], 53);
+}
+
+// Test container_concat with arrays
+TEST(ContainerConcat, ConcatTwoArrays)
+{
+    constexpr auto a1     = make_array(59, 61);
+    constexpr auto a2     = make_array(67, 71);
+    constexpr auto result = container_concat(a1, a2);
+
+    EXPECT_EQ(result.Size(), 4);
+    EXPECT_EQ(result[Number<0>{}], 59);
+    EXPECT_EQ(result[Number<1>{}], 61);
+    EXPECT_EQ(result[Number<2>{}], 67);
+    EXPECT_EQ(result[Number<3>{}], 71);
+}
+
+// Test make_uniform_tuple
+TEST(MakeUniformTuple, Size3)
+{
+    constexpr auto result = make_uniform_tuple<3>(Number<73>{});
+
+    EXPECT_EQ(result.Size(), 3);
+    EXPECT_EQ(result[Number<0>{}], 73);
+    EXPECT_EQ(result[Number<1>{}], 73);
+    EXPECT_EQ(result[Number<2>{}], 73);
+}
+
+TEST(MakeUniformTuple, Size1)
+{
+    constexpr auto result = make_uniform_tuple<1>(Number<79>{});
+
+    EXPECT_EQ(result.Size(), 1);
+    EXPECT_EQ(result[Number<0>{}], 79);
+}
+
+TEST(MakeUniformTuple, Size0)
+{
+    constexpr auto result = make_uniform_tuple<0>(Number<83>{});
+
+    EXPECT_EQ(result.Size(), 0);
+}
+
+TEST(MakeUniformTuple, Size5)
+{
+    constexpr auto result = make_uniform_tuple<5>(Number<89>{});
+
+    EXPECT_EQ(result.Size(), 5);
+    EXPECT_EQ(result[Number<0>{}], 89);
+    EXPECT_EQ(result[Number<1>{}], 89);
+    EXPECT_EQ(result[Number<2>{}], 89);
+    EXPECT_EQ(result[Number<3>{}], 89);
+    EXPECT_EQ(result[Number<4>{}], 89);
+}
+
+// Test make_tuple_functor (used internally by container_concat)
+TEST(MakeTupleFunctor, CreatesTuple)
+{
+    make_tuple_functor functor;
+    auto result = functor(Number<97>{}, Number<101>{}, Number<103>{});
+
+    EXPECT_EQ(result.Size(), 3);
+    EXPECT_EQ(result[Number<0>{}], 97);
+    EXPECT_EQ(result[Number<1>{}], 101);
+    EXPECT_EQ(result[Number<2>{}], 103);
+}
+
+// Test container_push_front and container_push_back
+TEST(ContainerPush, PushFront)
+{
+    constexpr auto t      = make_tuple(Number<109>{}, Number<113>{});
+    constexpr auto result = container_push_front(t, Number<107>{});
+
+    EXPECT_EQ(result.Size(), 3);
+    EXPECT_EQ(result[Number<0>{}], 107);
+    EXPECT_EQ(result[Number<1>{}], 109);
+    EXPECT_EQ(result[Number<2>{}], 113);
+}
+
+TEST(ContainerPush, PushBack)
+{
+    constexpr auto t      = make_tuple(Number<127>{}, Number<131>{});
+    constexpr auto result = container_push_back(t, Number<137>{});
+
+    EXPECT_EQ(result.Size(), 3);
+    EXPECT_EQ(result[Number<0>{}], 127);
+    EXPECT_EQ(result[Number<1>{}], 131);
+    EXPECT_EQ(result[Number<2>{}], 137);
+}
+
+// Test container_product
+TEST(ContainerProduct, TupleOfNumbers)
+{
+    constexpr auto t      = make_tuple(Number<2>{}, Number<3>{}, Number<5>{});
+    constexpr auto result = container_product(t);
+
+    EXPECT_EQ(result, 30); // 2 * 3 * 5 = 30
+}
+
+TEST(ContainerProduct, ArrayOfIntegers)
+{
+    constexpr auto a      = make_array(7, 11, 13);
+    constexpr auto result = container_product(a);
+
+    EXPECT_EQ(result, 1001); // 7 * 11 * 13 = 1001
+}
+
+TEST(ContainerProduct, SingleElement)
+{
+    constexpr auto t      = make_tuple(Number<139>{});
+    constexpr auto result = container_product(t);
+
+    EXPECT_EQ(result, 139);
+}
+
+TEST(ContainerProduct, WithOne)
+{
+    constexpr auto t      = make_tuple(Number<1>{}, Number<17>{}, Number<19>{});
+    constexpr auto result = container_product(t);
+
+    EXPECT_EQ(result, 323); // 1 * 17 * 19 = 323
+}