[RELAY][PASS] FoldScaleAxis Forward #2020

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

include/tvm/relay/expr_functor.h

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

python/tvm/relay/op/transform.py

@@ -49,17 +49,17 @@ def transpose(data, axes=None):
    return _make.transpose(data, list(axes))


-def squeeze(data, axes=None):
+def squeeze(data, axis=None):
    """Squeeze axes in the array.

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


tvm.relay.Expr?

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

src/relay/pass/fold_scale_axis.cc

+  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

src/relay/pass/fold_scale_axis.cc

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


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

src/relay/pass/fold_scale_axis.cc

+  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

src/relay/pass/type_infer.cc

    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.


src/relay/pass/fold_scale_axis.cc

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

src/relay/pass/fold_scale_axis.cc

+
+// 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

src/relay/pass/fold_scale_axis.cc

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

python/tvm/relay/op/transform.py

        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

include/tvm/relay/op.h

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

src/relay/pass/pattern_util.h

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


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

src/relay/pass/fold_scale_axis.cc

+  }
+
+  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

src/relay/pass/fold_scale_axis.cc

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


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

src/relay/pass/fold_scale_axis.cc

+// 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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RELAY][PASS] FoldScaleAxis Forward #2020

[RELAY][PASS] FoldScaleAxis Forward #2020

tqchen commented Oct 28, 2018 •

edited

Loading

tqchen commented Oct 28, 2018

zhiics Oct 28, 2018

zhiics Oct 28, 2018

tqchen Oct 28, 2018

zhiics Oct 28, 2018

zhiics Oct 28, 2018

zhiics Oct 28, 2018

zhiics Oct 28, 2018

MarisaKirisame Oct 28, 2018

ZihengJiang Oct 28, 2018

tqchen Oct 28, 2018 •

edited

Loading

ZihengJiang Oct 28, 2018

tqchen Oct 28, 2018

tqchen commented Oct 28, 2018

MarisaKirisame Oct 28, 2018

MarisaKirisame Oct 28, 2018

tqchen Oct 28, 2018

MarisaKirisame Oct 28, 2018

ZihengJiang Oct 28, 2018

tqchen Oct 28, 2018

tqchen commented Oct 28, 2018

[RELAY][PASS] FoldScaleAxis Forward #2020

[RELAY][PASS] FoldScaleAxis Forward #2020

Conversation

tqchen commented Oct 28, 2018 • edited Loading

Goal

The Algorithm

tqchen commented Oct 28, 2018

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

tqchen Oct 28, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Oct 28, 2018

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

tqchen commented Oct 28, 2018

tqchen commented Oct 28, 2018 •

edited

Loading

tqchen Oct 28, 2018 •

edited

Loading