[Relay][Pass] Fix lambda lift pass for recursive call

[icemelon]

Thanks for contributing to TVM! Please refer to guideline https://docs.tvm.ai/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from Reviewers by @ them in the pull request thread.

cc @jroesch @zhiics @wweic

[zhiics]

LGTM. @jroesch PTAL

[MarisaKirisame]

Does this work for non-immediate recursion?
in particular, can you add this test case:
-- implement a * b using increment, in tail-recursive accumulator style
let rec multACC a b acc =
if b == 0
then acc
else
let rec addACC a oldB b acc =
if b == 0
then multACC a (oldB - 1) acc
else addACC a oldB (b - 1) (acc + 1)
in addACC a b b acc
let mult a b = multACC a b 0

[icemelon]

@MarisaKirisame No, we couldn't support non-immediate recursion. In this case, you need to support mutual recursion, which is still a pending PR(#2648).

[junrushao]

@MarisaKirisame I was wondering what the proper way is to do lambda lifting, where lambda could be returned or passed into a new function?

[MarisaKirisame]

@icemelon9 the program above do not rely on the mutual recursion feature.
The closer inside will capture the outside closure as free variable, While the outside closure create multiple inside closure across invocation.

[MarisaKirisame]

@junrushao1994 A global function returning that closure can be used to implement returning closure. (In this case, the old lambda lifting behavior is an optimization that do not return the closure, but call it directly).
As for closure being passed in - they are first class, so you do not have to do anything about it. treat it just like any other object.

[junrushao]

@MarisaKirisame Thank you for your reply! I am very ignorant in this field, and super interested in learning more. Would you mind suggesting better ideas if you have time?

Definition(?). I was thinking, lambda lifting should be a pass that eliminates free variables and lifting lambdas to global scope. Like what is described in Wikipedia.

Simplest case. If a lambda can be converted is only used inside relay.call, then it is safe to lift it to a tuple (global_var_lifted, captured_0, captured_1, ...). Mutual recursion works in this case as well.

More complicated case. However, in the general case, when lambda is first-class supported (like relay), it could be hard to lift them, because when adjusting call sites, type mismatch could happen (consider conditionally bind 2 lambdas to a variable, which might capture different things).

Your suggestion. As you suggested, in this case, make a "global function returning that closure can be used to implement returning closure." This is one of the solution that I was thinking before. It works, but I am not sure if it really helps (because closure still exists).

What Wikipedia says. In the section "algorithm", it says:

If the language has closures as first-class objects that can be passed as arguments or returned from other functions, the closure will need to be represented by a data structure that captures the bindings of the free variables.

Why I ask this If we are able to eliminate those lambdas, to "some data structure", then we are able to do autodiff without reference...

[MarisaKirisame]

@junrushao1994 No need to say that! DL compiler require many expertise and we are all ignorant in different things.
I should answer the general problem before the specific problem.
But before we start, in general, closure is the "killer" of ad. Some ad work do not support closure, like tangent or tensorflow-for-swift. Some work support it, but take lots of effort - Lambda The Ultimate BackPropagator and differentiate through a curry (by author of tf for swift) come to mind. Additionally, all work need some additional bookkeeping for captured value inside the closure, in order to backprop correctly to them. (If we do not use reference, the caller of a lambda must update the reference itself.) Some people use reflection, and some people use some dynamic typing and cast. Reference is a very lightweight bookkeeping technique while making stuff good for other optimization (Reflection will be a much bigger problem for opt/codegen then ref).

How to do autodiff of lambda with minimal amount of reference?

The method of turning lambdas into datastructure is well known and it is called defunctionalization. Basically, it try to make a new datatype, Func, where every single constructor is a "lambda" in the AST, and the constructor take in the closure number.

For example, suppose the AST contain the following lambda, all taking two int and returning int
(fun _ _ -> a),
(fun a b -> a * a + b * b)
(fun a b -> a + b + c)
(fun a b -> f (f a b) (f b a))
the following adt will be declared:
type FUN =
| CONST of int
| AddSquare
| AddThree of int
| Point of FUN.
There is some issue regarding making it type safe, but I imagine we can get by by using some cast here and there.

we can then turn some adt containing reference into reference of pure adt. After that we can then do store passing style to that single reference. But I do not think that we can in general remove it (I had tried like 10 times).

Off a side note, why are you so eager to remove reference? I need to write paper and I think this is a good paper story.

I will talk about lambda lifting below, but I think we should not mix up lambda lifting with defunctionalization.

[MarisaKirisame]

For lambda lifting, note that our main use of ll right now is to massage the code into a form suitable for VM Compiler. We can make the VM Compiler appreciate lambda-lifted closure returning code.

[kevinthesun]

LGTM

[kevinthesun]

Thanks @icemelon9 . This is now merged.