| 156 | |
| 157 | template <typename T> |
| 158 | xla::StatusOr<T> XRTMemoryManager::Run( |
| 159 | const std::function<xla::StatusOr<T>()>& runfn, xla::Backend* backend, |
| 160 | int device_ordinal, size_t requested_free_size) { |
| 161 | MemoryReclaimContext mrctx(backend, device_ordinal, requested_free_size); |
| 162 | while (true) { |
| 163 | // We assume that runfn is a relatively fast-fail function compared to the |
| 164 | // operations required to free up the required memory. Here we call into the |
| 165 | // TryFreeMemoryStep() API multiple times, which will run progressively more |
| 166 | // expensive operations. |
| 167 | auto result_or = runfn(); |
| 168 | if (result_or.status().code() != error::RESOURCE_EXHAUSTED) { |
| 169 | return result_or; |
| 170 | } |
| 171 | TF_RETURN_IF_ERROR(TryFreeMemoryStep(&mrctx, result_or.status())); |
| 172 | } |
| 173 | } |
| 174 | |
| 175 | } // namespace tensorflow |
| 176 | |