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. The rewritten_input 51 // stored in the core::RefCountPtr<DatasetBase>* output parameter is owned. 52 Status RewriteDataset(OpKernelContext* ctx, const DatasetBase* input, 53 std::function<RewriterConfig(void)> config_factory, 54 bool record_fingerprint, 55 core::RefCountPtr<DatasetBase>* rewritten_input); 56 57 // Creates a grappler item for `graph_def`, which is required for graph 58 // optimization. 59 // `dataset_node` is the name of the node corresponding to the dataset. 60 // If `add_fake_sinks` is true, it adds fake sink node to graph and functions to 61 // allow rewriting the actual sink nodes. 62 // TODO(b/118820916): When MetaOptimizer adds provisions for function retvals to 63 // be optimizable, we will no longer need to add fake nodes. 64 std::unique_ptr<tensorflow::grappler::GrapplerItem> GetGrapplerItem( 65 GraphDef* graph_def, std::string* dataset_node, bool add_fake_sinks); 66 67 // Returns the name of the node corresponding to the dataset. It is indicated by 68 // the symbolic `_Retval` node. 69 StatusOr<std::string> GetDatasetNode(const GraphDef& graph_def); 70 71 // Like `GetDatasetNode` above, but returns the entire node object. 72 StatusOr<NodeDef> GetDatasetNodeDef(const GraphDef& graph_def); 73 74 // Determines which optimizations should be applied. 75 // 76 // The result will contain any optimizations that are explicitly enabled, any 77 // default optimization that are not explicitly disabled, and any experiment 78 // that corresponds to an optimization as long as the optimization is not 79 // explicitly disabled. 80 absl::flat_hash_set<tstring> SelectOptimizations( 81 const absl::flat_hash_set<string>& experiments, 82 const absl::flat_hash_set<tstring>& optimizations_enabled, 83 const absl::flat_hash_set<tstring>& optimizations_disabled, 84 const absl::flat_hash_set<tstring>& optimizations_default); 85 86 } // namespace data 87 } // namespace tensorflow 88 #endif // !IS_MOBILE_PLATFORM 89 90 #endif // TENSORFLOW_CORE_DATA_REWRITE_UTILS_H_ 91