elementwise_sub_op.h 3.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
G
gongweibao 已提交
2

L
Luo Tao 已提交
3 4 5
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
G
gongweibao 已提交
6

7
    http://www.apache.org/licenses/LICENSE-2.0
G
gongweibao 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
14

F
fengjiayi 已提交
15
#pragma once
16

W
Wu Yi 已提交
17
#include "paddle/fluid/operators/elementwise/elementwise_op.h"
18
#include "paddle/fluid/platform/place.h"
G
gongweibao 已提交
19

20
#include "paddle/pten/kernels/elementwise_grad_kernel.h"
21
#include "paddle/pten/kernels/math_kernel.h"
G
gongweibao 已提交
22 23 24
namespace paddle {
namespace operators {

Q
QI JUN 已提交
25
template <typename DeviceContext, typename T>
Y
Yu Yang 已提交
26
class ElementwiseSubKernel : public framework::OpKernel<T> {
G
gongweibao 已提交
27 28
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
C
chengduo 已提交
29 30 31
    auto* x = ctx.Input<framework::LoDTensor>("X");
    auto* y = ctx.Input<framework::LoDTensor>("Y");
    auto* z = ctx.Output<framework::LoDTensor>("Out");
C
chengduoZH 已提交
32
    z->mutable_data<T>(ctx.GetPlace());
33

34 35
    auto& dev_ctx = ctx.device_context<DeviceContext>();
    int axis = ctx.Attr<int>("axis");
36
    pten::SubtractRawKernel<T>(
W
Wilber 已提交
37 38
        static_cast<const typename framework::ConvertToPtenContext<
            DeviceContext>::TYPE&>(dev_ctx),
39
        *x, *y, axis, z);
G
gongweibao 已提交
40 41 42
  }
};

Q
QI JUN 已提交
43
template <typename DeviceContext, typename T>
44
class ElementwiseSubGradKernel : public ElemwiseGradKernel<T> {
G
gongweibao 已提交
45 46
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
47
    ElemwiseGradKernel<T>::Compute(ctx);
C
chengduoZH 已提交
48 49
    using Tensor = framework::Tensor;

50 51
    auto* x = ctx.Input<Tensor>("X");
    auto* y = ctx.Input<Tensor>("Y");
C
chengduoZH 已提交
52 53 54
    auto* dout = ctx.Input<Tensor>(framework::GradVarName("Out"));
    auto* dx = ctx.Output<Tensor>(framework::GradVarName("X"));
    auto* dy = ctx.Output<Tensor>(framework::GradVarName("Y"));
55 56 57 58 59 60 61
    int axis = ctx.Attr<int>("axis");
    auto& dev_ctx = ctx.device_context<DeviceContext>();

    pten::SubtractGradKernel<T>(
        static_cast<const typename framework::ConvertToPtenContext<
            DeviceContext>::TYPE&>(dev_ctx),
        *x, *y, *dout, axis, dx, dy);
G
gongweibao 已提交
62 63
  }
};
64 65 66 67 68 69 70 71 72 73 74 75 76

template <typename DeviceContext, typename T>
class ElementwiseSubDoubleGradKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    using Tensor = framework::Tensor;

    auto* y = ctx.Input<Tensor>("Y");
    auto* dout = ctx.Input<Tensor>("DOut");
    auto* ddx = ctx.Input<Tensor>("DDX");
    auto* ddy = ctx.Input<Tensor>("DDY");

    auto* ddout = ctx.Output<Tensor>("DDOut");
77 78
    int axis = ctx.Attr<int>("axis");
    auto& dev_ctx = ctx.device_context<DeviceContext>();
79

80 81 82 83
    paddle::optional<const pten::DenseTensor&> ddx_optional = paddle::none;
    paddle::optional<const pten::DenseTensor&> ddy_optional = paddle::none;
    if (ddx != nullptr) {
      ddx_optional = *ddx;
84
    }
85 86 87 88 89 90 91
    if (ddy != nullptr) {
      ddy_optional = *ddy;
    }
    pten::SubtractDoubleGradKernel<T>(
        static_cast<const typename framework::ConvertToPtenContext<
            DeviceContext>::TYPE&>(dev_ctx),
        *y, ddx_optional, ddy_optional, *dout, axis, ddout);
92 93 94
  }
};

G
gongweibao 已提交
95 96
}  // namespace operators
}  // namespace paddle