Fix loop-carried Vector losing multi-dim shape (issue #734) - #759
Merged
Conversation
A Vector's multi-dim logical shape (e.g. (4,1)) lives only on the Python wrapper; the MLIR value is always a flat vector<Nxf32>. When a Vector is carried across an scf region boundary (for/while/if), the loop-carried value is rebuilt from the bare block-argument ir.Value, and reconstruction went through the classmethod __construct_from_ir_values__ which has no access to the exemplar instance's shape -- so it collapsed (4,1) -> (4,). The next vec_sum += vec then broadcast (4,) + (4,1) to (4,4) and aborted compilation with a vector<4xf32> vs vector<16xf32> type mismatch. Fix: add an optional instance-level reconstruction hook Vector.__reconstruct_from_ir_value__(self, value) that rebuilds the wrapper around the new ir.Value while copying the exemplar's shape/dtype, and prefer it in as_dsl_value before falling back to the classmethod protocol. Other types (Numeric, Constexpr, struct, BuiltinDslType) are unaffected -- they have no such hook and keep using __construct_from_ir_values__. This fixes Vector carries across for/while/if at once, since they share the same reconstruction path. Add end-to-end regression test tests/system/test_for_vector_carry_shape_e2e.py (fails pre-fix with the vector<16xf32> mismatch, passes post-fix). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace __reconstruct_from_ir_value__ with an optional exemplar argument threaded through __construct_from_ir_values__ and construct_from_ir_values, so a Vector rebuilt from a bare scf block-argument ir.Value preserves its shape/dtype. Keeps reconstruction within the DslType protocol and also fixes the jit/kernel argument rebuild path.
sjfeng1999
approved these changes
Jun 29, 2026
coderfeli
approved these changes
Jun 30, 2026
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
3 participants
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.
A Vector's multi-dim logical shape (e.g. (4,1)) lives only on the Python wrapper; the MLIR value is always a flat vector. When a Vector is carried across an scf region boundary (for/while/if), the loop-carried value is rebuilt from the bare block-argument ir.Value, and reconstruction went through the classmethod construct_from_ir_values which has no access to the exemplar instance's shape -- so it collapsed (4,1) -> (4,). The next vec_sum += vec then broadcast (4,) + (4,1) to (4,4) and aborted compilation with a vector<4xf32> vs vector<16xf32> type mismatch.
Fix: add an optional instance-level reconstruction hook Vector.reconstruct_from_ir_value(self, value) that rebuilds the wrapper around the new ir.Value while copying the exemplar's shape/dtype, and prefer it in as_dsl_value before falling back to the classmethod protocol. Other types (Numeric, Constexpr, struct, BuiltinDslType) are unaffected -- they have no such hook and keep using construct_from_ir_values. This fixes Vector carries across for/while/if at once, since they share the same reconstruction path.
Add end-to-end regression test tests/system/test_for_vector_carry_shape_e2e.py (fails pre-fix with the vector<16xf32> mismatch, passes post-fix).
close #734
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist