Skip to content

Commit 04e187a

Browse files
committed
[libc] Implement vector 'split' and 'concat' routines
Summary: This provides some helpers for the split and concatenation routines for changing the size of an existing vector. This includes a simple tuple type to do the splitting. The tuple doesn't support structured bindings yet. The concat function is more limited than what would be ideal, but the shufflevector builtin requires things of equivalent sizes and I didn't think it was worth wrangling with that just yet.
1 parent 43aa05b commit 04e187a

File tree

3 files changed

+102
-0
lines changed

3 files changed

+102
-0
lines changed

libc/src/__support/CPP/CMakeLists.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,4 +223,7 @@ add_header_library(
223223
simd
224224
HDRS
225225
simd.h
226+
DEPENDS
227+
.utility
228+
.tuple
226229
)

libc/src/__support/CPP/simd.h

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
#include "hdr/stdint_proxy.h"
1717
#include "src/__support/CPP/algorithm.h"
1818
#include "src/__support/CPP/limits.h"
19+
#include "src/__support/CPP/tuple.h"
1920
#include "src/__support/CPP/type_traits.h"
21+
#include "src/__support/CPP/utility/integer_sequence.h"
2022
#include "src/__support/macros/attributes.h"
2123
#include "src/__support/macros/config.h"
2224

@@ -51,6 +53,7 @@ template <typename T> inline constexpr size_t native_vector_size = 1;
5153
template <typename T> LIBC_INLINE constexpr T poison() {
5254
return __builtin_nondeterministic_value(T());
5355
}
56+
5457
} // namespace internal
5558

5659
// Type aliases.
@@ -267,6 +270,77 @@ LIBC_INLINE constexpr simd<T, N> select(simd<bool, N> m, simd<T, N> x,
267270
return m ? x : y;
268271
}
269272

273+
namespace internal {
274+
template <typename T, size_t N, size_t O, size_t... I>
275+
LIBC_INLINE constexpr static cpp::simd<T, sizeof...(I)>
276+
extend(cpp::simd<T, N> x, cpp::index_sequence<I...>) {
277+
return __builtin_shufflevector(x, x, (I < O ? static_cast<int>(I) : -1)...);
278+
}
279+
template <typename T, size_t N, size_t M, size_t O>
280+
LIBC_INLINE constexpr static auto extend(cpp::simd<T, N> x) {
281+
if constexpr (N == M)
282+
return x;
283+
else if constexpr (M <= 2 * N)
284+
return extend<T, N, M>(x, cpp::make_index_sequence<M>{});
285+
else
286+
return extend<T, 2 * N, M, O>(
287+
extend<T, N, 2 * N>(x, cpp::make_index_sequence<2 * N>{}));
288+
}
289+
template <typename T, size_t N, size_t M, size_t... I>
290+
LIBC_INLINE constexpr static cpp::simd<T, N + M>
291+
concat(cpp::simd<T, N> x, cpp::simd<T, M> y, cpp::index_sequence<I...>) {
292+
constexpr size_t L = (N > M ? N : M);
293+
294+
auto x_ext = extend<T, N, L, N>(x);
295+
auto y_ext = extend<T, M, L, M>(y);
296+
297+
auto remap = [](size_t idx) -> int {
298+
if (idx < N)
299+
return static_cast<int>(idx);
300+
if (idx < N + M)
301+
return static_cast<int>((idx - N) + L);
302+
return -1;
303+
};
304+
305+
return __builtin_shufflevector(x_ext, y_ext, remap(I)...);
306+
}
307+
308+
template <typename T, size_t N, size_t Count, size_t Offset, size_t... I>
309+
LIBC_INLINE constexpr static cpp::simd<T, Count>
310+
slice(cpp::simd<T, N> x, cpp::index_sequence<I...>) {
311+
return __builtin_shufflevector(x, x, (Offset + I)...);
312+
}
313+
template <typename T, size_t N, size_t Offset, size_t Head, size_t... Tail>
314+
LIBC_INLINE constexpr static auto split(cpp::simd<T, N> x) {
315+
auto first = cpp::make_tuple(
316+
slice<T, N, Head, Offset>(x, cpp::make_index_sequence<Head>{}));
317+
if constexpr (sizeof...(Tail) > 0)
318+
return cpp::tuple_cat(first, split<T, N, Offset + Head, Tail...>(x));
319+
else
320+
return first;
321+
}
322+
323+
} // namespace internal
324+
325+
// Shuffling helpers.
326+
template <typename T, size_t N, size_t M>
327+
LIBC_INLINE constexpr static auto concat(cpp::simd<T, N> x, cpp::simd<T, M> y) {
328+
return internal::concat(x, y, make_index_sequence<N + M>{});
329+
}
330+
template <typename T, size_t N, size_t M, typename... Rest>
331+
LIBC_INLINE constexpr static auto concat(cpp::simd<T, N> x, cpp::simd<T, M> y,
332+
Rest... rest) {
333+
auto xy = concat(x, y);
334+
if constexpr (sizeof...(Rest))
335+
return concat(xy, rest...);
336+
else
337+
return xy;
338+
}
339+
template <size_t... Sizes, typename T, size_t N> auto split(cpp::simd<T, N> x) {
340+
static_assert((... + Sizes) == N, "split sizes must sum to vector size");
341+
return internal::split<T, N, 0, Sizes...>(x);
342+
}
343+
270344
// TODO: where expressions, scalar overloads, ABI types.
271345

272346
} // namespace cpp

libc/test/src/__support/CPP/simd_test.cpp

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,3 +68,28 @@ TEST(LlvmLibcSIMDTest, MaskOperations) {
6868
EXPECT_EQ(find_first_set(mask), 0);
6969
EXPECT_EQ(find_last_set(mask), 2);
7070
}
71+
72+
TEST(LlvmLibcSIMDTest, SplitConcat) {
73+
simd<char, 8> v(1);
74+
auto [v1, v2, v3, v4] = split<2, 2, 2, 2>(v);
75+
static_assert(
76+
simd_size_v<decltype(v1)> == 2 && simd_size_v<decltype(v2)> == 2 &&
77+
simd_size_v<decltype(v3)> == 2 && simd_size_v<decltype(v4)> == 2,
78+
"invalid size");
79+
80+
v1 = simd<char, 2>(1);
81+
v2 = simd<char, 2>(2);
82+
v3 = simd<char, 2>(3);
83+
v4 = simd<char, 2>(4);
84+
simd<char, 8> m = concat(v1, v2, v3, v4);
85+
static_assert(simd_size_v<decltype(m)> == 8, "invalid size");
86+
87+
simd<char, 8> c = {1, 1, 2, 2, 3, 3, 4, 4};
88+
for (int i = 0; i < 8; ++i)
89+
EXPECT_EQ(c[i], m[i]);
90+
91+
simd<char, 1> c1('\0');
92+
simd<char, 8> c2('\0');
93+
simd<char, 9> c3 = concat(c1, c2);
94+
static_assert(simd_size_v<decltype(c3)> == 9, "invalid size");
95+
}

0 commit comments

Comments
 (0)