Add manual reduce

Author: davschneller

algorithms/sycl/Reduction.cpp

@@ -7,8 +7,106 @@
 #include "utils/logger.h"
 #include <device.h>
 
+#include <limits>
 #include <sycl/sycl.hpp>
 
+#if 1
+
+namespace {
+  using namespace device;
+
+  template<ReductionType Type, typename T>
+  constexpr T neutral() {
+    if constexpr(Type == ReductionType::Add) {
+      return T(0);
+    }
+    if constexpr(Type == ReductionType::Max) {
+      return std::numeric_limits<T>::min();
+    }
+    if constexpr(Type == ReductionType::Min) {
+      return std::numeric_limits<T>::max();
+    }
+    return T(0);
+  }
+
+  template <ReductionType Type, typename AccT, typename VecT, typename OpT> void launchReduction(AccT* result, const VecT *buffer, size_t size, OpT operation, bool overrideResult, void* streamPtr) {
+
+    constexpr auto DefaultValue = neutral<Type, AccT>();
+
+    ((sycl::queue *) streamPtr)->submit([&](sycl::handler &cgh) {
+      sycl::local_accessor<AccT, 1> shmem(256, cgh);
+      cgh.parallel_for(sycl::nd_range<1> { 1024, 1024 },
+        [=](sycl::nd_item<1> idx) {
+          const auto subgroup = idx.get_sub_group();
+          const auto sgSize = subgroup.get_local_range();
+
+          const auto warpCount = subgroup.get_group_range();
+          const auto currentWarp = subgroup.get_group_id();
+          const auto threadInWarp = subgroup.get_local_id();
+          const auto warpsNeeded = (size + sgSize - 1) / sgSize;
+
+          auto acc = DefaultValue;
+          
+          #pragma unroll 4
+          for (std::size_t i = currentWarp; i < warpsNeeded; i += warpCount) {
+            const auto id = threadInWarp + i * sgSize;
+            auto value = (id < size) ? static_cast<AccT>(ntload(&buffer[id])) : DefaultValue;
+
+            value = sycl::reduce_over_group(subgroup, value, operation);
+
+            acc = operation(acc, value);
+          }
+
+          if (threadInWarp == 0) {
+            shmem[currentWarp] = acc;
+          }
+
+          idx.barrier();
+
+          if (currentWarp == 0) {
+            const auto lastWarpsNeeded = (warpCount + sgSize - 1) / sgSize;
+            auto lastAcc = DefaultValue;
+            #pragma unroll 2
+            for (int i = 0; i < lastWarpsNeeded; ++i) {
+              const auto id = threadInWarp + i * sgSize;
+              auto value = (id < warpCount) ? shmem[id] : DefaultValue;
+
+              value = sycl::reduce_over_group(subgroup, value, operation);
+
+              lastAcc = operation(lastAcc, value);
+            }
+
+            if (threadInWarp == 0) {
+              if (overrideResult) {
+                ntstore(result, lastAcc);
+              }
+              else {
+                ntstore(result, operation(ntload(result), lastAcc));
+              }
+            }
+          }
+        });
+    });
+  }
+}
+
+namespace device {
+template <typename AccT, typename VecT> void Algorithms::reduceVector(AccT* result, const VecT *buffer, bool overrideResult, size_t size, ReductionType type, void* streamPtr) {
+  switch (type) {
+    case ReductionType::Add: {
+      return launchReduction<ReductionType::Add>(result, buffer, size, sycl::plus<AccT>(), overrideResult, streamPtr);
+    }
+    case ReductionType::Max: {
+      return launchReduction<ReductionType::Max>(result, buffer, size, sycl::maximum<AccT>(), overrideResult, streamPtr);
+    }
+    case ReductionType::Min: {
+      return launchReduction<ReductionType::Min>(result, buffer, size, sycl::minimum<AccT>(), overrideResult, streamPtr);
+    }
+  }
+}
+
+#else
+
 namespace {
   template <typename AccT, typename VecT, typename S> void launchReduction(AccT* result, const VecT *buffer, size_t size, S reducer, void* streamPtr) {
     ((sycl::queue *) streamPtr)->submit([&](sycl::handler &cgh) {
@@ -43,6 +141,8 @@ template <typename AccT, typename VecT> void Algorithms::reduceVector(AccT* resu
   }
 }
 
+#endif
+
 template void Algorithms::reduceVector(int* result, const int *buffer, bool overrideResult, size_t size, ReductionType type, void* streamPtr);
 template void Algorithms::reduceVector(unsigned* result, const unsigned *buffer, bool overrideResult, size_t size, ReductionType type, void* streamPtr);
 template void Algorithms::reduceVector(long* result, const int *buffer, bool overrideResult, size_t size, ReductionType type, void* streamPtr);