elementwise_mul_op.h 4.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

   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
16
#include "paddle/operators/elementwise_op_function.h"
17 18 19 20

namespace paddle {
namespace operators {

Q
QI JUN 已提交
21
template <typename DeviceContext, typename T>
Y
Yu Yang 已提交
22
class ElementwiseMulKernel : public framework::OpKernel<T> {
23 24
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
Q
QI JUN 已提交
25
    ElementwiseCompute<EigenMulFunctor, DeviceContext, T>(ctx);
G
gongweibao 已提交
26 27
  }
};
28

G
gongweibao 已提交
29 30 31 32 33
template <typename T>
struct ElementwiseMulGradFunctor {
  template <typename Device, typename X, typename Y, typename Z, typename dX,
            typename dY, typename dZ>
  void operator()(Device d, X x, Y y, Z z, dX dx, dY dy, dZ dz) {
34 35
    auto x_e = framework::EigenVector<T>::Flatten(*x);
    auto y_e = framework::EigenVector<T>::Flatten(*y);
G
gongweibao 已提交
36
    auto dz_e = framework::EigenVector<T>::Flatten(*dz);
37

G
gongweibao 已提交
38 39 40
    if (dx) {
      auto dx_e = framework::EigenVector<T>::Flatten(*dx);
      dx_e.device(d) = dz_e * y_e;
41 42
    }

G
gongweibao 已提交
43 44 45
    if (dy) {
      auto dy_e = framework::EigenVector<T>::Flatten(*dy);
      dy_e.device(d) = x_e * dz_e;
46 47 48 49
    }
  }
};

G
gongweibao 已提交
50 51 52 53 54
template <typename T>
struct ElementwiseMulBroadCastGradFunctor {
  template <typename Device, typename X, typename Y, typename Z, typename dX,
            typename dY, typename dZ, typename Pre, typename N>
  void operator()(Device d, X x, Y y, Z z, dX dx, dY dy, dZ dz, Pre pre, N n) {
55 56
    auto x_e = framework::EigenVector<T>::Flatten(*x);
    auto y_e = framework::EigenVector<T>::Flatten(*y);
G
gongweibao 已提交
57
    auto dz_e = framework::EigenVector<T>::Flatten(*dz);
58

G
gongweibao 已提交
59 60 61
    auto y_e_bcast = y_e.reshape(Eigen::DSizes<int, 2>(1, n))
                         .broadcast(Eigen::DSizes<int, 2>(pre, 1))
                         .reshape(Eigen::DSizes<int, 1>(x_e.size()));
62 63

    if (dx) {
G
gongweibao 已提交
64 65
      auto dx_e = framework::EigenVector<T>::Flatten(*dx);
      dx_e.device(d) = dz_e * y_e_bcast;
66
    }
G
gongweibao 已提交
67

68
    if (dy) {
G
gongweibao 已提交
69 70 71 72
      auto dy_e = framework::EigenVector<T>::Flatten(*dy);
      dy_e.device(d) = (x_e * dz_e)
                           .reshape(Eigen::DSizes<int, 2>(pre, n))
                           .sum(Eigen::array<int, 1>{{0}});
73
    }
G
gongweibao 已提交
74 75
  }
};
76

G
gongweibao 已提交
77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92
template <typename T>
struct ElementwiseMulBroadCast2GradFunctor {
  template <typename Device, typename X, typename Y, typename Z, typename dX,
            typename dY, typename dZ, typename Pre, typename N, typename Post>
  void operator()(Device d, X x, Y y, Z z, dX dx, dY dy, dZ dz, Pre pre, N n,
                  Post post) {
    auto x_e = framework::EigenVector<T>::Flatten(*x);
    auto y_e = framework::EigenVector<T>::Flatten(*y);
    auto dz_e = framework::EigenVector<T>::Flatten(*dz);

    auto y_e_bcast = y_e.reshape(Eigen::DSizes<int, 3>(1, n, 1))
                         .broadcast(Eigen::DSizes<int, 3>(pre, 1, post))
                         .reshape(Eigen::DSizes<int, 1>(x_e.size()));
    if (dx) {
      auto dx_e = framework::EigenVector<T>::Flatten(*dx);
      dx_e.device(d) = dz_e * y_e_bcast;
93 94
    }

G
gongweibao 已提交
95 96 97 98 99
    if (dy) {
      auto dy_e = framework::EigenVector<T>::Flatten(*dy);
      dy_e.device(d) = (x_e * dz_e)
                           .reshape(Eigen::DSizes<int, 3>(pre, n, post))
                           .sum(Eigen::array<int, 2>{{0, 2}});
100 101 102 103
    }
  }
};

Q
QI JUN 已提交
104
template <typename DeviceContext, typename T>
Y
Yu Yang 已提交
105
class ElementwiseMulGradKernel : public framework::OpKernel<T> {
G
gongweibao 已提交
106 107
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
Q
QI JUN 已提交
108
    ElementwiseGradCompute<DeviceContext, T, ElementwiseMulGradFunctor<T>,
G
gongweibao 已提交
109 110 111 112 113 114
                           ElementwiseMulGradFunctor<T>,
                           ElementwiseMulBroadCastGradFunctor<T>,
                           ElementwiseMulBroadCast2GradFunctor<T>>(ctx);
  }
};

115 116
}  // namespace operators
}  // namespace paddle