Repository navigation
fix[next]: lower function definitions to let bindings - #2905
tehrengruber wants to merge 5 commits into
Conversation
`InlineFundefs` replaced every `SymRef` matching a program-level function definition, regardless of any binder of the same name in between. A lambda parameter named like a function definition was therefore silently replaced by that function, producing an ill-typed program (e.g. an `AssertionError` in `type_synthesizer._canonicalize_nb_fields`). Instead of substituting references, bind all function definitions in a `let` wrapping every statement expression. Regular scoping rules then apply, i.e. an inner binder shadows the function definition, and unused bindings are removed by dead code elimination, which makes the separate `prune_unreferenced_fundefs` pass obsolete.
|
|
The traced `flux` function definition is polymorphic in its parameter `d`, it is called with an `I` and a `J` offset. Type inference can not type a `let` bound function used at several types, which the function definitions are lowered to now. Dropping the `fundef` decorator makes tracing inline `flux` as a lambda instead, which keeps the test unchanged otherwise.
|
The 1. Regression: @gtx.field_operator
def inner(a: F) -> F:
return a + 1.0
@gtx.field_operator
def outer(a: F, b: F) -> F:
return concat_where(I >= 1, inner(a), b)
@gtx.program
def p(a: F, b: F, o1: F, o2: F):
outer(a, b, out=o1)
inner(a, out=o2)This fails in both Every statement binds every function definition, so in the second statement the lambda for 2. False
@gtx.field_operator
def foo(foo: F) -> F:
return foo + 1.0
# CycleError: ('nodes are in a cycle', ['foo', 'foo'])
3. Other behaviour lost compared to
Suggestion: substitute at each reference, but respect binders @dataclasses.dataclass(frozen=True)
class InlineFundefs(PreserveLocationVisitor, NodeTranslator):
fundefs: dict[str, itir.FunctionDefinition]
def visit_SymRef(self, node, *, shadowed):
if node.id in self.fundefs and node.id not in shadowed:
fundef = self.fundefs[node.id]
return self.visit(
itir.Lambda(params=fundef.params, expr=fundef.expr), shadowed=frozenset()
)
return self.generic_visit(node, shadowed=shadowed)
def visit_Lambda(self, node, *, shadowed):
return self.generic_visit(node, shadowed=shadowed | {p.id for p in node.params})
def inline_fundefs(program: itir.Program) -> itir.Program:
visitor = InlineFundefs({str(f.id): f for f in program.function_definitions})
return itir.Program(
id=program.id,
function_definitions=[],
params=program.params,
declarations=visitor.visit(program.declarations, shadowed=frozenset()),
body=visitor.visit(program.body, shadowed=frozenset()),
)
Tested with this patched into the pass manager:
I haven't compiled anything with Trade-offs of this approach It keeps a capture bug from @fundef
def times2(x):
return multiplies(deref(x), 2.0)
@fundef
def app(multiplies, x):
return times2(x)
# stencil: app(lambda p, q: minus(p, q), x)
# main and this suggestion: [-1.0, 0.0, 1.0, 2.0]
# this PR, and expected: [2.0, 4.0, 6.0, 8.0]That wrong result is silent, but it needs a callable bound to a builtin's name, which I only reached from tracing or hand-built IR. From the frontend, a parameter named It also gives up the smaller IR. For Your comment above says all lambdas are force-inlined anyway, so I don't see what
The behaviour in 3 would then remain, and Comparison on icon4py I compared the IR from
The three unresolved programs are 🤖 Written by an agent on behalf of @havogt |
Missing information propagates as `DeferredType` throughout the iterator type inference, but `visit_FunCall` called whatever `visit` returned for `node.fun`, raising `TypeError: 'DeferredType' object is not callable` for an undeclared symbol. This is hit when `reinfer` is run on an expression without its enclosing scope, e.g. by `concat_where.expand_tuple_args` on a not yet typed function body.
Three fixes to `inline_fundefs`: - The dependency order was computed from references collected without the definition's own parameters in scope, so a parameter named like a function definition was taken for a dependency. A parameter named like the function itself, or two functions with parameters named like each other, resulted in a bogus `graphlib.CycleError`. - Only the expression of a statement was wrapped, so a function definition referenced from a statement or temporary domain was left dangling once the function definitions were dropped. All of them are wrapped now. - To keep these expressions small, every position only binds the function definitions it references, directly or transitively. Unreferenced definitions are not bound at all, which also keeps domain inference from seeing bodies that are never called.
|
All three reproduce. 1. and 2. are addressed, 3. only partly. 1. Missing information propagates as 2. The dependency collection now has the definition's own parameters in scope, so a parameter named like a function definition is no longer taken for a dependency. Added unit tests for a parameter named like its own function, for two functions with parameters named like each other, and for a definition that shadows one function but genuinely calls another. 3. Only the function definition in a domain is addressed. Statement domains and temporary domains are now wrapped like statement expressions. To keep these expressions small, each position only binds the function definitions it actually references, directly or transitively. Unreferenced definitions are not bound at all. As a side effect this also makes the case in 1. unreachable, since the dead binding is never created. The polymorphic |
InlineFundefsreplaced everySymRefmatching a program-level function definition, regardless of any binder of the same name in between. A lambda parameter named like a field operator was therefore silently replaced by that function, producing an ill-typed program (e.g. anAssertionErrorintype_synthesizer._canonicalize_nb_fields).Instead of substituting references, all function definitions are now bound in nested
lets wrapping every statement expression, andProgram.function_definitionsis cleared. Regular scoping rules then apply, i.e. an inner binder shadows the function definition, and unused bindings are removed by dead code elimination, which makes the separateprune_unreferenced_fundefspass obsolete. Thelets are nested in topological order since function definitions may call each other.Follow-up to #2833 (comment), where tuple-comprehension targets made this easy to hit.
AI disclaimer: This PR has been created with the help of AI tools. The code has been reviewed, but tests were only briefly sighted.