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

This reverts commit 71c2eab.

Author: georgepaw

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

@@ -28,16 +28,18 @@ namespace xla {
 namespace poplarplugin {
 namespace {
 
-void RegisterSharedExtensions(HloOpcode opcode) {
+void RegisterAllocatingOutputExtensions(HloOpcode opcode) {
   auto allocating_output = [](const HloInstruction*) { return true; };
   RegisterHloInstructionExtension<AllocatingOutputExtension>(opcode,
                                                              allocating_output);
 }
-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);
+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);
 
 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 false; }
+bool HloInterIpuCopy::AllocatingOutput() const { return true; }
 
 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 false; }
+bool HloStatelessRandom::AllocatingOutput() const { return true; }
 
 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.*',
-        'Copy_*vs/conv2d/Conv2D/convolution', 'vs/conv2d_1/Conv2D/convolution'
+        'vs/conv2d_1/Conv2D/convolution'
     ]
     self.assert_all_compute_sets_and_list(report, ok)