microsoft / knossos-ksc

Compiler with automatic differentiation
Other
45 stars 10 forks source link

Make all binders unique ABU #807

Open awf opened 3 years ago

awf commented 3 years ago

There are numerous delicacies caused across the codebase because we allow shadowing of the form

(let (FOO (lam (arg) ....outer_x....)) 
   (let (outer_x 77)  ; shadow, boo!
       (FOO outer_x))

and then when we try to inline FOO, we get

(let (FOO (lam (arg) ....outer_x....)) 
   (let (outer_x 77) 
       ((lam (arg) ....outer_x....) outer_x))

Oh no! If we hadn't allowed shadowing, we would have had

(let (FOO (lam (arg) ....outer_x....)) 
   (let (inner_x 77)  
       (FOO inner_x))

rewriting to

(let (FOO (lam (arg) ....outer_x....)) 
   (let (inner_x 77) 
       ((lam (arg) ....outer_x....) inner_x)); phew!

Whereas:

  1. It is easy to automatically transform a non-ABU expression to an ABU one
  2. We need ABU anyway for hashing modulo alpha equivalence

We should just do it.

simonpj commented 3 years ago

There are two variants

  1. No shadowing
  2. All binders everywhere are unique

(2) is harder than (1); consider f (\x.e1) (\x.e2), where the lambda binders are in sibling expressions. Of course you can achieve 2 by having a single, infinite supply of binders and always picking a fresh one -- you just have to everything in a state monad.

(1) is significantly easier, and is pretty much what we have.

acl33 commented 3 years ago

I have no technical objection to ABU. We opted for non-unique-names / capture-avoiding-substitution, because we thought it would be faster as it involved less renaming(=copying); I don't see that anything has changed there. My main concern is just how high priority is this?

AWF asks whether there are any other hairy cases. The case discussed above in the OP is (by far) the hairiest one that I remember from RLO. Note that the let-bound lambdas are not something we support elsewhere yet (lam can only occur immediately inside build/fold/map/sumbuild/etc.). The other "hairy case" that ABU helps with is a rule like

(rule "lift_let_from_if_true" <...args...> (if p (let (x e) body) y) (let (x e) (if p body y)))`

where we must rename x (within body) if there is another x free in p or y. See this snippet for one way of dealing with this.

Whether there are any hairy cases with ABU is a different question. I suspect that any such are hairly only in terms of what's necessary to make them go fast, which is something we didn't manage to do in RLO, so I can't draw on experience there to answer. Clearly there would be some design work required. (And perhaps infrequent slow cases are less problematic if we can be sure they are infrequent.)

acl33 commented 3 years ago

We might also want to look at https://www.schoolofhaskell.com/user/edwardk/bound. And/or perhaps our PLDI paper may suggest other schemes (tho I've not yet figured out anything that looks practical myself).

acl33 commented 3 years ago

Our previous discussion on the amount of copying required concluded:

SO: it boils down to this? let v = e1 in e2 [and then consider inlining some occurrence of 'v' in e2] ABU: need to deep-copy e1 whenever rhs contains lets (not necessary, but sufficient)

[I'm not sure why "not necessary, but sufficient" - it looks necessary to me but maybe I've misunderstood the context. Deep-copying means every node in e1 that's on a path to a usage of a variable bound within e1.]

Non-unique names: need to deep-copy e2 whenever free_vars(e1) intersect bound_vars(e2) != 0

[actually only need bound_vars() along the path (in e2) to the usage of v that's being inlined, and only as much deep-copying as necessary to propagate those renames]

awf commented 3 years ago

Plan:

  1. We don't try to ensure unique names in ctor.
    • That would require storage of set of bound vars.
    • Could be micro-bloom filter: bool => "no bound vars"
    • But we won't do that now
  2. We provide O(n * (log n for map lookups)) expr = ensure_all_names_unique(expr)
  3. We do have a O(n log n) check are_all_names_unique(e : Expr)
  4. Also deep-copying
    def copy_renaming_binders(e : Expr) -> Expr:
    if e.has_no_bound_vars:
    return e
    # Now do real work - recursively copy_renaming_binders
  5. Changes:
    • cav_subst: replace(path, root, replacement). Might need to delay-copy replacement.
    • rewrites: inline_var needs to copy_renaming_binders; also any parsed rule where a template_var appears >1 time in the replacement
awf commented 3 years ago

Some examples, and what we need to make happen

root: let x = build 7  (\ i . foo i y) in add x x
path: ----------------------------------------^
replacement: {build 7  (\ i . foo i y)}  !! Alert?  This should have been a copy!

root: let x = build 7  (\ i . foo i y) in add x x
path: ----------------------------------------^
replacement: {build 7  (\ i1 . foo i1 y)}
answer: let x = build 7  (\ i . foo i y) in add (build 7  (\ i1 . foo i1 y)) x

Now simple.

--------------

root: let foo = \ i . bar i y in Call (compose foo foo) x
path: ------------------------------------^^^
replacement: \ i2 . foo i2 y
answer: let foo = \ i . bar i y in Call (compose (\ i2 . bar i2 y) foo) x

As above

root: Call (\ x . foo x y) (add y z)
path: ^
replacement: let x = add y z in foo x y
answer: let x = add y z in foo x y