[CINN] Upgrade ReductionReduction rule (#57963)

BiynXu · web-flow · commit e565b0fc2f26 · 2023-10-10T16:06:49.000+08:00
diff --git a/paddle/cinn/auto_schedule/search_space/auto_gen_rule/reduction_factoring.cc b/paddle/cinn/auto_schedule/search_space/auto_gen_rule/reduction_factoring.cc
@@ -144,45 +144,39 @@ void ReductionFactoring::Apply(const std::string& block_name,
     }
   }
 
-  // TODO(BiynXu): After implementing the factorize_reduction schedule
-  // primitive, restore the following annotations. The factorize_reduction
-  // schedule primitive needs to support complex subscripts to support pre
-  // schedule transformations.
-
-  // // 4. Fuse all reduction loops
-  // ir::Expr fused_reduce_loop;
-  // if (num_reduction_loops > 1) {
-  //   std::vector<int> reduction_loop_indices;
-  //   for (int i = num_spatial_loops - 1; i < all_loops.size(); ++i) {
-  //     reduction_loop_indices.push_back(i);
-  //   }
-  //   CHECK_EQ(reduction_loop_indices.size(), num_reduction_loops);
-  //   fused_reduce_loop = ir_schedule->Fuse(block_name,
-  //   reduction_loop_indices);
-  // } else {
-  //   all_loops = ir_schedule->GetLoops(block_name);
-  //   fused_reduce_loop = all_loops.back();
-  // }
-  // // 5. Split the reduction loop into 2 part
-  // int factor = 1;
-  // int extent = ir::GetLoopExtent(fused_reduce_loop);
-  // for (int i = ceil(sqrt(extent)); i >= 1; --i) {
-  //   if (extent % i == 0) {
-  //     factor = i;
-  //     break;
-  //   }
-  // }
-  // std::vector<cinn::ir::Expr> splited_reduction_loops =
-  // ir_schedule->Split(fused_reduce_loop, {-1, factor});
-  // // Apply FactorizeReduction
-  // LOG(INFO) << "before FactorizeReduction: " <<
-  // ir_schedule->GetModule().GetExprs()[0];
-  // ir_schedule->FactorizeReduction(splited_reduction_loops[0],
-  // num_spatial_loops);
-
-  // Apply rfactor
-  all_loops = ir_schedule->GetLoops(block_name);
-  ir_schedule->Rfactor(all_loops[num_spatial_loops], num_spatial_loops);
+  // 4. Fuse all reduction loops
+  ir::Expr fused_reduce_loop;
+  VLOG(6) << "before Fuse: " << ir_schedule->GetModule().GetExprs()[0];
+  if (num_reduction_loops > 1) {
+    std::vector<int> reduction_loop_indices;
+    for (int i = num_spatial_loops; i < all_loops.size(); ++i) {
+      reduction_loop_indices.push_back(i);
+    }
+    CHECK_EQ(reduction_loop_indices.size(), num_reduction_loops);
+    fused_reduce_loop = ir_schedule->Fuse(block_name, reduction_loop_indices);
+  } else {
+    all_loops = ir_schedule->GetLoops(block_name);
+    fused_reduce_loop = all_loops.back();
+  }
+  // 5. Split the reduction loop into 2 part
+  VLOG(6) << "before Split: " << ir_schedule->GetModule().GetExprs()[0];
+  int factor = 1;
+  int extent = ir::GetLoopExtent(fused_reduce_loop);
+  for (int i = ceil(sqrt(extent)); i >= 1; --i) {
+    if (extent % i == 0) {
+      factor = i;
+      break;
+    }
+  }
+  std::vector<cinn::ir::Expr> splited_reduction_loops =
+      ir_schedule->Split(fused_reduce_loop, {-1, factor});
+  // 6.  Apply FactorizeReduction
+  VLOG(6) << "before FactorizeReduction: "
+          << ir_schedule->GetModule().GetExprs()[0];
+  ir_schedule->FactorizeReduction(splited_reduction_loops[0],
+                                  num_spatial_loops);
+  VLOG(6) << "after FactorizeReduction: "
+          << ir_schedule->GetModule().GetExprs()[0];
 }
 
 }  // namespace auto_schedule
diff --git a/paddle/cinn/auto_schedule/search_space/auto_gen_rule/reduction_factoring_test.cc b/paddle/cinn/auto_schedule/search_space/auto_gen_rule/reduction_factoring_test.cc
@@ -77,28 +77,26 @@ TEST_F(TestReductionFactoring, AnalyseApplyType) {
             RuleApplyType::kCannotApply);
 }
 
-TEST_F(TestReductionFactoring, ApplyOnBlock) {
+TEST_F(TestReductionFactoring, ApplyOnBlock1ReduceDim) {
   std::string expected_ir = R"({
   ScheduleBlock(root)
   {
     {
       serial for (i, 0, 32)
       {
-        serial for (rf_reduce_k_0, 0, 64)
+        serial for (reduce_k_0_0, 0, 8)
         {
-          ScheduleBlock(rf_var_0__reduce_init)
+          ScheduleBlock(var_0_rf__reduce_init)
           {
-            i0, i1_0 = axis.bind(i, rf_reduce_k_0)
-            rf_var_0__reduce_init[i0, i1_0] = 0.00000000f
+            vreduce_k_0_0, i0_0 = axis.bind(reduce_k_0_0, i)
+            var_0_rf__reduce_init[i0_0, vreduce_k_0_0] = 0.00000000f
           }
-          serial for (reduce_k_1, 0, 128)
+          serial for (reduce_k_0_1, 0, 8)
           {
-            ScheduleBlock(rf_var_0)
+            ScheduleBlock(var_0_rf)
             {
-              i0_0, i1, i2 = axis.bind(i, rf_reduce_k_0, reduce_k_1)
-              read_buffers(_var_0[i0_0(0:32)], _X[i0_0(0:32), i1(0:64), i2(0:128)])
-              write_buffers(_var_0[i0_0(0:32)])
-              rf_var_0[i0_0, i1] = (rf_var_0[i0_0, i1] + X[i0_0, i1, i2])
+              vreduce_k_0_0, i0_0, vreduce_k_0_1 = axis.bind(reduce_k_0_0, i, reduce_k_0_1)
+              var_0_rf[i0_0, vreduce_k_0_0] = (var_0_rf[i0_0, vreduce_k_0_0] + X[i0_0, ((8 * vreduce_k_0_0) + vreduce_k_0_1)])
             }
           }
         }
@@ -107,17 +105,61 @@ TEST_F(TestReductionFactoring, ApplyOnBlock) {
       {
         ScheduleBlock(var_0__reduce_init)
         {
-          i0 = axis.bind(i)
-          var_0__reduce_init[i0] = 0.00000000f
+          i0_0 = axis.bind(i)
+          var_0__reduce_init[i0_0] = 0.00000000f
         }
-        serial for (reduce_k_0, 0, 64)
+        serial for (reduce_k_0_0, 0, 8)
         {
           ScheduleBlock(var_0)
           {
-            i0_0, i1 = axis.bind(i, reduce_k_0)
-            read_buffers(_var_0[i0_0(0:32)], _X[i0_0(0:32), i1(0:64), i2(0:128)])
-            write_buffers(_var_0[i0_0(0:32)])
-            var_0[i0_0] = (var_0[i0_0] + rf_var_0[i0_0, i1])
+            vreduce_k_0_0, i0_0 = axis.bind(reduce_k_0_0, i)
+            var_0[i0_0] = (var_0[i0_0] + var_0_rf[i0_0, vreduce_k_0_0])
+          }
+        }
+      }
+    }
+  }
+})";
+  TestApplyOnReduce({32, 64}, {1}, "var_0", expected_ir);
+}
+
+TEST_F(TestReductionFactoring, ApplyOnBlock2ReduceDim) {
+  std::string expected_ir = R"({
+  ScheduleBlock(root)
+  {
+    {
+      serial for (i, 0, 32)
+      {
+        serial for (reduce_k_0_reduce_k_1_fused, 0, 128)
+        {
+          ScheduleBlock(var_0_rf__reduce_init)
+          {
+            vreduce_k_0_reduce_k_1_fused, i0_0 = axis.bind(reduce_k_0_reduce_k_1_fused, i)
+            var_0_rf__reduce_init[i0_0, vreduce_k_0_reduce_k_1_fused] = 0.00000000f
+          }
+          serial for (reduce_k_0_reduce_k_1_fused_0, 0, 64)
+          {
+            ScheduleBlock(var_0_rf)
+            {
+              vreduce_k_0_reduce_k_1_fused, i0_0, vreduce_k_0_reduce_k_1_fused_0 = axis.bind(reduce_k_0_reduce_k_1_fused, i, reduce_k_0_reduce_k_1_fused_0)
+              var_0_rf[i0_0, vreduce_k_0_reduce_k_1_fused] = (var_0_rf[i0_0, vreduce_k_0_reduce_k_1_fused] + X[i0_0, (((64 * vreduce_k_0_reduce_k_1_fused) + vreduce_k_0_reduce_k_1_fused_0) / 128), (((64 * vreduce_k_0_reduce_k_1_fused) + vreduce_k_0_reduce_k_1_fused_0) % 128)])
+            }
+          }
+        }
+      }
+      serial for (i, 0, 32)
+      {
+        ScheduleBlock(var_0__reduce_init)
+        {
+          i0_0 = axis.bind(i)
+          var_0__reduce_init[i0_0] = 0.00000000f
+        }
+        serial for (reduce_k_0_reduce_k_1_fused, 0, 128)
+        {
+          ScheduleBlock(var_0)
+          {
+            vreduce_k_0_reduce_k_1_fused, i0_0 = axis.bind(reduce_k_0_reduce_k_1_fused, i)
+            var_0[i0_0] = (var_0[i0_0] + var_0_rf[i0_0, vreduce_k_0_reduce_k_1_fused])
           }
         }
       }
@@ -127,5 +169,51 @@ TEST_F(TestReductionFactoring, ApplyOnBlock) {
   TestApplyOnReduce({32, 64, 128}, {1, 2}, "var_0", expected_ir);
 }
 
+TEST_F(TestReductionFactoring, ApplyOnBlock3ReduceDim) {
+  std::string expected_ir = R"({
+  ScheduleBlock(root)
+  {
+    {
+      serial for (i, 0, 32)
+      {
+        serial for (reduce_k_0_reduce_k_1_reduce_k_2_fused, 0, 512)
+        {
+          ScheduleBlock(var_0_rf__reduce_init)
+          {
+            vreduce_k_0_reduce_k_1_reduce_k_2_fused, i0_0 = axis.bind(reduce_k_0_reduce_k_1_reduce_k_2_fused, i)
+            var_0_rf__reduce_init[i0_0, vreduce_k_0_reduce_k_1_reduce_k_2_fused] = 0.00000000f
+          }
+          serial for (reduce_k_0_reduce_k_1_reduce_k_2_fused_0, 0, 512)
+          {
+            ScheduleBlock(var_0_rf)
+            {
+              vreduce_k_0_reduce_k_1_reduce_k_2_fused, i0_0, vreduce_k_0_reduce_k_1_reduce_k_2_fused_0 = axis.bind(reduce_k_0_reduce_k_1_reduce_k_2_fused, i, reduce_k_0_reduce_k_1_reduce_k_2_fused_0)
+              var_0_rf[i0_0, vreduce_k_0_reduce_k_1_reduce_k_2_fused] = (var_0_rf[i0_0, vreduce_k_0_reduce_k_1_reduce_k_2_fused] + X[i0_0, ((((512 * vreduce_k_0_reduce_k_1_reduce_k_2_fused) + vreduce_k_0_reduce_k_1_reduce_k_2_fused_0) / 64) / 64), ((((512 * vreduce_k_0_reduce_k_1_reduce_k_2_fused) + vreduce_k_0_reduce_k_1_reduce_k_2_fused_0) / 64) % 64), (((512 * vreduce_k_0_reduce_k_1_reduce_k_2_fused) + vreduce_k_0_reduce_k_1_reduce_k_2_fused_0) % 64)])
+            }
+          }
+        }
+      }
+      serial for (i, 0, 32)
+      {
+        ScheduleBlock(var_0__reduce_init)
+        {
+          i0_0 = axis.bind(i)
+          var_0__reduce_init[i0_0] = 0.00000000f
+        }
+        serial for (reduce_k_0_reduce_k_1_reduce_k_2_fused, 0, 512)
+        {
+          ScheduleBlock(var_0)
+          {
+            vreduce_k_0_reduce_k_1_reduce_k_2_fused, i0_0 = axis.bind(reduce_k_0_reduce_k_1_reduce_k_2_fused, i)
+            var_0[i0_0] = (var_0[i0_0] + var_0_rf[i0_0, vreduce_k_0_reduce_k_1_reduce_k_2_fused])
+          }
+        }
+      }
+    }
+  }
+})";
+  TestApplyOnReduce({32, 64, 64, 64}, {1, 2, 3}, "var_0", expected_ir);
+}
+
 }  // namespace auto_schedule
 }  // namespace cinn
