1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
43namespace tensorflow {
44namespace data {
45
46RewriterConfig 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.
52Status 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.
64std::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.
69StatusOr<std::string> GetDatasetNode(const GraphDef& graph_def);
70
71// Like `GetDatasetNode` above, but returns the entire node object.
72StatusOr<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.
80absl::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