fix: handle co_freevars mismatch in ASTRewriter for closure kernels#474
fix: handle co_freevars mismatch in ASTRewriter for closure kernels#474ZJLi2013 wants to merge 1 commit into
Conversation
…ure refs Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
Pull request overview
This PR fixes a failure mode in ASTRewriter.transform() when rewriting closure-defined @flyc.kernel functions: after AST transforms inject helper name references, the rewritten code object can end up with a different co_freevars shape than the original function’s closure, causing f.__code__ = new_code to raise a ValueError. The change adds a fallback path that constructs a new types.FunctionType with a synthesized closure matching the rewritten co_freevars.
Changes:
- Detect
co_freevarsmismatches for closure functions and build a replacement function with a matching closure. - Reuse existing closure cells where possible; synthesize new cells for newly-introduced freevars using values from
f.__globals__. - Keep the non-closure / no-mismatch path unchanged (still assigns
f.__code__directly).
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| new_closure.append(old_cells[var]) | ||
| else: | ||
| # Create a cell whose value comes from globals | ||
| cell = (lambda v: lambda: v)(f.__globals__.get(var)).__closure__[0] |
| if f.__closure__ and new_f_code_o.co_freevars != f.__code__.co_freevars: | ||
| old_cells = {name: cell for name, cell | ||
| in zip(f.__code__.co_freevars, f.__closure__)} | ||
| new_closure = [] | ||
| for var in new_f_code_o.co_freevars: | ||
| if var in old_cells: | ||
| new_closure.append(old_cells[var]) | ||
| else: | ||
| # Create a cell whose value comes from globals | ||
| cell = (lambda v: lambda: v)(f.__globals__.get(var)).__closure__[0] | ||
| new_closure.append(cell) | ||
| new_f = types.FunctionType( | ||
| new_f_code_o, f.__globals__, f.__name__, | ||
| f.__defaults__, tuple(new_closure), | ||
| ) | ||
| new_f.__kwdefaults__ = f.__kwdefaults__ | ||
| return new_f |
| new_f_code_o, f.__globals__, f.__name__, | ||
| f.__defaults__, tuple(new_closure), | ||
| ) | ||
| new_f.__kwdefaults__ = f.__kwdefaults__ |
| # AST transformers may inject references to new names (e.g. | ||
| # scf_if_dispatch, scf_if_collect_results) inside the kernel or | ||
| # its rewriter-generated sub-functions (__then_N, __else_N). | ||
| # Because enclosing_mod creates a closure layer, these unresolved | ||
| # names become free vars (LOAD_DEREF) rather than globals. This | ||
| # causes new_f_code_o.co_freevars to have more entries than the | ||
| # original f.__closure__. Direct f.__code__ assignment would | ||
| # raise ValueError, so we build a new function with a matching | ||
| # closure instead. | ||
| if f.__closure__ and new_f_code_o.co_freevars != f.__code__.co_freevars: | ||
| old_cells = {name: cell for name, cell | ||
| in zip(f.__code__.co_freevars, f.__closure__)} | ||
| new_closure = [] | ||
| for var in new_f_code_o.co_freevars: | ||
| if var in old_cells: | ||
| new_closure.append(old_cells[var]) | ||
| else: | ||
| # Create a cell whose value comes from globals | ||
| cell = (lambda v: lambda: v)(f.__globals__.get(var)).__closure__[0] | ||
| new_closure.append(cell) | ||
| new_f = types.FunctionType( | ||
| new_f_code_o, f.__globals__, f.__name__, | ||
| f.__defaults__, tuple(new_closure), | ||
| ) | ||
| new_f.__kwdefaults__ = f.__kwdefaults__ | ||
| return new_f |
|
@xudoyuan take a look? |
|
Hi, Please provide a use case for the scenario "new_f_code_o.co_freevars to have more entries than the original f.closure". Thanks. |
|
You can put the minimum reproducible end-to-end test cases under ‘FlyDSL_Path/tests/system‘. |
| in zip(f.__code__.co_freevars, f.__closure__)} | ||
| new_closure = [] | ||
| for var in new_f_code_o.co_freevars: | ||
| if var in old_cells: |
There was a problem hiding this comment.
Assuming that 'var' is in 'f.globals' simply because it's not in 'new_f_code_o.co_freevars' is incorrect, as this can lead to situations where it can't be found even if 'f.globals' is ignored. 'f.code.co_freevars' should be directly replaced with 'new_f_code_o.co_freevars' . Due to AST processing, some 'co_freevars' can be optimized away.
|
|
||
| f.__code__ = new_f_code_o | ||
|
|
||
| for name, val in cls.rewrite_globals.items(): |
There was a problem hiding this comment.
The line 'f.globals[name] = val' can be moved before the 'f.closure' check, which makes it easier to merge the two 'f.closure' code sections.
|
Thank you for your report. There is indeed a problem. Currently, I can only find scenarios where |
Summary
When a
@flyc.kernelfunction is a closure (defined inside another function),AST transformers like
ReplaceIfWithDispatchinject new name references(
scf_if_dispatch,scf_if_collect_results) into the kernel body. Becauseenclosing_modwraps the kernel in an extra function scope, Python's compilerresolves these names as
LOAD_DEREF(free variables) rather thanLOAD_GLOBAL.This causes
new_f_code_o.co_freevarsto have more entries thanf.__closure__,and
f.__code__ = new_f_code_oraises:ValueError: kernel() requires a code object with N free vars, not M
Fix: When
co_freevarsmismatch is detected, build a new function viatypes.FunctionTypewith a closure that maps original free vars to theirexisting cells and creates new cells (from
f.__globals__) forrewriter-injected names. The non-closure path is unchanged.
Test plan