Lml/add prim ops

paddle/fluid/operators/CMakeLists.txt

@@ -22,6 +22,7 @@ add_subdirectory(reduce_ops)
 add_subdirectory(sequence_ops)
 add_subdirectory(string)
 add_subdirectory(jit)
+add_subdirectory(prim_ops)
 if(WITH_MKLDNN)
     add_subdirectory(mkldnn)
 endif()

paddle/fluid/operators/prim_ops/CMakeLists.txt

@@ -0,0 +1,28 @@
+include(operators)
+if(WITH_UNITY_BUILD)
+    # Load Unity Build rules for operators in paddle/fluid/operators/prim_ops.
+    include(unity_build_rule.cmake)
+endif()
+register_operators()
+
+SET(PRIM_OP_SRCS
+    reshape_p_op.cc
+    broadcast_p_op.cc
+    reduce_p_op.cc
+    transpose_p_op.cc
+    split_p_op.cc
+    concat_p_op.cc
+    slice_select_p_op.cc
+    slice_assign_p_op.cc
+    gather_p_op.cc
+    scatter_add_p_op.cc
+    add_p_op.cc
+    sub_p_op.cc
+    mul_p_op.cc
+    div_p_op.cc
+    sqrt_p_op.cc
+    tanh_p_op.cc
+    matmul_p_op.cc
+    fill_constant_p_op.cc)
+
+cc_test(prim_op_test SRCS prim_op_test.cc ${PRIM_OP_SRCS} DEPS op_registry)

paddle/fluid/operators/prim_ops/add_p_op.cc

@@ -0,0 +1,116 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// 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.
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+
+namespace paddle {
+namespace framework {
+class InferShapeContext;
+class VarDesc;
+}  // namespace framework
+}  // namespace paddle
+
+namespace paddle {
+namespace operators {
+class AddPrimOp : public framework::OperatorBase {
+ public:
+  AddPrimOp(const std::string &type, const framework::VariableNameMap &inputs,
+            const framework::VariableNameMap &outputs,
+            const framework::AttributeMap &attrs)
+      : framework::OperatorBase(type, inputs, outputs, attrs) {}
+  void RunImpl(const framework::Scope &scope,
+               const platform::Place &dev_place) const override {
+    PADDLE_THROW(platform::errors::Unimplemented(
+        "Prim operator add_p should not be excuted directly"));
+  }
+};
+
+class AddPrimOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X", "(Tensor), The input tensor of add_p op.");
+    AddInput("Y", "(Tensor), The input tensor of add_p op.");
+    AddOutput("Z", "(Tensor), The output tensor of add_p op.");
+    AddComment(R"DOC(
+Autograd primitive add_p operator.
+)DOC");
+  }
+};
+
+class AddPrimOpShapeInference : public framework::InferShapeBase {
+ public:
+  void operator()(framework::InferShapeContext *ctx) const override {
+    framework::InferShapeVarPtr x_var_ptr = ctx->GetInputVarPtrs("X")[0];
+    framework::InferShapeVarPtr y_var_ptr = ctx->GetInputVarPtrs("Y")[0];
+    framework::InferShapeVarPtr z_var_ptr = ctx->GetOutputVarPtrs("Z")[0];
+
+    framework::VarDesc *x_var = BOOST_GET(framework::VarDesc *, x_var_ptr);
+    framework::VarDesc *y_var = BOOST_GET(framework::VarDesc *, y_var_ptr);
+    auto x_shape = x_var->GetShape();
+    auto y_shape = y_var->GetShape();
+    size_t x_rank = x_shape.size();
+    size_t y_rank = y_shape.size();
+    PADDLE_ENFORCE_EQ(x_rank, y_rank,
+                      platform::errors::InvalidArgument(
+                          "The dimensions of two input tensor should be same, "
+                          "but get %d and %d",
+                          x_rank, y_rank));
+    for (size_t i = 0; i < x_rank; ++i) {
+      PADDLE_ENFORCE_EQ(
+          x_shape[i], y_shape[i],
+          platform::errors::InvalidArgument(
+              "The shape of two input tensor at dimension %d should be same, "
+              "but get %d and %d",
+              i, x_rank, y_rank));
+    }
+
+    BOOST_GET(framework::VarDesc *, z_var_ptr)->SetShape(x_shape);
+  }
+};
+
+class AddPrimOpVarTypeInference
+    : public framework::StaticGraphVarTypeInference {
+ public:
+  void operator()(framework::InferVarTypeContext *ctx) const override {
+    auto x_name = Input(ctx, "X")[0];
+    auto y_name = Input(ctx, "Y")[0];
+    auto z_name = Output(ctx, "Z")[0];
+    auto x_type = GetType(ctx, x_name);
+    auto y_type = GetType(ctx, y_name);
+    auto x_dtype = GetDataType(ctx, x_name);
+    auto y_dtype = GetDataType(ctx, y_name);
+    PADDLE_ENFORCE_EQ(x_type, y_type,
+                      platform::errors::InvalidArgument(
+                          "The type of two input tensor should be same, "
+                          "but get %d and %d",
+                          x_type, y_type));
+    PADDLE_ENFORCE_EQ(x_dtype, y_dtype,
+                      platform::errors::InvalidArgument(
+                          "The datatype of two input tensor should be same, "
+                          "but get %d and %d",
+                          x_dtype, y_dtype));
+
+    SetType(ctx, z_name, x_type);
+    SetDataType(ctx, z_name, x_dtype);
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+REGISTER_OPERATOR(add_p, paddle::operators::AddPrimOp,
+                  paddle::operators::AddPrimOpMaker,
+                  paddle::operators::AddPrimOpShapeInference,
+                  paddle::operators::AddPrimOpVarTypeInference);

paddle/fluid/operators/prim_ops/broadcast_p_op.cc

@@ -0,0 +1,110 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// 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.
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+
+namespace paddle {
+namespace framework {
+class InferShapeContext;
+class VarDesc;
+}  // namespace framework
+}  // namespace paddle
+
+namespace paddle {
+namespace operators {
+class BroadcastPrimOp : public framework::OperatorBase {
+ public:
+  BroadcastPrimOp(const std::string &type,
+                  const framework::VariableNameMap &inputs,
+                  const framework::VariableNameMap &outputs,
+                  const framework::AttributeMap &attrs)
+      : framework::OperatorBase(type, inputs, outputs, attrs) {}
+  void RunImpl(const framework::Scope &scope,
+               const platform::Place &dev_place) const override {
+    PADDLE_THROW(platform::errors::Unimplemented(
+        "Prim operator broadcast_p should not be excuted directly"));
+  }
+};
+
+class BroadcastPrimOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X", "(Tensor), The input tensor of broadcast_p op.");
+    AddOutput("Y", "(Tensor), The output tensor of broadcast_p op.");
+    AddAttr<std::vector<int64_t>>(
+        "shape",
+        "(std::vector<int64_t>) Target shape of broadcast_p operator.");
+    AddComment(R"DOC(
+Autograd primitive broadcast_p operator.
+)DOC");
+  }
+};
+
+static void CheckShapeValid(const std::vector<int64_t> &x_shape,
+                            const std::vector<int64_t> &target_shape) {
+  size_t x_rank = x_shape.size();
+  size_t target_rank = target_shape.size();
+  PADDLE_ENFORCE_GE(target_rank, x_rank,
+                    platform::errors::InvalidArgument(
+                        "The rank of target shape should be greater than or "
+                        "equal to input tensor's dimensions, "
+                        "but received %d and %d",
+                        target_rank, x_rank));
+  std::vector<int64_t>::const_iterator it = target_shape.begin();
+  for (size_t i = 0; i < x_rank; i++, it++) {
+    if (x_shape[i] != 1) {
+      it = std::find(it, target_shape.end(), x_shape[i]);
+    }
+    PADDLE_ENFORCE_EQ(
+        it != target_shape.end(), true,
+        platform::errors::InvalidArgument(
+            "Invalid shape, can not broadcast input tensor into target shape,"
+            "the first dismatching shape  %d is shape of input tensor at "
+            "dimension %d",
+            x_shape[i], i));
+  }
+}
+
+class BroadcastPrimOpShapeInference : public framework::InferShapeBase {
+ public:
+  void operator()(framework::InferShapeContext *ctx) const override {
+    framework::InferShapeVarPtr x_var_ptr = ctx->GetInputVarPtrs("X")[0];
+    framework::InferShapeVarPtr y_var_ptr = ctx->GetOutputVarPtrs("Y")[0];
+    framework::VarDesc *x_var = BOOST_GET(framework::VarDesc *, x_var_ptr);
+    auto x_shape = x_var->GetShape();
+    auto target_shape = ctx->Attrs().Get<std::vector<int64_t>>("shape");
+    CheckShapeValid(x_shape, target_shape);
+    BOOST_GET(framework::VarDesc *, y_var_ptr)->SetShape(target_shape);
+  }
+};
+
+class BroadcastPrimOpVarTypeInference
+    : public framework::StaticGraphVarTypeInference {
+ public:
+  void operator()(framework::InferVarTypeContext *ctx) const override {
+    auto x_name = Input(ctx, "X")[0];
+    auto y_name = Output(ctx, "Y")[0];
+    SetType(ctx, y_name, GetType(ctx, x_name));
+    SetDataType(ctx, y_name, GetDataType(ctx, x_name));
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+REGISTER_OPERATOR(broadcast_p, paddle::operators::BroadcastPrimOp,
+                  paddle::operators::BroadcastPrimOpMaker,
+                  paddle::operators::BroadcastPrimOpShapeInference,
+                  paddle::operators::BroadcastPrimOpVarTypeInference);

paddle/fluid/operators/prim_ops/concat_p_op.cc

@@ -0,0 +1,134 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// 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.
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+
+namespace paddle {
+namespace framework {
+class InferShapeContext;
+class VarDesc;
+}  // namespace framework
+}  // namespace paddle
+
+namespace paddle {
+namespace operators {
+class ConcatPrimOp : public framework::OperatorBase {
+ public:
+  ConcatPrimOp(const std::string &type,
+               const framework::VariableNameMap &inputs,
+               const framework::VariableNameMap &outputs,
+               const framework::AttributeMap &attrs)
+      : framework::OperatorBase(type, inputs, outputs, attrs) {}
+  void RunImpl(const framework::Scope &scope,
+               const platform::Place &dev_place) const override {
+    PADDLE_THROW(platform::errors::Unimplemented(
+        "Prim operator concat_p should not be excuted directly"));
+  }
+};
+
+class ConcatPrimOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("XS", "(Tensor), The input tensors of concat_p op.")
+        .AsDuplicable();
+    AddOutput("Y", "(Tensor), The output tensor of concat_p op.");
+    AddAttr<int64_t>("axis", "(int64_t), The axis along which to concat.");
+    AddComment(R"DOC(
+Autograd primitive concat_p operator.
+)DOC");
+  }
+};
+
+class ConcatPrimOpShapeInference : public framework::InferShapeBase {
+ public:
+  void operator()(framework::InferShapeContext *ctx) const override {
+    auto x_var_ptrs = ctx->GetInputVarPtrs("XS");
+    framework::InferShapeVarPtr y_var_ptr = ctx->GetOutputVarPtrs("Y")[0];
+    auto axis = ctx->Attrs().Get<int64_t>("axis");
+    int64_t cnt_along_axis = 0;
+    framework::VarDesc *first_x_var =
+        BOOST_GET(framework::VarDesc *, x_var_ptrs[0]);
+    auto first_x_shape = first_x_var->GetShape();
+    cnt_along_axis += first_x_shape[axis];
+    size_t first_x_rank = first_x_shape.size();
+    for (size_t i = 1; i < x_var_ptrs.size(); ++i) {
+      framework::VarDesc *x_var =
+          BOOST_GET(framework::VarDesc *, x_var_ptrs[i]);
+      auto x_shape = x_var->GetShape();
+      cnt_along_axis += x_shape[axis];
+      size_t x_rank = x_shape.size();
+      PADDLE_ENFORCE_EQ(
+          x_rank, first_x_rank,
+          platform::errors::InvalidArgument("The dimensions of %d input tensor "
+                                            "should be same as the dimensions "
+                                            "of 1st input tensor's, "
+                                            "but get %d and %d",
+                                            i + 1, x_rank, first_x_rank));
+      for (size_t j = 0; j < x_rank; ++j) {
+        if (j != size_t(axis)) {
+          PADDLE_ENFORCE_EQ(x_shape[j], first_x_shape[j],
+                            platform::errors::InvalidArgument(
+                                "The shape of %d input tensor at dimension %d "
+                                "should be same as the 1st input tensor's, "
+                                "but get %d and %d",
+                                i + 1, j, x_shape[j], first_x_shape[j]));
+        }
+      }
+    }
+
+    std::vector<int64_t> y_shape(first_x_shape);
+    y_shape[axis] = cnt_along_axis;
+    BOOST_GET(framework::VarDesc *, y_var_ptr)->SetShape(y_shape);
+  }
+};
+
+class ConcatPrimOpVarTypeInference
+    : public framework::StaticGraphVarTypeInference {
+ public:
+  void operator()(framework::InferVarTypeContext *ctx) const override {
+    auto x_names = Input(ctx, "XS");
+    auto y_name = Output(ctx, "Y")[0];
+    auto first_x_name = x_names[0];
+    auto first_x_type = GetType(ctx, first_x_name);
+    auto first_x_dtype = GetDataType(ctx, first_x_name);
+    for (size_t i = 1; i < x_names.size(); ++i) {
+      auto x_name = x_names[i];
+      auto x_type = GetType(ctx, x_name);
+      auto x_dtype = GetDataType(ctx, x_name);
+      PADDLE_ENFORCE_EQ(x_type, first_x_type,
+                        platform::errors::InvalidArgument(
+                            "The type of %d input tensor should be same as the "
+                            "first input tensor's, "
+                            "but get %d and %d",
+                            i + 1, x_type, first_x_type));
+      PADDLE_ENFORCE_EQ(x_dtype, first_x_dtype,
+                        platform::errors::InvalidArgument(
+                            "The datatype of %d input tensor should be same as "
+                            "the first input tensor's, "
+                            "but get %d and %d",
+                            i + 1, x_dtype, first_x_dtype));
+    }
+    SetType(ctx, y_name, GetType(ctx, first_x_name));
+    SetDataType(ctx, y_name, GetDataType(ctx, first_x_name));
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+REGISTER_OPERATOR(concat_p, paddle::operators::ConcatPrimOp,
+                  paddle::operators::ConcatPrimOpMaker,
+                  paddle::operators::ConcatPrimOpShapeInference,
+                  paddle::operators::ConcatPrimOpVarTypeInference);