1 /* Copyright 2017 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_DATA_REWRITE_UTILS_H_ 16 #define TENSORFLOW_CORE_DATA_REWRITE_UTILS_H_ 17 18 #include "tensorflow/core/platform/platform.h" 19 20 // On mobile we do not provide this functionality because not all of its 21 // dependencies are available there. 22 #if !defined(IS_MOBILE_PLATFORM) 23 24 #include <functional> 25 #include <memory> 26 #include <string> 27 28 #include "absl/container/flat_hash_set.h" 29 #include "tensorflow/core/common_runtime/function.h" 30 #include "tensorflow/core/framework/dataset.h" 31 #include "tensorflow/core/framework/function.h" 32 #include "tensorflow/core/framework/graph.pb.h" 33 #include "tensorflow/core/framework/op_kernel.h" 34 #include "tensorflow/core/framework/resource_mgr.h" 35 #include "tensorflow/core/framework/tensor.h" 36 #include "tensorflow/core/grappler/grappler_item.h" 37 #include "tensorflow/core/platform/status.h" 38 #include "tensorflow/core/platform/statusor.h" 39 #include "tensorflow/core/platform/tstring.h" 40 #include "tensorflow/core/platform/types.h" 41 #include "tensorflow/core/protobuf/rewriter_config.pb.h" 42 43 namespace tensorflow { 44 namespace data { 45 46 RewriterConfig CreateRewriterConfig( 47 const absl::flat_hash_set<tstring>& optimizations, 48 const absl::flat_hash_set<tstring>& optimizations_configs); 49 50 // Rewrites the input dataset using the given config. 51 Status RewriteDataset(OpKernelContext* ctx, const DatasetBase* input, 52 std::function<RewriterConfig(void)> config_factory, 53 bool record_fingerprint, DatasetBase** rewritten_input); 54 55 // Creates a grappler item for `graph_def`, which is required for graph 56 // optimization. 57 // `dataset_node` is the name of the node corresponding to the dataset. 58 // If `add_fake_sinks` is true, it adds fake sink node to graph and functions to 59 // allow rewriting the actual sink nodes. 60 // TODO(b/118820916): When MetaOptimizer adds provisions for function retvals to 61 // be optimizable, we will no longer need to add fake nodes. 62 std::unique_ptr<tensorflow::grappler::GrapplerItem> GetGrapplerItem( 63 GraphDef* graph_def, std::string* dataset_node, bool add_fake_sinks); 64 65 // Returns the name of the node corresponding to the dataset. It is indicated by 66 // the symbolic `_Retval` node. 67 StatusOr<std::string> GetDatasetNode(const GraphDef& graph_def); 68 69 } // namespace data 70 } // namespace tensorflow 71 #endif // !IS_MOBILE_PLATFORM 72 73 #endif // TENSORFLOW_CORE_DATA_REWRITE_UTILS_H_ 74