Allocate input/output buffers at the correct memory space.

PiperOrigin-RevId: 317761146
Change-Id: Ia27b6765a547451575240e76ea419f705e29bd32

tensorflow/compiler/xla/service/interpreter/executor.cc

@@ -35,7 +35,6 @@ XlaInterpreterExecutor::~XlaInterpreterExecutor() {}
 
 DeviceMemoryBase XlaInterpreterExecutor::Allocate(uint64 size,
                                                   int64 memory_space) {
-  CHECK_EQ(memory_space, 0);
   return DeviceMemoryBase(new char[size], size);
 }
 

tensorflow/compiler/xla/service/transfer_manager.cc

@@ -355,7 +355,9 @@ StatusOr<ScopedShapedBuffer> TransferManager::AllocateScopedShapedBuffer(
         ShapeUtil::GetSubshape(shaped_buffer.on_device_shape(), index);
     TF_ASSIGN_OR_RETURN(auto memory,
                         allocator->Allocate(shaped_buffer.device_ordinal(),
-                                            GetByteSizeRequirement(subshape)));
+                                            GetByteSizeRequirement(subshape),
+                                            /*retry_on_failure=*/true,
+                                            subshape.layout().memory_space()));
     // Move the allocated buffer into the ScopedShapedBuffer, which owns it.
     memory_base = memory.Release();
   }