diff --git a/src/s_tir/analysis/conditional_bounds.cc b/src/s_tir/analysis/conditional_bounds.cc index 9221bb8414f2..aeccf4e9c01c 100644 --- a/src/s_tir/analysis/conditional_bounds.cc +++ b/src/s_tir/analysis/conditional_bounds.cc @@ -34,6 +34,7 @@ #include #include +#include #include #include "../../sym/int_operator.h" @@ -614,6 +615,66 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) { #pragma optimize("g", on) #endif +std::optional> GetUnsignedRange(const PrimExpr& e) { + auto match = [&e](const auto* op) -> std::optional> { + if (!op) return std::nullopt; + for (bool reverse : {false, true}) { + PrimExpr value = reverse ? op->b : op->a; + PrimExpr bound = reverse ? op->a : op->b; + const auto* var = value.as(); + PrimType dtype = value.ty(); + // Only direct comparisons are safe; unsigned arithmetic may wrap. + if (!var || !dtype.IsScalar() || !dtype.MatchesCode(DLDataTypeCode::kDLUInt) || + dtype.bits() > 64) { + continue; + } + const auto* constant = bound.as(); + if (!constant || constant->value < 0) continue; + const ffi::BigInt& c = constant->value; + ffi::BigInt maximum = (ffi::BigInt(1) << dtype.bits()) - 1; + ffi::BigInt lower = 0, upper = maximum; + if (e->IsInstance()) { + lower = upper = c; + } else if (e->IsInstance()) { + // Only an excluded endpoint can be represented by a single interval. + if (c == 0) { + lower = 1; + } else if (c == maximum) { + upper = maximum - 1; + } else { + return std::nullopt; + } + } else { + bool is_lower = e->IsInstance() || e->IsInstance(); + bool strict = e->IsInstance() || e->IsInstance(); + if (reverse) is_lower = !is_lower; + // Leave impossible endpoint comparisons unresolved, rather than wrap. + if (strict && ((is_lower && c == maximum) || (!is_lower && c == 0))) { + return std::nullopt; + } + if (is_lower) { + lower = c + strict; + } else { + upper = c - strict; + } + } + // The full type domain has no representable unsigned extent and adds no bound. + if (lower == 0 && upper == maximum) return std::nullopt; + return std::make_pair(ffi::GetRef(var), + Range::FromMinExtent(prim::MakeConst(dtype, lower), + prim::MakeConst(dtype, upper - lower + 1))); + } + return std::nullopt; + }; + if (const auto* op = e.as()) return match(op); + if (const auto* op = e.as()) return match(op); + if (const auto* op = e.as()) return match(op); + if (const auto* op = e.as()) return match(op); + if (const auto* op = e.as()) return match(op); + if (const auto* op = e.as()) return match(op); + return std::nullopt; +} + } // namespace ffi::Optional> ConditionalBoundsContext::TrySolveCondition() { @@ -626,10 +687,15 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( } ffi::Array equations; ffi::Array vars; - std::function fvisit = [&equations, &vars, &fvisit](const PrimExpr& e) { + std::vector> unsigned_ranges; + std::function fvisit = [&](const PrimExpr& e) { if (e->IsInstance() || e->IsInstance() || e->IsInstance() || e->IsInstance() || e->IsInstance() || e->IsInstance()) { + if (auto bound = GetUnsignedRange(e)) { + unsigned_ranges.push_back(*bound); + return; + } bool is_simple = true; std::vector cand_vars; auto walk_fn = [&cand_vars, &is_simple, @@ -638,8 +704,13 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( return ffi::WalkResult::Advance(); } else if (const VarNode* var = obj.as()) { PrimType var_ty = var->ty.as_or_throw(); - if (var_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) { + if (var_ty.MatchesCode(DLDataTypeCode::kDLInt)) { cand_vars.push_back(ffi::GetRef(var).as_or_throw()); + } else { + // The inequality solver constructs signed coefficients in the + // variable's type. Unsigned arithmetic cannot be treated as + // ordered integer arithmetic; leave such conditions unresolved. + is_simple = false; } } else { is_simple &= obj->IsInstance() || obj->IsInstance() || @@ -670,7 +741,7 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( } }; fvisit(condition); - if (equations.empty() || vars.empty()) { + if (equations.empty() && unsigned_ranges.empty()) { return std::nullopt; } // build dom ranges for related vars @@ -690,13 +761,31 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( ranges.Set(v, Range::FromMinExtent(dom.min(), analyzer->Simplify(dom.max() - dom.min() + 1))); } } - // solve constraints + // Keep unsigned comparisons out of signed-coefficient elimination. IntConstraints constraint(vars, ranges, equations); - IntConstraints result = SolveInequalitiesToRange(constraint); - if (!result.relations.empty()) { - return std::nullopt; + IntConstraints result = vars.empty() ? constraint : SolveInequalitiesToRange(constraint); + if (result.relations.empty()) { + ranges = result.ranges; + } else { + ranges.clear(); + } + // Intersect the static unsigned bounds collected during condition traversal. + for (auto [var, range] : unsigned_ranges) { + if (auto previous = ranges.Get(var)) { + const ffi::BigInt& min = range->min.as_or_throw()->value; + const ffi::BigInt& extent = range->extent.as_or_throw()->value; + const ffi::BigInt& previous_min = previous.value()->min.as_or_throw()->value; + const ffi::BigInt& previous_extent = previous.value()->extent.as_or_throw()->value; + ffi::BigInt lower = std::max(min, previous_min); + ffi::BigInt upper = std::min(min + extent - 1, previous_min + previous_extent - 1); + if (lower > upper) return std::nullopt; + range = Range::FromMinExtent(prim::MakeConst(range->min.ty(), lower), + prim::MakeConst(range->min.ty(), upper - lower + 1)); + } + ranges.Set(var, range); } - return result.ranges; + if (ranges.empty()) return std::nullopt; + return ranges; } ConditionalBoundsContext::ConditionalBoundsContext( @@ -719,7 +808,18 @@ void ConditionalBoundsContext::EnterWithScope() { // update solved var ranges for (const auto& kv : constraints.value()) { const VarNode* var = kv.first.get(); - sym::IntSet new_dom = sym::IntSet::FromRange(kv.second); + sym::IntSet new_dom; + if (var->ty.as_or_throw().MatchesCode(DLDataTypeCode::kDLUInt)) { + // These static ranges are nonempty. Compute the endpoint in BigInt + // before constructing the unsigned expression to avoid wraparound. + const ffi::BigInt& min = kv.second->min.as_or_throw()->value; + const ffi::BigInt& extent = kv.second->extent.as_or_throw()->value; + new_dom = sym::IntSet::Interval( + kv.second->min, + extent == 1 ? kv.second->min : prim::MakeConst(kv.second->min.ty(), min + (extent - 1))); + } else { + new_dom = sym::IntSet::FromRange(kv.second); + } auto relax_it = relax_map_->find(var); if (relax_it != relax_map_->end()) { // this is a bound for relaxed var diff --git a/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py index d42ab45ea397..f3ac0154f1b6 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. # ruff: noqa: E501 +import pytest + import tvm import tvm.testing from tvm import s_tir, tirx @@ -1430,5 +1432,25 @@ def expected(p_output0: T.handle, n: T.int32): ] +@pytest.mark.parametrize("dtype, value", [("uint32", 0), ("uint64", 2**63), ("uint64", 2**64 - 1)]) +def test_unsigned_condition(dtype, value): + x = tirx.Var("x", dtype) + # Both branches must handle unsigned constants, including BigInt values. + func = tirx.PrimFunc([x], tirx.Evaluate(tirx.if_then_else(x != tirx.const(value, dtype), 1, 0))) + before = tvm.IRModule.from_expr(func) + # Exercise ConditionalBoundsContext without any buffer accesses. + after = s_tir.transform.CompactBufferAllocation()(before) + tvm.ir.assert_structural_equal(after, before) + + +def test_unsigned_wraparound_condition(): + x = tirx.Var("x", "uint32") + # Do not treat modular arithmetic as a signed linear inequality. + func = tirx.PrimFunc([x], tirx.Evaluate(tirx.if_then_else(x + 1 < x, 1, 0))) + before = tvm.IRModule.from_expr(func) + after = s_tir.transform.CompactBufferAllocation()(before) + tvm.ir.assert_structural_equal(after, before) + + if __name__ == "__main__": tvm.testing.main()