llvm::SmallPtrSet<Block*, 4> control_flow_blocks;
  func::FuncOp clone;
};

// 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) {
  // Find functions containing communication ops.
  SmallVector<func::FuncOp, 4> funcs_to_visit;