[XLA:SchedulingGroups][2/N] Do not combine all-reduces that have scheduling annotations.

PiperOrigin-RevId: 931253508

Author: seherellis

third_party/xla/xla/hlo/transforms/collectives/BUILD

@@ -370,6 +370,7 @@ cc_library(
     hdrs = ["all_reduce_combiner.h"],
     deps = [
         "//xla:shape_util",
+        "//xla:side_effect_util",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",

third_party/xla/xla/hlo/transforms/collectives/all_reduce_combiner.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/service/hlo_domain_map.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/side_effect_util.h"
 #include "xla/status_macros.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
@@ -169,7 +170,10 @@ absl::StatusOr<bool> AllReduceCombiner::RunWithKeyCombiner(
 
     auto key_fn = [&domain_map, &combine_key](const HloInstruction* instruction)
         -> std::optional<AllReduceCombiner::GroupKey> {
-      if (instruction->opcode() != HloOpcode::kAllReduce) {
+      if (instruction->opcode() != HloOpcode::kAllReduce ||
+          (instruction->has_frontend_attributes() &&
+           instruction->frontend_attributes().map().contains(
+               kXlaSchedulingGroupIdAttr))) {
         return std::nullopt;
       }
       return combine_key(instruction, *domain_map);

third_party/xla/xla/hlo/transforms/collectives/all_reduce_combiner_test.cc

@@ -542,5 +542,34 @@ TEST_F(AllReduceCombinerTest, PreservesMetadata) {
                         op::GetTupleElement(combined_all_reduce, 1)));
 }
 
+TEST_F(AllReduceCombinerTest, DoNotCombineOpWithSchedulingAnnotation) {
+  absl::string_view hlo_text = R"(
+    HloModule Module
+
+    %add (x: f32[], y: f32[]) -> f32[] {
+      %x = f32[] parameter(0)
+      %y = f32[] parameter(1)
+      ROOT %add = f32[] add(f32[] %x, f32[] %y)
+    }
+
+    ENTRY entry {
+      %param.0 = f32[32] parameter(0)
+      %param.1 = f32[32] parameter(1)
+      %all-reduce.0 = f32[32] all-reduce(%param.0), replica_groups={},
+          to_apply=%add, frontend_attributes={_scheduling_group_id="0"}
+      %all-reduce.1 = f32[32] all-reduce(%param.1), replica_groups={},
+          to_apply=%add
+      ROOT tuple = (f32[32], f32[32]) tuple(%all-reduce.0, %all-reduce.1)
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(hlo_text));
+  AllReduceCombiner combine(1024 * 1024, kMaxCombineCount);
+  EXPECT_THAT(combine.Run(module.get()), absl_testing::IsOkAndHolds(false));
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              op::Tuple(op::AllReduce(op::Parameter(0)),
+                        op::AllReduce(op::Parameter(1))));
+}
+
 }  // namespace
 }  // namespace xla