Run control dep and DCE on repeat

Summary:
Run them until a fixpoint to remove as many instructions as possible.
Also remove control deps from parameters.

Ref T48127

Test Plan: CI, added new tests

Reviewers: vladimirm, markf, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

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

Maniphest Tasks: T48127

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

Author: georgepaw

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

@@ -35,7 +35,7 @@ StatusOr<bool> DeadControlDependenciesElimination::Run(HloModule* module) {
 
     for (HloInstruction* inst : comp->MakeInstructionPostOrder()) {
       if (inst->user_count() == 0 && comp->root_instruction() != inst &&
-          inst->opcode() != HloOpcode::kParameter && !inst->HasSideEffect()) {
+          !inst->HasSideEffect()) {
         VLOG(2) << "Dropping control dependencies for " << inst->ToString();
         // To preserve topological constraints, make sure that control
         // predecessors are copied to control successors.
@@ -44,9 +44,13 @@ StatusOr<bool> DeadControlDependenciesElimination::Run(HloModule* module) {
             TF_RETURN_IF_ERROR(predecessor->AddControlDependencyTo(successor));
           }
         }
+        changed = inst->control_successors().size() ||
+                  inst->control_predecessors().size();
+
         TF_RETURN_IF_ERROR(inst->DropAllControlDeps());
-        CHECK(comp->IsSafelyRemovable(inst));
-        changed = true;
+        if (inst->opcode() != HloOpcode::kParameter) {
+          CHECK(comp->IsSafelyRemovable(inst));
+        }
       }
     }
   }

tensorflow/compiler/plugin/poplar/driver/poplar_compiler.cc

@@ -1574,8 +1574,13 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
       if (resources.recomputation_enabled) {
         pipeline.AddPass<FlattenCallGraph>();
       }
-      pipeline.AddPass<DeadControlDependenciesElimination>();
-      pipeline.AddPass<HloDCE>();
+      {
+        auto& dce_pass = pipeline.AddPass<HloPassFix<HloPassPipeline>>(
+            "dead-code-and-control-deps-elimination",
+            pipeline_compiler_stats.get());
+        dce_pass.AddPass<DeadControlDependenciesElimination>();
+        dce_pass.AddPass<HloDCE>();
+      }
       // Beyond this point non of the passes in the pipeline are allowed to
       // modify the instructions in the HloModule.
 

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

@@ -85,6 +85,28 @@ ENTRY c1 {
   EXPECT_THAT(log->control_successors(), ::testing::ElementsAre(add));
 }
 
+TEST_F(DeadControlDependenciesEliminationTest, TestParameters) {
+  std::string hlo = R"(
+HloModule top
+
+ENTRY c1 {
+  arg0 = f32[] parameter(0)
+  arg1 = f32[] parameter(1)
+  ROOT log = f32[] log(arg0)
+  arg2 = f32[] parameter(2), control-predecessors={log, arg1}
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto module, ParseAndReturnVerifiedModule(hlo, GetModuleConfigForTest()));
+  EXPECT_TRUE(
+      DeadControlDependenciesElimination().Run(module.get()).ValueOrDie());
+
+  HloInstruction* arg2 = FindInstruction(module.get(), "arg2");
+
+  EXPECT_TRUE(arg2->control_predecessors().empty());
+  EXPECT_TRUE(arg2->control_successors().empty());
+}
+
 }  // namespace
 }  // namespace poplarplugin
 }  // namespace xla