Micky774 · Micky774 · Aug 28, 2023 · Aug 28, 2023 · Aug 29, 2023 · Aug 29, 2023
diff --git a/.gitignore b/.gitignore
@@ -105,3 +105,6 @@ sklearn/neighbors/_kd_tree.pyx
 
 # Default JupyterLite content
 jupyterlite_contents
+
+#SIMD
+!sklearn/metrics/_simd/*
diff --git a/.gitmodules b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "highway"]
+	path = highway
+	url = https://github.com/google/highway.git
diff --git a/highway b/highway
diff --git a/setup.py b/setup.py
@@ -11,6 +11,7 @@
 import sys
 import traceback
 from os.path import join
+from pathlib import Path
 
 from setuptools import Command, Extension, setup
 from setuptools.command.build_ext import build_ext
@@ -28,7 +29,6 @@
 # away from numpy.distutils?
 builtins.__SKLEARN_SETUP__ = True
 
-
 DISTNAME = "scikit-learn"
 DESCRIPTION = "A set of python modules for machine learning and data mining"
 with open("README.rst") as f:
@@ -43,13 +43,36 @@
     "Documentation": "https://scikit-learn.org/stable/documentation.html",
     "Source Code": "https://github.com/scikit-learn/scikit-learn",
 }
+SIMD_DIRECTORY = join("sklearn", "metrics", "_simd")
 
 # We can actually import a restricted version of sklearn that
 # does not need the compiled code
 import sklearn  # noqa
 import sklearn._min_dependencies as min_deps  # noqa
 from sklearn._build_utils import _check_cython_version  # noqa
 from sklearn.externals._packaging.version import parse as parse_version  # noqa
+from sklearn._build_utils.pre_build_helpers import compile_test_program  # noqa
+
+# runtime_check_program = dedent("""
+#     #include <highway.h>
+#     #include <iostream>
+#     int main(){
+#          std::cout << xsimd::available_architectures().avx;
+#          return 0;
+#     }
+#     """)
+# # XXX: Can we do this without an absolute path?
+dir_path = os.path.dirname(os.path.realpath(__file__))
+HWY_INCLUDE_PATH = Path(dir_path, "highway")
+HAS_AVX_RUNTIME = True
+# HAS_AVX_RUNTIME = int(
+#     compile_test_program(
+#         runtime_check_program,
+#         extra_preargs=["-lstdc++"],
+#         extra_postargs=[f"-I{HWY_INCLUDE_PATH}"],
+#         extension="cpp",
+#     )[0]
+# )
 
 
 VERSION = sklearn.__version__
@@ -252,8 +275,16 @@ def check_package_status(package, min_version):
     "metrics": [
         {"sources": ["_pairwise_fast.pyx"], "include_np": True},
         {
-            "sources": ["_dist_metrics.pyx.tp", "_dist_metrics.pxd.tp"],
+            "sources": [
+                "_dist_metrics.pyx.tp",
+                "_dist_metrics.pxd.tp",
+                join("_simd", "_dist_optim.cpp"),
+            ],
             "include_np": True,
+            "language": "c++",
+            "extra_compile_args": ["-std=c++11"],
+            "define_macros": [("DIST_METRICS", None)],
+            "include_dirs": [".", join("..", "..", HWY_INCLUDE_PATH)],
         },
     ],
     "metrics.cluster": [
@@ -471,6 +502,25 @@ def configure_extension_modules():
     build_with_debug_symbols = (
         os.environ.get("SKLEARN_BUILD_ENABLE_DEBUG_SYMBOLS", "0") != "0"
     )
+    BUILD_WITH_SIMD = int(
+        os.environ.get("SKLEARN_NO_SIMD", "0") == "0" and HAS_AVX_RUNTIME
+    )
+    if BUILD_WITH_SIMD:
+        libraries.append(
+            (
+                "avx_dist_metrics",
+                {
+                    "language": "c++",
+                    "sources": [join(SIMD_DIRECTORY, "simd.cpp")],
+                    "cflags": ["-std=c++14", "-mavx"],
+                    "extra_link_args": ["-std=c++14"],
+                    "include_dirs": [HWY_INCLUDE_PATH],
+                },
+            ),
+        )
+        extension_config["metrics"][1]["define_macros"].append(
+            ("WITH_SIMD", BUILD_WITH_SIMD)
+        )
     if os.name == "posix":
         if build_with_debug_symbols:
             default_extra_compile_args.append("-g")
@@ -538,13 +588,13 @@ def configure_extension_modules():
             optimization_level = extension.get(
                 "optimization_level", default_optimization_level
             )
+            define_macros = extension.get("define_macros", [])
             if os.name == "posix":
                 extra_compile_args.append(f"-{optimization_level}")
             else:
                 extra_compile_args.append(f"/{optimization_level}")
 
             libraries_ext = extension.get("libraries", []) + default_libraries
-
             new_ext = Extension(
                 name=name,
                 sources=sources,
@@ -554,6 +604,7 @@ def configure_extension_modules():
                 depends=depends,
                 extra_link_args=extension.get("extra_link_args", None),
                 extra_compile_args=extra_compile_args,
+                define_macros=define_macros,
             )
             cython_exts.append(new_ext)
 

diff --git a/sklearn/metrics/_dist_metrics.pxd.tp b/sklearn/metrics/_dist_metrics.pxd.tp
@@ -16,6 +16,11 @@ from ..utils._typedefs cimport float64_t, float32_t, int32_t, intp_t
 cdef class DistanceMetric:
     pass
 
+cdef extern from "_simd/_dist_optim.cpp":
+    cdef Type simd_manhattan_dist[Type](Type * x, Type * y, intp_t size) nogil
+    cdef Type simd_manhattan_dist_scalar[Type](Type * x, Type * y, intp_t size) nogil
+    int WITH_SIMD
+
 {{for name_suffix, INPUT_DTYPE_t, INPUT_DTYPE in implementation_specific_values}}
 
 ######################################################################
@@ -70,6 +75,7 @@ cdef class DistanceMetric{{name_suffix}}(DistanceMetric):
     cdef intp_t size
     cdef object func
     cdef object kwargs
+    cdef bint simd_runtime
 
     cdef {{INPUT_DTYPE_t}} dist(
         self,

diff --git a/sklearn/metrics/_dist_metrics.pyx.tp b/sklearn/metrics/_dist_metrics.pyx.tp
@@ -24,7 +24,6 @@ from scipy.sparse import csr_matrix, issparse
 from ..utils._typedefs cimport float64_t, float32_t, int32_t, intp_t
 from ..utils import check_array
 from ..utils.fixes import parse_version, sp_base_version
-
 cdef inline double fmax(double a, double b) noexcept nogil:
     return max(a, b)
 
@@ -837,7 +836,7 @@ cdef class DistanceMetric{{name_suffix}}(DistanceMetric):
 
             intp_t i1, i2
             intp_t x1_start, x1_end
-            {{INPUT_DTYPE_t}} * x2_data
+            const {{INPUT_DTYPE_t}} * x2_data
 
         with nogil:
             # Use the exact same adaptation for CSR than in SparseDenseDatasetsPair
@@ -901,7 +900,7 @@ cdef class DistanceMetric{{name_suffix}}(DistanceMetric):
             {{INPUT_DTYPE_t}}[:, ::1] Darr = np.empty((n_X, n_Y), dtype={{INPUT_DTYPE}}, order='C')
 
             intp_t i1, i2
-            {{INPUT_DTYPE_t}} * x1_data
+            const {{INPUT_DTYPE_t}} * x1_data
 
             intp_t x2_start, x2_end
 
@@ -1083,7 +1082,8 @@ cdef class EuclideanDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const intp_t size,
     ) except -1 nogil:
         return sqrt(
-            self.rdist_csr(
+            EuclideanDistance{{name_suffix}}.rdist_csr(
+            self,
             x1_data,
             x1_indices,
             x2_data,
@@ -1132,7 +1132,7 @@ cdef class SEuclideanDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const {{INPUT_DTYPE_t}}* x2,
         intp_t size,
     ) except -1 nogil:
-        return sqrt(self.rdist(x1, x2, size))
+        return sqrt(SEuclideanDistance{{name_suffix}}.rdist(self, x1, x2, size))
 
     cdef inline {{INPUT_DTYPE_t}} _rdist_to_dist(self, {{INPUT_DTYPE_t}} rdist) except -1 nogil:
         return sqrt(rdist)
@@ -1212,7 +1212,8 @@ cdef class SEuclideanDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const intp_t size,
     ) except -1 nogil:
         return sqrt(
-            self.rdist_csr(
+            SEuclideanDistance{{name_suffix}}.rdist_csr(
+            self,
             x1_data,
             x1_indices,
             x2_data,
@@ -1235,18 +1236,18 @@ cdef class ManhattanDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
     """
     def __init__(self):
         self.p = 1
+        self.simd_runtime = True
 
     cdef inline {{INPUT_DTYPE_t}} dist(
         self,
         const {{INPUT_DTYPE_t}}* x1,
         const {{INPUT_DTYPE_t}}* x2,
         intp_t size,
     ) except -1 nogil:
-        cdef float64_t d = 0
-        cdef intp_t j
-        for j in range(size):
-            d += fabs(x1[j] - x2[j])
-        return d
+        if self.simd_runtime:
+            return simd_manhattan_dist[{{INPUT_DTYPE_t}}](x1, x2, size)
+        else:
+            return simd_manhattan_dist_scalar[{{INPUT_DTYPE_t}}](x1, x2, size)
 
     cdef inline {{INPUT_DTYPE_t}} dist_csr(
         self,
@@ -1463,7 +1464,7 @@ cdef class MinkowskiDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const {{INPUT_DTYPE_t}}* x2,
         intp_t size,
     ) except -1 nogil:
-        return pow(self.rdist(x1, x2, size), 1. / self.p)
+        return pow(MinkowskiDistance{{name_suffix}}.rdist(self, x1, x2, size), 1. / self.p)
 
     cdef inline {{INPUT_DTYPE_t}} _rdist_to_dist(self, {{INPUT_DTYPE_t}} rdist) except -1 nogil:
         return pow(rdist, 1. / self.p)
@@ -1570,7 +1571,8 @@ cdef class MinkowskiDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const intp_t size,
     ) except -1 nogil:
         return pow(
-            self.rdist_csr(
+            MinkowskiDistance{{name_suffix}}.rdist_csr(
+                self,
                 x1_data,
                 x1_indices,
                 x2_data,
@@ -1655,7 +1657,7 @@ cdef class MahalanobisDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const {{INPUT_DTYPE_t}}* x2,
         intp_t size,
     ) except -1 nogil:
-        return sqrt(self.rdist(x1, x2, size))
+        return sqrt(MahalanobisDistance{{name_suffix}}.rdist(self, x1, x2, size))
 
     cdef inline {{INPUT_DTYPE_t}} _rdist_to_dist(self, {{INPUT_DTYPE_t}} rdist) except -1 nogil:
         return sqrt(rdist)
@@ -1736,7 +1738,8 @@ cdef class MahalanobisDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const intp_t size,
     ) except -1 nogil:
         return sqrt(
-            self.rdist_csr(
+            MahalanobisDistance{{name_suffix}}.rdist_csr(
+            self,
             x1_data,
             x1_indices,
             x2_data,
@@ -2641,7 +2644,7 @@ cdef class HaversineDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const {{INPUT_DTYPE_t}}* x2,
         intp_t size,
     ) except -1 nogil:
-        return 2 * asin(sqrt(self.rdist(x1, x2, size)))
+        return 2 * asin(sqrt(HaversineDistance{{name_suffix}}.rdist(self, x1, x2, size)))
 
     cdef inline {{INPUT_DTYPE_t}} _rdist_to_dist(self, {{INPUT_DTYPE_t}} rdist) except -1 nogil:
         return 2 * asin(sqrt(rdist))
@@ -2669,7 +2672,8 @@ cdef class HaversineDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const int32_t x2_end,
         const intp_t size,
     ) except -1 nogil:
-        return 2 * asin(sqrt(self.rdist_csr(
+        return 2 * asin(sqrt(HaversineDistance{{name_suffix}}.rdist_csr(
+            self,
             x1_data,
             x1_indices,
             x2_data,
@@ -2779,7 +2783,7 @@ cdef class PyFuncDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
         const {{INPUT_DTYPE_t}}* x2,
         intp_t size,
     ) except -1 nogil:
-        return self._dist(x1, x2, size)
+        return PyFuncDistance{{name_suffix}}._dist(self, x1, x2, size)
 
     cdef inline {{INPUT_DTYPE_t}} _dist(
         self,

diff --git a/sklearn/metrics/_simd/_dist_optim.cpp b/sklearn/metrics/_simd/_dist_optim.cpp
@@ -0,0 +1,43 @@
+/*
+Only provide a body upon inclusion if we are being included in _dist_metrics.pyx
+otherwise we remain empty so as to bypass Cython's forced inclusion of
+this file due to cimporting _dist_metrics
+*/
+#ifdef DIST_METRICS
+
+/* If built with SIMD support, include the compiled library code */
+#if WITH_SIMD == 1
+#include "simd.hpp"
+#else
+#include <cstddef>
+
+/* Else, we provide trivial functions for compilation */
+template <typename Type>
+Type simd_manhattan_dist(
+    const Type* x,
+    const Type* y,
+    const size_t size
+){return -1;}
+#endif
+
+/*
+In case of a runtime machine without AVX, we need to
+provide alternative scalar implementations.
+*/
+template <typename Type>
+Type simd_manhattan_dist_scalar(
+    const Type* x,
+    const Type* y,
+    const size_t size
+){
+    double scalar_sum = 0;
+
+    for(std::size_t idx = 0; idx < size; ++idx) {
+        scalar_sum += fabs(x[idx] - y[idx]);
+    }
+    return (Type) scalar_sum;
+}
+
+#else
+/* Empty body */
+#endif
diff --git a/sklearn/metrics/_simd/simd.cpp b/sklearn/metrics/_simd/simd.cpp
@@ -0,0 +1,12 @@
+#include "simd.hpp"
+
+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
+);