PaddlePaddle
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions b/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/op/reduction.cc
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎paddle/cinn/hlir/op/reduction.cc
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/pe/reduction.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/cinn/hlir/pe/reduction.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/pe/reduction.h
Copy file name to clipboard
+21Lines changed: 21 additions & 0 deletions b/‎paddle/cinn/hlir/pe/reduction.h
Copy file name to clipboard
+21Lines changed: 21 additions & 0 deletions
diff --git a/‎paddle/cinn/ir/ir.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/ir/ir.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/cinn/ir/schedule/factorize_reduction.h
Copy file name to clipboard
+15-15Lines changed: 15 additions & 15 deletions b/‎paddle/cinn/ir/schedule/factorize_reduction.h
Copy file name to clipboard
+15-15Lines changed: 15 additions & 15 deletions
diff --git a/‎paddle/cinn/lang/builtin.h
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions b/‎paddle/cinn/lang/builtin.h
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions
diff --git a/‎paddle/cinn/optim/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/optim/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/cinn/optim/optimize.cc
Copy file name to clipboard
+9-5Lines changed: 9 additions & 5 deletions b/‎paddle/cinn/optim/optimize.cc
Copy file name to clipboard
+9-5Lines changed: 9 additions & 5 deletions
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/cinn/ast_gen_ius/ast_gen.h"
+#include "paddle/cinn/hlir/pe/reduction.h"
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/ir_base.h"
 #include "paddle/cinn/ir/ir_printer.h"
@@ -67,6 +68,14 @@ StmtRef ConvertReduceBody(ir::Expr body,
       return Store(tensor, tensor(axis_exprs) && reduce_node->body, axis_exprs);
     case ir::Reduce::kAny:
       return Store(tensor, tensor(axis_exprs) || reduce_node->body, axis_exprs);
+    case ir::Reduce::kVariance:
+      return Store(tensor,
+                   ir::Call::Make(tensor->type(),
+                                  hlir::pe::kVarianceFuncName,
+                                  {tensor(axis_exprs), reduce_node->body},
+                                  {},
+                                  ir::CallType::Intrinsic),
+                   axis_exprs);
     default:
       CINN_NOT_IMPLEMENTED
   }
 
@@ -372,6 +372,7 @@ STRATEGY_FOR_REDUCE(reduce_any,
                     pe::TwoStepBlockReduceAny,
                     pe::BlockShuffleReduceAny,
                     pe::ReduceAny);
+STRATEGY_FOR_REDUCE(variance, Variance, nullptr, nullptr, pe::Variance);
 
 STRATEGY_FOR_REDUCE_SYMBOLIC(reduce_sum,
                              ReduceSum,
@@ -403,6 +404,8 @@ STRATEGY_FOR_REDUCE_SYMBOLIC(reduce_any,
                              pe::TwoStepBlockReduceAny,
                              pe::BlockShuffleReduceAny,
                              pe::ReduceAny);
+STRATEGY_FOR_REDUCE_SYMBOLIC(
+    variance, Variance, nullptr, nullptr, pe::Variance);
 
 #undef STRATEGY_FOR_REDUCE
 #undef STRATEGY_FOR_REDUCE_SYMBOLIC
@@ -431,7 +434,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(variance, Variance);
   CINN_REGISTER_REDUCTION(reduce_max, ReduceMax);
   CINN_REGISTER_REDUCTION(reduce_min, ReduceMin);
 
 
@@ -310,6 +310,14 @@ Tensor ReduceAny(const Tensor& A,
   return Reduce(A, axes, lang::ReduceAny, keep_dims, Expr(false), output_name);
 }
 
+Tensor Variance(const Tensor& A,
+                const std::vector<int>& axes,
+                const bool keep_dims,
+                const std::string& output_name) {
+  return Reduce(
+      A, axes, lang::Variance, keep_dims, lang::Zero(A->type()), output_name);
+}
+
 std::vector<Tensor> WarpReduce(const ir::Tensor& A,
                                const int last_reduce_dim_num,
                                const bool keep_dim,
 
@@ -143,6 +143,25 @@ ir::Tensor ReduceAny(const ir::Tensor& A,
                      const bool keep_dims = false,
                      const std::string& output_name = "T_Reduce_Any_out");
 
+/**
+ * @brief compute the variance of array elements over the given axis
+ *
+ * @param A The input Tensor
+ * @param axis Axis or axes to compute the variance over. If axis is empty, the
+ * operation will compute over all elements of the input array. If axis is
+ * negative it counts from the last to the first axis.
+ * @param keep_dims If it is set true, the axes which are reduced are left in
+ * the result as dimensions with size one. With this option, the result will
+ * broadcast correctly against the input array.
+ * @param output_name The name of the output Tensor
+ *
+ * @return The result Tensor.
+ */
+ir::Tensor Variance(const ir::Tensor& A,
+                    const std::vector<int>& axis,
+                    const bool keep_dims = false,
+                    const std::string& output_name = "T_Variance_out");
+
 /**
  * @brief find the max of array elements over the last dimension
  *
@@ -468,6 +487,8 @@ std::vector<ir::Tensor> TwoStepBlockReduceAny(
     const bool keep_dim,
     const std::string& output_name = "T_Reduce_Any_out");
 
+constexpr char* kVarianceFuncName = "cinn_reduce_variance";
+
 std::string CrossThreadReduceExternalFuncName(const ir::Expr& op,
                                               const ir::Expr& tensor);
 
 
@@ -494,6 +494,7 @@ struct Reduce : public ExprNode<Reduce> {
     kMin,
     kAll,
     kAny,
+    kVariance,
   };
 
   //! The initial value.
 
@@ -451,18 +451,12 @@ class RBBlockCreator : public ReduceBlockCreator {
   void CreateUpdateStmt() override {
     Expr original_store_body = original_update_stmt_.As<ir::Store>()->value;
     Expr new_store_body = ir_utils::IRCopy(original_store_body);
-    std::string original_store_name =
-        original_update_stmt_.As<ir::Store>()->tensor.as_tensor()->name;
 
-#define REPLACE_RF_TENSOR(Op)                                          \
-  if (new_store_body.As<Op>()) {                                       \
-    auto* node = new_store_body.As<Op>();                              \
-    PADDLE_ENFORCE_NOT_NULL(node,                                      \
-                            ::common::errors::InvalidArgument(         \
-                                "The conversion of new_store_body to " \
-                                "Op* failed, node is nullptr."));      \
-    auto& operand = node->b();                                         \
-    operand = Load::Make(rf_tensor_, rf_tensor_access_indices_);       \
+#define REPLACE_RF_TENSOR(Op)                                    \
+  if (new_store_body.As<Op>()) {                                 \
+    auto* node = new_store_body.As<Op>();                        \
+    auto& operand = node->b();                                   \
+    operand = Load::Make(rf_tensor_, rf_tensor_access_indices_); \
   }
 
     REPLACE_RF_TENSOR(Add)
@@ -471,12 +465,18 @@ class RBBlockCreator : public ReduceBlockCreator {
     REPLACE_RF_TENSOR(Min)
     REPLACE_RF_TENSOR(And)
     REPLACE_RF_TENSOR(Or)
-    REPLACE_RF_TENSOR(LT)
-    REPLACE_RF_TENSOR(LE)
-    REPLACE_RF_TENSOR(GT)
-    REPLACE_RF_TENSOR(GE)
 #undef REPLACE_RF_TENSOR
 
+    if (new_store_body.As<ir::Call>()) {
+      auto* node = new_store_body.As<ir::Call>();
+      PADDLE_ENFORCE_EQ(node->read_args.size(),
+                        2UL,
+                        ::common::errors::InvalidArgument(
+                            "The reduction Call op must have exactly two "
+                            "arguments."));
+      node->read_args[1] = Load::Make(rf_tensor_, rf_tensor_access_indices_);
+    }
+
     Expr original_store_tensor = original_update_stmt_.As<ir::Store>()->tensor;
     std::vector<Expr> original_store_indices =
         original_update_stmt_.As<ir::Store>()->indices;
 
@@ -176,6 +176,15 @@ inline Expr ReduceAny(Expr e,
   return ir::Reduce::Make(ir::Reduce::kAny, initial, e, reduce_axis);
 }
 
+inline Expr Variance(Expr e,
+                     const std::vector<Var>& reduce_axis,
+                     Expr initial = Expr()) {
+  if (!initial.defined()) {
+    initial = Zero(e->type());
+  }
+  return ir::Reduce::Make(ir::Reduce::kVariance, initial, e, reduce_axis);
+}
+
 Expr IsNan(Expr e);
 
 Expr Infinity(const Type& type);
 
@@ -24,6 +24,7 @@ gather_srcs(
   lower_intrin.cc
   cast_bool_to_int8.cc
   var_mod_simplify.cc
+  realize_welford_pass.cc
   remove_schedule_block_pass.cc
   replace_cross_block_reduction.cc
   replace_cross_thread_reduction.cc
 
@@ -31,6 +31,7 @@
 #include "paddle/cinn/optim/lower_function_call_bind_vars.h"
 #include "paddle/cinn/optim/lower_intrin.h"
 #include "paddle/cinn/optim/map_extern_call.h"
+#include "paddle/cinn/optim/realize_welford_pass.h"
 #include "paddle/cinn/optim/rearrange_load_instruction_pass.h"
 #include "paddle/cinn/optim/reindex_transpose_buffer_pass.h"
 #include "paddle/cinn/optim/remove_schedule_block_pass.h"
@@ -68,18 +69,21 @@ ir::LoweredFunc Optimize(ir::LoweredFunc fn,
   Simplify(&copied->body);
   EliminateInvariantLoop(&copied->body);
   VLOG(4) << "After Optimize EliminateInvariantLoop:" << copied;
-  ReplaceCrossThreadReduction(copied);
-  VLOG(4) << "After Optimize ReplaceCrossThreadReduction:" << copied;
-  ReplaceCrossBlockReduction(copied);
-  VLOG(4) << "After Optimize ReplaceCrossBlockReduction:" << copied;
 
   {
     FuncPassManager func_pass_manager;
+    func_pass_manager.AddPass(CreateRealizeWelfordPass());
     func_pass_manager.AddPass(CreateReindexTransposeBufferPass());
     func_pass_manager.Run(copied);
-    VLOG(4) << "After Optimize ReindexTransposeBuffer:" << copied;
+    VLOG(4) << "After Optimize RealizeWelford and ReindexTransposeBuffer: "
+            << copied;
   }
 
+  ReplaceCrossThreadReduction(copied);
+  VLOG(4) << "After Optimize ReplaceCrossThreadReduction:" << copied;
+  ReplaceCrossBlockReduction(copied);
+  VLOG(4) << "After Optimize ReplaceCrossBlockReduction:" << copied;
+
   target.arch.Match(
       [&](common::NVGPUArch) {
 #ifdef CINN_WITH_CUDA