Method Run

tensorflow/core/nccl/nccl_rewrite.cc:246–270

class NcclReplacePass : public GraphOptimizationPass {
 public:
  Status Run(const GraphOptimizationPassOptions& options) override {
    if (options.graph == nullptr) {
      return Status::OK();
    }
    Graph* graph = options.graph->get();
    if (graph == nullptr) {
      return errors::Internal(
          "NCCL replacement should happen before partitioning and a "
          "graph should be available.");
    }
    // Find reduction and broadcast ops and replace them with Send/Recv ops.
    for (Node* node : graph->op_nodes()) {
      StringPiece type = node->type_string();
      if (!absl::StartsWith(type, "Nccl")) {
        continue;
      }
      if (type == "NcclReduce") {
        TF_RETURN_IF_ERROR(ReplaceReduce(graph, node));
      }
      if (type == "NcclBroadcast") {
        TF_RETURN_IF_ERROR(ReplaceBroadcast(graph, node));
      }
    }
    return Status::OK();
  }
};
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_PLACEMENT, 0,
                      NcclReplacePass);