| 191 | } |
| 192 | |
| 193 | xla::StatusOr<RefPtr<XRTTupleAllocation>> ExecuteComputation( |
| 194 | OpKernelContext* context, XRTMemoryManager* memory_manager, |
| 195 | XRTGenericDeviceAccessor::ScopedRef* device_ref, |
| 196 | xla::LocalExecutable* executable, const InputBuffers& input_buffers, |
| 197 | se::Stream* stream, int rng_seed) { |
| 198 | auto runfn = [&]() { |
| 199 | return RunExecutable(context, device_ref, executable, input_buffers, stream, |
| 200 | rng_seed); |
| 201 | }; |
| 202 | |
| 203 | // We pass zero as requested_free_size as there is no simple way to get the |
| 204 | // peak heap size. Upon zero, the Run() API will try to free chunks of device |
| 205 | // memory, until either the runfn can run, or we run out of freeable memory. |
| 206 | return memory_manager->Run<RefPtr<XRTTupleAllocation>>( |
| 207 | runfn, device_ref->backend(), device_ref->device_ordinal(), |
| 208 | /*requested_free_size=*/0); |
| 209 | } |
| 210 | |
| 211 | xla::StatusOr<RefPtr<XRTTupleAllocation>> ExecuteComputation( |
| 212 | OpKernelContext* context, const RefPtr<XRTMemoryManager>& memory_manager, |
no test coverage detected