tensorflow/compiler/xla/service/flatten_call_graph.cc
// Flatten a single call graph node. Expects to visit nodes in postorder.
Flatten a single call graph node. Expects to visit nodes in postorder.
Flatten a single call graph node . Expects to visit nodes in postorder .
Status FlattenNode(const CallGraphNode& node) {
  HloComputation* computation = node.computation();
  HloModule* module = computation->parent();
  for (int i = 0; i < node.caller_callsites().size(); ++i) {
    CallSite call_site = node.caller_callsites()[i];
    if (call_site.context() == CallContext::kEmbedded) {
      continue;
    }
    CHECK_EQ(call_site.context(), CallContext::kControlFlow);
    if (node.context() != CallContext::kBoth && i == 0) {
      continue;
    }
    HloComputation* clone =
        module->AddEmbeddedComputation(computation->Clone());
    ReplaceCalledComputation(call_site.instruction(), computation, clone);
    std::vector<HloComputation*> worklist;
    worklist.push_back(clone);
    while (!worklist.empty()) {
      auto current = worklist.back();
      worklist.pop_back();
      for (auto* instruction : current->instructions()) {
        if (GetInstructionCallContext(instruction->opcode()) !=
            CallContext::kControlFlow) {
          continue;
        }
        for (auto callee : instruction->called_computations()) {
          HloComputation* callee_clone =
              module->AddEmbeddedComputation(callee->Clone());
          ReplaceCalledComputation(instruction, callee, callee_clone);
          worklist.push_back(callee_clone);
        }
      }
    }
  }
  return Status::OK();
}
Status FlattenNode ( const CallGraphNode & node ) { HloComputation * computation = node . computation ( ) ; HloModule * module = computation -> parent ( ) ; for ( int i = 0 ; i < node . caller_callsites ( ) . size ( ) ; ++ i ) { CallSite call_site = node . caller_callsites ( ) [ i ] ; if ( call_site . context ( ) == CallContext :: kEmbedded ) { continue ; } CHECK_EQ ( call_site . context ( ) , CallContext :: kControlFlow ) ; if ( node . context ( ) != CallContext :: kBoth && i == 0 ) { continue ; } HloComputation * clone = module -> AddEmbeddedComputation ( computation -> Clone ( ) ) ; ReplaceCalledComputation ( call_site . instruction ( ) , computation , clone ) ; std :: vector < HloComputation * > worklist ; worklist . push_back ( clone ) ; while ( ! worklist . empty ( ) ) { auto current = worklist . back ( ) ; worklist . pop_back ( ) ; for ( auto * instruction : current -> instructions ( ) ) { if ( GetInstructionCallContext ( instruction -> opcode ( ) ) != CallContext :: kControlFlow ) { continue ; } for ( auto callee : instruction -> called_computations ( ) ) { HloComputation * callee_clone = module -> AddEmbeddedComputation ( callee -> Clone ( ) ) ; ReplaceCalledComputation ( instruction , callee , callee_clone ) ; worklist . push_back ( callee_clone ) ; } } } } return Status :: OK ( ) ; }
Flatten a single call graph node.
Flatten a single call graph node .
// Clone callee for all call-sites except the first one. // Only consider sequential call contexts. // Skip first element if this computation is only called from a sequential // context. // Clone computation for the remaining sequential context call sites. // Clone the sub-tree of all computations called from this node.