selected_rows_functor.h 9.1 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. */
#pragma once
15

S
sneaxiy 已提交
16
#include <map>
17 18
#include <vector>

Y
Yi Wang 已提交
19 20
#include "paddle/fluid/framework/eigen.h"
#include "paddle/fluid/framework/selected_rows.h"
21 22
#include "paddle/fluid/operators/math/blas.h"
#include "paddle/fluid/operators/math/math_function.h"
Y
Yi Wang 已提交
23
#include "paddle/fluid/platform/device_context.h"
24

T
wip  
typhoonzero 已提交
25 26 27 28
#define INLINE_FOR2(sizei, sizej)     \
  for (int64_t i = 0; i < sizei; i++) \
    for (int64_t j = 0; j < sizej; j++)

29 30 31 32 33 34
namespace paddle {
namespace operators {
namespace math {

// SelectedRows + SelectedRows will simplely concat value and rows.
// The real computation happens in dealing with LoDTensor.
Q
QI JUN 已提交
35
template <typename DeviceContext, typename T>
36
struct SelectedRowsAdd {
Q
QI JUN 已提交
37
  void operator()(const DeviceContext& context,
38 39 40 41 42
                  const framework::SelectedRows& input1,
                  const framework::SelectedRows& input2,
                  framework::SelectedRows* output);
};

Q
QI JUN 已提交
43
template <typename DeviceContext, typename T>
44
struct SelectedRowsAddTensor {
Q
QI JUN 已提交
45
  void operator()(const DeviceContext& context,
46 47 48 49
                  const framework::SelectedRows& input1,
                  const framework::Tensor& input2, framework::Tensor* output);
};

50
// input2 = input1 + input2
Q
QI JUN 已提交
51
template <typename DeviceContext, typename T>
52
struct SelectedRowsAddTo {
Q
QI JUN 已提交
53
  void operator()(const DeviceContext& context,
54 55 56 57
                  const framework::SelectedRows& input1,
                  const int64_t input2_offset, framework::SelectedRows* input2);
};

58 59 60 61 62 63 64 65 66
// input2 = [all input in input1] + input2
template <typename DeviceContext, typename T>
struct SelectedRowsSumTo {
  void operator()(const DeviceContext& context,
                  const std::vector<framework::SelectedRows*>& input1,
                  const std::vector<int64_t>& input2_offsets,
                  framework::SelectedRows* input2);
};

67
// input2 = input1 + input2
Q
QI JUN 已提交
68
template <typename DeviceContext, typename T>
69
struct SelectedRowsAddToTensor {
Q
QI JUN 已提交
70
  void operator()(const DeviceContext& context,
71 72 73 74
                  const framework::SelectedRows& input1,
                  framework::Tensor* input2);
};

75 76 77 78 79 80
namespace scatter {
// functors for manuplating SelectedRows data
template <typename DeviceContext, typename T>
struct MergeAdd {
  // unary functor, merge by adding duplicated rows in
  // the input SelectedRows object.
T
wip  
typhoonzero 已提交
81 82
  framework::SelectedRows operator()(const DeviceContext& context,
                                     const framework::SelectedRows& input);
S
sneaxiy 已提交
83 84 85
  void operator()(const DeviceContext& context,
                  const framework::SelectedRows& input,
                  framework::SelectedRows* output);
86 87
};

88 89 90 91 92 93 94 95 96 97 98 99 100
template <>
struct MergeAdd<platform::CPUDeviceContext, float> {
  framework::SelectedRows operator()(const platform::CPUDeviceContext& context,
                                     const framework::SelectedRows& input) {
    framework::SelectedRows out;
    (*this)(context, input, &out);
    return out;
  }

  void operator()(const platform::CPUDeviceContext& context,
                  const framework::SelectedRows& input,
                  framework::SelectedRows* output) {
    framework::SelectedRows& out = *output;
S
sneaxiy 已提交
101 102 103 104 105
    std::vector<int64_t> input_rows(input.rows());

    std::map<int64_t, std::vector<int64_t>> merge_row_map;
    for (size_t i = 0; i < input_rows.size(); ++i) {
      merge_row_map[input_rows[i]].push_back(i);
106 107
    }

S
sneaxiy 已提交
108 109 110
    std::vector<int64_t> merge_rows(merge_row_map.size());
    size_t idx = 0;
    int64_t input_width = input.value().dims()[1];
111
    out.set_height(input.height());
S
sneaxiy 已提交
112 113

    auto* out_data = out.mutable_value()->mutable_data<float>(
114 115 116
        framework::make_ddim(
            {static_cast<int64_t>(merge_rows.size()), input_width}),
        context.GetPlace());
S
sneaxiy 已提交
117
    auto* in_data = input.value().data<float>();
118 119

    auto blas = GetBlas<platform::CPUDeviceContext, float>(context);
S
sneaxiy 已提交
120 121 122 123 124 125 126 127 128 129 130
    for (auto& row_pair : merge_row_map) {
      auto* out_ptr = out_data + idx * input_width;
      auto& rows = row_pair.second;
      merge_rows[idx] = row_pair.first;
      ++idx;
      // rows.size() is always larger than 0
      blas.VCOPY(input_width, in_data + rows[0] * input_width, out_ptr);

      for (size_t i = 1; i < rows.size(); ++i) {
        blas.AXPY(input_width, 1., in_data + rows[i] * input_width, out_ptr);
      }
131
    }
S
sneaxiy 已提交
132 133

    out.set_rows(merge_rows);
134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149
  }
};

template <>
struct MergeAdd<platform::CPUDeviceContext, double> {
  framework::SelectedRows operator()(const platform::CPUDeviceContext& context,
                                     const framework::SelectedRows& input) {
    framework::SelectedRows out;
    (*this)(context, input, &out);
    return out;
  }

  void operator()(const platform::CPUDeviceContext& context,
                  const framework::SelectedRows& input,
                  framework::SelectedRows* output) {
    framework::SelectedRows& out = *output;
S
sneaxiy 已提交
150 151 152 153 154
    std::vector<int64_t> input_rows(input.rows());

    std::map<int64_t, std::vector<int64_t>> merge_row_map;
    for (size_t i = 0; i < input_rows.size(); ++i) {
      merge_row_map[input_rows[i]].push_back(i);
155 156
    }

S
sneaxiy 已提交
157 158 159
    std::vector<int64_t> merge_rows(merge_row_map.size());
    size_t idx = 0;
    int64_t input_width = input.value().dims()[1];
160
    out.set_height(input.height());
S
sneaxiy 已提交
161 162

    auto* out_data = out.mutable_value()->mutable_data<double>(
163 164 165
        framework::make_ddim(
            {static_cast<int64_t>(merge_rows.size()), input_width}),
        context.GetPlace());
S
sneaxiy 已提交
166
    auto* in_data = input.value().data<double>();
167 168

    auto blas = GetBlas<platform::CPUDeviceContext, double>(context);
S
sneaxiy 已提交
169 170 171 172 173 174 175 176 177 178 179
    for (auto& row_pair : merge_row_map) {
      auto* out_ptr = out_data + idx * input_width;
      auto& rows = row_pair.second;
      merge_rows[idx] = row_pair.first;
      ++idx;
      // rows.size() is always larger than 0
      blas.VCOPY(input_width, in_data + rows[0] * input_width, out_ptr);

      for (size_t i = 1; i < rows.size(); ++i) {
        blas.AXPY(input_width, 1., in_data + rows[i] * input_width, out_ptr);
      }
180
    }
S
sneaxiy 已提交
181 182

    out.set_rows(merge_rows);
183 184 185
  }
};

186 187
template <typename DeviceContext, typename T>
struct Add {
T
wip  
typhoonzero 已提交
188 189 190 191 192 193 194 195 196
  framework::SelectedRows operator()(const DeviceContext& context,
                                     const framework::SelectedRows& input1,
                                     const framework::SelectedRows& input2) {
    framework::SelectedRows out;
    out.set_rows(input1.rows());
    out.set_height(input1.height());
    out.mutable_value()->mutable_data<T>(input1.value().dims(),
                                         context.GetPlace());
    auto e_out = framework::EigenVector<T>::Flatten(*(out.mutable_value()));
T
wip  
typhoonzero 已提交
197 198
    auto e_in1 = framework::EigenVector<T>::Flatten(input1.value());
    auto e_in2 = framework::EigenVector<T>::Flatten(input2.value());
199
    e_out.device(*context.eigen_device()) = e_in1 + e_in2;
T
wip  
typhoonzero 已提交
200
    return out;
201 202 203 204 205
  }
};

template <typename DeviceContext, typename T>
struct Mul {
T
wip  
typhoonzero 已提交
206 207 208 209 210 211 212 213 214 215
  // multiply two SelectedRows
  framework::SelectedRows operator()(const DeviceContext& context,
                                     const framework::SelectedRows& input1,
                                     const framework::SelectedRows& input2) {
    framework::SelectedRows out;
    out.set_rows(input1.rows());
    out.set_height(input1.height());
    out.mutable_value()->mutable_data<T>(input1.value().dims(),
                                         context.GetPlace());
    auto e_out = framework::EigenVector<T>::Flatten(*(out.mutable_value()));
T
wip  
typhoonzero 已提交
216 217
    auto e_in1 = framework::EigenVector<T>::Flatten(input1.value());
    auto e_in2 = framework::EigenVector<T>::Flatten(input2.value());
218
    e_out.device(*context.eigen_device()) = e_in1 * e_in2;
T
wip  
typhoonzero 已提交
219 220 221 222 223 224 225 226 227 228 229 230 231 232 233
    return out;
  }
  // multiply scalar to SelectedRows
  framework::SelectedRows operator()(const DeviceContext& context,
                                     const framework::SelectedRows& input1,
                                     const T input2) {
    framework::SelectedRows out;
    out.set_rows(input1.rows());
    out.set_height(input1.height());
    out.mutable_value()->mutable_data<T>(input1.value().dims(),
                                         context.GetPlace());
    auto e_out = framework::EigenVector<T>::Flatten(*(out.mutable_value()));
    auto e_in1 = framework::EigenVector<T>::Flatten(input1.value());
    e_out.device(*context.eigen_device()) = input2 * e_in1;
    return out;
234 235 236
  }
};

T
wip  
typhoonzero 已提交
237 238 239 240 241
enum class ScatterOps { ASSIGN, ADD, SUB, SUBBY, MUL, DIV, DIVBY };

// out = seleted_rows_in / tensor
template <typename DeviceContext, typename T>
struct UpdateToTensor {
T
typhoonzero 已提交
242 243 244
  void operator()(const DeviceContext& context, const ScatterOps& op,
                  const framework::SelectedRows& input1,
                  framework::Tensor* input2);
T
wip  
typhoonzero 已提交
245 246
};

247
}  // namespace scatter
248 249 250
}  // namespace math
}  // namespace operators
}  // namespace paddle
反馈
建议
客服 返回
顶部