Skip to content

Commit 978d222

Browse files
DiamonDinoiaDiamonDinoia
committed
Add AVX512VL gather scatter kernels
1 parent 7778faa commit 978d222

2 files changed

Lines changed: 132 additions & 0 deletions

File tree

‎include/xsimd/arch/xsimd_avx512vl_128.hpp‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,72 @@ namespace xsimd
335335
detail::maskstore128(mem, src, mask.mask(), Mode {});
336336
}
337337

338+
// gather
339+
template <class T, class A, class U, detail::enable_sized_integral_t<T, 4> = 0, detail::enable_sized_integral_t<U, 4> = 0>
340+
XSIMD_INLINE batch<T, A> gather(batch<T, A> const&, T const* src, batch<U, A> const& index,
341+
kernel::requires_arch<avx512vl_128>) noexcept
342+
{
343+
return _mm_mmask_i32gather_epi32(_mm_setzero_si128(), (__mmask8)0xF, index, static_cast<const void*>(src), sizeof(T));
344+
}
345+
346+
template <class T, class A, class U, detail::enable_sized_integral_t<T, 8> = 0, detail::enable_sized_integral_t<U, 8> = 0>
347+
XSIMD_INLINE batch<T, A> gather(batch<T, A> const&, T const* src, batch<U, A> const& index,
348+
kernel::requires_arch<avx512vl_128>) noexcept
349+
{
350+
return _mm_mmask_i64gather_epi64(_mm_setzero_si128(), (__mmask8)0x3, index, static_cast<const void*>(src), sizeof(T));
351+
}
352+
353+
template <class A, class U, detail::enable_sized_integral_t<U, 4> = 0>
354+
XSIMD_INLINE batch<float, A> gather(batch<float, A> const&, float const* src,
355+
batch<U, A> const& index,
356+
kernel::requires_arch<avx512vl_128>) noexcept
357+
{
358+
return _mm_mmask_i32gather_ps(_mm_setzero_ps(), (__mmask8)0xF, index, src, sizeof(float));
359+
}
360+
361+
template <class A, class U, detail::enable_sized_integral_t<U, 8> = 0>
362+
XSIMD_INLINE batch<double, A> gather(batch<double, A> const&, double const* src,
363+
batch<U, A> const& index,
364+
requires_arch<avx512vl_128>) noexcept
365+
{
366+
return _mm_mmask_i64gather_pd(_mm_setzero_pd(), (__mmask8)0x3, index, src, sizeof(double));
367+
}
368+
369+
// scatter
370+
template <class A, class T,
371+
class = std::enable_if_t<std::is_same<uint32_t, T>::value || std::is_same<int32_t, T>::value>>
372+
XSIMD_INLINE void scatter(batch<T, A> const& src, T* dst,
373+
batch<int32_t, A> const& index,
374+
kernel::requires_arch<avx512vl_128>) noexcept
375+
{
376+
_mm_i32scatter_epi32(dst, index, src, sizeof(T));
377+
}
378+
379+
template <class A, class T,
380+
class = std::enable_if_t<std::is_same<uint64_t, T>::value || std::is_same<int64_t, T>::value>>
381+
XSIMD_INLINE void scatter(batch<T, A> const& src, T* dst,
382+
batch<int64_t, A> const& index,
383+
kernel::requires_arch<avx512vl_128>) noexcept
384+
{
385+
_mm_i64scatter_epi64(dst, index, src, sizeof(T));
386+
}
387+
388+
template <class A>
389+
XSIMD_INLINE void scatter(batch<float, A> const& src, float* dst,
390+
batch<int32_t, A> const& index,
391+
kernel::requires_arch<avx512vl_128>) noexcept
392+
{
393+
_mm_i32scatter_ps(dst, index, src, sizeof(float));
394+
}
395+
396+
template <class A>
397+
XSIMD_INLINE void scatter(batch<double, A> const& src, double* dst,
398+
batch<int64_t, A> const& index,
399+
kernel::requires_arch<avx512vl_128>) noexcept
400+
{
401+
_mm_i64scatter_pd(dst, index, src, sizeof(double));
402+
}
403+
338404
// max
339405
template <class A>
340406
XSIMD_INLINE batch<int64_t, A> max(batch<int64_t, A> const& self, batch<int64_t, A> const& other, requires_arch<avx512vl_128>) noexcept

‎include/xsimd/arch/xsimd_avx512vl_256.hpp‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,72 @@ namespace xsimd
334334
detail::maskstore256(mem, src, mask.mask(), Mode {});
335335
}
336336

337+
// gather
338+
template <class T, class A, class U, detail::enable_sized_integral_t<T, 4> = 0, detail::enable_sized_integral_t<U, 4> = 0>
339+
XSIMD_INLINE batch<T, A> gather(batch<T, A> const&, T const* src, batch<U, A> const& index,
340+
kernel::requires_arch<avx512vl_256>) noexcept
341+
{
342+
return _mm256_mmask_i32gather_epi32(_mm256_setzero_si256(), (__mmask8)0xFF, index, static_cast<const void*>(src), sizeof(T));
343+
}
344+
345+
template <class T, class A, class U, detail::enable_sized_integral_t<T, 8> = 0, detail::enable_sized_integral_t<U, 8> = 0>
346+
XSIMD_INLINE batch<T, A> gather(batch<T, A> const&, T const* src, batch<U, A> const& index,
347+
kernel::requires_arch<avx512vl_256>) noexcept
348+
{
349+
return _mm256_mmask_i64gather_epi64(_mm256_setzero_si256(), (__mmask8)0xF, index, static_cast<const void*>(src), sizeof(T));
350+
}
351+
352+
template <class A, class U, detail::enable_sized_integral_t<U, 4> = 0>
353+
XSIMD_INLINE batch<float, A> gather(batch<float, A> const&, float const* src,
354+
batch<U, A> const& index,
355+
kernel::requires_arch<avx512vl_256>) noexcept
356+
{
357+
return _mm256_mmask_i32gather_ps(_mm256_setzero_ps(), (__mmask8)0xFF, index, src, sizeof(float));
358+
}
359+
360+
template <class A, class U, detail::enable_sized_integral_t<U, 8> = 0>
361+
XSIMD_INLINE batch<double, A> gather(batch<double, A> const&, double const* src,
362+
batch<U, A> const& index,
363+
requires_arch<avx512vl_256>) noexcept
364+
{
365+
return _mm256_mmask_i64gather_pd(_mm256_setzero_pd(), (__mmask8)0xF, index, src, sizeof(double));
366+
}
367+
368+
// scatter
369+
template <class A, class T,
370+
class = std::enable_if_t<std::is_same<uint32_t, T>::value || std::is_same<int32_t, T>::value>>
371+
XSIMD_INLINE void scatter(batch<T, A> const& src, T* dst,
372+
batch<int32_t, A> const& index,
373+
kernel::requires_arch<avx512vl_256>) noexcept
374+
{
375+
_mm256_i32scatter_epi32(dst, index, src, sizeof(T));
376+
}
377+
378+
template <class A, class T,
379+
class = std::enable_if_t<std::is_same<uint64_t, T>::value || std::is_same<int64_t, T>::value>>
380+
XSIMD_INLINE void scatter(batch<T, A> const& src, T* dst,
381+
batch<int64_t, A> const& index,
382+
kernel::requires_arch<avx512vl_256>) noexcept
383+
{
384+
_mm256_i64scatter_epi64(dst, index, src, sizeof(T));
385+
}
386+
387+
template <class A>
388+
XSIMD_INLINE void scatter(batch<float, A> const& src, float* dst,
389+
batch<int32_t, A> const& index,
390+
kernel::requires_arch<avx512vl_256>) noexcept
391+
{
392+
_mm256_i32scatter_ps(dst, index, src, sizeof(float));
393+
}
394+
395+
template <class A>
396+
XSIMD_INLINE void scatter(batch<double, A> const& src, double* dst,
397+
batch<int64_t, A> const& index,
398+
kernel::requires_arch<avx512vl_256>) noexcept
399+
{
400+
_mm256_i64scatter_pd(dst, index, src, sizeof(double));
401+
}
402+
337403
// max
338404
template <class A>
339405
XSIMD_INLINE batch<int64_t, A> max(batch<int64_t, A> const& self, batch<int64_t, A> const& other, requires_arch<avx512vl_256>) noexcept

0 commit comments

Comments
 (0)