[Relay][AlterOp] Improving support for broadcast layout alteration. #4040

anijain2305 · 2019-10-01T05:57:40Z

Improve broadcast layout support while Altering layouts

Support scenarios where original operands were of type [N, H, W, C] and [N, H, W, 1]. In this case, we might have NCHW16c coming for 1 operand. However, the other operand does not have enough C dimension. To reuse broadcasting, we would want to use NCHW1c for the second operand.
Support scenarios where original operands were of type [N, H, W, C] and [C]. In this case, while transforming the layout, we expand dims to make C go to NHWC, and then use the modified layout of the first operator to call the layout transform.

@yzhliu @vinx13 @tqchen @zhiics @ZihengJiang

yzhliu · 2019-10-01T16:42:58Z

include/tvm/data_layout.h

@@ -210,6 +210,20 @@ class Layout : public NodeRef {
    return ct;
  }

+  /*! \return Concatenation of all primal axes */
+  inline std::string get_primal_axes() const {


suggest to wrap it as Layout so later we can deal with Layout instead of std::string

vinx13 · 2019-10-01T23:20:42Z

shouldn't bias_add be expanded to expand_dims and add in CanonicalizeOps pass before AlterOpLayout?

anijain2305 · 2019-10-01T23:27:30Z

shouldn't bias_add be expanded to expand_dims and add in CanonicalizeOps pass before AlterOpLayout?

That will solve bias_add problem. But, the problem will still exist if we directly use the add operator that uses broadcast.

I thought of doing it in canocalizeOps (where bias_add is converted to add with expand_dims), that will definitely be simpler. But, will miss cases where we create network with just add.

vinx13 · 2019-10-01T23:44:24Z

@anijain2305 Do you mean scenarios like adding a 1-D bias vector to NCHW tensor? This doesn't follow the broadcast semantic.
Can you also explain why NHWC is a problem? If you add a NHWC tensor with a C vector, it should be fine since the vector can be broadcasted and no layout change or expand_dims should be needed.

anijain2305 · 2019-10-02T05:28:30Z

NCHW (Currently working)

Original

bias -> [C]
data -> [NCHW]
conv = conv(data, weight) -> [NCHW]
bias_add(conv, bias)

After CanonicalizeOps

bias -> [C]
data -> [NCHW]
conv = conv(data, weight) -> [NCHW]
new_bias = expand_dims(bias, 1, 2) --> [C,1,1] or CHW
add(conv, new_bias) -> [NCHW]

After AlterOpLayout

bias -> [C]
data -> [NCHW]
data_nchwc = LT(data, NCHW, NCHW16c) --> [NCHW16c]
conv = conv(data, weight_transformed) -> [NCHW16c]
new_bias = expand_dims(bias, 1, 2) --> [C,1,1] or CHW
new_bias_chwc = LT(new_bias, CHW, CHW16c) -> [CHW16c]
add(conv, new_bias_chwc) -> [NCHW16c]

This is good for now. Now, lets look at NHWC

NHWC (Does not work - Happens in TFLite)

Original

bias -> [C]
data -> [NHWC]
conv = conv(data, weight) -> [NHWC]
bias_add(conv, bias)

After CanonicalizeOps

bias -> [C]
data -> [NHWC]
conv = conv(data, weight) -> [NHWC]
add(conv, bias) -> [NHWC]
Nothing happens here because the bias is already at the last C axis.

After AlterOpLayout

bias -> [C]
data -> [NHWC]
data_nchwc = LT(data, NHWC, NCHW16c) --> [NCHW16c]
conv = conv(data, weight_transformed) -> [NCHW16c]
new_bias_nchwc = LT(new_bias, C, NCHWC16c) -> [NCHW16c] // The problem
add(conv, bias) -> [NCHW16c]

The problem is that LT does not support going from C to NCHW16c. This PR inserts an expand dims to make bias go to NHWC, and then call layout transform. So, this PR bring following transformation
bias -> [C]
data -> [NHWC]
data_nchwc = LT(data, NHWC, NCHW16c) --> [NCHW16c]
conv = conv(data, weight_transformed) -> [NCHW16c]
new_bias = expand_dims(bias, 0, 3) --> [NHWC]
new_bias_nchwc = LT(new_bias, NHWC, NCHWC16c) -> [NCHW16c] // This is supported now
add(conv, bias) -> [NCHW16c]

Questions to ask

Should LT (layout transform) support C to NCHW16c conversion internally? This would avoid expand_dims but make LT implementation more complex.

anijain2305 · 2019-10-03T22:59:39Z

@yzhliu @vinx13 @tqchen @zhiics @ZihengJiang Ping for review. I am done from my side.

kevinthesun · 2019-10-04T17:09:44Z

@anijain2305 Thank you for adding this! I have the same issue when trying to internally transfer conv2d to NCHWc layout for tf model.

One question is that can we directly use BinaryBroadcastLayout for bias_add?

kevinthesun · 2019-10-04T17:51:20Z

After discussion, I think this change is necessary. LGTM

yzhliu · 2019-10-04T18:19:02Z

include/tvm/data_layout.h

+   * \param dst_layout The dst layout to which current layout has to be expanded.
+   * \return The expanded Layout.
+   */
+  inline Layout ExpandLayout(const Layout& dst_layout) {


suggest to name it ExpandPrimal

yzhliu · 2019-10-04T18:22:15Z

include/tvm/data_layout.h

+    // 2) Now, add the primal axis of the current layout.
+    for (auto src_axis : operator->()->axes) {
+      if (LayoutAxis::Get(src_axis).IsPrimal()) {
+        new_src_layout_str += src_axis->var->name_hint;
+      }
+    }


can we simply do new_src_layout_str += this->name() ?

yzhliu · 2019-10-04T18:24:58Z

include/tvm/data_layout.h

+    for (auto axis : operator->()->axes) {
+      if (!LayoutAxis::Get(axis).IsPrimal()) {
+        // 1) Find the corresponding dual axis
+        auto dual_axis = std::toupper(LayoutAxis::Get(axis).name()[0]);


LayoutAxis::Get(axis).ToPrimal() or ToDual()

The function is private. Can't do that.

yzhliu · 2019-10-04T18:26:16Z

include/tvm/data_layout.h

+    std::string new_src_layout_str = "";
+    for (auto dst_axis : dst_layout->axes) {
+      if (LayoutAxis::Get(dst_axis).IsPrimal()) {
+        if (this->IndexOf(LayoutAxis::Get(dst_axis)) == -1) {


suggest to use this->Contains

yzhliu · 2019-10-04T18:31:43Z

include/tvm/data_layout.h

+        bool is_shape_one = false;
+        if (auto* shape_int = shape_val.as<IntImm>()) {
+          if (shape_int->value == 1) {
+            new_small_layout += "1";
+            is_shape_one = true;
+          }
+        }
+
+        // 4) b) If shape is not 1, retain the factor.
+        if (!is_shape_one) {
+          auto new_shape_val = FactorOf(LayoutAxis::Get(dual_axis));
+          new_small_layout += std::to_string(new_shape_val);
+        }


it might be better to have another helper function ReplaceFactor implement the functionality above, and move this AdjustSubordinateFactors to alter_op_layout, as this is not as general for tvm::Layout.

yzhliu · 2019-10-06T04:22:06Z

src/relay/pass/alter_op_layout.h

+      auto dual_axis = LayoutAxis::Get(axis).ToPrimal().name()[0];
+
+      // 2) Find the index of this dual axis in old_layout
+      int old_axis = old_layout.IndexOf(LayoutAxis::Get(dual_axis));


actually you don't need to do .name() first then ::Get, ToPrimal() already returns LayoutAxis

I tried that. But it does not work because some of the methods are private. I didn't want to change the class description.

…pache#4040)

* master: (21 commits) [Fix][VM] Fix VM invoke with set_params (apache#4079) [QNN] Refactor fixed point multiplication in requantize (apache#4073) Fix match case in Python-side expr functor (apache#4037) Hide symbols from dependent libraries if HIDE_PRIVATE_SYMBOLS is ON. (apache#4041) Add gradient for log-softmax (apache#4069) [DOC] Fix typos in tutorials (apache#4066) dicrease the complexity of CalcDep from exponential to linear (apache#4053) [Relay][AlterOp] Minor refactor. (apache#4064) [Relay][AlterOp] Improving support for broadcast layout alteration. (apache#4040) Add parses support for zeros_like tflite operator (apache#4042) [Bugfix][TF] reset graph after getting tag of savedmodel (apache#4055) [Relay][VM] Add more passes to VMCompiler (apache#4058) [Relay][VM] Add autotvm context when compile (apache#4062) [Bugfix] Fix target host for vm compiler (apache#4057) [Relay][Training] Add gradient for Crossentropy (apache#3925) [llvm] switch to use Align for llvm trunk (apache#4051) [Relay][TopHub] Add switch to disable TopHub download (apache#4015) [Relay][Op] Add instance norm op (apache#4004) [QNN][Relay] Calling Dialect passes from inside Relay Build API. (apache#3971) [RELAY/PASS] Fix the extent for the post_stmt in the loop partition (apache#3734) ...

anijain2305 force-pushed the broadcast branch from 385c7f8 to 6185291 Compare October 1, 2019 06:02

yzhliu reviewed Oct 1, 2019

View reviewed changes

anijain2305 mentioned this pull request Oct 3, 2019

[Relay][CanonicalizeOps] Make Bias_add shape same as the other operand #4050

Closed

anijain2305 force-pushed the broadcast branch 4 times, most recently from e57b1f9 to 821c402 Compare October 3, 2019 21:34

kevinthesun approved these changes Oct 4, 2019

View reviewed changes

yzhliu reviewed Oct 4, 2019

View reviewed changes

anijain2305 force-pushed the broadcast branch from 821c402 to b81fb16 Compare October 5, 2019 00:28

[Relay][AlterOp] Improving support for broadcast layout alteration.

28aa9ff

anijain2305 force-pushed the broadcast branch from b81fb16 to 28aa9ff Compare October 5, 2019 21:19

vinx13 approved these changes Oct 6, 2019

View reviewed changes

yzhliu approved these changes Oct 6, 2019

View reviewed changes

yzhliu merged commit d703fb4 into apache:master Oct 6, 2019

yzhliu reviewed Oct 6, 2019

View reviewed changes

anijain2305 mentioned this pull request Oct 7, 2019

[AlterOpLayout][x86] NHWC to NCHWc conv support. #4080

Merged

anijain2305 added a commit to anijain2305/tvm that referenced this pull request Oct 17, 2019

[Relay][AlterOp] Improving support for broadcast layout alteration. (a…

bff8b88

…pache#4040)

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

[Relay][AlterOp] Improving support for broadcast layout alteration. (a…

47b3459

…pache#4040)

tqchen mentioned this pull request Nov 8, 2019

[RELEASE][DRAFT] TVM v0.6 Release candidate #4259

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Relay][AlterOp] Improving support for broadcast layout alteration. #4040

[Relay][AlterOp] Improving support for broadcast layout alteration. #4040

anijain2305 commented Oct 1, 2019 •

edited

Loading

yzhliu Oct 1, 2019

vinx13 commented Oct 1, 2019

anijain2305 commented Oct 1, 2019

vinx13 commented Oct 1, 2019 •

edited

Loading

anijain2305 commented Oct 2, 2019 •

edited

Loading

anijain2305 commented Oct 3, 2019

kevinthesun commented Oct 4, 2019 •

edited

Loading

kevinthesun commented Oct 4, 2019

yzhliu Oct 4, 2019

yzhliu Oct 4, 2019

yzhliu Oct 4, 2019

anijain2305 Oct 5, 2019

yzhliu Oct 4, 2019

yzhliu Oct 4, 2019

yzhliu Oct 6, 2019

anijain2305 Oct 6, 2019

[Relay][AlterOp] Improving support for broadcast layout alteration. #4040

[Relay][AlterOp] Improving support for broadcast layout alteration. #4040

Conversation

anijain2305 commented Oct 1, 2019 • edited Loading

Choose a reason for hiding this comment

vinx13 commented Oct 1, 2019

anijain2305 commented Oct 1, 2019

vinx13 commented Oct 1, 2019 • edited Loading

anijain2305 commented Oct 2, 2019 • edited Loading

NCHW (Currently working)

Original

After CanonicalizeOps

After AlterOpLayout

NHWC (Does not work - Happens in TFLite)

Original

After CanonicalizeOps

After AlterOpLayout

Questions to ask

anijain2305 commented Oct 3, 2019

kevinthesun commented Oct 4, 2019 • edited Loading

kevinthesun commented Oct 4, 2019

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

anijain2305 commented Oct 1, 2019 •

edited

Loading

vinx13 commented Oct 1, 2019 •

edited

Loading

anijain2305 commented Oct 2, 2019 •

edited

Loading

kevinthesun commented Oct 4, 2019 •

edited

Loading