[GitHub] [tvm] lazycal commented on a change in pull request #10118: Fix LayoutRewriter

[GitBox <gi...@apache.org>]

lazycal commented on a change in pull request #10118:
URL: https://github.com/apache/tvm/pull/10118#discussion_r796811838



##########
File path: src/relay/transforms/transform_layout.h
##########
@@ -371,6 +375,17 @@ Expr LayoutRewriter(const Call& ref_call, const Array<Expr>& new_args, const Obj
   ICHECK_EQ(new_in.size(), new_in2.size())
       << "The number of input nodes should keep the same during alter_op_layout";
 
+  auto transform_layout = [&memorizer](Expr arg_item, const Layout& old_in, const Layout& old_in2,
+                                       const Layout& new_in, const Layout& new_in2) {
+    if (old_in2.Equals(old_in)) {  // the two transforms can be fused to one
+      arg_item = memorizer.Transform(arg_item, new_in, new_in2);
+    } else {
+      if (old_in.defined()) arg_item = memorizer.Transform(arg_item, new_in, old_in);
+      arg_item = memorizer.Transform(arg_item, old_in2, new_in2);

Review comment:
       Good catch. I thought that `old_in` and `old_in2` should be isomorphic, i.e., having the same structure (rank and subcoordinate factors, etc.) and only differing in how each axis is named (e.g., `NW` vs `NC`), given that they are describing the same tensor's layout. In this case, the transform can be applied. A concrete example: `new_in=NW8w`, `old_in=NW`, `old_in2=NC`, `new_in2=NC16c`, we will apply `NW8w->NW` and `NC->NC16c`, which is valid since layout_transform will work as long as the layout structure match the tensor shape. The net outcome is equivalent to a single transform `NC8c->NC16c`.
   
   However, I just hit a bizare case in BroadcastInferLayout that does not give isomorphic layouts:
   https://github.com/apache/tvm/blob/6a274af9cb4e7baf6d9dd06f0fb38bea5ac3154a/src/relay/transforms/infer_layout_utils.h#L224-L234
   This code path may expand the tensor's layout and assign `old_in2` to something with larger rank. For example, if the op is `a+b`, and originally `a`'s layout is `NCHW` and `b` is `W`, then its consumer (the broadcast node) will infer `old_in2=NCHW` for `b`. Now `W` and `NCHW` are not really isomorphic... I'm working on a fix, but this does sound pretty werid for me when you say a tensor with rank 1 is inferred with a layout with rank 4....




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org