ConvertLayout: Dealing with depthwise-convolutions

Development
#1

This a followup discussion for a question which came up in https://github.com/apache/tvm/pull/14272.

The ConvertLayout pass allows to alter the data- and kernel-layouts of specific operators in the relay graph by providing a mapping as follows:

mod = relay.transform.ConvertLayout(
  {
    "nn.conv2d": ["NCHW", "HWIO"],
    "nn.avg_pool2d": ["NCHW", "default"],
  }
)(mod)

I ran into a situation, where I would like to different kernel layouts for regaular and depthwise convolutions to achieve optimal performance. Due to the non-existant nn.depthwise_conv2d operator in favour of the generic nn.conv2d relay op supporting groups!=1 this can not easily described using the aformentioned mapping.

While there of course exists workarounds for this, by overwriting FTVMConvertOpLayout manually (see https://github.com/apache/tvm/blob/main/python/tvm/relay/op/contrib/arm_compute_lib.py#L115) we should establish a way to deal with this in a consistant manner. This would be especially useful for the tvmc compile ... --desired-layouts ... feature but should also be applicable when using relay.transform.ConvertLayout directly.

Talking about special cases (such as depthwise convolutions) there are of course additional canditates worth mentioning:

  • Pointwise (1x1) Convolutions prefer channels-last (NHWC) layouts (because height/width never interact) for improved locality, while
  • Depthwise Convolutions prefer channels-first (NCHW) layouts (because channels never interact) for improved locality
  • Optional: Dealing with custom constraints, such as stride_h=stride_w=2, which might cause failures when no operator strategies are registered for the provided layouts.

While I am a aware that these details can be handled automatically i.e. during AlterOpLayout, I would love to get more user-level control on the layout transformation passes.

I am looking forward to any ideas on how to deal with this in the future!

CC @lhutton1 @guberti

1 Like
#2

Thanks for creating this @PhilippvK, apologies for only just following up.

While there of course exists workarounds for this, by overwriting FTVMConvertOpLayout manually (see https://github.com/apache/tvm/blob/main/python/tvm/relay/op/contrib/arm_compute_lib.py#L115 ) we should establish a way to deal with this in a consistant manner. This would be especially useful for the tvmc compile ... --desired-layouts ... feature but should also be applicable when using relay.transform.ConvertLayout directly.

I agree that it would be useful to distinguish between different types of convolution from the ConvertLayout pass itself, although this seems more difficult to fix since it is tightly coupled to how operators are declared in Relay - I’ve had a think about how to work around this but haven’t come up with a sensible solution yet.

Perhaps, for now, it’s simpler to consider working around the issue in the TVMC front-end by making use of pass context as you mentioned. From the user POV, ‘special’ operator names could be provided to the --desired-layout-ops option, such as nn.conv2d_pointwise, nn.conv2d_depthwise which override the layout(s) specified by nn.conv2d when the conditions their names refer to are met.