[RELAY][PASS] FoldScaleAxis Forward by tqchen · Pull Request #2020 · apache/tvm

tqchen · 2018-10-28T02:14:46Z

This is a first serious attempt to implement an NN related optimization pass on relay. This PR contains the following changes:

Make squeeze's attrs numpy compatible
Add option to type resolver to write back checked_type to type_annotation if the original one is None
- There could be some debate on whether we should do so, but it is in general helpful in NN optimizations, where later phase of optimization will rely on knowing types of all arguments @jroesch
Support a general FoldScaleAxis Forward
- This pass assumes that the program is in dataflow form and inlines all lets
- So far only forward direction is implemented.

Hopefully, this can serve as an example of how optimizations can be done in NNVMv2(relay) IR, and how the infrastructure of relay makes writing optimization in a more principled fashion.

Goal

Fold the scaling of axis(usually caused by BatchNorm) into weight of conv2d in the future. For example

Old:

%1 = multiply(%x, %scale)
%2 = conv2d(%1, %w, data_layout="NHWC")

Transformed:

# scale weight's input channel
%1 = multiply(%w, expand_dims(%scale, axis=1, num_newaxis=2))
%2 = conv2d(%x, %1, data_layout="NHWC")

Further constant folding can fold the multiplication and we remove the scaling in the network.

The Algorithm

While so far only the forward direction is implemented. The general idea is that we transform Expr to tuple of (value, axes, scale), where the final result satisfies:

result = value
for i, k in enumerate(axes):
   k-ith dimension of result *= i-th dimension of scale

Then we can propagate this signal along and fold the scale if necessary. However, it is possible that certain scale may never be consumed if there is no dense/conv2d that follow multiplication.

In order to make sure all the scale we sent out can be consumed eventually, we run a backward "preparation phase", which propagates the demand of the potential axes scaling back to its input.

The new pass is more general than the FoldScaleAxis in nnvm

The new pass support arbitrary scaling of multiple axes(although the further implementation is necessary), which could be helpful in NHWCc case,
The new pass support folding forward into two separate path of conv2d, which was not possible in NNVM.

tqchen · 2018-10-28T02:15:49Z

@jroesch @ZihengJiang @yzhliu @zhiics @FrozenGene @merrymercy @srkreddy1238 @masahi please review

zhiics · 2018-10-28T04:10:28Z


    Parameters
    ----------
    data : relay.Expr


tvm.relay.Expr?

zhiics · 2018-10-28T04:18:07Z

+  if (!rhs.defined()) return rhs;
+  AxesSet ret;
+  size_t i = 0, j = 0;
+  while (i < lhs.size() && j < rhs.size()) {


Out of curiosity, are both lhs and rhs always sorted?

Yes, this is the requirement of axis set, Thanks for pointing this out, will add a comment block about it

zhiics · 2018-10-28T04:20:52Z

+
+/*!
+ * \brief The transform function, transform an old call to
+ *  new one given the new args.


zhiics · 2018-10-28T04:45:33Z

+  std::unordered_map<const Node*, AxesSet> message_;
+  // Update the message stored at node.
+  void Update(const Expr& node, const AxesSet& axes) {
+    // We run interection of messages:


s/interection/intersection

zhiics · 2018-10-28T05:04:29Z

    Expr new_e = ExprMutator::VisitExpr_(op);
-    if (!checked_type.same_as(new_e->checked_type_)) {
+    // new_call and new_var's code is only going to be valid for VarNode/CallNode.
+    // Compiler optimization will likely fold the these away for other nodes.


zhiics · 2018-10-28T18:26:58Z

+
+// Conv2D consumes the scale axis during transformation.
+STuple Conv2DForwardTransform(const Call& ref_call,
+                                const AxesSet& expected_axes,


expected_axes is unused. Should it be attached to rnode?

also indent

ZihengJiang · 2018-10-28T19:04:48Z

+    axis : None or List[int]
        Axes to remove.
-        If axes = [] or = None, remove all axis of dimensions 1.
+        If axes = None, remove all axis of dimensions 1.


change the comment also

I do not get what do you mean

axes should be changed to axis here

tqchen · 2018-10-28T20:10:48Z

Thanks @zhiics @ZihengJiang for helpful reviews, please check again

MarisaKirisame · 2018-10-28T20:05:39Z

+      }
+      ++j;
+    } else {
+      if (i >= base) {


MarisaKirisame · 2018-10-28T20:10:35Z

+  }
+
+  void VisitExpr_(const TupleGetItemNode* op) {
+    // pass, do nothing


why dont you visit inside? maybe there is opt ability inside

MarisaKirisame · 2018-10-28T20:11:58Z

+
+  void VisitExpr_(const IfNode* op) {
+    ExprVisitor::VisitExpr_(op);
+    // do pass through condition.


ZihengJiang · 2018-10-28T21:08:11Z

+// AddSub
+Array<AxesSet> AddSubForwardPrep(const Call& call, AxesSet out_axes) {
+  const auto* tlhs = call->args[0]->type_as<TensorTypeNode>();
+  const auto* trhs = call->args[1]->type_as<TensorTypeNode>();


Do we need to add check here?

type_as does the check

tqchen · 2018-10-28T21:24:09Z

Thanks @MarisaKirisame @ZihengJiang , I have make followup changes, please check again

* [RELAY][PASS] FoldScaleAxis Forward * Introduce helper function type_as * Update per review comment * Fix according to comments

[RELAY][PASS] FoldScaleAxis Forward

259546b

tqchen force-pushed the master branch from 0930cb4 to 259546b Compare October 28, 2018 03:18

Introduce helper function type_as

50f6abf

zhiics reviewed Oct 28, 2018

View reviewed changes

ZihengJiang reviewed Oct 28, 2018

View reviewed changes

Update per review comment

4c93f0e

MarisaKirisame reviewed Oct 28, 2018

View reviewed changes

ZihengJiang reviewed Oct 28, 2018

View reviewed changes

Fix according to comments

6e3c3d6

zhiics approved these changes Oct 28, 2018

View reviewed changes

ZihengJiang approved these changes Oct 28, 2018

View reviewed changes

MarisaKirisame approved these changes Oct 28, 2018

View reviewed changes

ZihengJiang merged commit 05ea601 into apache:master Oct 28, 2018

tqchen mentioned this pull request Oct 29, 2018

[RELAY][PASS] FoldScaleAxis Backward #2024

Merged

eqy pushed a commit to eqy/tvm that referenced this pull request Oct 29, 2018

[RELAY][PASS] FoldScaleAxis Forward (apache#2020)

7430467

* [RELAY][PASS] FoldScaleAxis Forward * Introduce helper function type_as * Update per review comment * Fix according to comments

eqy pushed a commit to eqy/tvm that referenced this pull request Oct 29, 2018

[RELAY][PASS] FoldScaleAxis Forward (apache#2020)

38d8e07

* [RELAY][PASS] FoldScaleAxis Forward * Introduce helper function type_as * Update per review comment * Fix according to comments

eqy pushed a commit to eqy/tvm that referenced this pull request Oct 29, 2018

[RELAY][PASS] FoldScaleAxis Forward (apache#2020)

38d867b

* [RELAY][PASS] FoldScaleAxis Forward * Introduce helper function type_as * Update per review comment * Fix according to comments

FrozenGene pushed a commit to FrozenGene/tvm that referenced this pull request Dec 27, 2018

[RELAY][PASS] FoldScaleAxis Forward (apache#2020)

8949716

* [RELAY][PASS] FoldScaleAxis Forward * Introduce helper function type_as * Update per review comment * Fix according to comments

ZihengJiang mentioned this pull request Feb 2, 2019

TVM 0.5 Release Note #2448

Closed

wweic pushed a commit to neo-ai/tvm that referenced this pull request Feb 20, 2019

[RELAY][PASS] FoldScaleAxis Forward (apache#2020)

d7d3e86

* [RELAY][PASS] FoldScaleAxis Forward * Introduce helper function type_as * Update per review comment * Fix according to comments

wweic pushed a commit to neo-ai/tvm that referenced this pull request Feb 20, 2019

[RELAY][PASS] FoldScaleAxis Forward (apache#2020)

ee7238e

* [RELAY][PASS] FoldScaleAxis Forward * Introduce helper function type_as * Update per review comment * Fix according to comments

wzh99 mentioned this pull request Aug 25, 2022

[TOPI][Bugfix] Make semantics of empty axis in squeeze consistent with Relay #12596

Merged

Conversation

tqchen commented Oct 28, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Goal

The Algorithm

Uh oh!

tqchen commented Oct 28, 2018

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tqchen Oct 28, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tqchen commented Oct 28, 2018

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tqchen commented Oct 28, 2018

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tqchen commented Oct 28, 2018 •

edited

Loading

tqchen Oct 28, 2018 •

edited

Loading