Skip to content

Fix numel constraints for invariant extents - #3194

Merged
jansel merged 2 commits into
mainfrom
helion-autotune-fix
Jul 31, 2026
Merged

Fix numel constraints for invariant extents#3194
jansel merged 2 commits into
mainfrom
helion-autotune-fix

Conversation

@ethche

@ethche ethche commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Fix tensor-numel constraint extraction for kernels whose register tiles include fixed or persistent reduction dimensions.

Problem

This prevents Helion’s default and randomly generated autotune configurations from exceeding Triton’s hard per-tensor limit of 1,048,576 elements. The same constraint handling applies to TileIR, which uses Triton codegen and inherits the same tensor size restrictions. Issue was observed in: #3176

When a kernel tiles over block-count dimensions alongside fixed or persistent reduction dimensions, each tile point carries a fixed register footprint:

Tensor shape: x_blk = [block_cb, 4, 32, block_nb, 128]
Total numel: block_cb * block_nb * 16,384

Because Triton limits each tensor to 1,048,576 elements, valid configurations must satisfy:
block_cb * block_nb <= 64

In the failing environment, setting autotune_effort="none" selected [32, 32], resulting in 16,777,216 elements and triggering a compilation failure.

The previous numel extractor missed this constraint because it only emitted constraints when every symbolic dimension mapped directly to a tunable block size. The non-tunable dimensions (4, 32, and 128) were represented as reduction block symbols rather than tunable block sizes, causing the entire constraint expression to be discarded. While newer heuristics might happen to select [1, 1] and mask the issue, the underlying omission remained a bug for default and randomly generated autotune candidates.

Fix

For Triton-codegen backends, the constraint extractor now resolves non-tunable dimensions when their per-program extent is invariant across the generated search space:

  • FixedBlockSizeSource dimensions with concrete extents.
  • Reduction dimensions with no loop spec.
  • Reduction dimensions whose entire generated search range decodes to persistent execution.

Genuinely rollable reductions remain symbolic, and the extractor conservatively skips those expressions rather than assuming an incorrect extent. Reduction-loop specs are matched by block ID rather than allocation ordinal because only rollable reductions receive configuration entries.

Once fixed dimensions are substituted, the constraint machinery sees: 16,384 * block_cb * block_nb <= 1,048,576

The existing default/random-config shrinker then properly enforces a maximum block product of 64. This check targets codegen_name == "triton", correctly covering both Triton and TileIR without altering behavior for other backends.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 28, 2026
@ethche
ethche requested a review from jansel July 29, 2026 01:53
@ethche
ethche requested a review from yushangdi July 29, 2026 18:59
@jansel
jansel merged commit 22d2e95 into main Jul 31, 2026
24 checks passed
@jansel
jansel deleted the helion-autotune-fix branch July 31, 2026 18:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants