[Pass] Canonicalizing Bindings

[slyubomirsky]

It may be useful for some passes to collapse chains of definitions, particularly after other compiler transformations that may reduce or simplify some expressions.

This pass will take chains of definitions and replace references to later definitions to the original one. It works by checking LookupBinding for each var use-site and replacing the var with its definition if the definition was another var. (Note: This required updating BlockBuilder to also update its binding map for MatchShape nodes; that was arguably a bug.) Additionally, MatchShape bindings where the LHS and the RHS are guaranteed to match at compile time are canonicalized into ordinary VarBindings.

Example:

y = x
z = y
w = z
o = w
p = o

Will be replaced with

y = x
z = x
w = x
o = x
p = x

Later, unused bindings can be removed by another pass. Eventually, we can consider other kinds of identity operations like multiplying by 1 or adding 0 in this canonicalization, but this pass does not handle it for now because such operations may affect aliasing (if they allocate new tensors). We can revisit that issue once our position on aliasing is clear.

Some tricky parts of this pass:

1. Var bindings can escape DataflowBlocks, so if a Var is assigned to a DataflowVar, we must take care never to replace a Var with a DataflowVar.
2. ExprMutator's default behavior is to update type annotations based on the checked_type_ of the body. This can result in erasing user-provided type annotations, which we should respect.
3. We face a similar issue for the shape_ field, with the added complexity that the shape_ for some operators can feature calls to PackedFuncs, so we will have to ensure variable replacements are performed within those.

I would greatly appreciate review or proposals for improvements, especially for test cases that may be problematic. I think removing the unused bindings should be handled by other passes. I imagine this as a pass that could be useful when combined with others. (If we don't feel it's useful, we can also drop the PR and revisit down the line).

[sunggg]

Thank you for the great work!
I haven't reviewed the code yet (will do!), but have a general comment I'd like to mention.
Can we design this pass to consider future extension with these relay passes? It would be great if we can reuse the basic structure.
https://github.com/apache/tvm/blob/main/src/relay/transforms/canonicalize_ops.cc
https://github.com/apache/tvm/blob/main/src/relay/transforms/canonicalize_cast.cc

[slyubomirsky]

Openness to extension is certainly a good idea. What might that entail? Naming this CanonicalizeBindings so as to open the door to other canonicalization passes in the future?

[LeshengJin]

Thanks for the great work! 🎉
I love the comments in your code. They helped me a lot in understanding this pr. Really learned a lot.

I have two small questions for you:)

[sunggg]

That's a certainly one option. But it may result in many passes that does different kinds of canonicalization.
Or we can implement a group of benign canonicalizations (applying them is always good, like your binding pass) inside the single pass and just call it canonicalization pass.
Personally, I prefer the latter option since it is easier to configure at pipeline perspective.

[psrivas2]

Thank you for implementing this important optimization!

I like that this pass composes well with dead code elimination - nice application of "separation of concerns" here :)

I have a few high-level comments:

1. It seems we can further break down this optimization into two independent canonicalization patterns:

• Patterns:
  • Identity folding: Replace y=x; z=y with y=x; z=x.
  • match_shape folding: Replace y = match_shape(x, shape_expr) with y = x if x.shape is equivalent to shape_expr.
• The two patterns above would compose well and can be applied greedily until fixed point to achieve the same result. We can add more such patterns in the future.

2. Cast Ops: y: Object = x needs special handling because of implicit casting. If we have explicit casting then this would be a cast op and we do not need special handling for these. Since the decision on that is still pending, just wanted to highlight this case. For future optimizations too implicit casting might complicate things.
3. Usually with canonicalization patterns we might want to apply many of them greedily until fixed point. For example, the two patterns above. This way we can keep the system easily extensible.

[psrivas2]

I think match_shape can have an independent canonicalization pattern. It is trivially removed if for R.match_shape(a, shape_expr) shape_expr is equivalent to a.shape_ because we know it is always true.
This way we can have two simple canonicalization patterns: identity and match_shape which can be composed and applied greedily to get the same result.

[slyubomirsky]

I like the idea of structuring this pass in terms of different patterns instead of handling everything together, but I'm having a little trouble seeing how these patterns would be represented in the code. Could you please elaborate on how you think that should be implemented?

[psrivas2]

This might be a good motivation to have explicit cast ops. This is not an identity operation, but a cast operation (Tensor -> Object) and hence we need special handling for this case.

[slyubomirsky]

I would not call this a cast per se because Object is a supertype of Tensor; a strongly typed language like Java would not require an explicit cast if assigning a variable to one with a supertype. Do you think the pass requires any additional logic for cases where the RHS and LHS have different types? (Edit: I suppose it can be called an implicit cast, and it is implicit precisely because there is no specific expression associated with it.)

[sunggg]

Thanks for the great contribution!
My high-level feedback is here: #233 (comment)
LGTM. Just a few nits.

[slyubomirsky]

@psrivas2 I like your idea about having multiple canonicalization patterns and applying them to fixpoint. Could you elaborate on how that might be structured in the code? What would be the interface for a canonicalization pattern?

[slyubomirsky]

I could imagine a fixpoint approach that does not make use of a union-find. Namely, for every var, look up its binding site. If the binding site is another var, replace that var with the parent (unless the checked_type_ or shape_ differ). Repeat until fixpoint. In principle, we could check for other canonicalization patterns when we do the variable lookup. Is that what you had in mind, @psrivas2?

[psrivas2]

I could imagine a fixpoint approach that does not make use of a union-find. Namely, for every var, look up its binding site. If the binding site is another var, replace that var with the parent (unless the checked_type_ or shape_ differ). Repeat until fixpoint. In principle, we could check for other canonicalization patterns when we do the variable lookup. Is that what you had in mind, @psrivas2?

Yes, that is what I had in mind. It would also be nice (your call to do it in this PR or leave it to follow up PRs) if the canonicalizer can do both -- apply patterns and also constant fold aggressively. Constant folding can produce new canonicalization possibilities.

[slyubomirsky]

I will try the fixpoint approach. I will also see if I can use constant folding inside this pass as well.

[slyubomirsky]

@psrivas2 Thank you for suggesting another approach--using LookupBinding without a union-find vastly simplified the code. Interestingly, I found that I did not even need to run multiple iterations of the pass to achieve fixpoint. Can you (or others) think of any cases where multiple passes are needed?

[psrivas2]

Is my understanding correct that this block of code is not needed if we allow type refinement. So z: Object = x can be refined to z: Tensor = x. So do we still need this?

[slyubomirsky]

I am under the impression that we should respect user annotations when they appear, which is why I've done it

[psrivas2]

I went back and checked the notes in August 16, 2022 meeting. Seems like this topic is under discussion. Would be great if can reach consensus on this soon.

[psrivas2]

When is this required? None of the tests check for this.

[slyubomirsky]

The tests do fail if you leave this out. This is because visiting shape_ can cause variables inside the shape_ field to change. This comes up in the test cases with relax.add, since the shape_ for that is a PackedFunc call that uses variables in the program.

[psrivas2]

using LookupBinding without a union-find vastly simplified the code.

Indeed! Thank you for making the changes. Left some comments in the code.

Can you (or others) think of any cases where multiple passes are needed

Good point. For single pattern, a single pass should be fine.

In case of multiple patterns, it depends on how different patterns are being applied. While visiting bindings in RPO, if we apply all patterns to the value of a binding till none of the patterns are applicable, and only then move to next binding, it does seem like a single pass would be enough.