selected_rows_functor.cc 31.3 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/math/selected_rows_functor.h"
16
#include "paddle/fluid/platform/device/device_wrapper.h"
17

L
lidanqing 已提交
18 19 20 21
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/operators/mkldnn/axpy_handler.h"
#endif

22 23 24 25
namespace paddle {
namespace operators {
namespace math {
template <typename T>
Q
QI JUN 已提交
26 27
struct SelectedRowsAdd<platform::CPUDeviceContext, T> {
  void operator()(const platform::CPUDeviceContext& context,
28 29
                  const phi::SelectedRows& input1,
                  const phi::SelectedRows& input2, phi::SelectedRows* output) {
30
    auto in1_height = input1.height();
31 32 33 34 35 36
    PADDLE_ENFORCE_EQ(
        in1_height, input2.height(),
        platform::errors::InvalidArgument("The two inputs height must be equal."
                                          "But recieved first input height  = "
                                          "[%d], second input height = [%d]",
                                          in1_height, input2.height()));
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53
    output->set_height(in1_height);

    auto& in1_rows = input1.rows();
    auto& in2_rows = input2.rows();
    std::vector<int64_t> out_rows;
    out_rows.reserve(in1_rows.size() + in2_rows.size());

    // concat rows
    out_rows.insert(out_rows.end(), in1_rows.begin(), in1_rows.end());
    out_rows.insert(out_rows.end(), in2_rows.begin(), in2_rows.end());
    output->set_rows(out_rows);

    auto* out_value = output->mutable_value();
    auto& in1_value = input1.value();
    auto& in2_value = input2.value();

    auto in1_row_numel = in1_value.numel() / in1_rows.size();
54 55 56 57 58 59 60 61 62 63 64 65
    PADDLE_ENFORCE_EQ(
        in1_row_numel, in2_value.numel() / in2_rows.size(),
        platform::errors::InvalidArgument(
            "The two inputs width must be equal."
            "But recieved first input width = [%d], second input width = [%d]",
            in1_row_numel, in2_value.numel() / in2_rows.size()));
    PADDLE_ENFORCE_EQ(
        in1_row_numel, out_value->numel() / out_rows.size(),
        platform::errors::InvalidArgument(
            "The input and oupput width must be equal."
            "But recieved input width = [%d], output width = [%d]",
            in1_row_numel, out_value->numel() / out_rows.size()));
66 67

    auto in1_place = input1.place();
68 69 70
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(in1_place), true,
                      platform::errors::InvalidArgument(
                          "The running enviroment is not on the CPU place."));
71
    auto in2_place = input2.place();
72 73 74
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(in2_place), true,
                      platform::errors::InvalidArgument(
                          "The running enviroment is not on the CPU place."));
75
    auto out_place = context.GetPlace();
76 77 78
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(out_place), true,
                      platform::errors::InvalidArgument(
                          "The running enviroment is not on the CPU place."));
79 80 81

    auto* out_data = out_value->data<T>();
    auto* in1_data = in1_value.data<T>();
82
    memory::Copy(out_place, out_data, in1_place, in1_data,
83 84 85
                 in1_value.numel() * sizeof(T));

    auto* in2_data = in2_value.data<T>();
86
    memory::Copy(out_place, out_data + in1_value.numel(), in2_place, in2_data,
87 88 89 90
                 in2_value.numel() * sizeof(T));
  }
};

Q
QI JUN 已提交
91 92
template struct SelectedRowsAdd<platform::CPUDeviceContext, float>;
template struct SelectedRowsAdd<platform::CPUDeviceContext, double>;
93 94

template <typename T>
Q
QI JUN 已提交
95 96
struct SelectedRowsAddTensor<platform::CPUDeviceContext, T> {
  void operator()(const platform::CPUDeviceContext& context,
97
                  const phi::SelectedRows& input1,
98 99 100 101
                  const framework::Tensor& input2, framework::Tensor* output) {
    auto in1_height = input1.height();
    auto in2_dims = input2.dims();
    auto out_dims = output->dims();
102 103 104 105 106 107 108 109 110 111 112 113
    PADDLE_ENFORCE_EQ(
        in1_height, in2_dims[0],
        platform::errors::InvalidArgument("The two inputs height must be equal."
                                          "But recieved first input height = "
                                          "[%d], second input height = [%d]",
                                          in1_height, in2_dims[0]));
    PADDLE_ENFORCE_EQ(
        in1_height, out_dims[0],
        platform::errors::InvalidArgument(
            "The input and output height must be equal."
            "But recieved input height = [%d], output height = [%d]",
            in1_height, out_dims[0]));
114 115 116 117 118

    auto& in1_value = input1.value();
    auto& in1_rows = input1.rows();

    int64_t in1_row_numel = in1_value.numel() / in1_rows.size();
119 120 121 122 123 124 125 126 127 128 129 130
    PADDLE_ENFORCE_EQ(
        in1_row_numel, input2.numel() / in1_height,
        platform::errors::InvalidArgument(
            "The two inputs width must be equal."
            "But recieved first input width = [%d], second input width = [%d]",
            in1_row_numel, input2.numel() / in1_height));
    PADDLE_ENFORCE_EQ(
        in1_row_numel, output->numel() / in1_height,
        platform::errors::InvalidArgument(
            "The input and output width must be equal."
            "But recieved input width = [%d], output width = [%d]",
            in1_row_numel, output->numel() / in1_height));
131

132
    phi::funcs::SetConstant<platform::CPUDeviceContext, T> functor;
133 134 135 136 137 138 139 140 141 142 143 144 145 146
    functor(context, output, 0.0);

    auto* in1_data = in1_value.data<T>();
    auto* out_data = output->data<T>();

    for (size_t i = 0; i < in1_rows.size(); i++) {
      for (int64_t j = 0; j < in1_row_numel; j++) {
        out_data[in1_rows[i] * in1_row_numel + j] +=
            in1_data[i * in1_row_numel + j];
      }
    }

    auto out_eigen = framework::EigenVector<T>::Flatten(*output);
    auto in2_eigen = framework::EigenVector<T>::Flatten(input2);
Q
QI JUN 已提交
147
    out_eigen.device(*context.eigen_device()) = out_eigen + in2_eigen;
148 149 150
  }
};

Q
QI JUN 已提交
151 152
template struct SelectedRowsAddTensor<platform::CPUDeviceContext, float>;
template struct SelectedRowsAddTensor<platform::CPUDeviceContext, double>;
Q
QI JUN 已提交
153 154

template <typename T>
Q
QI JUN 已提交
155 156
struct SelectedRowsAddTo<platform::CPUDeviceContext, T> {
  void operator()(const platform::CPUDeviceContext& context,
157 158
                  const phi::SelectedRows& input1, const int64_t input2_offset,
                  phi::SelectedRows* input2) {
Q
QI JUN 已提交
159
    auto in1_height = input1.height();
160 161 162 163 164 165
    PADDLE_ENFORCE_EQ(
        in1_height, input2->height(),
        platform::errors::InvalidArgument("The two inputs height must be equal."
                                          "But recieved first input height = "
                                          "[%d], second input height = [%d]",
                                          in1_height, input2->height()));
Q
QI JUN 已提交
166 167 168 169 170 171 172 173

    auto& in1_rows = input1.rows();
    auto& in2_rows = *(input2->mutable_rows());

    auto& in1_value = input1.value();
    auto* in2_value = input2->mutable_value();

    // concat rows
174 175
    paddle::framework::MixVector<int64_t> mixv_in2_rows(&in2_rows);
    mixv_in2_rows.Extend(in1_rows.begin(), in1_rows.end());
Q
QI JUN 已提交
176 177

    auto in1_place = input1.place();
178 179 180
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(in1_place), true,
                      platform::errors::InvalidArgument(
                          "The running enviroment is not on the CPU place."));
Q
QI JUN 已提交
181
    auto in2_place = input2->place();
182 183 184
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(in2_place), true,
                      platform::errors::InvalidArgument(
                          "The running enviroment is not on the CPU place."));
Q
QI JUN 已提交
185 186 187

    auto* in1_data = in1_value.data<T>();
    auto* in2_data = in2_value->data<T>();
188
    memory::Copy(in2_place, in2_data + input2_offset, in1_place, in1_data,
Q
QI JUN 已提交
189 190 191 192
                 in1_value.numel() * sizeof(T));
  }
};

Q
QI JUN 已提交
193 194 195 196
template struct SelectedRowsAddTo<platform::CPUDeviceContext, float>;
template struct SelectedRowsAddTo<platform::CPUDeviceContext, double>;
template struct SelectedRowsAddTo<platform::CPUDeviceContext, int>;
template struct SelectedRowsAddTo<platform::CPUDeviceContext, int64_t>;
Q
QI JUN 已提交
197

M
minqiyang 已提交
198 199 200
template <typename T>
struct SelectedRowsSumTo<platform::CPUDeviceContext, T> {
  void operator()(const platform::CPUDeviceContext& context,
201
                  const std::vector<phi::SelectedRows*>& input1,
M
minqiyang 已提交
202
                  const std::vector<int64_t>& input2_offsets,
203
                  phi::SelectedRows* input2) {
M
minqiyang 已提交
204 205 206 207 208 209
    // Ensure all selected rows have the same height
    size_t size = 0u;
    for (auto iter = input1.begin(); iter != input1.end(); ++iter) {
      auto& in_rows = (*iter)->rows();
      size += in_rows.end() - in_rows.begin();
      auto in1_height = (*iter)->height();
210 211 212 213 214 215
      PADDLE_ENFORCE_EQ(in1_height, input2->height(),
                        platform::errors::InvalidArgument(
                            "The two inputs height must be equal."
                            "But recieved first input height = [%d], second "
                            "input height = [%d]",
                            in1_height, input2->height()));
M
minqiyang 已提交
216 217 218 219 220 221 222 223 224 225 226 227
    }
    // concat rows
    std::vector<int64_t> in2_rows;
    in2_rows.reserve(in2_rows.size() + size);
    for (auto iter = input1.begin(); iter != input1.end(); ++iter) {
      const framework::Vector<int64_t>& in_rows = (*iter)->rows();
      in2_rows.insert(in2_rows.end(), in_rows.begin(), in_rows.end());
    }
    input2->set_rows(in2_rows);

    auto* in2_value = input2->mutable_value();
    auto* in2_data = in2_value->data<T>();
228
    auto blas = phi::funcs::GetBlas<platform::CPUDeviceContext, T>(context);
M
minqiyang 已提交
229 230 231 232 233 234 235 236 237 238 239 240 241
    size_t offset = 0u;
    for (size_t i = 0u; i != input1.size(); ++i) {
      auto& in_value = input1[i]->value();
      const auto* in_data = in_value.data<T>();
      offset += input2_offsets[i];
      blas.VCOPY(in_value.numel(), in_data, in2_data + offset);
    }
  }
};

template struct SelectedRowsSumTo<platform::CPUDeviceContext, float>;
template struct SelectedRowsSumTo<platform::CPUDeviceContext, double>;

Q
QI JUN 已提交
242
template <typename T>
Q
QI JUN 已提交
243 244
struct SelectedRowsAddToTensor<platform::CPUDeviceContext, T> {
  void operator()(const platform::CPUDeviceContext& context,
245
                  const phi::SelectedRows& input1, framework::Tensor* input2) {
Q
Qiao Longfei 已提交
246
    if (UNLIKELY(input1.rows().size() == 0)) {
247 248 249
      LOG(WARNING) << "input selected rows is empty!";
      return;
    }
Q
QI JUN 已提交
250 251
    auto in1_height = input1.height();
    auto in2_dims = input2->dims();
252 253 254 255 256 257
    PADDLE_ENFORCE_EQ(
        in1_height, in2_dims[0],
        platform::errors::InvalidArgument("The two inputs height must be equal."
                                          "But recieved first input height = "
                                          "[%d], second input height = [%d]",
                                          in1_height, in2_dims[0]));
Q
QI JUN 已提交
258 259 260 261 262

    auto& in1_value = input1.value();
    auto& in1_rows = input1.rows();

    int64_t in1_row_numel = in1_value.numel() / in1_rows.size();
263 264 265 266 267 268
    PADDLE_ENFORCE_EQ(
        in1_row_numel, input2->numel() / in1_height,
        platform::errors::InvalidArgument(
            "The two inputs width must be equal."
            "But recieved first input width = [%d], second input width = [%d]",
            in1_row_numel, input2->numel() / in1_height));
Q
QI JUN 已提交
269 270 271 272 273 274 275 276 277 278 279 280 281

    auto* in1_data = in1_value.data<T>();
    auto* input2_data = input2->data<T>();

    for (size_t i = 0; i < in1_rows.size(); i++) {
      for (int64_t j = 0; j < in1_row_numel; j++) {
        input2_data[in1_rows[i] * in1_row_numel + j] +=
            in1_data[i * in1_row_numel + j];
      }
    }
  }
};

Q
QI JUN 已提交
282 283 284 285
template struct SelectedRowsAddToTensor<platform::CPUDeviceContext, float>;
template struct SelectedRowsAddToTensor<platform::CPUDeviceContext, double>;
template struct SelectedRowsAddToTensor<platform::CPUDeviceContext, int>;
template struct SelectedRowsAddToTensor<platform::CPUDeviceContext, int64_t>;
286 287
template struct SelectedRowsAddToTensor<platform::CPUDeviceContext,
                                        platform::bfloat16>;
288

T
typhoonzero 已提交
289 290 291 292 293 294 295 296
// This is a separated namespace for manipulate SelectedRows typed
// data. Like merge duplicated rows, adding two SelectedRows etc.
//
// Another group of functors is called "scatter updates", which means
// use SelectedRows to update a dense tensor with different Ops, like
// add or mul.
namespace scatter {

L
lidanqing 已提交
297
template <typename T>
298
typename std::enable_if<!std::is_integral<T>::value>::type elementwise_add_to(
299
    phi::funcs::BlasT<platform::CPUDeviceContext, T>* blas, size_t data_len,
300
    const T* in, T* out) {
301
  blas->AXPY(data_len, T(1.f), in, out);
Q
Qiao Longfei 已提交
302 303
}

304 305
template <typename T>
typename std::enable_if<std::is_integral<T>::value>::type elementwise_add_to(
306
    phi::funcs::BlasT<platform::CPUDeviceContext, T>* blas, size_t data_len,
307
    const T* in, T* out) {
T
Tao Luo 已提交
308
  for (size_t i = 0; i < data_len; i++) {
Q
Qiao Longfei 已提交
309 310
    out[i] += in[i];
  }
T
typhoonzero 已提交
311 312
}

313 314
template <typename T>
typename std::enable_if<std::is_same<T, platform::bfloat16>::value>::type
315
add_sparse_inputs(const std::vector<const phi::SelectedRows*>& inputs,
316 317 318 319
                  const std::unordered_map<int64_t, size_t>& rows_to_id,
                  int64_t input_width,
                  const platform::CPUDeviceContext& context, T* out_data) {
#ifndef PADDLE_WITH_MKLDNN
320
  auto blas = phi::funcs::GetBlas<platform::CPUDeviceContext, T>(context);
321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348
#endif
  for (auto* input : inputs) {
    if (input->rows().size() == 0) {
      continue;
    }
    auto* input_data = input->value().data<T>();
    auto& input_rows = input->rows();

#ifdef PADDLE_WITH_MKLDNN
    OneDNNAXPYHandler<T> axpy_handler(input_width, T(1.f));
    for (size_t i = 0; i < input_rows.size(); i++) {
      size_t out_i = rows_to_id.at(input_rows[i]);
      axpy_handler(&input_data[i * input_width],
                   &out_data[out_i * input_width]);
    }
#else
    for (size_t i = 0; i < input_rows.size(); i++) {
      size_t out_i = rows_to_id.at(input_rows[i]);
      elementwise_add_to<T>(&blas, static_cast<size_t>(input_width),
                            &input_data[i * input_width],
                            &out_data[out_i * input_width]);
    }
#endif
  }
}

template <typename T>
typename std::enable_if<!std::is_same<T, platform::bfloat16>::value>::type
349
add_sparse_inputs(const std::vector<const phi::SelectedRows*>& inputs,
350 351 352 353
                  const std::unordered_map<int64_t, size_t>& rows_to_id,
                  int64_t input_width,
                  const platform::CPUDeviceContext& context, T* out_data) {
  VLOG(4) << "[CPU] add_sparse_inputs <" << typeid(T).name();
354
  auto blas = phi::funcs::GetBlas<platform::CPUDeviceContext, T>(context);
355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370
  for (auto* input : inputs) {
    if (input->rows().size() == 0) {
      continue;
    }
    auto* input_data = input->value().data<T>();
    auto& input_rows = input->rows();

    for (size_t i = 0; i < input_rows.size(); i++) {
      size_t out_i = rows_to_id.at(input_rows[i]);
      elementwise_add_to<T>(&blas, static_cast<size_t>(input_width),
                            &input_data[i * input_width],
                            &out_data[out_i * input_width]);
    }
  }
}

T
typhoonzero 已提交
371 372
template <typename T>
struct MergeAdd<platform::CPUDeviceContext, T> {
373 374 375 376
  phi::SelectedRows operator()(const platform::CPUDeviceContext& context,
                               const phi::SelectedRows& input,
                               const bool sorted_result = false) {
    phi::SelectedRows out;
377
    (*this)(context, input, &out, sorted_result);
S
sneaxiy 已提交
378 379 380 381
    return out;
  }

  void operator()(const platform::CPUDeviceContext& context,
382
                  const phi::SelectedRows& input, phi::SelectedRows* output,
383
                  const bool sorted_result = false) {
384
    std::vector<const phi::SelectedRows*> inputs;
385
    inputs.push_back(&input);
386
    (*this)(context, inputs, output, sorted_result);
387
  }
T
typhoonzero 已提交
388

389
  void operator()(const platform::CPUDeviceContext& context,
390 391
                  const std::vector<const phi::SelectedRows*>& inputs,
                  phi::SelectedRows* output, const bool sorted_result = false) {
Q
Qiao Longfei 已提交
392
    if (inputs.size() == 0) {
M
minqiyang 已提交
393
      VLOG(3) << "no input! return";
Q
Qiao Longfei 已提交
394 395
      return;
    }
396
    const phi::SelectedRows* has_value_input = nullptr;
Q
Qiao Longfei 已提交
397
    for (auto* in : inputs) {
Q
Qiao Longfei 已提交
398
      if (in->rows().size() > 0) {
Q
Qiao Longfei 已提交
399 400 401 402 403
        has_value_input = in;
        break;
      }
    }
    if (has_value_input == nullptr) {
M
minqiyang 已提交
404
      VLOG(3) << "no input has value! just return" << std::endl;
Q
Qiao Longfei 已提交
405 406 407 408
      return;
    }
    auto input_width = has_value_input->value().dims()[1];
    auto input_height = has_value_input->height();
409
    phi::SelectedRows& out = *output;
410
    std::set<int64_t> merged_row_set;
411
    size_t row_num = 0;
412
    for (auto* input : inputs) {
Q
Qiao Longfei 已提交
413
      if (input->rows().size() == 0) {
Q
Qiao Longfei 已提交
414 415
        continue;
      }
416
      PADDLE_ENFORCE_EQ(input_width, input->value().dims()[1],
417 418 419
                        platform::errors::InvalidArgument(
                            "All inputs should have same "
                            "dimension except for the first one."));
420
      PADDLE_ENFORCE_EQ(input_height, input->height(),
421 422
                        platform::errors::InvalidArgument(
                            "All inputs should have same height."));
423
      row_num += input->rows().size();
424 425
      merged_row_set.insert(input->rows().begin(), input->rows().end());
    }
426

427
    out.set_height(input_height);
T
wip  
typhoonzero 已提交
428
    out.mutable_value()->mutable_data<T>(
429
        phi::make_ddim(
430
            {static_cast<int64_t>(merged_row_set.size()), input_width}),
T
typhoonzero 已提交
431
        context.GetPlace());
432
    auto* out_data = out.mutable_value()->data<T>();
T
typhoonzero 已提交
433

434 435 436 437 438 439 440 441 442 443 444 445 446 447 448
    if (merged_row_set.size() == row_num && !sorted_result) {
      // no duplicated ids, just concat the result together
      std::vector<int64_t> merge_rows;
      merge_rows.reserve(row_num);
      // concat rows
      for (auto* in : inputs) {
        merge_rows.insert(merge_rows.end(), in->rows().begin(),
                          in->rows().end());
      }
      out.set_rows(merge_rows);
      auto in_place = inputs[0]->place();
      auto out_place = out.place();
      int64_t copied_numel = 0;
      for (auto* in : inputs) {
        auto* in_data = in->value().data<T>();
449
        auto in_numel = in->rows().size() * input_width;
450
        memory::Copy(out_place, out_data + copied_numel, in_place, in_data,
451 452 453 454 455 456
                     in_numel * sizeof(T));
        copied_numel += in_numel;
      }
    } else {
      std::vector<int64_t> merge_rows(merged_row_set.begin(),
                                      merged_row_set.end());
T
typhoonzero 已提交
457

458 459 460
      if (sorted_result) {
        std::sort(merge_rows.begin(), merge_rows.end());
      }
T
typhoonzero 已提交
461

462 463
      out.set_rows(merge_rows);

464
      phi::funcs::SetConstant<platform::CPUDeviceContext, T> constant_functor;
465
      constant_functor(context, out.mutable_value(), static_cast<T>(0.f));
466 467 468 469

      std::unordered_map<int64_t, size_t> rows_to_id;
      for (size_t i = 0; i < merge_rows.size(); ++i) {
        rows_to_id[merge_rows[i]] = i;
Q
Qiao Longfei 已提交
470
      }
471

472
      add_sparse_inputs<T>(inputs, rows_to_id, input_width, context, out_data);
T
typhoonzero 已提交
473
    }
T
wip  
typhoonzero 已提交
474 475 476
  }
};

477 478 479
#ifdef PADDLE_WITH_XPU
template <typename T>
struct MergeAdd<platform::XPUDeviceContext, T> {
480 481 482 483
  phi::SelectedRows operator()(const platform::XPUDeviceContext& context,
                               const phi::SelectedRows& input,
                               const bool sorted_result = false) {
    phi::SelectedRows out;
484 485 486 487 488
    (*this)(context, input, &out, sorted_result);
    return out;
  }

  void operator()(const platform::XPUDeviceContext& context,
489
                  const phi::SelectedRows& input, phi::SelectedRows* output,
490 491 492 493 494 495
                  const bool sorted_result = false) {
    framework::Vector<int64_t> input_rows(input.rows());
    if (input_rows.size() == 0) {
      return;
    }

496
    phi::SelectedRows& out = *output;
497 498 499 500 501 502 503
    std::set<int64_t> row_set(input_rows.begin(), input_rows.end());
    std::vector<int64_t> merge_rows(row_set.begin(), row_set.end());
    auto input_width = input.value().dims()[1];

    out.set_rows(merge_rows);
    out.set_height(input.height());
    out.mutable_value()->mutable_data<T>(
504
        phi::make_ddim({static_cast<int64_t>(merge_rows.size()), input_width}),
505 506 507 508 509 510 511
        context.GetPlace());

    std::unordered_map<int64_t, size_t> rows_to_id;
    for (size_t i = 0; i < merge_rows.size(); ++i) {
      rows_to_id[merge_rows[i]] = i;
    }

512 513 514 515
    auto* y_data = out.mutable_value()->data<T>();
    auto* x_data = input.value().data<T>();
    int xm = input_rows.size();
    int ym = merge_rows.size();
516
    int n = input_width;
517 518 519 520 521 522 523 524 525 526 527 528

    xpu::ctx_guard RAII_GUARD(context.x_context());
    int64_t* x_rows_data = RAII_GUARD.alloc_l3_or_gm<int64_t>(xm);
    int64_t* y_rows_data = RAII_GUARD.alloc_l3_or_gm<int64_t>(ym);
    memory::Copy(context.GetPlace(), y_rows_data, platform::CPUPlace(),
                 merge_rows.data(), ym * sizeof(int64_t));
    memory::Copy(context.GetPlace(), x_rows_data, platform::CPUPlace(),
                 input_rows.data(), xm * sizeof(int64_t));
    int r =
        xpu::merge_dup_rows<T, int64_t>(context.x_context(), x_data, y_data,
                                        x_rows_data, y_rows_data, xm, n, ym);
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "merge_dup_rows");
529 530 531
  }

  void operator()(const platform::XPUDeviceContext& context,
532 533
                  const std::vector<const phi::SelectedRows*>& inputs,
                  phi::SelectedRows* output, const bool sorted_result = false) {
534 535 536 537
    if (inputs.size() == 0) {
      VLOG(3) << "no input! return";
      return;
    }
538
    const phi::SelectedRows* has_value_input = nullptr;
539 540 541 542 543 544 545 546 547 548 549 550
    for (auto* in : inputs) {
      if (in->rows().size() > 0) {
        has_value_input = in;
        break;
      }
    }
    if (has_value_input == nullptr) {
      VLOG(3) << "no input has value! just return" << std::endl;
      return;
    }
    auto input_width = has_value_input->value().dims()[1];
    auto input_height = has_value_input->height();
551
    phi::SelectedRows& out = *output;
552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578
    std::set<int64_t> merged_row_set;
    size_t row_num = 0;
    for (auto* input : inputs) {
      if (input->rows().size() == 0) {
        continue;
      }
      PADDLE_ENFORCE_EQ(input_width, input->value().dims()[1],
                        platform::errors::InvalidArgument(
                            "All inputs should have same "
                            "dimension except for the first one."));
      PADDLE_ENFORCE_EQ(input_height, input->height(),
                        platform::errors::InvalidArgument(
                            "All inputs should have same height."));
      row_num += input->rows().size();
      merged_row_set.insert(input->rows().begin(), input->rows().end());
    }

    std::vector<int64_t> merge_rows(merged_row_set.begin(),
                                    merged_row_set.end());

    if (sorted_result) {
      std::sort(merge_rows.begin(), merge_rows.end());
    }

    out.set_rows(merge_rows);
    out.set_height(input_height);
    out.mutable_value()->mutable_data<T>(
579
        phi::make_ddim(
580 581 582
            {static_cast<int64_t>(merged_row_set.size()), input_width}),
        context.GetPlace());

583
    float* y_data = reinterpret_cast<float*>(out.mutable_value()->data<T>());
584 585 586 587 588 589 590 591 592 593 594 595

    std::unordered_map<int64_t, size_t> rows_to_id;
    for (size_t i = 0; i < merge_rows.size(); ++i) {
      rows_to_id[merge_rows[i]] = i;
    }

    for (auto* input : inputs) {
      if (input->rows().size() == 0) {
        continue;
      }
      auto& input_rows = input->rows();

596 597 598
      auto* x_data = input->value().data<T>();
      int xm = input_rows.size();
      int ym = merge_rows.size();
599
      int n = input_width;
600 601 602 603 604 605 606 607 608 609 610 611

      xpu::ctx_guard RAII_GUARD(context.x_context());
      int64_t* x_rows_data = RAII_GUARD.alloc_l3_or_gm<int64_t>(xm);
      int64_t* y_rows_data = RAII_GUARD.alloc_l3_or_gm<int64_t>(ym);
      memory::Copy(context.GetPlace(), y_rows_data, platform::CPUPlace(),
                   merge_rows.data(), ym * sizeof(int64_t));
      memory::Copy(context.GetPlace(), x_rows_data, platform::CPUPlace(),
                   input_rows.data(), xm * sizeof(int64_t));
      int r =
          xpu::merge_dup_rows<T, int64_t>(context.x_context(), x_data, y_data,
                                          x_rows_data, y_rows_data, xm, n, ym);
      PADDLE_ENFORCE_XDNN_SUCCESS(r, "merge_dup_rows");
612 613 614 615 616
    }
  }
};

#endif
617 618
template <typename T>
struct MergeAverage<platform::CPUDeviceContext, T> {
619 620 621
  phi::SelectedRows operator()(const platform::CPUDeviceContext& context,
                               const phi::SelectedRows& input) {
    phi::SelectedRows out;
622 623 624 625 626
    (*this)(context, input, &out);
    return out;
  }

  void operator()(const platform::CPUDeviceContext& context,
627 628
                  const phi::SelectedRows& input, phi::SelectedRows* output) {
    std::vector<const phi::SelectedRows*> inputs;
629 630 631 632 633
    inputs.push_back(&input);
    (*this)(context, inputs, output);
  }

  void operator()(const platform::CPUDeviceContext& context,
634 635
                  const std::vector<const phi::SelectedRows*>& inputs,
                  phi::SelectedRows* output) {
636 637 638 639
    if (inputs.size() == 0) {
      VLOG(3) << "no input! return";
      return;
    }
640
    const phi::SelectedRows* has_value_input = nullptr;
641 642 643 644 645 646 647 648 649 650 651 652
    for (auto* in : inputs) {
      if (in->rows().size() > 0) {
        has_value_input = in;
        break;
      }
    }
    if (has_value_input == nullptr) {
      VLOG(3) << "no input has value! just return" << std::endl;
      return;
    }
    auto input_width = has_value_input->value().dims()[1];
    auto input_height = has_value_input->height();
653
    phi::SelectedRows& out = *output;
654 655 656 657 658 659 660
    std::set<int64_t> merged_row_set;
    size_t row_num = 0;
    for (auto* input : inputs) {
      if (input->rows().size() == 0) {
        continue;
      }
      PADDLE_ENFORCE_EQ(input_width, input->value().dims()[1],
661 662 663
                        platform::errors::InvalidArgument(
                            "All inputs should have same "
                            "dimension except for the first one."));
664
      PADDLE_ENFORCE_EQ(input_height, input->height(),
665 666
                        platform::errors::InvalidArgument(
                            "All input should have same height."));
667 668 669 670 671 672
      row_num += input->rows().size();
      merged_row_set.insert(input->rows().begin(), input->rows().end());
    }

    out.set_height(input_height);
    out.mutable_value()->mutable_data<T>(
673
        phi::make_ddim(
674 675 676 677 678 679 680 681 682 683
            {static_cast<int64_t>(merged_row_set.size()), input_width}),
        context.GetPlace());
    auto* out_data = out.mutable_value()->data<T>();

    std::vector<int64_t> merge_rows(merged_row_set.begin(),
                                    merged_row_set.end());
    std::sort(merge_rows.begin(), merge_rows.end());

    out.set_rows(merge_rows);

684
    phi::funcs::SetConstant<platform::CPUDeviceContext, T> constant_functor;
685 686 687 688 689 690 691
    constant_functor(context, out.mutable_value(), 0.0);

    std::unordered_map<int64_t, size_t> rows_to_id;
    for (size_t i = 0; i < merge_rows.size(); ++i) {
      rows_to_id[merge_rows[i]] = i;
    }

692
    auto blas = phi::funcs::GetBlas<platform::CPUDeviceContext, T>(context);
693 694 695 696 697 698 699 700 701
    for (auto* input : inputs) {
      if (input->rows().size() == 0) {
        continue;
      }
      auto* input_data = input->value().data<T>();
      auto& input_rows = input->rows();

      for (size_t i = 0; i < input_rows.size(); i++) {
        size_t out_i = rows_to_id[input_rows[i]];
702 703 704
        elementwise_add_to<T>(&blas, static_cast<size_t>(input_width),
                              &input_data[i * input_width],
                              &out_data[out_i * input_width]);
705 706 707 708 709 710 711 712 713 714 715 716
      }
    }
    size_t input_width_cast = static_cast<size_t>(input_width);
    T count = static_cast<T>(inputs.size());
    for (size_t i = 0; i < merge_rows.size(); i++) {
      for (size_t j = 0; j < input_width_cast; j++) {
        out_data[i * input_width + j] = out_data[i * input_width + j] / count;
      }
    }
  }
};

T
wip  
typhoonzero 已提交
717 718
template struct MergeAdd<platform::CPUDeviceContext, int>;
template struct MergeAdd<platform::CPUDeviceContext, int64_t>;
Q
Qiao Longfei 已提交
719 720
template struct MergeAdd<platform::CPUDeviceContext, float>;
template struct MergeAdd<platform::CPUDeviceContext, double>;
721
template struct MergeAdd<platform::CPUDeviceContext,
722
                         paddle::platform::complex<float>>;
723
template struct MergeAdd<platform::CPUDeviceContext,
724
                         paddle::platform::complex<double>>;
725 726
template struct MergeAdd<platform::CPUDeviceContext,
                         paddle::platform::bfloat16>;
T
wip  
typhoonzero 已提交
727

728 729 730 731
#ifdef PADDLE_WITH_XPU
template struct MergeAdd<platform::XPUDeviceContext, float>;
#endif

732 733 734 735 736
template struct MergeAverage<platform::CPUDeviceContext, int>;
template struct MergeAverage<platform::CPUDeviceContext, int64_t>;
template struct MergeAverage<platform::CPUDeviceContext, float>;
template struct MergeAverage<platform::CPUDeviceContext, double>;

T
wip  
typhoonzero 已提交
737 738
template <typename T>
struct UpdateToTensor<platform::CPUDeviceContext, T> {
T
typhoonzero 已提交
739
  void operator()(const platform::CPUDeviceContext& context,
740
                  const ScatterOps& op, const phi::SelectedRows& input1,
T
typhoonzero 已提交
741
                  framework::Tensor* input2) {
T
wip  
typhoonzero 已提交
742 743
    auto in1_height = input1.height();
    auto in2_dims = input2->dims();
744 745 746 747 748 749
    PADDLE_ENFORCE_EQ(
        in1_height, in2_dims[0],
        platform::errors::InvalidArgument("The two inputs height must be equal."
                                          "But recieved first input height = "
                                          "[%d], second input height = [%d]",
                                          in1_height, in2_dims[0]));
T
wip  
typhoonzero 已提交
750 751 752 753 754

    auto& in1_value = input1.value();
    auto& in1_rows = input1.rows();

    int64_t in1_row_numel = in1_value.numel() / in1_rows.size();
755 756 757 758 759 760
    PADDLE_ENFORCE_EQ(
        in1_row_numel, input2->numel() / in1_height,
        platform::errors::InvalidArgument(
            "The two inputs width must be equal."
            "But recieved first input width = [%d], second input width = [%d]",
            in1_row_numel, input2->numel() / in1_height));
T
wip  
typhoonzero 已提交
761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804

    auto* in1_data = in1_value.data<T>();
    auto* input2_data = input2->data<T>();

    // FIXME(typhoonzero): use macro fix the below messy code.
    switch (op) {
      case ScatterOps::ASSIGN:
        INLINE_FOR2(in1_rows.size(), in1_row_numel)
        input2_data[in1_rows[i] * in1_row_numel + j] =
            in1_data[i * in1_row_numel + j];
        break;
      case ScatterOps::ADD:
        INLINE_FOR2(in1_rows.size(), in1_row_numel)
        input2_data[in1_rows[i] * in1_row_numel + j] +=
            in1_data[i * in1_row_numel + j];
        break;
      case ScatterOps::SUB:
        INLINE_FOR2(in1_rows.size(), in1_row_numel)
        input2_data[in1_rows[i] * in1_row_numel + j] -=
            in1_data[i * in1_row_numel + j];
        break;
      case ScatterOps::SUBBY:
        INLINE_FOR2(in1_rows.size(), in1_row_numel)
        input2_data[in1_rows[i] * in1_row_numel + j] =
            in1_data[i * in1_row_numel + j] -
            input2_data[in1_rows[i] * in1_row_numel + j];
        break;
      case ScatterOps::MUL:
        INLINE_FOR2(in1_rows.size(), in1_row_numel)
        input2_data[in1_rows[i] * in1_row_numel + j] *=
            in1_data[i * in1_row_numel + j];
        break;
      case ScatterOps::DIV:
        INLINE_FOR2(in1_rows.size(), in1_row_numel)
        input2_data[in1_rows[i] * in1_row_numel + j] /=
            in1_data[i * in1_row_numel + j];
        break;
      case ScatterOps::DIVBY:
        INLINE_FOR2(in1_rows.size(), in1_row_numel)
        input2_data[in1_rows[i] * in1_row_numel + j] =
            in1_data[i * in1_row_numel + j] /
            input2_data[in1_rows[i] * in1_row_numel + j];
        break;
    }
T
typhoonzero 已提交
805 806 807 808
  }
};

}  // namespace scatter
809 810 811
}  // namespace math
}  // namespace operators
}  // namespace paddle