fix reduction on bool (#63614)

BiynXu · web-flow · commit a3b51c7d8944 · 2024-04-22T10:48:01.000+08:00
diff --git a/paddle/cinn/ast_gen_ius/ast_gen.cc b/paddle/cinn/ast_gen_ius/ast_gen.cc
@@ -28,6 +28,10 @@ PD_DECLARE_bool(cinn_bucket_compile);
 namespace cinn {
 namespace ast_gen_ius {
 
+bool IsReduceBool(const ir::Expr& lhs, const ir::Expr& rhs) {
+  return lhs.type().is_bool() || rhs.type().is_bool();
+}
+
 ir::Expr ConvertReduceBody(ir::Expr body,
                            ir::Tensor tensor,
                            const std::vector<Expr>& axis_exprs) {
@@ -38,9 +42,17 @@ ir::Expr ConvertReduceBody(ir::Expr body,
 
   switch (reduce_node->reduce_type) {
     case ir::Reduce::kSum:
+      if (IsReduceBool(tensor(axis_exprs), reduce_node->body)) {
+        return ir::Store::Make(
+            tensor, tensor(axis_exprs) || reduce_node->body, axis_exprs);
+      }
       return ir::Store::Make(
           tensor, tensor(axis_exprs) + reduce_node->body, axis_exprs);
     case ir::Reduce::kMul:
+      if (IsReduceBool(tensor(axis_exprs), reduce_node->body)) {
+        return ir::Store::Make(
+            tensor, tensor(axis_exprs) && reduce_node->body, axis_exprs);
+      }
       return ir::Store::Make(
           tensor, tensor(axis_exprs) * reduce_node->body, axis_exprs);
     case ir::Reduce::kMax:
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.cc b/paddle/cinn/hlir/framework/pir/trivial_op_util.cc
@@ -417,14 +417,24 @@ ExprTransformer ChangeVarTransformer(const std::vector<ir::Var>& target_vars,
   return ExprTransformer(f);
 }
 
+bool IsReduceBool(const ir::Expr& lhs, const ir::Expr& rhs) {
+  return lhs.type().is_bool() || rhs.type().is_bool();
+}
+
 ExprTransformer WrapReduceOperation(const ir::Reduce::ReduceType& reduce_type,
                                     const ir::Tensor& tensor,
                                     const std::vector<ir::Expr>& axis_exprs) {
   const auto& f = [=](const ir::Expr& e) -> ir::Expr {
     switch (reduce_type) {
       case ir::Reduce::kSum:
+        if (IsReduceBool(tensor(axis_exprs), e)) {
+          return ir::Store::Make(tensor, tensor(axis_exprs) || e, axis_exprs);
+        }
         return ir::Store::Make(tensor, tensor(axis_exprs) + e, axis_exprs);
       case ir::Reduce::kMul:
+        if (IsReduceBool(tensor(axis_exprs), e)) {
+          return ir::Store::Make(tensor, tensor(axis_exprs) && e, axis_exprs);
+        }
         return ir::Store::Make(tensor, tensor(axis_exprs) * e, axis_exprs);
       case ir::Reduce::kMax:
         return ir::Store::Make(
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.cc b/paddle/cinn/ir/group_schedule/config/group_tile_config.cc
@@ -70,16 +70,29 @@ BuildPureStaticShapeConfig(
     const std::shared_ptr<ScheduleConfig::BaseInfo>& base_info,
     const common::Target& target) {
   if (base_info->spatial_numel == 1) {  // reduce all
-    BucketInfo bucket_info{/* sp_lower_bound = */ 1,
-                           /* sp_upper_bound = */ 1,
-                           /* rb_lower_bound = */ 1,
-                           /* rb_upper_bound = */ kMaxNumel};
-    ScheduleConfig::TileConfig tile_config{
-        /* warp_num = */ 8,
-        /* tree_reduce_num = */ 256,
-        /* spatial_inner_num = */ 1,
-        /* reduce_method = */ BlockReduceMethod()};
-    return {{bucket_info, tile_config}};
+    if (base_info->reduce_numel <= 256) {
+      BucketInfo bucket_info{/* sp_lower_bound = */ 1,
+                             /* sp_upper_bound = */ 1,
+                             /* rb_lower_bound = */ 1,
+                             /* rb_upper_bound = */ 256};
+      ScheduleConfig::TileConfig tile_config{
+          /* warp_num = */ (base_info->reduce_numel + 31) / 32,
+          /* tree_reduce_num = */ base_info->reduce_numel,
+          /* spatial_inner_num = */ 1,
+          /* reduce_method = */ BlockReduceMethod()};
+      return {{bucket_info, tile_config}};
+    } else {
+      BucketInfo bucket_info{/* sp_lower_bound = */ 1,
+                             /* sp_upper_bound = */ 1,
+                             /* rb_lower_bound = */ 257,
+                             /* rb_upper_bound = */ kMaxNumel};
+      ScheduleConfig::TileConfig tile_config{
+          /* warp_num = */ 8,
+          /* tree_reduce_num = */ 256,
+          /* spatial_inner_num = */ 1,
+          /* reduce_method = */ BlockReduceMethod()};
+      return {{bucket_info, tile_config}};
+    }
   } else if (base_info->reduce_numel == 1) {  // no reduce
     int64_t spatial_block = Next2Power(base_info->spatial_numel);
     if (spatial_block > 1024) {
