refine reduce op

paddle/fluid/framework/details/CMakeLists.txt

@@ -16,18 +16,16 @@ if(WITH_GPU)
     set(multi_devices_graph_builder_deps nccl_all_reduce_op_handle)
     nv_library(reduce_op_handle SRCS reduce_op_handle.cc DEPS op_handle_base variable_visitor scope ddim dynload_cuda)
     nv_library(broadcast_op_handle SRCS broadcast_op_handle.cc DEPS op_handle_base scope ddim memory variable_visitor dynload_cuda)
-    nv_library(reduce_blockop_handle SRCS reduce_blockop_handle.cc DEPS op_handle_base scope ddim memory variable_visitor dynload_cuda)
 else()
     set(multi_devices_graph_builder_deps)
     cc_library(reduce_op_handle SRCS reduce_op_handle.cc DEPS op_handle_base variable_visitor scope ddim)
     cc_library(broadcast_op_handle SRCS broadcast_op_handle.cc DEPS op_handle_base scope ddim memory variable_visitor)
-    cc_library(reduce_blockop_handle SRCS reduce_blockop_handle.cc DEPS op_handle_base scope ddim memory variable_visitor)
 endif()
 
 cc_library(gather_op_handle SRCS gather_op_handle.cc DEPS op_handle_base scope ddim memory variable_visitor)
 
 cc_library(multi_devices_graph_builder SRCS multi_devices_graph_builder.cc DEPS ssa_graph_builder computation_op_handle
-        scale_loss_grad_op_handle send_op_handle ${multi_devices_graph_builder_deps} reduce_op_handle broadcast_op_handle reduce_blockop_handle)
+        scale_loss_grad_op_handle send_op_handle ${multi_devices_graph_builder_deps} reduce_op_handle broadcast_op_handle)
 
 cc_library(ssa_graph_executor SRCS ssa_graph_executor.cc DEPS ssa_graph framework_proto)
 cc_library(threaded_ssa_graph_executor SRCS threaded_ssa_graph_executor.cc DEPS fetch_op_handle ssa_graph_executor scope

paddle/fluid/framework/details/multi_devices_graph_builder.cc

@@ -15,7 +15,6 @@
 #include <utility>
 #include "paddle/fluid/framework/details/broadcast_op_handle.h"
 #include "paddle/fluid/framework/details/computation_op_handle.h"
-#include "paddle/fluid/framework/details/reduce_blockop_handle.h"
 #include "paddle/fluid/framework/details/reduce_op_handle.h"
 #include "paddle/fluid/framework/details/scale_loss_grad_op_handle.h"
 #include "paddle/fluid/framework/details/send_op_handle.h"
@@ -278,7 +277,7 @@ void MultiDevSSAGraphBuilder::FuseReduceOpHandles(
 
     auto reduce_var_name = string::Sprintf("REDUCEBLOCK_DATA_%d", dev_id);
     for (size_t k = 0; k < this->places_.size(); ++k) {
-      // Allocate gradients space
+      // Allocate gradients memory
       auto reduce_t = this->local_scopes_.at(k)
                           ->Var(reduce_var_name)
                           ->GetMutable<LoDTensor>();
@@ -335,15 +334,16 @@ void MultiDevSSAGraphBuilder::RemoveOps(
 }
 
 void MultiDevSSAGraphBuilder::CreateReduceBlockOp(
-    SSAGraph *result, const int root_id, const std::string &reduce_var_name,
+    SSAGraph *result, const int dst_scope_id,
+    const std::string &reduce_var_name,
     const std::unordered_set<VarHandle *> &inputs,
     const std::unordered_set<VarHandle *> &outputs) const {
 #ifdef PADDLE_WITH_CUDA
-  result->ops_.emplace_back(new ReduceBlockOpHandle(
-      local_scopes_, places_, reduce_var_name, root_id, nccl_ctxs_));
+  result->ops_.emplace_back(new ReduceOpHandle(
+      local_scopes_, places_, nccl_ctxs_, dst_scope_id, reduce_var_name));
 #else
-  result->ops_.emplace_back(new ReduceBlockOpHandle(local_scopes_, places_,
-                                                    reduce_var_name, root_id));
+  result->ops_.emplace_back(new ReduceOpHandle(local_scopes_, places_,
+                                               dst_scope_id, reduce_var_name));
 #endif
   auto *op_handle = result->ops_.back().get();