abacus-aibox-842 fix some place

zmxdream · Nov 22, 2023 · da37757 · da37757
1 parent d54ca8b
commit da37757
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 18 deletions.
diff --git a/paddle/fluid/framework/fleet/box_wrapper_impl.h b/paddle/fluid/framework/fleet/box_wrapper_impl.h
@@ -740,22 +740,25 @@ void BoxWrapper::PushSparseGradCaseXPU(const paddle::platform::Place& place,
 //   }
   std::vector<int> slot_inner_offset(total_length);
   int out_count = 0;
-  for(int i=0;i<slot_num;i++) {
-      for(int j=0;j<slot_lengths[i];j++) {
-          slot_inner_offset[out_count++] = j;
-      }
+  for (int i = 0; i < slot_num; i++) {
+    for (int64_t j = 0; j < slot_lengths[i]; j++) {
+      slot_inner_offset[out_count++] = j;
+    }
   }
-  void *d_slot_inner_offset = nullptr;
-  xpu_malloc((void **)&d_slot_inner_offset, total_length * sizeof(int));
-  xpu_memcpy(d_slot_inner_offset, slot_inner_offset.data(), total_length * sizeof(int), XPU_HOST_TO_DEVICE);
+  auto d_slot_inner_offset_tmp = memory::Alloc(place, total_length * sizeof(int));
+  int* d_slot_inner_offset = reinterpret_cast<int*>(d_slot_inner_offset_tmp->ptr());
+  memory::Copy(place,
+               d_slot_inner_offset,
+               platform::CPUPlace(),
+               slot_inner_offset.data(),
+               total_length * sizeof(int));
 
   box_wrapper_kernel_->CopyForPush(place, xpu_values, total_grad_values_xpu,
       push_offset, total_length, slot_vector, (int*)d_slot_inner_offset, slot_lens, slot_num,
       hidden_size, batch_size, total_dims, skip_offset, key2slot,
       expand_embed_dim,
       push_float_num_,
       expand_only);
-  xpu_free(d_slot_inner_offset);
 
   push_boxps_timer.Resume();
 #ifdef TRACE_PROFILE

diff --git a/paddle/fluid/operators/pull_box_extended_sparse_op.h b/paddle/fluid/operators/pull_box_extended_sparse_op.h
@@ -65,7 +65,7 @@ static void PullBoxExtendedSparseFunctor(
     if(flags.empty()) {
       offset += outputs[i]->dims()[0];
     } else {
-      if(is_expand_slot_small==true){
+      if(is_expand_slot_small){
         if (flags[i] & 0x01) {
           offset += outputs[embedx_offset]->dims()[0];
           dims1 = outputs[embedx_offset]->dims()[1];
@@ -94,7 +94,7 @@ static void PullBoxExtendedSparseFunctor(
   }
 
   framework::LoDTensor total_values;
-  total_values.Resize(phi::make_ddim({max_total_dims0*(dims1+expand_dims1)}));
+  total_values.Resize(phi::make_ddim({max_total_dims0 * (dims1 + expand_dims1)}));
   total_values.mutable_data<T>(ctx.GetPlace());
 
   // BoxPS only supports float now
@@ -107,10 +107,10 @@ static void PullBoxExtendedSparseFunctor(
           reinterpret_cast<const uint64_t*>(slot->data<int64_t>());
       all_keys[i] = single_slot_keys;
       slot_lengths[i] = slot->numel();
-      if(outputs[embedx_offset]->numel()==0) {
+      if (outputs[embedx_offset]->numel() == 0) {
         outputs[embedx_offset]->set_layout(paddle::framework::DataLayout::UNDEFINED);
       } else {
-        int offset = slot_dims0_offset[i]*dims1* sizeof(T);
+        size_t offset = slot_dims0_offset[i] * dims1 * sizeof(T);
         total_values.set_offset(offset);
         outputs[i]->ShareBufferWith(total_values);
       }
@@ -119,8 +119,8 @@ static void PullBoxExtendedSparseFunctor(
       if(outputs_extend[expand_offset]->numel()==0) {
         outputs_extend[expand_offset]->set_layout(paddle::framework::DataLayout::UNDEFINED);
       } else {
-        int offset = slot_dims0_offset[i]*expand_dims1* sizeof(T);
-        total_values.set_offset(max_total_dims0*dims1* sizeof(T)+offset);
+        size_t offset = slot_dims0_offset[i] * expand_dims1 * sizeof(T);
+        total_values.set_offset(max_total_dims0 * dims1 * sizeof(T) + offset);
         outputs_extend[i]->ShareBufferWith(total_values);
       }
       auto *output_extend = outputs_extend[i]->mutable_data<T>(ctx.GetPlace());
@@ -136,10 +136,10 @@ static void PullBoxExtendedSparseFunctor(
       all_keys[i] = single_slot_keys;
       slot_lengths[i] = slot->numel();
       if (flags[i] & 0x01) {
-        if(outputs[embedx_offset]->numel()==0) {
+        if (outputs[embedx_offset]->numel() == 0) {
           outputs[embedx_offset]->set_layout(paddle::framework::DataLayout::UNDEFINED);
         } else {
-          int offset = slot_dims0_offset[i]*dims1* sizeof(T);
+          size_t offset = slot_dims0_offset[i] * dims1 * sizeof(T);
           total_values.set_offset(offset);
           outputs[embedx_offset]->ShareBufferWith(total_values);
         }
@@ -153,8 +153,8 @@ static void PullBoxExtendedSparseFunctor(
         if(outputs_extend[expand_offset]->numel()==0) {
           outputs_extend[expand_offset]->set_layout(paddle::framework::DataLayout::UNDEFINED);
         } else {
-          int offset = slot_dims0_offset[i]*expand_dims1* sizeof(T);
-          total_values.set_offset(max_total_dims0*dims1* sizeof(T)+offset);
+          size_t offset = slot_dims0_offset[i] * expand_dims1 * sizeof(T);
+          total_values.set_offset(max_total_dims0 * dims1 * sizeof(T) + offset);
           outputs_extend[expand_offset]->ShareBufferWith(total_values);
         }
         auto *output_extend = outputs_extend[expand_offset]->mutable_data<T>(ctx.GetPlace());