Initialize the following graph: (cpu0) (cpu1) a b | | c1 c1 | | id id / \ / \ // c2 c3 c2 c3 Here ci denotes a collective node with `instance_key` i. `a` and `b` are inputs, `id` is identity node.
| 108 | // Here ci denotes a collective node with `instance_key` i. `a` and `b` are |
| 109 | // inputs, `id` is identity node. |
| 110 | std::unique_ptr<Graph> InitGraph() { |
| 111 | GraphDefBuilder builder(GraphDefBuilder::kFailImmediately); |
| 112 | const string dev0 = "/job:localhost/replica:0/task:0/device:CPU:0"; |
| 113 | const string dev1 = "/job:localhost/replica:0/task:0/device:CPU:1"; |
| 114 | Node* a = ops::SourceOp("TestParams", |
| 115 | builder.opts().WithName("a").WithDevice(dev0)); |
| 116 | Node* b = ops::SourceOp("TestParams", |
| 117 | builder.opts().WithName("b").WithDevice(dev1)); |
| 118 | Node* c1_0 = CollectiveReduceNode(&builder, a, "c1_0", dev0, 1); |
| 119 | Node* c1_1 = CollectiveReduceNode(&builder, b, "c1_1", dev1, 1); |
| 120 | Node* id0 = ops::UnaryOp( |
| 121 | "Identity", c1_0, |
| 122 | builder.opts().WithName("id0").WithDevice(dev0).WithAttr("T", DT_FLOAT)); |
| 123 | Node* id1 = ops::UnaryOp( |
| 124 | "Identity", c1_1, |
| 125 | builder.opts().WithName("id1").WithDevice(dev1).WithAttr("T", DT_FLOAT)); |
| 126 | CollectiveReduceNode(&builder, id0, "c2_0", dev0, 2); |
| 127 | CollectiveReduceNode(&builder, id1, "c2_1", dev1, 2); |
| 128 | CollectiveReduceNode(&builder, id0, "c3_0", dev0, 3); |
| 129 | CollectiveReduceNode(&builder, id1, "c3_1", dev1, 3); |
| 130 | |
| 131 | std::unique_ptr<Graph> graph = absl::make_unique<Graph>(OpRegistry::Global()); |
| 132 | Status s = GraphDefBuilderToGraph(builder, graph.get()); |
| 133 | if (!s.ok()) { |
| 134 | LOG(FATAL) << "Error building graph " << s; |
| 135 | } |
| 136 | return graph; |
| 137 | } |
| 138 | |
| 139 | // Tests that in the graph created by `InitGraph`, exactly 2 control edges are |
| 140 | // added after calling `OrderCollectives`: c3_0 -> c2_0 and c3_1 -> c2_1. |
no test coverage detected