fix slot sage hang (PaddlePaddle#300)

Author: DesmonDay

paddle/fluid/framework/data_feed.cu

@@ -1459,7 +1459,8 @@ int GraphDataGenerator::FillSlotFeature(uint64_t *d_walk, size_t key_num, int te
                                                d_feature_size_list_buf_,
                                                d_feature_size_prefixsum_buf_,
                                                d_feature_list,
-                                               d_slot_list);
+                                               d_slot_list,
+                                               conf_.sage_mode);
   // num of slot feature
   int slot_num = conf_.slot_num - float_slot_num_;
   int64_t *slot_tensor_ptr_[slot_num];

paddle/fluid/framework/fleet/heter_ps/graph_gpu_ps_table.h

@@ -205,7 +205,8 @@ class GpuPsGraphTable
       std::shared_ptr<phi::Allocation> &size_list,
       std::shared_ptr<phi::Allocation> &size_list_prefix_sum,
       std::shared_ptr<phi::Allocation> &feature_list,  // NOLINT
-      std::shared_ptr<phi::Allocation> &slot_list);    // NOLINT
+      std::shared_ptr<phi::Allocation> &slot_list,
+      bool sage_mode = false);    // NOLINT
   int get_float_feature_info_of_nodes(
       int gpu_id,
       uint64_t *d_nodes,
@@ -231,7 +232,8 @@ class GpuPsGraphTable
        std::shared_ptr<phi::Allocation> &size_list,
        std::shared_ptr<phi::Allocation> &size_list_prefix_sum,
        std::shared_ptr<phi::Allocation> &feature_list,
-       std::shared_ptr<phi::Allocation> &slot_list);
+       std::shared_ptr<phi::Allocation> &slot_list,
+       bool sage_mode = false);
 
 
   NodeQueryResult query_node_list(int gpu_id,

paddle/fluid/framework/fleet/heter_ps/graph_gpu_ps_table_inl.cu

@@ -3073,7 +3073,8 @@ int GpuPsGraphTable::get_feature_info_of_nodes(
     std::shared_ptr<phi::Allocation> &size_list,
     std::shared_ptr<phi::Allocation> &size_list_prefix_sum,
     std::shared_ptr<phi::Allocation>& feature_list,
-    std::shared_ptr<phi::Allocation>& slot_list) {
+    std::shared_ptr<phi::Allocation>& slot_list,
+    bool sage_mode) {
    if (node_num == 0) {
      return 0;
    }
@@ -3086,7 +3087,8 @@ int GpuPsGraphTable::get_feature_info_of_nodes(
      } else {
        if (FLAGS_enable_graph_multi_node_sampling) {
          all_fea_num = get_feature_info_of_nodes_all2all(gpu_id, d_nodes, node_num, size_list,
-                                           size_list_prefix_sum, feature_list, slot_list);
+                                           size_list_prefix_sum, feature_list, slot_list,
+                                           sage_mode);
        }
      }
    } else {
@@ -3104,7 +3106,8 @@ int GpuPsGraphTable::get_feature_info_of_nodes_all2all(
      std::shared_ptr<phi::Allocation> &size_list,
      std::shared_ptr<phi::Allocation> &size_list_prefix_sum,
      std::shared_ptr<phi::Allocation>& feature_list,
-     std::shared_ptr<phi::Allocation>& slot_list) {
+     std::shared_ptr<phi::Allocation>& slot_list,
+     bool sage_mode) {
    if (node_num == 0) {
      return 0;
    }
@@ -3241,7 +3244,7 @@ int GpuPsGraphTable::get_feature_info_of_nodes_all2all(
                                         reinterpret_cast<uint64_t*>(feature_list_ptr),
                                         reinterpret_cast<uint64_t*>(inter_feature_list_ptr),
                                         stream,
-                                        false,
+                                        sage_mode,
                                         true);
    VLOG(2) << "end send feature list";
 
@@ -3260,7 +3263,7 @@ int GpuPsGraphTable::get_feature_info_of_nodes_all2all(
                                         reinterpret_cast<uint8_t*>(slot_list_ptr),
                                         reinterpret_cast<uint8_t*>(inter_slot_list_ptr),
                                         stream,
-                                        false,
+                                        sage_mode,
                                         true);
    VLOG(2) << "end send slot list";
 

paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.cu

@@ -997,7 +997,8 @@ int GraphGpuWrapper::get_feature_info_of_nodes(
     std::shared_ptr<phi::Allocation> &size_list,
     std::shared_ptr<phi::Allocation> &size_list_prefix_sum,
     std::shared_ptr<phi::Allocation> &feature_list,
-    std::shared_ptr<phi::Allocation> &slot_list) {
+    std::shared_ptr<phi::Allocation> &slot_list,
+    bool sage_mode) {
   platform::CUDADeviceGuard guard(gpu_id);
   PADDLE_ENFORCE_NOT_NULL(graph_table,
                           paddle::platform::errors::InvalidArgument(
@@ -1009,7 +1010,8 @@ int GraphGpuWrapper::get_feature_info_of_nodes(
                                   size_list,
                                   size_list_prefix_sum,
                                   feature_list,
-                                  slot_list);
+                                  slot_list,
+                                  sage_mode);
 }
 
 int GraphGpuWrapper::get_float_feature_info_of_nodes(

paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.h

@@ -186,7 +186,8 @@ class GraphGpuWrapper {
       std::shared_ptr<phi::Allocation>& size_list,
       std::shared_ptr<phi::Allocation>& size_list_prefix_sum,
       std::shared_ptr<phi::Allocation>& feature_list,  // NOLINT
-      std::shared_ptr<phi::Allocation>& slot_list);    // NOLINT
+      std::shared_ptr<phi::Allocation>& slot_list,
+      bool sage_mode = false);    // NOLINT
   int get_float_feature_info_of_nodes(
       int gpu_id,
       uint64_t *d_nodes,