Fix bug #40

Fix bug #40 (closed): add support for integer version of GlobalAveragePooling and ReduceMean.

include/aidge/backend/cpu/operator/GlobalAveragePoolingImpl_kernels.hpp

@@ -25,6 +25,40 @@
 
 
 namespace Aidge {
+
+template <typename T>
+typename std::enable_if<std::is_floating_point<T>::value, T>::type
+stableMean(const T* vec, size_t size) {
+  T mean = 0;
+  for (size_t i = 0; i < size; ++i) {
+    mean = std::fma<T>(vec[i] - mean, 1.0f / (i + 1), mean);
+  }
+  return mean;
+}
+
+// Specialization for integers: perform the mean computation in float
+template <typename T>
+typename std::enable_if<!std::is_floating_point<T>::value, T>::type
+stableMean(const T* vec, size_t size) {
+  double mean = 0;
+  for (size_t i = 0; i < size; ++i) {
+    mean = std::fma<double>(vec[i] - mean, 1.0f / (i + 1), mean);
+  }
+  return mean;
+}
+
+template <typename T>
+typename std::enable_if<std::is_floating_point<T>::value, T>::type
+castFromFloat(T value) {
+  return value;
+}
+
+template <typename T>
+typename std::enable_if<!std::is_floating_point<T>::value, T>::type
+castFromFloat(double value) {
+  return static_cast<T>(std::nearbyint(value));
+}
+
 template <class I, class O>
 void GlobalAveragePoolingImpl_cpu_forward_kernel(
     const std::vector<DimSize_t> &dims, const void *input_, void *output_) {
@@ -49,12 +83,7 @@ void GlobalAveragePoolingImpl_cpu_forward_kernel(
     for (DimSize_t channel = 0; channel < dims[1]; ++channel) {
       const I *filter_start = std::next(
           input, (batch * in_batch_nb_elems) + (channel * in_channel_nb_elems));
-      I mean = 0;
-      for (size_t i = 0; i < in_channel_nb_elems; ++i) {
-        // Single pass numerically stable mean, using the fmaf
-        mean = fmaf(filter_start[i] - mean, 1.0f/(i+1), mean);
-      }
-      output[batch * out_batch_nb_elems + channel] = mean;
+      output[batch * out_batch_nb_elems + channel] = castFromFloat<O>(stableMean<I>(filter_start, in_channel_nb_elems));
     }
   }
 }