Reverts 658f706

PiperOrigin-RevId: 644240260
tensorflow · Jun 18, 2024 · 223a3aa · 223a3aa
1 parent 42f555e
commit 223a3aa
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 15 deletions.
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
@@ -218,7 +218,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(false);
+  opts.set_xla_gpu_enable_cub_radix_sort(true);
   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,13 +46,7 @@ bool HloWasRewrittenToUseCubSort(const HloModule& module) {
 
 class CubSortKeysTest : public HloTestBase,
                         public ::testing::WithParamInterface<
-                            std::tuple<PrimitiveType, bool, int>> {
-  DebugOptions GetDebugOptionsForTest() override {
-    auto options = HloTestBase::GetDebugOptionsForTest();
-    options.set_xla_gpu_enable_cub_radix_sort(true);
-    return options;
-  }
-};
+                            std::tuple<PrimitiveType, bool, int>> {};
 
 TEST_P(CubSortKeysTest, CompareToReference) {
   int batch_size = std::get<2>(GetParam());
@@ -134,13 +128,7 @@ INSTANTIATE_TEST_SUITE_P(
 class CubSortPairsTest
     : public HloTestBase,
       public ::testing::WithParamInterface<
-          std::tuple<PrimitiveType, PrimitiveType, bool, int>> {
-  DebugOptions GetDebugOptionsForTest() override {
-    auto options = HloTestBase::GetDebugOptionsForTest();
-    options.set_xla_gpu_enable_cub_radix_sort(true);
-    return options;
-  }
-};
+          std::tuple<PrimitiveType, PrimitiveType, bool, int>> {};
 
 TEST_P(CubSortPairsTest, CompareToReference) {
   int batch_size = std::get<3>(GetParam());