PERF revert openmp use in csr_row_norms (#26275)

jeremiedbb · web-flow · commit 092caed407f3 · 2023-04-27T09:48:21.000-04:00
diff --git a/sklearn/metrics/_pairwise_distances_reduction/_base.pyx.tp b/sklearn/metrics/_pairwise_distances_reduction/_base.pyx.tp
@@ -5,7 +5,7 @@ from libcpp.vector cimport vector
 
 from ...utils._cython_blas cimport _dot
 from ...utils._openmp_helpers cimport omp_get_thread_num
-from ...utils._typedefs cimport intp_t, float32_t, float64_t
+from ...utils._typedefs cimport intp_t, float32_t, float64_t, int32_t
 
 import numpy as np
 
@@ -14,7 +14,6 @@ from numbers import Integral
 from sklearn import get_config
 from sklearn.utils import check_scalar
 from ...utils._openmp_helpers import _openmp_effective_n_threads
-from ...utils.sparsefuncs_fast import _sqeuclidean_row_norms_sparse
 
 #####################
 
@@ -84,6 +83,23 @@ cdef float64_t[::1] _sqeuclidean_row_norms32_dense(
     return squared_row_norms
 
 
+cdef float64_t[::1] _sqeuclidean_row_norms64_sparse(
+    const float64_t[:] X_data,
+    const int32_t[:] X_indptr,
+    intp_t num_threads,
+):
+    cdef:
+        intp_t n = X_indptr.shape[0] - 1
+        int32_t X_i_ptr, idx = 0
+        float64_t[::1] squared_row_norms = np.zeros(n, dtype=np.float64)
+
+    for idx in prange(n, schedule='static', nogil=True, num_threads=num_threads):
+        for X_i_ptr in range(X_indptr[idx], X_indptr[idx+1]):
+            squared_row_norms[idx] += X_data[X_i_ptr] * X_data[X_i_ptr]
+
+    return squared_row_norms
+
+
 {{for name_suffix in ["64", "32"]}}
 
 from ._datasets_pair cimport DatasetsPair{{name_suffix}}
@@ -98,7 +114,7 @@ cpdef float64_t[::1] _sqeuclidean_row_norms{{name_suffix}}(
         # by moving squared row norms computations in MiddleTermComputer.
         X_data = np.asarray(X.data, dtype=np.float64)
         X_indptr = np.asarray(X.indptr, dtype=np.int32)
-        return _sqeuclidean_row_norms_sparse(X_data, X_indptr, num_threads)
+        return _sqeuclidean_row_norms64_sparse(X_data, X_indptr, num_threads)
     else:
         return _sqeuclidean_row_norms{{name_suffix}}_dense(X, num_threads)
 
diff --git a/sklearn/utils/sparsefuncs_fast.pyx b/sklearn/utils/sparsefuncs_fast.pyx
@@ -11,9 +11,6 @@ from libc.math cimport fabs, sqrt, isnan
 cimport numpy as cnp
 import numpy as np
 from cython cimport floating
-from cython.parallel cimport prange
-
-from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
 
 cnp.import_array()
 
@@ -28,14 +25,12 @@ def csr_row_norms(X):
     """Squared L2 norm of each row in CSR matrix X."""
     if X.dtype not in [np.float32, np.float64]:
         X = X.astype(np.float64)
-    n_threads = _openmp_effective_n_threads()
-    return _sqeuclidean_row_norms_sparse(X.data, X.indptr, n_threads)
+    return _sqeuclidean_row_norms_sparse(X.data, X.indptr)
 
 
 def _sqeuclidean_row_norms_sparse(
     const floating[::1] X_data,
     const integral[::1] X_indptr,
-    int n_threads,
 ):
     cdef:
         integral n_samples = X_indptr.shape[0] - 1
@@ -45,9 +40,10 @@ def _sqeuclidean_row_norms_sparse(
 
     cdef floating[::1] squared_row_norms = np.zeros(n_samples, dtype=dtype)
 
-    for i in prange(n_samples, schedule='static', nogil=True, num_threads=n_threads):
-        for j in range(X_indptr[i], X_indptr[i + 1]):
-            squared_row_norms[i] += X_data[j] * X_data[j]
+    with nogil:
+        for i in range(n_samples):
+            for j in range(X_indptr[i], X_indptr[i + 1]):
+                squared_row_norms[i] += X_data[j] * X_data[j]
 
     return np.asarray(squared_row_norms)
 
