chore: import upstream snapshot with attribution

tests/python/arith/test_arith_analyzer_object.py

@@ -0,0 +1,287 @@
+# 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 pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.arith.analyzer import CompareResult, Extension
+from tvm.runtime import Object
+
+
+def test_analyzer_is_ffi_object_with_persistent_state():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    assert isinstance(analyzer, Object)
+
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+    assert analyzer.const_int_bound_is_bound(x)
+    assert analyzer.can_prove(x < 8)
+    assert not analyzer.can_prove(x < 4)
+
+    bound = analyzer.const_int_bound(x + 1)
+    assert bound.min_value == 1
+    assert bound.max_value == 8
+
+
+def test_analyzer_object_constraint_scope_and_override_bind():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    with analyzer.constraint_scope(x % 3 == 0):
+        assert analyzer.modular_set(x).coeff == 3
+
+    assert analyzer.modular_set(x).coeff != 3
+
+    analyzer = tvm.arith.Analyzer()
+    y = tirx.Var("y", "int64")
+    analyzer.bind(y, tirx.const(4, "int64"))
+    tvm.ir.assert_structural_equal(analyzer.simplify(y + 1), tirx.const(5, "int64"))
+
+    analyzer.bind(y, tirx.const(8, "int64"), allow_override=True)
+    tvm.ir.assert_structural_equal(analyzer.simplify(y + 1), tirx.const(9, "int64"))
+
+
+def test_analyzer_object_update_const_int_bound():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    analyzer.update(x, tvm.arith.ConstIntBound(2, 5))
+
+    bound = analyzer.const_int_bound(x + 1)
+    assert bound.min_value == 3
+    assert bound.max_value == 6
+
+
+def test_analyzer_object_update_modular_set():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+
+    assert analyzer.modular_set(x).coeff == 1
+    analyzer.update(x, tvm.arith.ModularSet(4, 0))
+
+    result = analyzer.modular_set(x)
+    assert result.coeff == 4
+    assert result.base == 0
+
+
+def test_analyzer_object_update_int_set():
+    analyzer = tvm.arith.Analyzer()
+    y = tirx.Var("y", "int32")
+
+    analyzer.update(y, tvm.arith.IntervalSet(0, 8))
+
+    int_set = analyzer.int_set(y)
+    assert int_set.min_value.value == 0
+    assert int_set.max_value.value == 8
+
+
+def test_analyzer_object_update_rejects_unknown_info():
+    analyzer = tvm.arith.Analyzer()
+    y = tirx.Var("y", "int32")
+
+    with pytest.raises(TypeError):
+        analyzer.update(y, "not-an-info-object")
+
+
+def test_analyzer_object_can_prove_comparison_predicates():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+
+    assert analyzer.can_prove(x >= 0)
+    assert not analyzer.can_prove(x >= 1)
+    assert analyzer.can_prove(x < 8)
+    assert not analyzer.can_prove(x < 7)
+
+
+def test_analyzer_object_update_const_int_bound_half_space():
+    analyzer = tvm.arith.Analyzer()
+    n = tirx.Var("n", "int32")
+
+    assert not analyzer.can_prove(n >= 0)
+    analyzer.update(n, tvm.arith.ConstIntBound(0, tvm.arith.ConstIntBound.POS_INF))
+    assert analyzer.can_prove(n >= 0)
+
+
+def test_analyzer_object_int_set_from_bound_vars():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+
+    int_set = analyzer.int_set(x + 1)
+    assert int_set.min_value.value == 1
+    assert int_set.max_value.value == 8
+
+
+def test_analyzer_object_set_maximum_rewrite_steps():
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = (x + y) * 2 - x * 2 - y * 2 + tirx.max(x, y) - tirx.min(x, y)
+
+    capped = tvm.arith.Analyzer()
+    capped.set_maximum_rewrite_steps(1)
+    with pytest.raises(RuntimeError):
+        capped.rewrite_simplify(expr)
+
+    # A generous limit must not interfere with normal simplification.
+    relaxed = tvm.arith.Analyzer()
+    relaxed.set_maximum_rewrite_steps(1000)
+    relaxed.rewrite_simplify(expr)
+
+
+def test_analyzer_object_try_compare_transitive():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    z = tirx.Var("z", "int32")
+
+    assert analyzer.try_compare(x, y) == CompareResult.UNKNOWN
+
+    with analyzer.constraint_scope(x < y):
+        with analyzer.constraint_scope(y < z):
+            # Direct known comparison.
+            assert analyzer.try_compare(x, y) == CompareResult.LT
+            # Transitive chain x < y < z is found only when propagation is enabled.
+            assert analyzer.try_compare(x, z) == CompareResult.LT
+            assert analyzer.try_compare(x, z, propagate_inequalities=False) == CompareResult.UNKNOWN
+
+
+def test_analyzer_object_enabled_extensions_round_trip():
+    analyzer = tvm.arith.Analyzer()
+
+    assert analyzer.enabled_extensions == Extension.NoExtensions
+
+    analyzer.enabled_extensions = Extension.ComparisonOfProductAndSum
+    assert analyzer.enabled_extensions == Extension.ComparisonOfProductAndSum
+
+    analyzer.enabled_extensions = Extension.NoExtensions
+    assert analyzer.enabled_extensions == Extension.NoExtensions
+
+
+def test_analyzer_object_rewrite_simplify_stats():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+
+    analyzer.reset_rewrite_simplify_stats()
+    assert analyzer.rewrite_simplify_stats.nodes_visited == 0
+
+    analyzer.rewrite_simplify(x + 0)
+    assert analyzer.rewrite_simplify_stats.nodes_visited > 0
+
+    analyzer.reset_rewrite_simplify_stats()
+    assert analyzer.rewrite_simplify_stats.nodes_visited == 0
+
+
+def test_analyzer_object_state_persists_across_ffi_calls():
+    analyzer = tvm.arith.Analyzer()
+    tile = tirx.Var("tile", "int32")
+    i = tirx.Var("i", "int32")
+    analyzer.bind(tile, tvm.tirx.const(8, "int32"))
+
+    # The same analyzer object is borrowed by the C++ DetectIterMap entry point;
+    # its binding makes the otherwise-undetectable floormod recognizable.
+    result = tvm.arith.detect_iter_map([i % tile], {i: tvm.ir.Range(0, 32)}, analyzer=analyzer)
+    assert len(result.indices) == 1
+
+    # The binding still lives in the same stateful object after the FFI call.
+    tvm.ir.assert_structural_equal(analyzer.simplify(tile), tvm.tirx.const(8, "int32"))
+
+
+def test_analyzer_object_clone_is_independent():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+    y = tirx.Var("y", "int64")
+    z = tirx.Var("z", "int64")
+
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+
+    clone = analyzer.clone()
+    assert clone is not analyzer
+    assert clone.can_prove(x < 8)
+
+    clone.bind(y, tvm.ir.Range(0, 4))
+    assert clone.can_prove(y < 4)
+    assert not analyzer.can_prove(y < 4)
+
+    analyzer.bind(z, tvm.ir.Range(0, 4))
+    assert analyzer.can_prove(z < 4)
+    assert not clone.can_prove(z < 4)
+
+    assert analyzer.can_prove(x < 8)
+    assert clone.can_prove(x < 8)
+
+
+def test_analyzer_object_clone_copies_every_sub_analyzer():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+    w = tirx.Var("w", "int64")
+    v = tirx.Var("v", "int64")
+
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+    analyzer.update(x, tvm.arith.ModularSet(4, 0))
+    analyzer.bind(w, tirx.const(4, "int64"))
+    analyzer.update(v, tvm.arith.IntervalSet(2, 9))
+    analyzer.enabled_extensions = Extension.ComparisonOfProductAndSum
+
+    clone = analyzer.clone()
+
+    assert clone.can_prove(x < 8)
+    assert clone.modular_set(x).coeff == 4
+    tvm.ir.assert_structural_equal(clone.simplify(w + 1), tirx.const(5, "int64"))
+    assert clone.int_set(v).max_value.value == 9
+    assert clone.enabled_extensions == Extension.ComparisonOfProductAndSum
+    assert clone.try_compare(x, tirx.const(0, "int64")) == CompareResult.GE
+
+    t = tirx.Var("t", "int64")
+    clone.update(x, tvm.arith.ModularSet(8, 0), override=True)
+    clone.update(v, tvm.arith.IntervalSet(0, 3), override=True)
+    clone.bind(w, tirx.const(8, "int64"), allow_override=True)
+    clone.bind(t, tvm.ir.Range(0, 4))
+    clone.enabled_extensions = Extension.NoExtensions
+
+    assert analyzer.modular_set(x).coeff == 4
+    assert clone.modular_set(x).coeff == 8
+    assert analyzer.int_set(v).max_value.value == 9
+    assert clone.int_set(v).max_value.value == 3
+    tvm.ir.assert_structural_equal(analyzer.simplify(w + 1), tirx.const(5, "int64"))
+    tvm.ir.assert_structural_equal(clone.simplify(w + 1), tirx.const(9, "int64"))
+    assert analyzer.enabled_extensions == Extension.ComparisonOfProductAndSum
+    assert clone.enabled_extensions == Extension.NoExtensions
+    assert clone.try_compare(t, tirx.const(0, "int64")) == CompareResult.GE
+    assert analyzer.try_compare(t, tirx.const(0, "int64")) == CompareResult.UNKNOWN
+
+
+def test_analyzer_object_clone_resets_rewrite_stats():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+    y = tirx.Var("y", "int64")
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+    analyzer.bind(y, tvm.ir.Range(0, 8))
+    analyzer.simplify((x + y) * 2 - x - y)
+    source_attempts = analyzer.rewrite_simplify_stats.rewrites_attempted
+    assert source_attempts > 0
+
+    clone = analyzer.clone()
+    assert clone.rewrite_simplify_stats.rewrites_attempted == 0
+    assert analyzer.rewrite_simplify_stats.rewrites_attempted == source_attempts
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_canonical_simplify.py

@@ -0,0 +1,538 @@
+# 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.
+# ruff: noqa: E731, F841
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import te, tirx
+from tvm.script import tirx as T
+
+
+class CanonicalChecker:
+    def __init__(self):
+        self.analyzer = tvm.arith.Analyzer()
+
+    def _convert(self, expr):
+        # TODO(Lunderberg): Make utility functions `tirx.convert` and
+        # `relax.convert` that convert to their respective IR types.
+        # Implementation should be in C++, and should only consist of
+        # conversions that are applied automatically through FFI.
+        if isinstance(expr, int):
+            return T.int32(expr)
+        else:
+            return expr
+
+    def verify(self, data, expected):
+        res = self.analyzer.canonical_simplify(data)
+        expected = self._convert(expected)
+        assert tvm_ffi.structural_equal(res, expected), (
+            f"\ndata={data}\nres={res}\nexpected={expected}"
+        )
+
+
+def test_mul_sum_simplify():
+    ck = CanonicalChecker()
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+
+    ck.verify(2 + (3 * x + z + y + 1) * 4 + x, x * 13 + z * 4 + y * 4 + 6)
+    ck.verify(x * 3 - 4 * x + 1, 1 - x)
+    ck.verify(y + x * 3 - 5 * x + 1 + y, y * 2 + 1 - x * 2)
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    # trucdiv
+    ck.verify(tdiv(x + y + x + y * 3, 2), y * 2 + x)
+    ck.verify(tmod(x + y + x + y * 3, 2), 0)
+
+    # floordiv
+    fld = tvm.tirx.floordiv
+    flm = tvm.tirx.floormod
+    ck.verify(flm(x + x + y * 3, 2), flm(y * 3, 2))
+    ck.verify(fld(x + y + x + y * 3, 2), y * 2 + x)
+    ck.verify(flm(x + y + x + y * 3, 2), 0)
+    ck.verify(fld(x + x + y * 3, 2), fld(y * 3, 2) + x)
+
+
+def test_split_index_simplify():
+    ck = CanonicalChecker()
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+
+    # trucdiv
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+
+    # split div const
+    ck.verify(tdiv(x, 3) * 3 + tmod(x, 3), x)
+    ck.verify(tdiv(x, 6) * 6 + tmod(tdiv(x, 3), 2) * 3 + tmod(x, 3), x)
+    ck.verify(tdiv(tdiv(tmod(x, 16), 2) * 2, 4), tdiv(tmod(x, 16), 4))
+    ck.verify(tdiv(tmod(x, 2), 8), 0)
+    ck.verify(tdiv(tmod(x, 2), 7), 0)
+    ck.verify(tdiv(tdiv(tmod(x, 16), 2) * 2, 6), tdiv(tmod(x, 16), 6))
+
+    # split mod const
+    ck.verify(tmod((x * 8), 16), tmod(x, 2) * 8)
+    ck.verify(tmod(x * 8, 2), 0)
+
+    # simplify then fold
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 1000))
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1000))
+    ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y)
+    # complex fold
+    ck.verify(tdiv(z * 9 + y, 2) * 2 + tmod(z * 9 + y, 2), z * 9 + y)
+
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(-100, 1000), True)
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(-100, 1000), True)
+    ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y)
+
+    # floordiv
+    fld = tvm.tirx.floordiv
+    flm = tvm.tirx.floormod
+    ck.verify(fld(x * 5, 2), fld(x * 5, 2))
+    ck.verify(fld(x, 3) * 3 + flm(x, 3), x)
+    ck.verify(fld(x, 6) * 6 + flm(fld(x, 3), 2) * 3 + flm(x, 3), x)
+    ck.verify(fld(fld(flm(x, 16), 2) * 2, 4), fld(flm(x, 16), 4))
+    ck.verify(fld(flm(x, 2), 8), 0)
+    ck.verify(fld(flm(x, 2), 7), 0)
+    ck.verify(fld(fld(flm(x, 16), 2) * 2, 6), fld(flm(x, 16), 6))
+
+    # floordiv(floormod(sum, m*n), n) => floormod(floordiv(sum, n), m)
+    # when sum has parts divisible by n
+    d_tile = te.var("d_tile")
+    i = te.var("i")
+    v = te.var("v")
+    ck.analyzer.update(d_tile, tvm.arith.ConstIntBound(0, 7), True)
+    ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 1), True)
+    ck.analyzer.update(v, tvm.arith.ConstIntBound(0, 7), True)
+    ck.verify(fld(flm(d_tile * 16 + i * 8 + v, 64), 8), flm(d_tile * 2 + i, 8))
+
+    # cannot simplify mixed case, unless we canonicalize into one mode.
+    ck.verify(tdiv(x, 6) * 2 + tmod(fld(x, 3), 2), tdiv(x, 6) * 2 + tmod(fld(x, 3), 2))
+
+    ck.verify(tmod(-x, 2), tmod(x, -2) * -1)
+
+
+def test_div_simplify():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    tdiv = tvm.tirx.truncdiv
+
+    # truc div
+    ck.verify(tdiv(16 + 48 * x, 16), x * 3 + 1)
+    # (17+48*x)/16 is not simplifiable for arbitrary x because when 17+48*x<0
+    # (17+48*x)/16 != 1+3*x
+    ck.verify(tdiv(17 + 48 * x, 16), tdiv(x * 48 + 17, 16))
+    # However, when x >= 0, then 17+48*x >= 0 and (17+48*x)/16 can be simplified
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 10))
+    ck.verify(tdiv(17 + 48 * x, 16), x * 3 + 1)
+    # Trying expressions that are not simplifiable for any values of the variables
+    ck.verify(tdiv(17 + 47 * x, 16), tdiv(x * 47 + 17, 16))
+
+    # floordiv
+    fld = tvm.tirx.floordiv
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(-1000, 10000), True)
+    ck.verify(fld(16 + 48 * x, 16), x * 3 + 1)
+    ck.verify(fld(17 + 48 * x, 16), x * 3 + 1)
+    ck.verify(fld(17 + 47 * x, 16), fld(x * 47 + 17, 16))
+
+
+def test_fp16_const_fold():
+    ck = CanonicalChecker()
+    zero = tvm.tirx.const(0, "float16")
+    one = tvm.tirx.const(1, "float16")
+    half = tvm.tirx.const(0.5, "float16")
+
+    ck.verify(zero + half, half)
+    ck.verify(half - zero, half)
+
+    ck.verify(zero * half, zero)
+    ck.verify(half * one, half)
+
+    ck.verify(half / one, half)
+    ck.verify(zero / half, zero)
+
+
+def test_floormod_simplify():
+    ck = CanonicalChecker()
+    flm = tvm.tirx.floormod
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(flm(flm((x * 4) + y - 466036, 24528) - 24512, 16), flm((x * 4) + y + 12, 16))
+    ck.verify(flm(flm((x * 4), 16), 8), flm(x, 2) * 4)
+
+    ck.verify(flm(-x, 2), flm(x, -2) * -1)
+
+
+def test_canonical_mixed():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    z = tvm.tirx.const(3, "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    ck.verify(tdiv(x, (z * z)) - tdiv(x, (z * z)), 0)
+    ck.verify(tdiv(x, (z + z)) - tdiv(x, (z + z)), 0)
+    ck.verify(x - 2 < 3, x < 5)
+    ck.verify(tvm.tirx.max(x, 1) - tvm.tirx.max(x, 1), 0)
+    ck.verify(tvm.tirx.min(x, 1) - tvm.tirx.min(x, 1), 0)
+    ck.verify(x * x - x * x, 0)
+    ck.verify(tmod(tdiv(tmod(x, 20), 2) * 2, 4), tdiv(tmod(x, 4), 2) * 2)
+
+    fld = tvm.tirx.floordiv
+    ck.verify(fld(x, (z * z)) - fld(x, (z * z)), 0)
+    ck.verify(fld(x, (z + z)) - fld(x, (z + z)), 0)
+
+
+def test_reduce_combiner_simplify():
+    ck = CanonicalChecker()
+    dummy = tvm.tirx.Var("dummy", "int32")
+    comm_reducer = te.comm_reducer
+    prod = comm_reducer(lambda x, y: x * y, lambda t0: tvm.tirx.const(1, t0))
+
+    sum_or_prod = comm_reducer(
+        lambda x, y: tvm.tirx.Select(dummy < 0, x + y, x * y),
+        lambda t0: tvm.tirx.Select(dummy < 0, tvm.tirx.const(0, t0), tvm.tirx.const(1, t0)),
+    )
+    sum_and_prod = comm_reducer(
+        lambda x, y: (x[0] + y[0], x[1] * y[1]),
+        lambda t0, t1: (tvm.tirx.const(0, t0), tvm.tirx.const(5, t1) - tvm.tirx.const(4, t1)),
+    )
+    some_reducer1 = comm_reducer(
+        lambda x, y: (
+            x[0] + y[0],
+            x[0] + y[0] + x[1] + y[1],
+            x[0] * y[2] + y[0] * x[2],
+            x[1] + y[2],
+            4.0,
+        ),
+        lambda t0, t1, t2, t3, t4: (
+            tvm.tirx.const(0, t0),
+            tvm.tirx.const(1, t1),
+            tvm.tirx.const(2, t2),
+            tvm.tirx.const(3, t3),
+            tvm.tirx.const(4, t4),
+        ),
+    )
+
+    k = te.reduce_axis((0, 10), name="k")
+    A = te.placeholder((10,), name="A")
+    # Test that SimplifyCombiner makes use of vranges
+    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(-10, -4))
+    ck.verify(sum_or_prod(A[k], k), te.sum(A[k], k))
+    ck.verify(sum_or_prod(A[k], k, init=1), te.sum(A[k], k, init=1))
+    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(5, 9), True)
+    ck.verify(sum_or_prod(A[k], k), prod(A[k], k))
+    ck.verify(sum_or_prod(A[k], k, init=1), prod(A[k], k, init=1))
+    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(-10, 100), True)
+    ck.verify(sum_and_prod((A[k], A[10 - k]), k)[0], te.sum(A[k], k))
+    ck.verify(sum_and_prod((A[k], A[10 - k]), k)[1], prod(A[10 - k], k))
+
+    reference_simplified_sources = [
+        [A[0]],
+        [A[0], A[1]],
+        [A[0], A[2]],
+        [A[0], A[1], A[2], A[3]],
+        [A[4]],
+    ]
+    for j in range(5):
+        # Here we use the j-th component of the result, so only it and the components it
+        # depends on are left.
+        simplified = ck.analyzer.canonical_simplify(
+            some_reducer1((A[0], A[1], A[2], A[3], A[4]), k)[j]
+        )
+
+        # Check that the remaining components are the expected ones.
+        for lhs, rhs in zip(simplified.source, reference_simplified_sources[j]):
+            tvm.ir.assert_structural_equal(lhs, rhs)
+
+    # Test that components with side effects are not removed
+    dummy = tvm.ir.GlobalVar("dummy")
+    side_effect = lambda *xs: tvm.ir.Call(dummy, xs, ret_ty="int32")
+    ck.verify(
+        sum_and_prod((A[k], side_effect(A[10 - k])), k)[0],
+        sum_and_prod((A[k], side_effect(A[10 - k])), k)[0],
+    )
+    ck.verify(sum_and_prod((side_effect(A[k]), A[10 - k]), k)[0], te.sum(side_effect(A[k]), k))
+
+
+def test_reduce_simplify():
+    ck = CanonicalChecker()
+    k = te.reduce_axis((0, 10), name="k")
+    j = te.reduce_axis((-5, 3), name="j")
+    A = te.placeholder((10,), name="A")
+    ck.verify(te.sum(tvm.tirx.Select(k + j < 12, k + j, 0), [k, j]), te.sum(k + j, [k, j]))
+    ck.verify(te.sum(A[3], []), A[3])
+    ck.verify(te.sum(A[3], [], where=k > 12, init=1.0), tvm.tirx.const(1.0, dtype="float32"))
+    # The rule below is not typical, removed for now
+    ck.verify(te.sum(tvm.tirx.div(k, 10), k), te.sum(tvm.tirx.const(0, "int32"), k))
+
+
+def test_simplify_if_then_else():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    # simplification that takes condition into account.
+    res = tvm.tirx.if_then_else(
+        (x * 4 + y) >= 466036,
+        tvm.tirx.if_then_else(
+            24512 <= tmod(((x * 4) + y) - 466036, 24528),
+            tmod(tmod(((x * 4) + y) - 466036, 24528) - 24512, 16),
+            x,
+        ),
+        y,
+    )
+
+    res2 = tvm.tirx.if_then_else(
+        (x * 4) >= 466036 - y,
+        tvm.tirx.if_then_else(
+            24512 <= tmod(((x * 4) + y) - 466036, 24528),
+            tmod(tmod(((x * 4) + y) - 466036, 24528) - 24512, 16),
+            x,
+        ),
+        y,
+    )
+    expected = tvm.tirx.if_then_else(
+        tvm.tirx.LE(466036, (x * 4 + y)),
+        tvm.tirx.if_then_else(
+            tvm.tirx.LE(24512, tmod(((x * 4) + y) - 4, 24528)), tmod(((x * 4) + y) - 4, 16), x
+        ),
+        y,
+    )
+    ck.verify(res, expected)
+    ck.verify(res2, expected)
+    # can only simplify if condition
+    res = tvm.tirx.Select(tvm.tirx.all(x >= -1, y >= 0), tmod(x + y + 100, 3), tmod(x + 100, 3))
+    expected = tvm.tirx.Select(tvm.tirx.all(x >= -1, y >= 0), tmod(x + y + 1, 3), tmod(x + 100, 3))
+    ck.verify(res, ck.analyzer.canonical_simplify(expected))
+
+    res = tvm.tirx.Select(x >= 10, tvm.tirx.if_then_else(tdiv(x, 3) > 2, x, 0), 0)
+    expected = tvm.tirx.Select(x >= 10, x, 0)
+    ck.verify(res, ck.analyzer.canonical_simplify(expected))
+
+    res = tvm.tirx.Select(x >= 10, tvm.tirx.if_then_else(tdiv(x, 3) < 2, x, 0), 0)
+    ck.verify(res, 0)
+
+
+def test_complex_cases():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    res2 = (
+        tdiv(tdiv(tmod(x * 128 + y, 1296), 36) * 2 + 1, 2) * 36
+        + tdiv(tmod((x * 128) + y, 36) * 2 + 1, 2)
+        - tmod((x * 128) + y, 1296)
+        + 1
+    )
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 5))
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 127))
+    ck.verify(res2, 1)
+
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1024), True)
+    res3 = (
+        tdiv(x * 1024 + y, 65536)
+        + tdiv(tmod(x * 1024 + y, 65536), 256)
+        + tdiv(tmod(x * 1024 + y, 256), 16)
+        + tmod(x * 1024 + y, 16)
+        - tdiv(y, 256)
+        - tdiv(tmod(y, 256), 16)
+        - tmod(y, 16)
+        - (x * 4)
+    )
+    ck.verify(res3, tdiv((x * 1024) + y, 256) - tdiv(y, 256) - (x * 4))
+
+
+def test_simplify_cast():
+    ck = CanonicalChecker()
+    tcast = tvm.tirx.Cast
+    fld = tvm.tirx.floordiv
+    flm = tvm.tirx.floormod
+    # cast(i64, i + j + 1) - cast(i64, i)
+    i = tvm.tirx.Var("i", "int32")
+    j = tvm.tirx.Var("j", "int32")
+    res = tcast("int64", i + j + 1) - tcast("int64", i)
+    ck.verify(res, tcast("int64", j) + tvm.tirx.const(1, "int64"))
+    # cast(i32, i + j + 1) - cast(i32, i)
+    i = tvm.tirx.Var("i", "int64")
+    j = tvm.tirx.Var("j", "int64")
+    ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 10))
+    ck.analyzer.update(j, tvm.arith.ConstIntBound(0, 10))
+    res = tcast("int32", i + j + 1) - tcast("int32", i)
+    ck.verify(res, tcast("int32", j) + 1)
+    # cast(i32, i + j - 100)
+    i = tvm.tirx.Var("i", "int64")
+    j = tvm.tirx.Var("j", "int64")
+    ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 2**31 - 1))
+    ck.analyzer.update(j, tvm.arith.ConstIntBound(0, 10))
+    res = tcast("int32", i + j - 100)
+    ck.verify(res, res)
+    # cast(i32, flm(axis, 7i64) * 2i64 + 1i64) + 1i32
+    # - cast(i32, flm(axis, 7i64) * 2i64)
+    axis = tvm.tirx.Var("axis", "int64")
+    ck.analyzer.update(axis, tvm.arith.ConstIntBound(0, 42))
+    res = (
+        tcast(
+            "int32",
+            flm(axis, tvm.tirx.const(7, "int64")) * tvm.tirx.const(2, "int64")
+            + tvm.tirx.const(1, "int64"),
+        )
+        + tvm.tirx.const(1, "int32")
+        - tcast("int32", flm(axis, tvm.tirx.const(7, "int64")) * tvm.tirx.const(2, "int64"))
+    )
+    ck.verify(res, 2)
+
+
+def test_simplify_normalize_min_value_expr():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+
+    ck.verify(tvm.tirx.min_value("int32") - x == 0, x == tvm.tirx.min_value("int32"))
+    ck.verify(tvm.tirx.min_value("int32") + x == 0, tirx.const(False))
+    ck.verify(0 == tvm.tirx.min_value("int32") - x, x == tvm.tirx.min_value("int32"))
+    ck.verify(0 == tvm.tirx.min_value("int32") + x, tirx.const(False))
+    ck.verify(-x + tvm.tirx.min_value("int32") == 0, x == tvm.tirx.min_value("int32"))
+    ck.verify(x + tvm.tirx.min_value("int32") == 0, tirx.const(False))
+    ck.verify(0 == -x + tvm.tirx.min_value("int32"), x == tvm.tirx.min_value("int32"))
+    ck.verify(0 == x + tvm.tirx.min_value("int32"), tirx.const(False))
+
+
+def test_proddiv_simplify():
+    ck = CanonicalChecker()
+    flm = tvm.tirx.floormod
+    fld = tvm.tirx.floordiv
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("y", "int32")
+
+    ck.verify(flm(x * 32 * x, x), 0)
+    ck.verify(flm(z * x * 32 * x * y, x * z), 0)
+    ck.verify(flm(z * x * 32 * x * y, x * z * y * 8 * x), 0)
+    ck.verify(flm(z * x * 32 * (x * y), 6 * x * z), flm(x * y * 16, 3) * (x * z * 2))
+    ck.verify(flm(x * 32 * x, x * z), flm(x * 32, z) * x)
+
+    ck.verify(tmod(x * 32 * x, x), 0)
+    ck.verify(tmod(z * x * 32 * x * y, x * z), 0)
+    ck.verify(tmod(z * x * 32 * (x * y), 6 * x * z), tmod(x * y * 16, 3) * (x * z * 2))
+    ck.verify(tmod(x * 32 * x, x * z), tmod(x * 32, z) * x)
+
+    ck.verify(fld(x * 2 * x * z, 4 * x * x * x), fld(z, x * 2))
+    ck.verify(fld(x * (2 * y) * 3, 3 * y), x * 2)
+    ck.verify(fld(x * (2 * y) * 3, 3 * y * z), fld(x * 2, z))
+
+    ck.verify(tdiv(x * 2 * x * z, 4 * x * x * x), tdiv(z, x * 2))
+    ck.verify(tdiv(x * (2 * y) * 3, 3 * y), x * 2)
+    ck.verify(tdiv(x * (2 * y) * 3, 3 * y * z), tdiv(x * 2, z))
+
+
+def test_floormod_two():
+    ck = CanonicalChecker()
+    flm = tvm.tirx.floormod
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(flm(x * 10 + 1 + y * 2 + 2, 2), 1)
+
+
+def test_simplify_le():
+    ck = CanonicalChecker()
+    # Case 1. Ignore the extra expr if it's small than the division number
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    ck.analyzer.bind(y, tvm.ir.Range(0, 8))
+    ck.analyzer.bind(z, tvm.ir.Range(0, 2))
+    ck.verify(x * 8 + y < 16, x < 2)
+    ck.verify(x * 8 + z * 4 < 16, x < 2)
+    ck.verify(x * 8 + z * 4 < 16, x < 2)
+
+    # TODO: Not sure why `-2 < x` will be convert to `x > -2`, use a explicit simplify here.
+    ck.verify(x * -8 + y < 16, ck.analyzer.rewrite_simplify(-2 < x))
+    ck.verify(x * -8 + z * 4 < 16, ck.analyzer.rewrite_simplify(-2 < x))
+
+    ck.verify(x * 8 + y + z < 16, x * 8 + y + z < 16)
+
+    n = tvm.tirx.Var("n", "int32")
+    ck.verify(x * 8 + y < n, x * 8 + y < n)
+
+    # Case 2. Simplify the extra expr
+    x1, x2, ty, tx, vec = (
+        tvm.tirx.Var("x1", "int32"),
+        tvm.tirx.Var("x2", "int32"),
+        tvm.tirx.Var("ty", "int32"),
+        tvm.tirx.Var("tx", "int32"),
+        tvm.tirx.Var("vec", "int32"),
+    )
+    ck.analyzer.bind(x1, tvm.ir.Range(0, 2))
+    ck.analyzer.bind(x2, tvm.ir.Range(0, 3))
+    ck.analyzer.bind(ty, tvm.ir.Range(0, 8))
+    ck.analyzer.bind(tx, tvm.ir.Range(0, 32))
+    ck.analyzer.bind(vec, tvm.ir.Range(0, 8))
+    ck.verify(
+        x1 * 5632 + (((x2 * 8 + ty) * 32 + tx) * 8 + vec) % 5632 < 11008,
+        x1 * 22 + (x2 * 8 + ty) % 22 < 43,
+    )
+    ck.verify(tx // 2 % 8 + vec < 8, tx % 16 // 2 + vec < 8)
+
+    # Case 3. No failure
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    ck.analyzer.bind(y, tvm.ir.Range(0, 1024))
+    ck.verify(x * 1024 + y < z * 7168, x - z * 7 < 0)
+
+
+def test_simplify_le_negative_scale_extra():
+    """Regression: Case 2 of the LT-with-divisible-coeffs rewrite must not
+    fire when the leftover split term has a negative scale.
+
+    The rewrite ``S + xn < 0  ⇔  S/d + xn // d < 0`` is only sound when
+    the leftover ``xn`` has scale ``+1``. With scale ``-1`` the equivalence
+    becomes ``≤`` rather than ``<`` and the rewrite silently strengthens
+    the predicate. The original bug surfaced as ``row > col`` masks of
+    ``.16x*b`` tcgen05 readbacks collapsing to plain ``warp_id > k``
+    comparisons (lower-triangle writes were silently dropped on the
+    boundary warp).
+    """
+    ck = CanonicalChecker()
+    tx = tvm.tirx.Var("tx", "int32")
+    warp = tvm.tirx.Var("warp", "int32")
+    ck.analyzer.bind(tx, tvm.ir.Range(0, 128))
+    ck.analyzer.bind(warp, tvm.ir.Range(0, 4))
+
+    # Same-source joint projection: the comparison genuinely depends on tx
+    # at warp == 0 (e.g. tx == 4 ⇒ 0 < 1 = True; tx == 1 ⇒ 2 < 0 = False),
+    # so the simplifier must keep both sides.  Pre-fix this folded to
+    # ``0 < warp`` and dropped every True case in warp 0.
+    expr = (tx % 4) * 2 < warp * 16 + (tx % 32) // 4
+    ck.verify(expr, expr)
+
+    # The simpler ``scale = -1`` with ``lower_factor = 1`` shape.  Pre-fix
+    # this folded to ``False`` (drops all warp >= 1 cases where the rhs
+    # actually exceeds 8*warp).
+    expr = warp * 8 < (tx % 32)
+    ck.verify(expr, expr)
+
+    # The corresponding ``scale = +1`` Case 2 path (the rewrite this guards)
+    # must still optimize — verifies we did not over-restrict.
+    x1 = tvm.tirx.Var("x1", "int32")
+    y1 = tvm.tirx.Var("y1", "int32")
+    ck.verify(x1 * 64 + (y1 % 64) < 120, x1 * 8 + (y1 % 64) // 8 < 15)
+
+    # The truly-always-true comparison that arises from the same kernel
+    # (``r = 2 / va = 1`` in the tcgen05.ld.16x256b readback) must still
+    # fold to True so the masked store can be elided.
+    expr_true = (tx % 4) * 2 < warp * 16 + (tx % 32) // 4 + 8
+    ck.verify(expr_true, tvm.tirx.const(True, "bool"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_const_int_bound.py

@@ -0,0 +1,323 @@
+# 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.
+# ruff: noqa: RUF012
+
+import contextlib
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.arith import ConstIntBound
+
+NEG_INF = ConstIntBound.NEG_INF
+POS_INF = ConstIntBound.POS_INF
+
+
+class TestCase:
+    def __init__(self, expr, expected_bounds, known_bounds=None, constraint=None):
+        self.expr = expr
+        self.expected_bounds = expected_bounds
+        if known_bounds is None:
+            self.known_bounds = {}
+        else:
+            self.known_bounds = known_bounds
+
+        self.constraint = constraint
+
+    @property
+    def __name__(self):
+        return str(self.expr)
+
+
+class BaseCompare:
+    def test_const_bounds(self, test_case):
+        analyzer = tvm.arith.Analyzer()
+
+        for var, bounds in test_case.known_bounds.items():
+            analyzer.update(var, ConstIntBound(*bounds))
+            assert analyzer.const_int_bound_is_bound(var)
+
+        with contextlib.ExitStack() as stack:
+            if test_case.constraint is not None:
+                stack.enter_context(analyzer.constraint_scope(test_case.constraint))
+
+            bounds = analyzer.const_int_bound(test_case.expr)
+
+        if test_case.expected_bounds[0] is None:
+            assert bounds.max_value == test_case.expected_bounds[1]
+        elif test_case.expected_bounds[1] is None:
+            assert bounds.min_value == test_case.expected_bounds[0]
+        else:
+            assert (bounds.min_value, bounds.max_value) == test_case.expected_bounds
+
+
+class TestDataType(BaseCompare):
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Var("x", "int64"), (NEG_INF, POS_INF)),
+        TestCase(tvm.tirx.Var("x", "int8"), (-128, 127)),
+        TestCase(tvm.tirx.Var("x", "uint8"), (0, 255)),
+        TestCase(tvm.tirx.Var("x", "int32"), (-(2**31), 2**31 - 1)),
+    )
+
+
+def test_plain_var_non_negative_bound_requires_context():
+    var = tvm.tirx.Var("x", "int64")
+    analyzer = tvm.arith.Analyzer()
+
+    assert analyzer.const_int_bound(var).min_value == NEG_INF
+    with analyzer.constraint_scope(var >= 0):
+        assert analyzer.const_int_bound(var).min_value == 0
+    assert analyzer.const_int_bound(var).min_value == NEG_INF
+
+
+class TestCastBound(BaseCompare):
+    x = tvm.tirx.Var("x", "int8")
+    tmod = tvm.tirx.truncmod
+
+    test_case = tvm.testing.parameter(
+        TestCase(tmod(x, 3).astype("uint32"), (0, 2)),
+        TestCase(tmod(x, 3).astype("float32").astype("int32"), (-2, 2)),
+    )
+
+
+class TestAddSubBound(BaseCompare):
+    x = tvm.tirx.Var("x", "int64")
+    y = tvm.tirx.Var("y", "int64")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x + y, (NEG_INF, POS_INF)),
+        TestCase(x + y, (1, 14), known_bounds={x: (0, 4), y: (1, 10)}),
+        TestCase(x - y, (-10, 3), known_bounds={x: (0, 4), y: (1, 10)}),
+        TestCase(x - y, (-10, POS_INF), known_bounds={x: (0, POS_INF), y: (1, 10)}),
+        TestCase(1 - x, (NEG_INF, 1), known_bounds={x: (0, POS_INF), y: (1, 10)}),
+    )
+
+
+@pytest.mark.xfail(reason="Not currently supported")
+class TestBoundsUsingReciprocals(BaseCompare):
+    """Special handling for differences of reciprocals
+
+    These terms can appear when comparing the number of operations for
+    different orderings of matrix multiplications, with A, B, and C
+    known to be positive values.
+
+    In these cases, comparing `(A+B)*C < A*B` is equivalent to
+    `1/A + 1/B < 1/C`.  Working in terms of the reciprocals
+    allows the ConstIntBound analyzer to provide a tighter
+    bound for these differences than would otherwise be
+    available.
+
+    For `(A+B)*C - A*B`, the normal bottom-up integer bounds are unable to
+    provide the bounds required to provide these inequalities, because they
+    treat the terms as uncorrelated.  That is, they assume that `(A+B)*C` may
+    achieve its minimum while `A*B` simultaneously achieves its maximum.
+    """
+
+    A, B, C = [tvm.tirx.Var(letter, "int64") for letter in "ABC"]
+
+    symmetric_bounds = {A: (1, 4095), B: (1, 4095), C: (2048, 2048)}
+    asymmetric_bounds = {A: (1, 1024), B: (1, POS_INF), C: (2048, 2048)}
+
+    test_case = tvm.testing.parameter(
+        TestCase((A + B) * C - A * B, (2048, None), known_bounds=symmetric_bounds),
+        TestCase((A + B) * C - B * A, (2048, None), known_bounds=symmetric_bounds),
+        TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),
+        TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),
+        TestCase((A + B) * C - A * B, (2048, None), known_bounds=asymmetric_bounds),
+        TestCase((A + B) * C - B * A, (2048, None), known_bounds=asymmetric_bounds),
+        TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),
+        TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),
+    )
+
+
+class TestMulBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x * y + 20, (0, 60), {x: (-2, 4), y: (4, 10)}),
+        TestCase(x * y, (-32, 24), {x: (-3, 4), y: (-8, 2)}),
+        TestCase(x * y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-8, 2)}),
+    )
+
+
+class TestTruncDivBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    expr = tvm.tirx.truncdiv(x, y)
+
+    test_case = tvm.testing.parameter(
+        TestCase(expr, (-2, None), {x: (-9, 4), y: (4, 10)}),
+        TestCase(expr, (-4, 9), {x: (-9, 4), y: (-2, 0)}),
+        TestCase(expr, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),
+        TestCase(expr, (-9, 9), {x: (-9, 4), y: (-4, 12)}),
+    )
+
+
+class TestTruncModBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    expr = tvm.tirx.truncmod(x, y)
+
+    test_case = tvm.testing.parameter(
+        TestCase(expr, (-9, 4), {x: (-9, 4), y: (4, 10)}),
+        TestCase(expr, (-9, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(expr, (0, 9), {x: (1, POS_INF), y: (4, 10)}),
+    )
+
+
+class TestFloorDivBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ux = tvm.tirx.Var("x", "uint32")
+    uy = tvm.tirx.Var("y", "uint32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x // y, (-9 // 4, None), {x: (-9, 4), y: (4, 10)}),
+        TestCase(x // y, (-4, 9), {x: (-9, 4), y: (-2, 0)}),
+        TestCase(x // y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),
+        TestCase(x // y, (-9, 9), {x: (-9, 4), y: (-4, 12)}),
+        TestCase(ux // uy, (0, 4), {ux: (1, 4), uy: (0, 12)}),
+    )
+
+
+class TestFloorModBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x % y, (0, 9), {x: (-9, 4), y: (4, 10)}),
+        TestCase(x % y, (0, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(x % y, (0, 9), {x: (1, POS_INF), y: (4, 10)}),
+    )
+
+
+class TestMinMaxBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.min(x, y), (-9, 10), {x: (-9, 11), y: (4, 10)}),
+        TestCase(tvm.tirx.min(x, y), (NEG_INF, 10), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(tvm.tirx.max(x, y), (4, POS_INF), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(tvm.tirx.max(x, y), (4, POS_INF), {x: (1, POS_INF), y: (4, 10)}),
+    )
+
+
+class TestSelectBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(
+            tvm.tirx.Select(x > 1, (y < 0).astype("int32"), y + 1),
+            (0, 11),
+            {x: (-9, 11), y: (4, 10)},
+        ),
+    )
+
+
+class TestShiftAndBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x >> y, (-3, 2), {x: (-9, 11), y: (2, 10)}),
+        TestCase(x & y, (0, 10), {x: (-9, 11), y: (2, 10)}),
+        TestCase(x & y, (0, 10), {x: (10, 11), y: (2, 10)}),
+    )
+
+
+class TestMixIndexBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+
+    test_case = tvm.testing.parameter(
+        TestCase(tmod(x, 8) + tdiv(x, 8) * 8, (0, 24 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),
+        TestCase(y + x * 3, (0, 24 * 3 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),
+        TestCase(
+            tmod(x, 7) + tdiv(x, 7) * 7, (0, (23 // 7) * 7 + 6), {x: (0, 24 - 1), y: (0, 3 - 1)}
+        ),
+    )
+
+
+class TestLetBound(BaseCompare):
+    x = tvm.tirx.Var("x", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Let(x, 1, x + 1), (2, 2)),
+    )
+
+
+class TestFloorModNegativeDivisor(BaseCompare):
+    flm, fld = tvm.tirx.floormod, tvm.tirx.floordiv
+    a, b = tvm.tirx.Var("a", "int32"), tvm.tirx.Var("b", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(a % b, (-4, 6), {a: (0, 6), b: (-5, 7)}),
+    )
+
+
+class TestDivModAssumeNoZeroDivisor(BaseCompare):
+    """Divmod non negative expression makes assumption that divide by
+    zero won't occur this assumption is important to get best result
+    from symbolic shape programs
+    """
+
+    a, b = tvm.tirx.Var("a", "int32"), tvm.tirx.Var("b", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(a // b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),
+        TestCase(a % b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),
+    )
+
+
+class TestMultipleCondition(BaseCompare):
+    a = tvm.tirx.Var("a", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(
+            a % 58 - 1,
+            (0, None),
+            known_bounds={a: (0, 128)},
+            constraint=tvm.tirx.all(1 <= a % 58, a % 58 < 57),
+        ),
+    )
+
+
+class TestBroadcastBound(BaseCompare):
+    a = tvm.tirx.Var("a", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Broadcast(a, 4), (0, 128), {a: (0, 128)}),
+    )
+
+
+class TestRampBound(BaseCompare):
+    a = tvm.tirx.Var("a", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Ramp(a, 2, 4) + 2, (2, 128 + 2 * 3 + 2), {a: (0, 128)}),
+    )
+
+
+class TestModularSetBound(BaseCompare):
+    analyzer = tvm.arith.Analyzer()
+    tx = tvm.tirx.Var("tx", "int32")
+    bx = tvm.tirx.Var("bx", "int32")
+
+    expr = (bx * 2048 + tx * 16) % 7168
+
+    test_case = tvm.testing.parameter(
+        TestCase(expr, (0, 7152), {bx: (0, 3584), tx: (0, 128)}),
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_deduce_bound.py

@@ -0,0 +1,272 @@
+# 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.
+# ruff: noqa: F401, F841
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.tirx.buffer import decl_buffer
+
+
+def test_deduce():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = tvm.tirx.Var("c", "int32")
+    d = tvm.tirx.Var("d", "int32")
+
+    b_s = tvm.arith.IntervalSet(2, 3)
+    c_s = tvm.arith.IntervalSet(10, 15)
+    d_s = tvm.arith.IntervalSet(-3, -1)
+    zero = tvm.tirx.const(0, "int32")
+
+    fdiv = tvm.tirx.floordiv
+
+    e0 = (-b) * a + c - d
+    res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
+    ans0 = fdiv(d - c, b * -1)
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    # expression containing variable a is on rhs
+    res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    e0 = d * a + c - d
+    res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
+    ans0 = fdiv(d - c, d)
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    # expression containing variable a is on rhs
+    res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    e1 = a * 4 + b < c
+    res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
+    ans1 = fdiv(c - 1 - b, 4)
+    tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
+
+    # expression containing variable a is on rhs
+    e1 = c > a * 4 + b
+    res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
+
+    e2 = tvm.tirx.max(5, a * 4) < 0
+    res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {})
+    assert str(res2.max_value) == "neg_inf"
+    assert str(res2.min_value) == "pos_inf"
+
+    # expression containing variable a is on rhs
+    e2 = zero < tvm.tirx.max(5, a * 4)
+    res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {})
+    assert str(res2.max_value) == "neg_inf"
+    assert str(res2.min_value) == "pos_inf"
+
+    e3 = (-b) + a * c - d
+    res3 = tvm.arith.deduce_bound(a, e3 >= 0, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
+    ans3 = fdiv(2, c) + 1
+    tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
+
+    res3 = tvm.arith.deduce_bound(a, zero <= e3, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
+    tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
+
+    # tests for `EQ` op
+    res4 = tvm.arith.deduce_bound(a, a == b, {}, {})
+    tvm.testing.assert_prim_expr_equal(res4.max_value, b)
+    tvm.testing.assert_prim_expr_equal(res4.min_value, b)
+
+    # Unsatisfiable `EQ`, variable as one of the Operand
+    res5 = tvm.arith.deduce_bound(a, (a == b), {b: b_s}, {b: b_s})
+    assert str(res5.max_value) == "neg_inf"
+    assert str(res5.min_value) == "pos_inf"
+
+    # variable `a` on the RHS side
+    res6 = tvm.arith.deduce_bound(a, 10 == a, {}, {})
+    tvm.testing.assert_prim_expr_equal(res6.max_value, 10)
+    tvm.testing.assert_prim_expr_equal(res6.min_value, 10)
+
+    # Add, Sub in `EQ`
+    e4 = (a - c) == (b + d)
+    ans4 = b + d + c
+    res7 = tvm.arith.deduce_bound(a, e4, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res7.max_value, ans4)
+    tvm.testing.assert_prim_expr_equal(res7.min_value, ans4)
+
+    # Satisfiable Mul in `EQ` with negative sign
+    res8 = tvm.arith.deduce_bound(a, (5 * a == -10), {}, {})
+    tvm.testing.assert_prim_expr_equal(res8.max_value, -2)
+    tvm.testing.assert_prim_expr_equal(res8.min_value, -2)
+
+    # Unsatisfiable Mul in `EQ`
+    e5 = 4 * a == b
+    res9 = tvm.arith.deduce_bound(a, e5, {b: b_s}, {})
+    assert str(res9.max_value) == "neg_inf"
+    assert str(res9.min_value) == "pos_inf"
+
+    res10 = tvm.arith.deduce_bound(a, (b * a == b), {b: b_s}, {})
+    # simplifier is now able to prove symbolic relation (b * a % b == 0)
+    tvm.testing.assert_prim_expr_equal(res10.max_value, 1)
+    tvm.testing.assert_prim_expr_equal(res10.min_value, 1)
+
+
+def test_check():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = tvm.tirx.Var("c", "int32")
+    d = tvm.tirx.Var("d", "int32")
+
+    b_s = tvm.arith.IntervalSet(2, 3)
+    c_s = tvm.arith.IntervalSet(5, 7)
+    d_s = tvm.arith.IntervalSet(-3, -1)
+
+    # no compare operator
+    res1 = tvm.arith.deduce_bound(a, a + b, {b: b_s}, {})
+    assert res1.is_nothing()
+
+    # multiple compare operators
+    res2 = tvm.arith.deduce_bound(a, (a + b > 3).astype(c.ty) > c, {b: b_s, c: c_s}, {})
+    assert res2.is_nothing()
+
+    # multiple target variable
+    res2 = tvm.arith.deduce_bound(a, a * 2 - a > b, {b: b_s}, {})
+    assert res2.is_nothing()
+
+
+def test_deduce_basic():
+    def test_basic(a1, a2, coff):
+        a = tvm.tirx.Var("a", "int32")
+        b = tvm.tirx.Var("b", "int32")
+        b_s = tvm.arith.IntervalSet(a1, a2)
+        e0 = b + a * coff + 3
+
+        res1 = tvm.arith.deduce_bound(a, e0 < 17, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) < 17, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") < e0, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) > 17, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") >= e0, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) <= 17, True)
+
+        res1 = tvm.arith.deduce_bound(a, e0 >= 17, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) >= 17, True)
+
+    test_basic(0, 4, 4)
+    test_basic(1, 5, 4)
+    test_basic(2, 6, 4)
+    test_basic(0, 4, -4)
+    test_basic(1, 5, -4)
+    test_basic(2, 6, -4)
+
+
+def test_deduce_complex():
+    def test_complex(a1, a2, coff):
+        a = tvm.tirx.Var("a", "int32")
+        b = tvm.tirx.Var("b", "int32")
+        b_s = tvm.arith.IntervalSet(a1, a2)
+        e0 = (b * 3 + a * coff) * 4
+
+        res1 = tvm.arith.deduce_bound(a, e0 < 63, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) < 63, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") >= e0, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) <= 63, True)
+
+        res1 = tvm.arith.deduce_bound(a, e0 > 63, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) > 63, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") <= e0, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) >= 63, True)
+
+    test_complex(0, 4, 4)
+    test_complex(0, 4, -4)
+    test_complex(2, 6, 4)
+    test_complex(0, 4, -4)
+    test_complex(1, 5, -4)
+    test_complex(2, 6, -4)
+
+
+def test_deduce_non_support():
+    a = tvm.tirx.Var("a", "int32")
+
+    def test_non_support(lhs):
+        res = tvm.arith.deduce_bound(a, lhs < 10, {}, {})
+        assert res.is_nothing()
+
+    test_non_support(tvm.tirx.floormod(a, 16))
+    test_non_support(tvm.tirx.Min(a, 16))
+    test_non_support(tvm.tirx.Max(a, 16))
+    test_non_support(tvm.tirx.LE(a, 16))
+    test_non_support(tvm.tirx.LT(a, 16))
+    test_non_support(tvm.tirx.GE(a, 16))
+    test_non_support(tvm.tirx.GT(a, 16))
+    test_non_support(tvm.tirx.EQ(a, 16))
+    test_non_support(tvm.tirx.NE(a, 16))
+    test_non_support(tvm.tirx.log(a))
+    test_non_support(tvm.tirx.BufferLoad(decl_buffer([16], "int32"), [a]))
+
+
+def test_deduce_floordiv():
+    def do_test(gen_expr, dom_map, expect_min, expect_max):
+        a = tvm.tirx.Var("a", "int32")
+        expr = gen_expr(a)
+        res = tvm.arith.deduce_bound(a, expr, dom_map, dom_map)
+        if isinstance(expect_min, str):
+            assert str(res.min_value) == expect_min
+        else:
+            tvm.testing.assert_prim_expr_equal(res.min_value, expect_min)
+        if isinstance(expect_max, str):
+            assert str(res.max_value) == expect_max
+        else:
+            tvm.testing.assert_prim_expr_equal(res.max_value, expect_max)
+
+    # test basic cases
+    do_test(lambda a: a // 8 > 3, {}, 32, "pos_inf")
+    do_test(lambda a: a // 8 >= 3, {}, 24, "pos_inf")
+    do_test(lambda a: a // 8 < 3, {}, "neg_inf", 23)
+    do_test(lambda a: a // 8 <= 3, {}, "neg_inf", 31)
+    do_test(lambda a: a // 8 == 3, {}, "pos_inf", "neg_inf")
+    do_test(lambda a: a // 8 > -3, {}, -16, "pos_inf")
+    do_test(lambda a: a // 8 >= -3, {}, -24, "pos_inf")
+    do_test(lambda a: a // -8 > 3, {}, "neg_inf", -32)
+    do_test(lambda a: a // -8 >= 3, {}, "neg_inf", -24)
+    do_test(lambda a: a // -8 < 3, {}, -23, "pos_inf")
+    do_test(lambda a: a // -8 <= 3, {}, -31, "pos_inf")
+    do_test(lambda a: 8 // a >= 2, {}, "pos_inf", "neg_inf")
+
+    # test nested cases
+    b = tvm.tirx.Var("b", "int32")
+    bs = {b: tvm.arith.IntervalSet(2, 6)}
+    do_test(lambda a: b * 3 + a // 8 < 63, bs, "neg_inf", 359)
+    do_test(lambda a: b * 3 + a // 8 <= 63, bs, "neg_inf", 367)
+    do_test(lambda a: b * 3 + a // 8 > 63, bs, 464, "pos_inf")
+    do_test(lambda a: b * 3 + a // 8 >= 63, bs, 456, "pos_inf")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_detect_clip_bound.py

@@ -0,0 +1,55 @@
+# 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
+
+
+def test_basic():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = tvm.tirx.Var("c", "int32")
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a])
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6 - 1)
+    assert m[0].value == 2
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a, b])
+    assert len(m) == 0
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a + 10 * c <= 20, b - 1 > 0), [a, b])
+    tvm.testing.assert_prim_expr_equal(m[1], 20 - 10 * c)
+    tvm.testing.assert_prim_expr_equal(m[2], 2)
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(tvm.tirx.Not(a * 1 > b * 6), a - 1 > 0), [a])
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6)
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(tvm.tirx.Min(a, b) > 3, a - 10 < 0), [a, b])
+    tvm.testing.assert_prim_expr_equal(m[0], 4)
+    tvm.testing.assert_prim_expr_equal(m[1], 9)
+    tvm.testing.assert_prim_expr_equal(m[2], 4)
+
+
+def test_trivial_eq():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    m = tvm.arith.detect_clip_bound(b == 3, [a, b])
+    tvm.testing.assert_prim_expr_equal(m[2], 3)
+    tvm.testing.assert_prim_expr_equal(m[3], 3)
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a == 4, b == 3), [a, b])
+    tvm.testing.assert_prim_expr_equal(m[0], 4)
+    tvm.testing.assert_prim_expr_equal(m[1], 4)
+    tvm.testing.assert_prim_expr_equal(m[2], 3)
+    tvm.testing.assert_prim_expr_equal(m[3], 3)
+
+
+if __name__ == "__main__":
+    test_basic()

tests/python/arith/test_arith_detect_linear_equation.py

@@ -0,0 +1,82 @@
+# 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
+
+
+def test_basic():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    m = tvm.arith.detect_linear_equation(a * 4 + b * 6 + 7, [a])
+    assert m[0].value == 4
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7)
+
+    m = tvm.arith.detect_linear_equation(a * 4 * (a + 1) + b * 6 + 7, [a])
+    assert len(m) == 0
+
+    m = tvm.arith.detect_linear_equation(a * 4 + (a + 1) + b * 6 + 7, [a])
+    assert m[0].value == 5
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7 + 1)
+
+    m = tvm.arith.detect_linear_equation(a * b + 7, [a])
+    assert m[0] == b
+
+    m = tvm.arith.detect_linear_equation(b * 7, [a])
+    assert m[0].value == 0
+
+    m = tvm.arith.detect_linear_equation(b * 7, [])
+    assert len(m) == 1
+    tvm.testing.assert_prim_expr_equal(m[0], b * 7)
+
+    c = tvm.tirx.Var("c", "uint32")
+    m = tvm.arith.detect_linear_equation(128 - c, [c])
+    assert m[0].value == -1
+
+
+def test_multivariate():
+    v = [tvm.tirx.Var(f"v{i}", "int32") for i in range(4)]
+    b = tvm.tirx.Var("b", "int32")
+    m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8, v)
+
+    tvm.testing.assert_prim_expr_equal(m[0], b + 5)
+
+    assert m[1].value == 8
+
+    m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[2], v)
+    assert len(m) == 0
+
+    m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[1] + v[3], v)
+    assert len(m) == 0
+
+    m = tvm.arith.detect_linear_equation(((v[0] * b + v[1]) * 8 + v[2] + 1) * 2, v)
+    assert m[1].value == 16
+    assert m[2].value == 2
+    assert m[len(m) - 1].value == 2
+
+    m = tvm.arith.detect_linear_equation((v[0] - v[1]), [v[2]])
+    assert m[0].value == 0
+
+    tvm.testing.assert_prim_expr_equal(m[1], v[0] - v[1])
+
+    m = tvm.arith.detect_linear_equation((v[0] - v[1]), [])
+    assert len(m) == 1
+    tvm.testing.assert_prim_expr_equal(m[0], v[0] - v[1])
+
+
+if __name__ == "__main__":
+    test_basic()
+    test_multivariate()

tests/python/arith/test_arith_domain_touched.py

@@ -0,0 +1,95 @@
+# 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 pytest
+import tvm_ffi
+
+import tvm
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def scalar_func(a: T.handle, b: T.handle):
+    m = T.int32()
+    n = T.meta_var(100)
+    A = T.match_buffer(a, (n, m))
+    B = T.match_buffer(b, (n, m))
+
+    for i, j in T.grid(n, m):
+        A[i, j] = B[i - 1, j + 1] + A[i - 1, j - 1]
+
+
+def test_domain_touched():
+    func = scalar_func
+    a, b = [func.buffer_map[var] for var in func.params]
+    ir = func.body
+
+    a_domain_r = tvm.arith._ffi_api.DomainTouched(ir, a, True, False)
+
+    assert a_domain_r[0].min.value == -1
+    assert a_domain_r[0].extent.value == 100
+    assert a_domain_r[1].min.value == -1
+    assert a_domain_r[1].extent.name == "m"
+
+    a_domain_w = tvm.arith._ffi_api.DomainTouched(ir, a, False, True)
+    assert a_domain_w[0].min.value == 0
+    assert a_domain_w[0].extent.value == 100
+    assert a_domain_w[1].min.value == 0
+    assert a_domain_w[1].extent.name == "m"
+
+    a_domain_rw = tvm.arith._ffi_api.DomainTouched(ir, a, True, True)
+    assert a_domain_rw[0].min.value == -1
+    assert a_domain_rw[0].extent.value == 101
+    assert a_domain_rw[1].min.value == -1
+    assert isinstance(a_domain_rw[1].extent, tvm.tirx.Add)
+    assert a_domain_rw[1].extent.a.name == "m"
+    assert a_domain_rw[1].extent.b.value == 1
+
+    b_domain_r = tvm.arith._ffi_api.DomainTouched(ir, b, True, False)
+    assert b_domain_r
+    assert b_domain_r[0].min.value == -1
+    assert b_domain_r[0].extent.value == 100
+    assert b_domain_r[1].min.value == 1
+    assert b_domain_r[1].extent.name == "m"
+
+    b_domain_w = tvm.arith._ffi_api.DomainTouched(ir, b, False, True)
+    assert isinstance(b_domain_w, tvm_ffi.Array)
+    assert len(b_domain_w) == 0
+
+
+def test_domain_touched_vector():
+    pytest.skip("BufferRegion arithmetic in expressions not supported")
+    m = tvm.runtime.convert(128)
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle, n: T.int32):
+        A = T.match_buffer(a, (n * m,))
+        B = T.match_buffer(b, (n * m,))
+
+        for i in T.serial(n):
+            A[i * m : (i + 1) * m : 1] = A[i * m : (i + 1) * m : 1] + B[i * m : (i + 1) * m : 1]
+
+    a, b = [func.buffer_map[var] for var in func.params[:2]]
+
+    assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, False)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, False)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, True)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, b, True, False)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, b, True, False)[0].extent.value == 128
+
+
+if __name__ == "__main__":
+    test_domain_touched()

tests/python/arith/test_arith_intset.py

@@ -0,0 +1,456 @@
+# 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.
+# ruff: noqa: F841
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.arith.analyzer import Analyzer
+
+
+class IntSetChecker:
+    def __init__(self):
+        self.analyzer = tvm.arith.Analyzer()
+
+    def verify(self, data, dmap, expected):
+        res = self.analyzer.int_set(data, dmap)
+
+        def err_msg():
+            return f"\ndata={data}\ndmap={dmap}\nres={res}\nexpected={expected}"
+
+        assert self.analyzer.can_prove_equal(res.min_value, expected[0]), err_msg()
+        assert self.analyzer.can_prove_equal(res.max_value, expected[1]), err_msg()
+
+
+def test_basic():
+    s = tvm.arith.IntervalSet(2, 3)
+    assert s.min_value.value == 2
+    assert s.max_value.value == 3
+
+    s = tvm.arith.IntSet.single_point(2)
+    assert s.min_value.value == 2
+    assert s.max_value.value == 2
+
+
+def test_vector():
+    base = 10
+    stride = 3
+    lanes = 2
+    s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, stride, lanes))
+    assert s.min_value.value == base
+    assert s.max_value.value == base + stride * (lanes - 1)
+
+
+def test_scalable_vector():
+    base = 5
+    s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, 2, tvm.tirx.vscale() * 4))
+
+    assert s.min_value.value == base
+    assert s.max_value.same_as(tvm.arith.int_set.pos_inf())
+
+
+def test_add_sub():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10)}, (y, 10 + y))
+    ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (1, 21))
+    ck.verify(x - y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (-11, 9))
+
+
+def test_mul_div():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    tdiv = tvm.tirx.truncdiv
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True)
+    ck.verify(x * y, {x: tvm.arith.IntervalSet(0, 10)}, (0, 10 * y))
+    ck.verify(x * 2, {x: tvm.arith.IntervalSet(1, 10)}, (2, 20))
+    ck.verify(x * -2, {x: tvm.arith.IntervalSet(1, 10)}, (-20, -2))
+
+    ck.verify(tdiv(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, tdiv(10, y)))
+    ck.verify(tdiv(x, 2), {x: tvm.arith.IntervalSet(1, 10)}, (0, 5))
+
+    fld = tvm.tirx.floordiv
+    ck.verify(fld(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, fld(10, y)))
+    ck.verify(fld(x, 2), {x: tvm.arith.IntervalSet(-1, 10)}, (-1, 5))
+
+
+def test_mod():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    tmod = tvm.tirx.truncmod
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True)
+    ck.verify(tmod(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, y - 1))
+    ck.verify(tmod(x, 10), {x: tvm.arith.IntervalSet(1, 10)}, (0, 9))
+
+    flm = tvm.tirx.floormod
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(-10, 10)}, (0, 9))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 5)}, (3, 5))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(13, 15)}, (3, 5))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 15)}, (0, 9))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 11)}, (0, 9))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(1, 21)}, (0, 9))
+
+    fld = tvm.tirx.floordiv
+    z = tvm.tirx.Var("z", "int32")
+    ck.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 3))
+    ck.verify(
+        flm(y, 8),
+        {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)},
+        (
+            z * 8 + x * 4 - 8 * fld(z * 8 + x * 4, 8),
+            z * 8 + x * 4 + 3 - 8 * fld(z * 8 + x * 4, 8),
+        ),
+    )
+    ck1 = IntSetChecker()
+    ck1.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 2))
+    ck1.verify(
+        flm(y, 8), {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, (x * 4, x * 4 + 3)
+    )
+
+
+def test_max_min():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(tvm.tirx.max(x, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (1, 11))
+    ck.verify(tvm.tirx.min(x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 9))
+    ck.verify(tvm.tirx.min(x, y), {}, (tvm.tirx.min(x, y), tvm.tirx.min(x, y)))
+    ck.verify(tvm.tirx.max(x, y), {}, (tvm.tirx.max(x, y), tvm.tirx.max(x, y)))
+
+
+def test_select():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(tvm.tirx.Select(x > 0, x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 11))
+
+
+def check_region_bound(expect_region, var_dom, mode, predicate=None):
+    """Helper to check region bound estimation.
+
+    Parameters
+    ----------
+    expect_region: dict
+        The keys are of form (begin, end) or Expr as a single point. The values are
+        expected estimated region or region dict on different bindings.
+
+    var_dom: dict
+        Map var to iteration domain range.
+
+    mode: str
+        Specify "lowerbound", "upperbound" or else use strict bound estimation.
+
+    predicate: Expr
+        Extra predicate, defaults to True.
+    """
+    if predicate is None:
+        predicate = tvm.tirx.IntImm("bool", 1)
+    region = []
+    expect = []
+    for k, v in expect_region.items():
+        if not isinstance(k, tuple | list):
+            k = (k, k + 1)
+        region.append(tvm.ir.Range.from_min_extent(k[0], Analyzer().simplify(k[1] - k[0])))
+        expect.append(v)
+    if mode == "lowerbound":
+        result = tvm.arith.estimate_region_lower_bound(
+            region=region, var_dom=var_dom, predicate=predicate
+        )
+    elif mode == "upperbound":
+        result = tvm.arith.estimate_region_upper_bound(
+            region=region, var_dom=var_dom, predicate=predicate
+        )
+    else:
+        result = tvm.arith.estimate_region_strict_bound(
+            region=region, var_dom=var_dom, predicate=predicate
+        )
+    if result is None:
+        assert all([_ is None for _ in expect])
+        return
+    assert len(result) == len(expect)
+    for intset, expect_desc in zip(result, expect):
+        if isinstance(expect_desc, dict):
+            # check range on different free var bindings
+            for binding in expect_desc:
+                analyzer = Analyzer()
+                for k, v in binding:
+                    analyzer.bind(k, v)
+                expect_begin, expect_end = expect_desc[binding]
+                result_begin = analyzer.simplify(intset.min_value, 3)
+                result_end = analyzer.simplify(intset.max_value + 1, 3)
+                assert analyzer.can_prove_equal(result_begin - expect_begin, 0), (
+                    f"{result_begin} vs {expect_begin}"
+                )
+                assert analyzer.can_prove_equal(result_end - expect_end, 0), (
+                    f"{result_end} vs {expect_end}"
+                )
+        else:
+            # check range
+            expect_begin, expect_end = expect_desc
+            analyzer = Analyzer()
+            assert analyzer.can_prove_equal(intset.min_value - expect_begin, 0), (
+                f"{intset.min_value} vs {expect_begin}"
+            )
+            assert analyzer.can_prove_equal(intset.max_value - expect_end + 1, 0), (
+                f"{intset.max_value} vs {expect_end - 1}"
+            )
+
+
+def test_region_bound_not_independent():
+    # (i, i+2) and (i+2, i+4) are dependent, this the lowerbound is not available
+    i = tvm.tirx.Var("i", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=64),
+    }
+    check_region_bound({(i, i + 2): None, (i + 2, i + 4): None}, var_dom, mode="lowerbound")
+    check_region_bound({(i, i + 2): (0, 65), (i + 2, i + 4): (2, 67)}, var_dom, mode="upperbound")
+
+    # when only a subset of access indices are affine
+    i, j, k = tvm.tirx.Var("i", "int32"), tvm.tirx.Var("j", "int32"), tvm.tirx.Var("k", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=16),
+        j: tvm.ir.Range(begin=0, end=16),
+        k: tvm.ir.Range(begin=0, end=16),
+    }
+    check_region_bound(
+        {i // 4: None, j * 4 + i % 4: None, tirx.truncdiv(k, 2): None},
+        var_dom,
+        predicate=j * 4 + i % 4 > 3,
+        mode="lowerbound",
+    )
+    check_region_bound(
+        {i // 4: (0, 4), j * 4 + i % 4: (4, 64), tirx.truncdiv(k, 2): (0, 8)},
+        var_dom,
+        predicate=j * 4 + i % 4 > 3,
+        mode="upperbound",
+    )
+
+
+def test_region_bound_stride_too_wide():
+    i = tvm.tirx.Var("i", "int32")
+    var_dom = {i: tvm.ir.Range(begin=0, end=64)}
+    check_region_bound({(i * 4, i * 4 + 2): None}, var_dom, mode="lowerbound")
+    check_region_bound({(i * 4, i * 4 + 2): (0, 254)}, var_dom, mode="upperbound")
+
+
+def test_region_bound_small_stride():
+    i = tvm.tirx.Var("i", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=64),
+    }
+    check_region_bound({(i * 4, i * 4 + 8): (0, 260)}, var_dom, mode="lowerbound")
+
+
+def test_region_lower_bound_split_predicate():
+    x_o = tvm.tirx.Var("xo", "int32")
+    x_i = tvm.tirx.Var("xi", "int32")
+    x = x_o * 4 + x_i
+    var_dom = {
+        x_o: tvm.ir.Range(begin=0, end=16),
+        x_i: tvm.ir.Range(begin=0, end=4),
+    }
+    check_region_bound({(x * 4, x * 4 + 8): (0, 256)}, var_dom, predicate=x < 63, mode="lowerbound")
+
+    check_region_bound(
+        {(x * 4, x * 4 + 8): (0, 256), (x * 3, x * 3 + 5): (0, 191)},
+        var_dom,
+        predicate=x < 63,
+        mode="upperbound",
+    )
+
+
+def test_region_lower_bound_multiple_variables():
+    div = tvm.tirx.floordiv
+    mod = tvm.tirx.floormod
+    x = tvm.tirx.Var("x", "int32")
+    wid = tvm.tirx.Var("wid", "int32")
+    i = div(x, 16)
+    j = div(mod(x, 16), 4) * 8 + mod(x, 4) + div(wid, 32) * 4
+    k = wid % 32
+    var_dom = {
+        x: tvm.ir.Range(begin=0, end=32),
+        wid: tvm.ir.Range(begin=0, end=64),
+    }
+    check_region_bound({i: (0, 2), j: (0, 32), k: (0, 32)}, var_dom, mode="lowerbound")
+
+
+def test_region_lower_bound_negative_scale():
+    i = tvm.tirx.Var("i", "int32")
+    j = tvm.tirx.Var("j", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=4),
+        j: tvm.ir.Range(begin=0, end=4),
+    }
+    check_region_bound(
+        {(1 - i, 5 - i): (-2, 5), (20 - j * 4, 36 - j * 4): (8, 36)}, var_dom, mode="lowerbound"
+    )
+
+
+def test_region_lower_bound_for_non_perfect_tile():
+    h1 = tvm.tirx.Var("h1", "int32")
+    h2 = tvm.tirx.Var("h2", "int32")
+    h3 = tvm.tirx.Var("h3", "int32")
+
+    # non-uniform tiling, single inner variable
+    var_dom = {
+        h2: tvm.ir.Range(begin=0, end=10),
+    }
+    check_region_bound(
+        {
+            h3 * 8 + h2: {
+                (): (
+                    tvm.tirx.max(h3 * 8, 1),
+                    tvm.tirx.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)
+                ((h3, 27),): (h3 * 8, 224),  # h3 > 26: region is [h3 * 8, 224)
+            }
+        },
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2, h3 * 8 + h2 < 224),
+        mode="lowerbound",
+    )
+
+    # non-uniform tiling, two inner variables
+    var_dom = {
+        h1: tvm.ir.Range(begin=0, end=5),
+        h2: tvm.ir.Range(begin=0, end=2),
+    }
+    check_region_bound(
+        {
+            h3 * 8 + h2 * 5 + h1: {
+                (): (
+                    tvm.tirx.max(h3 * 8, 1),
+                    tvm.tirx.min(0, h3 * 8 - 214) + 224,
+                ),
+                ((h3, 0),): (1, 10),
+                ((h3, 10),): (h3 * 8, h3 * 8 + 10),
+                ((h3, 27),): (h3 * 8, 224),
+            }
+        },
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h2 * 5 + h1 < 224),
+        mode="lowerbound",
+    )
+
+    # lowerbound should fail on incompatible predicates
+    check_region_bound(
+        {h3 * 8 + h2 * 5 + h1: None},
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h1 * 2 + h2 < 224),
+        mode="lowerbound",
+    )
+    check_region_bound(
+        {h3 * 8 + h2 * 5 + h1: (h3 * 8, h3 * 8 + 10)},
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h1 * 2 + h2 < 224),
+        mode="upperbound",
+    )
+
+
+def test_region_lower_bound_unfusable():
+    var_dom = {
+        tvm.tirx.Var("i", "int32"): tvm.ir.Range(8),
+        tvm.tirx.Var("j", "int32"): tvm.ir.Range(4),
+    }
+    i, j = var_dom
+    check_region_bound({(i + j) // 2: (0, 6)}, var_dom, mode="lowerbound")
+
+
+def test_union_lower_bound():
+    neg_inf = tvm.arith.int_set.neg_inf()
+    pos_inf = tvm.arith.int_set.pos_inf()
+    set_0 = tvm.arith.IntervalSet(min_value=neg_inf, max_value=0)
+    set_1 = tvm.arith.IntervalSet(min_value=1, max_value=pos_inf)
+    result = tvm.arith.int_set.union_lower_bound([set_0, set_1])
+    assert result.min_value.same_as(neg_inf)
+    assert result.max_value.same_as(pos_inf)
+    set_2 = tvm.arith.IntervalSet(min_value=pos_inf, max_value=neg_inf)
+    result = tvm.arith.int_set.union_lower_bound([set_0, set_1, set_2])
+    assert result.min_value.same_as(neg_inf)
+    assert result.max_value.same_as(pos_inf)
+
+
+def test_modular_set():
+    ck = IntSetChecker()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    expr = (x * 2048 + y * 16) % 7168
+    ck.verify(
+        expr, {x: tvm.arith.IntervalSet(0, 128), y: tvm.arith.IntervalSet(0, 3584)}, (0, 7152)
+    )
+
+
+def test_relax_deep_variable_dependency_chain():
+    """Regression test for exponential variable-relaxation blowup.
+
+    When a variable's interval bound references another variable that is also in
+    the domain map, the evaluator relaxes it transitively. A diamond-shaped
+    chain -- where each variable's bound references the next one in *both* its
+    min and its max -- used to be re-expanded along every path, costing
+    O(2^depth) and hanging indefinitely. The relaxation is now memoized per
+    variable, so this completes in linear time.
+    """
+    ck = IntSetChecker()
+    n = 64  # 2^64 expansions without memoization; trivially fast with it.
+    xs = [tvm.tirx.Var(f"x{i}", "int32") for i in range(n + 1)]
+    dmap = {xs[i]: tvm.arith.IntervalSet(xs[i + 1] - 1, xs[i + 1] + 1) for i in range(n)}
+    dmap[xs[n]] = tvm.arith.IntervalSet(0, 100)
+    # x0 relaxes through the whole chain: [0 - n, 100 + n].
+    ck.verify(xs[0], dmap, (-n, 100 + n))
+
+
+def test_relax_cyclic_variable_dependency():
+    """A cyclic variable dependency must terminate (and stay symbolic)."""
+    ana = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    # x depends on y and y depends on x: relaxation must not loop forever.
+    dmap = {x: tvm.arith.IntervalSet(y, y), y: tvm.arith.IntervalSet(x, x)}
+    res = ana.int_set(x, dmap)
+    assert res is not None
+
+
+def test_estimate_region_accepts_external_analyzer():
+    i = tvm.tirx.Var("i", "int32")
+    tile = tvm.tirx.Var("tile", "int32")
+    region = [tvm.ir.Range.from_min_extent(i % tile, 1)]
+    dom = {i: tvm.ir.Range(0, 16)}
+
+    # Without knowing `tile`, the affine detection fails for exact bounds.
+    assert tvm.arith.estimate_region_lower_bound(region, dom, True) is None
+    assert tvm.arith.estimate_region_strict_bound(region, dom, True) is None
+    upper_without_analyzer = tvm.arith.estimate_region_upper_bound(region, dom, True)
+
+    analyzer = tvm.arith.Analyzer()
+    analyzer.bind(tile, tvm.tirx.const(4, "int32"))
+    # The external binding lets the affine detection succeed.
+    for estimate_region in [
+        tvm.arith.estimate_region_lower_bound,
+        tvm.arith.estimate_region_strict_bound,
+        tvm.arith.estimate_region_upper_bound,
+    ]:
+        result = estimate_region(region, dom, True, analyzer=analyzer)
+        assert result is not None
+        assert analyzer.can_prove_equal(result[0].min_value, 0)
+        assert analyzer.can_prove_equal(result[0].max_value, 3)
+
+    # The upper-bound fallback without analyzer is safe but much wider.
+    assert not analyzer.can_prove_equal(upper_without_analyzer[0].min_value, 0)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_modular_set.py

@@ -0,0 +1,234 @@
+# 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.
+# ruff: noqa: F841
+import tvm
+import tvm.testing
+from tvm import te
+
+
+def test_cast():
+    analyzer = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int8")
+    m = analyzer.modular_set((x * 3).astype("uint32"))
+    assert m.coeff == 3
+    assert m.base == 0
+    m = analyzer.modular_set((x * 3 + 1).astype("float32").astype("int32"))
+    assert m.coeff == 3
+    assert m.base == 1
+
+
+def test_add_sub():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int64"), tvm.tirx.Var("y", "int64")
+    m = analyzer.modular_set(x * 6 + y * 4)
+    assert m.coeff == 2
+    assert m.base == 0
+
+    analyzer.bind(y, x * 4 + 1)
+    m = analyzer.modular_set(1 - y)
+    assert m.coeff == 4
+    assert m.base == 0
+
+
+def test_mul():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set((x * 4 + 2) * (y * 6 + 1))
+    assert m.coeff == 4
+    assert m.base == 2
+
+
+def test_shift_left():
+    analyzer = tvm.arith.Analyzer()
+    x, y = te.var("x"), te.var("y")
+    m = analyzer.modular_set((x * 4 + 2) << 2)
+    assert m.coeff == 16
+    assert m.base == 8
+
+
+def test_floormod():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set(tvm.tirx.floormod(x * 128 + y * 4, 256))
+    assert m.coeff == 4
+    assert m.base == 0
+
+
+def test_div_shift():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    # not sure if x is non-negative
+    tdiv = tvm.tirx.truncdiv
+    m = analyzer.modular_set(tdiv(x * 4 + 2, 2))
+    assert m.coeff == 1
+    assert m.base == 0
+    # right shift always round down so it is fine
+    m = analyzer.modular_set((x * 4 + 2) >> 1)
+    assert m.coeff == 2
+    assert m.base == 1
+    fld = tvm.tirx.floordiv
+    m = analyzer.modular_set(fld(x * 4 + 2, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+    # x is non-negative
+    analyzer.update(x, tvm.arith.ConstIntBound(0, 100))
+    m = analyzer.modular_set(tdiv(x * 4 + 2, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+
+
+def test_mod():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    tmod = tvm.tirx.truncmod
+    fmod = tvm.tirx.floormod
+    # not sure if x is non-negative
+    m = analyzer.modular_set(tmod(x * 4 + 1, 4))
+    assert m.coeff == 1
+    assert m.base == 0
+    # no need to be positive if base == 0
+    m = analyzer.modular_set(tmod(x * 4, 4))
+    assert m.coeff == 4
+    assert m.base == 0
+    # floor mod tests
+    m = analyzer.modular_set(fmod(x * 4 + 3, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+    m = analyzer.modular_set(fmod(x * 4 + 3, 8))
+    assert m.coeff == 4
+    assert m.base == 3
+    # x is non-negative
+    analyzer.update(x, tvm.arith.ConstIntBound(0, 100))
+    m = analyzer.modular_set(tmod(x * 4 + 3, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+
+
+def test_min_max_select():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set(tvm.tirx.min(x * 3, y * 9))
+    assert m.coeff == 3
+    assert m.base == 0
+
+    m = analyzer.modular_set(tvm.tirx.max(x * 3 + 1, y * 9 + 4))
+    assert m.coeff == 3
+    assert m.base == 1
+
+    m = analyzer.modular_set(tvm.tirx.Select(x > 0, x * 3 + 1, y * 9 + 2))
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+def test_mix_index():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    analyzer = tvm.arith.Analyzer()
+    tdiv = tvm.tirx.truncdiv
+    m = analyzer.modular_set(a * 4 + b * 6 + 7)
+    assert m.coeff == 2
+    assert m.base == 1
+
+    m = analyzer.modular_set((a * 4 + 1) * (b * 8 + 3))
+    assert m.coeff == 4
+    assert m.base == 3
+
+    m = analyzer.modular_set(tdiv(a * 4 + 1, b * 8 + 3))
+    assert m.coeff == 1
+    assert m.base == 0
+
+    m = analyzer.modular_set((a * 4 + 1) * tdiv(b * 8, 4))
+    assert m.coeff == 2
+    assert m.base == 0
+
+    m = analyzer.modular_set((a * 12 + 1) - (b * 3 * 7 + 2))
+    assert m.coeff == 3
+    assert m.base == 2
+
+    m = analyzer.modular_set(a * 12 + tvm.tirx.min(b * 3 * 7, 2))
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+def test_constraint_scope():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    analyzer = tvm.arith.Analyzer()
+    tmod = tvm.tirx.truncmod
+
+    with analyzer.constraint_scope(tmod(b, 4) == 2):
+        m = analyzer.modular_set(b + 1)
+        assert m.coeff == 4
+        assert m.base == 3
+        with analyzer.constraint_scope(tmod(a, 2) == 1):
+            m = analyzer.modular_set(b + a * 2)
+            assert m.coeff == 4
+            assert m.base == 0
+        m = analyzer.modular_set(b + a * 2)
+        assert m.coeff == 2
+        assert m.base == 0
+
+    m = analyzer.modular_set(b + 1)
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+def test_intersect():
+    a = tvm.tirx.Var("a", "int32")
+    analyzer = tvm.arith.Analyzer()
+    tmod = tvm.tirx.truncmod
+    with analyzer.constraint_scope(tmod(a, 4) == 1):
+        with analyzer.constraint_scope(tmod(a, 3) == 1):
+            m = analyzer.modular_set(a)
+            assert m.coeff == 12
+            assert m.base == 1
+
+    with analyzer.constraint_scope(tmod(a, 3) == 2):
+        with analyzer.constraint_scope(tmod(a, 5) == 3):
+            with analyzer.constraint_scope(tmod(a, 7) == 2):
+                m = analyzer.modular_set(a)
+                assert m.coeff == 105
+                assert m.base == 23
+
+
+def test_let():
+    analyzer = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set(tvm.tirx.Let(x, y * 10, x + 1))
+    assert m.coeff == 10
+    assert m.base == 1
+
+
+def test_bitwise_and():
+    analyzer = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+
+    # RHS of bitwise_and is 2^p - 1
+    m = analyzer.modular_set((x * 16 + y * 4) & 31)
+    assert m.coeff == 4
+    assert m.base == 0
+
+    # arbitrary RHS
+    m = analyzer.modular_set((x * 16 + y * 4) & 17)
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_simplify.py

@@ -0,0 +1,195 @@
+# 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 pytest
+
+import tvm
+import tvm.ir
+import tvm.testing
+from tvm import tirx
+from tvm.script import tirx as T
+
+
+def test_simplify_reshape_flattened_index():
+    ana = tvm.arith.Analyzer()
+
+    i0 = tirx.Var("i0", "int64")
+    i1 = tirx.Var("i1", "int64")
+    ana.bind(i0, tvm.ir.Range(0, 8))
+    ana.bind(i1, tvm.ir.Range(0, 3))
+
+    i_flattened = i0 * 3 + i1
+    tvm.ir.assert_structural_equal(
+        ana.simplify((i_flattened) // 12 * 12 + (i_flattened) % 12 // 4 * 4 + (i_flattened) % 4),
+        i_flattened,
+    )
+
+
+dtype = tvm.testing.parameter(
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "float16",
+    "float32",
+    "float64",
+)
+
+
+def test_can_prove_self_identity(dtype):
+    ana = tvm.arith.Analyzer()
+
+    n = tirx.Var("n", dtype)
+    assert ana.can_prove(n == n)
+
+
+def test_can_prove_self_equal_to_self(dtype):
+    ana = tvm.arith.Analyzer()
+
+    n = tirx.Var("n", dtype)
+    assert ana.can_prove_equal(n, n)
+
+
+def test_simplify_symbolic_comparison():
+    ana = tvm.arith.Analyzer()
+
+    i0 = tirx.Var("i0", "int64")
+    i1 = tirx.Var("i1", "int64")
+    n, m = tvm.tirx.Var("n", "int64"), tvm.tirx.Var("m", "int64")
+    outer = (n + 31) // 32
+    PS = tvm.arith.ProofStrength
+
+    non_negative = tvm.arith.ConstIntBound(0, tvm.arith.ConstIntBound.POS_INF)
+    ana.update(n, non_negative)
+    ana.update(m, non_negative)
+    ana.bind(i0, tvm.ir.Range(0, outer))
+    ana.bind(i1, tvm.ir.Range(0, 32))
+    assert not ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32, PS.DEFAULT)
+    assert ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32 + m, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove(i0 * 32 + i1 + 1 <= (n + 31) // 32 * 32, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove((n + 31) // 32 * 32 >= i0 * 32 + i1 + 1, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove((n + 31) // 32 * 32 >= i0 * 32 + i1, PS.SYMBOLIC_BOUND)
+
+
+# These tests exercised arith::CanProve's substitution-based proof loop for
+# vscale-bearing expressions (iterating over known vscale values for a VLA target).
+# That loop has been removed -- arith no longer attempts target-dependent proofs
+# about scalable-vector lengths. The LOG(WARNING) for non-VLA targets is also gone.
+@pytest.mark.xfail(reason="arith no longer proves vscale-bearing inequalities via substitution")
+@pytest.mark.parametrize(
+    "expression",
+    [
+        T.vscale() * 32 < T.vscale() * 64,
+        T.vscale() * 2 * (T.vscale() * 2) >= T.vscale() * 4,
+        (T.vscale() * 4 + 114) // (T.vscale() * 4) * (T.vscale() * 4) >= 115,
+        64 % T.vscale() <= T.vscale(),
+    ],
+)
+def test_simplify_vscale_comparison_with_sve_target(expression):
+    ana = tvm.arith.Analyzer()
+
+    with tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}):
+        assert ana.can_prove(expression)
+
+
+@pytest.mark.xfail(
+    reason="arith no longer emits a LOG(WARNING) for vscale proofs on non-VLA targets"
+)
+def test_simplify_vscale_comparison_without_sve_target(capfd):
+    ana = tvm.arith.Analyzer()
+    vs = tvm.tirx.vscale()
+
+    with pytest.raises(AssertionError):
+        with tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"}):
+            assert ana.can_prove(vs * 32 < vs * 64)
+
+    warning_prefix = (
+        "Warning: The expression contains scalable values. An attempt to prove by substituting "
+        "with known values of vscale was not performed. This proof currently only supports "
+        "VLA targets, but the target was "
+    )
+    capture = capfd.readouterr().err
+    assert warning_prefix in capture
+    assert '"kind":"llvm"' in capture
+    assert '"mtriple":"aarch64-linux-gnu"' in capture
+
+
+def test_regression_simplify_inf_recursion():
+    ana = tvm.arith.Analyzer()
+    cond = tirx.Var("cond", "int32")
+
+    res = (tvm.tirx.NE(cond, 0).astype("int8") - tvm.tirx.NE(cond, 0).astype("int8")).astype(
+        "int32"
+    ) == 0
+    # regression in a previous case
+    # try compare and int set recursive call can cause infinite loop
+    ana.rewrite_simplify(res)
+
+
+def test_bind_allow_override():
+    ana = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    ana.bind(x, tvm.ir.Range(0, 10))
+    ana.bind(x, tvm.ir.Range(0, 5), allow_override=True)
+    assert ana.can_prove(x < 5)
+
+    with pytest.raises(RuntimeError, match="Trying to update var 'x' with a different const bound"):
+        ana.bind(x, tvm.ir.Range(0, 3))
+
+
+def test_simplify_floor_mod_with_linear_offset():
+    """
+    Test that the floor_mod is simplified correctly when the offset is linear.
+    """
+    ana = tvm.arith.Analyzer()
+    past_decoder_sequence_length = tirx.Var("past_decoder_sequence_length", "int64")
+    expr1 = (past_decoder_sequence_length + 1) * 64
+    divisor1 = (past_decoder_sequence_length + 1) * 32
+    assert ana.can_prove_equal(tvm.tirx.floormod(expr1, divisor1), 0)
+    divisor2 = 32 * (past_decoder_sequence_length + 1)
+    assert ana.can_prove_equal(tvm.tirx.floormod(expr1, divisor2), 0)
+
+
+def test_simplify_uint_floormod_const_scale_divisible():
+    """uint32 floormod(x * c1, c2) -> 0 when c1 % c2 == 0 (overflow-free)."""
+    ana = tvm.arith.Analyzer()
+    q = tirx.Var("q_stage_idx", "uint32")
+    expr = q * tirx.Cast("uint32", 128)
+    mod = expr % tirx.const(4, "uint32")
+    assert ana.can_prove_equal(mod, tirx.const(0, "uint32"))
+    tvm.ir.assert_structural_equal(ana.rewrite_simplify(mod), tirx.const(0, "uint32"))
+
+
+def test_simplify_float_division():
+    # Test for the discussion:
+    # https://discuss.tvm.apache.org/t/discuss-is-constant-division-to-multiplication-rewrite-in-tvm-necessary/18615
+    ana = tvm.arith.Analyzer()
+    x = tirx.Var("x", "float32")
+    ry = x / 27
+    # in old version, the division will be rewritten into x * T.float32(1 / 27)
+    sy = ana.rewrite_simplify(ry)
+    tvm.ir.assert_structural_equal(ry, sy)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_solve_linear_equations.py

@@ -0,0 +1,186 @@
+# 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.
+# ruff: noqa: F401
+import random
+import sys
+
+import pytest
+import tvm_ffi
+
+import tvm
+from tvm import arith, ir, testing, tirx
+from tvm.script import tirx as T
+
+
+def test_solution_consistency():
+    seed = random.randrange(sys.maxsize)
+    print(
+        "\nThis test is intentionally non-deterministic, "
+        f"if it fails please report it in GitHub issue together with this seed {seed}\n"
+    )
+    random.seed(seed)
+
+    def _check(num_vars, num_formulas, coef=(-5, 5), bounds=(-20, 20)):
+        variables = [tvm.tirx.Var("x" + str(i), "int32") for i in range(num_vars)]
+
+        relations = []
+        for i in range(num_formulas):
+            s1 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
+            s1 += random.randint(coef[0], coef[1])
+            s2 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
+            s2 += random.randint(coef[0], coef[1])
+            if random.random() < 0.7:
+                op = tvm.tirx.EQ
+            else:
+                # we also make sure it can correctly handle inequalities
+                op = random.choice([tvm.tirx.LE, tvm.tirx.LT, tvm.tirx.GE, tvm.tirx.GT])
+            relations.append(op(s1, s2))
+
+        vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in variables}
+        solution = arith.solve_linear_equations(relations, variables, vranges)
+
+        testing.check_int_constraints_trans_consistency(solution)
+
+        # leaving some variables as parameters should also be ok
+        for k in [1, 2]:
+            if len(variables) > k:
+                solution = arith.solve_linear_equations(relations, variables[:-k], vranges)
+                param_ranges = {v: vranges[v] for v in variables[-k:]}
+                testing.check_int_constraints_trans_consistency(solution, param_ranges)
+
+    for i in range(2):
+        _check(num_vars=1, num_formulas=1)
+    for i in range(2):
+        _check(num_vars=1, num_formulas=2)
+
+    for i in range(2):
+        _check(num_vars=2, num_formulas=1)
+    for i in range(2):
+        _check(num_vars=2, num_formulas=2)
+    for i in range(2):
+        _check(num_vars=2, num_formulas=3)
+
+    for i in range(3):
+        _check(num_vars=3, num_formulas=3, coef=(-2, 2))
+    for i in range(3):
+        _check(num_vars=3, num_formulas=4, coef=(-2, 2))
+
+    for i in range(3):
+        _check(num_vars=4, num_formulas=3, coef=(-1, 1))
+
+    for i in range(3):
+        _check(num_vars=10, num_formulas=2, coef=(-1, 1), bounds=(0, 4))
+    for i in range(3):
+        _check(num_vars=10, num_formulas=3, coef=(0, 1), bounds=(0, 4))
+
+
+def test_empty_var_to_solve():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    equations = [
+        tvm.tirx.EQ(x + y, 20),
+        tvm.tirx.EQ(x - y, 10),
+    ]
+    solution = arith.solve_linear_equations(equations)
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+    assert len(solution.src.variables) == 0
+    assert len(solution.src.ranges) == 0
+    assert tvm_ffi.structural_equal(solution.src.relations, equations)
+    assert tvm_ffi.structural_equal(solution.src, solution.dst)
+
+
+def test_unique_solution():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y, 20),
+            tvm.tirx.EQ(x - y, 10),
+        ],
+        [x, y],
+    )
+    assert list(solution.dst.variables) == []
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], T.int32(15))
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], T.int32(5))
+
+
+def test_low_rank():
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    ranges = {}
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y + z, 15),
+            tvm.tirx.EQ(x + y, 10),
+        ],
+        [x, y, z],
+        ranges,
+    )
+    [n0] = solution.dst.variables
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0 + 10)
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
+    assert tvm_ffi.structural_equal(solution.src_to_dst[z], T.int32(5))
+
+
+def test_infer_range():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ranges = {
+        x: tvm.ir.Range.from_min_extent(-5, 10),
+        y: tvm.ir.Range.from_min_extent(0, 10),
+    }
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y, 0),
+        ],
+        [x, y],
+        ranges,
+    )
+    [n0] = solution.dst.variables
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0)
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
+    # inferred from y's range
+    assert tvm_ffi.structural_equal(solution.dst.ranges[n0].min, T.int32(-9))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[n0].extent, T.int32(10))
+    # additional inequality is added into the system for x
+    [ineq] = solution.dst.relations
+    assert isinstance(ineq, tvm.tirx.LE)
+    assert tvm_ffi.structural_equal(ineq.a, T.int32(-5))
+    assert tvm_ffi.structural_equal(ineq.b, n0)
+
+
+def test_ill_formed():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y, 0),
+            tvm.tirx.EQ(x - y, 0),
+            tvm.tirx.EQ(x, 5),
+        ],
+        [x, y],
+        {},
+    )
+    assert list(solution.dst.variables) == []
+    [rel] = solution.dst.relations
+    ir.assert_structural_equal(rel, tirx.const(False))
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_solve_linear_inequality.py

@@ -0,0 +1,223 @@
+# 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 random
+import sys
+
+import pytest
+import tvm_ffi
+
+import tvm
+from tvm import arith, ir, testing, tirx
+from tvm.script import tirx as T
+
+
+@pytest.mark.skip(reason="See https://github.com/apache/tvm/issues/11458")
+def test_solution_consistency():
+    seed = random.randrange(sys.maxsize)
+    print(
+        "\nThis test is intentionally non-deterministic, "
+        f"if it fails please report it in GitHub issue together with this seed {seed}\n"
+    )
+    random.seed(seed)
+
+    def _check(variables, formulas, coef=(-5, 5), bounds=(-20, 20)):
+        vs = [tvm.tirx.Var("x" + str(i), "int32") for i in range(variables)]
+
+        fs = []
+        for i in range(formulas):
+            s1 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
+            s1 += random.randint(coef[0], coef[1])
+            s2 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
+            s2 += random.randint(coef[0], coef[1])
+            op = random.choice(
+                [tirx.expr.EQ, tirx.expr.LE, tirx.expr.LT, tirx.expr.GE, tirx.expr.GT]
+            )
+            fs.append(op(s1, s2))
+
+        vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in vs}
+        before = tvm.tirx.all(tirx.const(1, "bool"), *fs)
+        after = arith._ffi_api.SolveInequalitiesAsCondition(vs, vranges, fs)
+        after = tvm.tirx.all(tirx.const(1, "bool"), *after)
+        testing.check_bool_expr_is_true(before == after, vranges)
+
+        solution = arith.solve_linear_inequalities(fs, vs, vranges, deskew_range=True)
+        testing.check_int_constraints_trans_consistency(solution)
+
+    for i in range(3):
+        _check(1, 1)
+    for i in range(3):
+        _check(1, 2)
+
+    for i in range(3):
+        _check(2, 1)
+    for i in range(3):
+        _check(2, 2)
+    for i in range(3):
+        _check(2, 3)
+
+    # Somewhere here coefficients in the results become too large, leading to overflow,
+    # so we use smaller initial coefficients
+    for i in range(5):
+        _check(3, 3, coef=(-2, 2))
+    for i in range(5):
+        _check(3, 4, coef=(-2, 2))
+
+    for i in range(5):
+        _check(4, 3, coef=(-1, 1))
+
+    for i in range(5):
+        _check(10, 2, coef=(-1, 1), bounds=(0, 4))
+    for i in range(5):
+        _check(10, 3, coef=(0, 1), bounds=(0, 4))
+
+
+def test_dual_variable():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    variables = [x, y]
+    ranges = {
+        x: tvm.ir.Range(-100, 100),
+        y: tvm.ir.Range(0, 10),
+    }
+    problem = [
+        tvm.tirx.LE(x + y, 20),
+        tvm.tirx.GE(x - y, 10),
+    ]
+
+    # solution as conditions
+    solution = arith._ffi_api.SolveInequalitiesAsCondition(variables, ranges, problem)
+    assert tvm_ffi.structural_equal(solution[0], x >= (y + 10))
+    assert tvm_ffi.structural_equal(solution[1], x <= (20 - y))
+    assert tvm_ffi.structural_equal(solution[2], y >= 0)
+    assert tvm_ffi.structural_equal(solution[3], y <= 5)
+
+    # solve and get the ranges
+    solution = arith.solve_linear_inequalities(problem, variables, ranges)
+    # 0 <= y <=5
+    assert solution.ranges[y].min == 0
+    assert solution.ranges[y].extent == 6
+    # y + 10 <= x <= 20 - y
+    assert tvm_ffi.structural_equal(solution.ranges[x].min, y + 10)
+    assert solution.ranges[x].extent == 11  # max(10 - 2y)
+
+    # deskew the solved ranges to be starting from zero
+    solution = arith.solve_linear_inequalities(problem, variables, ranges, deskew_range=True)
+    [x_new, y_new] = solution.dst.variables
+    [rel] = solution.dst.relations
+    assert tvm_ffi.structural_equal(rel, (y_new * 2) + x_new <= 10)
+    assert tvm_ffi.structural_equal(solution.dst.ranges[x_new].min, T.int32(0))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[x_new].extent, T.int32(11))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[y_new].min, T.int32(0))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[y_new].extent, T.int32(6))
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], x_new + (y_new + 10))
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], y_new)
+    assert tvm_ffi.structural_equal(solution.dst_to_src[x_new], x - y - 10)
+    assert tvm_ffi.structural_equal(solution.dst_to_src[y_new], y)
+
+
+def test_equal():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    problem = [
+        tvm.tirx.GE(x + y, 10),
+        tvm.tirx.GE(x - y, 2),
+        tvm.tirx.LE(x, 6),
+    ]
+
+    solution = arith.solve_linear_inequalities(problem, [x, y])
+    assert solution.ranges[x].min == 6
+    assert solution.ranges[x].extent == 1
+    assert solution.ranges[y].min == 4
+    assert solution.ranges[y].extent == 1
+
+    solution = arith.solve_linear_inequalities(problem, [x, y], deskew_range=True)
+    assert len(solution.dst.variables) == 0
+    assert len(solution.dst.ranges) == 0
+    assert len(solution.dst.relations) == 0
+    assert solution.src_to_dst[x] == 6
+    assert solution.src_to_dst[y] == 4
+
+
+def test_multi_equal():
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    problem = [
+        tvm.tirx.LE(x, 6),
+        tvm.tirx.GE(x, 6),
+        tvm.tirx.GE(x - z * y, 0),
+        tvm.tirx.LE(x - z * y, 0),
+    ]
+
+    solution = arith.solve_linear_inequalities(problem, [x, y, z])
+    assert solution.ranges[x].min == 6
+    assert solution.ranges[x].extent == 1
+    assert len(solution.relations) == 3
+    assert tvm_ffi.structural_equal(solution.relations[0], x == z * y)
+
+    assert isinstance(solution.relations[1], tvm.tirx.LE)
+    assert solution.relations[1].b == 0
+    assert isinstance(solution.relations[2], tvm.tirx.LE)
+    assert solution.relations[2].b == 0
+    # (z*y - 6) <= 0 && (6 - z*y) <= 0
+    ana = tvm.arith.Analyzer()
+    assert ana.simplify(solution.relations[1].a + solution.relations[2].a) == 0
+    assert tvm_ffi.structural_equal(
+        solution.relations[1].a, (z * y - 6)
+    ) or tvm_ffi.structural_equal(solution.relations[2].a, (z * y - 6))
+
+    solution = arith.solve_linear_inequalities(problem, [x, y, z], deskew_range=True)
+    assert solution.src_to_dst[y] == y
+    assert solution.src_to_dst[z] == z
+    assert solution.src_to_dst[x] == 6
+
+
+def test_no_solution():
+    x = tvm.tirx.Var("x0", "int32")
+    vranges = {x: tvm.ir.Range.from_min_extent(-20, 41)}
+    problem = [-x - 4 <= -5 * x + 2, x * 4 + 5 <= x * 5]
+
+    solution = arith.solve_linear_inequalities(problem, [x], vranges, deskew_range=True)
+    assert list(solution.dst.variables) == []
+    [rel] = solution.dst.relations
+    ir.assert_structural_equal(rel, tirx.const(False))
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+
+    solution = arith.solve_linear_inequalities(problem, [x], vranges)
+    assert len(solution.variables) == 0
+    assert len(solution.ranges) == 0
+    [rel] = solution.relations
+    assert not rel
+
+
+def test_unbound_var_range():
+    x = tvm.tirx.Var("x0", "int32")
+    free_var = tvm.tirx.Var("fv", "int32")
+    vranges = {
+        x: tvm.ir.Range.from_min_extent(0, tvm.tirx.Cast("int32", 1 + tvm.tirx.log(free_var)))
+    }
+    problem = [x > 3]
+    solution = arith.solve_linear_inequalities(
+        problem,
+        [x],
+        vranges,
+    )
+    assert len(solution.variables) == 1
+    assert len(solution.ranges) == 0
+    assert len(solution.relations) == 3
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_z3.py

@@ -0,0 +1,756 @@
+# 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 gc
+import queue
+import threading
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.arith import Analyzer, ProofStrength
+
+# The Z3 prover is only consulted at the kSymbolicBound strength so the common
+# default path never pays the prover cost.
+SB = ProofStrength.SYMBOLIC_BOUND
+
+
+def _require_z3(analyzer):
+    if not analyzer.is_z3_enabled:
+        pytest.skip("Z3 prover is disabled in this build")
+
+
+def implies(x, y):
+    return tirx.Or(tirx.Not(x), y)
+
+
+# ---------------------------------------------------------------------------
+# API availability (works regardless of whether Z3 is built)
+# ---------------------------------------------------------------------------
+
+
+def test_z3_capability_query():
+    # `is_z3_enabled` is the supported way to detect the build configuration.
+    # The Z3-specific debug/config methods work only when it is True, and raise
+    # a clear error otherwise.
+    analyzer = Analyzer()
+    assert isinstance(analyzer.is_z3_enabled, bool)
+
+    if analyzer.is_z3_enabled:
+        assert isinstance(analyzer.get_smtlib2(), str)
+        assert isinstance(analyzer.get_z3_stats(), str)
+    else:
+        with pytest.raises(RuntimeError):
+            analyzer.get_smtlib2()
+        with pytest.raises(RuntimeError):
+            analyzer.get_z3_stats()
+        with pytest.raises(RuntimeError):
+            analyzer.set_z3_timeout_ms(1000)
+        with pytest.raises(RuntimeError):
+            analyzer.set_z3_rlimit(0)
+
+
+def test_z3_context_lifetime_outlives_worker_thread():
+    _require_z3(Analyzer())
+
+    result_queue = queue.Queue()
+
+    def worker():
+        try:
+            analyzer = Analyzer()
+            x = tirx.Var("x", "int32")
+            analyzer.bind(x, tvm.ir.Range(0, 16))
+            assert analyzer.can_prove(x >= 0, SB)
+            result_queue.put(("analyzer", analyzer))
+        except BaseException as err:  # pylint: disable=broad-exception-caught
+            result_queue.put(("error", err))
+
+    thread = threading.Thread(target=worker)
+    thread.start()
+    thread.join()
+
+    kind, payload = result_queue.get_nowait()
+    if kind == "error":
+        raise payload
+
+    del payload
+    gc.collect()
+
+
+# ---------------------------------------------------------------------------
+# Examples the native analyzer cannot prove but Z3 can.
+#
+# Each case asserts both that the native analyzers (kDefault, Z3 gated off)
+# fail and that Z3 (kSymbolicBound) succeeds. This demonstrates the added value
+# of the Z3 backend and that it is correctly gated behind kSymbolicBound.
+# ---------------------------------------------------------------------------
+
+
+def test_z3_floor_division_identity_constraint():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+
+    expr = ((b - a) // c) * c + a <= b
+    with analyzer.constraint_scope(tirx.all(a > 0, b > 0, c > 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_division_identity_via_bind_range():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+
+    analyzer.bind(a, tvm.ir.Range(1, 100000))
+    analyzer.bind(b, tvm.ir.Range(1, 100000))
+    analyzer.bind(c, tvm.ir.Range(1, 100000))
+
+    expr = ((b - a) // c) * c + a <= b
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_multiplication_monotonicity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    d = tirx.Var("d", "int32")
+
+    expr = implies(tirx.all(a < b, b < c, a * d < b * d), b * d < c * d)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_nested_floor_division_collapse():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    expr = implies(
+        tirx.all(a >= 0, a < 128),
+        a // 128 == (a // 64 * 32 + a % 32 // 16 * 8) // 64,
+    )
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_deeply_nested_floor_division_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    expr = implies(
+        tirx.all(a >= 0, a < 128),
+        (
+            a % 16 * 64
+            + a // 64 * 32
+            + a % 8 // 4 * 32
+            + (a % 32 // 16 + a % 2) % 2 * 8
+            + 16
+            - (a // 64 + a % 8 // 4) // 2 * 64
+        )
+        // 512
+        == (
+            a % 16 * 64
+            + a // 64 * 32
+            + a % 8 // 4 * 32
+            + (a % 32 // 16 + a % 2) % 2 * 8
+            - (a // 64 + a % 8 // 4) // 2 * 64
+        )
+        // 512,
+    )
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_min_max_sum_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = tirx.max(x, y) + tirx.min(x, y) == x + y
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_select_absolute_value_nonneg():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    expr = tirx.Select(x >= 0, x, -x) >= 0
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_transitive_inequality():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = implies(tirx.all(a <= b, b <= c), a <= c)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_square_expansion_nonneg():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = (a + b) * (a + b) >= a * a + b * b
+    with analyzer.constraint_scope(tirx.all(a >= 0, b >= 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_square_monotonicity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(tirx.all(0 <= a, a <= b), a * a <= b * b)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_strict_multiplication():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    d = tirx.Var("d", "int32")
+    expr = implies(tirx.all(a < b, d > 0), a * d < b * d)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_division_monotonicity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = implies(tirx.all(a <= b, c > 0), tirx.floordiv(a, c) <= tirx.floordiv(b, c))
+    assert not analyzer.can_prove(expr)
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_division_lower_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(b > 0, tirx.floordiv(a, b) * b <= a)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_modulo_range():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(b > 0, tirx.all(0 <= tirx.floormod(a, b), tirx.floormod(a, b) < b))
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_flattened_index_bound():
+    # Classic index-flattening bound used throughout TVM: for a row index i in
+    # [0, m) and a column index j in [0, n), the flattened index i * n + j stays
+    # within [0, m * n).
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    i = tirx.Var("i", "int32")
+    j = tirx.Var("j", "int32")
+    m = tirx.Var("m", "int32")
+    n = tirx.Var("n", "int32")
+    expr = tirx.all(0 <= i * n + j, i * n + j < m * n)
+    with analyzer.constraint_scope(tirx.all(0 <= i, i < m, 0 <= j, j < n, m > 0, n > 0)):
+        assert not analyzer.can_prove(expr)
+        analyzer.set_z3_rlimit(0)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_modular_combination():
+    # Native modular_set tracks single-variable moduli, but combining two
+    # independent modular facts to reason about their sum is left to Z3.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = tirx.floormod(x + y, 2) == 0
+    with analyzer.constraint_scope(tirx.all(tirx.floormod(x, 6) == 0, tirx.floormod(y, 6) == 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_square_non_negative():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert not analyzer.can_prove(a * a >= 0)
+    assert analyzer.can_prove(a * a >= 0, SB)
+
+
+def test_z3_min_max_average_bounds():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    assert not analyzer.can_prove(tirx.max(a, b) * 2 >= a + b)
+    assert analyzer.can_prove(tirx.max(a, b) * 2 >= a + b, SB)
+    assert analyzer.can_prove(tirx.min(a, b) * 2 <= a + b, SB)
+
+
+def test_z3_symbolic_bind_range_with_constraint():
+    # Combine a symbolic range binding (x in [0, n)) with a constraint on the
+    # extent to derive a concrete bound on x.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    n = tirx.Var("n", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, n))
+    with analyzer.constraint_scope(n <= 8):
+        assert not analyzer.can_prove(x < 8)
+        assert analyzer.can_prove(x < 8, SB)
+
+
+def test_z3_equality_congruence():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(a == b, a * a == b * b)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_integer_strict_transitivity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    # Over the integers, a < b and b < c implies a + 1 < c.
+    expr = implies(tirx.all(a < b, b < c), a + 1 < c)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_if_then_else_absolute_value():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    expr = tirx.if_then_else(x >= 0, x, -x) >= 0
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_unsigned_non_negative():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    u = tirx.Var("u", "uint32")
+    assert not analyzer.can_prove(u >= 0)
+    assert analyzer.can_prove(u >= 0, SB)
+
+
+def test_z3_unsigned64_non_negative():
+    # Exercises the special-cased uint64 range handling (UINT64_MAX bound).
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    u = tirx.Var("u", "uint64")
+    assert not analyzer.can_prove(u >= 0)
+    assert analyzer.can_prove(u >= 0, SB)
+
+
+def test_z3_int64_square_expansion():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int64")
+    b = tirx.Var("b", "int64")
+    expr = (a + b) * (a + b) >= a * a + b * b
+    with analyzer.constraint_scope(tirx.all(a >= 0, b >= 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_boolean_variable_reasoning():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    p = tirx.Var("p", "bool")
+    q = tirx.Var("q", "bool")
+    expr = implies(tirx.And(p, q), tirx.Or(p, q))
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_not_equal_from_strict_less():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = implies(x < y, tirx.NE(x, y))
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_let_expression():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    y = tirx.Var("y", "int32")
+    t = tirx.Var("t", "int32")
+    let = tirx.Let(t, y * 2, t)
+    assert not analyzer.can_prove(let == y * 2)
+    assert analyzer.can_prove(let == y * 2, SB)
+
+
+def test_z3_cast_preserves_bounds():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    s = tirx.Var("s", "int16")
+    widened = tirx.Cast("int32", s)
+    assert analyzer.can_prove(widened <= 32767, SB)
+    assert analyzer.can_prove(widened >= -32768, SB)
+
+
+def test_z3_bitwise_and_mask_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    assert analyzer.can_prove(tirx.bitwise_and(x, tirx.IntImm("int32", 7)) < 8, SB)
+
+
+def test_z3_bitwise_and_le_operand():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.bind(y, tvm.ir.Range(0, 256))
+    # Bit-vector reasoning over two variables exceeds the default deterministic
+    # rlimit; lift it (0 == unlimited, still deterministic) for this proof.
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.bitwise_and(x, y) <= x, SB)
+
+
+def test_z3_bitwise_or_ge_operand():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.bind(y, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.bitwise_or(x, y) >= x, SB)
+
+
+def test_z3_bitwise_xor_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.bind(y, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.bitwise_xor(x, y) < 256, SB)
+
+
+def test_z3_bitwise_not_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    # Two's complement: ~x == -x - 1.
+    assert analyzer.can_prove(tirx.bitwise_not(x) == -x - 1, SB)
+
+
+def test_z3_shift_right_halves():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    # For non-negative x, (x >> 1) * 2 <= x.
+    assert analyzer.can_prove(tirx.shift_right(x, tirx.IntImm("int32", 1)) * 2 <= x, SB)
+
+
+def test_z3_shift_left_lower_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    n = tirx.Var("n", "int32")
+    # Keep operands small so the 32-bit left shift cannot overflow; then
+    # x << n == x * 2 ** n >= x for x >= 1.
+    analyzer.bind(x, tvm.ir.Range(1, 16))
+    analyzer.bind(n, tvm.ir.Range(0, 4))
+    # Bit-vector shift reasoning exceeds the default deterministic rlimit.
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.shift_left(x, n) >= x, SB)
+
+
+# ---------------------------------------------------------------------------
+# Soundness / negative tests (Z3 must NOT prove false predicates)
+# ---------------------------------------------------------------------------
+
+
+def test_z3_negative_unprovable_inequality():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    # a < b does not hold for arbitrary a, b.
+    assert not analyzer.can_prove(a < b, SB)
+    # a * a > a is false (e.g. a == 0).
+    assert not analyzer.can_prove(a * a > a, SB)
+
+
+def test_z3_truncmod_can_be_negative():
+    # Regression test for truncated div/mod semantics: TVM Div/Mod round toward
+    # zero, so truncmod(a, 4) can be negative. A solver that modeled them as
+    # Euclidean would unsoundly "prove" truncmod(a, 4) >= 0.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert not analyzer.can_prove(tirx.truncmod(a, 4) >= 0, SB)
+
+
+def test_z3_truncdiv_truncmod_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = tirx.truncdiv(a, b) * b + tirx.truncmod(a, b) == a
+    with analyzer.constraint_scope(b != 0):
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floormod_nested_identities():
+    # Ported from TileLang's test_divmod. Here `%` is floormod: nested floormod
+    # by opposite-sign divisors collapses to the single-divisor result, while
+    # the mixed case does not.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert not analyzer.can_prove(a % 2 % -2 - a % 2 == 0, SB)
+    assert analyzer.can_prove(a % -2 % 2 - a % 2 == 0, SB)
+
+
+def test_z3_floormod_nonnegative():
+    # In contrast to truncmod, floormod with a positive divisor is always in
+    # [0, divisor), which Z3 should be able to prove.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert analyzer.can_prove(tirx.floormod(a, 4) >= 0, SB)
+    assert analyzer.can_prove(tirx.floormod(a, 4) < 4, SB)
+
+
+def test_z3_shift_does_not_poison_solver():
+    # Regression test: evaluating a shift expression must not add permanent
+    # assertions (such as `b >= 0` / `b < 64`) to the shared solver. Otherwise
+    # an unrelated, unbounded `b` would be wrongly provable to be < 100.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+
+    # Touch a shift expression so the prover visits the shift amount `b`.
+    analyzer.can_prove(tirx.shift_left(a, b) >= 0, SB)
+
+    # `b` is otherwise unconstrained, so this must remain unprovable.
+    assert not analyzer.can_prove(b < 100, SB)
+    assert not analyzer.can_prove(b >= 0, SB)
+
+
+def test_z3_constraint_scope_is_popped():
+    # Constraints entered through a scope must be removed once the scope exits,
+    # i.e. EnterConstraint's solver.push()/pop() must be balanced.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    with analyzer.constraint_scope(x > 5):
+        assert analyzer.can_prove(x > 0, SB)
+    # The constraint is gone; x is unconstrained again.
+    assert not analyzer.can_prove(x > 0, SB)
+
+
+def test_z3_opaque_call_is_safe():
+    # An opaque/unsupported sub-expression is modeled as a fresh free variable.
+    # It must neither crash nor be provable on its own, yet still be usable as a
+    # constraint.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    call = tirx.call_extern("int32", "foo", x)
+    assert not analyzer.can_prove(call > 0, SB)
+    with analyzer.constraint_scope(call > 0):
+        assert analyzer.can_prove(call > 0, SB)
+    assert not analyzer.can_prove(call > 0, SB)
+
+
+def test_z3_shift_overflow_is_not_proven():
+    # Z3 models fixed-width shifts via bit-vectors, so it correctly refuses to
+    # prove `x << n >= x` for an unbounded `x` (a large `x` overflows int32 and
+    # wraps to a negative value). This guards against unsound shift modeling.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    n = tirx.Var("n", "int32")
+    analyzer.set_z3_rlimit(0)
+    expr = implies(tirx.all(x >= 1, n >= 0, n < 8), tirx.shift_left(x, n) >= x)
+    assert not analyzer.can_prove(expr, SB)
+
+
+def test_z3_analyzers_are_isolated():
+    # Analyzers share a thread-local Z3 context but own separate solvers, so
+    # constraints and bindings in one must never leak into another.
+    analyzer_a = Analyzer()
+    analyzer_b = Analyzer()
+    _require_z3(analyzer_a)
+
+    x = tirx.Var("x", "int32")
+    with analyzer_a.constraint_scope(x > 100):
+        assert analyzer_a.can_prove(x > 50, SB)
+        assert not analyzer_b.can_prove(x > 50, SB)
+
+    analyzer_c = Analyzer()
+    analyzer_d = Analyzer()
+    analyzer_c.bind(x, tvm.ir.Range(0, 10))
+    assert analyzer_c.can_prove(x < 10, SB)
+    assert not analyzer_d.can_prove(x < 10, SB)
+
+
+def test_z3_repeated_can_prove_is_consistent():
+    # Repeated queries must be stateless: a CanProve call must not pollute the
+    # solver and change the result of a subsequent call.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    assert analyzer.can_prove(x > 0, SB) == analyzer.can_prove(x > 0, SB)
+
+    analyzer.bind(x, tvm.ir.Range(5, 10))
+    assert analyzer.can_prove(x >= 5, SB)
+    assert analyzer.can_prove(x >= 5, SB)
+
+
+def test_z3_is_gated_behind_symbolic_bound():
+    # The Z3 fallback must not run at the default strength.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = ((b - a) // c) * c + a <= b
+    with analyzer.constraint_scope(tirx.all(a > 0, b > 0, c > 0)):
+        assert not analyzer.can_prove(expr, ProofStrength.DEFAULT)
+        assert analyzer.can_prove(expr, SB)
+
+
+# ---------------------------------------------------------------------------
+# SMT-LIB2 export
+# ---------------------------------------------------------------------------
+
+
+def test_z3_smtlib2_roundtrip():
+    z3 = pytest.importorskip("z3")
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = ((b - a) // c) * c + a <= b
+
+    solver = z3.Solver()
+    with analyzer.constraint_scope(tirx.all(a > 0, b > 0, c > 0)):
+        solver.from_string(analyzer.get_smtlib2(expr))
+    assert solver.check() == z3.unsat
+
+
+def test_z3_smtlib2_roundtrip_with_timeout():
+    z3 = pytest.importorskip("z3")
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    analyzer.set_z3_timeout_ms(1000)
+
+    expr = implies(tirx.all(a > 0, b > 0, c > 0), ((b - a) // c) * c + a <= b)
+    solver = z3.Solver()
+    solver.from_string(analyzer.get_smtlib2(expr))
+    assert solver.check() == z3.unsat
+
+
+if __name__ == "__main__":
+    tvm.testing.main()