Do not constrain the layout of all-reduce

Summary:
We don't use it and it prevents DCE from removing such instructions/computations

TF2.4 Only.

Fix T43367.

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

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

Maniphest Tasks: T43367

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

Author: whoozle

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

@@ -288,7 +288,7 @@ ReplicatedResourceUpdateElementwiseClustering::AddClusterInput(
         builder->AddInstruction(HloInstruction::CreateAllReduce(
             scatter->shape(), std::vector<HloInstruction*>{scatter},
             sum_reduction, orthogonal_groups.ToXlaReplicaGroups(),
-            /*constrain_layout=*/true,
+            /*constrain_layout=*/false,
             /*channel_id=*/absl::nullopt, /*use_global_device_ids=*/false));
 
     context->MapInstruction(cluster_input, all_reduce);

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

@@ -204,7 +204,7 @@ struct PoplarReplicaGroupsHloTest : public HloTestBase {
     auto* all_reduce = builder.AddInstruction(HloInstruction::CreateAllReduce(
         shape, std::vector<HloInstruction*>{param}, reduction,
         groups.ToXlaReplicaGroups(),
-        /*constrain_layout=*/true,
+        /*constrain_layout=*/false,
         /*channel_id=*/absl::nullopt, /*use_global_device_ids=*/false));
     module->AddEntryComputation(builder.Build(all_reduce));