[Relay] Alter Op Layout #2150

merrymercy · 2018-11-21T08:37:16Z

Port alter_op_layout pass to relay. The new pass does forward rewrite in a single pass.

The major improvements over the old pass :

Enhance the layout inference for broadcast operators
- The old pass doesn't change the layout for broadcast operators to fit conv2d. It kind of hacks the batch_norm and relies on the bias being an attribute of conv2d. However, we cannot do this on relay. The new pass will make the layout of broadcast operators to fit conv2d correctly.
Memorize transformations so we won't create duplicated transformations.

This pr also did some cleaning to operators.

tqchen · 2018-11-21T17:53:55Z

src/relay/pass/simplify_bias_add.cc

+namespace tvm {
+namespace relay {
+
+class BiasAddSimplifier : public ExprMutator {


can we fold this into simplify inference?

No. Because it is not only for inference.

yzhliu · 2018-11-22T00:58:43Z

will review tonight

yzhliu · 2018-11-22T07:02:54Z

include/tvm/relay/op_attr_types.h

+ *  operator with other expressions.
+ *
+ * \param attrs The attribute of the node.
+ * \param inputs The arguments of this operator.


args & tinfos

yzhliu · 2018-11-22T07:40:57Z

src/relay/op/nn/pooling.cc

+      // if the new layout changes width or height dimension,
+      // fallback to old layout;
+      input = raw_layout;
+    }


since you change the param instead, then these checker & fallback should be no use.
but I guess we need to modify pool compute to let it know how to deal with w & h.

Sorry, there is a small mistake here. I fixed that. In fallback case, we don't change the param.

yzhliu · 2018-11-22T07:54:00Z

src/relay/op/tensor/transform.cc

+// relay.layout_transform
+std::pair<Layout, Layout> RemoveLeadingReduandantDimensions(
+    const Layout &src_layout, const Layout &dst_layout, size_t keep_size) {
+  // For a broadcast operator, if the input is a 3-dimensional tensor (64, 1, 1),


suggest to make it more clear, e.g., when broadcast (1, 64, 16, 16) with (64, 1, 1), we can still apply NCHW -> NCHW16c to the right tensor, by deleting N and applying normal CHW -> CHW16c layout transform. sth like that.

yzhliu · 2018-11-22T08:02:24Z

src/relay/op/tensor/transform.cc

+  // we can still apply 4-dimensional rule NCHW -> NCHW16c to it.
+  // In this case, we will delete leading redundant dimensions.
+  CHECK_GE(src_layout.ndim(), keep_size)
+    << "Can only apply layout transform rule for smaller tensor dimensions";


this error msg is not clear. suggest to print out src_layout and keep_size

yzhliu · 2018-11-22T09:20:27Z

src/relay/pass/alter_op_layout.cc

+
+  // old_in, old_out = op.infer(old_in)
+  bool success = false;
+  std::tie(old_in, old_out, success) = CallInfer(ref_call, old_in);


try to learn, how does batch_flatten's FInferCorrectLayout gets registered?

It is not defined. When we meet ops without FInferCorrectLayout such as batch_flatten. We will transform its inputs back to old layout.
https://github.com/merrymercy/tvm/blob/ad256f9d2fddf59c6c4224ab78be18437a13e6b2/src/relay/pass/alter_op_layout.cc#L99

yzhliu · 2018-11-23T01:56:40Z

src/relay/pass/alter_op_layout.h

+    lhs = rhs;
+  }
+
+  return Array<Array<Layout> > {{lhs, lhs}, {lhs}};


what if len(lhs) < len(rhs) ?

I agree that more checks and good fallback should be considered here.
Will let you know when it's ready for review.

yzhliu · 2018-11-23T02:09:30Z

src/relay/op/tensor/transform.cc

+    << "cannot convert from " << param->src_layout << " to " << param->dst_layout;
+
+  std::tie(src_layout, dst_layout) = RemoveLeadingReduandantDimensions(
+      src_layout, dst_layout, inputs[0]->shape.size());


RemoveLeadingReduandantDimensions has various checker, while it should still compile even though these restriction are not satisfied. e.g., broadcast(N16nCHW, CHW) -> broadcast(NCHW, CHW) will fail here, but actually we should be able to make it compile.

Also I'm a little worried about that, the approach allows layout and shape mismatch, and the mismatch is handled in layout_transform operator. In some sense, it's like other operators need to know how layout_transform works, then decide whether to infer a mismatch-shape layout. I would say, from the point view of system design, it's not quite good.

Also if later on there's another pass deal with layout, then it needs to carefully handle the mismatch stuff, and understand where does the mismatch come from.

Good points. I will clean this mismatch inside this pass.

jroesch · 2018-11-25T02:12:48Z

python/tvm/relay/op/transform.py

+
+
+def layout_transform(data, src_layout, dst_layout):
+    """Transform the layout of an tensor


Should be Transform the layout of a tensor.

tqchen · 2018-11-26T00:53:02Z

@merrymercy can you rebase the current PR, given the interest of this feature, it would be great to see if we can recover some of the end to end example benchmarks. #2163 should enable most of common networks already

merrymercy · 2018-11-26T16:40:27Z

@tqchen Some operators (winograd, NCHWc) are still missing. Will do it in follow up prs.

merrymercy · 2018-11-27T09:27:12Z

@yzhliu comments are addressed, please review again.
Now the layout transform for a broadcast operator is exactly "CHW" -> "CHW16c" rather than the old "NCHW" -> "NCHW16c".
The lhs-rhs length problem is fixed. Fallback mechanism is added.

kevinthesun · 2018-11-27T22:19:57Z

python/tvm/relay/op/_transform.py

+
+# concatenate
+@_reg.register_compute("concatenate")
+def concatenate_compute(attrs, inputs, output_type, target):


Do we also need a correct layout function for concatenate?

No. For concatenate, all its inputs will be transformed to old layouts. This is the default fallback case.

is it possible for concat to propagate layout?

Yes. We should add it

kevinthesun · 2018-11-28T03:59:08Z

src/relay/op/tensor/transform.cc

+                                       const Array<Layout> &in_layouts,
+                                       const Array<Array<IndexExpr>> &in_shapes) {
+  const ConcatenateAttrs* param = attrs.as<ConcatenateAttrs>();
+  CHECK_EQ(in_layouts.size(), 2);


Should allow more than two inputs?

This is wrong. Not ready for review..

merrymercy · 2018-11-28T06:22:35Z

@kevinthesun @tqchen please review again.

tqchen · 2018-11-28T18:03:54Z

python/tvm/relay/ir_pass.py

@@ -191,6 +191,23 @@ def simplify_inference(expr):
    return _ir_pass.simplify_inference(expr)


+def simplify_bias_add(expr):


Let us rename this to canonicalize_ops since bias_add lowering can be viewed as canonicalization step for ops, we might add more op canonicalizations later, (e.g. mean)..

tqchen · 2018-11-28T18:09:33Z

src/relay/pass/alter_op_layout.cc

+  TransformMemorizer() {}
+  explicit TransformMemorizer(NodePtr<Node> n) : NodeRef(n) {}
+
+  TransformMemorizerNode* operator->() {


You can just mark TransformMemorizerNode* operator->() const, to allow it to be accessed by the const

But I need to modify the content?

tqchen · 2018-11-28T18:11:02Z

src/relay/pass/alter_op_layout.cc

+
+// Make a transform CallNode
+Expr TransformLayout(Expr raw, Layout src_layout, Layout dst_layout) {
+  if (src_layout.Equals(dst_layout))


either keep the return in the single line, or enclose with { }

tqchen · 2018-11-28T18:12:06Z

src/relay/pass/alter_op_layout.cc

+}
+
+// Memorize layout transform so we can reuse internal transformed nodes
+class TransformMemorizerNode : public Node {


expand the comment a bit. what is the key of the map, what is the content

tqchen · 2018-11-28T18:12:21Z

src/relay/pass/alter_op_layout.cc

+  using TransformKey = std::tuple<const Node*, std::string, std::string>;
+  struct key_hash : public std::unary_function<TransformKey , std::size_t> {
+    std::size_t operator()(const TransformKey& k) const {
+      return std::hash<const Node*>()(std::get<0>(k)) ^


use dmlc::CombineHash

tqchen · 2018-11-28T18:12:48Z

src/relay/pass/alter_op_layout.cc

+    return static_cast<TransformMemorizerNode*>(node_.get());
+  }
+
+  // Transform layout with memorizer


document all the arguments and return value

I think the rule is functions in header files should be documented while functions in src files are not required to doc thoroughly, is it true?

that is true, but it is still helpful to document some of them to make code more readable

tqchen · 2018-11-28T18:13:00Z

src/relay/pass/alter_op_layout.cc

+  // Transform layout with memorizer
+  Expr Transform(Expr raw, const Layout& src_layout, const Layout& dst_layout) {
+    if (src_layout.Equals(dst_layout))
+      return raw;


keep raw in the same line as if

tqchen · 2018-11-28T18:14:05Z

src/relay/pass/alter_op_layout.cc

+
+// Call FInferCorrectLayout of an op.
+// Return inferred_input_layout, inferred_output_layout, success
+std::tuple<Array<Layout>, Array<Layout>, bool> CallInfer(


document the fields of the return value and the arguments

tqchen · 2018-11-28T18:17:11Z

src/relay/pass/alter_op_layout.cc

+                          const NodeRef& ctx) {
+  std::vector<LayoutAlternatedExpr> inputs;
+  std::vector<Expr> normal_new_args;
+  Array<Array<IndexExpr>> input_shapes;


space between > for compatibility to old compilers

tqchen · 2018-11-28T18:18:12Z

src/relay/pass/alter_op_layout.cc

+  // NOTE: discard the "const" qualifier
+  TransformMemorizer memorizer = Downcast<TransformMemorizer>(ctx);
+
+  // fill incomplete state and expand tuple


Explain that we always expect LayoutAlternatedExpr, and this is used to convert the arguments into LayoutAlternatedExpr.

tqchen · 2018-11-28T18:20:19Z

@yzhliu @jroesch @kevinthesun @ajtulloch can you please take a look at this PR if you have time?

merrymercy · 2018-11-29T02:28:05Z

comments are addressed!

yzhliu

good to me. ping @jroesch

tqchen · 2018-11-29T18:09:30Z

@merrymercy please rebase against the latest master

srkreddy1238 · 2018-12-01T08:46:42Z

Not sure if I am the only one see this issue in local build.
** submodule update is also done.

[ 52%] Building CXX object CMakeFiles/tvm.dir/src/relay/op/nn/convolution.cc.o
In file included from /home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/../op_common.h:14:0,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/resize.cc:11:
/home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/../../pass/alter_op_layout.h: In function ‘tvm::Array<tvm::Array<tvm::relay::Layout> > tvm::relay::ElemwiseArbitraryLayout(const tvm::Attrs&, const tvm::Array<tvm::relay::Layout>&, const tvm::Array<tvm::relay::Layout>&, const tvm::Array<tvm::Array<HalideIR::Expr> >&)’:
/home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/../../pass/alter_op_layout.h:55:72: error: no matching function for call to ‘tvm::Array<tvm::relay::Layout>::Array(size_t, tvm::relay::Layout&)’
   return Array<Array<Layout> >{Array<Layout>(old_in_layouts.size(), ret), {ret}};
                                                                        ^
In file included from /home/srk/work/NN/TF-DMLC/relay/tvm/3rdparty/HalideIR/src/ir/Expr.h:10:0,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/include/tvm/expr.h:9,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/include/tvm/packed_func_ext.h:17,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/include/tvm/api_registry.h:12,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/include/tvm/relay/base.h:9,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/include/tvm/relay/op.h:16,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/resize.cc:6:
/home/srk/work/NN/TF-DMLC/relay/tvm/3rdparty/HalideIR/src/tvm/node/container.h:180:3: note: candidate: tvm::Array<T, <template-parameter-1-2> >::Array(const std::vector<T>&) [with T = tvm::relay::Layout; <template-parameter-1-2> = void]
   Array(const std::vector<T>& init) { // NOLINT(*)
   ^~~~~
/home/srk/work/NN/TF-DMLC/relay/tvm/3rdparty/HalideIR/src/tvm/node/container.h:180:3: note:   candidate expects 1 argument, 2 provided
/home/srk/work/NN/TF-DMLC/relay/tvm/3rdparty/HalideIR/src/tvm/node/container.h:173:3: note: candidate: tvm::Array<T, <template-parameter-1-2> >::Array(std::initializer_list<_Tp>) [with T = tvm::relay::Layout; <template-parameter-1-2> = void]
   Array(std::initializer_list<T> init) { // NOLINT(*)
   ^~~~~
/home/srk/work/NN/TF-DMLC/relay/tvm/3rdparty/HalideIR/src/tvm/node/container.h:173:3: note:   candidate expects 1 argument, 2 provided
/home/srk/work/NN/TF-DMLC/relay/tvm/3rdparty/HalideIR/src/tvm/node/container.h:166:3: note: candidate: template<class IterType> tvm::Array<T, <template-parameter-1-2> >::Array(IterType, IterType)
   Array(IterType begin, IterType end) {
   ^~~~~
/home/srk/work/NN/TF-DMLC/relay/tvm/3rdparty/HalideIR/src/tvm/node/container.h:166:3: note:   template argument deduction/substitution failed:
In file included from /home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/../op_common.h:14:0,
                 from /home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/resize.cc:11:
/home/srk/work/NN/TF-DMLC/relay/tvm/src/relay/op/image/../../pass/alter_op_layout.h:55:72: note:   deduced conflicting types for parameter ‘IterType’ (‘long unsigned int’ and ‘tvm::relay::Layout’)
   return Array<Array<Layout> >{Array<Layout>(old_in_layouts.size(), ret), {ret}};

srkreddy1238 · 2018-12-01T11:28:27Z

Ignore above issue. looks like submodule update didn't sync all changes.

* [RELAY] Finish alter op pass * [RELAY] AlterOpLayout Pass * fix broadcast operators * fix broadcast operators * fix broadcast operators * Support concatenate * address comments * address comments * add comments * rebase

merrymercy force-pushed the relay-alter-op-layout branch from 16d6294 to 18d138a Compare November 21, 2018 10:02

tqchen self-assigned this Nov 21, 2018

tqchen added the status: need review label Nov 21, 2018

tqchen requested changes Nov 21, 2018

View reviewed changes

merrymercy force-pushed the relay-alter-op-layout branch from 49d22b3 to 1747a21 Compare November 21, 2018 20:47

yzhliu reviewed Nov 22, 2018

View reviewed changes

yzhliu reviewed Nov 23, 2018

View reviewed changes

jroesch reviewed Nov 25, 2018

View reviewed changes

merrymercy force-pushed the relay-alter-op-layout branch 2 times, most recently from ede6608 to 13338b9 Compare November 27, 2018 09:24

kevinthesun reviewed Nov 27, 2018

View reviewed changes

kevinthesun reviewed Nov 28, 2018

View reviewed changes

merrymercy force-pushed the relay-alter-op-layout branch from 81934f9 to eb299ac Compare November 28, 2018 06:17

merrymercy force-pushed the relay-alter-op-layout branch from 15abe1b to 51a20ba Compare November 28, 2018 06:24

tqchen requested changes Nov 28, 2018

View reviewed changes

merrymercy force-pushed the relay-alter-op-layout branch 2 times, most recently from 45f0754 to 8f9e396 Compare November 29, 2018 02:25

kevinthesun approved these changes Nov 29, 2018

View reviewed changes

tqchen approved these changes Nov 29, 2018

View reviewed changes

yzhliu approved these changes Nov 29, 2018

View reviewed changes

merrymercy added 2 commits November 29, 2018 13:39

[RELAY] Finish alter op pass

8a26c94

[RELAY] AlterOpLayout Pass

f55b89e

merrymercy added 7 commits November 29, 2018 13:41

fix broadcast operators

7972796

fix broadcast operators

627f568

fix broadcast operators

d918875

Support concatenate

dc098ca

address comments

70ac9d3

address comments

1f10fcd

add comments

e9cf2bb

merrymercy force-pushed the relay-alter-op-layout branch from f2af1ef to e9cf2bb Compare November 29, 2018 21:41

rebase

8849f9a

merrymercy merged commit 2a5656b into apache:master Nov 30, 2018

merrymercy deleted the relay-alter-op-layout branch November 30, 2018 01:13

ZihengJiang mentioned this pull request Feb 1, 2019

TVM 0.5 Release Note #2448

Closed



		def layout_transform(data, src_layout, dst_layout):
		"""Transform the layout of an tensor

		@@ -191,6 +191,23 @@ def simplify_inference(expr):
		return _ir_pass.simplify_inference(expr)


		def simplify_bias_add(expr):

[Relay] Alter Op Layout #2150

[Relay] Alter Op Layout #2150

Conversation

merrymercy commented Nov 21, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yzhliu commented Nov 22, 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

merrymercy Nov 26, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merrymercy Nov 26, 2018 • edited Loading

Choose a reason for hiding this comment

yzhliu Nov 23, 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 Nov 26, 2018

merrymercy commented Nov 26, 2018

merrymercy commented Nov 27, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merrymercy Nov 28, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merrymercy Nov 28, 2018 • edited Loading

Choose a reason for hiding this comment

merrymercy commented Nov 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

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 Nov 28, 2018

merrymercy commented Nov 29, 2018

yzhliu left a comment

Choose a reason for hiding this comment

tqchen commented Nov 29, 2018

srkreddy1238 commented Dec 1, 2018

srkreddy1238 commented Dec 1, 2018

merrymercy commented Nov 21, 2018 •

edited

Loading

merrymercy Nov 26, 2018 •

edited

Loading

merrymercy Nov 26, 2018 •

edited

Loading

yzhliu Nov 23, 2018 •

edited

Loading

merrymercy commented Nov 27, 2018 •

edited

Loading

merrymercy Nov 28, 2018 •

edited

Loading

merrymercy Nov 28, 2018 •

edited

Loading