| 97 | }; |
| 98 | |
| 99 | TEST_F(FlattenCallGraphTest, ComplexGraph) { |
| 100 | // Test a call graph of a module with several computation called in various |
| 101 | // contexts. The call graph looks like: |
| 102 | // |
| 103 | // entry |
| 104 | // / | |
| 105 | // a | |
| 106 | // / | \ | |
| 107 | // b | cond |
| 108 | // \ | |
| 109 | // c |
| 110 | // |
| 111 | // Calls are made via kCall, kWhile, and kMap instructions. |
| 112 | auto module = CreateNewVerifiedModule(); |
| 113 | HloComputation* cond_computation = |
| 114 | module->AddEmbeddedComputation(MakeConditionComputation()); |
| 115 | HloComputation* c_computation = |
| 116 | module->AddEmbeddedComputation(MakeScalarComputation()); |
| 117 | HloComputation* b_computation = module->AddEmbeddedComputation( |
| 118 | MakeMappingComputation(c_computation, /*callsites=*/1)); |
| 119 | |
| 120 | HloComputation* a_computation; |
| 121 | { |
| 122 | HloComputation::Builder builder(TestName() + ".a"); |
| 123 | HloInstruction* param0 = builder.AddInstruction( |
| 124 | HloInstruction::CreateParameter(0, kScalarShape, "param0")); |
| 125 | HloInstruction* call = builder.AddInstruction( |
| 126 | HloInstruction::CreateCall(kScalarShape, {param0}, c_computation)); |
| 127 | builder.AddInstruction(HloInstruction::CreateWhile( |
| 128 | kScalarShape, cond_computation, b_computation, call)); |
| 129 | a_computation = module->AddEmbeddedComputation(builder.Build()); |
| 130 | } |
| 131 | |
| 132 | HloComputation* entry_computation; |
| 133 | { |
| 134 | HloComputation::Builder builder(TestName() + ".entry"); |
| 135 | HloInstruction* param0 = builder.AddInstruction( |
| 136 | HloInstruction::CreateParameter(0, kScalarShape, "param0")); |
| 137 | builder.AddInstruction(HloInstruction::CreateWhile( |
| 138 | kScalarShape, cond_computation, a_computation, param0)); |
| 139 | entry_computation = module->AddEntryComputation(builder.Build()); |
| 140 | } |
| 141 | |
| 142 | { |
| 143 | TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module.get())); |
| 144 | EXPECT_TRUE(result); |
| 145 | std::unique_ptr<CallGraph> flat_call_graph = CallGraph::Build(module.get()); |
| 146 | const CallGraphNode& c_node = flat_call_graph->GetNode(c_computation); |
| 147 | EXPECT_EQ(1, c_node.caller_callsites().size()); |
| 148 | } |
| 149 | } |
| 150 | |
| 151 | // Test corner case of a computation used as a body and a loop condition. |
| 152 | TEST_F(FlattenCallGraphTest, SharedWhileConditionAndBody) { |
nothing calls this directly
no test coverage detected