| 257 | } |
| 258 | |
| 259 | xla::StatusOr<se::OwningDeviceMemory> XRTMemoryManager::Allocate( |
| 260 | xla::Backend* backend, int device_ordinal, size_t size) { |
| 261 | se::DeviceMemoryAllocator* allocator = backend->memory_allocator(); |
| 262 | auto memory_or = |
| 263 | allocator->Allocate(device_ordinal, size, /*retry_on_failure=*/false); |
| 264 | if (memory_or.status().code() == error::RESOURCE_EXHAUSTED) { |
| 265 | VLOG(4) << "Allocate of " << size << " bytes failed on device " |
| 266 | << device_ordinal; |
| 267 | |
| 268 | DeviceContext* device_context = |
| 269 | GetDeviceContext(device_ordinal, |
| 270 | /*create_if_missing=*/false); |
| 271 | if (device_context != nullptr) { |
| 272 | Status status = device_context->TryFreeMemory(backend, size).status(); |
| 273 | if (status.ok()) { |
| 274 | // As long as there is no error, we still try again the allocation, even |
| 275 | // if the TryFreeMemory() call ended up freeing less memory than the |
| 276 | // required size. Fragmentation could make the memory allocation succeed |
| 277 | // even if the freed memory is indeed lower. |
| 278 | memory_or = allocator->Allocate(device_ordinal, size, |
| 279 | /*retry_on_failure=*/false); |
| 280 | } else if (status.code() != error::RESOURCE_EXHAUSTED) { |
| 281 | VLOG(4) << "Allocate of " << size << " bytes on device " |
| 282 | << device_ordinal << ": " << status; |
| 283 | return status; |
| 284 | } |
| 285 | } |
| 286 | } |
| 287 | return memory_or; |
| 288 | } |
| 289 | |
| 290 | string XRTMemoryManager::DebugString() const { |
| 291 | // We might want to emit more detailed information here, like per device |
nothing calls this directly
no test coverage detected