[Fix][Arith] Analyzer simplification starts with canonical

This PR updates the order of arithmetic analyzer simplification, by
adding a stage of canonical simplification at the very beginning so
that every simplification always starts with a canonical round. This
is because the rewrite simplification may destroy some PrimExpr property
that the canonical simplification can make use of. Therefore, adding
the canonical one in the front can maximize the use of canonical
simplification.

Author: MasterJH5574

src/arith/analyzer.cc

@@ -129,6 +129,10 @@ bool Analyzer::CanProve(const PrimExpr& expr) {
 PrimExpr Analyzer::Simplify(const PrimExpr& expr, int steps) {
   PrimExpr res = expr;
 
+  // Always starts with a canonical simplification, as some structural property
+  // of an expression might be destroyed by rewrite simplification.
+  res = this->canonical_simplify(res);
+
   for (int i = 0; i < steps; ++i) {
     if (tir::is_const_int(res)) {
       return res;

tests/python/unittest/test_arith_intset.py

@@ -182,7 +182,6 @@ def check_region_bound(expect_region, var_dom, mode, predicate=None):
                 expect_begin, expect_end = expect_desc[binding]
                 result_begin = analyzer.simplify(intset.min_value, 3)
                 result_end = analyzer.simplify(intset.max_value + 1, 3)
-                print(result_end)
                 assert analyzer.can_prove_equal(
                     result_begin - expect_begin, 0
                 ), f"{result_begin} vs {expect_begin}"
@@ -306,10 +305,7 @@ def test_region_lower_bound_for_non_perfect_tile():
             + h2: {
                 (): (
                     tvm.tir.max(h3 * 8, 1),
-                    tvm.tir.max(h3 * 8, 1)
-                    - tvm.tir.max(h3 * 8, 214)
-                    - tvm.tir.max(1 - h3 * 8, 0)
-                    + 224,
+                    tvm.tir.min(0, h3 * 8 - 214) + 224,
                 ),
                 ((h3, 0),): (1, 10),  # h3 == 0: region is [1, 10)
                 ((h3, 10),): (h3 * 8, h3 * 8 + 10),  # 0 < h3 <= 26: region is [h3 * 8, h3 * 8 + 10)
@@ -333,10 +329,7 @@ def test_region_lower_bound_for_non_perfect_tile():
             + h1: {
                 (): (
                     tvm.tir.max(h3 * 8, 1),
-                    tvm.tir.max(h3 * 8, 1)
-                    - tvm.tir.max(h3 * 8, 214)
-                    - tvm.tir.max(1 - h3 * 8, 0)
-                    + 224,
+                    tvm.tir.min(0, h3 * 8 - 214) + 224,
                 ),
                 ((h3, 0),): (1, 10),
                 ((h3, 10),): (h3 * 8, h3 * 8 + 10),

tests/python/unittest/test_arith_simplify.py

@@ -0,0 +1,38 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+import tvm.testing
+from tvm import tir
+
+
+def test_simplify_reshape_flattened_index():
+    ana = tvm.arith.Analyzer()
+
+    i0 = tir.Var("i0", "int64")
+    i1 = tir.Var("i1", "int64")
+    ana.bind(i0, tvm.ir.Range(0, 8))
+    ana.bind(i1, tvm.ir.Range(0, 3))
+
+    i_flattened = i0 * 3 + i1
+    assert tvm.ir.structural_equal(
+        ana.simplify((i_flattened) // 12 * 12 + (i_flattened) % 12 // 4 * 4 + (i_flattened) % 4),
+        i_flattened,
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/unittest/test_tir_schedule_analysis.py

@@ -126,7 +126,7 @@ def test_suggest_index_map_winograd():
             floordiv(i0, 2),
             floordiv(i1, 2),
             floormod(i0, 2),
-            floormod(((i1 * 4) + floordiv(i2, 32)), 8),
+            floormod(i1, 2) * 4 + floordiv(i2, 32),
             floormod(i2, 32),
             floordiv(i3, 32),
             floormod(i3, 32),
@@ -137,8 +137,8 @@ def test_suggest_index_map_winograd():
     expected_inverse_index_map = IndexMap.from_func(
         lambda i0, i1, i2, i3, i4, i5, i6: (
             ((i0 * 2) + i2),
-            ((i1 * 2) + floordiv(((i3 * 32) + i4), 128)),
-            floormod(((i3 * 32) + i4), 128),
+            i1 * 2 + floordiv(i3, 4),
+            floormod(i3, 4) * 32 + i4,
             ((i5 * 32) + i6),
         )
     )