add some check

paddle/fluid/operators/distributed/parameter_recv.cc

@@ -80,7 +80,9 @@ void ParameterRecv<T>::operator()(const RpcContext &rpc_ctx,
     framework::Tensor *recv_tensor =
         recv_var->GetMutable<framework::LoDTensor>();
     auto dev_ctx = paddle::platform::CPUDeviceContext();
+    int64_t recv_numel = 0;
     for (auto *in : recved_tensors) {
+      recv_numel += in->numel();
       auto in_stride = framework::stride_numel(in->dims());
       auto out_stride = framework::stride_numel(recv_tensor->dims());
       StridedNumelCopyWithAxis<T>(
@@ -88,6 +90,7 @@ void ParameterRecv<T>::operator()(const RpcContext &rpc_ctx,
           in->data<T>(), in_stride, in_stride[0]);
       output_offset += in_stride[0];
     }
+    PADDLE_ENFORCE_EQ(recv_numel, recv_tensor->numel());
   }
 
   delete local_scope;

paddle/fluid/operators/distributed/parameter_send.cc

@@ -99,7 +99,7 @@ void ParameterSend<T>::operator()(const RpcContext &rpc_ctx,
 
     // split rows index into output sparse vars
     for (size_t i = 0; i < send_rows.size(); ++i) {
-      int out_idx = FindOutIdx(send_rows[i], abs_sections);
+      int out_idx = GetSectionIndex(send_rows[i], abs_sections);
       outs_rows_idx[out_idx].push_back(send_rows[i]);
       outs_dense_idx[out_idx].push_back(i);
     }

paddle/fluid/operators/distributed_ops/send_recv_util.h

@@ -48,16 +48,6 @@ inline bool NeedSend(const framework::Scope& scope,
   return false;
 }
 
-inline int FindOutIdx(int row, const std::vector<int64_t>& abs_sections) {
-  for (size_t i = 1; i < abs_sections.size(); ++i) {
-    if (row < abs_sections[i]) {
-      return i - 1;
-    }
-  }
-  PADDLE_ENFORCE_LT(row, abs_sections.back(), "row should be less then max id");
-  return abs_sections.size() - 1;
-}
-
 inline std::vector<int64_t> ToAbsoluteSection(
     const std::vector<int64_t>& height_sections) {
   std::vector<int64_t> abs_sections;

paddle/fluid/operators/split_selected_rows_op.h

@@ -32,7 +32,8 @@ class SplitSelectedRowsOpKernel : public framework::OpKernel<T> {
 
     auto abs_sections = ToAbsoluteSection(height_sections);
 
-    auto x_rows = x->rows();
+    auto& x_rows = x->rows();
+    auto height = x->height();
     std::vector<std::vector<int>> outs_rows_idx;
     std::vector<std::vector<int>> outs_dense_idx;
 
@@ -44,8 +45,10 @@ class SplitSelectedRowsOpKernel : public framework::OpKernel<T> {
 
     // split rows index into output sparse vars
     for (size_t i = 0; i < x_rows.size(); ++i) {
-      int out_idx = FindOutIdx(x_rows[i], abs_sections);
-      outs_rows_idx[out_idx].push_back(x_rows[i]);
+      auto& id = x_rows[i];
+      PADDLE_ENFORCE_LT(id, height);
+      int out_idx = GetSectionIndex(id, abs_sections);
+      outs_rows_idx[out_idx].push_back(id);
       outs_dense_idx[out_idx].push_back(i);
     }
     auto place = ctx.GetPlace();
@@ -59,7 +62,9 @@ class SplitSelectedRowsOpKernel : public framework::OpKernel<T> {
       outs[i]->mutable_rows()->clear();
       if (rows_idx.size() > 0) {
         for (auto idx : rows_idx) {
-          outs[i]->mutable_rows()->push_back(idx - abs_sections[i]);
+          auto id_offset = idx - abs_sections[i];
+          PADDLE_ENFORCE_LT(id_offset, height_sections[i]);
+          outs[i]->mutable_rows()->push_back(id_offset);
         }
         auto dst = outs[i]->mutable_value()->mutable_data<T>(ctx.GetPlace());
         for (size_t j = 0; j < rows_idx.size(); j++) {