tensorflow/compiler/mlir/xla/transforms/legalize_tf_communication.cc
// Finds all functions that need to be rewritten with communication ops and
// and associated tokens.
Finds all functions that need to be rewritten with communication ops and
and associated tokens.
Finds all functions that need to be rewritten with communication ops and and associated tokens .
LogicalResult GetFunctionsToRewrite(
    ModuleOp module,
    llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs_to_rewrite) {
  SmallVector<FuncOp, 4> funcs_to_visit;
  for (FuncOp func : module.getOps<FuncOp>()) {
    FuncToRewrite func_to_rewrite{func, {},
                                  {},
                                  nullptr};
    bool has_communication_ops = false;
    if (failed(FindCommunicationOps(func, func_to_rewrite.control_flow_ops,
                                    func_to_rewrite.control_flow_blocks,
                                    has_communication_ops)))
      return failure();
    if (!has_communication_ops) continue;
    funcs_to_rewrite.insert({func.getName(), func_to_rewrite});
    funcs_to_visit.push_back(func);
  }
  while (!funcs_to_visit.empty()) {
    SmallVector<FuncOp, 4> new_funcs_to_visit;
    for (FuncOp& func : funcs_to_visit) {
      auto uses = func.getSymbolUses(module);
      if (!uses) continue;
      for (auto& use : *uses) {
        if (!isa<mlir::func::CallOp>(use.getUser())) continue;
        auto caller_parent_func = use.getUser()->getParentOfType<FuncOp>();
        if (!caller_parent_func) continue;
        FuncToRewrite func_to_rewrite{caller_parent_func,
                                      {},
                                      {},
                                      nullptr};
        if (failed(GetControlFlowAncestors(
                use.getUser(), func_to_rewrite.control_flow_ops,
                func_to_rewrite.control_flow_blocks)))
          return failure();
        auto it = funcs_to_rewrite.insert(
            {caller_parent_func.getName(), func_to_rewrite});
        if (it.second) {
          new_funcs_to_visit.push_back(caller_parent_func);
        } else {
          it.first->getSecond().control_flow_ops.insert(
              func_to_rewrite.control_flow_ops.begin(),
              func_to_rewrite.control_flow_ops.end());
          it.first->getSecond().control_flow_blocks.insert(
              func_to_rewrite.control_flow_blocks.begin(),
              func_to_rewrite.control_flow_blocks.end());
        }
      }
    }
    funcs_to_visit.swap(new_funcs_to_visit);
  }
  SymbolTable symbol_table(module);
  for (auto& func : funcs_to_rewrite) {
    if (func.getSecond().original.isPublic() &&
        !func.getSecond().original.symbolKnownUseEmpty(module)) {
      auto clone = func.getSecond().original.clone();
      clone.setPrivate();
      symbol_table.insert(clone);
      func.getSecond().clone = clone;
    }
  }
  return success();
}
LogicalResult GetFunctionsToRewrite ( ModuleOp module , llvm :: SmallDenseMap < StringRef , FuncToRewrite > & funcs_to_rewrite ) { SmallVector < FuncOp , 4 > funcs_to_visit ; for ( FuncOp func : module . getOps < FuncOp > ( ) ) { FuncToRewrite func_to_rewrite { func , { } , { } , nullptr } ; bool has_communication_ops = false ; if ( failed ( FindCommunicationOps ( func , func_to_rewrite . control_flow_ops , func_to_rewrite . control_flow_blocks , has_communication_ops ) ) ) return failure ( ) ; if ( ! has_communication_ops ) continue ; funcs_to_rewrite . insert ( { func . getName ( ) , func_to_rewrite } ) ; funcs_to_visit . push_back ( func ) ; } while ( ! funcs_to_visit . empty ( ) ) { SmallVector < FuncOp , 4 > new_funcs_to_visit ; for ( FuncOp & func : funcs_to_visit ) { auto uses = func . getSymbolUses ( module ) ; if ( ! uses ) continue ; for ( auto & use : * uses ) { if ( ! isa < mlir :: func :: CallOp > ( use . getUser ( ) ) ) continue ; auto caller_parent_func = use . getUser ( ) -> getParentOfType < FuncOp > ( ) ; if ( ! caller_parent_func ) continue ; FuncToRewrite func_to_rewrite { caller_parent_func , { } , { } , nullptr } ; if ( failed ( GetControlFlowAncestors ( use . getUser ( ) , func_to_rewrite . control_flow_ops , func_to_rewrite . control_flow_blocks ) ) ) return failure ( ) ; auto it = funcs_to_rewrite . insert ( { caller_parent_func . getName ( ) , func_to_rewrite } ) ; if ( it . second ) { new_funcs_to_visit . push_back ( caller_parent_func ) ; } else { it . first -> getSecond ( ) . control_flow_ops . insert ( func_to_rewrite . control_flow_ops . begin ( ) , func_to_rewrite . control_flow_ops . end ( ) ) ; it . first -> getSecond ( ) . control_flow_blocks . insert ( func_to_rewrite . control_flow_blocks . begin ( ) , func_to_rewrite . control_flow_blocks . end ( ) ) ; } } } funcs_to_visit . swap ( new_funcs_to_visit ) ; } SymbolTable symbol_table ( module ) ; for ( auto & func : funcs_to_rewrite ) { if ( func . getSecond ( ) . original . isPublic ( ) && ! func . getSecond ( ) . original . symbolKnownUseEmpty ( module ) ) { auto clone = func . getSecond ( ) . original . clone ( ) ; clone . setPrivate ( ) ; symbol_table . insert ( clone ) ; func . getSecond ( ) . clone = clone ; } } return success ( ) ; }
Finds all functions that need to be rewritten with communication ops and
and associated tokens.
Finds all functions that need to be rewritten with communication ops and and associated tokens .
// Find functions containing communication ops. /*original=*/ /*control_flow_ops=*/ /*control_flow_blocks=*/ /*clone=*/ // Find functions that call functions with communication ops, transitively. // Only `mlir::func::CallOp` is supported as this requires knowing how // to rewrite arguments and results to a function. /*original=*/ /*control_flow_ops=*/ /*control_flow_blocks=*/ /*clone=*/ // Clone public functions that need to be rewritten. Function calls to this // function will be replaced with the cloned function.