[Unity][Transform] Improved canonicalization of non-dataflow Var #15941

Lunderberg · 2023-10-17T15:26:35Z

Prior to this commit, relax.transform.CanonicalizeBindings removed trivial bindings var_y = var_x where a var_y: relax.DataflowVar and var_x: relax.Var, but did not remove trivial bindings when var_y: relax.Var and var_x: relax.DataflowVar. This was to avoid invalid use of a relax.DataflowVar outside of a dataflow block.

This commit updates CanonicalizeBindings to handle this type of binding as well. To ensure that no relax.DataflowVar instances are used outside of a dataflow block, this is done by replacing var_y: relax.DataflowVar at its point of definition, instead of replacing var_x: relax.Var at its point of use.

This commit also canonicalizes relax.Var definitions to relax.DataflowVar, if the binding occurs within a dataflow block, and the variable is never used outside of a dataflow block.

Prior to this commit, `relax.transform.CanonicalizeBindings` removed trivial bindings `var_y = var_x` where a `var_y: relax.DataflowVar` and `var_x: relax.Var`, but did not remove trivial bindings when `var_y: relax.Var` and `var_x: relax.DataflowVar`. This was to avoid invalid use of a `relax.DataflowVar` outside of a dataflow block. This commit updates `CanonicalizeBindings` to handle this type of binding as well. To ensure that no `relax.DataflowVar` instances are used outside of a dataflow block, this is done by replacing `var_y: relax.DataflowVar` at its point of definition, instead of replacing `var_x: relax.Var` at its point of use. This commit also canonicalizes `relax.Var` definitions to `relax.DataflowVar`, if the binding occurs within a dataflow block, and the variable is never used outside of a dataflow block.

Lunderberg · 2023-10-24T19:33:57Z

Rebased onto unity head and resolved conflicts with the #15791.

slyubomirsky

This is an excellent refactor and is very well documented. I made a couple of comments that I think would be worth having responses to before merging, but I think this leaves the transformation in great shape.

slyubomirsky · 2023-10-24T23:09:10Z

src/relax/transform/canonicalize_bindings.cc

+      value = ptr->value;
+    } else if (auto ptr = binding.as<MatchCastNode>()) {
+      has_same_struct_info =
+          StructuralEqual()(GetStructInfo(binding->var), GetStructInfo(ptr->value));


Might we want to consider using the arithmetic solver to check if shapes are statically provable to be equal? The struct info can be equivalent without being structurally equal. (We should factor out that checking into a utility function if it isn't already, as this comes up a lot.)

I'd like to do so, yes, though I think that would be best isolated in an independent PR. Here, the check for having identical struct info is the same check as was done previously in the CanCanonicalizeVar method, and was kept the same in order to minimize the extent of changes.

Okay. We can revisit this issue separately (I bet there are other passes where we're checking structural equality of struct info instead of static equivalence).

slyubomirsky · 2023-10-24T23:11:31Z

src/relax/transform/canonicalize_bindings.cc

+      while (true) {
+        if (auto tuple_var = tuple.as<Var>()) {


Might it be cleaner to have while (auto* tuple_var = tuple.as<VarNode>()) { as the condition? I believe this should work and it would be correct.

Agreed, and updated.

slyubomirsky · 2023-10-24T23:17:00Z

src/relax/transform/canonicalize_bindings.cc

+      while (true) {
+        if (auto it = visitor.trivial_bindings_.find(bound_to);
+            it != visitor.trivial_bindings_.end()) {
+          // This may be a trivial binding into a trivial binding.  In
+          // that case, unwrap the bindings until we find the earliest
+          // non-trivial binding.
+          bound_to = it->second;
+        } else {
+          break;
+        }
+      }


I think this might work instead and might be clearer:

for (auto it = visitor.trivial_bindings_.find(bound_to); it != visitor.trivial_bindings_.end(); it = visitor.trivial_bindings_.find(bound_to)) { bound_to = it->second; }

(Or the first clause can be moved to before the loop.)

Matter of preference, perhaps, if you consider it simpler than the above.

The same applies to the loop below as well.

I've gone back and forth on this one a bit. On the one hand, the current version has an extra level of nesting, but on the other hand, using for would require repeating the .find clause. In a perfect world, we would have while(init_statement; condition), but unfortunately (according to this stackoverflow post), that variation was left out for sake of simplicity.

I made another update to use Map<Id, Expr> instead of std::unordered_map<Id, Expr>, so that the loops could be rewritten as while(auto opt = lookup.Get(key)), avoiding the while(true) altogether.

Yeah, I checked to see if while (stmt1; condition) would work and it doesn't. I'm not against having any while(true), but avoiding a level of nesting is usually a good change.

slyubomirsky

Thank you for addressing the comments.

Lunderberg · 2023-10-25T14:13:28Z

No problem!

I'm going to merge this in, as it resolves a few test failures in unity head that resulted from a conflict between PR#15791 and PR#15904. Each PR passed CI on its own, but the combination of the two changes caused a few minor test failures. Since the test failures were resolved in this PR, that avoids needing to make a hotfix PR.

Lunderberg requested a review from slyubomirsky October 17, 2023 15:26

Lunderberg force-pushed the unity_canonicalize_var_to_dataflowvar_inside_dataflowblock branch 2 times, most recently from 86c3f0e to ddf30fb Compare October 17, 2023 19:28

Lunderberg mentioned this pull request Oct 20, 2023

[MultiGPU] Support pre-sharded model weights mlc-ai/mlc-llm#1096

Merged

Lunderberg force-pushed the unity_canonicalize_var_to_dataflowvar_inside_dataflowblock branch from c673f16 to b3f0a94 Compare October 20, 2023 17:12

Lunderberg force-pushed the unity_canonicalize_var_to_dataflowvar_inside_dataflowblock branch from 5ada41f to a488ac4 Compare October 24, 2023 19:29

slyubomirsky requested changes Oct 24, 2023

View reviewed changes

Lunderberg added 2 commits October 24, 2023 18:54

Simplify unwrapping of known bindings

85d512a

Updated to use Map<Id,Var>, to avoid while(true) loops

ba475d1

slyubomirsky approved these changes Oct 25, 2023

View reviewed changes

Lunderberg merged commit 4d19c8a into apache:unity Oct 25, 2023

Lunderberg deleted the unity_canonicalize_var_to_dataflowvar_inside_dataflowblock branch October 25, 2023 14:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Unity][Transform] Improved canonicalization of non-dataflow Var #15941

[Unity][Transform] Improved canonicalization of non-dataflow Var #15941

Lunderberg commented Oct 17, 2023

Lunderberg commented Oct 24, 2023

slyubomirsky left a comment

slyubomirsky Oct 24, 2023

Lunderberg Oct 24, 2023

slyubomirsky Oct 25, 2023

slyubomirsky Oct 24, 2023

Lunderberg Oct 24, 2023

slyubomirsky Oct 24, 2023

slyubomirsky Oct 24, 2023

Lunderberg Oct 24, 2023

Lunderberg Oct 25, 2023

slyubomirsky Oct 25, 2023

slyubomirsky left a comment

Lunderberg commented Oct 25, 2023

[Unity][Transform] Improved canonicalization of non-dataflow Var #15941

[Unity][Transform] Improved canonicalization of non-dataflow Var #15941

Conversation

Lunderberg commented Oct 17, 2023

Lunderberg commented Oct 24, 2023

slyubomirsky left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

slyubomirsky left a comment

Choose a reason for hiding this comment

Lunderberg commented Oct 25, 2023