microsoft / knossos-ksc

Compiler with automatic differentiation
Other
45 stars 10 forks source link

Improve call-site appearance of replace_subtree #784

Open awf opened 3 years ago

awf commented 3 years ago
 replace_subtree(parent, path_to_child, Const(0.0), lambda *_: if_node.t_body)

This Const(0.0) is grotty. We should at least allow "None" to be passed.

Ideally, we should be able to make

replace_subtree(parent, path_to_child, if_node.t_body)

work well.

acl33 commented 3 years ago

Indeed. The replace_subtree interface was quite complex to allow inlining calls to 'lam's that might have free variables, as happens in RLO. With first-order ksc we can simplify and get that "ideal", because calls are to defs, not lams; def names are not "free"; and defs have no free variables.

Ideally I'd like to get #779 in first as that makes concrete what we do need to do to support inline_call, and then cleanup there and elsewhere.

awf commented 3 years ago

I'm against the restriction on calling lams. We should be able to tidy this without that restriction.

Let's post some examples here of what needs to be achieved. If the payload has free vars, are they in the parent scope, destination scope, or a mixture? If the latter, how is that signalled?

acl33 commented 3 years ago

The only case where there's a mixture is the "future ks language" case of of inlining calls to lams.

For inline_var, the payload is in the parent scope, i.e. we must rename binders to avoid capture.

For all other rewrites, the new subtree is created in the destination scope - we do that by building the new subtree inside a lambda (function from the old subtree to new), which in some cases is necessary (the lambda needs to be given the existing subtree to compute the new one), and in other cases, we have the new subtree already, and the lambda merely serves to hide it from the capture avoidance.

For inline_lam_into_call, thus, we pass the bit in the parent scope (the Lam) as the "payload", and then compute the new subtree in the python lambda, which is passed the call arguments (in the child scope).

If we only wanted capture-avoidance or no-capture-avoidance, but never a mixture, we could pass a bool to replace_subtree (or have two versions: replace_subtree_avoiding_capture, and replace_subtree_with_capture).

See #803 for a take on that where the underlying replace_subtrees method supports a mixture, but the usual interface (for all existing rewrites) use replace_subtree_avoid_capture and apply_func_to_subtree which do not support mixing the two.

I've closed #803 for now, but could come back to that after #805, which seems a smaller (but smoother) step.

awf commented 3 years ago

See #807