From ecc6fe906649c3d76473c38ecd5295790ad59d79 Mon Sep 17 00:00:00 2001 From: Brian Zhao Date: Wed, 26 Aug 2020 23:33:17 -0700 Subject: [PATCH] Continue splitting SignatureDefFunction from ConcreteFunction. PiperOrigin-RevId: 328682987 Change-Id: I4d32d48bbfb5f461c5e3edc50513057808c8c47f --- .../saved_model/core/revived_types/BUILD | 23 +++++ .../tf_signature_def_function.cc | 96 +++++++++++++++++++ .../revived_types/tf_signature_def_function.h | 90 +++++++++++++++++ 3 files changed, 209 insertions(+) create mode 100644 tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.cc create mode 100644 tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h diff --git a/tensorflow/c/experimental/saved_model/core/revived_types/BUILD b/tensorflow/c/experimental/saved_model/core/revived_types/BUILD index 2b883618c87..21a8dae3d93 100644 --- a/tensorflow/c/experimental/saved_model/core/revived_types/BUILD +++ b/tensorflow/c/experimental/saved_model/core/revived_types/BUILD @@ -81,3 +81,26 @@ cc_library( "@com_google_absl//absl/types:span", ], ) + +cc_library( + name = "tf_signature_def_function", + srcs = [ + "tf_signature_def_function.cc", + ], + hdrs = [ + "tf_signature_def_function.h", + ], + deps = [ + ":tensorhandle_convertible", + "//tensorflow/c/eager:abstract_tensor_handle", + "//tensorflow/c/eager:immediate_execution_context", + "//tensorflow/c/eager:immediate_execution_operation", + "//tensorflow/c/eager:immediate_execution_tensor_handle", + "//tensorflow/c/experimental/saved_model/core:signature_def_function", + "//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core/common_runtime/eager:context", + "@com_google_absl//absl/types:span", + ], +) diff --git a/tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.cc b/tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.cc new file mode 100644 index 00000000000..2a9f23b63cd --- /dev/null +++ b/tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.cc @@ -0,0 +1,96 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h" + +#include +#include + +#include "absl/types/span.h" +#include "tensorflow/c/eager/abstract_tensor_handle.h" +#include "tensorflow/c/eager/immediate_execution_operation.h" +#include "tensorflow/c/eager/immediate_execution_tensor_handle.h" +#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h" +#include "tensorflow/core/common_runtime/eager/context.h" +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/status.h" +#include "tensorflow/core/protobuf/saved_object_graph.pb.h" +#include "tensorflow/core/protobuf/struct.pb.h" + +namespace tensorflow { + +TFSignatureDefFunction::TFSignatureDefFunction( + const std::string& name, + std::vector captures, + SignatureDefFunctionMetadata metadata, ImmediateExecutionContext* ctx) + : name_(name), + captures_(std::move(captures)), + metadata_(std::move(metadata)), + ctx_(ctx) {} + +TFSignatureDefFunction::~TFSignatureDefFunction() { + Status status = ctx_->RemoveFunction(name_); + if (!status.ok()) { + LOG(ERROR) << "Failed to remove functiondef " << name_ << ". " + << status.error_message(); + } +} + +Status TFSignatureDefFunction::Create( + const FunctionDef* function_def, + std::vector captures, + SignatureDefFunctionMetadata metadata, ImmediateExecutionContext* ctx, + std::unique_ptr* out) { + TF_RETURN_IF_ERROR(ctx->AddFunctionDef(*function_def)); + out->reset(new TFSignatureDefFunction(function_def->signature().name(), + std::move(captures), + std::move(metadata), ctx)); + return Status(); +} + +const SignatureDefFunctionMetadata& +TFSignatureDefFunction::GetFunctionMetadata() const { + return metadata_; +} + +Status TFSignatureDefFunction::MakeCallOp( + absl::Span inputs, ImmediateOpPtr* out) const { + out->reset(ctx_->CreateOperation()); + // In eager mode, TF2 python executes functions by constructing an op with + // the name of the functiondef: + // https://github.com/tensorflow/tensorflow/blob/66668ec0ca432e2f38a575b814f45b6d299d01ed/tensorflow/python/eager/function.py#L545 + // In graph mode, we create a PartitionedCallOp instead: + // https://github.com/tensorflow/tensorflow/blob/66668ec0ca432e2f38a575b814f45b6d299d01ed/tensorflow/python/eager/function.py#L573 + + // TODO(bmzhao): After discussing with Allen, we should execute this via a + // PartitionedCallOp for compatibility with "tooling that assumes functions in + // graphs are PartitionedCallOps". + TF_RETURN_IF_ERROR((*out)->Reset(name_.c_str(), nullptr)); + + // Adding the user-provided inputs to the function. + TF_RETURN_IF_ERROR((*out)->AddInputList(inputs)); + + absl::Span captures( + reinterpret_cast(captures_.data()), + captures_.size()); + + // Adding the captures of the function. + TF_RETURN_IF_ERROR((*out)->AddInputList(captures)); + return Status(); +} + +} // namespace tensorflow diff --git a/tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h b/tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h new file mode 100644 index 00000000000..523560685c3 --- /dev/null +++ b/tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h @@ -0,0 +1,90 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SIGNATURE_DEF_FUNCTION_H_ +#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SIGNATURE_DEF_FUNCTION_H_ + +#include +#include +#include +#include +#include + +#include "tensorflow/c/eager/immediate_execution_context.h" +#include "tensorflow/c/eager/immediate_execution_operation.h" +#include "tensorflow/c/eager/immediate_execution_tensor_handle.h" +#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h" +#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h" +#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h" +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/protobuf/saved_object_graph.pb.h" + +namespace tensorflow { + +// This is the TF eager runtime implementation of SignatureDefFunction (separate +// from the TFRT implementation). The user-facing API of SignatureDefFunctions +// and their semantic differences from ConcreteFunction are described here: +// https://github.com/tensorflow/tensorflow/blob/e2db60c9d9598ebae0b7741587ce6f5d473584d9/tensorflow/cc/saved_model/experimental/public/signature_def_function.h#L30-L59 +// Additional implementation notes are available here: +// https://github.com/tensorflow/tensorflow/blob/e2db60c9d9598ebae0b7741587ce6f5d473584d9/tensorflow/c/experimental/saved_model/core/signature_def_function.h#L31-L48 +class TFSignatureDefFunction : public SignatureDefFunction { + public: + // Factory function for creating a TFSignatureDefFunction. + // + // Params: + // function_def - The function_def associated with the created + // TFSignatureDefFunction. TFSignatureDefFunction will + // register this function_def with `ctx` on creation, and + // de-register it on destruction. function_def must be + // non-null, but otherwise has no lifetime requirements. + // captures - The captured TensorHandles associated with this + // TFConcreteFunction. + // metadata - FunctionMetadata associated with this TFSignatureDefFunction. + // ctx - A handle to the Tensorflow runtime. This MUST be non-null and + // outlive TFSignatureDefFunction. + // out - The output TFSignatureDefFunction. + static Status Create(const FunctionDef* function_def, + std::vector captures, + SignatureDefFunctionMetadata metadata, + ImmediateExecutionContext* ctx, + std::unique_ptr* out); + + // This method creates a "Call" Op used to execute the function. + Status MakeCallOp(absl::Span inputs, + ImmediateOpPtr* out) const override; + + const SignatureDefFunctionMetadata& GetFunctionMetadata() const override; + + ~TFSignatureDefFunction() override; + + private: + TFSignatureDefFunction(const std::string& name, + std::vector captures, + SignatureDefFunctionMetadata metadata, + ImmediateExecutionContext* ctx); + + TFSignatureDefFunction(const TFSignatureDefFunction&) = delete; + TFSignatureDefFunction& operator=(const TFSignatureDefFunction&) = delete; + + // Name of the FunctionDef corresponding to this TFSignatureDefFunction + std::string name_; + std::vector captures_; + SignatureDefFunctionMetadata metadata_; + ImmediateExecutionContext* ctx_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SIGNATURE_DEF_FUNCTION_H_