Revert "Revert "Revert "Add AllocatingOutput extensions for the instructions which use HasTensorAllocation target."""

This reverts commit e4cc26b.

Author: georgepaw

tensorflow/compiler/plugin/poplar/driver/tools/custom_ops/hlo_extensions.cc

@@ -28,18 +28,16 @@ namespace xla {
 namespace poplarplugin {
 namespace {
 
-void RegisterAllocatingOutputExtensions(HloOpcode opcode) {
+void RegisterSharedExtensions(HloOpcode opcode) {
   auto allocating_output = [](const HloInstruction*) { return true; };
   RegisterHloInstructionExtension<AllocatingOutputExtension>(opcode,
                                                              allocating_output);
 }
-REGISTER_HLO_INST_EXTENSIONS(kConstant, RegisterAllocatingOutputExtensions);
-REGISTER_HLO_INST_EXTENSIONS(kInfeed, RegisterAllocatingOutputExtensions);
-REGISTER_HLO_INST_EXTENSIONS(kParameter, RegisterAllocatingOutputExtensions);
-REGISTER_HLO_INST_EXTENSIONS(kReduceWindow, RegisterAllocatingOutputExtensions);
-REGISTER_HLO_INST_EXTENSIONS(kRng, RegisterAllocatingOutputExtensions);
-REGISTER_HLO_INST_EXTENSIONS(kSelectAndScatter,
-                             RegisterAllocatingOutputExtensions);
+REGISTER_HLO_INST_EXTENSIONS(kConstant, RegisterSharedExtensions);
+REGISTER_HLO_INST_EXTENSIONS(kInfeed, RegisterSharedExtensions);
+REGISTER_HLO_INST_EXTENSIONS(kParameter, RegisterSharedExtensions);
+REGISTER_HLO_INST_EXTENSIONS(kReduceWindow, RegisterSharedExtensions);
+REGISTER_HLO_INST_EXTENSIONS(kRng, RegisterSharedExtensions);
 
 void RegisterReduceExtensions(HloOpcode opcode) {
   auto allocating_output = [](const HloInstruction*) { return true; };

tensorflow/compiler/plugin/poplar/driver/tools/custom_ops/inter_ipu_copy.cc

@@ -29,7 +29,7 @@ absl::flat_hash_set<int64_t> HloInterIpuCopy::AllocatingIndices() const {
   return {};
 }
 
-bool HloInterIpuCopy::AllocatingOutput() const { return true; }
+bool HloInterIpuCopy::AllocatingOutput() const { return false; }
 
 absl::flat_hash_map<int64_t, int64_t> HloInterIpuCopy::LayoutDependencies()
     const {

tensorflow/compiler/plugin/poplar/driver/tools/custom_ops/stateless_random.cc

@@ -31,7 +31,7 @@ absl::flat_hash_set<int64_t> HloStatelessRandom::AllocatingIndices() const {
   return {};
 }
 
-bool HloStatelessRandom::AllocatingOutput() const { return true; }
+bool HloStatelessRandom::AllocatingOutput() const { return false; }
 
 absl::flat_hash_map<int64_t, int64_t> HloStatelessRandom::LayoutDependencies()
     const {

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

@@ -68,7 +68,7 @@ def testConvolutionsDontMatchDifferentDevices(self):
     # Note how there are two convolutions
     ok = [
         '*OnTileCopy*', 'vs/conv2d/Conv2D/convolution.*',
-        'vs/conv2d_1/Conv2D/convolution'
+        'Copy_*vs/conv2d/Conv2D/convolution', 'vs/conv2d_1/Conv2D/convolution'
     ]
     self.assert_all_compute_sets_and_list(report, ok)
 

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

@@ -208,8 +208,8 @@ def my_graph(pa, pb, pc):
 
       for c in comps:
         tiles = tm.tile_ids(c)
-        if len(tiles) == 5:
-          self.assertEqual(tiles, set((0, 1, 1472, 1473, 1474)))
+        if len(tiles) == 3:
+          self.assertEqual(tiles, set((0, 1, 1472)))
         else:
           self.assertEqual(len(tiles), 0)