From bb3e8693fce32da9160265cb061fa7a0576efdbf Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Tue, 15 Sep 2026 21:16:30 +0800 Subject: [PATCH 1/3] [FIX][TIR] Conservatively analyze unsigned branch conditions --- src/tirx/transform/ir_utils.cc | 7 +- .../analysis/test_sblock_access_region.py | 87 +++++++++++++++++++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 7c5c717caa9f..3c31a6ec6e89 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -807,8 +807,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() || diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py index b9266be09ae4..d19438234f7a 100644 --- a/tests/python/s_tir/analysis/test_sblock_access_region.py +++ b/tests/python/s_tir/analysis/test_sblock_access_region.py @@ -457,5 +457,92 @@ def func( tvm.ir.assert_structural_equal(block.writes, ret[1]) +@pytest.mark.parametrize("dtype", ["int32", "uint8", "uint32", "uint64", "bool"]) +@pytest.mark.parametrize("equal", [False, True]) +def test_conditional_region_zero_comparison(dtype, equal): + """The implicit else condition must not introduce unsigned coefficients.""" + tir = tvm.tirx + mask_buffer = tir.decl_buffer((1,), dtype, name="mask_buffer") + data = tir.decl_buffer((4,), "int32", name="data") + output = tir.decl_buffer((1,), "int32", name="output") + mask = tir.Var("mask", dtype) + condition = mask == 0 if equal else mask != 0 + body = tir.SeqStmt( + [ + tir.Bind(mask, mask_buffer[0]), + tir.BufferStore(output, tir.if_then_else(condition, data[0], data[3]), [0]), + ] + ) + block = tir.SBlock([], [], [], "conditional_read", body) + buffers = {buf: buf for buf in (mask_buffer, data, output)} + for analyze in ( + s_tir.analysis.get_sblock_access_region, + s_tir.analysis.get_sblock_read_write_region, + ): + regions = analyze(block, buffers) + reads, writes = regions[:2] + data_region = next(region for region in reads if region.buffer.same_as(data)) + tvm.ir.assert_structural_equal(data_region.region, [Range(0, 4)]) + tvm.ir.assert_structural_equal(writes, [tir.BufferRegion(output, [Range(0, 1)])]) + + +@pytest.mark.parametrize("dtype", ["uint32", "uint64"]) +@pytest.mark.parametrize("condition_kind", ["upper", "high_bit", "wraparound"]) +def test_unsigned_conditional_region_preserves_both_branches(dtype, condition_kind): + """Fallback must not prune accesses using signed or non-wrapping semantics.""" + tir = tvm.tirx + mask_buffer = tir.decl_buffer((1,), dtype, name="mask_buffer") + data = tir.decl_buffer((4,), "int32", name="data") + output = tir.decl_buffer((1,), "int32", name="output") + mask = tir.Var("mask", dtype) + if condition_kind == "upper": + condition = mask < 8 + elif condition_kind == "high_bit": + condition = mask >= tir.const(1 << (int(dtype[4:]) - 1), dtype) + else: + condition = mask + tir.const(1, dtype) < mask + body = tir.SeqStmt( + [ + tir.Bind(mask, mask_buffer[0]), + tir.BufferStore(output, tir.if_then_else(condition, data[0], data[3]), [0]), + ] + ) + block = tir.SBlock([], [], [], "unsigned_read", body) + buffers = {buf: buf for buf in (mask_buffer, data, output)} + reads, _, _ = s_tir.analysis.get_sblock_access_region(block, buffers) + data_region = next(region for region in reads if region.buffer.same_as(data)) + tvm.ir.assert_structural_equal(data_region.region, [Range(0, 4)]) + + +@pytest.mark.parametrize("dtype", ["uint32", "uint64"]) +def test_unsigned_condition_keeps_independent_signed_bound(dtype): + tir = tvm.tirx + mask_buffer = tir.decl_buffer((1,), dtype, name="mask_buffer") + data = tir.decl_buffer((8,), "int32", name="data") + output = tir.decl_buffer((8,), "int32", name="output") + mask = tir.Var("mask", dtype) + i = tir.Var("i", "int32") + body = tir.SeqStmt( + [ + tir.Bind(mask, mask_buffer[0]), + tir.For( + i, + 0, + 8, + tir.ForKind.SERIAL, + tir.IfThenElse( + tir.And(i < 4, mask == 0), tir.BufferStore(output, data[i], [i]), None + ), + ), + ] + ) + block = tir.SBlock([], [], [], "mixed_condition", body) + buffers = {buf: buf for buf in (mask_buffer, data, output)} + reads, writes, _ = s_tir.analysis.get_sblock_access_region(block, buffers) + data_region = next(region for region in reads if region.buffer.same_as(data)) + tvm.ir.assert_structural_equal(data_region.region, [Range(0, 4)]) + tvm.ir.assert_structural_equal(writes, [tir.BufferRegion(output, [Range(0, 4)])]) + + if __name__ == "__main__": tvm.testing.main() From d6b8e75273c670b27222dd69d916a4f7f6eb75c2 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:46:27 +0800 Subject: [PATCH 2/3] [TEST][TIR] Minimize unsigned conditional regression in Python --- .../analysis/test_sblock_access_region.py | 87 ------------------- ...t_s_tir_transform_compact_buffer_region.py | 9 ++ 2 files changed, 9 insertions(+), 87 deletions(-) diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py index d19438234f7a..b9266be09ae4 100644 --- a/tests/python/s_tir/analysis/test_sblock_access_region.py +++ b/tests/python/s_tir/analysis/test_sblock_access_region.py @@ -457,92 +457,5 @@ def func( tvm.ir.assert_structural_equal(block.writes, ret[1]) -@pytest.mark.parametrize("dtype", ["int32", "uint8", "uint32", "uint64", "bool"]) -@pytest.mark.parametrize("equal", [False, True]) -def test_conditional_region_zero_comparison(dtype, equal): - """The implicit else condition must not introduce unsigned coefficients.""" - tir = tvm.tirx - mask_buffer = tir.decl_buffer((1,), dtype, name="mask_buffer") - data = tir.decl_buffer((4,), "int32", name="data") - output = tir.decl_buffer((1,), "int32", name="output") - mask = tir.Var("mask", dtype) - condition = mask == 0 if equal else mask != 0 - body = tir.SeqStmt( - [ - tir.Bind(mask, mask_buffer[0]), - tir.BufferStore(output, tir.if_then_else(condition, data[0], data[3]), [0]), - ] - ) - block = tir.SBlock([], [], [], "conditional_read", body) - buffers = {buf: buf for buf in (mask_buffer, data, output)} - for analyze in ( - s_tir.analysis.get_sblock_access_region, - s_tir.analysis.get_sblock_read_write_region, - ): - regions = analyze(block, buffers) - reads, writes = regions[:2] - data_region = next(region for region in reads if region.buffer.same_as(data)) - tvm.ir.assert_structural_equal(data_region.region, [Range(0, 4)]) - tvm.ir.assert_structural_equal(writes, [tir.BufferRegion(output, [Range(0, 1)])]) - - -@pytest.mark.parametrize("dtype", ["uint32", "uint64"]) -@pytest.mark.parametrize("condition_kind", ["upper", "high_bit", "wraparound"]) -def test_unsigned_conditional_region_preserves_both_branches(dtype, condition_kind): - """Fallback must not prune accesses using signed or non-wrapping semantics.""" - tir = tvm.tirx - mask_buffer = tir.decl_buffer((1,), dtype, name="mask_buffer") - data = tir.decl_buffer((4,), "int32", name="data") - output = tir.decl_buffer((1,), "int32", name="output") - mask = tir.Var("mask", dtype) - if condition_kind == "upper": - condition = mask < 8 - elif condition_kind == "high_bit": - condition = mask >= tir.const(1 << (int(dtype[4:]) - 1), dtype) - else: - condition = mask + tir.const(1, dtype) < mask - body = tir.SeqStmt( - [ - tir.Bind(mask, mask_buffer[0]), - tir.BufferStore(output, tir.if_then_else(condition, data[0], data[3]), [0]), - ] - ) - block = tir.SBlock([], [], [], "unsigned_read", body) - buffers = {buf: buf for buf in (mask_buffer, data, output)} - reads, _, _ = s_tir.analysis.get_sblock_access_region(block, buffers) - data_region = next(region for region in reads if region.buffer.same_as(data)) - tvm.ir.assert_structural_equal(data_region.region, [Range(0, 4)]) - - -@pytest.mark.parametrize("dtype", ["uint32", "uint64"]) -def test_unsigned_condition_keeps_independent_signed_bound(dtype): - tir = tvm.tirx - mask_buffer = tir.decl_buffer((1,), dtype, name="mask_buffer") - data = tir.decl_buffer((8,), "int32", name="data") - output = tir.decl_buffer((8,), "int32", name="output") - mask = tir.Var("mask", dtype) - i = tir.Var("i", "int32") - body = tir.SeqStmt( - [ - tir.Bind(mask, mask_buffer[0]), - tir.For( - i, - 0, - 8, - tir.ForKind.SERIAL, - tir.IfThenElse( - tir.And(i < 4, mask == 0), tir.BufferStore(output, data[i], [i]), None - ), - ), - ] - ) - block = tir.SBlock([], [], [], "mixed_condition", body) - buffers = {buf: buf for buf in (mask_buffer, data, output)} - reads, writes, _ = s_tir.analysis.get_sblock_access_region(block, buffers) - data_region = next(region for region in reads if region.buffer.same_as(data)) - tvm.ir.assert_structural_equal(data_region.region, [Range(0, 4)]) - tvm.ir.assert_structural_equal(writes, [tir.BufferRegion(output, [Range(0, 4)])]) - - if __name__ == "__main__": tvm.testing.main() 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() From a543f2b695c2772a78e2876346919f1fbb267c11 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Wed, 16 Sep 2026 02:12:25 +0800 Subject: [PATCH 3/3] [FIX][TIR] Extract unsigned equality bounds during condition traversal --- src/tirx/transform/ir_utils.cc | 133 +++++++++++++++++++++++++++++++-- 1 file changed, 125 insertions(+), 8 deletions(-) diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 3c31a6ec6e89..b01b287f09e5 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -785,6 +785,82 @@ Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region return result; } +namespace { + +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(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(MakeConst(dtype, lower), 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() { // extract equations and related vars from condition expression. // currently only extract simple integral equations which could be solvable. @@ -799,6 +875,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, @@ -844,7 +924,7 @@ ffi::Optional ConditionalBoundsContext::TrySolveCondition } }; fvisit(condition); - if (equations.empty() || vars.empty()) { + if (equations.empty()) { return std::nullopt; } // build dom ranges for related vars @@ -864,13 +944,39 @@ ffi::Optional ConditionalBoundsContext::TrySolveCondition ranges.Set(v, Range::FromMinExtent(dom.min(), analyzer->Simplify(dom.max() - dom.min() + 1))); } } - // solve constraints - arith::IntConstraints constraint(vars, ranges, equations); - arith::IntConstraints result = arith::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; + arith::IntConstraints constraint(vars, ranges, signed_equations); + arith::IntConstraints result = + vars.empty() ? constraint : arith::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(MakeConst(range->min.ty(), lower), + MakeConst(range->min.ty(), upper - lower + 1)); + } + ranges.Set(var, range); + } + } + if (ranges.empty()) return std::nullopt; + return arith::IntConstraints(vars, ranges, {}); } ConditionalBoundsContext::ConditionalBoundsContext( @@ -893,7 +999,18 @@ void ConditionalBoundsContext::EnterWithScope() { // update solved var ranges for (const auto& kv : constraints.value()->ranges) { 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 : 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