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 bb07a8c94 [Fix][Sym] Guard IntSet evaluation against floating-point
subexpressions (#216)
bb07a8c94 is described below
commit bb07a8c94571ee339f28af4ca705bc661816dcd9
Author: Shushi Hong <[email protected]>
AuthorDate: Fri Nov 6 17:19:30 2026 -0500
[Fix][Sym] Guard IntSet evaluation against floating-point subexpressions
(#216)
This fixes a crash in IntSet analysis when floating-point expressions
appear in bindings or inside integer casts. The analyzer could apply
integer interval rules to these expressions and try to create an
`IntImm` with a floating-point type.
The relevant handlers now check the expression’s result type and return
`IntervalSet::Everything()` for unsupported types. The checks use
borrowed type pointers to avoid extra reference-count operations.
---
src/sym/int_set.cc | 14 ++++++++++++++
tests/python/sym/test_intset.py | 24 ++++++++++++++++++++++++
2 files changed, 38 insertions(+)
diff --git a/src/sym/int_set.cc b/src/sym/int_set.cc
index 13285df90..864416232 100644
--- a/src/sym/int_set.cc
+++ b/src/sym/int_set.cc
@@ -428,6 +428,7 @@ class IntervalSetEvaluator : public
tvm::ExprFunctor<IntervalSet(const Expr&)> {
}
IntervalSet Dispatch_(const VarObj* op) final {
+ if (!CheckIntegerOrBoolType(op)) return IntervalSet::Everything();
Var var = ffi::GetRef<Var>(op);
auto it = dom_map_.find(var);
@@ -557,17 +558,20 @@ class IntervalSetEvaluator : public
tvm::ExprFunctor<IntervalSet(const Expr&)> {
}
IntervalSet Dispatch_(const prim::BroadcastObj* op) final {
+ if (!CheckIntegerOrBoolType(op)) return IntervalSet::Everything();
TVM_FFI_ICHECK(eval_vec_);
return Dispatch(op->value);
}
IntervalSet Dispatch_(const prim::SelectObj* op) final {
+ if (!CheckIntegerOrBoolType(op)) return IntervalSet::Everything();
IntervalSet true_set = this->Eval(op->true_value);
IntervalSet false_set = this->Eval(op->false_value);
return Union(analyzer_, false_set, true_set);
}
IntervalSet Dispatch_(const prim::CastObj* op) final {
+ if (!CheckIntegerOrBoolType(op)) return IntervalSet::Everything();
IntervalSet value_set = this->Eval(op->value);
// short cut for the int set.
if (value_set->min_value.same_as(value_set->max_value)) {
@@ -614,6 +618,15 @@ class IntervalSetEvaluator : public
tvm::ExprFunctor<IntervalSet(const Expr&)> {
}
private:
+ // Integer expressions can contain floating-point subexpressions. Only apply
+ // integer interval rules to nodes with an integer or boolean result type.
+ static bool CheckIntegerOrBoolType(const ExprObj* op) {
+ const auto* ty = op->ty.as<PrimTypeObj>();
+ return ty != nullptr &&
+ (ty->dtype.code == DLDataTypeCode::kDLInt || ty->dtype.code ==
DLDataTypeCode::kDLUInt ||
+ ty->dtype.code == DLDataTypeCode::kDLBool);
+ }
+
// whether set is exactly single point that equals value.
bool MatchPoint(const IntervalSet& set, const PrimExpr& value) const {
return set->min_value.same_as(value) && set->max_value.same_as(value);
@@ -622,6 +635,7 @@ class IntervalSetEvaluator : public
tvm::ExprFunctor<IntervalSet(const Expr&)> {
template <typename TOp, typename T>
inline IntervalSet VisitBinaryExpr_(const T* op) {
static_assert(std::is_same<typename TOp::ContainerType, T>::value,
"constraint");
+ if (!CheckIntegerOrBoolType(op)) return IntervalSet::Everything();
IntervalSet a = this->Eval(op->a);
IntervalSet b = this->Eval(op->b);
if (MatchPoint(a, op->a) && MatchPoint(b, op->b)) {
diff --git a/tests/python/sym/test_intset.py b/tests/python/sym/test_intset.py
index 573fe996e..183095101 100644
--- a/tests/python/sym/test_intset.py
+++ b/tests/python/sym/test_intset.py
@@ -154,6 +154,30 @@ def test_select():
)
+def test_float_subexpressions():
+ analyzer = Analyzer()
+ i = tirx.Var("i", "int32")
+ scale = tirx.Var("scale", "float32")
+ coordinate = tirx.Cast("float32", i) * scale
+ domain = {i: tvm_tirx.sym.IntervalSet(0, 4)}
+ for expr in [
+ tvm_tirx.ir.FloatImm("float32", 3.5),
+ scale,
+ tirx.Cast("float32", i),
+ coordinate,
+ tirx.call_pure_extern("returns_float", ty="float32"),
+ ]:
+ assert analyzer.int_set(expr, domain).is_everything()
+ assert analyzer.int_set(tirx.Cast("int64", expr),
domain).is_everything()
+
+ context = tvm_tirx.sym.ConstraintContext()
+ context.append(tvm_tirx.sym.RangeConstraint(i, tvm_tirx.ir.Range(0, 5)))
+ value = tirx.Var("coordinate", "float32")
+ context.append(tvm_tirx.sym.BindingConstraint(value, coordinate))
+ analyzer.sync_constraint_context(context)
+ assert analyzer.int_set(value).is_everything()
+
+
def check_region_bound(expect_region, var_dom, mode, predicate=None):
"""Helper to check region bound estimation.
