| 285 | } |
| 286 | |
| 287 | Status GpuExecutable::ExecuteThunks( |
| 288 | const ServiceExecutableRunOptions* run_options, |
| 289 | const BufferAllocations& buffer_allocations, bool block_host_until_done, |
| 290 | HloExecutionProfile* hlo_execution_profile) { |
| 291 | TF_RETURN_IF_ERROR( |
| 292 | CheckCompatibilityWithServiceExecutableRunOptions(run_options)); |
| 293 | GpuDebugInfoManager::Get()->OnModuleStart(module().name()); |
| 294 | auto cleanup = MakeCleanup( |
| 295 | [&]() { GpuDebugInfoManager::Get()->OnModuleStop(module().name()); }); |
| 296 | |
| 297 | se::Stream* main_stream = run_options->stream(); |
| 298 | se::StreamExecutor* executor = main_stream->parent(); |
| 299 | SetExecutor(executor->implementation()); |
| 300 | |
| 301 | bool do_profile = hlo_execution_profile != nullptr; |
| 302 | if (do_profile) { |
| 303 | LOG(WARNING) << "PROFILING: profiling is enabled"; |
| 304 | } |
| 305 | |
| 306 | // Stream 0 indicates `main_stream` and substreams start from stream 1. |
| 307 | std::vector<StreamPool::Ptr> sub_streams; |
| 308 | sub_streams.reserve(thunk_schedule_->StreamCount() - 1); |
| 309 | while (sub_streams.size() + 1 < thunk_schedule_->StreamCount()) { |
| 310 | sub_streams.emplace_back(); |
| 311 | TF_ASSIGN_OR_RETURN(sub_streams.back(), |
| 312 | run_options->BorrowStream(executor->device_ordinal())); |
| 313 | } |
| 314 | |
| 315 | se::Stream* capture_stream = main_stream; |
| 316 | StreamPool::Ptr private_capture_stream; |
| 317 | bool use_gpu_graph_capture = |
| 318 | GpuGraphCaptureEnabled(module().name()) && CanUseGpuGraphCapture() && |
| 319 | executor->platform_kind() == stream_executor::PlatformKind::kCuda && |
| 320 | !is_graph_capture_costly_; |
| 321 | if (use_gpu_graph_capture) { |
| 322 | // We need a private stream for capturing to avoid interference from other |
| 323 | // threads. |
| 324 | TF_ASSIGN_OR_RETURN(private_capture_stream, |
| 325 | run_options->BorrowStream(executor->device_ordinal())); |
| 326 | capture_stream = private_capture_stream.get(); |
| 327 | } |
| 328 | |
| 329 | HloExecutionProfiler profiler(do_profile, hlo_execution_profile, main_stream, |
| 330 | sub_streams, hlo_module_->entry_computation()); |
| 331 | uint64 start_micros = tensorflow::Env::Default()->NowMicros(); |
| 332 | |
| 333 | tensorflow::profiler::TraceMe hlo_module_activity( |
| 334 | [&] { return absl::StrCat(hlo_module_->name(), ":XLA GPU module"); }, |
| 335 | tensorflow::profiler::TraceMeLevel::kInfo); |
| 336 | std::vector<std::function<void()>> deferred_host_callbacks; |
| 337 | if (!use_gpu_graph_capture) { |
| 338 | TF_RETURN_IF_ERROR(ExecuteThunkSequence( |
| 339 | run_options, buffer_allocations, profiler, main_stream, capture_stream, |
| 340 | sub_streams, deferred_host_callbacks)); |
| 341 | } else { |
| 342 | auto bufs_key = buffer_allocations.Key(); |
| 343 | auto temp_buf_key = buffer_allocations.TempBufferKey(); |
| 344 | VLOG(3) << "For gpu executable " << this << " tmp buffer base: " |
nothing calls this directly
no test coverage detected