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!