[Relay][topi/nn] Ways to get layout of each operator

I posted this as an issue in GitHub before, and reopen it here according to @tqchen’s suggestion.

I’m using TVM as front end, and need to know the layout of each operator in implement. Operators like conv2d have “data_layout”, “kernel_layout” in it’s attrs field, which I could use directly, but many operators like tvm.relay.nn.batch_norm/tvm.relay.nn.relu/tvm.topi.add etc. don’t have the layout field in attrs. Is there a way to quickly get the layout for those operators? Thanks very much.

Since we don’t define layout attribute for such operators, nor the layout attribute in the tensor, there’s no symbolic representation for their tensor layouts. In other words, you cannot get something like

%1 = nn.dense(...); /* output layout = NC */

The way I can think of is just infer the layout manually from input tensor or the op with the layout attribute. Maybe this is also another opportunity to revisit the topic about whether to put layout attribute to the tensor, so that we could have a pass to calculate the input/output layouts of each op in symbolic way, and can further improve the layout conversion mechanism.

cc @yzhliu @anijain2305 @tqchen

We had discussion before regarding whether to put layout as part of type system. At the time the conclusion was not to complicate type, but have layout a separate property during the layout inference pass. Part of the reasons is layout is not a must have, and is not quite well-defined, e.g., why the name “NCHW” but not “ABCD”. but maybe we shall expose an API to allow users get the inferred layout after running layout inference pass.

3 Likes

Thanks @comaniac, that’s what I’m doing right now.