Skip to content

Add support for signed indexes in gather/scatter #36

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danieltowner merged 3 commits into from
Oct 5, 2026
Merged

Add support for signed indexes in gather/scatter #36

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a648dd5
Fixed namespace issue for simd-iterator test
danieltowner Oct 2, 2026
6fd7dcf
Fixed gather/scatter to work with signed indexes
danieltowner Oct 1, 2026
ea8bd76
Update TestSimdGatherScatter.cpp
danieltowner Oct 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 50 additions & 55 deletions include/xvec/generic/generic_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -686,66 +686,65 @@ compress_mask_by_mask(compact_mask_tag, const _M& v, const _M& selector, bool fi
/// Gather the contents of memory from the supplied index positions.
// :TODO: :COMPILER: Expose the gather llvm primitive and use it directly?
/// @tparam _Range The type of the range
/// @tparam _Idx The index type
/// @tparam _IdxAbi The ABI of the index simd::vec
/// @tparam _Idx The simd index type
/// @tparam _Flags The type of the flags
/// @param r The range from which to gather
/// @param indexes The simd of indexes
/// @param flags The flags controlling the gather
/// @return The gathered simd value
template<detail::contiguous_sized_range _Range, std::integral _Idx, typename _IdxAbi, typename... _Flags>
constexpr auto gather_from(generic_tag, const _Range& r, const basic_vec<_Idx, _IdxAbi>& indexes, flags<_Flags...> flags)
template<detail::contiguous_sized_range _Range, vec_integral _Idx, typename... _Flags>
constexpr auto gather_from(generic_tag, const _Range& r, const _Idx& indexes, flags<_Flags...> flags)
{
using _Tp = std::ranges::range_value_t<_Range>;
const auto rmax = _Idx(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));
const auto rmax = r.size();

if constexpr (contains_flag<flag_unchecked>(flags))
checkStaticMemoryBounds("gather", indexes, rmax);

constexpr int numIndexes = basic_vec<_Idx, _IdxAbi>::size();
constexpr int numIndexes = _Idx::size();

// Put the indexes and the result into memory, as this generates better code
// than trying to pull out individual elements of a simd register to pass to a
// separate load instruction, and to put the results back into a register.
const auto indexPtr = reinterpret_cast<const _Idx*>(&indexes);
const auto indexPtr = reinterpret_cast<const typename _Idx::value_type*>(&indexes);
_Tp result[numIndexes];

for (int i=0; i<numIndexes; ++i)
{
auto idx = indexPtr[i];
result[i] = contains_flag<flag_unchecked>(flags) || (idx < rmax) ? r[idx] : _Tp();
bool validIndex = (idx >= 0 && std::cmp_less(+idx, rmax));
result[i] = contains_flag<flag_unchecked>(flags) || validIndex ? r[idx] : _Tp();
}

return vec<_Tp, numIndexes>(result);
return rebind_t<_Tp, _Idx>(result);
}

/// Gather the contents of memory from the supplied index positions, using a mask.
/// @tparam _Range The type of the range
/// @tparam _Idx The index type
/// @tparam _IdxAbi The ABI of the index simd::vec
/// @tparam _Idx The simd index type
/// @tparam _Flags The type of the flags
/// @param r The range from which to gather
/// @param mask The mask to use for gathering
/// @param indexes The simd of indexes
/// @param flags The flags controlling the gather
/// @return The gathered simd value
template<detail::contiguous_sized_range _Range, std::integral _Idx, typename _IdxAbi, typename... _Flags>
template<detail::contiguous_sized_range _Range, vec_integral _Idx, typename... _Flags>
constexpr auto gather_from(generic_tag, const _Range& r,
const typename basic_vec<_Idx, _IdxAbi>::mask_type& mask,
const basic_vec<_Idx, _IdxAbi>& indexes, flags<_Flags...> flags)
const typename _Idx::mask_type& mask,
const _Idx& indexes, flags<_Flags...> flags)
{
using _Tp = std::ranges::range_value_t<_Range>;
const auto rmax = _Idx(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));
const auto rmax = r.size();

if constexpr (contains_flag<flag_unchecked>(flags))
checkStaticMemoryBounds("gather", indexes, rmax, mask);

constexpr int numIndexes = basic_vec<_Idx, _IdxAbi>::size();
constexpr int numIndexes = _Idx::size();

// Put the indexes and the result into memory, as this generates better code
// than trying to pull out individual elements of a simd register to pass to a
// separate load instruction, and to put the results back into a register.
const auto indexPtr = reinterpret_cast<const _Idx*>(&indexes);
const auto indexPtr = reinterpret_cast<const typename _Idx::value_type*>(&indexes);
_Tp result[numIndexes];

// Need simple mask which can be queried by index.
Expand All @@ -754,83 +753,79 @@ constexpr auto gather_from(generic_tag, const _Range& r,
for (int i=0; i<numIndexes; ++i)
{
auto idx = indexPtr[i];
bool validIndex = contains_flag<flag_unchecked>(flags) || (idx < rmax);
bool validIndex = contains_flag<flag_unchecked>(flags) || (idx >= 0 && std::cmp_less(+idx, rmax));
result[i] = validIndex && mb[i] ? r[idx] : _Tp();
}

return vec<_Tp, numIndexes>(result);
return rebind_t<_Tp, _Idx>(result);
}

/// Scatter the contents of a simd to the supplied index positions.
/// @tparam _Range The type of the range
/// @tparam _Tp The element type
/// @tparam _TpAbi The ABI of the vec
/// @tparam _Idx The index type
/// @tparam _IdxAbi The ABI of the index simd::vec
/// @tparam _Vp The values to scatter to the range.
/// @tparam _Idx The indexes type
/// @tparam _Flags The type of the flags
/// @param values The simd of values to scatter
/// @param r The range to which to scatter
/// @param indexes The simd of indexes
/// @param flags The flags controlling the scatter
template<typename _Range, typename _Tp, typename _TpAbi, typename _Idx, typename _IdxAbi, typename... _Flags>
template<typename _Range, vec_type _Vp, vec_integral _Idx, typename... _Flags>
constexpr void
scatter_to(generic_tag, const basic_vec<_Tp, _TpAbi>& values, _Range&& r,
const basic_vec<_Idx, _IdxAbi>& indexes, flags<_Flags...> flags)
scatter_to(generic_tag, const _Vp& values, _Range&& r, const _Idx& indexes, flags<_Flags...> flags)
{
static_assert(basic_vec<_Tp, _TpAbi>::size() == basic_vec<_Idx, _IdxAbi>::size(),
"Caller should ensure sizes match");
using _Out = std::ranges::range_value_t<_Range>;

const auto rmax = _Idx(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));
const auto mb = (indexes < rmax).to_bitset();
const auto rmax = r.size();
auto data_ptr = std::to_address(r.begin());

if constexpr (contains_flag<flag_unchecked>(flags))
checkStaticMemoryBounds("scatter", indexes, rmax);

[=]<std::size_t... _Iota>(std::index_sequence<_Iota...>) {

if constexpr (detail::contains_flag<detail::flag_unchecked>(flags))
((data_ptr[indexes[_Iota]] = values[_Iota]), ...);
else
((mb[_Iota] ? (data_ptr[indexes[_Iota]] = values[_Iota]) : _Out()), ...);

}(std::make_index_sequence<basic_vec<_Idx, _IdxAbi>::size>());
auto doWrite = [&](auto i) {
const auto idx = indexes[i];
const bool isValidIndex = idx >= 0 && std::cmp_less(+idx, rmax);
if (contains_flag<flag_unchecked>(flags) || isValidIndex)
data_ptr[idx] = values[i];
};

[&]<std::size_t... _Iota>(std::index_sequence<_Iota...>) {
(doWrite(_Iota), ...);
}(std::make_index_sequence<_Idx::size>());
}

/// Scatter the contents of a simd to the supplied index positions, using a mask.
/// @tparam _Range The type of the range
/// @tparam _Tp The element type
/// @tparam _TpAbi The ABI of the vec
/// @tparam _Idx The index type
/// @tparam _IdxAbi The ABI of the index simd::vec
/// @tparam _Vp The values to scatter to the range
/// @tparam _Idx The indexes type
/// @tparam _Flags The type of the flags
/// @param values The simd of values to scatter
/// @param r The range to which to scatter
/// @param mask The mask to use for scattering
/// @param indexes The simd of indexes
/// @param flags The flags controlling the scatter
template<typename _Range, typename _Tp, typename _TpAbi, typename _Idx, typename _IdxAbi, typename... _Flags>
template<typename _Range, vec_type _Vp, vec_integral _Idx, typename... _Flags>
constexpr void
scatter_to(generic_tag, const basic_vec<_Tp, _TpAbi>& values, _Range&& r,
const typename basic_vec<_Idx, _IdxAbi>::mask_type& mask,
const basic_vec<_Idx, _IdxAbi>& indexes, flags<_Flags...> flags)
scatter_to(generic_tag, const _Vp& values, _Range&& r,
const typename _Idx::mask_type& mask,
const _Idx& indexes, flags<_Flags...> flags)
{
using _Out = std::ranges::range_value_t<_Range>;
const auto rmax = _Idx(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));
auto rangeCheckedMask =
contains_flag<detail::flag_unchecked>(flags) ? mask : (indexes < rmax) && mask;
const auto rmax = r.size();
auto data_ptr = std::to_address(r.begin());

if constexpr (contains_flag<flag_unchecked>(flags))
checkStaticMemoryBounds("scatter", indexes, rmax, mask);

const auto mb = rangeCheckedMask.to_bitset();
auto doWrite = [&](auto i) {
if (!mask[i])
return;

const auto idx = indexes[i];
const bool isValidIndex = idx >= 0 && std::cmp_less(+idx, rmax);
if (contains_flag<flag_unchecked>(flags) || isValidIndex)
data_ptr[idx] = values[i];
};

[=]<std::size_t... _Iota>(std::index_sequence<_Iota...>) {
((mb[_Iota] ? (data_ptr[indexes[_Iota]] = values[_Iota]) : _Out()), ...);
}(std::make_index_sequence<basic_vec<_Idx, _IdxAbi>::size>());
(doWrite(_Iota), ...);
}(std::make_index_sequence<_Idx::size>());
}

/// :TODO: FMA - Don't provide generic - should always match a real instruction to get the right precision?
Expand Down
47 changes: 23 additions & 24 deletions include/xvec/x86/x86_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -690,50 +690,49 @@ constexpr auto gather_from(x86_avx2_tag, const _Range& r,
{
using _Tp = std::ranges::range_value_t<_Range>;
auto data_ptr = std::to_address(r.begin());
const auto rmax = _Idx(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));

target_overloads impl {

#if defined(__AVX512F__)
// 32-bit Intel AVX-512
[=](xmm_register<uint32_t> auto i, auto m) { return _mm_mmask_i32gather_epi32(__m128i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](ymm_register<uint32_t> auto i, auto m) { return _mm256_mmask_i32gather_epi32(__m256i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](xmm_register<int32_t> auto i, auto m) { return _mm_mmask_i32gather_epi32(__m128i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](ymm_register<int32_t> auto i, auto m) { return _mm256_mmask_i32gather_epi32(__m256i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },

[=](xmm_register<uint64_t> auto i, auto m) { return _mm_mmask_i64gather_epi64(__m128i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](ymm_register<uint64_t> auto i, auto m) { return _mm256_mmask_i64gather_epi64(__m256i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](xmm_register<int64_t> auto i, auto m) { return _mm_mmask_i64gather_epi64(__m128i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](ymm_register<int64_t> auto i, auto m) { return _mm256_mmask_i64gather_epi64(__m256i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },

[=](zmm_register<uint32_t> auto i, auto m) { return _mm512_mask_i32gather_epi32(__m512i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](zmm_register<uint64_t> auto i, auto m) { return _mm512_mask_i64gather_epi64(__m512(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](zmm_register<int32_t> auto i, auto m) { return _mm512_mask_i32gather_epi32(__m512i(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },
[=](zmm_register<int64_t> auto i, auto m) { return _mm512_mask_i64gather_epi64(__m512(), m.to_register(), i.to_register(), data_ptr, sizeof(_Tp)); },

#elif defined(__AVX2__)
// Note the explicit growth of the mask to the correct register size to
// avoid inadvertently leaving uninitialised elements active.

// 32-bit Intel AVX2
[=](xmm_register<uint32_t> auto i, auto m)
[=](xmm_register<int32_t> auto i, auto m)
{ return _mm_mask_i32gather_epi32(__m128i(), data_ptr, i.to_register(), grow<4>(m).to_register(), sizeof(_Tp)); },
[=](ymm_register<uint32_t> auto i, auto m)
[=](ymm_register<int32_t> auto i, auto m)
{ return _mm256_mask_i32gather_epi32(__m256i(), data_ptr, i.to_register(), grow<8>(m).to_register(), sizeof(_Tp)); },

// 64-bit Intel AVX2
[=](xmm_register<uint64_t> auto i, auto m)
[=](xmm_register<int64_t> auto i, auto m)
{ return _mm_mask_i64gather_epi64(__m128i(), data_ptr, i.to_register(), grow<2>(m).to_register(), sizeof(_Tp)); },
[=](ymm_register<uint64_t> auto i, auto m)
[=](ymm_register<int64_t> auto i, auto m)
{ return _mm256_mask_i64gather_epi64(__m256i(), data_ptr, i.to_register(), grow<4>(m).to_register(), sizeof(_Tp)); },

#endif

[=](auto unhandled, auto) { static_assert(dependent_false<decltype(unhandled)>); }
};

// Convert the indexes to at least 32-bit, since no instruction exists
// operating at a smaller granularity.
constexpr auto idxNumBytes = std::max({sizeof(uint32_t), sizeof(_Tp), sizeof(_Idx)});
using _IW = container_for_num_bytes<idxNumBytes>;
// The indexes and the values to gether must be the same size, and at least
// 32-bit signed to match the hardware expectation.
constexpr auto idxNumBytes = std::max({sizeof(int32_t), sizeof(_Tp), sizeof(_Idx)});
using _IW = std::make_signed_t<container_for_num_bytes<idxNumBytes>>;
const auto rmax = std::min<_IW>(r.size(), std::numeric_limits<_Idx>::max());

auto wrapper = [=]<typename _VecIdx>(_VecIdx idx, auto m) {
// Bump the indexes and mask up to the preferred size if necessary.
using RetypedIndex = vec<_IW, _VecIdx::size>;
using RetypedIndex = rebind_t<_IW, _VecIdx>;
auto ri = RetypedIndex(idx);

// Note that the mask is actively resized to zero out unwanted positions.
Expand All @@ -744,7 +743,7 @@ constexpr auto gather_from(x86_avx2_tag, const _Range& r,

// Add a bounds check if necessary.
if constexpr (!contains_flag<flag_unchecked>(flags))
rm = rm && (ri < rmax);
rm = rm && (ri >= _IW(0) && ri < rmax);

// Do the gather, reading back elements which are at least as big as the required type, and then downcasting to the types size if necessary.
auto t = RetypedIndex(impl(ri, rm));
Expand Down Expand Up @@ -774,7 +773,7 @@ constexpr auto gather_from(x86_avx2_tag, const _Range& r,

inline auto scatterImpl =
[]<typename _Tp, typename _VT, typename _IT>(basic_vec<_Tp, _VT> v,
basic_vec<uint32_t, _IT> idx, auto ptr, uint64_t m)
basic_vec<int32_t, _IT> idx, auto ptr, uint64_t m)
{
constexpr uint64_t justMaskBits = (1ULL << basic_vec<_Tp, _VT>::size()) - 1;
const uint64_t nm = justMaskBits & m;
Expand Down Expand Up @@ -807,23 +806,23 @@ scatter_to(x86_avx512_tag, const basic_vec<_Tp, _TpAbi>& values, _Range&& r,
const basic_vec<_Idx, _IdxAbi>& indexes, flags<_Flags...> flags)
{
auto ptr = std::to_address(r.begin());
const auto rmax = _Idx(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));
const auto rmax = int32_t(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));

// Bind extra pointer and default mask parameters to the scatter lambda. Note
// that a range check is added if necessary too.
auto impl = [=](auto v, auto i) {
if constexpr (contains_flag<detail::flag_unchecked>(flags))
scatterImpl(v, i, ptr, ~0ULL);
else
scatterImpl(v, i, ptr, (i < rmax).to_register());
scatterImpl(v, i, ptr, (i >= 0 && i < rmax).to_register());
};

if constexpr (contains_flag<flag_unchecked>(flags))
checkStaticMemoryBounds("scatter", indexes, rmax);

// All scatter instructions can be made to use 32-bit indexes.
// :TODO: A special case might be 64-bit values using 64-bit indexes
auto indexAs32 = vec<uint32_t, basic_vec<_Idx, _IdxAbi>::size>(indexes);
auto indexAs32 = vec<int32_t, basic_vec<_Idx, _IdxAbi>::size>(indexes);
chunked_invoke<(int)basic_vec<_Tp>::size>(impl, vec_as_container(values), indexAs32);
}

Expand All @@ -840,7 +839,7 @@ scatter_to(x86_avx512_tag, const basic_vec<_Tp, _TpAbi>& values, _Range&& r,
// Range check if required.
const auto rmax = _Idx(std::min(r.size(), size_t(std::numeric_limits<_Idx>::max())));
auto rangeCheckedMask =
contains_flag<detail::flag_unchecked>(flags) ? mask : (indexes < rmax) && mask;
contains_flag<detail::flag_unchecked>(flags) ? mask : (indexes >= _Idx(0) && indexes < rmax) && mask;

if constexpr (contains_flag<flag_unchecked>(flags))
checkStaticMemoryBounds("scatter", indexes, rmax, mask);
Expand All @@ -850,7 +849,7 @@ scatter_to(x86_avx512_tag, const basic_vec<_Tp, _TpAbi>& values, _Range&& r,

// All scatter instructions can be made to use 32-bit indexes.
// :TODO: A special case might be 64-bit values using 64-bit indexes
auto indexAs32 = vec<uint32_t, basic_vec<_Idx, _IdxAbi>::size>(indexes);
auto indexAs32 = vec<int32_t, basic_vec<_Idx, _IdxAbi>::size>(indexes);
chunked_invoke<(int)basic_vec<_Tp>::size()>(impl, vec_as_container(values), indexAs32, rangeCheckedMask);
}
#endif
Expand Down
Loading
Loading