// and associated tokens.
LogicalResult GetFunctionsToRewrite(
    ModuleOp module,
    llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs_to_rewrite) {
  // Find functions containing communication ops.
  SmallVector<func::FuncOp, 4> funcs_to_visit;
  for (func::FuncOp func : module.getOps<func::FuncOp>()) {
    FuncToRewrite func_to_rewrite{/*original=*/func, /*control_flow_ops=*/{},
                                  /*control_flow_blocks=*/{},