[XLA] Enable rewriting sparse gradient accumulation with reshapes.

ukoxyz · tensorflower-gardener · commit e2c9d554a077 · 2021-07-29T16:27:20.000-07:00
A change in AlgebraicSimplifier. Moving the reshapes can help enable the
existing rewrite rule for sparse gradient accumulation.

PiperOrigin-RevId: 387692854
Change-Id: I0ec3c8f0bfa1f685a1718309fffaa21ac80e5ecd
diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
@@ -4055,6 +4055,60 @@ Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) {
     }
   }
 
+  // Moves the reshape in reshape(dus(...), x, ...)) before dus so that it can
+  // enable other optimizations, e.g., merging with broadcast, and sparse update
+  // (add(x, dus(broadcast(0), y, ...)) -> dus(x, add(ds(x), y), ...)).
+  if (!options_.is_layout_sensitive()) {
+    bool trivial_reshape;
+    std::vector<int64> deleted_dims;
+    std::vector<int64> inserted_dims;
+
+    HloInstruction* dus;
+    HloInstruction* slice;
+    std::tie(trivial_reshape, deleted_dims, inserted_dims) =
+        reshape->ReshapeMerelyInsertsOrDeletes1SizedDimensions();
+    // 1-sized dimensions added and removed will be one sized in both the update
+    // slice and the dynamic-update-slice result.
+    if (trivial_reshape &&
+        Match(reshape->mutable_operand(0),
+              m::Op(&dus)
+                  .WithOpcode(HloOpcode::kDynamicUpdateSlice)
+                  .WithOperand(1, m::Op(&slice))) &&
+        !dus->has_sharding() && !dus->operand(0)->has_sharding()) {
+      auto new_operand =
+          computation_->AddInstruction(HloInstruction::CreateReshape(
+              reshape->shape(), dus->mutable_operand(0)));
+      std::vector<int64> new_slice_shape;
+      std::vector<HloInstruction*> new_dus_operands;
+      new_dus_operands.push_back(new_operand);
+      new_dus_operands.push_back(nullptr);
+      auto zero = MakeScalarLike(dus->mutable_operand(2), 0);
+      const Shape& old_slice_shape = dus->operand(1)->shape();
+      for (int64 i = 0; i <= old_slice_shape.rank(); ++i) {
+        if (absl::c_linear_search(deleted_dims, i)) {
+          continue;
+        }
+        if (absl::c_linear_search(inserted_dims, new_slice_shape.size())) {
+          new_slice_shape.push_back(1);
+          new_dus_operands.push_back(zero);
+        }
+        if (i < old_slice_shape.rank()) {
+          new_slice_shape.push_back(old_slice_shape.dimensions(i));
+          new_dus_operands.push_back(dus->mutable_operand(2 + i));
+        }
+      }
+      auto new_slice =
+          computation_->AddInstruction(HloInstruction::CreateReshape(
+              ShapeUtil::MakeShape(old_slice_shape.element_type(),
+                                   new_slice_shape),
+              slice));
+      new_dus_operands[1] = new_slice;
+      auto new_dus =
+          dus->CloneWithNewOperands(reshape->shape(), new_dus_operands);
+      return ReplaceWithNewInstruction(reshape, std::move(new_dus));
+    }
+  }
+
   // Make this a bitcast if possible.
   if (HloInstruction* bitcast_operand =
           BitcastingOperandOfReshapeOrCopyChain(reshape, options_)) {
diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc
@@ -5562,6 +5562,41 @@ ENTRY AddBroadcastZeroWithDynamicSlice {
   EXPECT_THAT(root->operand(1)->opcode(), HloOpcode::kPad);
 }
 
+TEST_F(AlgebraicSimplifierTest, AddDynamicUpdateSliceToAddSlice) {
+  const char* hlo_string = R"(
+HloModule AddDynamicUpdateSliceToAddSlice
+
+ENTRY AddDynamicUpdateSliceToAddSlice {
+  param0 = f32[1, 4,12,512] parameter(0)
+  constant = f32[] constant(0)
+  broadcast = f32[4,12,512] broadcast(constant), dimensions={}
+  param1 = f32[1,12,512] parameter(1)
+  param2 = s32[] parameter(2)
+  constant.1 = s32[] constant(0)
+  dynamic-update-slice = f32[4,12,512] dynamic-update-slice(
+    broadcast, param1, param2, constant.1, constant.1)
+  reshape = f32[1,4,12,512] reshape(dynamic-update-slice)
+  ROOT add = f32[1,4,12,512] add(param0, reshape)
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(hlo_string));
+  VLOG(2) << "Before rewrite reshape\n" << module->ToString();
+  AlgebraicSimplifier simplifier(default_options_);
+  ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie());
+  ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie());
+  VLOG(2) << "After rewrite to add slice\n" << module->ToString();
+  auto root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(
+      root,
+      GmockMatch(m::DynamicUpdateSlice(
+          m::Parameter(0),
+          m::Add(m::DynamicSlice(m::Parameter(0), m::Constant(),
+                                 m::Parameter(2), m::Constant(), m::Constant()),
+                 m::Reshape(m::Parameter(1))),
+          m::Constant(), m::Parameter(2), m::Constant(), m::Constant())));
+}
+
 TEST_F(AlgebraicSimplifierTest, ScalarMultiplyReduction) {
   const char* hlo_string = R"(
 HloModule ConstScalarMultiply
diff --git a/tensorflow/compiler/xla/service/pattern_matcher.h b/tensorflow/compiler/xla/service/pattern_matcher.h
@@ -2227,6 +2227,7 @@ XLA_VARIADIC_OP_PATTERN(AfterAll);
 XLA_VARIADIC_OP_PATTERN(Concatenate);
 XLA_VARIADIC_OP_PATTERN(Conditional);
 XLA_VARIADIC_OP_PATTERN(DynamicSlice)
+XLA_VARIADIC_OP_PATTERN(DynamicUpdateSlice)
 XLA_VARIADIC_OP_PATTERN(Fusion);
 XLA_VARIADIC_OP_PATTERN(Map)
 XLA_VARIADIC_OP_PATTERN(Reduce);
