PaddlePaddle
diff --git a/‎paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
Copy file name to clipboard
+11-1Lines changed: 11 additions & 1 deletion b/‎paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
Copy file name to clipboard
+11-1Lines changed: 11 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
Copy file name to clipboard
+9-1Lines changed: 9 additions & 1 deletion b/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
Copy file name to clipboard
+9-1Lines changed: 9 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.h
Copy file name to clipboard
+17-2Lines changed: 17 additions & 2 deletions b/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.h
Copy file name to clipboard
+17-2Lines changed: 17 additions & 2 deletions
@@ -117,8 +117,18 @@ BucketLoweredFuncsWrapper OpLowererImpl::BucketLower(
 
   // VLOG(4) << "Bucket Lower output values is : " << group->output_values();
   func_bodies = OperationFusion(ops, func_bodies, group->fusion_tracker_ptr);
+
+  std::unordered_set<std::string> fusion_group_args;
+  for (auto value : group->GetInputOpValues()) {
+    fusion_group_args.insert(ValueName(value));
+  }
+
+  for (auto value : group->GetGroupOutputValues()) {
+    fusion_group_args.insert(ValueName(value));
+  }
+
   std::shared_ptr<FusionGroupInfo> fusion_group_info =
-      GetFusionGroupInfo(func_bodies);
+      GetFusionGroupInfo(func_bodies, fusion_group_args);
   // TODO(liangshuhao): grid reduce is disabled for broadcast-leaf group now,
   // because grid reduce introduces extra func args that currently cannot be
   // unified with other broadcast-leaf groups.
 
@@ -26,6 +26,7 @@
 #include "paddle/cinn/ir/dim.h"
 #include "paddle/cinn/ir/group_schedule/base_group_scheduler.h"
 #include "paddle/cinn/ir/group_schedule/config/group_tile_util.h"
+#include "paddle/cinn/ir/ir_analyzer/ir_analyzer.h"
 #include "paddle/cinn/ir/schedule/ir_schedule.h"
 #include "paddle/cinn/ir/schedule/ir_schedule_util.h"
 #include "paddle/cinn/lang/placeholder.h"
@@ -37,6 +38,7 @@
 #include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
 
 PD_DECLARE_bool(cinn_enable_grid_reduce);
+PD_DECLARE_bool(cinn_enable_vectorize);
 
 namespace cinn {
 namespace hlir {
@@ -539,7 +541,8 @@ std::vector<ir::Var> GetAllForIters(const ir::Expr& expr) {
 }  // namespace trivial_fusion_detail
 
 std::shared_ptr<FusionGroupInfo> GetFusionGroupInfo(
-    const std::vector<ir::Expr>& op_compute_bodies) {
+    const std::vector<ir::Expr>& op_compute_bodies,
+    const std::unordered_set<std::string>& group_args) {
   using trivial_fusion_detail::AppendBound;
   using trivial_fusion_detail::GetAllForIters;
   using trivial_fusion_detail::IsReduceBody;
@@ -599,6 +602,11 @@ std::shared_ptr<FusionGroupInfo> GetFusionGroupInfo(
         GetCanApplyGridReduce(op_compute_bodies, group_info->reduce_axis);
   }
 
+  if (FLAGS_cinn_enable_vectorize) {
+    group_info->vectorize_info =
+        GetCanApplyVectorize(op_compute_bodies, group_args);
+  }
+
   VLOG(4) << group_info->DebugPrint();
   return group_info;
 }
 
@@ -33,6 +33,7 @@
 #include "paddle/common/ddim.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
+#
 
 namespace cinn {
 namespace hlir {
@@ -158,26 +159,40 @@ std::vector<ir::Var> GetAllForIters(const ir::Expr& expr);
 
 }  // namespace trivial_fusion_detail
 
+struct GroupVectorizeInfo {
+  bool can_apply_vectorize{false};
+  bool has_if_else_op{false};
+  bool has_select_op{false};
+  int continuous_arg_nums{0};
+  int fusion_group_arg_nums{0};
+};
+
 struct FusionGroupInfo {
   std::vector<int64_t> loop_ranges;
   std::vector<int64_t> loop_strides;
   std::vector<int64_t> reduce_axis;
   std::vector<std::string> reduce_var_name;
   bool can_apply_grid_reduce;
+  GroupVectorizeInfo vectorize_info;
 
   std::string DebugPrint() {
     std::stringstream ss;
     ss << "GroupInfo\nloop_ranges: " << cinn::utils::Join(loop_ranges, " ")
        << "\nloop_strides: " << cinn::utils::Join(loop_strides, ", ")
        << "\nreduce_axis: " << cinn::utils::Join(reduce_axis, " ")
        << "\nreduce_var_name: " << cinn::utils::Join(reduce_var_name, " ")
-       << "\ncan_apply_grid_reduce: " << can_apply_grid_reduce;
+       << "\ncan_apply_grid_reduce: " << can_apply_grid_reduce
+       << "\ncan_apply_vectorize: " << vectorize_info.can_apply_vectorize
+       << "\nhas_select_op: " << vectorize_info.has_select_op
+       << "\ncontinuous_arg_nums: " << vectorize_info.continuous_arg_nums
+       << "\nfusion_group_arg_nums: " << vectorize_info.fusion_group_arg_nums;
     return ss.str();
   }
 };
 
 std::shared_ptr<FusionGroupInfo> GetFusionGroupInfo(
-    const std::vector<ir::Expr>& op_compute_bodies);
+    const std::vector<ir::Expr>& op_compute_bodies,
+    const std::unordered_set<std::string>& group_args);
 
 }  // namespace pir
 }  // namespace framework