There are a lot of references to devices in this function and around. Here is what they mean: EagerOperation::Device(): The device on which the user requested the op be executed, except if we had to change the device due to resource inputs or CPU pinning. If the user did not request a device, the op does not take resources, and we did not pin it to CPU, the device can be nullptr. KernelAndDevice::
| 453 | // a function with explicitly requested device has different behavior than |
| 454 | // running without an explicitly requested device. |
| 455 | Status EagerLocalExecute(EagerOperation* op, TensorHandle** retvals, |
| 456 | int* num_retvals) { |
| 457 | profiler::TraceMe activity( |
| 458 | [&] { return absl::StrCat("EagerLocalExecute: ", op->Name()); }, |
| 459 | profiler::TraceMeLevel::kInfo); |
| 460 | EagerContext* ctx = op->EagerContext(); |
| 461 | auto* executor = op->Executor(); |
| 462 | TF_RETURN_IF_ERROR(executor->status()); |
| 463 | Device* device = op->Device(); |
| 464 | |
| 465 | Fprint128 cache_key = op->MutableAttrs()->CacheKey( |
| 466 | DeviceNameOrUnspecified(op->GetDeviceName())); |
| 467 | |
| 468 | bool is_multi_device_function = |
| 469 | IsMultiDevice(ctx->FindFunctionDef(op->Name())); |
| 470 | |
| 471 | std::vector<Device*> input_dev_ptrs; |
| 472 | // `input_tensor_shapes` contains (potentially a subset of) non DT_RESOURCE |
| 473 | // arguments, and `input_resource_variable_dtypes_and_shapes` contains shapes |
| 474 | // and underlying types for (potentially a subset) of DT_RESOURCE arguments. |
| 475 | std::unordered_map<int, TensorShape> input_tensor_shapes; |
| 476 | std::unordered_map<int, DtypeAndPartialTensorShape> |
| 477 | input_resource_variable_dtypes_and_shapes; |
| 478 | if (is_multi_device_function) { |
| 479 | profiler::TraceMe activity("EagerCopyToDeviceAndAddCacheKey", |
| 480 | profiler::TraceMeLevel::kInfo); |
| 481 | input_dev_ptrs.reserve(op->Inputs().size()); |
| 482 | // All inputs need to be on local devices. |
| 483 | // TODO(b/122851476): This is a limitation of the current code base (but |
| 484 | // should be possible to get around). |
| 485 | // Code changes will need to be made to pass input objects to the |
| 486 | // function library runtime instead of just "Tensor"s. |
| 487 | // Once that is the case, we will be able to write a thin wrapper layer over |
| 488 | // the EagerService that behaves similar to the current |
| 489 | // ClusterFunctionLibraryRuntime/DistributedFunctionLibraryRuntime. |
| 490 | for (int i = 0; i < op->Inputs().size(); i++) { |
| 491 | TensorHandle* input = op->Inputs()[i]; |
| 492 | if (input->IsRemote()) { |
| 493 | TensorHandle* handle = nullptr; |
| 494 | TF_RETURN_IF_ERROR(EagerCopyToDevice( |
| 495 | input, ctx, executor, device == nullptr ? ctx->HostCPU() : device, |
| 496 | ctx->MirrorTensors(), &handle)); |
| 497 | op->UpdateInput(i, handle); |
| 498 | // Unref handle since it has a ref as an input now |
| 499 | handle->Unref(); |
| 500 | input = handle; |
| 501 | } |
| 502 | |
| 503 | // Get device for this input, and add it to 'cache_key'. |
| 504 | Device* input_device; |
| 505 | TF_RETURN_IF_ERROR(GetDeviceForInput(ctx, input, &input_device)); |
| 506 | input_dev_ptrs.push_back(input_device); |
| 507 | cache_key = |
| 508 | FingerprintCat128(cache_key, Fingerprint128(input_device->name())); |
| 509 | |
| 510 | // If input is normal tensor, get its shape and add it to 'cache_key'; |
| 511 | // If input is a ResourceHandle, get its resource handle dtypes and shapes |
| 512 | // and add them to 'cache_key'. |
no test coverage detected