diff --git a/src/s_tir/analysis/conditional_bounds.cc b/src/s_tir/analysis/conditional_bounds.cc index 07c5570c1b17..ed487b1c4f92 100644 --- a/src/s_tir/analysis/conditional_bounds.cc +++ b/src/s_tir/analysis/conditional_bounds.cc @@ -32,7 +32,9 @@ #include #include +#include #include +#include #include #include "../../arith/int_operator.h" @@ -611,6 +613,78 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) { #pragma optimize("g", on) #endif +std::optional GetConstUInt(const PrimExpr& value) { + if (const auto* imm = value.as()) { + if (imm->value >= 0) return static_cast(imm->value); + } else if (const auto* call = value.as()) { + if (call->op.same_as(tirx::builtin::large_uint_imm())) { + return static_cast(call->args[0].as_or_throw()->value) | + (static_cast(call->args[1].as_or_throw()->value) << 32); + } + } + return std::nullopt; +} + +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; + } + auto constant = GetConstUInt(bound); + if (!constant) continue; + uint64_t c = *constant; + uint64_t maximum = UINT64_MAX >> (64 - dtype.bits()); + uint64_t 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() { @@ -627,6 +701,10 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( if (e->IsInstance() || e->IsInstance() || e->IsInstance() || e->IsInstance() || e->IsInstance() || e->IsInstance()) { + if (GetUnsignedRange(e)) { + equations.push_back(e); + return; + } bool is_simple = true; std::vector cand_vars; auto walk_fn = [&cand_vars, &is_simple, @@ -635,8 +713,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() || @@ -667,7 +750,7 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( } }; fvisit(condition); - if (equations.empty() || vars.empty()) { + if (equations.empty()) { return std::nullopt; } // build dom ranges for related vars @@ -687,13 +770,38 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( ranges.Set(v, Range::FromMinExtent(dom.min(), analyzer->Simplify(dom.max() - dom.min() + 1))); } } - // solve constraints - IntConstraints constraint(vars, ranges, equations); - IntConstraints result = SolveInequalitiesToRange(constraint); - if (!result.relations.empty()) { - return std::nullopt; + // Keep unsigned comparisons out of signed-coefficient elimination. + ffi::Array signed_equations; + for (const PrimExpr& e : equations) { + if (!GetUnsignedRange(e)) signed_equations.push_back(e); } - return result.ranges; + IntConstraints constraint(vars, ranges, signed_equations); + IntConstraints result = vars.empty() ? constraint : SolveInequalitiesToRange(constraint); + if (result.relations.empty()) { + ranges = result.ranges; + } else { + ranges.clear(); + } + // Reuse the same range map for directly solved unsigned comparisons. + for (const PrimExpr& e : equations) { + if (auto bound = GetUnsignedRange(e)) { + auto [var, range] = *bound; + if (auto previous = ranges.Get(var)) { + uint64_t min = GetConstUInt(range->min).value(); + uint64_t extent = GetConstUInt(range->extent).value(); + uint64_t previous_min = GetConstUInt(previous.value()->min).value(); + uint64_t previous_extent = GetConstUInt(previous.value()->extent).value(); + uint64_t lower = std::max(min, previous_min); + uint64_t 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); + } + } + if (ranges.empty()) return std::nullopt; + return ranges; } ConditionalBoundsContext::ConditionalBoundsContext( @@ -716,7 +824,18 @@ void ConditionalBoundsContext::EnterWithScope() { // update solved var ranges for (const auto& kv : constraints.value()) { const VarNode* var = kv.first.get(); - arith::IntSet new_dom = arith::IntSet::FromRange(kv.second); + arith::IntSet new_dom; + if (var->ty.as_or_throw().MatchesCode(DLDataTypeCode::kDLUInt)) { + // These static ranges are nonempty. Compute the endpoint without unsigned + // wraparound or signed-int64 constant folding in IntSet::FromRange. + uint64_t min = GetConstUInt(kv.second->min).value(); + uint64_t extent = GetConstUInt(kv.second->extent).value(); + new_dom = arith::IntSet::Interval( + kv.second->min, + extent == 1 ? kv.second->min : prim::MakeConst(kv.second->min.ty(), min + (extent - 1))); + } else { + new_dom = arith::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 f0000581e1f1..f3ba28c151a9 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 @@ -1430,5 +1430,14 @@ def expected(p_output0: T.handle, n: T.int32): ] +def test_unsigned_condition(): + x = tirx.Var("x", "uint32") + func = tirx.PrimFunc([x], tirx.Evaluate(tirx.if_then_else(x != 0, 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) + + if __name__ == "__main__": tvm.testing.main()