Reverts changelist 578813627

PiperOrigin-RevId: 643924789
tensorflow · Jun 17, 2024 · 658f706 · 658f706
1 parent 5cd6ae2
commit 658f706
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
@@ -217,7 +217,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_filter_kernels_spilling_registers_on_autotuning(true);
   opts.set_xla_gpu_llvm_verification_level(0);
   opts.set_xla_gpu_target_config_filename("");
-  opts.set_xla_gpu_enable_cub_radix_sort(true);
+  opts.set_xla_gpu_enable_cub_radix_sort(false);
   opts.set_xla_gpu_enable_cudnn_layer_norm(false);
   opts.set_xla_gpu_threshold_for_windowed_einsum_mib(100000);
 

diff --git a/third_party/xla/xla/service/gpu/tests/gpu_cub_sort_test.cc b/third_party/xla/xla/service/gpu/tests/gpu_cub_sort_test.cc
@@ -46,7 +46,13 @@ bool HloWasRewrittenToUseCubSort(const HloModule& module) {
 
 class CubSortKeysTest : public HloTestBase,
                         public ::testing::WithParamInterface<
-                            std::tuple<PrimitiveType, bool, int>> {};
+                            std::tuple<PrimitiveType, bool, int>> {
+  DebugOptions GetDebugOptionsForTest() override {
+    auto options = HloTestBase::GetDebugOptionsForTest();
+    options.set_xla_gpu_enable_cub_radix_sort(true);
+    return options;
+  }
+};
 
 TEST_P(CubSortKeysTest, CompareToReference) {
   int batch_size = std::get<2>(GetParam());
@@ -96,7 +102,13 @@ INSTANTIATE_TEST_SUITE_P(
 class CubSortPairsTest
     : public HloTestBase,
       public ::testing::WithParamInterface<
-          std::tuple<PrimitiveType, PrimitiveType, bool, int>> {};
+          std::tuple<PrimitiveType, PrimitiveType, bool, int>> {
+  DebugOptions GetDebugOptionsForTest() override {
+    auto options = HloTestBase::GetDebugOptionsForTest();
+    options.set_xla_gpu_enable_cub_radix_sort(true);
+    return options;
+  }
+};
 
 TEST_P(CubSortPairsTest, CompareToReference) {
   int batch_size = std::get<3>(GetParam());