mean axis support int_array (#70926)

wanghuancoder · web-flow · commit 6be721c586e6 · 2025-02-06T10:41:26.000+08:00
* mean axis support int_array
diff --git a/paddle/fluid/pir/transforms/gpu/add_norm_fuse_pass.cc b/paddle/fluid/pir/transforms/gpu/add_norm_fuse_pass.cc
@@ -42,8 +42,9 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     paddle::drr::SourcePattern pat = ctx->SourcePattern();
     const auto &pow = pat.Op(paddle::dialect::PowOp::name());
-    const auto &mean =
-        pat.Op(paddle::dialect::MeanOp::name(), {{"axis", pat.Attr("axis")}});
+    const auto &full_int_array = pat.Op(paddle::dialect::FullIntArrayOp::name(),
+                                        {{"value", pat.Attr("value")}});
+    const auto &mean = pat.Op(paddle::dialect::MeanOp::name());
     const auto &full = pat.Op(paddle::dialect::FullOp::name());
     const auto &scale =
         pat.Op(paddle::dialect::ScaleOp::name(), {{"bias", pat.Attr("bias")}});
@@ -55,7 +56,9 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
                                  {{"dtype", pat.Attr("cast_type_1")}});
       pat.Tensor("cast_1_out") = cast1(pat.Tensor("x"));
       pat.Tensor("pow_out") = pow(pat.Tensor("cast_1_out"));
-      pat.Tensor("mean_out") = mean(pat.Tensor("pow_out"));
+      pat.Tensor("full_int_array_out") = full_int_array();
+      pat.Tensor("mean_out") =
+          mean(pat.Tensor("pow_out"), pat.Tensor("full_int_array_out"));
       pat.Tensor("scale_out") = scale(pat.Tensor("mean_out"), full());
       pat.Tensor("rsqrt_out") = rsqrt(pat.Tensor("scale_out"));
       pat.Tensor("multiply_out1") =
@@ -67,7 +70,9 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
           multiply2(pat.Tensor("cast_2_out"), pat.Tensor("w"));
     } else {
       pat.Tensor("pow_out") = pow(pat.Tensor("x"));
-      pat.Tensor("mean_out") = mean(pat.Tensor("pow_out"));
+      pat.Tensor("full_int_array_out") = full_int_array();
+      pat.Tensor("mean_out") =
+          mean(pat.Tensor("pow_out"), pat.Tensor("full_int_array_out"));
       pat.Tensor("scale_out") = scale(pat.Tensor("mean_out"), full());
       pat.Tensor("rsqrt_out") = rsqrt(pat.Tensor("scale_out"));
       pat.Tensor("multiply_out1") =
@@ -76,7 +81,7 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
           multiply2(pat.Tensor("multiply_out1"), pat.Tensor("w"));
     }
     pat.AddConstraint([this](const paddle::drr::MatchContext &match_ctx) {
-      auto axis = match_ctx.Attr<std::vector<int64_t>>("axis");
+      auto axis = match_ctx.Attr<std::vector<int64_t>>("value");
       if (axis.size() > 1) {
         return false;
       }
@@ -107,7 +112,7 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
     paddle::drr::ResultPattern res = pat.ResultPattern();
     const auto &begin_norm_axis =
         res.ComputeAttr([](const paddle::drr::MatchContext &match_ctx) -> int {
-          const auto &axis = match_ctx.Attr<std::vector<int64_t>>("axis");
+          const auto &axis = match_ctx.Attr<std::vector<int64_t>>("value");
           auto pow_out_shape =
               pir::GetShapeFromValue(match_ctx.Tensor("pow_out"));
           return axis[0] == -1 ? static_cast<int>(pow_out_shape.size()) - 1
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -138,8 +138,7 @@ class ReduceCommonOpPattern : public pir::OpRewritePattern<OpType> {
       return false;
     }
 
-    if constexpr (std::is_same_v<OpType, paddle::dialect::MeanOp> ||
-                  std::is_same_v<OpType, paddle::dialect::AnyOp> ||
+    if constexpr (std::is_same_v<OpType, paddle::dialect::AnyOp> ||
                   std::is_same_v<OpType, paddle::dialect::AllOp>) {
       if (!op->HasAttribute("axis")) {
         VLOG(3) << "The axis attribute does not exist";
diff --git a/paddle/fluid/pir/transforms/xpu/rms_norm_xpu_fuse.cc b/paddle/fluid/pir/transforms/xpu/rms_norm_xpu_fuse.cc
@@ -71,8 +71,9 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     paddle::drr::SourcePattern pat = ctx->SourcePattern();
     const auto &pow = pat.Op(paddle::dialect::PowOp::name());
-    const auto &mean =
-        pat.Op(paddle::dialect::MeanOp::name(), {{"axis", pat.Attr("axis")}});
+    const auto &full_int_array = pat.Op(paddle::dialect::FullIntArrayOp::name(),
+                                        {{"value", pat.Attr("value")}});
+    const auto &mean = pat.Op(paddle::dialect::MeanOp::name());
     const auto &full = pat.Op(paddle::dialect::FullOp::name());
     const auto &scale =
         pat.Op(paddle::dialect::ScaleOp::name(), {{"bias", pat.Attr("bias")}});
@@ -87,7 +88,9 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
       pat.Tensor("cast_1_out") = cast1(pat.Tensor("x"));
       pat.Tensor("cast_3_out") = cast3(pat.Tensor("x"));
       pat.Tensor("pow_out") = pow(pat.Tensor("cast_1_out"));
-      pat.Tensor("mean_out") = mean(pat.Tensor("pow_out"));
+      pat.Tensor("full_int_array_out") = full_int_array();
+      pat.Tensor("mean_out") =
+          mean(pat.Tensor("pow_out"), pat.Tensor("full_int_array_out"));
       pat.Tensor("scale_out") = scale(pat.Tensor("mean_out"), full());
       pat.Tensor("rsqrt_out") = rsqrt(pat.Tensor("scale_out"));
       pat.Tensor("multiply_out1") =
@@ -99,7 +102,9 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
           multiply2(pat.Tensor("cast_2_out"), pat.Tensor("w"));
     } else {
       pat.Tensor("pow_out") = pow(pat.Tensor("x"));
-      pat.Tensor("mean_out") = mean(pat.Tensor("pow_out"));
+      pat.Tensor("full_int_array_out") = full_int_array();
+      pat.Tensor("mean_out") =
+          mean(pat.Tensor("pow_out"), pat.Tensor("full_int_array_out"));
       pat.Tensor("scale_out") = scale(pat.Tensor("mean_out"), full());
       pat.Tensor("rsqrt_out") = rsqrt(pat.Tensor("scale_out"));
       pat.Tensor("multiply_out1") =
@@ -109,7 +114,7 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
     }
 
     pat.AddConstraint([this](const paddle::drr::MatchContext &match_ctx) {
-      auto axis = match_ctx.Attr<std::vector<int64_t>>("axis");
+      auto axis = match_ctx.Attr<std::vector<int64_t>>("value");
       if (axis.size() > 1) {
         return false;
       }
@@ -141,7 +146,7 @@ class RmsNormFusePattern : public paddle::drr::DrrPatternBase {
 
     const auto &begin_norm_axis =
         res.ComputeAttr([](const paddle::drr::MatchContext &match_ctx) -> int {
-          const auto &axis = match_ctx.Attr<std::vector<int64_t>>("axis");
+          const auto &axis = match_ctx.Attr<std::vector<int64_t>>("value");
           auto pow_out_shape =
               pir::GetShapeFromValue(match_ctx.Tensor("pow_out"));
           return axis[0] == -1 ? static_cast<int>(pow_out_shape.size()) - 1
diff --git a/paddle/phi/kernels/reduce_kernel_impl.cc b/paddle/phi/kernels/reduce_kernel_impl.cc
@@ -20,16 +20,11 @@ namespace phi {
 // oneDNN's reduction kernel is optimized only for reducing throughout the
 // most outer dims, so in case of another type of reduction, it would be
 // better to fallback to native implementation
-inline bool HasOptimizedOneDNNKernel(const KernelContext* ctx,
-                                     const bool mean_op) {
+inline bool HasOptimizedOneDNNKernel(const KernelContext* ctx) {
   const DenseTensor& x = ctx->InputAt<phi::DenseTensor>(0);
   IntArray dims_array;
-  if (mean_op) {
-    dims_array = ctx->AttrAt<IntArray>(0);
-  } else {
-    const TensorRef& dims_tmp = ctx->AttrAt<TensorRef>(0);
-    dims_array = IntArray(*dims_tmp.Get());
-  }
+  const TensorRef& dims_tmp = ctx->AttrAt<TensorRef>(0);
+  dims_array = IntArray(*dims_tmp.Get());
   int ndims = x.dims().size();
   const bool reduce_all = recompute_reduce_all(x, dims_array);
   auto dims = dims_array.GetData();
@@ -61,15 +56,16 @@ inline bool HasOptimizedOneDNNKernel(const KernelContext* ctx,
 
 bool ReduceCheckIfOneDNNSupport(const KernelContext* ctx) {
   if (ctx->InputAt<phi::DenseTensor>(0).dims().size() > 5 ||
-      !HasOptimizedOneDNNKernel(ctx, false)) {
+      !HasOptimizedOneDNNKernel(ctx)) {
     return false;
   }
   return true;
 }
 
 bool ReduceMeanCheckIfOneDNNSupport(const KernelContext* ctx) {
+  std::cout << "ReduceMeanCheckIfOneDNNSupport" << std::endl;
   if (ctx->InputAt<phi::DenseTensor>(0).dims().size() > 5 ||
-      !HasOptimizedOneDNNKernel(ctx, true)) {
+      !HasOptimizedOneDNNKernel(ctx)) {
     return false;
   }
   return true;
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -2537,6 +2537,10 @@
     {axis : dim, keepdim : keep_dim}
   extra :
     attrs : [bool use_mkldnn = false]
+  int_array:
+      axis :
+        data_type : int
+        support_tensor : true
 
 - op : mean_all (mean)
   backward : mean_all_grad (mean_grad)
diff --git a/python/paddle/distributed/auto_parallel/static/tuner/to_distributed_api_patterns.py b/python/paddle/distributed/auto_parallel/static/tuner/to_distributed_api_patterns.py
@@ -864,13 +864,13 @@ def build(self):
         down_linear_dist_infos = MpDistInfos("row")
         # # # build ops dist infos # # #
         ops_dist_infos = {
-            (21,): qkv_linear_dist_infos,
             (22,): qkv_linear_dist_infos,
             (23,): qkv_linear_dist_infos,
-            (88,): out_linear_dist_infos,
-            (97,): up_linear_dist_infos,
-            (98,): up_linear_dist_infos,
-            (100,): down_linear_dist_infos,
+            (24,): qkv_linear_dist_infos,
+            (89,): out_linear_dist_infos,
+            (99,): up_linear_dist_infos,
+            (100,): up_linear_dist_infos,
+            (102,): down_linear_dist_infos,
         }
         self.ops_dist_infos = ops_dist_infos
 
diff --git a/python/paddle/tensorrt/impls/math.py b/python/paddle/tensorrt/impls/math.py
@@ -278,6 +278,11 @@ def sum_converter(network, paddle_op, inputs):
     return add_reduce_layer(network, paddle_op, inputs, trt.ReduceOperation.SUM)
 
 
+@converter_registry.register("pd_op.mean", trt_version="8.x")
+def mean_converter(network, paddle_op, inputs):
+    return add_reduce_layer(network, paddle_op, inputs, trt.ReduceOperation.AVG)
+
+
 @converter_registry.register("pd_op.any", trt_version="8.x")
 def any_converter(network, paddle_op, inputs):
     return add_cast_reduce_layer(
diff --git a/python/paddle/tensorrt/impls/stat.py b/python/paddle/tensorrt/impls/stat.py
@@ -12,22 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import tensorrt as trt
+# import tensorrt as trt
 
-from paddle.tensorrt.converter_utils import get_axes_for_reduce_op
-from paddle.tensorrt.register import converter_registry
+# from paddle.tensorrt.converter_utils import get_axes_for_reduce_op
+# from paddle.tensorrt.register import converter_registry
 
 
-@converter_registry.register("pd_op.mean", trt_version="trt_version_ge=8.0")
-def mean_converter(network, paddle_op, inputs):
-    input_tensor = inputs[0]
-    keep_dim = paddle_op.attrs().get("keepdim")
-    dim = paddle_op.attrs().get("axis")
+# @converter_registry.register("pd_op.mean", trt_version="trt_version_ge=8.0")
+# def mean_converter(network, paddle_op, inputs):
+#     input_tensor = inputs[0]
+#     keep_dim = paddle_op.attrs().get("keepdim")
+#     dim = paddle_op.attrs().get("axis")
 
-    mean_layer = network.add_reduce(
-        input_tensor,
-        trt.ReduceOperation.AVG,
-        axes=get_axes_for_reduce_op(dim, network.has_implicit_batch_dimension),
-        keep_dims=keep_dim,
-    )
-    return mean_layer.get_output(0)
+#     mean_layer = network.add_reduce(
+#         input_tensor,
+#         trt.ReduceOperation.AVG,
+#         axes=get_axes_for_reduce_op(dim, network.has_implicit_batch_dimension),
+#         keep_dims=keep_dim,
+#     )
+#     return mean_layer.get_output(0)
diff --git a/test/auto_parallel/pir/auto_parallel_recompute_pir_pass_unittest.py b/test/auto_parallel/pir/auto_parallel_recompute_pir_pass_unittest.py
@@ -228,6 +228,7 @@ def run_test_cases(self):
         segment_num_base, bwd_rc_op_num_base = self.get_recompute_message(
             base_prog
         )
+
         segment_num_1, bwd_rc_op_num_1 = self.get_recompute_message(prog_1)
         segment_num_2, bwd_rc_op_num_2 = self.get_recompute_message(prog_2)
         segment_num_3, bwd_rc_op_num_3 = self.get_recompute_message(prog_3)
@@ -240,7 +241,7 @@ def run_test_cases(self):
 
         # check recompute op number
         assert bwd_rc_op_num_base == 0
-        assert bwd_rc_op_num_1 == 288
+        assert bwd_rc_op_num_1 == 296
         assert bwd_rc_op_num_2 == 204
         assert bwd_rc_op_num_3 == 60
 
diff --git a/test/auto_parallel/pir/auto_parallel_refined_recompute_pir_pass_unittest.py b/test/auto_parallel/pir/auto_parallel_refined_recompute_pir_pass_unittest.py
@@ -66,6 +66,7 @@ def run_test_cases(self):
         segment_num_base, bwd_rc_op_num_base = self.get_recompute_message(
             base_prog
         )
+
         segment_num_0, bwd_rc_op_num_0 = self.get_recompute_message(prog_0)
         segment_num_1, bwd_rc_op_num_1 = self.get_recompute_message(prog_1)
         segment_num_2, bwd_rc_op_num_2 = self.get_recompute_message(prog_2)
@@ -78,9 +79,9 @@ def run_test_cases(self):
 
         # check recompute op number
         assert bwd_rc_op_num_base == 0
-        assert bwd_rc_op_num_0 == 288
-        assert bwd_rc_op_num_1 == 284
-        assert bwd_rc_op_num_2 == 280
+        assert bwd_rc_op_num_0 == 296
+        assert bwd_rc_op_num_1 == 292
+        assert bwd_rc_op_num_2 == 288
 
         # memory check
         assert max_mem_reserved_0 < max_mem_reserved_1
diff --git a/test/auto_parallel/pir/mlp_demo_3d.py b/test/auto_parallel/pir/mlp_demo_3d.py
@@ -139,6 +139,7 @@ def test_to_static_program(self):
             'pd_op.relu',
             'pd_op.subtract',
             'pd_op.square',
+            'pd_op.full_int_array',
             'pd_op.mean',
             'builtin.shadow_output',
             'pd_op.full',
diff --git a/test/auto_parallel/static_reshard_api_cross_mesh.py b/test/auto_parallel/static_reshard_api_cross_mesh.py
@@ -190,6 +190,7 @@ def test_reshard_mesh(self):
                 'pd_op.add',
                 'pd_op.subtract',
                 'pd_op.square',
+                'pd_op.full_int_array',
                 'pd_op.mean',
                 'builtin.shadow_output',
                 'pd_op.full',
diff --git a/test/collective/test_communication_api_base.py b/test/collective/test_communication_api_base.py
@@ -111,3 +111,7 @@ def gen_product_envs_list(default_envs, changeable_envs):
         envs.update(default_envs)
         envs_list.append(envs)
     return envs_list
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/cpp/pir/core/program_translator_test.cc b/test/cpp/pir/core/program_translator_test.cc
@@ -71,7 +71,7 @@ TEST(OperatorDialectTest, MainProgram) {
   // ops.size() = op size in BlockDesc + parameter_op + combine op + int
   // array op + full op (Note: p already has a full)
   EXPECT_EQ(program->block()->size(),
-            p.Block(0).OpSize() + program->parameters_num() + 20 + 5 + 9);
+            p.Block(0).OpSize() + program->parameters_num() + 20 + 5 + 9 + 2);
   EXPECT_GT(ss.str().size(), 0u);
 }
 
diff --git a/test/cpp/prim/test_vjp.cc b/test/cpp/prim/test_vjp.cc
@@ -180,7 +180,8 @@ TEST(VJP, MeanBackwardTest) {
       std::vector<int64_t>{}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
 
   std::vector<std::vector<bool>> stop_gradients{{false}};
-  std::vector<std::vector<pir::Value>> inputs{{op1.out()}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()},
+                                              {op2.operand_source(1)}};
   std::vector<std::vector<pir::Value>> outputs{{op2.out()}};
   std::vector<std::vector<pir::Value>> out_grads{{op3.out()}};
 
diff --git a/test/ir/pir/fused_pass/python/test_py_add_norm_fuse_pass.py b/test/ir/pir/fused_pass/python/test_py_add_norm_fuse_pass.py
@@ -45,18 +45,22 @@ def fused_rms_norm_pattern(self):
         pat = ctx.SourcePattern()
 
         def constraint_function(match_ctx):
-            axis = match_ctx.VectorInt64Attr("axis")
+            axis = match_ctx.VectorInt64Attr("value")
             if len(axis) > 1:
                 return False
             return True
 
         # Source Pattern
         pow = pat.Op("pd_op.pow")
 
-        mean = pat.Op("pd_op.mean", {"axis": pat.Attr("axis")})
+        mean = pat.Op("pd_op.mean")
 
         full = pat.Op("pd_op.full")
 
+        full_int_array = pat.Op(
+            "pd_op.full_int_array", {"value": pat.Attr("value")}
+        )
+
         scale = pat.Op("pd_op.scale", {"bias": pat.Attr("bias")})
 
         rsqrt = pat.Op("pd_op.rsqrt")
@@ -66,7 +70,12 @@ def constraint_function(match_ctx):
         # Operation connections
         pow([pat.Tensor("x")], [pat.Tensor("pow_out")])
 
-        mean([pat.Tensor("pow_out")], [pat.Tensor("mean_out")])
+        full_int_array([], [pat.Tensor("full_int_array_out")])
+
+        mean(
+            [pat.Tensor("pow_out"), pat.Tensor("full_int_array_out")],
+            [pat.Tensor("mean_out")],
+        )
 
         full([], [pat.Tensor("full_out")])
 
@@ -93,7 +102,7 @@ def constraint_function(match_ctx):
         res = pat.ResultPattern()
 
         def compute_begin_norm_axis(match_ctx):
-            axis = match_ctx.VectorInt64Attr("axis")
+            axis = match_ctx.VectorInt64Attr("value")
             pow_out_shape = match_ctx.Tensor("pow_out").shape
             return (
                 len(pow_out_shape) - 1 if axis[0] == -1 else axis[0],
diff --git a/test/ir/pir/test_build_op.py b/test/ir/pir/test_build_op.py
@@ -200,14 +200,14 @@ def false_func():
             )
             pred = paddle.less_than(y, x)
             out = paddle.static.nn.cond(pred, true_func, false_func)
-            value1 = main_program.get_value_by_op_id(65)
+            value1 = main_program.get_value_by_op_id(69)
             self.assertEqual(
                 out.get_defining_op().id(),
                 value1[0].get_defining_op().id(),
             )
-            value2 = main_program.get_value_by_op_id([58, 65])
+            value2 = main_program.get_value_by_op_id([58, 69])
             self.assertEqual(
-                58,
+                69,
                 value2[0].get_defining_op().id(),
             )
 
diff --git a/test/ir/pir/test_ir_vjp.py b/test/ir/pir/test_ir_vjp.py
diff --git a/test/ir/pir/test_static_pylayer_api.py b/test/ir/pir/test_static_pylayer_api.py
diff --git a/test/legacy_test/test_prune.py b/test/legacy_test/test_prune.py
diff --git a/test/prim/pir_prim/test_auto_recompute_dy2static.py b/test/prim/pir_prim/test_auto_recompute_dy2static.py
diff --git a/test/prim/pir_prim/test_decomp_op.py b/test/prim/pir_prim/test_decomp_op.py

Original file line number	Diff line number	Diff line change
`@@ -138,8 +138,7 @@ class ReduceCommonOpPattern : public pir::OpRewritePattern<OpType> {`
`138`	`138`	`return false;`
`139`	`139`	`}`
`140`	`140`
`141`		`- if constexpr (std::is_same_v<OpType, paddle::dialect::MeanOp> \|\|`
`142`		`- std::is_same_v<OpType, paddle::dialect::AnyOp> \|\|`
	`141`	`+ if constexpr (std::is_same_v<OpType, paddle::dialect::AnyOp> \|\|`
`143`	`142`	`std::is_same_v<OpType, paddle::dialect::AllOp>) {`
`144`	`143`	`if (!op->HasAttribute("axis")) {`
`145`	`144`	`VLOG(3) << "The axis attribute does not exist";`