1 /* Copyright 2021 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 16 #ifndef TENSORFLOW_COMPILER_MLIR_TFRT_SAVED_MODEL_SAVED_MODEL_H_ 17 #define TENSORFLOW_COMPILER_MLIR_TFRT_SAVED_MODEL_SAVED_MODEL_H_ 18 19 #include <string> 20 #include <unordered_set> 21 #include <vector> 22 23 #include "absl/container/flat_hash_map.h" 24 #include "absl/strings/string_view.h" 25 #include "mlir/IR/BuiltinTypes.h" // from @llvm-project 26 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h" 27 #include "tensorflow/core/framework/tensor_shape.h" 28 #include "tensorflow/core/framework/types.pb.h" 29 #include "tensorflow/core/platform/status.h" 30 #include "tfrt/bef/bef_buffer.h" // from @tf_runtime 31 #include "tfrt/core_runtime/tensor_handle.h" // from @tf_runtime 32 33 namespace tfrt { 34 class CoreRuntime; 35 } 36 37 namespace mlir { 38 class ModuleOp; 39 } 40 41 namespace tensorflow { 42 43 // TFRTSavedModelSignatureInfo contains the metadata for a signature in the 44 // savedmodel such as function name, inputs/outputs' names and types. This can 45 // be used to retrieve these information in a tf_saved_model module. 46 struct TFRTSavedModelSignatureInfo { 47 llvm::StringRef func_name; 48 49 // The following are metadata for inputs. 50 llvm::ArrayRef<llvm::StringRef> input_names; 51 llvm::ArrayRef< 52 std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>> 53 input_specs; 54 llvm::ArrayRef<llvm::StringRef> input_devices; 55 56 // The following are metadata for outputs. 57 llvm::ArrayRef<llvm::StringRef> output_names; 58 llvm::ArrayRef< 59 std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>> 60 output_specs; 61 62 // The following are metadata for bound_inputs, ie. captures. 63 llvm::ArrayRef<mlir::Operation*> bound_inputs; 64 }; 65 66 // Apply `map_fn` on every exported function in the module with the 67 // corresponding signature metadata populated in TFRTSavedModelSignatureInfo for 68 // the function. 69 Status MapFunctionSignaturesFromTFSavedModelMLIR( 70 mlir::ModuleOp module, 71 llvm::function_ref<void(const TFRTSavedModelSignatureInfo&)> map_fn); 72 73 } // namespace tensorflow 74 75 #endif // TENSORFLOW_COMPILER_MLIR_TFRT_SAVED_MODEL_SAVED_MODEL_H_ 76