Simplify all-reduce(replication-normalise(all-reduce))

Summary:
This expression is effectively just all-reduce(). This simplification significantly reduces RTS memory footprint.

Fix T43594

Test Plan: CI, new test

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, jakeh, georgep, hakons

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep

Maniphest Tasks: T43594

Differential Revision: https://phabricator.sourcevertex.net/D49231

Author: whoozle

tensorflow/compiler/plugin/poplar/driver/passes/poplar_algebraic_simplifier.cc

@@ -171,6 +171,20 @@ StatusOr<HloInstruction*> PreserveFrontendAttributesIfNeeded(
   return new_inst;
 }
 
+bool IsGlobalAllReduceWithSum(const HloInstruction* all_reduce) {
+  if (all_reduce->opcode() != HloOpcode::kAllReduce ||
+      !all_reduce->replica_groups().empty()) {
+    return false;
+  }
+  auto& called_computations = all_reduce->called_computations();
+  if (called_computations.size() != 1) {
+    return false;
+  }
+  const HloComputation* comp = called_computations.front();
+  return Match(comp->root_instruction(),
+               m::Add(m::Parameter(), m::Parameter()));
+}
+
 // AlgebraicSimplifierVisitor traverses the HLO computation and reduces certain
 // algebraic expressions to simplified forms. Note: This only supports
 // simplifications that simply look at the operands of an instruction. For the
@@ -183,6 +197,8 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor {
 
   Status HandleAdd(HloInstruction* add) override;
 
+  Status HandleAllReduce(HloInstruction* all_reduce) override;
+
   Status HandleAnd(HloInstruction* logical_and) override;
 
   Status HandleBitcast(HloInstruction* bitcast) override;
@@ -3413,6 +3429,27 @@ Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice(
   return Status::OK();
 }
 
+Status AlgebraicSimplifierVisitor::HandleAllReduce(HloInstruction* all_reduce) {
+  /// Replace all-reduce(replication-normalise(all-reduce(arg))) with
+  /// all-reduce(arg)
+  if (all_reduce->operand_count() != 1 ||
+      !IsGlobalAllReduceWithSum(all_reduce)) {
+    return Status::OK();
+  }
+  HloInstruction* normalise = all_reduce->mutable_operand(0);
+  if (!pp::IsPoplarInstruction(PoplarOp::ReplicationNormalise, normalise) ||
+      normalise->operand_count() != 1) {
+    return Status::OK();
+  }
+  HloInstruction* top_all_reduce = normalise->mutable_operand(0);
+  if (top_all_reduce->opcode() == HloOpcode::kAllReduce &&
+      top_all_reduce->operand_count() == 1 &&
+      IsGlobalAllReduceWithSum(top_all_reduce)) {
+    return ReplaceInstruction(all_reduce, top_all_reduce);
+  }
+  return Status::OK();
+}
+
 Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* hlo) {
   HloReduceInstruction* reduce = Cast<HloReduceInstruction>(hlo);
   bool multi_output_reduce = reduce->shape().IsTuple();

tensorflow/compiler/plugin/poplar/tests/poplar_algebraic_simplifier_test.cc

@@ -5779,6 +5779,29 @@ ENTRY main {
   EXPECT_FALSE(PoplarAlgebraicSimplifier().Run(m.get()).ValueOrDie());
 }
 
+TEST_F(PoplarAlgebraicSimplifierTest, SimplifyAllReduceNormaliseAllReduce) {
+  const char* kModuleStr = R"(
+    HloModule m
+    sum {
+      y = f32[] parameter(1)
+      x = f32[] parameter(0), control-predecessors={y}
+      ROOT add = f32[] add(x, y), backend_config="{\"isInplace\":true}"
+    }
+
+    ENTRY main {
+      arg0 = f32[1000] parameter(0)
+      all-reduce0 = f32[1000] all-reduce(arg0), to_apply=sum
+      normalise = f32[1000] custom-call(all-reduce0), custom_call_target="ReplicationNormalise"
+      ROOT all-reduce1 = f32[1000] all-reduce(normalise), to_apply=sum
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  ASSERT_TRUE(CustomOpReplacer().Run(m.get()).ValueOrDie());
+  ASSERT_TRUE(PoplarAlgebraicSimplifier().Run(m.get()).ValueOrDie());
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::AllReduce(m::Parameter(0))));
+}
+
 }  // namespace
 }  // namespace poplarplugin
 }  // namespace xla