chore: import upstream snapshot with attribution

python/tvm/arith/__init__.py

@@ -0,0 +1,40 @@
+# isort: skip_file
+# 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.
+"""Integer bound analysis, simplification and pattern detection."""
+
+from .int_set import (
+    IntSet,
+    IntervalSet,
+    PresburgerSet,
+    estimate_region_lower_bound,
+    estimate_region_strict_bound,
+    estimate_region_upper_bound,
+)
+from .analyzer import ModularSet, ConstIntBound, Analyzer, ProofStrength, Extension, CompareResult
+from .bound import deduce_bound
+from .pattern import detect_linear_equation, detect_clip_bound
+from .int_solver import solve_linear_equations, solve_linear_inequalities
+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,
+    normalize_to_iter_sum,
+    subspace_divide,
+    inverse_affine_iter_map,
+)

python/tvm/arith/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs for tvm.arith"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("arith", __name__)

python/tvm/arith/analyzer.py

@@ -0,0 +1,529 @@
+# 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.
+# pylint: disable=invalid-name
+"""Arithmetic data structure and utility"""
+
+import enum
+
+import tvm_ffi
+
+from tvm import ir, tirx
+from tvm.arith import IntSet
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+class ProofStrength(enum.IntEnum):
+    """Proof strength of the analysis"""
+
+    DEFAULT = 0
+    SYMBOLIC_BOUND = 1
+
+
+class CompareResult(enum.IntEnum):
+    """Result of a transitive comparison.
+
+    Values must match the C++ ``arith::CompareResult`` enum.
+    """
+
+    INCONSISTENT = 0
+    EQ = 1
+    LT = 2
+    LE = 3
+    GT = 4
+    GE = 5
+    NE = 6
+    UNKNOWN = 7
+
+
+class Extension(enum.Flag):
+    """Extensions enabled for RewriteSimplifier
+
+    Values should match `RewriteSimplifier::Extensions`
+    """
+
+    NoExtensions = 0
+    TransitivelyProveInequalities = 1 << 0
+    ConvertBooleanToAndOfOrs = 1 << 1
+    ApplyConstraintsToBooleanBranches = 1 << 2
+    ComparisonOfProductAndSum = 1 << 3
+
+
+@tvm_ffi.register_object("arith.ModularSet")
+class ModularSet(Object):
+    """Represent range of (coeff * x + base) for x in Z"""
+
+    def __init__(self, coeff, base):
+        self.__init_handle_by_constructor__(_ffi_api.ModularSet, coeff, base)
+
+
+@tvm_ffi.register_object("arith.ConstIntBound")
+class ConstIntBound(Object):
+    """Represent constant integer bound
+
+    Parameters
+    ----------
+    min_value : int
+        The minimum value of the bound.
+
+    max_value : int
+        The maximum value of the bound.
+    """
+
+    POS_INF = (1 << 63) - 1
+    NEG_INF = -POS_INF
+
+    def __init__(self, min_value, max_value):
+        self.__init_handle_by_constructor__(_ffi_api.ConstIntBound, min_value, max_value)
+
+
+class ConstraintScope:
+    """Constraint scope.
+
+    Parameters
+    ----------
+    fenter : function
+        A function that will be called to create an enter context.
+
+    Note
+    ----
+    Do not create object directly, use Analyzer.constraint_scope
+    """
+
+    def __init__(self, fenter):
+        self._fenter = fenter
+        self._fexit = None
+
+    def __enter__(self):
+        self._fexit = self._fenter()
+
+    def __exit__(self, ptype, value, trace):
+        self._fexit()
+
+
+@tvm_ffi.register_object("arith.Analyzer")
+class Analyzer(Object):
+    """Integer arithmetic analyzer
+
+    This is a stateful analyzer class that can be used to perform
+    various symbolic integer analysis. The same analyzer instance can
+    be passed to FFI APIs to share accumulated facts across calls.
+    """
+
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_api.Analyzer)
+
+    @property
+    def is_z3_enabled(self) -> bool:
+        """Whether this build includes the Z3 backend (``USE_Z3=ON``).
+
+        The Z3-specific methods (:py:meth:`get_smtlib2`, :py:meth:`get_z3_stats`,
+        :py:meth:`set_z3_timeout_ms`, :py:meth:`set_z3_rlimit`) only work when
+        this is ``True``.
+        """
+        return bool(_ffi_api.AnalyzerIsZ3Enabled(self))
+
+    def _check_z3_enabled(self) -> None:
+        if not self.is_z3_enabled:
+            raise RuntimeError(
+                "The Z3 backend is not available in this build. "
+                "Rebuild TVM with USE_Z3=ON to use Z3-specific Analyzer APIs."
+            )
+
+    def get_smtlib2(self, expr: tirx.Expr | None = None) -> str:
+        """Get the current Z3 problem in SMT-LIB2 format.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``), since there is no
+            solver state to export. Use :py:attr:`is_z3_enabled` to check first.
+
+        Parameters
+        ----------
+        expr : Optional[Expr]
+            The expression to prove. If provided, its negation is added to the problem.
+        """
+        self._check_z3_enabled()
+        return _ffi_api.AnalyzerGetSMTLIB2(self, expr)
+
+    def set_z3_timeout_ms(self, timeout_ms: int) -> None:
+        """Set Z3 timeout in milliseconds.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
+        Parameters
+        ----------
+        timeout_ms : int
+            The timeout in milliseconds.
+        """
+        self._check_z3_enabled()
+        _ffi_api.AnalyzerSetZ3TimeoutMs(self, timeout_ms)
+
+    def set_z3_rlimit(self, rlimit: int) -> None:
+        """Set Z3 resource limit.
+
+        The resource limit gives deterministic solver budgeting (unlike a wall
+        clock timeout). A value of ``0`` disables the limit.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
+        Parameters
+        ----------
+        rlimit : int
+            The resource limit.
+        """
+        self._check_z3_enabled()
+        _ffi_api.AnalyzerSetZ3RLimit(self, rlimit)
+
+    def get_z3_stats(self) -> str:
+        """Get Z3 solver statistics.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
+        Returns
+        -------
+        stats : str
+            The Z3 statistics.
+        """
+        self._check_z3_enabled()
+        return _ffi_api.AnalyzerGetZ3Stats(self)
+
+    def const_int_bound(self, expr: tirx.Expr) -> ConstIntBound:
+        """Find constant integer bound for expr.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        bound : ConstIntBound
+            The result bound
+        """
+        return _ffi_api.AnalyzerConstIntBound(self, expr)
+
+    def const_int_bound_is_bound(self, var: tirx.Var) -> bool:
+        """Check if a variable is bound to a range.
+
+        Parameters
+        ----------
+        var : tvm.tirx.Var
+            The variable.
+
+        Returns
+        -------
+        result : bool
+            Whether the variable is bound to a range.
+        """
+        return _ffi_api.AnalyzerConstIntBoundIsBound(self, var)
+
+    def modular_set(self, expr: tirx.Expr) -> ModularSet:
+        """Find a modular set that expr belongs to.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        result : ModularSet
+            The result.
+        """
+        return _ffi_api.AnalyzerModularSet(self, expr)
+
+    def simplify(self, expr: tirx.Expr, steps: int = 2) -> tirx.Expr:
+        """Simplify expression via both rewrite and canonicalization.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+        steps : The simplification runs in the order of
+                rewrite_simplify (step 1) -> canonical_simplify (step 2) ->
+                rewrite_simplify (step 3) -> canonical_simplify (step 4) -> ...
+                param steps controls how many steps to run.
+                Default is 2, i.e., rewrite_simplify + canonical_simplify.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerSimplify(self, expr, steps)
+
+    def clone(self) -> "Analyzer":
+        """Return a deep copy of this analyzer with independent state.
+
+        The returned analyzer carries the same accumulated facts (variable
+        bounds, modular sets, bindings, integer-set domains, literal
+        constraints and transitive comparisons) as this one, but owns its own
+        state: binding or simplifying on either analyzer afterwards does not
+        affect the other. Unlike copying the handle, this is a true deep copy.
+
+        Do not call this while a constraint scope is active on this analyzer.
+
+        Returns
+        -------
+        result : Analyzer
+            A new analyzer holding an independent copy of the facts.
+        """
+        return _ffi_api.AnalyzerClone(self)
+
+    def rewrite_simplify(self, expr: tirx.Expr) -> tirx.Expr:
+        """Simplify expression via rewriting rules.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerRewriteSimplify(self, expr)
+
+    @property
+    def rewrite_simplify_stats(self):
+        return _ffi_api.AnalyzerGetRewriteSimplifyStats(self)
+
+    def reset_rewrite_simplify_stats(self):
+        _ffi_api.AnalyzerResetRewriteSimplifyStats(self)
+
+    def canonical_simplify(self, expr: tirx.Expr) -> tirx.Expr:
+        """Simplify expression via canonicalization.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerCanonicalSimplify(self, expr)
+
+    def int_set(self, expr: tirx.Expr, dom_map: dict[tirx.Var, IntSet] | None = None) -> IntSet:
+        """Compute a symbolic IntSet that covers expr for all values in dom_map.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        dom_map : Optional[Dict[tvm.tirx.Var, tvm.arith.IntSet]]
+            The domain for variables to be relaxed.  When omitted, the analyzer
+            uses the domains of the variables already bound to it.
+
+        Returns
+        -------
+        result : IntSet
+            The result.
+        """
+        return _ffi_api.AnalyzerIntSet(self, expr, dom_map)
+
+    def can_prove(self, expr: tirx.Expr, strength: ProofStrength = ProofStrength.DEFAULT) -> bool:
+        """Check whether we can prove expr to be true.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        strength: ProofStrength
+            The proof strength. When TVM is built with Z3 (``USE_Z3=ON``), the
+            optional Z3 fallback is only consulted at ``SYMBOLIC_BOUND`` or
+            higher, after the native analyzers fail to prove the predicate.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerCanProve(self, expr, strength)
+
+    def set_maximum_rewrite_steps(self, maximum: int) -> None:
+        """Set the maximum allowed number of rewrite-simplify steps.
+
+        When a positive limit is set, the simplifier raises an exception once
+        it exceeds that number of rewrite steps.  This is useful for guarding
+        against performance regressions in tests.
+
+        Parameters
+        ----------
+        maximum : int
+            The maximum number of rewrite steps, or a non-positive value to
+            allow an unlimited number of steps.
+        """
+        _ffi_api.AnalyzerSetMaximumRewriteSteps(self, maximum)
+
+    def bind(
+        self,
+        var: tirx.Var,
+        expr: tirx.Expr | ir.Range,
+        allow_override: bool = False,
+    ) -> None:
+        """Bind a variable to the expression.
+
+        Parameters
+        ----------
+        var : tvm.tirx.Var
+            The variable.
+
+        expr : Union[tirx.Expr, ir.Range]
+            The expression or the range to bind to.
+
+        allow_override : bool
+            Whether to allow overriding an existing binding for the variable.
+        """
+        return _ffi_api.AnalyzerBind(self, var, expr, allow_override)
+
+    def constraint_scope(self, constraint: tirx.Expr) -> ConstraintScope:
+        """Create a constraint scope.
+
+        Parameters
+        ----------
+        constraint : Expr
+            The constraint expression.
+
+        returns
+        -------
+        scope : ConstraintScope
+            The constraint scope
+
+        Examples
+        --------
+        .. code-block:: python
+
+          x = te.var("x")
+          analyzer = tvm.arith.Analyzer()
+          with analyzer.constraint_scope(x % 3 == 0):
+              # constraint in effect
+              assert analyzer.modular_set(x).coeff == 3
+          # constraint no longer in effect
+          assert analyzer.modular_set(x).coeff != 3
+        """
+
+        def _fenter():
+            return _ffi_api.AnalyzerEnterConstraintContext(self, constraint)
+
+        return ConstraintScope(_fenter)
+
+    def update(
+        self, var: tirx.Var, info: ConstIntBound | ModularSet | IntSet, override: bool = False
+    ) -> None:
+        """Update information about var.
+
+        Parameters
+        ----------
+        var : tvm.tirx.Var
+            The variable.
+
+        info : Union[ConstIntBound, ModularSet, IntSet]
+            Related information.  A ``ConstIntBound`` updates the constant
+            integer bound, a ``ModularSet`` updates the modular set, and an
+            ``IntSet`` updates the integer-set domain of ``var``.
+
+        override : bool
+            Whether allow override.
+        """
+        if isinstance(info, ConstIntBound):
+            _ffi_api.AnalyzerConstIntBoundUpdate(self, var, info, override)
+        elif isinstance(info, ModularSet):
+            _ffi_api.AnalyzerModularSetUpdate(self, var, info, override)
+        elif isinstance(info, IntSet):
+            _ffi_api.AnalyzerIntSetUpdate(self, var, info, override)
+        else:
+            raise TypeError(f"Do not know how to handle type {type(info)}")
+
+    def can_prove_equal(self, lhs: tirx.Expr, rhs: tirx.Expr) -> bool:
+        """Whether we can prove that lhs == rhs
+
+        Parameters
+        ----------
+        lhs: Expr
+            The left-hand side of the comparison
+
+        rhs: Expr
+            The right-hand side of the comparison
+
+        Returns
+        -------
+        result: bool
+            Whether we can prove that lhs == rhs
+        """
+        return _ffi_api.AnalyzerCanProveEqual(self, lhs, rhs)
+
+    def try_compare(
+        self, lhs: tirx.Expr, rhs: tirx.Expr, propagate_inequalities: bool = True
+    ) -> CompareResult:
+        """Compare lhs and rhs using previously provided known comparisons.
+
+        Parameters
+        ----------
+        lhs : Expr
+            The left-hand side of the comparison.
+
+        rhs : Expr
+            The right-hand side of the comparison.
+
+        propagate_inequalities : bool
+            If true, attempt to find a sequence of transitive inequalities that
+            allow lhs and rhs to be compared.
+
+        Returns
+        -------
+        result : CompareResult
+            The most specific result that can be proven about the comparison.
+            Returns ``CompareResult.UNKNOWN`` when nothing can be proven.
+        """
+        return CompareResult(_ffi_api.AnalyzerTryCompare(self, lhs, rhs, propagate_inequalities))
+
+    @property
+    def enabled_extensions(self) -> Extension:
+        """Return the currently enabled extensions"""
+        value = _ffi_api.AnalyzerGetEnabledExtensions(self)
+        return Extension(value)
+
+    @enabled_extensions.setter
+    def enabled_extensions(self, flags: int | Extension):
+        """Enable extensions for the analyzer
+
+        Parameters
+        ----------
+        flags: Union[int,Extension]
+
+            The extensions to enable.
+        """
+        flags = Extension(flags).value
+        _ffi_api.AnalyzerSetEnabledExtensions(self, flags)

python/tvm/arith/bound.py

@@ -0,0 +1,40 @@
+# 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.
+"""Bound deduction."""
+
+from . import _ffi_api
+
+
+def deduce_bound(var, cond, hint_map, relax_map):
+    """Deduce the bound of the target variable in the cond.
+
+    Parameters
+    ----------
+    var : tvm.tirx.Var
+        The target variable to be deduced.
+
+    cond : Expr
+        The condition
+
+    hint_map : Map[tvm.tirx.Var, IntSet]
+        Domain of variables used to help deduction.
+
+    relax_map : Map[tvm.tirx.Var, IntSet]
+        The fomain of the variables to be relaxed
+        using the provided domain.
+    """
+    return _ffi_api.DeduceBound(var, cond, hint_map, relax_map)

python/tvm/arith/int_set.py

@@ -0,0 +1,213 @@
+# 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.
+"""Integer set."""
+
+import tvm_ffi
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("ir.IntSet")
+class IntSet(Object):
+    """Represent a set of integer in one dimension."""
+
+    def is_nothing(self):
+        """Whether the set represent nothing"""
+        return _ffi_api.IntSetIsNothing(self)
+
+    def is_everything(self):
+        """Whether the set represent everything"""
+        return _ffi_api.IntSetIsEverything(self)
+
+    @staticmethod
+    def vector(vec):
+        """Construct an integer set that covers the vector expr
+
+        Parameters
+        ----------
+        vec : Expr
+            The vector expression.
+
+        Returns
+        -------
+        rset : IntSet
+            The result set.
+        """
+        return _ffi_api.intset_vector(vec)
+
+    @staticmethod
+    def single_point(point):
+        """Construct a point set.
+
+        Parameters
+        ----------
+        point : Expr
+            The vector expression.
+
+        Returns
+        -------
+        rset : IntSet
+            The result set.
+        """
+        return _ffi_api.intset_single_point(point)
+
+
+@tvm_ffi.register_object("arith.IntervalSet")
+class IntervalSet(IntSet):
+    """Represent set of continuous interval [min_value, max_value]
+
+    Parameters
+    ----------
+    min_value : Expr
+        The minimum value in the interval.
+
+    max_value : Expr
+        The maximum value in the interval.
+    """
+
+    def __init__(self, min_value, max_value):
+        self.__init_handle_by_constructor__(_ffi_api.IntervalSet, min_value, max_value)
+
+
+@tvm_ffi.register_object("arith.PresburgerSet")
+class PresburgerSet(IntSet):
+    """Represent of Presburger Set"""
+
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_api.PresburgerSet)
+
+
+def estimate_region_lower_bound(region, var_dom, predicate, analyzer=None):
+    """Analyze the region with affine map, given the domain of variables and their predicate
+    Some subregion may be discarded during the lower-bound analysis.
+
+    Parameters
+    ----------
+    region : List[Range]
+        The region to be analyzed.
+
+    var_dom : Dict[tvm.tirx.Var, Range]
+        The ranges of the variables
+
+    predicate : Expr
+        The predicate for the affine map
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    ----------
+    region_int_set : Optional[List[IntSet]]
+        None if the detection fails, or an array of IntSets as the result of analysis
+    """
+    return _ffi_api.EstimateRegionLowerBound(region, var_dom, predicate, analyzer)
+
+
+def estimate_region_strict_bound(region, var_dom, predicate, analyzer=None):
+    """Analyze the region with affine map, given the domain of variables and their predicate
+    The result should be strict, i.e. no region is discarded or relaxed.
+
+    Parameters
+    ----------
+    region : List[Range]
+        The region to be analyzed.
+
+    var_dom : Dict[tvm.tirx.Var, Range]
+        The ranges of the variables
+
+    predicate : Expr
+        The predicate for the affine map
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    ----------
+    region_int_set : Optional[List[IntSet]]
+        None if the detection fails, or an array of IntSets as the result of analysis
+    """
+    return _ffi_api.EstimateRegionStrictBound(region, var_dom, predicate, analyzer)
+
+
+def estimate_region_upper_bound(region, var_dom, predicate, analyzer=None):
+    """Analyze the region with affine map, given the domain of variables and their predicate
+    Relaxation of the region may be used in upper-bound analysis,
+    i.e. some extra region may be added to the result.
+
+    Parameters
+    ----------
+    region : List[Range]
+        The region to be analyzed.
+
+    var_dom : Dict[tvm.tirx.Var, Range]
+        The ranges of the variables
+
+    predicate : Expr
+        The predicate for the affine map
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    ----------
+    region_int_set : List[IntSet]
+        an array of IntSets as the result of analysis
+    """
+    return _ffi_api.EstimateRegionUpperBound(region, var_dom, predicate, analyzer)
+
+
+def pos_inf():
+    """Returns the symbolic positive infinity
+
+    Returns
+    ----------
+    pos_inf : tvm.tirx.Var
+        A symbolic var that indicates positive infinity
+    """
+    return _ffi_api.PosInf()
+
+
+def neg_inf():
+    """Returns the symbolic positive infinity
+
+    Returns
+    ----------
+    neg_inf : tvm.tirx.Var
+        A symbolic var that indicates positive infinity
+    """
+    return _ffi_api.NegInf()
+
+
+def union_lower_bound(sets):
+    """Create a lower-bound of union set, where some of the segments may be dropped
+
+    Parameters
+    ----------
+    sets : List[IntSet]
+        The sets to be combined
+
+    Returns
+    ----------
+    union_lower_bound : List[IntSet]
+        An N-dimensional integer set, the lower bound of the union
+    """
+    return _ffi_api.UnionLowerBound(sets)

python/tvm/arith/int_solver.py

@@ -0,0 +1,183 @@
+# 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.
+"""integer constraints data structures and solvers"""
+
+import tvm_ffi
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("arith.IntGroupBounds")
+class IntGroupBounds(Object):
+    """Represent integer grouped bounds which are classified into
+       lower bounds (include), upper bounds (include) and equalities.
+
+    Parameters
+    ----------
+    coef : tvm.ir.Expr
+        The coefficient. Must be integer type.
+        coef * var >= lower
+        coef * var == equal
+        coef * var >= upper
+    lower : List[tvm.ir.Expr]
+        the lower bounds (include)
+    equal : List[tvm.ir.Expr]
+        equalities
+    upper : List[tvm.ir.Expr]
+        the upper bounds (include)
+    """
+
+    def __init__(self, coef, lower, equal, upper):
+        self.__init_handle_by_constructor__(_ffi_api.IntGroupBounds, coef, lower, equal, upper)
+
+    @staticmethod
+    def from_range(rng):
+        """Construct a IntGroupedBounds by Range.
+
+        Parameters
+        ----------
+        rng : tvm.ir.Range
+
+
+        Returns
+        -------
+        ret : Range
+            The constructed range.
+        """
+        return _ffi_api.IntGroupBounds_from_range(rng)
+
+    def find_best_range(self):
+        """Return the best range from the grouped bounds.
+        None if (-inf, +inf).
+        """
+        return _ffi_api.IntGroupBounds_FindBestRange(self)
+
+
+@tvm_ffi.register_object("arith.IntConstraints")
+class IntConstraints(Object):
+    """Represent a set of integer constraints including variables, their ranges and
+       the relations between them (either equations or inequalities)
+
+    Parameters
+    ----------
+    variables : List[tvm.tirx.Var]
+        The variables in the constraints. Must be integers
+    ranges    : Map[tvm.tirx.Var, tvm.ir.Range]
+        The ranges of the variables.
+    relations : List[tvm.ir.Expr]
+        The relations between the variables (either equations or inequalities)
+    """
+
+    def __init__(self, variables, ranges, relations):
+        self.__init_handle_by_constructor__(_ffi_api.IntConstraints, variables, ranges, relations)
+
+
+@tvm_ffi.register_object("arith.IntConstraintsTransform")
+class IntConstraintsTransform(Object):
+    """We can have different set of variables to represent the same integer constraints.
+       For example, the following two constrains are equivalent,
+       {a + b = 0 | a >= 0, b >= 0} and
+       {m - n = 0 | m >= 0, n <= 0}
+       This data structure represents the transformation
+       between two equivalent integer constraints.
+       In the above example,
+       src        : {a + b = 0 | a >= 0, b >= 0}
+       dst        : {m - n = 0 | m >= 0, n <= 0}
+       src_to_dst : {a -> m, b -> -n}
+       dst_to_src : {m -> a, n -> -b}
+
+    Parameters
+    ----------
+    src : arith.IntConstraints
+        source integer constraints, e.g., {a + b = 0 | a >= 0, b >= 0}
+    dst : arith.IntConstraints
+        integer constraints equivalent to the source, e.g., {m - n = 0 | m >= 0, n <= 0}
+    src_to_dst : Map[tvm.tirx.Var, tvm.ir.Expr]
+        mapping from variables in the src to the variables in the dst,
+                e.g., {a -> m, b -> -n}
+    dst_to_src : Map[tvm.tirx.Var, tvm.ir.Expr]
+        mapping from variables in the dst to the variables in the src,
+        e.g., {m -> a, n -> -b}
+    """
+
+    def __init__(self, src, dst, src_to_dst, dst_to_src):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IntConstraintsTransform, src, dst, src_to_dst, dst_to_src
+        )
+
+
+def solve_linear_equations(equations, variables=None, ranges=None):
+    """Solve linear equations.
+
+    Parameters
+    ----------
+    equations: List[tvm.ir.Expr] or IntConstraints
+        The equations of the variables
+    variables : Optional[List[tvm.tirx.Var]]
+        The variables in the system.
+    ranges    : Optional[Map[tvm.tirx.Var, tvm.ir.Range]]
+        The ranges of the variables.
+
+    Returns
+    -------
+    int_constraints_transform : IntConstraintsTransform
+        New integer constraints, with less variables (if the problem is NOT of full rank),
+        or no variable (if the problem is of full rank),
+        or an empty integer constraints (if the problem is unsolvable).
+        It also provides the ranges of the variables in the new system,
+        as well as inequalities inferred from the problem.
+        You can get the mapping from the original variables to the solution via
+        int_constraints_transform.src_to_dst.
+    """
+    if isinstance(equations, IntConstraints):
+        return _ffi_api.SolveLinearEquations(equations)
+    return _ffi_api.SolveLinearEquations(variables, ranges, equations)
+
+
+def solve_linear_inequalities(equations, variables=None, ranges=None, deskew_range=False):
+    """Solve linear inequalities.
+
+    Parameters
+    ----------
+    equations   : List[tvm.ir.Expr] or IntConstraints
+        The inequalities of the variables
+    variables   : Optional[List[tvm.tirx.Var]]
+        The variables in the system.
+    ranges      : Optional[Map[tvm.tirx.Var, tvm.ir.Range]]
+        The ranges of the variables.
+    deskew_range: Optional[bool]
+        Whether deskew the result ranges to be started from zero.
+        Default false.
+
+    Returns
+    -------
+    ret_ranges: IntConstraints or IntConstraintsTransform
+        The result ranges for each variables.
+        Constrains that cannot be transformed to Range will be stored in IntConstraints.relations.
+        If deskew_range is set (=True), the result ranges will be deskewed to be started from zero.
+        New variables are created accordingly therefore IntConstraintsTransform is returned.
+    """
+    solver = (
+        _ffi_api.SolveInequalitiesDeskewRange if deskew_range else _ffi_api.SolveInequalitiesToRange
+    )
+    if isinstance(equations, IntConstraints):
+        assert variables is None
+        assert ranges is None
+        return solver(equations)
+    return solver(variables, ranges, equations)

python/tvm/arith/iter_affine_map.py

@@ -0,0 +1,375 @@
+# 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.
+"""Iterator (quasi)affine mapping patterns."""
+
+from enum import IntEnum
+
+import tvm_ffi
+
+from tvm.ir import Expr
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("arith.IterMapExpr")
+class IterMapExpr(Expr):
+    """Base class of all IterMap expressions."""
+
+
+@tvm_ffi.register_object("arith.IterMark")
+class IterMark(Object):
+    """Mark the source as an iterator in [0, extent).
+
+    Parameters
+    ----------
+    source : Expr.
+        The source expression.
+
+    extent : Expr
+        The extent of the iterator.
+    """
+
+    def __init__(self, source, extent):
+        self.__init_handle_by_constructor__(_ffi_api.IterMark, source, extent)
+
+
+@tvm_ffi.register_object("arith.IterSplitExpr")
+class IterSplitExpr(IterMapExpr):
+    """Split of an iterator.
+
+    result = floormod(floordiv(source, lower_factor), extent) * scale
+
+    Parameters
+    ----------
+    source : IterMark
+        The source marked iterator.
+
+    lower_factor : Expr
+        The lower factor to split the domain.
+
+    extent : Expr
+        The extent of the split.
+
+    scale : Expr
+        Additional scale to the split.
+    """
+
+    def __init__(self, source, lower_factor, extent, scale):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IterSplitExpr, source, lower_factor, extent, scale
+        )
+
+
+@tvm_ffi.register_object("arith.IterSumExpr")
+class IterSumExpr(IterMapExpr):
+    """Fuse multiple iterators by summing them with scaling.
+
+    result = sum(args) + base
+
+    Parameters
+    ----------
+    args : List[IterSplitExpr]
+        The input to the sum expression.
+
+    base : Expr
+        The base offset.
+    """
+
+    def __init__(self, args, base):
+        self.__init_handle_by_constructor__(_ffi_api.IterSumExpr, args, base)
+
+
+@tvm_ffi.register_object("arith.IterMapResult")
+class IterMapResult(Object):
+    """Result of iter map detection."""
+
+
+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,
+    check_level=IterMapLevel.Surjective,
+    simplify_trivial_iterators=True,
+    analyzer=None,
+):
+    """Detect if indices can be written as mapped iters from input iters
+
+    Parameters
+    ----------
+    indices : List[Expr]
+        The input indices
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of each input iterators.
+
+    predicate : Expr
+        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.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    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.DetectIterMap(
+        indices, input_iters, predicate, check_level, simplify_trivial_iterators, analyzer
+    )
+
+
+def normalize_to_iter_sum(index, input_iters, analyzer=None):
+    """Normalize expr to iter sum.
+
+    The normalized result ensures that
+    each scale is in the form of (symbol_prod) * cscale
+    It will also sort in desc order by cscale then len(symbol_prod).
+
+    Parameters
+    ----------
+    index : Expr
+        The input index
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of each input iterators.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    -------
+    iter_sum: IterSumExpr
+        The result iter sum
+
+    Note
+    ----
+    This function does best effort detection, so some undetected
+    part can go into iter_sum.base
+
+    This function is useful to decide the stride multiplier and
+    division factor in buffer access patterns.
+    """
+    return _ffi_api.NormalizeToIterSum(index, input_iters, analyzer)
+
+
+def iter_map_simplify(
+    indices,
+    input_iters,
+    predicate=True,
+    check_level=IterMapLevel.Surjective,
+    simplify_trivial_iterators=True,
+    analyzer=None,
+):
+    """Simplify the indices using iter map detection.
+
+    Parameters
+    ----------
+    indices : List[Expr]
+        The input indices
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of each input iterators.
+
+    predicate : Expr
+        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.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    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, analyzer
+    )
+
+
+def normalize_iter_map_to_expr(expr):
+    """Given an IterMapExpr, transform it to normal Expr
+
+    Parameters
+    ----------
+    expr : IterMapExpr
+        the input IterMapExpr
+
+    Returns
+    -------
+    result : Expr
+        the corresponding normal Expr
+    """
+    return _ffi_api.NormalizeIterMapToExpr(expr)
+
+
+def subspace_divide(
+    bindings,
+    input_iters,
+    sub_iters,
+    predicate=True,
+    check_level=IterMapLevel.Surjective,
+    simplify_trivial_iterators=True,
+    analyzer=None,
+):
+    """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)
+        b = some-quasi-affine-iter-map(sub_iters)
+        c is constant symbols
+        e is the extent of b
+
+    For example::
+
+        z*12 + y*3 + x + c = (z*4+y)*3 + x
+        bindings = [z*12 + y*3 + x + c]
+        input_iters = [z, y, x]
+        sub_iter = [x]
+        Then the result will be [a, b] where
+        a = [z*4 + y]
+        b = [x]
+
+    Parameters
+    ----------
+    bindings : List[Expr]
+        The input bindings
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of input iterator, which is the basis of the whole space
+
+    sub_iters : Array[tvm.tirx.Var]
+        The subset of input_iters, which is the basis of the subspace
+
+    predicate : Expr
+        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.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    -------
+    results : List[List[Expr]]
+        The result list has length ``len(bindings) + 1``.
+
+        - ``[0, len(bindings))``: The iter map matching result.
+          The inner list is of length 2. The first expr is the basis
+          of the quotient space. The second expr is the basis of the subspace.
+        - ``len(bindings)``: the predicate of outer space and inner space.
+        - Empty array if no match can be found.
+    """
+    if isinstance(check_level, str):
+        check_level = IterMapLevel.from_str(check_level)
+    return _ffi_api.SubspaceDivide(
+        bindings,
+        input_iters,
+        sub_iters,
+        predicate,
+        check_level,
+        simplify_trivial_iterators,
+        analyzer,
+    )
+
+
+def inverse_affine_iter_map(iter_map, outputs):
+    """Apply the inverse of the affine transformation to the outputs.
+    Similar to the back-propagation, starting from the outputs, it visits the DAG of the expressions
+    in reverse topology order and applies the inverse of the affine transformation until it reaches
+    the input. The affine iter map is required to be bijective.
+
+    For example, iter_map = [l0 // 16, l0 % 16], outputs = [output_0, output_1],
+    the affine transformation specified by `iter_map` will be applied to `outputs` and the result
+    will be {l0: ((output_0*16) + output_1)}.
+
+    See also :any:`detect_iter_map`.
+
+    Parameters
+    ----------
+    iter_map : List[IterSumExpr]
+        The bijective affine iter map.
+    outputs : List[Expr]
+        The outputs of the affine transformation.
+
+    Returns
+    -------
+    results : Map[tvm.tirx.Var, Expr]
+        The map from the input to the transformed result.
+    """
+    return _ffi_api.InverseAffineIterMap(iter_map, outputs)

python/tvm/arith/pattern.py

@@ -0,0 +1,61 @@
+# 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.
+"""Detect common patterns."""
+
+from . import _ffi_api
+
+
+def detect_linear_equation(expr, var_list):
+    """Match `expr = sum_{i=0}^{n-1} var[i] * coeff[i] + coeff[n]`
+
+    Where coeff[i] and base are invariant of var[j] for all i and j.
+
+    Parameters
+    ----------
+    expr : Expr
+        The expression to be matched.
+
+    var_list : List[tvm.tirx.Var]
+        A list of variables.
+
+    Returns
+    -------
+    coeff : List[Expr]
+        A list of co-efficients if the match is successful.
+        An empty list if the match failed.
+    """
+    return _ffi_api.DetectLinearEquation(expr, var_list)
+
+
+def detect_clip_bound(expr, var_list):
+    """Detect if expression corresponds to clip bound of the vars
+
+    Parameters
+    ----------
+    expr : Expr
+        The expression to be matched.
+
+    var_list : List[tvm.tirx.Var]
+        A list of variables.
+
+    Returns
+    -------
+    coeff : List[Expr]
+        `concat([min_value[i], max_value[i]] for i, v in enumerate(var_list))`
+        An empty list if the match failed.
+    """
+    return _ffi_api.DetectClipBound(expr, var_list)