Reverts 2e3fce3

PiperOrigin-RevId: 688856199
openxla · Oct 23, 2024 · 534189b · 534189b
1 parent 37635cf
commit 534189b
Show file tree

Hide file tree

Showing 6 changed files with 3 additions and 28 deletions.
diff --git a/xla/service/gpu/BUILD b/xla/service/gpu/BUILD
@@ -1547,7 +1547,6 @@ cc_library(
         "//xla/service/gpu/runtime:thunk",
         "//xla/service/gpu/transforms:algebraic_simplifier",
         "//xla/service/gpu/transforms:algorithm_checker",
-        "//xla/service/gpu/transforms:all_gather_dynamic_slice_simplifier",
         "//xla/service/gpu/transforms:all_gather_optimizer",
         "//xla/service/gpu/transforms:all_reduce_blueconnect",
         "//xla/service/gpu/transforms:all_reduce_splitter",
@@ -1693,7 +1692,6 @@ xla_test(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/ir:hlo_module_group",
-        "//xla/hlo/testlib:verified_hlo_module",
         "//xla/service:compiler",
         "//xla/service:executable",
         "//xla/service:hlo_module_config",

diff --git a/xla/service/gpu/gpu_compiler.cc b/xla/service/gpu/gpu_compiler.cc
@@ -176,7 +176,6 @@ limitations under the License.
 #include "xla/service/gpu/stream_executor_util.h"
 #include "xla/service/gpu/transforms/algebraic_simplifier.h"
 #include "xla/service/gpu/transforms/algorithm_checker.h"
-#include "xla/service/gpu/transforms/all_gather_dynamic_slice_simplifier.h"
 #include "xla/service/gpu/transforms/all_gather_optimizer.h"
 #include "xla/service/gpu/transforms/all_reduce_blueconnect.h"
 #include "xla/service/gpu/transforms/all_reduce_splitter.h"
@@ -861,7 +860,6 @@ absl::Status RunCollectiveOptimizationPasses(
   collectives_pipeline.AddPass<AllReduceFolder>();
   collectives_pipeline.AddPass<AllReduceSplitter>();
   collectives_pipeline.AddPass<AllGatherOptimizer>();
-  collectives_pipeline.AddPass<AllGatherDynamicSliceSimplifier>();
   collectives_pipeline.AddPass<AllReduceReassociate>(
       debug_options.xla_gpu_enable_reassociation_for_converted_ar());
   collectives_pipeline.AddPass<ReduceScatterReassociate>();

diff --git a/xla/service/gpu/gpu_compiler_test.cc b/xla/service/gpu/gpu_compiler_test.cc
@@ -41,7 +41,6 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_module_group.h"
 #include "xla/hlo/ir/hlo_opcode.h"
-#include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/primitive_util.h"
@@ -1304,16 +1303,6 @@ TEST_F(PassOrderTest, CollectivePipelinerRunsAfterCollectiveQuantizer) {
                   /*last_pass_regex=*/"collective-pipeliner.*");
 }
 
-TEST_F(PassOrderTest,
-       AllGatherDynamicSliceSimplifierRunsAfterAllGatherOptimizer) {
-  DebugOptions options = GetDebugOptionsForTest();
-  SetDebugOptions(options);
-
-  VerifyPassOrder(
-      /*first_pass_regex=*/".*all-gather-optimizer.*",
-      /*last_pass_regex=*/".*all-gather-dynamic-slice-simplifier.*");
-}
-
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/xla/service/gpu/transforms/BUILD b/xla/service/gpu/transforms/BUILD
@@ -342,11 +342,9 @@ cc_library(
     hdrs = ["all_gather_dynamic_slice_simplifier.h"],
     deps = [
         "//xla/hlo/ir:hlo",
-        "//xla/hlo/transforms:op_expander_pass",
         "//xla/service:collective_opt_utils",
-        "//xla/service:hlo_module_config",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:string_view",
+        "//xla/service:hlo_creation_utils",
+        "//xla/service:op_expander_pass",
     ],
 )
 

diff --git a/xla/service/gpu/transforms/all_gather_dynamic_slice_simplifier.cc b/xla/service/gpu/transforms/all_gather_dynamic_slice_simplifier.cc
@@ -15,13 +15,8 @@ limitations under the License.
 
 #include "xla/service/gpu/transforms/all_gather_dynamic_slice_simplifier.h"
 
-#include "absl/status/statusor.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
-#include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/ir/hlo_instructions.h"
-#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/collective_opt_utils.h"
-#include "xla/service/hlo_module_config.h"
 
 namespace xla {
 bool AllGatherDynamicSliceSimplifier::InstructionMatchesPattern(

diff --git a/xla/service/gpu/transforms/all_gather_dynamic_slice_simplifier.h b/xla/service/gpu/transforms/all_gather_dynamic_slice_simplifier.h
@@ -16,10 +16,7 @@ limitations under the License.
 #ifndef XLA_SERVICE_GPU_TRANSFORMS_ALL_GATHER_DYNAMIC_SLICE_SIMPLIFIER_H_
 #define XLA_SERVICE_GPU_TRANSFORMS_ALL_GATHER_DYNAMIC_SLICE_SIMPLIFIER_H_
 
-#include "absl/status/statusor.h"
-#include "absl/strings/string_view.h"
-#include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/transforms/expanders/op_expander_pass.h"
+#include "xla/service/op_expander_pass.h"
 
 namespace xla {