Speed up bounds inference for funcs with many pure update dims

[abadams]

Problem

Computation bounds inference dominates lowering time for the bgu app. Its solve is expressed as a Func with ~100 update stages (unrolled Cholesky), but almost every dimension (x, y, z) is pure in all stages — only the two matrix-index dimensions vary per update. On bgu, computation bounds inference alone took ~381ms out of ~578ms total lowering.

Cause

Stage::populate_scope creates a distinct per-stage bound variable (f.sK.x.min/.max) for every dimension, including dimensions that are pure in every stage. When a many-staged Func is a consumer, merge_boxes then has to combine one structurally distinct term per stage for those pure dimensions, building large nested min/max expressions and calling simplify on each — O(stages) work per pure dimension, per consumed producer.

define_bounds already contained the observation that a dimension pure in every update has a single bound shared by all stages (it aliases non-final stages to the last stage), but it did so only after the expensive merge had already built and simplified the per-stage expressions, then threw that work away.

Fix

Key the required region of an always-pure dimension off the last stage's bound variables directly in populate_scope. The per-stage boxes are then structurally identical in that dimension, so they collapse on merge instead of growing one term per stage. always_pure_dims is computed once per Func (shared across its stages) rather than recomputed inside define_bounds.

Impact

On bgu (measured with HL_TIME_LOWERING_PASSES=1):

	before	after
computation bounds inference	~381 ms	~23 ms
total lowering	~578 ms	~210 ms

Generated code is unchanged.

Testing

make test_correctness passes except for the pre-existing simd_op_check_sve2 failure (an unrelated LLVM 23 backend crash on this tree). No new regression test is included since this is a compile-time optimization with no change to generated code; the existing suite covers correctness of multi-stage / pure-dimension bounds inference.

🤖 Generated with Claude Code

[mcourteaux]

This PR makes correctness_gpu_mixed_dimensionality stop terminating on any GPU target. It's visible in CI: buildbot/halide-llvm_21-x86-64-linux and -llvm_22 both failed after ~13h, and arm-64 / ci-llvm-main was cancelled with Test (host-metal) having run for 4h48m (57 of 60 tests completed; correctness_gpu_mixed_dimensionality started and never finished, leaving an orphaned correctness_gpu... process). Test (host) passed in 21 minutes on the same job, because the test [SKIP]s without a GPU target — so none of the x86 jobs can see it.

I bisected it to a single line, and I think the rest of the PR is fine. Details below, plus a fix that keeps the speedup.

The assumption that breaks

define_bounds already contained the observation that a dimension pure in every update has a single bound shared by all stages

That observation is true of the bounds required of a Func at root, but not at inner loop levels, and that's the gap.

In BoundsInference::visit(const For *), the block that consumes boxes_for_fused_group re-defines f.sK.<dim>.min/.max to the sub-box actually being provided at that point in the nest — and it does so only for the stage currently being produced (plus its fused partners). Sibling stages keep whatever value they had further out. So partway down stage 0's loop nest, f.s0.z.min/.max is the current tile while f.s1.z.min/.max is still the root-level extent.

The old code respected this. It aliased f.sk.<dim>.min = f.sN.<dim>.min inside define_bounds, which is emitted at a specific loop level, so both sides denote that level's bindings — and define_bounds is applied in increasing stage order, so f.sN's lets land outermost and the reference resolves. That aliasing is a statement about the func's own stages, made at a point where all of its per-stage bound lets are written together in one sweep. It's safe.

The new populate_scope makes a different claim: the region I require, phrased via my last stage's variables. populate_scope runs once in the constructor and is loop-level-agnostic, so that reference to f.sN.* escapes into the required-region expressions of every producer of stage k — and those get emitted wherever the producer is computed, including deep inside stage k's own nest, where f.sN.* was never narrowed.

Why this test in particular

gpu_mixed_dimensionality is built to hit exactly this. Every Func is f(x,y,z) = ...; f(x,y,z) += 1;, so all three dimensions are "always pure" and all get keyed off the last stage. And the schedule deliberately gives the two stages different loop nests:

out.gpu_tile(x, y, z, xi, yi, zi, 4, 4, 4);   // stage 0: 3D tiled
out.update().gpu_tile(x, y, xi, yi, 4, 4);    // stage 1: 2D only -- no z loop at all
h.compute_at(out, x).gpu_threads(x, y);       // h lives inside stage 0's nest

Inside out's stage-0 nest, out.s0.z.min/max is narrowed to the current 4-wide z tile, but out.s1.z.min/max is still 0..63 — stage 1 has no z loop to narrow it. So h's required region in z becomes 64 instead of 4, and it cascades through h → g → f.

Same pipeline, only that line differing, HL_TARGET=host-cuda:

	.maxntid	per-thread local depot
this PR	64, 64, 4	__local_depot0[2097152]
that line reverted	4, 4, 4	none

2097152 B is two 1 MiB buffers = 2 × 64³ × 4 B — g and f each allocated the whole volume per thread. 64×64×4 is also 16384 threads/block, well past the 1024 limit. It isn't an infinite loop; it drowns in redundant recompute and spilled local memory.

Suggested fix

The collapse is worth having — it just has to happen somewhere that knows the loop level. So populate_scope goes back to each stage's own bound variables, and define_bounds does the collapsing instead.

Each entry of a stage's bounds carries the region in both forms:

struct RequiredRegion {
    Box per_stage;   // consumer stage's own bound vars -- correct at any loop level
    Box canonical;   // always-pure dims phrased via the last stage's vars
};

and define_bounds picks per entry:

if (owns_loop_nest) {   // == producing_stage_index, or fused with it
    merge_boxes(b, i.second.per_stage);
} else if (inner_productions.count(func_name)) {
    merge_boxes(b, i.second.canonical.empty() ? i.second.per_stage : i.second.canonical);
}

The safety argument for the canonical branch: bounds_needed is set per Func, so if a consumer is produced further in, every one of its stages gets define_bounds called at this level, and the pre-existing aliasing is precisely what makes the two forms interchangeable there. In the owns_loop_nest case only that one stage's variables were narrowed, so the per-stage box is mandatory — that's the gpu_mixed_dimensionality case. Fused stages share the nest via boxes_for_fused_group, which also narrows only the group's stages, so they take the per-stage path too. The correct form is the default and canonical is opt-in, so getting the condition wrong costs compile time rather than correctness.

One implementation note: the last stage's bound variable Exprs are built once per Func in compute_always_pure_dims() and shared by all of its stages (which works because that runs on the stage-0 Stage before the update-stage copies are made). That makes merge_boxes' same_as pointer-identity fast path fire, which is why the numbers below come out a bit better than the PR's.

Numbers

bgu, HL_TIME_LOWERING_PASSES=1, LLVM 22, 3 runs each, run sequentially on one machine:

	computation bounds inference
58868f0c7 (pre-PR baseline)	745.8 ms
this PR as submitted	39.2 / 39.7 / 40.3 ms
with the fix	25.0 / 24.8 / 24.4 ms

Verification:

• correctness_gpu_mixed_dimensionality on host-cuda: timeout → Success!
• The lowered IR for that pipeline is byte-identical to the known-correct build, so generated code is unchanged as the PR claims
• ctest -L correctness on host: 469/469 passed
• All 469 correctness binaries under HL_JIT_TARGET=host-cuda (run directly, since ctest bakes HL_JIT_TARGET from configure-time Halide_TARGET): 451 pass, 18 skip, 0 failures or timeouts

No new regression test seemed warranted, since gpu_mixed_dimensionality already covers this exactly. The thing that might be worth a separate look is that this class of bug is only reachable from a GPU-target CI configuration, so it can sail through every x86 job green.

🤖 Generated with Claude Code

[mcourteaux]

@abadams @alexreinking Please review. This is a fully automated fix by Claude for the CI hang. This was Opus 5 on medium effort. I'm still not familiar enough with BoundsInference to be able to review this meaningfully. All CI should go green after this though, but perhaps there is a different more elegant fix you can come up with.

[codecov]

Codecov Report

❌ Patch coverage is 96.82540% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.01%. Comparing base (58868f0) to head (eae5a5d).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
src/BoundsInference.cpp	96.82%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #9442      +/-   ##
==========================================
+ Coverage   69.95%   70.01%   +0.06%     
==========================================
  Files         261      261              
  Lines       79596    79663      +67     
  Branches    19400    19421      +21     
==========================================
+ Hits        55678    55776      +98     
+ Misses      18004    17997       -7     
+ Partials     5914     5890      -24     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[abadams]

That's quite a complicated additional commit. I suspect there's a simpler approach. Removing Alex's approval so I don't absent-mindedly merge this.