| 199 | } |
| 200 | |
| 201 | StatusOr<ScopedShapedBuffer> HloRunner::ExecuteWithDeviceBuffers( |
| 202 | std::unique_ptr<HloModule> module, |
| 203 | absl::Span<const ShapedBuffer* const> arguments, bool run_hlo_passes, |
| 204 | ExecutionProfile* profile) { |
| 205 | // Get service run options. |
| 206 | se::Stream stream(backend().default_stream_executor()); |
| 207 | stream.Init(); |
| 208 | ServiceExecutableRunOptions service_run_options = |
| 209 | GetServiceRunOptionsForDevice(backend().default_device_ordinal(), &stream, |
| 210 | nullptr, RunId()); |
| 211 | service_run_options.mutable_run_options()->set_execution_profile(profile); |
| 212 | |
| 213 | TF_ASSIGN_OR_RETURN(std::unique_ptr<Executable> executable, |
| 214 | CreateExecutable(std::move(module), run_hlo_passes)); |
| 215 | TF_ASSIGN_OR_RETURN( |
| 216 | ScopedShapedBuffer retval, |
| 217 | executable->ExecuteOnStreamWrapper(&service_run_options, arguments)); |
| 218 | TF_RETURN_IF_ERROR(stream.BlockHostUntilDone()); |
| 219 | return std::move(retval); |
| 220 | } |
| 221 | |
| 222 | StatusOr<ScopedShapedBuffer> HloRunner::ExecuteWithDeviceBuffers( |
| 223 | std::unique_ptr<HloModule> module, |
nothing calls this directly
no test coverage detected