Do we still need relay.nn.contrib_conv2d_NCHWc op?

masahi · January 11, 2022, 8:28pm

Anyone knows why we have relay.nn.contrib_conv2d_NCHWc op, while there is no op specific for other layouts such as NHWC? Since layout_transform op is inserted and precomputed if possible automatically, I don’t see a need for this op.

The only difference between the regular conv2d op is that it is using a specific type relation to bypass the shape checking on weight:

github.com

apache/tvm/blob/813136401a11a49d6c15e6013c34dd822a5c4ff6/src/relay/op/nn/convolution.cc#L508


- **data**: Input is 5D packed tensor.
- **weight**: 6D packed tensor.


- **out**:  Output is 5D packed tensor
)code" TVM_ADD_FILELINE)
    .set_attrs_type<Conv2DAttrs>()
    .set_num_inputs(2)
    .add_argument("data", "Tensor", "The input tensor.")
    .add_argument("weight", "Tensor", "The weight tensor.")
    .set_support_level(10)
    .add_type_rel("Conv2DNCHWc", Conv2DWinogradRel<Conv2DAttrs>)
    .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ConvInferCorrectLayout<Conv2DAttrs>);


// Positional relay function to create depthwise conv2d NCHWc operator
// used by frontend FFI.
TVM_REGISTER_GLOBAL("relay.op.nn._make.contrib_depthwise_conv2d_NCHWc")
    .set_body_typed([](Expr data, Expr weight, Array<IndexExpr> strides, Array<IndexExpr> padding,
                       Array<IndexExpr> dilation, int groups, IndexExpr channels,
                       Array<IndexExpr> kernel_size, String data_layout, String kernel_layout,
                       String out_layout, DataType out_dtype) {
      return MakeConv<Conv2DAttrs>(data, weight, strides, padding, dilation, groups, channels,

The same question applies to the dense_prepack op introduced relatively recently in https://github.com/apache/tvm/pull/7404. If we add weight_layout attribute to the dense op and update its type relation accordingly, I don’t see a need for introducing the new op.

@yzhliu @vinx13 @comaniac

masahi · January 11, 2022, 8:51pm

More data point:

CUDA conv2d int8 kernel uses the NCHWc layout, but it uses the regular relay.nn.conv2d op https://github.com/apache/tvm/blob/d1dafbd6d050ddcd673f567c5ff720a6b6419cfe/python/tvm/topi/cuda/conv2d_alter_op.py#L89-L121
On the other hand, the x86 int8 kernel uses a custom implementation of weight-layout transform and relay.nn.contrib_conv2d_nchwc op: https://github.com/apache/tvm/blob/d1dafbd6d050ddcd673f567c5ff720a6b6419cfe/python/tvm/topi/x86/conv2d_alter_op.py#L171-L181

If the weight-layout transform cannot be expressed by relay.layout_transform op, I can see the reason why we need a layout-specific op. But I’d rather update relay.layout_transform to support more transform types and not to introduce another ops.