This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 7bcc2e51c [FEAT] Infer constant bounds from affine predicates (#69)
7bcc2e51c is described below
commit 7bcc2e51ce7da4532918a277fd70e5b51aa9e524
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 25 13:38:55 2026 -0400
[FEAT] Infer constant bounds from affine predicates (#69)
Constant-bound analysis now derives variable bounds from predicates on
positive-scale affine expressions such as `3*x + 2 == -7`.
Extract one signed index variable with constant scale and offset, invert
its bounds using exact integer floor and ceiling division, and retain
conservative bounds for unsupported expressions or unrepresentable
endpoints. Update the existing cyclic-shift regression to use context
bindings with the current Analyzer API.
---
src/sym/const_int_bound.cc | 60 ++++++++++++++++++++++++++++++++
tests/python/sym/test_const_int_bound.py | 14 ++++++++
tests/python/sym/test_iter_affine_map.py | 4 ++-
3 files changed, 77 insertions(+), 1 deletion(-)
diff --git a/src/sym/const_int_bound.cc b/src/sym/const_int_bound.cc
index 37cd96989..03ac9d566 100644
--- a/src/sym/const_int_bound.cc
+++ b/src/sym/const_int_bound.cc
@@ -25,6 +25,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/op.h>
+#include <tvm/ir/prim/const_fold.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/ir/prim/pattern_match.h>
@@ -37,6 +38,7 @@
#include "constraint_extract.h"
#include "constraint_helpers.h"
#include "int_operator.h"
+#include "product_normal_form.h"
namespace tvm {
namespace sym {
@@ -851,6 +853,64 @@ class ConstIntBoundAnalyzer::Impl final
}
}
+ // A bound on an affine expression does not by itself bound queries for its
+ // variable. For scale > 0, lo <= scale * var + offset <= hi implies
+ // ceil((lo - offset) / scale) <= var <= floor((hi - offset) / scale).
+ // For example, -6 <= 3*x + 2 <= 13 implies -2 <= x <= 3.
+ // Restrict this to signed index arithmetic, where overflow is undefined;
+ // the same inference would be unsound for wrapping unsigned arithmetic.
+ const size_t raw_count = info.size();
+ for (size_t i = 0; i < raw_count; ++i) {
+ const BoundInfo fact = info[i];
+ if (!prim::detail::IsIndexTypedExpr(fact.expr) ||
fact.expr.as<PrimVar>()) continue;
+
+ // Match one nonconstant sum term, then one variable in its product. Do
+ // not simplify using the current context while extracting new facts.
+ ffi::BigInt offset = 0, scale = 1;
+ ffi::Optional<PrimExpr> product, symbol;
+ bool valid = true;
+ UnpackSum(fact.expr, [&](PrimExpr term, int sign) {
+ if (const auto* imm = term.as<IntImmObj>()) {
+ offset += imm->value * sign;
+ } else if (!product.has_value()) {
+ product = term;
+ scale = sign;
+ } else {
+ valid = false;
+ }
+ });
+ if (!valid || !product.has_value()) continue;
+ UnpackReduction<prim::MulObj>(*product, [&](PrimExpr factor) {
+ if (const auto* imm = factor.as<IntImmObj>()) {
+ scale *= imm->value;
+ } else if (factor.as<PrimVar>() && !symbol.has_value()) {
+ symbol = factor;
+ } else {
+ valid = false;
+ }
+ });
+ if (!valid || !symbol.has_value() || scale <= 0) continue;
+
+ // ceil(n / scale) = -floor(-n / scale), including negative n.
+ // BigInt avoids intermediate overflow. Missing or unrepresentable
+ // endpoints remain conservatively infinite, including sentinel values.
+ int64_t lower = kNegInf, upper = kPosInf;
+ if (fact.bound.min_value != kNegInf) {
+ ffi::BigInt bound = -ffi::floordiv(offset - fact.bound.min_value,
scale);
+ if (auto value = bound.as<int64_t>(); value && *value > kNegInf &&
*value < kPosInf) {
+ lower = *value;
+ }
+ }
+ if (fact.bound.max_value != kPosInf) {
+ ffi::BigInt bound = ffi::floordiv(ffi::BigInt(fact.bound.max_value) -
offset, scale);
+ if (auto value = bound.as<int64_t>(); value && *value > kNegInf &&
*value < kPosInf) {
+ upper = *value;
+ }
+ }
+ if (lower != kNegInf || upper != kPosInf) {
+ info.emplace_back(*symbol, MakeBound(lower, upper));
+ }
+ }
return info;
}
diff --git a/tests/python/sym/test_const_int_bound.py
b/tests/python/sym/test_const_int_bound.py
index c4f6fd5e4..ae3fcf924 100644
--- a/tests/python/sym/test_const_int_bound.py
+++ b/tests/python/sym/test_const_int_bound.py
@@ -104,6 +104,20 @@ def test_plain_var_non_negative_bound_requires_context():
assert analyzer.const_int_bound(var).min_value == NEG_INF
+def test_affine_predicate_bound():
+ x = tvm_tirx.tirx.Var("x", "int64")
+ analyzer = tvm_tirx.sym.Analyzer()
+ parent = tvm_tirx.sym.ConstraintContext()
+ context = parent.snapshot()
+ context.append(tvm_tirx.sym.PredicateConstraint(3 * x + 2 == -7))
+ analyzer.sync_constraint_context(context)
+ bound = analyzer.const_int_bound(x)
+ assert (bound.min_value, bound.max_value) == (-3, -3)
+ analyzer.sync_constraint_context(parent)
+ bound = analyzer.const_int_bound(x)
+ assert (bound.min_value, bound.max_value) == (NEG_INF, POS_INF)
+
+
class TestCastBound(BaseCompare):
x = tvm_tirx.tirx.Var("x", "int8")
tmod = tvm_tirx.tirx.truncmod
diff --git a/tests/python/sym/test_iter_affine_map.py
b/tests/python/sym/test_iter_affine_map.py
index 9c91e5aa2..ca0764b80 100644
--- a/tests/python/sym/test_iter_affine_map.py
+++ b/tests/python/sym/test_iter_affine_map.py
@@ -1319,7 +1319,9 @@ def test_iter_map_simplify_cyclic_shift():
simplified = tvm_tirx.sym.iter_map_simplify([index], var_dom([(fused,
42)]))[0]
for value in range(42):
analyzer = tvm_tirx.sym.Analyzer()
- analyzer.bind(fused, tvm_tirx.tirx.const(value, "int32"))
+ context = tvm_tirx.sym.ConstraintContext()
+ context.append(tvm_tirx.sym.BindingConstraint(fused,
tvm_tirx.tirx.const(value, "int32")))
+ analyzer.sync_constraint_context(context)
expected = ((value // 7 + 1) % 6) * 7 + (value % 7 + 2) % 7
assert int(analyzer.simplify(simplified)) == expected
