| 416 | } |
| 417 | |
| 418 | void StartBody() { |
| 419 | Status s; |
| 420 | if (rets_.size() != 1) { |
| 421 | s = errors::InvalidArgument( |
| 422 | "Expected a single scalar return value from WhileOp cond, got ", |
| 423 | rets_.size(), " tensors."); |
| 424 | return Finish(s); |
| 425 | } |
| 426 | Tensor cond_t; |
| 427 | #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM |
| 428 | const DeviceBase::GpuDeviceInfo* gpu_device_info = |
| 429 | ctx_->device()->tensorflow_gpu_device_info(); |
| 430 | const bool is_hostmem_dtype = |
| 431 | rets_[0].dtype() == DT_INT32 || rets_[0].dtype() == DT_INT64; |
| 432 | if (!is_hostmem_dtype && gpu_device_info && |
| 433 | (opts_.rets_alloc_attrs.empty() || |
| 434 | !opts_.rets_alloc_attrs[0].on_host())) { |
| 435 | // Copy the ret value to host if it's allocated on device. |
| 436 | Device* device = down_cast<Device*>(ctx_->device()); |
| 437 | DeviceContext* device_ctx = ctx_->op_device_context(); |
| 438 | cond_t = Tensor(rets_[0].dtype(), rets_[0].shape()); |
| 439 | Notification done_copy; |
| 440 | device_ctx->CopyDeviceTensorToCPU( |
| 441 | &rets_[0], /*tensor_name=*/"", device, &cond_t, |
| 442 | [&done_copy, &s](const Status& status) { |
| 443 | s = status; |
| 444 | done_copy.Notify(); |
| 445 | }); |
| 446 | done_copy.WaitForNotification(); |
| 447 | if (!s.ok()) { |
| 448 | return Finish(s); |
| 449 | } |
| 450 | } else { |
| 451 | cond_t = rets_[0]; |
| 452 | } |
| 453 | #else |
| 454 | cond_t = rets_[0]; |
| 455 | #endif |
| 456 | bool cond; |
| 457 | s = ToBool({cond_t}, &cond); |
| 458 | |
| 459 | if (!s.ok()) { |
| 460 | return Finish(s); |
| 461 | } |
| 462 | if (!cond) { |
| 463 | return Finish(Status::OK()); |
| 464 | } |
| 465 | rets_.clear(); |
| 466 | profiler::TraceMe trace_me( |
| 467 | [&] { |
| 468 | return absl::StrCat( |
| 469 | "WhileOp-StartBody #parent_step_id=", ctx_->step_id(), |
| 470 | ",function_step_id=", opts_.step_id, "#"); |
| 471 | }, |
| 472 | /*level=*/2); |
| 473 | lib_->Run( |
| 474 | // Evaluate the body. |
| 475 | opts_, body_handle_, args_, &rets_, |
nothing calls this directly
no test coverage detected