Reallocate tensors used for device to host copy

Summary:
This commit fixes a regression introduced in D55136.
If broadcast was passed to device to host copy, it ends on the single tile, rebalance output tensors if they have an aliases.

Ref T51875.
Ref T56588.

Test Plan: CI, regressed yolo3 model

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep, samuelh, babakk

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

Maniphest Tasks: T56588, T51875

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

Author: whoozle

tensorflow/compiler/plugin/poplar/driver/visitors/entry_visitor.cc

@@ -287,6 +287,14 @@ Status EntryVisitor::FinishDeferedAllocationVisit(HloInstruction* root) {
            ++tuple_index) {
         poplar::Tensor out = ConvertFromDeviceLayout(
             layout_sub_shapes[tuple_index], out_tensors[tuple_index]);
+        if (!out.isParallelWriteable()) {
+          poplar::Tensor out_copy = TensorCloneAndRebalanceAliasing(
+              graph, resources_, out, debug_name_and_id);
+
+          seq.add(
+              poplar::program::Copy(out, out_copy, false, debug_name_and_id));
+          out = out_copy;
+        }
 
         const std::string handle = out_info.Handles().at(tuple_index);
         auto fifo = graph.addDeviceToHostFIFO(handle, out.elementType(),

tensorflow/compiler/plugin/poplar/tests/size_speed_tests/matmul_size_test.py

@@ -94,6 +94,36 @@ def testInference(self):
     report = pva.openReport(report_helper.find_report())
     self.assert_total_tile_memory(report, 12181182)
 
+  def testInferenceManyLayersRebalancedBroadcastInInitialisation(self):
+    cfg = ipu.utils.IPUConfig()
+    report_helper = tu.ReportHelper()
+    report_helper.set_autoreport_options(cfg)
+    cfg.ipu_model.compile_ipu_code = True
+    cfg.ipu_model.tiles_per_ipu = 1472
+    cfg.optimizations.math.dot_strength = False
+    cfg.configure_ipu_system()
+
+    with self.session() as sess:
+      x = array_ops.placeholder(datatype, shape=[2, 112 * 112 * 4])
+      y = array_ops.placeholder(datatype, shape=[2, 4])
+
+      with ipu.scopes.ipu_scope("/device:IPU:0"):
+        logits = x
+        for idx in range(20):
+          # Any value less than tiles per ipu would work here,
+          # see ShouldRebalanceTensor().
+          logits = layer(logits, 1400, "l%d" % idx)
+        logits = layer(logits, 4, "out")
+
+        math_ops.reduce_mean(
+            nn_ops.softmax_cross_entropy_with_logits_v2(logits=logits,
+                                                        labels=y))
+
+      sess.run(variables.global_variables_initializer())
+      report = pva.openReport(report_helper.find_report())
+      self.assert_max_tile_memory(report, 162410)
+    # Skip the rest of the test, broadcast are returned for global initialisers.
+
   def testTrainingBs1(self):
     cfg = ipu.utils.IPUConfig()
     report_helper = tu.ReportHelper()