[ARITH] Enhance IterMapSimplify for symbolic

This PR refactors and enhances DetectIterMap and IterMapSimplify
to enable symbolic shape simplification. Specifically, we add
a routine to combine multiple IterSplitExpr into one if they
come from the same source.

It is helpful to distinguish iterator from normal constants
in the simplification process. IterMapSimplify takes advantage
of these information.

This improvements is helpful to simplify the indices in flattened buffer
when there is symbolic shape involved and normal simplifier.

Also updated FlattenBuffer to take benefit of the enhanced simplifier.
Test cases are added.

Author: tqchen

include/tvm/arith/iter_affine_map.h

@@ -349,12 +349,13 @@ IterMapResult DetectIterMap(const Array<PrimExpr>& indices, const Map<Var, Range
  * \param input_iters Map from variable to iterator's range.
  * \param input_pred The predicate constraints on the input iterators
  * \param check_level The iter mapping checking level.
+ * \param analyzer Analyzer used to get context information.
  * \param simplify_trivial_iterators If true, iterators with unit extents are simplified
  * \return The indices after rewrite
  */
 Array<PrimExpr> IterMapSimplify(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
                                 const PrimExpr& input_pred, IterMapLevel check_level,
-                                bool simplify_trivial_iterators = true);
+                                arith::Analyzer* analyzer, bool simplify_trivial_iterators = true);
 
 /*!
  * \brief Apply the inverse of the affine transformation to the outputs.

python/tvm/arith/__init__.py

@@ -30,6 +30,7 @@
 from .iter_affine_map import IterMapExpr, IterMark, IterSplitExpr, IterSumExpr
 from .iter_affine_map import (
     detect_iter_map,
+    iter_map_simplify,
     normalize_iter_map_to_expr,
     subspace_divide,
     inverse_affine_iter_map,

python/tvm/arith/iter_affine_map.py

@@ -156,6 +156,49 @@ def detect_iter_map(
     )
 
 
+def iter_map_simplify(
+    indices,
+    input_iters,
+    predicate=True,
+    check_level=IterMapLevel.Surjective,
+    simplify_trivial_iterators=True,
+):
+    """Simplify the indices using iter map detection.
+
+    Parameters
+    ----------
+    indices : List[PrimExpr]
+        The input indices
+
+    input_iters : Map[Var, Range]
+        The domain of each input iterators.
+
+    predicate : PrimExpr
+        The predicate constraints on the input iterators
+
+    check_level : Union[str, IterMapLevel]
+        Checking level of iteration mapping
+
+    simplify_trivial_iterators: bool
+        If true, iterators with extent of 1 will be replaced with a
+        constant value.
+
+    Returns
+    -------
+    results : IterMapResult
+        The iter map matching result.
+        The result's .indices is empty array if no match can be found.
+
+    """
+    if isinstance(check_level, str):
+        check_level = IterMapLevel.from_str(check_level)
+    elif check_level is None:
+        check_level = IterMapLevel.NoCheck
+    return _ffi_api.IterMapSimplify(
+        indices, input_iters, predicate, check_level, simplify_trivial_iterators
+    )
+
+
 def normalize_iter_map_to_expr(expr):
     """Given an IterMapExpr, transform it to normal PrimExpr
 

src/arith/canonical_simplify.cc

@@ -27,6 +27,7 @@
 
 #include "const_fold.h"
 #include "pattern_match.h"
+#include "product_normal_form.h"
 #include "rewrite_simplify.h"
 
 namespace tvm {
@@ -808,12 +809,17 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const MulNode* op) {
   }
 
   // normal path.
+  // this only happens when b is symbolic
   a = Normalize(a);
   b = Normalize(b);
-  if (op->a.same_as(a) && op->b.same_as(b)) {
+
+  PrimExpr ret = MulAndNormalize(a, b);
+  const MulNode* mul = ret.as<MulNode>();
+
+  if (mul && mul->a.same_as(op->a) && mul->b.same_as(op->b)) {
     return GetRef<PrimExpr>(op);
   } else {
-    return Mul(a, b);
+    return ret;
   }
 }