[XLA:ALGEBRAIC_SIMPLIFIER] Turn constant all-gather into broadcast

PiperOrigin-RevId: 698114685
openxla · Nov 19, 2024 · 3f032e1 · 3f032e1
1 parent 1edbb92
commit 3f032e1
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 0 deletions.
diff --git a/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -1221,6 +1221,18 @@ absl::StatusOr<bool> AlgebraicSimplifierVisitor::TrySimplifyTautologicalCompare(
   return false;
 }
 
+absl::Status AlgebraicSimplifierVisitor::HandleAllGather(
+    HloInstruction* all_gather) {
+  if (all_gather->shape().IsArray() &&
+      Match(all_gather->mutable_operand(0),
+            m::Broadcast(m::ConstantScalar()))) {
+    return ReplaceWithNewInstruction(
+        all_gather,
+        all_gather->mutable_operand(0)->CloneWithNewShape(all_gather->shape()));
+  }
+  return absl::OkStatus();
+}
+
 absl::Status AlgebraicSimplifierVisitor::HandleAllToAll(
     HloInstruction* all_to_all) {
   if (all_to_all->shape().IsArray() &&

diff --git a/xla/hlo/transforms/simplifiers/algebraic_simplifier.h b/xla/hlo/transforms/simplifiers/algebraic_simplifier.h
@@ -411,6 +411,8 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor {
 
   absl::Status HandleAdd(HloInstruction* add) override;
 
+  absl::Status HandleAllGather(HloInstruction* all_gather) override;
+
   absl::Status HandleAllToAll(HloInstruction* all_to_all) override;
 
   absl::Status HandleAnd(HloInstruction* logical_and) override;

diff --git a/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -12586,6 +12586,20 @@ TEST_F(AlgebraicSimplifierTest, BitcastBroadcastDifferentLayout) {
   EXPECT_FALSE(simplifier.Run(module.get()).value());
 }
 
+TEST_F(AlgebraicSimplifierTest, AllGatherOfBroadcast) {
+  const char* kModuleStr = R"(
+    HloModule m
+    test {
+      z = f32[] constant(0)
+      b = f32[4,4] broadcast(z), dimensions={}
+      ROOT ag = f32[16,4] all-gather(b), dimensions={0}, replica_groups={{0, 1, 2, 3}}
+    })";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::Broadcast(m::Constant())));
+}
+
 TEST_F(AlgebraicSimplifierTest, TrivialMin) {
   const char* kModuleStr = R"(
     HloModule m