[CINN]add variance op in frontend (#71184)

Hongqing-work · web-flow · commit 5ffeee82987e · 2025-02-20T11:01:21.000+08:00
* [CINN]add variance op in frontend

* fix-comment
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/single_op_fallback_to_phi.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/single_op_fallback_to_phi.cc
@@ -58,6 +58,10 @@ class FusionOpPattern : public pir::OpRewritePattern<cinn::dialect::FusionOp> {
             "The last operator of fusion_op must be YieldOp, but got %s",
             fusion_op.GetOperators()[2]->name()));
 
+    if (fusion_op.GetOperators()[0]->isa<paddle::dialect::VarianceOp>()) {
+      return false;
+    }
+
     std::optional<pir::Operation*> paddle_op =
         FallBackOp(fusion_op.GetOperators()[0], rewriter);
     if (!paddle_op.has_value()) {
diff --git a/paddle/cinn/hlir/op/reduction.cc b/paddle/cinn/hlir/op/reduction.cc
@@ -431,6 +431,7 @@ CINN_REGISTER_HELPER(reduce_ops) {
 
   CINN_REGISTER_REDUCTION(reduce_sum, ReduceSum);
   CINN_REGISTER_REDUCTION(reduce_prod, ReduceProd);
+  CINN_REGISTER_REDUCTION(variance, ReduceProd);
   CINN_REGISTER_REDUCTION(reduce_max, ReduceMax);
   CINN_REGISTER_REDUCTION(reduce_min, ReduceMin);
 
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -4464,6 +4464,16 @@ bool UnstackOpInferSymbolicShape(
   return true;
 }
 
+bool VarianceOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &axis = details::GetVectorAttr(op, "axis");
+  return details::ReduceInferDim(op,
+                                 infer_context,
+                                 axis,
+                                 GetBoolAttr(op, "keepdim"), /*keepdim*/
+                                 axis.size() == 0 /*reduce_all*/);
+}
+
 bool WeightQuantizeOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -166,6 +166,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unsqueeze)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unsqueeze_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unfold)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unstack)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Variance)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightQuantize)
 
 }  // namespace paddle::dialect
diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -56,3 +56,4 @@
 - isnan
 - isfinite
 - take_along_axis
+- variance
diff --git a/paddle/fluid/primitive/base/primitive_ops.h b/paddle/fluid/primitive/base/primitive_ops.h
@@ -107,6 +107,7 @@ const std::set<std::string>& GetPrimitiveOpNames() {
       "pd_op.sigmoid",
       "pd_op.reduce_as",
       /* skip some special ops */
+      "pd_op.variance",
       "pd_op.conv2d",
       "pd_op.pad3d",
       "pd_op.nearest_interp",
diff --git a/paddle/fluid/primitive/primitive/primitive.yaml b/paddle/fluid/primitive/primitive/primitive.yaml
@@ -126,3 +126,4 @@
 - sigmoid
 - reduce_as
 - pow
+- variance
diff --git a/paddle/phi/kernels/reduce_variance_kernel.cc b/paddle/phi/kernels/reduce_variance_kernel.cc
@@ -0,0 +1,45 @@
+// Copyright (c) 2025 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
+//
+//     https://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/phi/kernels/reduce_variance_kernel.h"
+#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
+#include "paddle/phi/kernels/elementwise_subtract_kernel.h"
+#include "paddle/phi/kernels/reduce_mean_kernel.h"
+
+#include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+template <typename T, typename Context>
+void VarianceKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const IntArray& dims,
+                    bool keep_dim,
+                    DenseTensor* out) {
+  DenseTensor temp_mean = Mean<T, Context>(dev_ctx, x, dims, true);
+  DenseTensor temp_differences = Subtract<T, Context>(dev_ctx, x, temp_mean);
+  DenseTensor temp_pow =
+      Multiply<T, Context>(dev_ctx, temp_differences, temp_differences);
+
+  MeanKernel<T, Context>(dev_ctx, temp_pow, dims, keep_dim, out);
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    variance, CPU, ALL_LAYOUT, phi::VarianceKernel, float, double) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(
+    variance, GPU, ALL_LAYOUT, phi::VarianceKernel, float, double) {}
+#endif
diff --git a/paddle/phi/kernels/reduce_variance_kernel.h b/paddle/phi/kernels/reduce_variance_kernel.h
@@ -0,0 +1,40 @@
+// Copyright (c) 2025 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
+//
+//     https://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.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/infermeta/unary.h"
+
+namespace phi {
+template <typename T, typename Context>
+void VarianceKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const IntArray& dims,
+                    bool keep_dim,
+                    DenseTensor* out);
+
+template <typename T, typename Context>
+DenseTensor Variance(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const IntArray& axis,
+                     bool keep_dim) {
+  DenseTensor dense_out;
+  MetaTensor meta_out(&dense_out);
+  SumRawInferMeta(x, axis, keep_dim, false, x.dtype(), &meta_out);
+  VarianceKernel<T, Context>(dev_ctx, x, axis, keep_dim, &dense_out);
+  return dense_out;
+}
+
+}  // namespace phi
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -5295,6 +5295,15 @@
   interfaces : paddle::dialect::InferSymbolicShapeInterface
   traits : paddle::dialect::ForwardOnlyTrait
 
+- op : variance
+  args : (Tensor x, int64_t[] axis={}, bool keepdim=false)
+  output : Tensor(out)
+  infer_meta :
+    func : ReduceInferMeta
+  kernel :
+    func : variance
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
+
 - op : view_dtype
   args : (Tensor input, DataType dtype)
   output : Tensor(out)
