fix conflicts with upstream

paddle/common/ddim.h

@@ -64,7 +64,7 @@ inline void dynamic_dim_assign(const T1* in, T2* out, int n) {
  *
  * The number of dimensions must be between [1, 9].
  */
-class DDim {
+class TEST_API DDim {
  public:
   constexpr static int kMaxRank = 9;
 

paddle/fluid/framework/details/CMakeLists.txt

@@ -317,8 +317,8 @@ cc_test_old(
   var_handle
   op_handle_base
   scope
-  phi
   memory
+  phi
   device_context
   gather_op_handle)
 

paddle/fluid/pir/dialect/op_generator/api_gen.py

@@ -21,9 +21,7 @@
     PD_MANUAL_OP_LIST,
     OpCompatParser,
     OpInfoParser,
-    check_need_update_ops,
     to_pascal_case,
-    update_ops,
 )
 
 PD_MANUAL_API_LIST = {
@@ -164,18 +162,12 @@ def __init__(self) -> None:
 
     def _parse_yaml(self, op_yaml_files, op_compat_yaml_file):
         op_compat_parser = OpCompatParser(op_compat_yaml_file)
-        need_update_ops, update_yaml_file = check_need_update_ops(op_yaml_files)
 
         op_yaml_items = []
         for yaml_file in op_yaml_files:
-            if update_yaml_file == yaml_file:
-                continue
             with open(yaml_file, "r") as f:
                 ops = yaml.safe_load(f)
                 op_yaml_items = op_yaml_items + ops
-        # replace old ir ops with pir ops
-        if need_update_ops:
-            update_ops(op_yaml_items, update_yaml_file)
 
         op_info_items = []
         for op in op_yaml_items:

paddle/fluid/pir/dialect/op_generator/op_gen.py

@@ -1055,18 +1055,12 @@ def OpGenerator(
 
     # (2) Prepare: Get all op item in all op_yaml_files
     op_compat_parser = OpCompatParser(op_compat_yaml_file)
-    need_update_ops, update_yaml_file = check_need_update_ops(op_yaml_files)
 
     op_yaml_items = []
     for yaml_file in op_yaml_files:
-        if update_yaml_file == yaml_file:
-            continue
         with open(yaml_file, "r") as f:
             ops = yaml.safe_load(f)
             op_yaml_items = op_yaml_items + ops
-    # replace old ir ops with pir ops
-    if need_update_ops:
-        update_ops(op_yaml_items, update_yaml_file)
 
     op_info_items = {}
     for op in op_yaml_items:

paddle/fluid/pir/dialect/operator/ir/ops.yaml

@@ -1,5 +1,7 @@
-# The apis in this file are unstandardized that may caused by a variety of reasons,
-# we are trying to fix these apis and will move standardized apis into ops.yaml.
+# The operators included in this file are:
+# 1) Operators defined only in PIR, dynamic graphs do not exist;
+# 2) The definitions of static graphs and dynamic graphs are inconsistent, but the final definition plan has not yet been clarified.
+# After the definition is clearly defined, migrate to paddle /fluid/pir/dialect/operator/ir/update_ops.yaml or paddle/phi/api/yaml/ops.yaml
 
 - op : adadelta_
   args : (Tensor param, Tensor grad, Tensor avg_squared_grad, Tensor avg_squared_update, Tensor learning_rate, Tensor master_param, float rho, float epsilon, bool multi_precision)
@@ -83,19 +85,6 @@
   kernel :
     func : any
 
-- op : arange
-  args : (Scalar start, Scalar end, Scalar step, DataType dtype=DataType::FLOAT64, Place place=CPUPlace())
-  output : Tensor(out)
-  infer_meta :
-    func : ArangeInferMeta
-    param : [start, end, step, dtype]
-  kernel :
-    func : arange
-    param : [start, end, step]
-    data_type : dtype
-    backend : place
-  support_tensor : [start, end, step]
-
 - op : assert
   args : (Tensor cond, Tensor[] data, int64_t summarize = -1)
   output :

paddle/fluid/pir/dialect/operator/ir/ops_backward.yaml

@@ -1,3 +1,8 @@
+# The operators included in this file are:
+# 1) Operators defined only in PIR, dynamic graphs do not exist;
+# 2) The definitions of static graphs and dynamic graphs are inconsistent, but the final definition plan has not yet been clarified.
+# After the definition is clearly defined, migrate to paddle /fluid/pir/dialect/operator/ir/update_ops.yaml or paddle/phi/api/yaml/ops.yaml
+
 - backward_op : add_double_grad
   forward : add_grad (Tensor x, Tensor y, Tensor grad_out, int axis = -1) -> Tensor(grad_x), Tensor(grad_y)
   args : (Tensor y, Tensor grad_out, Tensor grad_x_grad, Tensor grad_y_grad, int axis = -1)

paddle/fluid/pir/dialect/operator/ir/update_ops.yaml

@@ -1,4 +1,7 @@
-# Ops in this file is only used for pir currently and will replace ops of legacy_ops.yaml/ops.yaml of PHI in future.
+# The operators contained in this file are:
+# Operators that are inconsistent with the dynamic graph definition currently,
+# but the final definition scheme of the static graph has been determined, after
+# the dynamic graph is simultaneously upgraded, the operators in this file will be migrated to paddle/phi/api/yaml/ops.yaml.
 
 - op : arange
   args : (Scalar start, Scalar end, Scalar step, DataType dtype=DataType::FLOAT64, Place place=CPUPlace())

paddle/fluid/pir/drr/drr_rewrite_pattern.cc

@@ -401,7 +401,7 @@ MatchContextImpl DrrRewritePattern::CreateOperations(
           max_index_op = ir_input_op;
         } else if (max_input_op_index == op_2_temp_program_index[ir_input_op]) {
           const auto& ops_vec = temp_program[max_input_op_index];
-          for (auto it = ops_vec.rbegin(); it != ops_vec.rend(); it++) {
+          for (auto it = ops_vec.begin(); it != ops_vec.end(); it++) {
             if (*it == max_index_op) {
               break;
             } else if (*it == ir_input_op) {

paddle/fluid/pir/transforms/fusion/fused_linear_param_grad_add_pass.cc

@@ -26,15 +26,22 @@ class FusedMatmulAddGradAddPattern
  public:
   void operator()(pir::drr::DrrPatternContext *ctx) const override {
     pir::drr::SourcePattern pat = ctx->SourcePattern();
+    const auto &matmul0 = pat.Op(paddle::dialect::MatmulOp::name(),
+                                 {{"transpose_x", pat.Attr("trans_x")},
+                                  {"transpose_y", pat.Attr("trans_y")}});
+    const auto &add0 = pat.Op(paddle::dialect::AddOp::name());
     const auto &add_grad = pat.Op(paddle::dialect::AddGradOp::name());
     const auto &matmul_grad = pat.Op(paddle::dialect::MatmulGradOp::name(),
                                      {{"transpose_x", pat.Attr("trans_x")},
                                       {"transpose_y", pat.Attr("trans_y")}});
     const auto &add_ = pat.Op(paddle::dialect::Add_Op::name());
 
-    add_grad(
-        {&pat.Tensor("out"), &pat.Tensor("bias"), &pat.Tensor("addout_grad")},
-        {&pat.Tensor("out_grad"), &pat.Tensor("dbias")});
+    pat.Tensor("out") = matmul0(pat.Tensor("x"), pat.Tensor("weight"));
+    pat.Tensor("fwd_add_out") = add0(pat.Tensor("out"), pat.Tensor("bias"));
+    add_grad({&pat.Tensor("out"),
+              &pat.Tensor("bias"),
+              &pat.Tensor("fwd_add_out_grad")},
+             {&pat.Tensor("out_grad"), &pat.Tensor("dbias")});
     matmul_grad(
         {&pat.Tensor("x"), &pat.Tensor("weight"), &pat.Tensor("out_grad")},
         {&pat.Tensor("x_grad"), &pat.Tensor("weight_grad")});
@@ -47,7 +54,7 @@ class FusedMatmulAddGradAddPattern
       return (match_ctx.Tensor("weight_grad").Shape() ==
                   match_ctx.Tensor("dweight").Shape() &&
               match_ctx.Tensor("out").Shape() ==
-                  match_ctx.Tensor("addout_grad").Shape() &&
+                  match_ctx.Tensor("fwd_add_out_grad").Shape() &&
               x_trans == false && y_trans == false);
     });
 
@@ -70,10 +77,10 @@ class FusedMatmulAddGradAddPattern
         paddle::dialect::FusedLinearParamGradAddOp::name(),
         {{{"multi_precision", muti_precision_attr}, {"has_bias", true_attr}}});
 
-    matmul({&res.Tensor("addout_grad"), &res.Tensor("weight")},
+    matmul({&res.Tensor("fwd_add_out_grad"), &res.Tensor("weight")},
            {&res.Tensor("x_grad")});
     fused_linear_param_grad_add({&res.Tensor("x"),
-                                 &res.Tensor("addout_grad"),
+                                 &res.Tensor("fwd_add_out_grad"),
                                  &res.Tensor("dweight"),
                                  &res.NoneTensor()},
                                 {&res.Tensor("add_out"), &res.Tensor("dbias")});
@@ -232,15 +239,25 @@ class FusedMatmulAddGradAddaPattern
  public:
   void operator()(pir::drr::DrrPatternContext *ctx) const override {
     pir::drr::SourcePattern pat = ctx->SourcePattern();
-    const auto &add_grad = pat.Op(paddle::dialect::AddGradOp::name());
     const auto &matmul = pat.Op(paddle::dialect::MatmulOp::name(),
                                 {{"transpose_x", pat.Attr("trans_x")},
                                  {"transpose_y", pat.Attr("trans_y")}});
+    const auto &add = pat.Op(paddle::dialect::AddOp::name());
+    const auto &add_grad = pat.Op(paddle::dialect::AddGradOp::name());
+    const auto &matmul_g0 = pat.Op(paddle::dialect::MatmulOp::name(),
+                                   {{"transpose_x", pat.Attr("trans_xg0")},
+                                    {"transpose_y", pat.Attr("trans_yg0")}});
+    const auto &matmul_g1 = pat.Op(paddle::dialect::MatmulOp::name(),
+                                   {{"transpose_x", pat.Attr("trans_xg1")},
+                                    {"transpose_y", pat.Attr("trans_yg1")}});
     const auto &add_ = pat.Op(paddle::dialect::Add_Op::name());
+
+    pat.Tensor("out") = matmul(pat.Tensor("x"), pat.Tensor("weight"));
+    pat.Tensor("fwd_add_out") = add(pat.Tensor("out"), pat.Tensor("bias"));
     add_grad({&pat.Tensor("out"), &pat.Tensor("bias"), &pat.Tensor("dadd_out")},
              {&pat.Tensor("dout"), &pat.Tensor("dbias")});
-    matmul({&pat.Tensor("x"), &pat.Tensor("dout")},
-           {&pat.Tensor("weight_grad")});
+    pat.Tensor("dx") = matmul_g0(pat.Tensor("dout"), pat.Tensor("weight"));
+    pat.Tensor("weight_grad") = matmul_g1(pat.Tensor("x"), pat.Tensor("dout"));
     pat.Tensor("dweight_out") =
         add_(pat.Tensor("dweight"), pat.Tensor("weight_grad"));
 
@@ -277,15 +294,25 @@ class FusedMatmulAddGradAddbPattern
  public:
   void operator()(pir::drr::DrrPatternContext *ctx) const override {
     pir::drr::SourcePattern pat = ctx->SourcePattern();
-    const auto &add_grad = pat.Op(paddle::dialect::AddGradOp::name());
     const auto &matmul = pat.Op(paddle::dialect::MatmulOp::name(),
                                 {{"transpose_x", pat.Attr("trans_x")},
                                  {"transpose_y", pat.Attr("trans_y")}});
+    const auto &add = pat.Op(paddle::dialect::AddOp::name());
+    const auto &add_grad = pat.Op(paddle::dialect::AddGradOp::name());
+    const auto &matmul_g0 = pat.Op(paddle::dialect::MatmulOp::name(),
+                                   {{"transpose_x", pat.Attr("trans_xg0")},
+                                    {"transpose_y", pat.Attr("trans_yg0")}});
+    const auto &matmul_g1 = pat.Op(paddle::dialect::MatmulOp::name(),
+                                   {{"transpose_x", pat.Attr("trans_xg1")},
+                                    {"transpose_y", pat.Attr("trans_yg1")}});
     const auto &add_ = pat.Op(paddle::dialect::Add_Op::name());
+
+    pat.Tensor("out") = matmul(pat.Tensor("x"), pat.Tensor("weight"));
+    pat.Tensor("fwd_add_out") = add(pat.Tensor("out"), pat.Tensor("bias"));
     add_grad({&pat.Tensor("out"), &pat.Tensor("bias"), &pat.Tensor("dadd_out")},
              {&pat.Tensor("dout"), &pat.Tensor("dbias")});
-    matmul({&pat.Tensor("x"), &pat.Tensor("dout")},
-           {&pat.Tensor("weight_grad")});
+    pat.Tensor("dx") = matmul_g0(pat.Tensor("dout"), pat.Tensor("weight"));
+    pat.Tensor("weight_grad") = matmul_g1(pat.Tensor("x"), pat.Tensor("dout"));
     pat.Tensor("dweight_out") =
         add_(pat.Tensor("weight_grad"), pat.Tensor("dweight"));
 

paddle/fluid/primitive/codegen/gen.py

@@ -14,7 +14,6 @@
 
 import argparse
 import hashlib
-import os
 import pathlib
 import sys
 
@@ -353,20 +352,19 @@ def gen(
         compats,
         ir_fwds,
         ir_revs,
+        ir_update_fwds,
     ) = (
         load(prim_path),
         load(fwd_path),
         load(rev_path),
         load(compat_path),
         load(fwd_pd_op_path),
         load(rev_pd_op_path),
+        load(update_fwd_pd_op_path),
     )
     filter_compat_info(compats)
 
-    fwd_apis = fwds + ir_fwds
-    # replace old ir ops with pir ops
-    if os.path.exists(update_fwd_pd_op_path):
-        update_apis(fwd_apis, update_fwd_pd_op_path)
+    fwd_apis = fwds + ir_fwds + ir_update_fwds
 
     apis = [{**api, **{'is_fwd': True}} for api in fwd_apis]
     apis = apis + [{**api, **{'is_fwd': False}} for api in revs + ir_revs]

paddle/phi/api/lib/api_custom_impl.cc

@@ -264,12 +264,17 @@ void embedding_grad_impl(const Tensor& x,
       auto meta_dist_input_x = MakeDistMetaTensor(*x.impl());
       auto meta_dist_input_weight = MakeDistMetaTensor(*weight.impl());
       auto meta_dist_input_out_grad = MakeDistMetaTensor(*out_grad.impl());
-      auto spmd_info = phi::distributed::VariadicReplicatedInferSpmdDynamic(
-          meta_dist_input_weight, meta_dist_input_x, meta_dist_input_out_grad);
+      auto spmd_info =
+          phi::distributed::EmbeddingGradInferSpmd(meta_dist_input_x,
+                                                   meta_dist_input_weight,
+                                                   meta_dist_input_out_grad,
+                                                   padding_idx,
+                                                   sparse);
 
       // 2. Create Temporary Output & Prepare Dist and Dense Output
       std::shared_ptr<phi::distributed::DistTensor> shared_dist_out =
-          CreateKernelDistOutput(weight_grad, !rank_is_in_current_mesh);
+          CreateKernelDistOutput(
+              weight_grad, !rank_is_in_current_mesh, spmd_info.second[0]);
       phi::distributed::DistTensor* dist_out = shared_dist_out.get();
       phi::DenseTensor* dense_out = dist_out->unsafe_mutable_value();
       if (dense_out && !rank_is_in_current_mesh && !dist_out->defined()) {
@@ -284,37 +289,29 @@ void embedding_grad_impl(const Tensor& x,
       UnchangedInferMeta(MakeMetaTensor(*weight.impl()), &meta_dist_out);
 
       // 4. Set Output Dist Attr For Default Impl
-      auto current_process_mesh =
-          paddle::holds_alternative<phi::distributed::TensorDistAttr>(
-              spmd_info.first[0])
-              ? paddle::get<0>(spmd_info.first[0]).process_mesh()
-              : paddle::get<1>(spmd_info.first[0]).at(0).process_mesh();
-      SetReplicatedDistAttrForOutput(dist_out, current_process_mesh);
 
       if (rank_is_in_current_mesh) {
         // 5. Reshard Input
-        auto dist_input_weight =
-            ReshardApiInputToKernelInput(dev_ctx, weight, spmd_info.first[0]);
         auto dist_input_x =
-            ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[1]);
+            ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[0]);
+        auto dist_input_weight =
+            ReshardApiInputToKernelInput(dev_ctx, weight, spmd_info.first[1]);
         auto dist_input_out_grad =
             ReshardApiInputToKernelInput(dev_ctx, out_grad, spmd_info.first[2]);
 
         // 6. PrepareData (DataTransform & Prepare Dense Input)
-        dist_input_weight = PrepareDataForDistTensor(
-            dist_input_weight,
+        dist_input_x = PrepareDataForDistTensor(
+            dist_input_x,
             GetKernelInputArgDef(kernel.InputAt(0), kernel_key.backend()),
             {},
             kernel_result.is_stride_kernel);
-        auto input_weight = &dist_input_weight->value();
-
-        dist_input_x = PrepareDataForDistTensor(
-            dist_input_x,
+        auto input_x = &dist_input_x->value();
+        dist_input_weight = PrepareDataForDistTensor(
+            dist_input_weight,
             GetKernelInputArgDef(kernel.InputAt(1), kernel_key.backend()),
             {},
             kernel_result.is_stride_kernel);
-        auto input_x = &dist_input_x->value();
-
+        auto input_weight = &dist_input_weight->value();
         dist_input_out_grad = PrepareDataForDistTensor(
             dist_input_out_grad,
             GetKernelInputArgDef(kernel.InputAt(2), kernel_key.backend()),

paddle/phi/api/yaml/legacy_ops.yaml

@@ -34,6 +34,7 @@
   output : Tensor(out)
   infer_meta :
     func : ReduceInferMeta
+    spmd_rule : ReductionAllInferSpmdDynamic
   kernel :
     func : all
 
@@ -361,6 +362,7 @@
   infer_meta :
     func : EmbeddingInferMeta
     param : [x, weight, padding_idx]
+    spmd_rule: EmbeddingInferSpmd
   kernel :
     func : embedding {dense, dense -> dense}
            sparse_weight_embedding {dense, selected_rows -> dense}

paddle/phi/backends/xpu/xpu2_op_list.cc

@@ -595,7 +595,10 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"mean_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"mean", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"mean",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"merged_adam", XPUKernelSet({phi::DataType::FLOAT32})},
       {"merged_momentum",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},

paddle/phi/backends/xpu/xpu3_op_list.cc

@@ -571,7 +571,10 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"mean_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"mean", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"mean",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"merged_momentum",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"mish_grad", XPUKernelSet({phi::DataType::FLOAT32})},

paddle/phi/core/ddim.h

@@ -248,7 +248,7 @@ DDim slice_ddim(const DDim& dim, int begin, int end);
 
 int arity(const DDim& ddim);
 
-TEST_API std::ostream& operator<<(std::ostream&, const DDim&);
+std::ostream& operator<<(std::ostream&, const DDim&);
 
 /**
  * \brief Flatten dim to 3d