chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -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()