[RFC] [Relay]A Normal Form Conversion

[MarisaKirisame]

#2251 implement a pass which converts expression with sharing (different occurrence of Expr might be the same ptr) into expression without sharing, by inserting let all over the place.

It is useful because it allows subsequent passes to ignore sharing and only deal with let (which might also reside in the original program anyway).

However, to do so, we must define precisely what is the scope of a node.
It is determined by the Expr that immediately depends on it.
For example in the Graph a + b, a + b immediate depends on both a and b.
We will call 'immediately depends on' 'point to' from now on.

Let's start from the degenerate case that there is only one graph a that points to b.
In that case, GetScope(a, b) = scope(a) if (b is not in curly brace position of a) else new_subscope(scope(a))

Now the other case.
If nothing points to node b:
b is the outmost scope

If there are multiple pointer pointing to b:
b's scope is the least common ancestor of all GetScope.

[tqchen]

How about we define the scope of an expression to be the least common ancestor of all the scopes that directly refers the expression?

[MarisaKirisame]

@tqchen we cant, it will lift outside of if else.
Beside special casing for singly scoped object it is what we do.

[tqchen]

Can you give a code example on why we cannot use the LCA scope as the scope ?

[MarisaKirisame]

@tqchen suppose we have
sum(x) = if (x == 0) {0} else {x + sum(x - 1)}
if we try to use the LCA scope, we will lift the 'x + sum(x - 1)' out of the if else, and in the same scope as the parent if, which is incorrect and make the whole thing an infinite loop. It is one reason the definition is complex.

[tqchen]

@MarisaKirisame in your particular case

 if (x == 0) {
    0
 } else {
   %1 = x + sum(x - 1)
   %1
}

Because %1 is only used in the scope of else, all its usage is only in the scope of else. So the LCA of %1 is the else scope. However, if the code is:

%1 = %x -1
 if (%1 == 0) {
    0
 } else {
   %2 = x + sum(%1)
   %1
}

Then because %1 get used in both the condition of if, and else scope of if, its LCA is outside the if scope.

[MarisaKirisame]

@tqchen this is exactly our algorithm - we cant always take the parent scope, if it is in special location (inside any branch of if/else, inside function definition, inside pattern matching), new scope have to be created.

It is indeed the LCA of all scope reference though.

[tqchen]

@MarisaKirisame Sounds good, can you list a few examples(in text format) to explain the behavior of this algorithm. In particular, we also specify the order of evaluation. Let me give an example:

%2 = %y + 1
%1 = %x
if (%1 == 0) {
   %2
} else {
   %1 + %2
}

Should be translated to:

let %v1 = %x
let %v2 = %y + 1
if (%v1 == 0) {
  %v2
} else {
   %v1 + %v2
}

The mapping from node to temp vars are : %1->%v1, %2->%v2. Notice how the order of assignments is listed as v1 first then v2. This is because we need to have a way to uniquely define the order of the assignments if they are in the same scope level.

One way to define that is via post-DFS visit order. If we post DFS visit the graph, %1 get visited first (because of the condition). So our key is that we need deterministically maps a graph to ANF and give this mapping a clear semantics.

I see you also commented on what might happen to lambda lifting, please give an example on that as well. These examples, along with a description of your algorithm, will help more reviewers to understand this pass, which is quite important. cc @jroesch

[MarisaKirisame]

In particular, we also specify the order of evaluation.

@tqchen I am against this. Either the order does not matter, or, if the order matter in a graph, something is wrong. This behavior should be implementation-defined, simply because no code should depend on how it is order - something is wrong in the semantic level.

[tqchen]

It is true that order does not matter in a graph. We essentially need to impose additional semantics when converting from graph to ANF to make sure such conversion deterministic.

Having such deterministic conversion makes developers easy to reason about the consequences of the operation. Post-DFS order is a frequently used lowering rule already for graph code, so I propose that we adopt this semantics in the ToANF pass

[MarisaKirisame]

@tqchen i just checked and it is post order. I will update the comment.

[junrushao]

@tqchen in your example, %1 = %x -1 should be %1 = %x, which is a typo

[junrushao]

Just curious, have you guys decided to do opt based on anf or cps? What does the time complexity looks like when it comes to crazy neural nets as deep as ~10k layers?