Skip to content

[Arith] Add opt-in lazily evaluated bounds for bound variables - #72

Open
LeiWang1999 wants to merge 1 commit into
tilelang_mainfrom
wl/lazy-bind-bounds
Open

LeiWang1999 wants to merge 1 commit into
tilelang_mainfrom
wl/lazy-bind-bounds

Conversation

@LeiWang1999

Copy link
Copy Markdown
Member

Implements the scoped variant of option 3 from tile-ai/tilelang#3220.

Problem

Analyzer::Bind eagerly snapshots the value's integer bounds at bind time and never re-evaluates them:

a.Bind(tx, Range(0, 256));
a.Bind(v, tx);                          // snapshot: v ∈ [0,255]
With<ConstraintContext> ctx(&a, tx < 64);
a.const_int_bound(tx);                  // [0,63]  — on demand, fresh
a.const_int_bound(v);                   // [0,255] — bind-time snapshot, stale

v ≡ tx, yet their bounds disagree under the same constraint stack. Any consumer that replays collected binds/constraints in an order unrelated to program order (merged constraint sets in tile-ai/tilelang#3220) can therefore lose tightening permanently, and every ordering workaround downstream is a symptom of this premature evaluation.

Change

Analyzer::BindLazyBounds(var, expr, allow_override) — identical to Bind for all sub-analyzers except const_int_bound, which stores the definition (ConstIntBoundAnalyzer::UpdateDefinition) instead of a snapshot:

  • the bound is derived on demand under the constraints active at query time, memoized per constraint epoch (a counter bumped on EnterConstraint entry and exit), so bounds tighten and relax correctly across nested scopes;
  • definition chains (w = v * 2 with v itself lazily bound) recurse, with a cycle guard degrading to Everything;
  • results whose subtree consulted a lazy definition are excluded from the caller-provided memo table (operator()(expr, bound)), whose consistency ICHECK assumes constraint-independent entries;
  • IsBound reports definition-backed vars; Clone/CopyFrom carries them.

Opt-in only: eager Bind is untouched and remains the default; nothing inside this repository opts in, so behavior and performance are unchanged unless a caller explicitly asks. Exposed to Python as Analyzer.bind_lazy_bounds via the analyzer factory.

Validation

  • New tests/python/arith/test_arith_lazy_bind.py (7 tests): documents the eager staleness, then pins lazy behavior — later constraints visible, bind-order insensitivity, definition chains, memo invalidation across nested scopes (tighten → tighter → relax), CanProve integration, override discipline.
  • Full tilelang build on top of this branch: constraint-machinery consumer suites (thread_sync, verify_parallel_loop, auto_schedule, analysis/, constr_set_merge) — 106 + 12 passed, zero behavior change with nothing opted in.

Notes for reviewers

  • One honest scoping observation: CanProve-routed queries are usually rescued from the stale snapshot by the rewrite path (the bound var is substituted by its definition before bounds are consulted), which is why [Transform][Arith] ConstrSet::Populate ordering is load-bearing in two incompatible directions (eager Analyzer::Bind snapshots) tilelang#3220 found upstream tests indistinguishable under reordering. The staleness bites consumers that query const_int_bound on bound vars directly — that is the surface this PR fixes. Switching tilelang's ConstrSet::Populate to the lazy mode is therefore optional and deferred until a workload demonstrates a need.
  • Scope is const_int_bound only; modular_set/int_set keep their eager snapshots (their staleness has no demonstrated consumer yet and each would need its own epoch memo). Can extend in a follow-up if wanted.

Analyzer::Bind eagerly snapshots the value's integer bounds at bind
time, so a bind replayed before the predicates that constrain its value
keeps the wide bound forever: v = tx inside 'if (tx < 64)' stays at
[0, 255] when the bind is entered first, even though tx itself tightens
on demand. Analyses that replay collected constraint sets in an order
unrelated to program order (tile-ai/tilelang#3220) hit exactly this.

Add Analyzer::BindLazyBounds: identical to Bind except const_int_bound
stores the definition instead of a snapshot. The bound is derived on
demand under the constraints active at query time and memoized per
constraint epoch (bumped on EnterConstraint entry/exit); definition
chains recurse with a cycle guard. Results whose subtree consulted a
lazy definition stay out of the caller-provided memo table, whose
consistency check assumes constraint-independent entries.

Eager Bind is untouched and remains the default everywhere; nothing
opts in inside this repository, so behavior is unchanged unless a
caller asks for the new mode.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant