Minor refactor

Micky774 · Micky774 · Aug 28, 2023 · Aug 28, 2023 · Aug 29, 2023 · Aug 29, 2023
commit bf6c2521534d771e4ae87c31793a752c222d9a57
diff --git a/setup.py b/setup.py
@@ -514,7 +514,7 @@ def configure_extension_modules():
                     "sources": [join(SIMD_DIRECTORY, "simd.cpp")],
                     "cflags": ["-std=c++14", "-mavx"],
                     "extra_link_args": ["-std=c++14"],
-                    "include_dirs": [HWY_INCLUDE_PATH],
+                    "include_dirs": [SIMD_DIRECTORY, HWY_INCLUDE_PATH],
                 },
             ),
         )

diff --git a/sklearn/metrics/_simd/_dist_optim.cpp b/sklearn/metrics/_simd/_dist_optim.cpp
@@ -8,6 +8,10 @@ this file due to cimporting _dist_metrics
 /* If built with SIMD support, include the compiled library code */
 #if WITH_SIMD == 1
 #include "simd.hpp"
+#include "hwy/highway.h"
+
+template<typename T>
+auto simd_manhattan_dist = manhattan::_simd_manhattan_dist;
 #else
 #include <cstddef>
 

diff --git a/sklearn/metrics/_simd/simd.cpp b/sklearn/metrics/_simd/simd.cpp
@@ -1,12 +1,79 @@
 #include "simd.hpp"
+#include <cmath>
 
-template float simd_manhattan_dist(
-    const float* x,
-    const float* y,
-    const size_t size
-);
-template double simd_manhattan_dist(
-    const double* x,
-    const double* y,
-    const size_t size
-);
+#undef HWY_TARGET_INCLUDE
+#define HWY_TARGET_INCLUDE "simd.cpp"
+#include "hwy/foreach_target.h"
+#include "hwy/highway.h"
+
+HWY_BEFORE_NAMESPACE();
+namespace manhattan {
+
+    namespace HWY_NAMESPACE {
+        namespace hn = hwy::HWY_NAMESPACE;
+
+
+        template <typename Type>
+        HWY_ATTR Type manhattan_dist(
+            const Type* x,
+            const Type* y,
+            const size_t size
+        ) {
+            const hn::ScalableTag<Type> d;
+            auto simd_sum_1 = hn::Zero(d);
+            auto simd_sum_2 = hn::Zero(d);
+
+            auto lane_size = hn::Lanes(d);
+            size_t loop_iter = lane_size * 2;
+            size_t vec_size = size - size % loop_iter;
+            size_t vec_remainder_size = size - size % lane_size;
+
+            for (size_t i = 0; i < vec_size; i += loop_iter) {
+                const auto simd_x_1 = hn::LoadU(d, x + i);
+                const auto simd_y_1 = hn::LoadU(d, y + i);
+                simd_sum_1 += hn::AbsDiff(simd_x_1, simd_y_1);
+
+                const auto simd_x_2 = hn::LoadU(d, x + i + lane_size);
+                const auto simd_y_2 = hn::LoadU(d, y + i + lane_size);
+                simd_sum_2 += hn::AbsDiff(simd_x_2, simd_y_2);
+            }
+            for (size_t i = vec_size; i < vec_remainder_size; i += loop_iter) {
+                const auto simd_x_1 = hn::LoadU(d, x + i);
+                const auto simd_y_1 = hn::LoadU(d, y + i);
+                simd_sum_1 += hn::AbsDiff(simd_x_1, simd_y_1);
+            }
+            simd_sum_1 += simd_sum_2;
+            Type scalar_sum = hn::ReduceSum(d, simd_sum_1);
+            for (size_t i = vec_remainder_size; i < size; i += 1) {
+                scalar_sum += fabs(x[i] - y[i]);
+            }
+            return scalar_sum;
+        }
+    auto manhattan_dist_float = manhattan_dist<float>;
+    auto manhattan_dist_double = manhattan_dist<double>;
+    }
+}
+HWY_AFTER_NAMESPACE();
+
+#if HWY_ONCE
+
+namespace manhattan {
+
+    HWY_EXPORT(manhattan_dist_float);
+    HWY_EXPORT(manhattan_dist_double);
+
+    template <typename Type>
+    HWY_DLLEXPORT Type _simd_manhattan_dist(
+        const Type* x,
+        const Type* y,
+        const size_t size
+    ){
+        if(std::is_same<Type, float>::value){
+            return HWY_DYNAMIC_DISPATCH(manhattan_dist_float)(x,  y, size);
+        }
+        else{
+            return HWY_DYNAMIC_DISPATCH(manhattan_dist_double)(x,  y, size);
+        }
+    }
+}
+#endif  // HWY_ONCE
diff --git a/sklearn/metrics/_simd/simd.hpp b/sklearn/metrics/_simd/simd.hpp
@@ -1,53 +1,13 @@
-#include <hwy/highway.h>
-#include <iostream>
-#include <cmath>
-namespace hn = hwy::HWY_NAMESPACE;
+#ifndef SIMD_HPP
+#define SIMD_HPP
+#include <hwy/base.h>
 
-template <typename Type>
-Type simd_manhattan_dist(
-    const Type* x,
-    const Type* y,
-    const size_t size
-) {
-    const hn::ScalableTag<Type> d;
-    auto simd_sum_1 = hn::Zero(d);
-    auto simd_sum_2 = hn::Zero(d);
+namespace manhattan{
 
-    auto lane_size = hn::Lanes(d);
-    size_t loop_iter = lane_size * 2;
-    size_t vec_size = size - size % loop_iter;
-    size_t vec_remainder_size = size - size % lane_size;
-
-    for (size_t i = 0; i < vec_size; i += loop_iter) {
-        const auto simd_x_1 = hn::LoadU(d, x + i);
-        const auto simd_y_1 = hn::LoadU(d, y + i);
-        simd_sum_1 += hn::AbsDiff(simd_x_1, simd_y_1);
-
-        const auto simd_x_2 = hn::LoadU(d, x + i + lane_size);
-        const auto simd_y_2 = hn::LoadU(d, y + i + lane_size);
-        simd_sum_2 += hn::AbsDiff(simd_x_2, simd_y_2);
-    }
-    for (size_t i = vec_size; i < vec_remainder_size; i += loop_iter) {
-        const auto simd_x_1 = hn::LoadU(d, x + i);
-        const auto simd_y_1 = hn::LoadU(d, y + i);
-        simd_sum_1 += hn::AbsDiff(simd_x_1, simd_y_1);
-    }
-    simd_sum_1 += simd_sum_2;
-    Type scalar_sum = hn::ReduceSum(d, simd_sum_1);
-    for (size_t i = vec_remainder_size; i < size; i += 1) {
-        scalar_sum += fabs(x[i] - y[i]);
-    }
-    return scalar_sum;
+    HWY_DLLEXPORT float _simd_manhattan_dist(
+        const float* x,
+        const float* y,
+        const size_t size
+    );
 }
-
-/*Extern declarations to later be mapped to simd.cpp declarations at link time*/
-extern template float simd_manhattan_dist(
-    const float* x,
-    const float* y,
-    const size_t size
-);
-extern template double simd_manhattan_dist(
-    const double* x,
-    const double* y,
-    const size_t size
-);
+#endif /*SIMD_HPP*/