Reverts 25ee973

PiperOrigin-RevId: 839253089

Author: Google-ML-Automation

xla/service/gpu/build_defs.bzl

@@ -27,7 +27,6 @@ def get_cub_sort_kernel_types(name = ""):
         "u16_b32",
         "u16_b64",
         "u32_b16",
-        "s32_b32",
         "u32_b32",
         "u32_b64",
         "u64_b16",

xla/service/gpu/gpu_compiler.cc

@@ -709,9 +709,10 @@ absl::Status RunOptimizationPasses(
 
   // DynamicPadder creates a stable KeyValue sort for dynamic reshapes.
   pipeline.AddPass<DynamicPadder>(dynamic_padder_options);
-  // SortRewriter needs to run before StableSortExpander.
-  pipeline.AddPass<SortRewriter>(gpu_target_config.device_description,
-                                 gpu_target_config.platform_name);
+
+  // TODO(b/407909195): Add SortRewriter here once it supports S32 keys for
+  // KeyValueSort. It needs to run before StableSortExpander, otherwise we will
+  // not match the comparison computation.
 
   // Expand the sort op to support stable sorting if required.
   pipeline.AddPass<StableSortExpander>();

xla/service/gpu/gpu_compiler_test.cc

@@ -1710,9 +1710,8 @@ TEST_F(PassOrderTest,
        SortRewriterRunsBeforeStableSortExpanderAndComparisonExpander) {
   VerifyPassOrder(/*first_pass_regex=*/"sort-rewriter",
                   /*last_pass_regex=*/"stable-sort-expander");
-  VerifyPassRunsAtLeastOnceBefore(
-      /*first_pass_regex=*/"sort-rewriter",
-      /*other_pass_regex=*/"comparison-expander");
+  VerifyPassOrder(/*first_pass_regex=*/"sort-rewriter",
+                  /*last_pass_regex=*/"comparison-expander");
 }
 
 TEST_F(PassOrderTest,

xla/service/gpu/transforms/sort_rewriter_test.cc

@@ -190,33 +190,6 @@ ENTRY %main {
                                   m::GetTupleElement(m::CustomCall(), 1))));
 }
 
-// Sort a pair of S32 tensors, keys go first.
-TEST_F(SortRewriterTest, SortS32Pairs) {
-  constexpr char kHlo[] = R"(
-HloModule TestModule
-
-%compare {
-  %lhs_key = s32[] parameter(0)
-  %rhs_key = s32[] parameter(1)
-  %lhs_value = s32[] parameter(2)
-  %rhs_value = s32[] parameter(3)
-  ROOT %lt = pred[] compare(%lhs_key, %rhs_key), direction=LT
-}
-
-ENTRY %main {
-  %input_keys = s32[1000] parameter(0)
-  %input_values = s32[1000] parameter(1)
-  ROOT %sort = (s32[1000], s32[1000]) sort(%input_keys, %input_values),
-      dimensions={0}, is_stable=true, to_apply=%compare
-})";
-
-  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
-  EXPECT_TRUE(RunModuleAndPass(module.get()));
-  EXPECT_THAT(module->entry_computation()->root_instruction(),
-              GmockMatch(m::Tuple(m::GetTupleElement(m::CustomCall(), 0),
-                                  m::GetTupleElement(m::CustomCall(), 1))));
-}
-
 // Sort a pair of tensors, keys go last.
 TEST_F(SortRewriterTest, SortPairsSwapped) {
   constexpr char kHlo[] = R"(

xla/stream_executor/cuda/cub_sort_kernel_cuda.cc

@@ -186,9 +186,6 @@ XLA_CUB_DEFINE_SORT_PAIRS(u16_b64, uint16_t, uint64_t)
 #endif
 
 // Pairs with 32-bit key.
-#ifdef CUB_TYPE_S32_B32
-XLA_CUB_DEFINE_SORT_PAIRS(s32_b32, int32_t, uint32_t)
-#endif
 #ifdef CUB_TYPE_U32_B16
 XLA_CUB_DEFINE_SORT_PAIRS(u32_b16, uint32_t, uint16_t)
 #endif

xla/stream_executor/cuda/cub_sort_kernel_cuda_impl.cu.cc

@@ -198,9 +198,6 @@ XLA_CUB_DEFINE_SORT_PAIRS(uint16_t, uint64_t)
 #endif
 
 // Pairs with 32-bit key.
-#ifdef CUB_TYPE_S32_B32
-XLA_CUB_DEFINE_SORT_PAIRS(int32_t, uint32_t)
-#endif
 #ifdef CUB_TYPE_U32_B16
 XLA_CUB_DEFINE_SORT_PAIRS(uint32_t, uint16_t)
 #endif

xla/stream_executor/rocm/cub_sort_kernel_rocm.cu.cc

@@ -348,9 +348,6 @@ XLA_CUB_DEFINE_SORT_PAIRS(u16_b64, uint16_t, uint64_t)
 #endif
 
 // Pairs with 32-bit key.
-#ifdef CUB_TYPE_S32_B32
-XLA_CUB_DEFINE_SORT_PAIRS(s32_b32, int32_t, uint32_t)
-#endif
 #ifdef CUB_TYPE_U32_B16
 XLA_CUB_DEFINE_SORT_PAIRS(u32_b16, uint32_t, uint16_t)
 #endif