BLD: Add openmp flags to build x86-simd-sort

r-devulap · r-devulap · commit 7fd938b91d56 · 2025-04-30T21:16:34.000-07:00
Also adds a simple unit test to stress the openmp code paths
diff --git a/numpy/_core/meson.build b/numpy/_core/meson.build
@@ -128,6 +128,16 @@ if use_intel_sort and not fs.exists('src/npysort/x86-simd-sort/README.md')
   error('Missing the `x86-simd-sort` git submodule! Run `git submodule update --init` to fix this.')
 endif
 
+# Setup openmp flags for x86-simd-sort:
+omp_cflags = []
+omp = []
+if use_intel_sort and cpp.has_argument('-fopenmp')
+  omp = dependency('openmp', required : false)
+  if omp.found()
+    omp_cflags = ['-fopenmp', '-DXSS_USE_OPENMP']
+  endif
+endif
+
 if not fs.exists('src/common/pythoncapi-compat')
   error('Missing the `pythoncapi-compat` git submodule! ' +
         'Run `git submodule update --init` to fix this.')
@@ -867,14 +877,17 @@ foreach gen_mtargets : [
     ] : []
   ],
 ]
+
+
+
   mtargets = mod_features.multi_targets(
     gen_mtargets[0], multiarray_gen_headers + gen_mtargets[1],
     dispatch: gen_mtargets[2],
     # baseline: CPU_BASELINE, it doesn't provide baseline fallback
     prefix: 'NPY_',
     dependencies: [py_dep, np_core_dep],
     c_args: c_args_common + max_opt,
-    cpp_args: cpp_args_common + max_opt,
+    cpp_args: cpp_args_common + max_opt + omp_cflags,
     include_directories: [
       'include',
       'src/common',
@@ -1286,7 +1299,7 @@ py.extension_module('_multiarray_umath',
     'src/umath',
     'src/highway'
   ],
-  dependencies: [blas_dep],
+  dependencies: [blas_dep, omp],
   link_with: [
     npymath_lib,
     unique_hash_so,
diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
@@ -10292,6 +10292,12 @@ def test_argsort_int(N, dtype):
     arr[N - 1] = maxv
     assert_arg_sorted(arr, np.argsort(arr, kind='quick'))
 
+# Test large arrays that leverage openMP implementations from x86-simd-sort:
+@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
+def test_sort_largearrays(dtype):
+    N = 1000000
+    arr = np.random.rand(N)
+    assert_equal(np.sort(arr, kind='quick'), np.sort(arr, kind='heap'))
 
 @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
 def test_gh_22683():
