[RELAY] Basic block normal form #6152

eric-haibin-lin · 2020-07-27T22:37:41Z

Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/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.

As mentioned in https://discuss.tvm.ai/t/basic-block-normal-form/5908/5, the requirements for basic block normal form are:

D0, Define a block as a group of expressions implied by the scope structure.
D1: Each graph node can only belong to a single block
D2: For values that are being used in multiple blocks, they have to be referred by a Var which is defined in a block

We rely on analysis::DependencyGraph to compute the dependency among nodes, and use LCA to compute the scope of each node. As we traverse the dependency graph, we record the expression with a "lifted scope" if its original scope is different from the result of LCA of nodes depending on it. For instance, a variable shared by the true_branch and false_branch in an If statement. If any expression with lifted scope is found, the function violates the basic block normal form.

The expressions with lifted scope need to be bound by let binding. It is done with a pass almost the same as the ANF pass. Same as ANF we use a let list per scope to help construct let bindings of expressions. The main difference with ANF pass is that we only push expressions to the let list if the expression's scope is lifted, otherwise we return the original expression. When we come across If/function node, we get the expression generated from the let list of the corresponding scope.

@MarisaKirisame would you suggest we merge the two passes into a single class?

jroesch · 2020-07-27T23:11:26Z

cc @MarisaKirisame

tqchen · 2020-07-28T16:24:37Z

Thanks @eric-haibin-lin , it would also be useful to have a verification analysis util that verifies the IR is in the basic block form

ANSHUMAN87 · 2020-07-29T05:35:03Z

@eric-haibin-lin : Thanks for the PR! I believe this feature will be really helpful for certain detection & mutation process.
Please let me know, once your PR is ready for review. Would like to get a hold on it!

tqchen · 2020-07-29T17:05:17Z

cc @MarisaKirisame @mbrookhart @electriclilies @jroesch @junrushao1994 would be great if you can help to take a look at the PR

MarisaKirisame · 2020-08-01T11:41:34Z

src/relay/transforms/pass_util.h

+ *  \param lifted_exprs the output set of expressions whose scope is lifted due to dependency
+ */
+void CalcScope(const DependencyGraph& dg,
+               std::unordered_map<DependencyGraph::Node*, Scope>* expr_scope,


why use pointer instead of returning?

@MarisaKirisame thanks for the review! For BBlock I want to get a set of expressions whose scope should be lifted and rewritten with let bindings. I found CalcScope already performs finding the LCA for all nodes in the dependency graph, so i just added a few lines of code so that the set of corresponding expressions are also returned by this function. However, setting the return type to std::pair<std::unordered_map<node, scope> std::unordered_set<expr>> causes cpplint to fails due to the long length. So I instead used pointers for both outputs. Please let me know the preferred way for returning both outputs. Thanks!

I think we should disable cppint and use tuple, or just do a few typedef to shorten stuff.

Ok I'll update the code

MarisaKirisame · 2020-08-03T03:41:15Z

tests/python/relay/test_analysis_basic_block_normal_form.py

+    }
+    """
+    print(body)
+    assert not check_basic_block_normal_form(body)


is_basic_block_normal_form test if something is or is not in bbnf.
check_basic_block_normal_form should assert it, and perhaps give more useful error message

👍 Agreed.

MarisaKirisame · 2020-08-03T04:14:37Z

src/relay/transforms/to_basic_block_normal_form.cc

+/* Fill expressions based on each scope's let list. Different from FillANF,
+ * only expressions with lifted scope will be pushed to the let list.
+ */
+class FillBasicBlock : ExprFunctor<Expr(const Expr&, const Var&)> {


can you reuse the Fill in to_anf? the code is almost identically the same, and duplicating it will make maintenance significantly harder.

mbrookhart

Mostly this looks great, thank you!

I few nitpicks, a bunch of print statements that need to be removed, a question on naming, and one OOP design question.

Feel free to argue against my suggestions, nothing but removing the prints is really required.

mbrookhart · 2020-08-03T15:38:20Z

src/relay/transforms/pass_util.h

+  // the parent scope
+  Scope parent;
+  // the corresponding let list which holds all let bindings in the scope
+  std::shared_ptr<LetList> ll = std::make_shared<LetList>();


Could we use a more expressive variable name, maybe let_list?

mbrookhart · 2020-08-03T15:45:27Z

src/relay/transforms/to_basic_block_normal_form.cc

+namespace tvm {
+namespace relay {
+
+Expr ToBasicBlockNormalFormAux(const Expr& e) {


Why Basic Block normal form? Is there a more advanced version coming? Or should we just call it BlockNormalForm?

I guess it's because of https://en.wikipedia.org/wiki/Basic_block ?

:) I missed this naming influence in the conversations @tqchen and I had a few months ago. Thanks! That makes things a lot clearer for me.

mbrookhart · 2020-08-03T15:48:14Z

tests/python/relay/test_analysis_basic_block_normal_form.py

+    body = relay.If(v2, true_branch, false_branch)
+    body = relay.Let(v1, relay.add(x, one), body)
+    func = relay.Function([x], body)
+    print(func)


remove print

mbrookhart · 2020-08-03T15:48:31Z

tests/python/relay/test_analysis_basic_block_normal_form.py

+    }
+    """
+    print(body)
+    assert not check_basic_block_normal_form(body)


👍 Agreed.

mbrookhart · 2020-08-03T15:49:34Z

tests/python/relay/test_pass_to_basic_block_normal_form.py

+    y = op.add(x, x)
+    z = op.add(y, y)
+    f = relay.Function([], op.add(z, z))
+    print(f)


remove print

mbrookhart · 2020-08-03T15:54:56Z

tests/python/relay/test_pass_to_basic_block_normal_form.py

+    print(top)
+    print('bblock=')
+    print(bblock)
+    assert check_basic_block_normal_form(bblock)


I'd love an expected graph on this test, but I don't think it's critical

mbrookhart · 2020-08-03T15:55:11Z

tests/python/relay/test_analysis_basic_block_normal_form.py

+      add(%shared, %shared)
+    }
+    """
+    print(body)


Remove print

mbrookhart · 2020-08-03T15:55:17Z

tests/python/relay/test_analysis_basic_block_normal_form.py

+      add(%shared, %shared)
+    }
+    """
+    print(body)


remove print

mbrookhart · 2020-08-03T16:01:19Z

src/relay/transforms/pass_util.h

+  static Expr ToANormalForm(const Expr& e, const DependencyGraph& dg, NodeScopeMap* node_scope);
+
+  // For basic block normal form, bind expressions only if the original expression's
+  // scope should be lifted
+  static Expr ToBasicBlockNormalForm(const Expr& e, const DependencyGraph& dg,
+                                     NodeScopeMap* node_scope, ExprSet* lifted);


I really like that you're reusing the infrastructure from A-Normal form to do this, there's a lot of overlap. I don't love that this class violates the Single Responsiblity Principle and that you have to introduce branches in a couple of the recursive functions to handle the two different cases. Given that this class has a private constructor, I'm not sure it 100% matters, but there might be a cleaner solution?

thanks for bringing this up. I was also thinking about how to best leverage the existing code without too much ad-hoc code or duplication. I think adding an inclusion/exclusion argument to Fill is still acceptable, and the Fill constructor is private anyway

mbrookhart · 2020-08-03T16:02:05Z

src/relay/transforms/to_a_normal_form.cc

-      args.push_back(VisitExpr(a));
-    }
-    return Compound(e, Call(VisitExpr(c->op), args, c->attrs, c->type_args), v);
+Expr Fill::VisitExpr_(const CallNode* c, const Var& v) {


Why remove final from all of these functions?

final is still used in the header at declaration https://en.cppreference.com/w/cpp/language/final#Syntax

eric-haibin-lin

thanks @mbrookhart for the review. I have updated some code as suggested

eric-haibin-lin · 2020-08-03T18:26:24Z

src/relay/transforms/to_basic_block_normal_form.cc

+namespace tvm {
+namespace relay {
+
+Expr ToBasicBlockNormalFormAux(const Expr& e) {


I guess it's because of https://en.wikipedia.org/wiki/Basic_block ?

eric-haibin-lin · 2020-08-03T18:27:37Z

src/relay/transforms/to_a_normal_form.cc

-      args.push_back(VisitExpr(a));
-    }
-    return Compound(e, Call(VisitExpr(c->op), args, c->attrs, c->type_args), v);
+Expr Fill::VisitExpr_(const CallNode* c, const Var& v) {


final is still used in the header at declaration https://en.cppreference.com/w/cpp/language/final#Syntax

eric-haibin-lin · 2020-08-03T18:50:34Z

src/relay/transforms/pass_util.h

+  static Expr ToANormalForm(const Expr& e, const DependencyGraph& dg, NodeScopeMap* node_scope);
+
+  // For basic block normal form, bind expressions only if the original expression's
+  // scope should be lifted
+  static Expr ToBasicBlockNormalForm(const Expr& e, const DependencyGraph& dg,
+                                     NodeScopeMap* node_scope, ExprSet* lifted);


thanks for bringing this up. I was also thinking about how to best leverage the existing code without too much ad-hoc code or duplication. I think adding an inclusion/exclusion argument to Fill is still acceptable, and the Fill constructor is private anyway

mbrookhart

LGTM

MarisaKirisame · 2020-08-03T20:44:23Z

Thanks @eric-haibin-lin @mbrookhart @junrushao1994

MarisaKirisame · 2020-08-03T20:48:36Z

I dont think it break single responsibility - the code is doing conversion and a single configuration param denote which conversion it is. Another way to think about it is that that visitor is declaring a 'scoped mutator' and toanf/tobbnf is two subclass of it.
We can refactor to the latter later if needed.

mbrookhart · 2020-08-03T20:56:53Z

I dont think it break single responsibility - the code is doing conversion and a single configuration param denote which conversion it is. Another way to think about it is that that visitor is declaring a 'scoped mutator' and toanf/tobbnf is two subclass of it.
We can refactor to the latter later if needed.

That is exactly breaking the SRP, it's a class that does two mutually-exclusive conversions :) The SRP-compliant way to do this would be exactly what you said, a base class and two subclasses. That being said, since we're not allowing anyone to inherit from this class, and the constructor is private, it's probably not the end of the world, I'm cool with leaving it the way it is.

MarisaKirisame · 2020-08-03T20:58:40Z

We can allow it if needed!

MarisaKirisame · 2020-08-04T00:20:28Z

@eric-haibin-lin please at me when it is done. I will merge it then.

MarisaKirisame · 2020-08-04T04:38:48Z

Thanks @eric-haibin-lin @mbrookhart @junrushao1994

* initial commit * refactor utils * add util * revert anf test * update test * fix logging * fix scope bug * complete tests * remove logging * revert refactoring * add one more test case * fix missing var binding * fix test * fix lint * fix lint * fix clang-format * fix lint * fix lint * commit missing code * add analysis api * fix lint * fix lint * lint * add test for func * address CR * fix typo * fix return type * fix lint * refactor classes * fix lint * remove prints * address comments Co-authored-by: Ubuntu <ubuntu@ip-172-31-42-138.ec2.internal>

Ubuntu added 13 commits July 26, 2020 01:08

initial commit

1b2e476

refactor utils

39c23bc

add util

3cb49b4

revert anf test

ce6acb6

update test

204fac2

fix logging

1eaa9ae

fix scope bug

4b2c2b9

complete tests

90e89d6

remove logging

4b23546

revert refactoring

30e9a0e

add one more test case

04eba08

fix missing var binding

2ce94fc

fix test

6b974ef

tqchen added the status: need review label Jul 27, 2020

fix lint

940265c

eric-haibin-lin marked this pull request as draft July 27, 2020 22:45

Ubuntu added 4 commits July 27, 2020 22:48

fix lint

2b8b813

fix clang-format

b2e785b

fix lint

8834614

fix lint

ce05edd

commit missing code

b9aaf5a

Merge remote-tracking branch 'origin/master' into bblock

6f06c79

Ubuntu added 4 commits July 29, 2020 18:47

add analysis api

d150a6c

fix lint

5947334

fix lint

29f6738

lint

98c723f

tqchen self-assigned this Jul 31, 2020

address CR

a478e6b

MarisaKirisame requested changes Aug 1, 2020

View reviewed changes

Ubuntu added 3 commits August 1, 2020 17:39

fix typo

2e1c9a5

fix return type

99502bd

fix lint

9e00305

MarisaKirisame requested changes Aug 3, 2020

View reviewed changes

Ubuntu added 2 commits August 3, 2020 05:58

refactor classes

aebc621

fix lint

d066986

mbrookhart requested changes Aug 3, 2020

View reviewed changes

tqchen changed the title ~~Basic block normal form~~ [RELAY] Basic block normal form Aug 3, 2020

Ubuntu added 2 commits August 3, 2020 17:18

remove prints

7a23134

address comments

5ad16fd

eric-haibin-lin commented Aug 3, 2020

View reviewed changes

mbrookhart approved these changes Aug 3, 2020

View reviewed changes

MarisaKirisame approved these changes Aug 3, 2020

View reviewed changes

MarisaKirisame merged commit b6db7e3 into apache:master Aug 4, 2020

eric-haibin-lin deleted the bblock branch August 4, 2020 04:52

ZihengJiang mentioned this pull request Sep 25, 2020

TVM v0.7 Release Note Candidate #6486

Closed

[RELAY] Basic block normal form #6152

[RELAY] Basic block normal form #6152

Conversation

eric-haibin-lin commented Jul 27, 2020 • edited Loading

jroesch commented Jul 27, 2020

tqchen commented Jul 28, 2020

ANSHUMAN87 commented Jul 29, 2020

tqchen commented Jul 29, 2020 • edited Loading

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

mbrookhart 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

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

eric-haibin-lin 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

mbrookhart left a comment

Choose a reason for hiding this comment

MarisaKirisame commented Aug 3, 2020

MarisaKirisame commented Aug 3, 2020

mbrookhart commented Aug 3, 2020 • edited Loading

MarisaKirisame commented Aug 3, 2020

MarisaKirisame commented Aug 4, 2020

MarisaKirisame commented Aug 4, 2020

eric-haibin-lin commented Jul 27, 2020 •

edited

Loading

tqchen commented Jul 29, 2020 •

edited

Loading

mbrookhart commented Aug 3, 2020 •

edited

Loading