Simplify multi update instruction

Summary:
TF2.4 Only

Ref T45278

Test Plan: CI

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, vladimirm, jakeh, georgew

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

Maniphest Tasks: T45278

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

Author: georgepaw

tensorflow/compiler/plugin/poplar/BUILD

@@ -226,7 +226,6 @@ cc_library(
         "driver/passes/multi_conv_fixer.cc",
         "driver/passes/multi_slice_combiner.cc",
         "driver/passes/multi_update_apply.cc",
-        "driver/passes/multi_update_canonicalize.cc",
         "driver/passes/multi_update_combiner.cc",
         "driver/passes/multi_update_scale_apply.cc",
         "driver/passes/multi_use_feeds_finder.cc",
@@ -435,7 +434,6 @@ cc_library(
         "driver/passes/multi_conv_fixer.h",
         "driver/passes/multi_slice_combiner.h",
         "driver/passes/multi_update_apply.h",
-        "driver/passes/multi_update_canonicalize.h",
         "driver/passes/multi_update_combiner.h",
         "driver/passes/multi_update_scale_apply.h",
         "driver/passes/multi_use_feeds_finder.h",
@@ -3121,22 +3119,6 @@ xla_test(
     ],
 )
 
-xla_test(
-    name = "multi_update_canonicalize_test",
-    size = "small",
-    srcs = ["tests/multi_update_canonicalize_test.cc"],
-    backends = ["poplar"],
-    copts = ["-fexceptions"],
-    deps = [
-        ":optimizers",
-        "//tensorflow/compiler/xla/service:hlo_matchers",
-        "//tensorflow/compiler/xla/service:hlo_query",
-        "//tensorflow/compiler/xla/service:shape_inference",
-        "//tensorflow/compiler/xla/tests:hlo_test_base",
-        "//tensorflow/compiler/xla/tests:xla_internal_test_main",
-    ],
-)
-
 xla_test(
     name = "resource_update_fixer_test",
     size = "small",
@@ -5848,7 +5830,6 @@ test_suite(
         "multi_run_test",
         "multi_slice_combiner_test",
         "multi_update_apply_test",
-        "multi_update_canonicalize_test",
         "multi_update_combiner_test",
         "multi_update_scale_apply_test",
         "multi_use_feeds_finder_test",

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

@@ -161,9 +161,6 @@ static Status ApplyTransformation(HloMatcherMatched& match,
 
   HloMultiUpdateInstruction* multi_update =
       Cast<HloMultiUpdateInstruction>(match.instruction_mapping[2]);
-  // This pass should have been run after MultiUpdateCanonicalize.
-  CHECK_EQ(multi_update->GetIndexVectorDimension(), 1);
-  CHECK_EQ(multi_update->GetUpdateSliceDimension(), 1);
 
   HloInstruction* broadcast = match.instruction_mapping[3];
   HloInstruction* indices = match.instruction_mapping[5];
@@ -199,14 +196,14 @@ static Status ApplyTransformation(HloMatcherMatched& match,
     HloInstruction* scale = match.instruction_mapping[7];
     output = comp->AddInstruction(CreateMultiUpdateAdd(
         broadcast->shape(),
-        {broadcast, reduced_indices, normalized_updates, scale}, 1, 1,
+        {broadcast, reduced_indices, normalized_updates, scale},
         serialization_factor));
   } else {
     // Create MultiUpdate.
     CHECK_EQ(match.pattern_idx, 1);
     output = comp->AddInstruction(CreateMultiUpdate(
-        broadcast->shape(), {broadcast, reduced_indices, normalized_updates}, 1,
-        1, serialization_factor));
+        broadcast->shape(), {broadcast, reduced_indices, normalized_updates},
+        serialization_factor));
   }
 
   // Replace with the new output.

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

@@ -29,9 +29,11 @@ limitations under the License.
 #include "tensorflow/compiler/plugin/poplar/driver/tools/matcher_predicates.h"
 #include "tensorflow/compiler/plugin/poplar/driver/tools/util.h"
 #include "tensorflow/compiler/plugin/poplar/kernels/ops.pb.h"
+
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/service/call_graph.h"
 #include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
+#include "tensorflow/compiler/xla/service/hlo_creation_utils.h"
 #include "tensorflow/compiler/xla/shape_util.h"
 
 namespace xla {
@@ -280,8 +282,6 @@ HloComputation* ReplaceResourceUpdateFunction(
       CreateMultiUpdateAdd(old_sink_shape,
                            {zero_broadcast.get(), new_indices_reshape.get(),
                             new_grads_arg.get(), const_1.get()},
-                           multi_update_add->GetIndexVectorDimension(),
-                           multi_update_add->GetUpdateSliceDimension(),
                            multi_update_add->GetSerializationFactor());
 
   for (auto old_sink_user : old_sink_arg->users()) {
@@ -355,8 +355,15 @@ StatusOr<HloComputation*> ReplaceAccumulatorCaller(
       ShapeUtil::MakeShape(update_index->shape().element_type(),
                            {plan.mini_batch_size}),
       update_index.get(), {});
-  auto update_index_broadcast_indices = HloInstruction::CreateBroadcast(
-      ShapeUtil::MakeShape(update_index->shape().element_type(), {1}),
+
+  auto update_index_broadcast_grads_reshaped = HloInstruction::CreateReshape(
+      ShapeUtil::MakeShape(update_index->shape().element_type(),
+                           /*dimensions=*/{plan.mini_batch_size, 1}),
+      update_index_broadcast_grads.get(), {});
+
+  auto update_index_broadcast_indices_reshaped = HloInstruction::CreateReshape(
+      ShapeUtil::MakeShape(update_index->shape().element_type(),
+                           /*dimensions=*/{1, 1}),
       update_index.get(), {});
 
   auto const_int_1 = HloInstruction::CreateConstant(LiteralUtil::One(S32));
@@ -369,15 +376,19 @@ StatusOr<HloComputation*> ReplaceAccumulatorCaller(
       plan.accum_grads_shape,
       {pipeline_stage_accum_grads_param ? pipeline_stage_accum_grads_param.get()
                                         : accum_grads.get(),
-       update_index_broadcast_grads.get(), grads, scale},
-      1, 1, multi_update_add->GetSerializationFactor());
-  auto indices_update = CreateMultiUpdateAdd(
-      plan.accum_indices_shape,
-      {pipeline_stage_accum_indices_param
-           ? pipeline_stage_accum_indices_param.get()
-           : accum_indices.get(),
-       update_index_broadcast_indices.get(), indices, const_int_1.get()},
-      0, 0, multi_update_add->GetSerializationFactor());
+       update_index_broadcast_grads_reshaped.get(), grads, scale},
+      multi_update_add->GetSerializationFactor());
+
+  TF_ASSIGN_OR_RETURN(auto indices_transpose,
+                      MakeTransposeHlo(indices, {1, 0}));
+  auto indices_update =
+      CreateMultiUpdateAdd(plan.accum_indices_shape,
+                           {pipeline_stage_accum_indices_param
+                                ? pipeline_stage_accum_indices_param.get()
+                                : accum_indices.get(),
+                            update_index_broadcast_indices_reshaped.get(),
+                            indices_transpose, const_int_1.get()},
+                           multi_update_add->GetSerializationFactor());
 
   std::unique_ptr<HloInstruction> grads_update_gte, indices_update_gte;
   if (pipeline_stage) {
@@ -525,6 +536,11 @@ absl::optional<Candidate> CheckEmbeddingsCandidate(HloInstruction* inst) {
                "found.";
     return absl::nullopt;
   }
+  // Updates need to be 2D.
+  if (inst->shape().rank() != 2) {
+    VLOG(2) << "The shape needs to be 2D.";
+    return absl::nullopt;
+  }
   return Candidate{Cast<HloGradientAccumulatorSink>(inst), {}};
 }
 
@@ -627,6 +643,11 @@ absl::optional<PipelineCandidate> CheckPipelineEmbeddingsCandidate(
     return absl::nullopt;
   }
 
+  if (grad_create->shape().rank() != 2) {
+    VLOG(2) << "The shape needs to be 2D.";
+    return absl::nullopt;
+  }
+
   PipelineCandidate candidate{grad_create,    grad_add,      multi_update_add,
                               grad_sink,      grad_sink_gte, pipeline_stage,
                               resource_update};

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

@@ -242,8 +242,6 @@ StatusOr<HloInstruction*> CreateNewMultiUpdate(
   // Create a new multi update add instruction.
   HloInstruction* new_multi_update_add = comp->AddInstruction(
       CreateMultiUpdateAdd(operand->shape(), {operand, indices, updates, scale},
-                           multi_update->GetIndexVectorDimension(),
-                           multi_update->GetUpdateSliceDimension(),
                            multi_update->GetSerializationFactor()));
 
   if (shard) {