Skip to content

Reject data-dependent output shapes in infer_output_spec#1722

Merged
gmagogsfm merged 1 commit intomainfrom
reject-data-dependent-output-shapes
Mar 17, 2026
Merged

Reject data-dependent output shapes in infer_output_spec#1722
gmagogsfm merged 1 commit intomainfrom
reject-data-dependent-output-shapes

Conversation

@gmagogsfm
Copy link
Copy Markdown
Contributor

When a helion kernel's output contains an unbacked SymInt (indicating a data-dependent shape) that is produced inside Helion Kernel (via .item() calls), raise DataDependentOutputShapeNotSupported.

This is a defensive check: currently, type propagation catches known data-dependent ops (nonzero, boolean indexing) before infer_output_spec runs, making this path unreachable in practice. It is left untested for this reason, but serves as a safety net for future ops that may produce unbacked SymInts without a corresponding type propagation check.

When a helion kernel's output contains an unbacked SymInt (indicating a
data-dependent shape), raise DataDependentOutputShapeNotSupported in
_remap_or_resolve instead of silently concretizing via size_hint.

This is a defensive check: currently, type propagation catches known
data-dependent ops (nonzero, boolean indexing) before infer_output_spec
runs, making this path unreachable in practice. It is left untested for
this reason, but serves as a safety net for future ops that may produce
unbacked SymInts without a corresponding type propagation check.

Also adds test coverage for infer_output_spec symbol remapping with
single and multi-output kernels.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
@gmagogsfm gmagogsfm requested a review from yf225 March 16, 2026 17:49
@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 16, 2026
@gmagogsfm gmagogsfm merged commit 3b236ad into main Mar 17, 2026
45 of 47 checks passed
yf225 added a commit that referenced this pull request Mar 17, 2026
The previous code from PR #1722 treated ALL unbacked symbols as
data-dependent, but compound expressions like `flag * 2` (= 2*u0)
have unbacked symbols that come from known kernel parameters.

Fix: check whether unbacked symbols are a subset of sym_remap keys.
If yes, substitute concrete values and evaluate. If any symbol is
unknown, raise DataDependentOutputShapeNotSupported.

Un-skip 3 tests that were broken by the overly broad check:
test_kernel_returns_none_in_tuple,
test_kernel_returns_none_first_in_tuple,
test_kernel_returns_tuple_of_scalars.
yf225 added a commit that referenced this pull request Mar 19, 2026
The previous code from PR #1722 treated ALL unbacked symbols as
data-dependent, but compound expressions like `flag * 2` (= 2*u0)
have unbacked symbols that come from known kernel parameters.

Fix: check whether unbacked symbols are a subset of sym_remap keys.
If yes, substitute concrete values and evaluate. If any symbol is
unknown, raise DataDependentOutputShapeNotSupported.

Un-skip 3 tests that were broken by the overly broad check:
test_kernel_returns_none_in_tuple,
test_kernel_returns_none_first_in_tuple,
test_kernel_returns_tuple_of_scalars.
yf225 added a commit that referenced this pull request Mar 19, 2026
The previous code from PR #1722 treated ALL unbacked symbols as
data-dependent, but compound expressions like `flag * 2` (= 2*u0)
have unbacked symbols that come from known kernel parameters.

Fix: check whether unbacked symbols are a subset of sym_remap keys.
If yes, substitute concrete values and evaluate. If any symbol is
unknown, raise DataDependentOutputShapeNotSupported.

Un-skip 3 tests that were broken by the overly broad check:
test_kernel_returns_none_in_tuple,
test_kernel_returns_none_first_in_tuple,
test_kernel_returns_tuple_of_scalars.
yf225 added a commit that referenced this pull request Mar 20, 2026
The previous code from PR #1722 treated ALL unbacked symbols as
data-dependent, but compound expressions like `flag * 2` (= 2*u0)
have unbacked symbols that come from known kernel parameters.

Fix: check whether unbacked symbols are a subset of sym_remap keys.
If yes, substitute concrete values and evaluate. If any symbol is
unknown, raise DataDependentOutputShapeNotSupported.

Un-skip 3 tests that were broken by the overly broad check:
test_kernel_returns_none_in_tuple,
test_kernel_returns_none_first_in_tuple,
test_kernel_returns_tuple_of_scalars.
hinriksnaer pushed a commit to hinriksnaer/helion that referenced this pull request Mar 20, 2026
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
yf225 added a commit that referenced this pull request Mar 23, 2026
The previous code from PR #1722 treated ALL unbacked symbols as
data-dependent, but compound expressions like `flag * 2` (= 2*u0)
have unbacked symbols that come from known kernel parameters.

Fix: check whether unbacked symbols are a subset of sym_remap keys.
If yes, substitute concrete values and evaluate. If any symbol is
unknown, raise DataDependentOutputShapeNotSupported.

Un-skip 3 tests that were broken by the overly broad check:
test_kernel_returns_none_in_tuple,
test_kernel_returns_none_first_in_tuple,
test_kernel_returns_tuple_of_scalars.
yf225 added a commit that referenced this pull request Mar 23, 2026
The previous code from PR #1722 treated ALL unbacked symbols as
data-dependent, but compound expressions like `flag * 2` (= 2*u0)
have unbacked symbols that come from known kernel parameters.

Fix: check whether unbacked symbols are a subset of sym_remap keys.
If yes, substitute concrete values and evaluate. If any symbol is
unknown, raise DataDependentOutputShapeNotSupported.

Un-skip 3 tests that were broken by the overly broad check:
test_kernel_returns_none_in_tuple,
test_kernel_returns_none_first_in_tuple,
test_kernel_returns_tuple_of_scalars.
umechand-amd pushed a commit to umechand-amd/helion that referenced this pull request Mar 23, 2026
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
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.

2 participants