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
16// The builtin inputs provide a mechanism to generate simple TensorFlow graphs
17// and feed them as inputs to Grappler. This can be used for quick experiments
18// or to derive small regression tests.
19
20#include "tensorflow/cc/ops/standard_ops.h"
21
22#include "tensorflow/core/framework/graph.pb.h"
23#include "tensorflow/core/grappler/grappler_item.h"
24#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
25
26namespace tensorflow {
27namespace grappler {
28
29// Make a program with specified number of stages and "width" ops per stage.
30namespace {
31GraphDef CreateGraphDef(int num_stages, int width, int tensor_size,
32                        bool use_multiple_devices, bool insert_queue,
33                        const std::vector<string>& device_names) {
34  using namespace ::tensorflow::ops;  // NOLINT(build/namespaces)
35
36  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
37
38  // x is from the feed.
39  const int batch_size = tensor_size < 0 ? 1 : tensor_size;
40  Output x = RandomNormal(s.WithOpName("x").WithDevice("/CPU:0"),
41                          {batch_size, 1}, DataType::DT_FLOAT);
42
43  // Create stages.
44  std::vector<Output> last_stage;
45  last_stage.push_back(x);
46  for (int i = 0; i < num_stages; i++) {
47    std::vector<Output> this_stage;
48    for (int j = 0; j < width; j++) {
49      if (last_stage.size() == 1) {
50        Output unary_op = Square(
51            s.WithDevice(
52                device_names[use_multiple_devices ? j % device_names.size()
53                                                  : 0]),
54            last_stage[0]);
55        this_stage.push_back(unary_op);
56      } else {
57        Output combine =
58            AddN(s.WithDevice(
59                     device_names[use_multiple_devices ? j % device_names.size()
60                                                       : 0]),
61                 last_stage);
62        this_stage.push_back(combine);
63      }
64    }
65    last_stage = this_stage;
66  }
67
68  if (insert_queue) {
69    FIFOQueue queue(s.WithOpName("queue").WithDevice("/CPU:0"),
70                    {DataType::DT_FLOAT});
71    QueueEnqueue enqueue(s.WithOpName("enqueue").WithDevice("/CPU:0"), queue,
72                         last_stage);
73    QueueDequeue dequeue(s.WithOpName("dequeue").WithDevice("/CPU:0"), queue,
74                         {DataType::DT_FLOAT});
75    QueueClose cancel(s.WithOpName("cancel").WithDevice("/CPU:0"), queue,
76                      QueueClose::CancelPendingEnqueues(true));
77    last_stage = {dequeue[0]};
78  }
79
80  // Create output.
81  AddN output(s.WithOpName("y").WithDevice("/CPU:0"), last_stage);
82
83  GraphDef def;
84  TF_CHECK_OK(s.ToGraphDef(&def));
85  return def;
86}
87}  // namespace
88
89TrivialTestGraphInputYielder::TrivialTestGraphInputYielder(
90    int num_stages, int width, int tensor_size, bool insert_queue,
91    const std::vector<string>& device_names)
92    : num_stages_(num_stages),
93      width_(width),
94      tensor_size_(tensor_size),
95      insert_queue_(insert_queue),
96      device_names_(device_names) {}
97
98bool TrivialTestGraphInputYielder::NextItem(GrapplerItem* item) {
99  GrapplerItem r;
100  r.id = strings::StrCat("ns:", num_stages_, "/",  // wrap
101                         "w:", width_, "/",        // wrap
102                         "ts:", tensor_size_);
103  r.graph = CreateGraphDef(num_stages_, width_, tensor_size_,
104                           true /*use_multiple_devices*/, insert_queue_,
105                           device_names_);
106  // If the batch size is variable, we need to choose a value to create a feed
107  const int batch_size = tensor_size_ < 0 ? 1 : tensor_size_;
108  Tensor x(DT_FLOAT, TensorShape({batch_size, 1}));
109  r.feed.push_back(std::make_pair("x", x));
110  r.fetch.push_back("y");
111
112  if (insert_queue_) {
113    QueueRunnerDef queue_runner;
114    queue_runner.set_queue_name("queue");
115    queue_runner.set_cancel_op_name("cancel");
116    *queue_runner.add_enqueue_op_name() = "enqueue";
117    r.queue_runners.push_back(queue_runner);
118  }
119
120  *item = std::move(r);
121  return true;
122}
123
124}  // end namespace grappler
125}  // end namespace tensorflow
126