Specialize gather/scatter for xsimd in n-body

alpaka-group · Jan 1, 2024 · 1c2e499 · 1c2e499
1 parent abf44ef
commit 1c2e499
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 0 deletions.
diff --git a/examples/alpaka/nbody/nbody.cpp b/examples/alpaka/nbody/nbody.cpp
@@ -83,6 +83,22 @@ struct llama::SimdTraits<Batch, std::enable_if_t<xsimd::is_batch<Batch>::value>>
     {
         batch.store_unaligned(mem);
     }
+
+    static LLAMA_FORCE_INLINE auto indicesToReg(std::array<int, lanes> indices)
+    {
+        return [&]<std::size_t... Is>(std::index_sequence<Is...>)
+        { return xsimd::batch<int, typename Batch::arch_type>(indices[Is]...); }(std::make_index_sequence<lanes>{});
+    }
+
+    static LLAMA_FORCE_INLINE auto gather(const value_type* mem, std::array<int, lanes> indices) -> Batch
+    {
+        return Batch::gather(mem, indicesToReg(indices));
+    }
+
+    static LLAMA_FORCE_INLINE void scatter(Batch batch, value_type* mem, std::array<int, lanes> indices)
+    {
+        batch.scatter(mem, indicesToReg(indices));
+    }
 };
 
 template<typename T>

diff --git a/examples/nbody/nbody.cpp b/examples/nbody/nbody.cpp
@@ -66,6 +66,22 @@ struct llama::SimdTraits<Batch, std::enable_if_t<xsimd::is_batch<Batch>::value>>
     {
         batch.store_unaligned(mem);
     }
+
+    static LLAMA_FORCE_INLINE auto indicesToReg(std::array<int, lanes> indices)
+    {
+        return [&]<std::size_t... Is>(std::index_sequence<Is...>)
+        { return xsimd::batch<int, typename Batch::arch_type>(indices[Is]...); }(std::make_index_sequence<lanes>{});
+    }
+
+    static LLAMA_FORCE_INLINE auto gather(const value_type* mem, std::array<int, lanes> indices) -> Batch
+    {
+        return Batch::gather(mem, indicesToReg(indices));
+    }
+
+    static LLAMA_FORCE_INLINE void scatter(Batch batch, value_type* mem, std::array<int, lanes> indices)
+    {
+        batch.scatter(mem, indicesToReg(indices));
+    }
 };
 
 template<typename T>