From 3b14e0769030441a0c97c47068eec5755b06447e Mon Sep 17 00:00:00 2001
From: xs_courtesy <xs1580802568@gmail.com>
Date: Wed, 3 Jan 2024 09:20:47 +0000
Subject: [PATCH] simplify extent of loop after fuse and add corresponding test
 case

---
 .../ir/schedule/impl/loop_transformation.cc   |  2 +-
 test/cinn/ir/test_llir_schedule_fuse_split.py | 55 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/paddle/cinn/ir/schedule/impl/loop_transformation.cc b/paddle/cinn/ir/schedule/impl/loop_transformation.cc
index 1d66697f43d13..daa453d85a744 100644
--- a/paddle/cinn/ir/schedule/impl/loop_transformation.cc
+++ b/paddle/cinn/ir/schedule/impl/loop_transformation.cc
@@ -264,7 +264,7 @@ Expr DyScheduleImpl::Fuse(const std::vector<Expr>& loops) {
   for (int i = 0; i < loops_number; ++i) {
     fused_extent = fused_extent * for_nodes[i]->extent;
   }
-
+  fused_extent = cinn::common::AutoSimplify(fused_extent);
   if (!fused_body.As<ir::Block>()) fused_body = Block::Make({fused_body});
   Expr new_stmt = For::Make(fused_var,
                             Expr(0),
diff --git a/test/cinn/ir/test_llir_schedule_fuse_split.py b/test/cinn/ir/test_llir_schedule_fuse_split.py
index 362cb81f87b96..b4722a1a02434 100644
--- a/test/cinn/ir/test_llir_schedule_fuse_split.py
+++ b/test/cinn/ir/test_llir_schedule_fuse_split.py
@@ -37,7 +37,7 @@ def elementwise_fuse_assign_loop(
     def elementwise_fuse_assign_loop_gt(
         X: DataArray((128, 128, 128)), Y: DataArray((128, 128, 128))
     ):
-        for i in range(((1 * 128) * 128) * 128):
+        for i in range(2097152):
             with ir.ScheduleBlockContext("Y") as block_y:
                 i1_1, j1_1, k1_1 = ir.AxisMap(
                     "SSS", [(i / 128) / 128, (i / 128) % 128, i % 128]
@@ -148,7 +148,7 @@ def elementwise_fuse_assign_loop(
             Y: DataArray((-1, 128, 128)),
             N: ir.Var(),
         ):
-            for i_j_k_fused in range(((1 * N) * 128) * 128):
+            for i_j_k_fused in range(16384 * N):
                 with ir.ScheduleBlockContext("Y") as block_y:
                     i1, j1, k1 = ir.AxisMap(
                         "SSS",
@@ -207,9 +207,60 @@ def elementwise_split(
     assert_llir_equal(origin.elementwise_split, expected.elementwise_split)
 
 
+def test_fuse_split():
+    @to_cinn_llir
+    def elementwise_fuse_split_origin(
+        X: DataArray((64, 128, 128)), Y: DataArray((64, 128, 128))
+    ):
+        for i in range(64):
+            for j in range(128):
+                for k in range(128):
+                    with ir.ScheduleBlockContext("Y") as Y_block:
+                        i1, j1, k1 = ir.AxisMap("SSS", [i, j, k])
+                        fused = sch.fuse([i, j])
+                        sch.split(fused, factors=[2, 512, -1])
+                        Y[i1, j1, k1] = X[i1, j1, k1] * 2.0
+
+    @to_cinn_llir
+    def elementwise_fuse_split_expected(
+        X: DataArray((64, 128, 128)), Y: DataArray((64, 128, 128))
+    ):
+        for i_j_fused in range(2):
+            for i_j_fused_0 in range(512):
+                for i_j_fused_1 in range(8):
+                    for k in range(128):
+                        with ir.ScheduleBlockContext("Y") as Y_block:
+                            i1, j1, k1 = ir.AxisMap(
+                                "SSS",
+                                [
+                                    (
+                                        (
+                                            (4096 * i_j_fused)
+                                            + ((8 * i_j_fused_0) + i_j_fused_1)
+                                        )
+                                        / 128
+                                    ),
+                                    (
+                                        (
+                                            (4096 * i_j_fused)
+                                            + ((8 * i_j_fused_0) + i_j_fused_1)
+                                        )
+                                        % 128
+                                    ),
+                                    k,
+                                ],
+                            )
+                            Y[i1, j1, k1] = X[i1, j1, k1] * 2.0
+
+    assert_llir_equal(
+        elementwise_fuse_split_origin, elementwise_fuse_split_expected
+    )
+
+
 if __name__ == "__main__":
     test_fuse()
     test_split()
+    test_fuse_split()
     test_split_predicate()
     test_fuse_dynamic()
     test_split_dynamic()