| 120 | } |
| 121 | |
| 122 | TEST(LowerIfWhileTest, CondInWhile) { |
| 123 | // Tests the value of `a` for different values of args after the following |
| 124 | // program: |
| 125 | // |
| 126 | // Args: |
| 127 | // counter = Placeholder(type = int32) |
| 128 | // pred = Placeholder(type = bool) |
| 129 | // a = Placeholder(type = int32) |
| 130 | // N = 3 |
| 131 | // while (counter < N) { |
| 132 | // counter += 1; |
| 133 | // if (pred) { |
| 134 | // a *= 2; |
| 135 | // } else { |
| 136 | // a *= 4; |
| 137 | // } |
| 138 | // } |
| 139 | |
| 140 | std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global())); |
| 141 | |
| 142 | FunctionDefLibrary f_lib_proto; |
| 143 | // Cond functions. |
| 144 | *f_lib_proto.add_function() = test::function::XTimesTwo(); |
| 145 | *f_lib_proto.add_function() = test::function::XTimesFour(); |
| 146 | // While functions. |
| 147 | *f_lib_proto.add_function() = WhileWithIfCond(3); |
| 148 | *f_lib_proto.add_function() = WhileWithIfBody(); |
| 149 | |
| 150 | Scope root = Scope::NewRootScope().ExitOnError(); |
| 151 | TF_ASSERT_OK(root.graph()->AddFunctionLibrary(f_lib_proto)); |
| 152 | auto counter = ops::Placeholder(root.WithOpName("counter"), DT_INT32); |
| 153 | auto pred = ops::Placeholder(root.WithOpName("pred"), DT_BOOL); |
| 154 | auto a = ops::Placeholder(root.WithOpName("A"), DT_INT32); |
| 155 | std::vector<NodeBuilder::NodeOut> inputs( |
| 156 | {NodeBuilder::NodeOut(counter.node()), NodeBuilder::NodeOut(pred.node()), |
| 157 | NodeBuilder::NodeOut(a.node())}); |
| 158 | Node* while_node; |
| 159 | AttrValue cond_func; |
| 160 | cond_func.mutable_func()->set_name("WhileWithIfCond"); |
| 161 | AttrValue body_func; |
| 162 | body_func.mutable_func()->set_name("WhileWithIfBody"); |
| 163 | TF_ASSERT_OK(NodeBuilder("while", "While", &root.graph()->flib_def()) |
| 164 | .Input(inputs) |
| 165 | .Attr("T", {DT_INT32, DT_BOOL, DT_INT32}) |
| 166 | .Attr("cond", cond_func) |
| 167 | .Attr("body", body_func) |
| 168 | .Attr(kLowerUsingSwitchMergeAttr, true) |
| 169 | .Finalize(root.graph(), &while_node)); |
| 170 | TF_ASSERT_OK(root.DoShapeInference(while_node)); |
| 171 | TF_ASSERT_OK(root.ToGraph(graph.get())); |
| 172 | TF_ASSERT_OK(Rewrite(&graph)); |
| 173 | |
| 174 | // Lowered graph has no While and If ops. |
| 175 | for (const auto* op : graph->op_nodes()) { |
| 176 | ASSERT_NE(op->type_string(), "While"); |
| 177 | ASSERT_NE(op->type_string(), "If"); |
| 178 | } |
| 179 |
nothing calls this directly
no test coverage detected