Fixing RTS for LAMB

Summary:
These changes are an update to the elementwise clustering algorithm to allow RTS to work with the LAMB optimiser.

RTS was not working with LAMB because the clusters in LAMB had intermediate scalar values which were not included.
In the old algorithm, all scalars were excluded by default to prevent otherwise unrelated clusters from being joined via hyperparameters.
There was a mechanism to allow certain scalar values into clusters:
  - Recursively search operands for all non-scalar values which the scalar depends on.
  - If all of the non-scalar values have the same shape as the top of the cluster, add the scalar to the cluster.
The problem with this was that in LAMB with BERT, there are scalars which depend on a large number of non-scalar values. The algorithm doesn't consider if the ops it's looking at are actually clusterable (replica-identical and supported op type) so it ends up checking lots of irrelevant ops and failing.

With the new algorithm, all ops in a given cluster are the same shape when the clusters are first created.
Then the merging step is used to allow for clusters with differently shaped intermediate values.
This guarantees that only clusterable ops get considered.
Since the clusters are already constructed, we can more robustely determine if one cluster contains intermediates for another.

Changes to the algorithm:
  - Renaming some functions to be less confusing.
  - Using `ReplicaIdenticalDataflowAnalysis` to identify more accurately which ops are replica identical. This allowed the the removal of code which precalculated which fusion computations are clusterable.
  - Moving all logic which determines if an op is clusterable into one function (IsClusterable).
  - Tightening the requirements for which ops can be added to a cluster during cluster creation. The shape of the op must match the shape of the top of the cluster. This means initially, all ops in a cluster have the same shape.
  - Changing the merging logic so that normally clusters will only be merged if their top elements have the same shape. The exception to this is when a cluster `b` is surrounded by another cluster `a` (`a` directly uses the outputs of `b` but `b`'s inputs are reachable from `a`), in which case they are merged even if the shapes don't match. This essentially allows clusters to have differently shaped intermediate values.

Test Plan:
There are existing tests (`resource_update_elementwise_clustering_test.cc` and `replicated_resource_update_elementwise_clustering_test.cc`).
I have added a test based on the issue which was preventing RTS working with LAMB. The test fails before the changes and passes after.

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

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

Subscribers: gauthamg

Maniphest Tasks: T65700

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

Author: Alfie-Edwards

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

@@ -219,7 +219,7 @@ ReplicatedResourceUpdateElementwiseClustering::AddClusterInput(
   // Lower the all reduce into the cluster if all its users will be in the
   // cluster too.
   const bool lower_all_reduce = IsGlobalAllReduce(cluster_input) &&
-                                cluster.AllUsersIn(cluster_input) &&
+                                cluster.ContainsAllUsersOf(cluster_input) &&
                                 status_or_collective_op.ok();
 
   if (lower_all_reduce) {

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

@@ -110,8 +110,7 @@ ResourceUpdateElementwiseClustering::CreateValidator(
 
 StatusOr<std::vector<ElementwiseCluster>>
 ResourceUpdateElementwiseClustering::GetClustersIn(
-    HloInstruction* const call,
-    const absl::flat_hash_set<const HloComputation*>& elementwise_comps) const {
+    HloInstruction* const call) const {
   CHECK(IsRepeatLoop(call) || IsPipelineOp(call));
   HloComputation* call_comp = call->to_apply();
   // Make sure that the root of the call op is a tuple instruction.
@@ -149,9 +148,9 @@ ResourceUpdateElementwiseClustering::GetClustersIn(
   }
 
   auto validator = CreateValidator(resource_update_comp);
-  TF_ASSIGN_OR_RETURN(std::vector<ElementwiseCluster> clusters,
-                      ElementwiseCluster::GetClustersIn(
-                          resource_update, elementwise_comps, *validator));
+  TF_ASSIGN_OR_RETURN(
+      std::vector<ElementwiseCluster> clusters,
+      ElementwiseCluster::GetClustersIn(resource_update, *validator));
   // Try to print some helpful warnings if things don't look right.
   TF_RETURN_IF_ERROR(
       ValidateResourceUpdateAndClusters(resource_update, clusters));
@@ -239,7 +238,7 @@ StatusOr<HloInstruction*> ResourceUpdateElementwiseClustering::AddClusterInput(
   // Lower the all reduce into the cluster if all its users will be in the
   // cluster too.
   const bool lower_all_reduce =
-      IsAllReduce(cluster_input) && cluster.AllUsersIn(cluster_input);
+      IsAllReduce(cluster_input) && cluster.ContainsAllUsersOf(cluster_input);
 
   if (lower_all_reduce) {
     HloInstruction* input = cluster_input->mutable_operand(0);
@@ -398,10 +397,9 @@ Status ResourceUpdateElementwiseClustering::UpdateClusterBackendConfig(
 }
 
 StatusOr<bool> ResourceUpdateElementwiseClustering::RewriteCall(
-    HloModule* module, HloInstruction* call,
-    const absl::flat_hash_set<const HloComputation*>& elementwise_comps) const {
+    HloModule* module, HloInstruction* call) const {
   TF_ASSIGN_OR_RETURN(std::vector<ElementwiseCluster> clusters,
-                      GetClustersIn(call, elementwise_comps));
+                      GetClustersIn(call));
 
   if (clusters.empty()) {
     VLOG(3) << "No clusters found.";
@@ -466,13 +464,9 @@ StatusOr<bool> ResourceUpdateElementwiseClustering::Run(HloModule* module) {
 
   TF_RETURN_IF_ERROR(RunDataflowAnalysis(module));
 
-  const absl::flat_hash_set<const HloComputation*> elementwise_comps =
-      ElementwiseCluster::GetElementwiseClusterableComputations(module);
-
   bool module_changed = false;
   for (auto call : to_optimize) {
-    TF_ASSIGN_OR_RETURN(auto changed,
-                        RewriteCall(module, call, elementwise_comps));
+    TF_ASSIGN_OR_RETURN(auto changed, RewriteCall(module, call));
     if (changed) {
       module_changed = true;
     }

tensorflow/compiler/plugin/poplar/driver/passes/resource_update_elementwise_clustering.h

@@ -55,9 +55,7 @@ class ResourceUpdateElementwiseClustering : public HloModulePass {
   // Get clusters inside of the call, where the call has to be a repeat loop or
   // a pipeline.
   StatusOr<std::vector<ElementwiseCluster>> GetClustersIn(
-      HloInstruction* const call,
-      const absl::flat_hash_set<const HloComputation*>& elementwise_comps)
-      const;
+      HloInstruction* const call) const;
 
   // Outline the provided cluster - returns the call instruction to the cluster.
   StatusOr<HloInstruction*> OutlineCluster(ElementwiseCluster& cluster) const;
@@ -103,9 +101,7 @@ class ResourceUpdateElementwiseClustering : public HloModulePass {
       const HloInstruction* ru, std::vector<ElementwiseCluster> clusters) const;
 
  private:
-  StatusOr<bool> RewriteCall(HloModule* module, HloInstruction* call,
-                             const absl::flat_hash_set<const HloComputation*>&
-                                 elementwise_comps) const;
+  StatusOr<bool> RewriteCall(HloModule* module, HloInstruction* call) const;
 
   StatusOr<HloInstruction*> AddClusterInputToOutlinedComputation(
       int64_t param_idx, const ElementwiseCluster& cluster,

tensorflow/compiler/plugin/poplar/driver/tools/BUILD

@@ -1015,6 +1015,7 @@ poplar_cc_library(
         ":matcher_predicates",
         ":offloading_util",
         ":pipeline_util",
+        ":replica_identical_dataflow_analysis",
         ":util",
         "//tensorflow/compiler/plugin/poplar/driver/tools/custom_ops:all_gather",
         "//tensorflow/compiler/plugin/poplar/driver/tools/custom_ops:reduce_scatter",