[compiler] carry Python containers through dynamic if/for/while/ifexp#874
Open
xudoyuan wants to merge 20 commits into
Open
[compiler] carry Python containers through dynamic if/for/while/ifexp#874xudoyuan wants to merge 20 commits into
xudoyuan wants to merge 20 commits into
Conversation
xudoyuan
force-pushed
the
xudoyuan/dyn-cf-list-carry
branch
2 times, most recently
from
July 24, 2026 09:07
cea9d05 to
e6a56ab
Compare
xudoyuan
marked this pull request as ready for review
July 24, 2026 09:11
Dynamic scf.if / scf.for / scf.while / ifexp previously required every carried variable to be a single ir.Value, so a reassigned list/tuple/dict local raised "state variable '...' is list, not an MLIR Value". Route the four dispatchers through an explode/assemble step that reuses the existing DSL<->ir.Value protocol (protocol.extract/construct): each carried value (possibly a nested container) is exploded to per-element iter_args/ results and assembled back from the pre-region value used as the structural template on exit. Scalars stay the degenerate single-slot case so existing behaviour is unchanged; a branch/body that changes a container's shape or an element's dtype now fails with a clear error. Adds MLIR-level unit tests and device system tests for list/tuple/dict/nested carry across if/for/while/ifexp. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
The device/launch test belongs with the other end-to-end tests under tests/system (renamed to the *_e2e.py convention); the MLIR-level unit test stays under tests/unit. Aligns the header with the existing e2e tests (direct torch import + skip). Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Rename the test class to *E2E, reword the module docstring, and use the direct `import torch` + skip style shared by the other tests/system e2e files. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Replace the bespoke explode/assemble recursion in the dynamic control-flow dispatchers with the existing protocol.extract_to_ir_values / construct_from_ir_values, so there is a single source of truth for the DSL-value <-> flat ir.Value round trip. Removes the duplicate _explode_carried / _assemble_carried / _carried_slot_count helpers (and the now-dead _unwrap_mlir_values / _pack_dispatch_results); keeps only the three shared unpack/pack/verify helpers reused by if/for/while/ifexp. protocol.py: add dict handling to get_ir_types / extract_to_ir_values / construct_from_ir_values (parallel to the SimpleNamespace branch); make the construct list/tuple branch dispatch on the exemplar so a top-level container works while staying backward compatible with the type-sequence call; add a bare-scalar coercion fallback to extract so python scalars still promote. Net -96 lines in ast_rewriter.py. unit + struct + dispatch + device e2e green. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Add unit cases carrying pure dict, tuple and SimpleNamespace through a dynamic
scf.if, plus mixed nesting (tuple-of-dict, dict-of-list, SimpleNamespace of
list+tuple+scalar, and a 3-deep list[dict{tuple}, list]). Also carry a dict
through scf.for and a tuple-of-dict through scf.while. A small _carry_through_if
helper removes the repeated MLIR scaffolding for the if cases.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
…_values extract_to_ir_values now raises for anything it cannot flatten, matching its original contract; the promotion of a bare python scalar to an MLIR constant lives in _explode_states, where the pre-refactor control-flow code handled it. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Normalize a carried python scalar to an MLIR constant with a direct type check before flattening, mirroring how the pre-refactor control flow ran every value through as_ir_value. Drops the try/except-on-TypeError fallback so scalar coercion is a deliberate control-flow step rather than an error path. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Carry a dict through if, a tuple through for, a SimpleNamespace through while, and a tuple-holding-a-dict through if, each in a real kernel launched on device, matching the container variety already covered at the MLIR-unit level. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
The whole tuple of pre-region values is the structure template, so unpack it in one extract_to_ir_values / get_ir_types call and rebuild it in one construct_from_ir_values call. Replaces the per-name _explode_states / _assemble_states / _explode_branch_outputs (with hand-maintained counts and per-name exemplars) with three thin helpers: _unpack_states, _pack_states, and _unpack_branch_outputs. Structure verification collapses to a single flat-type equality check; its message unifies the former shape/dtype errors, so the matching tests are updated. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
…tched var - Condense the module comment and the three helper docstrings to what they do. - Drop the outer list() at _pack_states call sites; the helper already wraps slots. - _unpack_branch_outputs now checks each carried variable against its entry value per name, so the error names the offending variable and shows its entry vs exit types instead of one flattened whole-region comparison. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
scf.for infers its result types from iter_args, so the flat types from _unpack_states are unused there; unpack them into _ instead of a named var. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Two defects in carrying state through dynamic control flow: 1. A bare python scalar entry (acc = 0) was promoted to a raw ir.Value, which construct_from_ir_values cannot rebuild (no __construct_from_ir_values__), so packing threw 'Cannot construct DSL value for ArithValue'. Promote scalars to a DSL numeric via as_dsl_value(as_ir_value(v)) instead, so the exemplar stays reconstructable. This restores the acc = 0 / acc = 0.0 accumulator pattern. 2. Branch-output verification compared only the flattened leaf ir.Types, so a branch reordering dict/namespace keys yielded values into the wrong slots with no diagnostic (silent miscompile). Compare the full pytree structure (container kind + keys/length + order + leaf dtypes) via _tree_structure. Adds unit guards (bare-scalar for-accumulator, dict key-reorder error) and an e2e bare-scalar accumulator; error-message match strings updated. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
A dict's key order is not semantic ({'a':1,'b':2} == {'b':2,'a':1}), so a branch
listing keys in a different order than the entry should carry correctly, not
error. Replace the structure-equality check with _flatten_like, which flattens a
branch output guided by the entry template: dict/namespace are aligned by key
(order-insensitive), list/tuple by position. Reordered keys now land in the
right slots; genuine mismatches (different key set, length, or dtype) still
raise a clear per-variable error.
Guards: unit tests for key reorder (accepted) and key-set mismatch (errors);
e2e for a reordered else branch keeping a=100, b=200.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Covers acc = 0.0 alongside the existing acc = 0 guard; both go through the same scalar-promotion path but the review asked for 0/0.0 explicitly. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
…evel
_promote_scalars only mapped the top-level carried states, so a bare literal
nested in a container (e.g. d = {'count': 0, 'vec': [0, x]}) survived unpromoted
and crashed in extract_to_ir_values ('Cannot extract IR values from 0'). Make
promotion recursive (_promote walks dict/namespace/list/tuple), matching the
leaf-level promotion _flatten_like already does for branch outputs. Adds a
nested bare-scalar carry guard.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Structure {"a": i32, "b": [i32, (i32, i32)], "c": {"d": i32, "e": [i32]}}
(6 leaves, dict/list/tuple mutually nested) carried through dynamic for, if, and
while, with every leaf rewritten inside the region, launched on device and
checked against expected values.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
fx.buffer_ops (attribute on the flydsl.expr package) is gone on current main; import the submodule directly (from flydsl.expr import buffer_ops) like the other system tests do. The old form only worked against a stale local build. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
PR #880 moved buffer_ops to kernels/common and dropped the flydsl.expr alias, so neither fx.buffer_ops nor 'from flydsl.expr import buffer_ops' resolves on main. Write outputs with plain tensor subscripting (Out[i] = val), matching the other system tests (e.g. test_while_e2e) and dropping the buffer_ops dependency. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
xudoyuan
force-pushed
the
xudoyuan/dyn-cf-list-carry
branch
from
July 24, 2026 16:31
9d49c82 to
70b0af7
Compare
_promote is already recursive, so the _promote_scalars wrapper (its only caller being _unpack_states) is just [_promote(v) for v in values]; inline it. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Kernels may carry a raw ir.Value (e.g. an ArithValue) as control-flow state, not a DSL numeric -- flash-attn's lazy_rescale_o does. _promote only handled bare python scalars, so the raw ir.Value stayed in the exemplar and construct threw 'Cannot construct DSL value for ArithValue'. Promote raw ir.Values via as_dsl_value too, restoring the old dispatcher's behavior. Adds a regression test carrying a raw ir.Value through scf.if. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dynamic scf.if / scf.for / scf.while / ifexp previously required every carried variable to be a single ir.Value, so a reassigned list/tuple/dict local raised "state variable '...' is list, not an MLIR Value".
Route the four dispatchers through an explode/assemble step that reuses the existing DSL<->ir.Value protocol (protocol.extract/construct): each carried value (possibly a nested container) is exploded to per-element iter_args/ results and assembled back from the pre-region value used as the structural template on exit. Scalars stay the degenerate single-slot case so existing behaviour is unchanged; a branch/body that changes a container's shape or an element's dtype now fails with a clear error.
Adds MLIR-level unit tests and device system tests for list/tuple/dict/nested carry across if/for/while/ifexp.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist