1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_CORE_RUNTIME_FALLBACK_KERNEL_KERNEL_FALLBACK_UTILS_H_ 16 #define TENSORFLOW_CORE_RUNTIME_FALLBACK_KERNEL_KERNEL_FALLBACK_UTILS_H_ 17 18 #include "llvm/ADT/ArrayRef.h" 19 #include "tensorflow/core/framework/device.h" 20 #include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h" 21 #include "tensorflow/core/tfrt/fallback/op_kernel_runner.h" 22 #include "tensorflow/core/tfrt/utils/fallback_tensor.h" 23 #include "tfrt/host_context/async_value.h" // from @tf_runtime 24 #include "tfrt/host_context/sync_kernel_utils.h" // from @tf_runtime 25 #include "tfrt/host_context/value.h" // from @tf_runtime 26 #include "tfrt/support/variant.h" // from @tf_runtime 27 28 namespace tensorflow { 29 namespace tfd { 30 31 using TfInputs = 32 tfrt::Variant<tfrt::Monostate, llvm::ArrayRef<tfrt::AsyncValue*>, 33 tfrt::RepeatedSyncArguments<tfrt_stub::FallbackTensor>&>; 34 35 // Sets up the OpKernelcontext::Params in `run_state` with the objects and data 36 // in `runner`, `fallback_request_state` and `device`. 37 void SetUpParams(const tensorflow::tfrt_stub::OpKernelRunner& runner, 38 const KernelFallbackCompatRequestState& fallback_request_state, 39 tensorflow::Device* device, 40 tensorflow::tfrt_stub::OpKernelRunState& run_state); 41 42 // Return the device to be used for the fallback kernel execution. The device is 43 // guaranteed to be alive during the graph execution. 44 tensorflow::Device* GetDeviceFromFallbackState( 45 const KernelFallbackCompatRequestState& fallback_request_state, 46 const tfrt_stub::OpKernelRunner& kernel_runner); 47 48 } // namespace tfd 49 } // namespace tensorflow 50 51 #endif // TENSORFLOW_CORE_RUNTIME_FALLBACK_KERNEL_KERNEL_FALLBACK_UTILS_H_ 52