Where does the layout transform of each op happen during alter_op_layout pass?

comaniac · November 6, 2020, 1:50am

When you see the tensors changed from 4D to 5D, the corresponding conv2d op has already been changed from NCHW to NCHWc; otherwise the type won’t match. This is called “alter op layout”. Specifically, the function you pointed returns the altered NCHWc op:

github.com

apache/incubator-tvm/blob/main/python/tvm/topi/x86/conv2d_alter_op.py#L114


                new_attrs["data_layout"],
                new_attrs["out_layout"],
                out_dtype,
            ],
            topi_tmpl,
        )
        dispatch_ctx.update(target, new_workload, cfg)
    else:
        assert _NCHWc_matcher.match(data_layout)
        assert _OIHWio_matcher.match(kernel_layout)
    return relay.nn.contrib_conv2d_nchwc(*inputs, **new_attrs)

if topi_tmpl == "conv2d_NCHWc_int8.x86":
    # TODO(@icemelon9, @anijain2305): Need to support data layout NHWC with kernel layout HWIO
    assert data_layout == "NCHW" and kernel_layout == "OIHW"
    if cfg.is_fallback:
        _get_default_config_int8(
            cfg, data_tensor, kernel_tensor, strides, padding, out_dtype, False, data_layout
        )

    batch_size, in_channel, height, width = get_const_tuple(data_tensor.shape)

Accordingly, your graph changed from 4D -> conv2d_NCHW to 4D -> layout_tranform -> 5D -> conv2d_NCHWc.