Fix selection of inputs to clusters

Summary: Ref T40083

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

Reviewed By: vladimirm, hakons, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Maniphest Tasks: T40083

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

Author: georgepaw

tensorflow/compiler/plugin/poplar/driver/tools/elementwise_cluster.cc

@@ -130,19 +130,19 @@ ElementwiseClusterValidator::Inputs ElementwiseClusterValidator::GetValidInputs(
   Inputs valid_inputs;
   for (const HloInstruction* inst : comp->MakeInstructionPostOrder()) {
     bool valid_input = false;
-    if (IsParameter(inst) && parameter_filter(inst->parameter_number())) {
+    if (IsParameter(inst)) {
+      valid_input = parameter_filter(inst->parameter_number());
+    } else if (IsGlobalAllReduce(inst)) {
       valid_input = true;
+    } else if (IsBroadcast(inst)) {
+      valid_input =
+          IsScalar(inst->operand(0)) && valid_inputs.contains(inst->operand(0));
     } else if (!IsPoplarInstruction(PoplarOp::ExecutionCounter, inst) &&
                !inst->HasSideEffect()) {
       valid_input = absl::c_all_of(
           inst->operands(), [&valid_inputs](const HloInstruction* operand) {
             return valid_inputs.contains(operand);
           });
-    } else if (IsBroadcast(inst)) {
-      valid_input =
-          IsScalar(inst->operand(0)) && valid_inputs.contains(inst->operand(0));
-    } else if (IsGlobalAllReduce(inst)) {
-      valid_input = true;
     }
 
     if (valid_input) {

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

@@ -1666,26 +1666,8 @@ TEST_F(TestPartitionReplicationFactor, TestNonGlobalAllReduce) {
       ElementwiseCluster::GetElementwiseClusterableComputations(module.get());
   TF_ASSERT_OK_AND_ASSIGN(auto clusters,
                           pass.GetClustersIn(loop, elementwise_comps));
-  ASSERT_THAT(clusters.size(), 1);
-
-  const int64 shard_size = 128 / partition_replication_factor;
-  auto& cluster = *std::begin(clusters);
-  EXPECT_THAT(cluster.GetClusterSize(), 128);
-  EXPECT_THAT(cluster.GetAlignedClusterSize(), 128);
-  EXPECT_THAT(cluster.GetShardSize(), shard_size);
-
-  // Convert the cluster.
-  TF_ASSERT_OK(pass.OutlineCluster(cluster).status());
-  TF_ASSERT_OK_AND_ASSIGN(bool eliminated, HloDCE().Run(module.get()));
-
-  EXPECT_THAT(arg0->user_count(), 1);
-  HloInstruction* all_reduce = GetNextUser(arg0).ValueOrDie();
-
-  // The non-global all-reduce should be left alone.
-  EXPECT_THAT(all_reduce->opcode(), HloOpcode::kAllReduce);
-  EXPECT_THAT(all_reduce->replica_groups().size(), 1);
-  EXPECT_THAT(all_reduce->replica_groups()[0].replica_ids(),
-              ::testing::ElementsAre(0));
+  // No clusters because the all-reduce is not global.
+  ASSERT_THAT(clusters.size(), 0);
 }
 
 TEST_F(TestPartitionReplicationFactor, IgnoreImplicit2ScalarBroadcast) {