[Cherry-Pick] Split Macros and Add modeling unittest (#31266)

* [CustomOp] Add Modeling with Custom op unittest (#31218)

* add unittest for static/dygraph/dy2stat

* add PE unittet

* remove usless code

* add unittest in CMakeList.txt

* [CustomOp] Split build op marco & polish details (#31229)

* split build op marco & polish details

* revert register api del

* fix other unittest

* [CustomOP]Support Incremental compilation and Add Version management (#31228)

* Support Incremental compilation and Add Version management

* replace hash with hashlib

* fix test_op_num unittest

* Revert "fix test_op_num unittest"

This reverts commit 2f78de9.

Co-authored-by: Chen Weihang <[email protected]>

Author: Aurelius84

paddle/fluid/extension/include/op_meta_info.h

@@ -38,6 +38,8 @@ class PD_DLL_DECL OpMetaInfoHelper;
 
 using Tensor = paddle::Tensor;
 
+///////////////// Util Marco Define ////////////////
+
 #define PD_DISABLE_COPY_AND_ASSIGN(classname)      \
  private:                                          \
   classname(const classname&) = delete;            \
@@ -65,6 +67,12 @@ using Tensor = paddle::Tensor;
     END_HANDLE_THE_ERROR               \
   } while (0)
 
+#define STATIC_ASSERT_GLOBAL_NAMESPACE(uniq_name, msg)                        \
+  struct __test_global_namespace_##uniq_name##__ {};                          \
+  static_assert(std::is_same<::__test_global_namespace_##uniq_name##__,       \
+                             __test_global_namespace_##uniq_name##__>::value, \
+                msg)
+
 ///////////////// Util Define and Function ////////////////
 
 inline std::string Grad(const std::string& var_name) {
@@ -288,9 +296,9 @@ class PD_DLL_DECL OpMetaInfo {
   std::vector<std::string> attrs_;
 
   // 2. func info
-  KernelFunc kernel_fn_;
-  InferShapeFunc infer_shape_fn_;
-  InferDtypeFunc infer_dtype_fn_;
+  KernelFunc kernel_fn_{nullptr};
+  InferShapeFunc infer_shape_fn_{nullptr};
+  InferDtypeFunc infer_dtype_fn_{nullptr};
 };
 
 //////////////// Op Meta Info Map /////////////////
@@ -321,20 +329,22 @@ class PD_DLL_DECL OpMetaInfoMap {
 
 class PD_DLL_DECL OpMetaInfoBuilder {
  public:
-  explicit OpMetaInfoBuilder(std::string&& name);
+  explicit OpMetaInfoBuilder(std::string&& name, size_t index);
   OpMetaInfoBuilder& Inputs(std::vector<std::string>&& inputs);
   OpMetaInfoBuilder& Outputs(std::vector<std::string>&& outputs);
   OpMetaInfoBuilder& Attrs(std::vector<std::string>&& attrs);
   OpMetaInfoBuilder& SetKernelFn(KernelFunc func);
   OpMetaInfoBuilder& SetInferShapeFn(InferShapeFunc func);
   OpMetaInfoBuilder& SetInferDtypeFn(InferDtypeFunc func);
-  OpMetaInfoBuilder& SetBackwardOp(const std::string& bwd_op_name);
 
  private:
   // Forward Op name
   std::string name_;
-  // Point to the currently constructed op meta info
+  // ref current info ptr
   OpMetaInfo* info_ptr_;
+  // The current op meta info index in vector
+  // - 0: op, 1: grad_op, 2: grad_grad_op
+  size_t index_;
 };
 
 /////////////////////// Op register API /////////////////////////
@@ -350,14 +360,25 @@ void LoadCustomOperatorLib(const std::string& dso_name);
 
 /////////////////////// Op register Macro /////////////////////////
 
-#define PD_BUILD_OP_WITH_COUNTER(op_name, counter)                  \
-  static ::paddle::OpMetaInfoBuilder __op_meta_info_##counter##__ = \
-      ::paddle::OpMetaInfoBuilder(op_name)
-
-#define PD_BUILD_OP_INNER(op_name, counter) \
-  PD_BUILD_OP_WITH_COUNTER(op_name, counter)
-
-#define PD_BUILD_OP(op_name) PD_BUILD_OP_INNER(op_name, __COUNTER__)
+#define PD_BUILD_OP(op_name)                                                   \
+  STATIC_ASSERT_GLOBAL_NAMESPACE(                                              \
+      __reg_op__##op_name, "PD_BUILD_OP must be called in global namespace."); \
+  static ::paddle::OpMetaInfoBuilder __op_meta_info_##op_name##__ =            \
+      ::paddle::OpMetaInfoBuilder(#op_name, 0)
+
+#define PD_BUILD_GRAD_OP(op_name)                                        \
+  STATIC_ASSERT_GLOBAL_NAMESPACE(                                        \
+      __reg_grad_op__##op_name,                                          \
+      "PD_BUILD_GRAD_OP must be called in global namespace.");           \
+  static ::paddle::OpMetaInfoBuilder __grad_op_meta_info_##op_name##__ = \
+      ::paddle::OpMetaInfoBuilder(#op_name, 1)
+
+#define PD_BUILD_DOUBLE_GRAD_OP(op_name)                                      \
+  STATIC_ASSERT_GLOBAL_NAMESPACE(                                             \
+      __reg_grad_grad_op__##op_name,                                          \
+      "PD_BUILD_DOUBLE_GRAD_OP must be called in global namespace.");         \
+  static ::paddle::OpMetaInfoBuilder __grad_grad_op_meta_info_##op_name##__ = \
+      ::paddle::OpMetaInfoBuilder(#op_name, 2)
 
 }  // namespace paddle
 

paddle/fluid/extension/src/op_meta_info.cc

@@ -19,6 +19,7 @@ limitations under the License. */
 #include <vector>
 
 #include "paddle/fluid/framework/custom_operator.h"
+#include "paddle/fluid/platform/enforce.h"
 
 namespace paddle {
 
@@ -62,11 +63,38 @@ OpMetaInfoMap::GetMap() const {
 
 //////////////// Op Meta Info Builder /////////////////
 
-OpMetaInfoBuilder::OpMetaInfoBuilder(std::string&& name) {
+OpMetaInfoBuilder::OpMetaInfoBuilder(std::string&& name, size_t index) {
+  // 1. member assign
   name_ = std::forward<std::string>(name);
+  index_ = index;
+
+  // 2. check and meta info build
   auto& info_vector = OpMetaInfoMap::Instance()[name_];
+  // index check
+  PADDLE_ENFORCE_EQ(
+      info_vector.size(), index_,
+      platform::errors::PreconditionNotMet(
+          "The operator %s's meta info register failed. "
+          "Please make sure you call marcos as order `PD_BUILD_OP`, "
+          "`PD_BUILD_GRAD_OP`, `PD_BUILD_DOUBLE_GRAD_OP`.",
+          name_));
+  switch (index_) {
+    case 0:
+      break;
+    case 1:
+      name_ = name_ + "_grad";
+      break;
+    case 2:
+      name_ = name_ + "_grad_grad";
+    default:
+      PADDLE_THROW(platform::errors::InvalidArgument(
+          "Not support index `%d` when construct OpMetaInfoBuilder, "
+          "now only support `0, 1, 2`.",
+          index_));
+  }
   auto op_meta = OpMetaInfo(name_);
   info_vector.emplace_back(std::move(op_meta));
+  // 3. get current info ptr
   info_ptr_ = &(info_vector.back());
 }
 
@@ -93,24 +121,27 @@ OpMetaInfoBuilder& OpMetaInfoBuilder::SetKernelFn(KernelFunc func) {
 }
 
 OpMetaInfoBuilder& OpMetaInfoBuilder::SetInferShapeFn(InferShapeFunc func) {
+  PADDLE_ENFORCE_EQ(
+      index_, 0UL,
+      platform::errors::Unimplemented(
+          "Currently, the InferShapeFn setting of Grad Op is not supported, "
+          "And backward Tensor `X@GRAD` will use the shape of forward Tensor "
+          "`X` by default."));
   info_ptr_->SetInferShapeFn(std::forward<InferShapeFunc>(func));
   return *this;
 }
 
 OpMetaInfoBuilder& OpMetaInfoBuilder::SetInferDtypeFn(InferDtypeFunc func) {
+  PADDLE_ENFORCE_EQ(
+      index_, 0UL,
+      platform::errors::Unimplemented(
+          "Currently, the InferDtypeFn setting of Grad Op is not supported, "
+          "And backward Tensor `X@GRAD` will use the dtype of forward Tensor "
+          "`X` by default."));
   info_ptr_->SetInferDtypeFn(std::forward<InferDtypeFunc>(func));
   return *this;
 }
 
-OpMetaInfoBuilder& OpMetaInfoBuilder::SetBackwardOp(
-    const std::string& bwd_op_name) {
-  auto& info_vector = OpMetaInfoMap::Instance()[name_];
-  auto op_meta = OpMetaInfo(bwd_op_name);
-  info_vector.emplace_back(std::move(op_meta));
-  info_ptr_ = &(info_vector.back());
-  return *this;
-}
-
 /////////////////////// Op register API /////////////////////////
 
 void RegisterAllCustomOperator() {

paddle/fluid/framework/custom_operator.cc

@@ -153,12 +153,21 @@ static void RunKernelFunc(const framework::ExecutionContext& ctx,
   }
 
   VLOG(1) << "Run ComputeFunc.";
-  auto outs = func(custom_ins, custom_attrs);
+  try {
+    auto outs = func(custom_ins, custom_attrs);
 
-  VLOG(1) << "Custom Operator: Share outputs into ExecutionContext.";
-  for (size_t i = 0; i < outputs.size(); ++i) {
-    auto* true_out = ctx.Output<Tensor>(outputs[i]);
-    CustomTensorUtils::ShareDataTo(outs.at(i), true_out);
+    VLOG(1) << "Custom Operator: Share outputs into ExecutionContext.";
+    for (size_t i = 0; i < outputs.size(); ++i) {
+      auto* true_out = ctx.Output<Tensor>(outputs[i]);
+      CustomTensorUtils::ShareDataTo(outs.at(i), true_out);
+    }
+  } catch (platform::EnforceNotMet& exception) {
+    throw std::move(exception);
+  } catch (std::exception& ex) {
+    PADDLE_THROW(platform::errors::External("%s", ex.what()));
+  } catch (...) {
+    PADDLE_THROW(platform::errors::Fatal(
+        "Custom operator raises an unknown exception in rumtime."));
   }
 }
 
@@ -475,58 +484,108 @@ void RegisterOperatorWithMetaInfo(
           op_name, info.proto_->InitializationErrorString()));
 
   // InferShape
-  PADDLE_ENFORCE_NOT_NULL(
-      infer_shape_func,
-      platform::errors::PreconditionNotMet(
-          "InferShapeFn is nullptr. Need to set the InferShapeFn of custom "
-          "operator by .SetInferShapeFn(PD_INFER_SHAPE(...))"));
-  info.infer_shape_ = [op_inputs, op_outputs,
-                       infer_shape_func](InferShapeContext* ctx) {
-    std::vector<std::vector<int64_t>> input_shapes;
-
-    VLOG(1) << "Custom Operator: InferShape - get input ddim.";
-    for (auto& in_name : op_inputs) {
-      OP_INOUT_CHECK(ctx->HasInput(in_name), "Input", in_name, "Custom");
-      auto ddim = ctx->GetInputDim(in_name);
-      input_shapes.emplace_back(framework::vectorize(ddim));
-    }
+  if (infer_shape_func == nullptr) {
+    // use default InferShape
+    info.infer_shape_ = [op_inputs, op_outputs](InferShapeContext* ctx) {
+      PADDLE_ENFORCE_EQ(
+          op_inputs.size(), 1UL,
+          platform::errors::Unavailable(
+              "Your custom operator contains multiple inputs. "
+              "We only allow a custom operator that contains only one input "
+              "and "
+              "only one output without setting the InferShapeFn. At this time, "
+              "the input shape will be directly set to the output shape.\n"
+              "Please set the InferShapeFn of custom "
+              "operator by .SetInferShapeFn(PD_INFER_SHAPE(...))"));
+      PADDLE_ENFORCE_EQ(
+          op_outputs.size(), 1UL,
+          platform::errors::Unavailable(
+              "Your custom operator contains multiple outputs. "
+              "We only allow a custom operator that contains only one input "
+              "and "
+              "only one output without setting the InferShapeFn. At this time, "
+              "the input shape will be directly set to the output shape.\n"
+              "Please set the InferShapeFn of custom "
+              "operator by .SetInferShapeFn(PD_INFER_SHAPE(...))"));
+
+      VLOG(1) << "Custom Operator: Default InferShape - share ddim.";
+      ctx->ShareDim(op_inputs[0], op_outputs[0]);
+    };
+  } else {
+    info.infer_shape_ = [op_inputs, op_outputs,
+                         infer_shape_func](InferShapeContext* ctx) {
+      std::vector<std::vector<int64_t>> input_shapes;
+
+      VLOG(1) << "Custom Operator: InferShape - get input ddim.";
+      for (auto& in_name : op_inputs) {
+        OP_INOUT_CHECK(ctx->HasInput(in_name), "Input", in_name, "Custom");
+        auto ddim = ctx->GetInputDim(in_name);
+        input_shapes.emplace_back(framework::vectorize(ddim));
+      }
 
-    VLOG(1) << "Custom Operator: InferShape - calc output ddim.";
-    auto output_shapes = infer_shape_func(input_shapes);
+      VLOG(1) << "Custom Operator: InferShape - calc output ddim.";
+      auto output_shapes = infer_shape_func(input_shapes);
 
-    VLOG(1) << "Custom Operator: InferShape - set output ddim.";
-    for (size_t i = 0; i < op_outputs.size(); ++i) {
-      ctx->SetOutputDim(op_outputs[i], framework::make_ddim(output_shapes[i]));
-    }
-  };
+      VLOG(1) << "Custom Operator: InferShape - set output ddim.";
+      for (size_t i = 0; i < op_outputs.size(); ++i) {
+        ctx->SetOutputDim(op_outputs[i],
+                          framework::make_ddim(output_shapes[i]));
+      }
+    };
+  }
 
   // Infer Dtype
-  PADDLE_ENFORCE_NOT_NULL(
-      infer_dtype_func,
-      platform::errors::PreconditionNotMet(
-          "InferDtypeFn is nullptr. Need to set the InferDtypeFn of custom "
-          "operator by .SetInferDtypeFn(PD_INFER_DTYPE(...))"));
-  info.infer_var_type_ = [op_inputs, op_outputs,
-                          infer_dtype_func](InferVarTypeContext* ctx) {
-    std::vector<DataType> input_dtypes;
-
-    VLOG(1) << "Custom Operator: InferDtype - get input dtype.";
-    for (auto& in_name : op_inputs) {
-      auto dtype = ctx->GetInputDataType(in_name);
-      input_dtypes.emplace_back(
-          CustomTensorUtils::ConvertInnerDTypeToEnumDType(dtype));
-    }
+  if (infer_dtype_func == nullptr) {
+    // use defalut InferDtype
+    info.infer_var_type_ = [op_inputs, op_outputs](InferVarTypeContext* ctx) {
+      PADDLE_ENFORCE_EQ(
+          op_inputs.size(), 1UL,
+          platform::errors::Unavailable(
+              "Your custom operator contains multiple inputs. "
+              "We only allow a custom operator that contains only one input "
+              "and "
+              "only one output without setting the InferDtypeFn. At this time, "
+              "the input dtype will be directly set to the output dtype.\n"
+              "Please set the InferDtypeFn of custom "
+              "operator by .SetInferDtypeFn(PD_INFER_DTYPE(...))"));
+      PADDLE_ENFORCE_EQ(
+          op_outputs.size(), 1UL,
+          platform::errors::Unavailable(
+              "Your custom operator contains multiple outputs. "
+              "We only allow a custom operator that contains only one input "
+              "and "
+              "only one output without setting the InferDtypeFn. At this time, "
+              "the input dtype will be directly set to the output dtype.\n"
+              "Please set the InferDtypeFn of custom "
+              "operator by .SetInferDtypeFn(PD_INFER_DTYPE(...))"));
+
+      VLOG(1) << "Custom Operator: InferDtype - share dtype.";
+      auto dtype = ctx->GetInputDataType(op_inputs[0]);
+      ctx->SetOutputDataType(op_outputs[0], dtype);
+    };
+  } else {
+    info.infer_var_type_ = [op_inputs, op_outputs,
+                            infer_dtype_func](InferVarTypeContext* ctx) {
+      std::vector<DataType> input_dtypes;
+
+      VLOG(1) << "Custom Operator: InferDtype - get input dtype.";
+      for (auto& in_name : op_inputs) {
+        auto dtype = ctx->GetInputDataType(in_name);
+        input_dtypes.emplace_back(
+            CustomTensorUtils::ConvertInnerDTypeToEnumDType(dtype));
+      }
 
-    VLOG(1) << "Custom Operator: InferDtype - infer output dtype.";
-    auto output_dtypes = infer_dtype_func(input_dtypes);
+      VLOG(1) << "Custom Operator: InferDtype - infer output dtype.";
+      auto output_dtypes = infer_dtype_func(input_dtypes);
 
-    VLOG(1) << "Custom Operator: InferDtype - set output dtype.";
-    for (size_t i = 0; i < op_outputs.size(); ++i) {
-      ctx->SetOutputDataType(
-          op_outputs[i],
-          CustomTensorUtils::ConvertEnumDTypeToInnerDType(output_dtypes[i]));
-    }
-  };
+      VLOG(1) << "Custom Operator: InferDtype - set output dtype.";
+      for (size_t i = 0; i < op_outputs.size(); ++i) {
+        ctx->SetOutputDataType(
+            op_outputs[i],
+            CustomTensorUtils::ConvertEnumDTypeToInnerDType(output_dtypes[i]));
+      }
+    };
+  }
 
   // Kernel func
   RegisterOperatorKernel(op_name, kernel_fn, op_inputs, op_outputs, op_attrs);

python/paddle/fluid/tests/custom_op/CMakeLists.txt

@@ -3,10 +3,12 @@ if(WITH_GPU)
     # 'test_custom_relu_op_setup/jit' compile .cc and .cu file
     py_test(test_custom_relu_op_setup SRCS test_custom_relu_op_setup.py)
     py_test(test_custom_relu_op_jit SRCS test_custom_relu_op_jit.py)
+    py_test(test_custom_relu_model SRCS test_custom_relu_model.py)
 
     # Compiling shared library will cost some time, but running process is very fast.
     set_tests_properties(test_custom_relu_op_setup PROPERTIES TIMEOUT 250)
     set_tests_properties(test_custom_relu_op_jit PROPERTIES TIMEOUT 180)
+    set_tests_properties(test_custom_relu_model PROPERTIES TIMEOUT 180)
 endif()
 
 py_test(test_sysconfig SRCS test_sysconfig.py)