adjust related interfaces for IterMapLevel

Author: wrongtest-intellif

include/tvm/arith/iter_affine_map.h

@@ -263,10 +263,10 @@ class IterSumExpr : public IterMapExpr {
 enum IterMapLevel {
   // Require the mapping to be bijective.
   Bijective = 0,
-  // Require the mapping to be subjective.
+  // Require the mapping to be surjective.
   Surjective = 1,
-  // Require the mapping to be injective.
-  Injective = 2
+  // No mapping safety check.
+  NoCheck = 3
 };
 
 /*!
@@ -327,7 +327,7 @@ class IterMapResult : public ObjectRef {
  * \param indices The indices to detect pattern for.
  * \param input_iters Map from variable to iterator's range.
  * \param predicate The predicate constraints on the input iterators
- * \param check_level The iter mapping check level.
+ * \param check_level The iter mapping checking level.
  * \param analyzer Analyzer used to get context information.
  * \param simplify_trivial_iterators If true, iterators with extent of
  *           1 will be replaced with a constant value.
@@ -345,12 +345,12 @@ IterMapResult DetectIterMap(const Array<PrimExpr>& indices, const Map<Var, Range
  * \param indices The indices to detect pattern for.
  * \param input_iters Map from variable to iterator's range.
  * \param input_pred The predicate constraints on the input iterators
- * \param require_bijective A boolean flag that indicates whether the mapping should be bijective.
+ * \param check_level The iter mapping checking level.
  *
  * \return The indices after rewrite
  */
 Array<PrimExpr> IterMapSimplify(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
-                                const PrimExpr& input_pred, bool require_bijective);
+                                const PrimExpr& input_pred, IterMapLevel check_level);
 
 /*!
  * \brief Apply the inverse of the affine transformation to the outputs.
@@ -390,7 +390,7 @@ Map<Var, PrimExpr> InverseAffineIterMap(const Array<IterSumExpr>& iter_map,
  * \param input_iters Map from variable to iterator's range.
  * \param sub_iters Iterators of subspace.
  * \param predicate The predicate constraints on the input iterators
- * \param require_bijective A boolean flag that indicates whether the mapping should be bijective.
+ * \param check_level The iter mapping checking level.
  * \param analyzer Analyzer used to get context information.
  *
  * \return The result list has length len(bindings) + 1
@@ -403,7 +403,7 @@ Map<Var, PrimExpr> InverseAffineIterMap(const Array<IterSumExpr>& iter_map,
 Array<Array<IterMark>> SubspaceDivide(const Array<PrimExpr>& bindings,
                                       const Map<Var, Range>& input_iters,
                                       const Array<Var>& sub_iters, const PrimExpr& predicate,
-                                      bool require_bijective, arith::Analyzer* analyzer);
+                                      IterMapLevel check_level, arith::Analyzer* analyzer);
 
 /*!
  * \brief Given an expression that may contain IterMapExpr, transform it to normal PrimExpr.

python/tvm/arith/iter_affine_map.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """ Iterator (quasi)affine mapping patterns."""
+from enum import IntEnum
 import tvm._ffi
 from tvm.runtime import Object
 from tvm.ir import PrimExpr
@@ -88,11 +89,35 @@ def __init__(self, args, base):
         self.__init_handle_by_constructor__(_ffi_api.IterSumExpr, args, base)
 
 
+class IterMapLevel(IntEnum):
+    """Possible kinds of iter mapping check level."""
+
+    Bijective = 0
+    Surjective = 1
+    NoCheck = 3
+
+    @staticmethod
+    def from_str(name: str):
+        """Helper to create level enum from string"""
+        if name is None:
+            return IterMapLevel.NoCheck
+        name = name.lower()
+        if name == "bijective":
+            check_level = IterMapLevel.Bijective
+        elif name == "surjective":
+            check_level = IterMapLevel.Surjective
+        elif name == "nocheck":
+            check_level = IterMapLevel.NoCheck
+        else:
+            raise ValueError(f"Unknown check level {name}")
+        return check_level
+
+
 def detect_iter_map(
     indices,
     input_iters,
     predicate=True,
-    require_bijective=False,
+    check_level=IterMapLevel.Surjective,
     simplify_trivial_iterators=True,
 ):
     """Detect if indices can be written as mapped iters from input iters
@@ -108,8 +133,8 @@ def detect_iter_map(
     predicate : PrimExpr
         The predicate constraints on the input iterators
 
-    require_bijective : bool
-        A boolean flag that indicates whether the mapping should be bijective
+    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
@@ -122,9 +147,13 @@ def detect_iter_map(
         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.DetectIterMap(
-        indices, input_iters, predicate, require_bijective, simplify_trivial_iterators
-    ).indices
+        indices, input_iters, predicate, check_level, simplify_trivial_iterators
+    )
 
 
 def normalize_iter_map_to_expr(expr):
@@ -143,7 +172,9 @@ def normalize_iter_map_to_expr(expr):
     return _ffi_api.NormalizeIterMapToExpr(expr)
 
 
-def subspace_divide(bindings, input_iters, sub_iters, predicate=True, require_bijective=False):
+def subspace_divide(
+    bindings, input_iters, sub_iters, predicate=True, check_level=IterMapLevel.Surjective
+):
     """Detect if bindings can be written as
     [a_0*e_0 + b_0 + c_0, a_1*e_1 + b_1, ..., a_n*e_n + b_n]
     where a = some-quasi-affine-iter-map(input_iters set_minus sub_iters)
@@ -172,8 +203,8 @@ def subspace_divide(bindings, input_iters, sub_iters, predicate=True, require_bi
     predicate : PrimExpr
         The predicate constraints on the input iterators
 
-    require_bijective : bool
-        A boolean flag that indicates whether the bindings should be bijective
+    check_level : Union[str, IterMapLevel]
+        Checking level of iteration mapping
 
     Returns
     -------
@@ -185,7 +216,9 @@ def subspace_divide(bindings, input_iters, sub_iters, predicate=True, require_bi
         len(bindings): the predicate of outer space and inner space
         Empty array if no match can be found.
     """
-    return _ffi_api.SubspaceDivide(bindings, input_iters, sub_iters, predicate, require_bijective)
+    if isinstance(check_level, str):
+        check_level = IterMapLevel.from_str(check_level)
+    return _ffi_api.SubspaceDivide(bindings, input_iters, sub_iters, predicate, check_level)
 
 
 def inverse_affine_iter_map(iter_map, outputs):