Reverts 029d50d

PiperOrigin-RevId: 706770797
openxla · Dec 16, 2024 · 0954ab3 · 0954ab3
1 parent 68ebed6
commit 0954ab3
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 8 deletions.
diff --git a/xla/debug_options_flags.cc b/xla/debug_options_flags.cc
@@ -250,7 +250,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_enable_bf16_3way_gemm(false);
   opts.set_xla_gpu_nccl_collective_max_nchannels(0);
   opts.set_xla_gpu_nccl_p2p_max_nchannels(0);
-  opts.set_xla_gpu_multi_streamed_windowed_einsum(true);
+  opts.set_xla_gpu_multi_streamed_windowed_einsum(false);
 
   opts.set_xla_gpu_experimental_stream_annotation(false);
   // Minimum combined size of matrices in matrix multiplication to

diff --git a/xla/service/gpu/transforms/stream_attribute_annotator.cc b/xla/service/gpu/transforms/stream_attribute_annotator.cc
@@ -199,17 +199,15 @@ absl::StatusOr<bool> StreamAttributeAnnotator::Run(
                             AnnotateStreamAttributesForInstruction(
                                 instr, instr_gpu_config.value()));
         changed |= comp_result;
-      } else if (instr->opcode() == HloOpcode::kCopyStart &&
-                 module->has_schedule()) {
+      } else if (instr->opcode() == HloOpcode::kCopyStart) {
         TF_ASSIGN_OR_RETURN(bool comp_result,
                             AnnotateStreamAttributesForCopyStart(
                                 instr, channel_id, instr_gpu_config.value()));
         changed |= comp_result;
         continue;
       } else if (comp->IsAsyncComputation() &&
                  (instr->opcode() == HloOpcode::kDynamicSlice ||
-                  instr->opcode() == HloOpcode::kDynamicUpdateSlice) &&
-                 module->has_schedule()) {
+                  instr->opcode() == HloOpcode::kDynamicUpdateSlice)) {
         TF_ASSIGN_OR_RETURN(bool comp_result,
                             WrapIntoFusionAndAnnotateStreamAttributes(
                                 instr, channel_id, instr_gpu_config.value(),

diff --git a/xla/service/gpu/transforms/stream_attribute_annotator_test.cc b/xla/service/gpu/transforms/stream_attribute_annotator_test.cc
@@ -185,7 +185,7 @@ TEST_F(StreamAttributeAnnotatorTest, FusionIsAnnotated) {
 
 TEST_F(StreamAttributeAnnotatorTest, CopyStartIsAnnotated) {
   constexpr absl::string_view kHloString = R"(
-  HloModule offloading, is_scheduled=true
+  HloModule offloading
     ENTRY %main (param_0: f32[1024], param_1: f32[1024]) -> f32[1024] {
     %param_1 = f32[1024]{0} parameter(1)
     %param_0 = f32[1024]{0} parameter(0)
@@ -250,7 +250,7 @@ TEST_F(StreamAttributeAnnotatorTest, DynamicUpdateSliceWrappedAndAnnotated) {
 
   TF_ASSERT_OK_AND_ASSIGN(
       bool changed,
-      StreamAttributeAnnotator(device_description()).Run(module.get()));
+      StreamAttributeAnnotator{device_description()}.Run(module.get()));
   EXPECT_TRUE(changed);
 
   // Check that the dynamic-update-slice instruction is wrapped in a fusion
@@ -314,7 +314,7 @@ TEST_F(StreamAttributeAnnotatorTest, DynamicSliceWrappedAndAnnotated) {
   EXPECT_TRUE(module->has_schedule());
   TF_ASSERT_OK_AND_ASSIGN(
       bool changed,
-      StreamAttributeAnnotator(device_description()).Run(module.get()));
+      StreamAttributeAnnotator{device_description()}.Run(module.get()));
   EXPECT_TRUE(changed);
 
   // Check that the dynamic-slice instruction is wrapped in a fusion