compare_op.cc 5.8 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Y
Yu Yang 已提交
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
Y
Yu Yang 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yu Yang 已提交
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. */
Y
Yu Yang 已提交
14

W
Wu Yi 已提交
15
#include "paddle/fluid/operators/controlflow/compare_op.h"
S
Siddharth Goyal 已提交
16
#include <string>
Y
Yi Wang 已提交
17
#include "paddle/fluid/framework/op_registry.h"
18

Y
Yu Yang 已提交
19 20
namespace paddle {
namespace operators {
Y
Yiqun Liu 已提交
21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44

template <typename Functor>
class CompareOpKernel<platform::CPUDeviceContext, Functor>
    : public framework::OpKernel<typename Functor::ELEM_TYPE> {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    using T = typename Functor::ELEM_TYPE;
    using Tensor = framework::Tensor;

    auto* x = context.Input<Tensor>("X");
    auto* y = context.Input<Tensor>("Y");
    auto* z = context.Output<Tensor>("Out");
    int axis = context.Attr<int>("axis");

    if (x->numel() == 1 && y->numel() == 1) {
      bool* z_data = z->mutable_data<bool>(context.GetPlace());
      z_data[0] = Functor()(x->data<T>()[0], y->data<T>()[0]);
    } else {
      ElementwiseComputeEx<Functor, platform::CPUDeviceContext, T, bool>(
          context, x, y, axis, Functor(), z);
    }
  }
};

Y
Yu Yang 已提交
45 46 47
template <typename OpComment>
class CompareOpProtoMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
48
  void Make() override {
Y
Yu Yang 已提交
49
    OpComment comment;
Y
yuyang18 已提交
50 51 52 53
    AddInput("X", string::Sprintf("the left hand operand of %s operator",
                                  comment.type));
    AddInput("Y", string::Sprintf("the right hand operand of %s operator",
                                  comment.type));
54 55 56 57 58
    AddAttr<int>(
        "axis",
        "The start dimension index for broadcasting Y onto X. [default -1]")
        .SetDefault(-1)
        .EqualGreaterThan(-1);
J
JiayiFeng 已提交
59
    AddAttr<bool>("force_cpu",
Y
yuyang18 已提交
60
                  "Force fill output variable to cpu "
J
JiayiFeng 已提交
61
                  "memory. Otherwise, fill output variable to the running "
Y
yuyang18 已提交
62 63 64 65
                  "device [default true].")
        .SetDefault(true);
    AddOutput("Out", string::Sprintf("n-dim bool tensor. Each element is %s",
                                     comment.equation));
Y
yuyang18 已提交
66
    AddComment(string::Sprintf(R"DOC(
Y
Yu Yang 已提交
67 68
It operates element-wise on X and Y, and returns the Out. Each of them is a
N-dim tensor. X and Y could be any type.  The each element of the Out tensor is
Y
yuyang18 已提交
69
calculated by $%s$
Y
Yu Yang 已提交
70
)DOC",
Y
yuyang18 已提交
71
                               comment.equation));
Y
Yu Yang 已提交
72 73 74 75
  }
};

template <typename OpComment>
Z
Zeng Jinle 已提交
76
class CompareOp : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
77
 public:
Z
Zeng Jinle 已提交
78 79 80 81
  using framework::OperatorWithKernel::OperatorWithKernel;

 protected:
  void InferShape(framework::InferShapeContext* context) const override {
Y
Yu Yang 已提交
82
    OpComment comment;
83 84
    OP_INOUT_CHECK(context->HasInput("X"), "Input", "X", comment.type);
    OP_INOUT_CHECK(context->HasInput("Y"), "Output", "Y", comment.type);
Y
Yu Yang 已提交
85 86
    auto dim_x = context->GetInputDim("X");
    auto dim_y = context->GetInputDim("Y");
87

88
    PADDLE_ENFORCE_GE(dim_x.size(), dim_y.size(),
89 90 91 92
                      platform::errors::InvalidArgument(
                          "The size of dim_y should not be greater than "
                          "dim_x's, but received dim_y: %d > dim_x: %d",
                          dim_y.size(), dim_x.size()));
Y
Yu Yang 已提交
93 94 95 96 97

    context->SetOutputDim("Out", context->GetInputDim("X"));
    context->ShareLoD("X", "Out");
  }

98
  framework::OpKernelType GetExpectedKernelType(
Y
Yiqun Liu 已提交
99
      const framework::ExecutionContext& ctx) const override {
100
    framework::OpKernelType kt = OperatorWithKernel::GetExpectedKernelType(ctx);
101
    // CompareOp kernel's device type is decided by input tensor place
J
JiayiFeng 已提交
102 103 104
    bool force_cpu = ctx.Attr<bool>("force_cpu");
    kt.place_ = force_cpu ? platform::CPUPlace()
                          : ctx.Input<framework::LoDTensor>("X")->place();
105 106 107 108
    return kt;
  }
};

Y
Yu Yang 已提交
109 110 111
}  // namespace operators
}  // namespace paddle

H
hong 已提交
112 113 114 115 116 117 118 119
#define REGISTER_COMPARE_OP(op_type, _equation)                         \
  struct _##op_type##Comment {                                          \
    static char type[];                                                 \
    static char equation[];                                             \
  };                                                                    \
  char _##op_type##Comment::type[]{#op_type};                           \
  char _##op_type##Comment::equation[]{_equation};                      \
  REGISTER_OPERATOR(                                                    \
Z
Zeng Jinle 已提交
120
      op_type, ::paddle::operators::CompareOp<_##op_type##Comment>,     \
H
hong 已提交
121 122 123
      ::paddle::operators::CompareOpProtoMaker<_##op_type##Comment>,    \
      ::paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>, \
      ::paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
Y
Yu Yang 已提交
124

Q
qiaolongfei 已提交
125 126 127 128
REGISTER_COMPARE_OP(less_than, "Out = X < Y");
REGISTER_COMPARE_KERNEL(less_than, CPU, paddle::operators::LessThanFunctor);
REGISTER_COMPARE_OP(less_equal, "Out = X <= Y");
REGISTER_COMPARE_KERNEL(less_equal, CPU, paddle::operators::LessEqualFunctor);
Q
qiaolongfei 已提交
129 130 131 132 133 134
REGISTER_COMPARE_OP(greater_than, "Out = X > Y");
REGISTER_COMPARE_KERNEL(greater_than, CPU,
                        paddle::operators::GreaterThanFunctor);
REGISTER_COMPARE_OP(greater_equal, "Out = X >= Y");
REGISTER_COMPARE_KERNEL(greater_equal, CPU,
                        paddle::operators::GreaterEqualFunctor);
Q
qiaolongfei 已提交
135 136 137 138
REGISTER_COMPARE_OP(equal, "Out = X == Y");
REGISTER_COMPARE_KERNEL(equal, CPU, paddle::operators::EqualFunctor);
REGISTER_COMPARE_OP(not_equal, "Out = X != Y");
REGISTER_COMPARE_KERNEL(not_equal, CPU, paddle::operators::NotEqualFunctor);