Add manual op ap_op.add. (#72701)

Author: Xreki

paddle/ap/include/paddle/hlir/manual_op.h

@@ -42,6 +42,23 @@ class IR_API FacadeOp
   bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
 };
 
+class IR_API AddOp
+    : public pir::Op<AddOp,
+                     pir::ImmutableLayoutTrait,
+                     ::paddle::dialect::InferSymbolicShapeInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "ap_op.add"; }
+  static constexpr uint32_t attributes_num = 0;
+  static constexpr const char **attributes_name = nullptr;
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value lhs,
+                    pir::Value rhs);
+  void VerifySig() const {}
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
+};
+
 class IR_API UpSpiderOp
     : public pir::Op<UpSpiderOp,
                      pir::SideEffectTrait,
@@ -152,6 +169,7 @@ class IR_API StoreToGlobalOp
 }  // namespace ap::dialect
 
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::FacadeOp);
+IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::AddOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::UpSpiderOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::DownSpiderOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::LoadFromRegisterOp);

paddle/ap/src/paddle/hlir/manual_op.cc

@@ -41,6 +41,34 @@ bool FacadeOp::InferSymbolicShape(
   return ApOpFacadeOpInferSymbolicShape(*this, infer_context);
 }
 
+void AddOp::Build(pir::Builder& builder,             // NOLINT
+                  pir::OperationArgument& argument,  // NOLINT
+                  pir::Value lhs,
+                  pir::Value rhs) {
+  argument.AddInput(lhs);
+  argument.AddInput(rhs);
+}
+
+bool AddOp::InferSymbolicShape(pir::InferSymbolicShapeContext* infer_context) {
+  const auto& lhs_shape_or_data =
+      infer_context->GetShapeOrDataForValue(operand_source(0));
+  const auto& rhs_shape_or_data =
+      infer_context->GetShapeOrDataForValue(operand_source(1));
+  PADDLE_ENFORCE_GT(lhs_shape_or_data.shape().size(),
+                    rhs_shape_or_data.shape().size(),
+                    phi::errors::InvalidArgument(
+                        "lhs and rhs of ap_op.add should have same rank"));
+  for (int i = 0; i < lhs_shape_or_data.shape().size(); ++i) {
+    const auto& lhs_dim_expr = lhs_shape_or_data.shape().at(i);
+    const auto& rhs_dim_expr = rhs_shape_or_data.shape().at(i);
+    if (lhs_dim_expr != rhs_dim_expr) {
+      infer_context->AddEqualCstr(lhs_dim_expr, rhs_dim_expr);
+    }
+  }
+  infer_context->SetShapeOrDataForValue(result(0), rhs_shape_or_data);
+  return true;
+}
+
 void UpSpiderOp::Build(pir::Builder& builder,             // NOLINT
                        pir::OperationArgument& argument,  // NOLINT
                        pir::Value lhs,
@@ -165,6 +193,7 @@ bool StoreToGlobalOp::InferSymbolicShape(
 }  // namespace ap::dialect
 
 IR_DEFINE_EXPLICIT_TYPE_ID(ap::dialect::FacadeOp);
+IR_DEFINE_EXPLICIT_TYPE_ID(ap::dialect::AddOp);
 IR_DEFINE_EXPLICIT_TYPE_ID(ap::dialect::UpSpiderOp);
 IR_DEFINE_EXPLICIT_TYPE_ID(ap::dialect::DownSpiderOp);
 IR_DEFINE_EXPLICIT_TYPE_ID(ap::dialect::LoadFromRegisterOp);

paddle/ap/src/paddle/hlir/op_dialect.cc

@@ -26,6 +26,7 @@ OperatorDialect::OperatorDialect(::pir::IrContext *context)
 
 void OperatorDialect::initialize() {
   RegisterOp<FacadeOp>();
+  RegisterOp<AddOp>();
   RegisterOp<UpSpiderOp>();
   RegisterOp<DownSpiderOp>();
   RegisterOp<LoadFromRegisterOp>();

paddle/ap/src/paddle/pass/op_factory.cc

@@ -36,6 +36,17 @@ adt::Result<pir::Operation*> ConstructPdOpSum(
   return op;
 }
 
+adt::Result<pir::Operation*> ConstructAddOp(
+    pir::Builder* builder,
+    const std::vector<pir::Value>& inputs,
+    const pir::AttributeMap& attrs) {
+  ADT_CHECK(inputs.size() == 2) << adt::errors::TypeError{
+      std::string() + "'ap_op.add' op takes 2 arguments, but " +
+      std::to_string(inputs.size()) + " were given"};
+  auto op = builder->Build<ap::dialect::AddOp>(inputs.at(0), inputs.at(1));
+  return op;
+}
+
 adt::Result<pir::Operation*> ConstructUpSpiderOp(
     pir::Builder* builder,
     const std::vector<pir::Value>& inputs,
@@ -186,6 +197,10 @@ adt::Result<std::optional<pir::Operation*>> CreateOperation(
     ADT_LET_CONST_REF(ret, ConstructShadowOutputOp(builder, inputs, attrs));
     return ret;
   }
+  if (op_name == "ap_op.add") {
+    ADT_LET_CONST_REF(ret, ConstructAddOp(builder, inputs, attrs));
+    return ret;
+  }
   if (op_name == "ap_op.up_spider") {
     ADT_LET_CONST_REF(ret, ConstructUpSpiderOp(builder, inputs, attrs));
     return ret;