[GPUPS] opt Push merge (PaddlePaddle#50)

* add merge grad atomic kernel * push merge grad Co-authored-by: yaoxuefeng6 <[email protected]>
zmxdream · Aug 15, 2022 · 0b424ce · 0b424ce
1 parent 184c631
commit 0b424ce
Show file tree

Hide file tree

Showing 2 changed files with 246 additions and 156 deletions.
diff --git a/paddle/fluid/framework/fleet/heter_ps/heter_comm.h b/paddle/fluid/framework/fleet/heter_ps/heter_comm.h
@@ -47,6 +47,12 @@ struct CustomGradMerger {
     return out;
   }
 
+  template <typename GPUAccessor>
+  __device__ __forceinline__
+  void copy_all_field(float* output, const float* input, GPUAccessor& gpu_accessor) {
+      gpu_accessor.PushValueFill(output, input);
+  }
+
   template <typename GPUAccessor>
   __device__ __forceinline__
   void copy_basic_field(float* output, const float* input, GPUAccessor& gpu_accessor) {
@@ -241,6 +247,7 @@ class HeterComm {
   std::vector<std::shared_ptr<cub::CachingDeviceAllocator>> allocators_;
   int multi_mf_dim_{8};
   int max_mf_dim_ = 8;
+  int use_merge_atomic_ = 1;
 };
 
 }  // end namespace framework