Rebalance tensors for inplace ops

Summary:
This commit removes poputil::duplicate from FindInplaceOutputs, properly remapping tensors with aliasing
Depends on D62682.
Ref T56588.

Test Plan: CI

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, babakk

Maniphest Tasks: T56588

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

Author: whoozle

tensorflow/compiler/plugin/poplar/driver/tensor.cc

@@ -1665,9 +1665,10 @@ StatusOr<TensorOrRemoteBufferVectors> FindInplaceOutputs(
         if (require_parallel_writeable && !tensor.isParallelWriteable()) {
           poplar::Graph& graph =
               GetGraphWithOutputIndex(res, inst->operand(inplace_index), j);
-          tensors[i][j] = poputil::duplicate(
-              graph, tensor, seq, {debug_name_and_id, "clone"},
-              poplar::TensorCloneMethod::PRESERVE_ORDER_UNLESS_ALIASES);
+          poplar::Tensor out = TensorCloneAndRebalanceAliasing(
+              graph, res, tensor, {debug_name_and_id, "clone"});
+          seq.add(poplar::program::Copy(tensor, out, false, debug_name_and_id));
+          tensors[i][j] = out;
         }
       }
     }

tensorflow/compiler/plugin/poplar/tests/wide_const_expansion_test.py

@@ -156,7 +156,7 @@ def my_graph(a, b):
       self.assertAllClose(np.zeros(pb.shape), out[0])
 
     report = pva.openReport(report_helper.find_report())
-    self.assert_max_tile_memory(report, 534373)
+    self.assert_max_tile_memory(report, 417947)
 
   def testCheckMaxTileSizePadding2(self):
     cfg = IPUConfig()