Parallelize sort for GCC build

annop-w · annop-w · commit 13eb689e0697 · 2025-03-31T12:39:17.000Z
Resolve #149977, #149979, #150094. Previously, #149505 used libstdc++ parallel mode by enabling -D_GLIBCXX_PARALLEL. However, mixing source files compiled with and without parallel mode can lead to undefined behavior (See https://gcc.gnu.org/onlinedocs/libstdc++/manual/parallel_mode_using.html) We switch to using the specific paralell sort from <parallel/algorithm> when compiled with GCC compiler. Note that use of std::execution policy has dependency on libtbb and we thus decide to avoid that.
diff --git a/aten/src/ATen/native/cpu/SortingKernel.cpp b/aten/src/ATen/native/cpu/SortingKernel.cpp
@@ -19,6 +19,9 @@
 #ifdef USE_FBGEMM
 #include <fbgemm/Utils.h>
 #endif
+#if __has_include(<parallel/algorithm>) && defined(_OPENMP)
+#include <parallel/algorithm>
+#endif
 
 namespace at::native {
 
@@ -146,6 +149,25 @@ static inline void sort_kernel_impl(const value_accessor_t& value_accessor,
   auto composite_accessor = CompositeRandomAccessorCPU<
     value_accessor_t, indices_accessor_t
   >(value_accessor, indices_accessor);
+#if __has_include(<parallel/algorithm>) && defined(_OPENMP)
+  if (descending) {
+    if (stable) {
+      __gnu_parallel::stable_sort(composite_accessor, composite_accessor + dim_size,
+        KeyValueCompDesc<scalar_t>());
+    } else {
+      __gnu_parallel::sort(composite_accessor, composite_accessor + dim_size,
+        KeyValueCompDesc<scalar_t>());
+    }
+  } else {
+    if (stable) {
+      __gnu_parallel::stable_sort(composite_accessor, composite_accessor + dim_size,
+        KeyValueCompAsc<scalar_t>());
+    } else {
+      __gnu_parallel::sort(composite_accessor, composite_accessor + dim_size,
+        KeyValueCompAsc<scalar_t>());
+    }
+  }
+#else
   if (descending) {
     if (stable) {
       std::stable_sort(composite_accessor, composite_accessor + dim_size,
@@ -163,6 +185,7 @@ static inline void sort_kernel_impl(const value_accessor_t& value_accessor,
         KeyValueCompAsc<scalar_t>());
     }
   }
+#endif
 }
 
 static void sort_kernel(
diff --git a/cmake/Codegen.cmake b/cmake/Codegen.cmake
@@ -416,11 +416,6 @@ if(INTERN_BUILD_ATEN_OPS)
         set(EXTRA_FLAGS "-DCPU_CAPABILITY=${CPU_CAPABILITY} -DCPU_CAPABILITY_${CPU_CAPABILITY}")
       endif(MSVC)
 
-      # Only parallelize the SortingKernel for now to avoid side effects
-      if(${NAME} STREQUAL "native/cpu/SortingKernel.cpp" AND NOT MSVC AND USE_OMP)
-        string(APPEND EXTRA_FLAGS " -D_GLIBCXX_PARALLEL")
-      endif()
-
       # Disable certain warnings for GCC-9.X
       if(CMAKE_COMPILER_IS_GNUCXX)
         if(("${NAME}" STREQUAL "native/cpu/GridSamplerKernel.cpp") AND ("${CPU_CAPABILITY}" STREQUAL "DEFAULT"))
