tensorflow/compiler/mlir/xla/transforms/legalize_tf_communication.cc
// Collects all control flow op ancestors of communication ops or function calls
// with communication ops (transitively).
Collects all control flow op ancestors of communication ops or function calls
with communication ops (transitively).
Collects all control flow op ancestors of communication ops or function calls with communication ops ( transitively ) .
void GetCommunicationControlFlowOps(
    FuncOp func,
    const llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs_to_rewrite,
    llvm::SmallPtrSetImpl<Operation*>& control_flow_ops,
    llvm::SmallPtrSetImpl<Block*>& control_flow_blocks) {
  func.walk([&](Operation* op) {
    if (IsCommunicationOp(op) ||
        IsFunctionCallWithCommunication(op, funcs_to_rewrite))
      if (failed(GetControlFlowAncestors(op, control_flow_ops,
                                         control_flow_blocks)))
        llvm_unreachable(
            "checking original function for control flow ancestors should have "
            "errored first");
  });
}
void GetCommunicationControlFlowOps ( FuncOp func , const llvm :: SmallDenseMap < StringRef , FuncToRewrite > & funcs_to_rewrite , llvm :: SmallPtrSetImpl < Operation * > & control_flow_ops , llvm :: SmallPtrSetImpl < Block * > & control_flow_blocks ) { func . walk ( [ & ] ( Operation * op ) { if ( IsCommunicationOp ( op ) || IsFunctionCallWithCommunication ( op , funcs_to_rewrite ) ) if ( failed ( GetControlFlowAncestors ( op , control_flow_ops , control_flow_blocks ) ) ) llvm_unreachable ( " " " " ) ; } ) ; }
Collects all control flow op ancestors of communication ops or function calls
with communication ops (transitively).
Collects all control flow op ancestors of communication ops or function calls with communication ops ( transitively ) .