[Bug] AlterLayout doesn't correctly wrap strided_slice with layout_transforms

[lazycal]

When doing AlterLayout pass on a conv followed by a strided_slice from NCHW4c to NCHW, the compiler does nothing to strided_slice, while (I think) the only correct behavior should be wrapping it with two layout_transforms. This leads to incorrect numerical result/crash at InferType, depending on the concrete input shape,
as shown in the following code snippet:

import tvm
from tvm import relay
from tvm.relay import transform
from tvm.relay.testing.temp_op_attr import TempOpAttr
import numpy as np

def test1(x_shape, w_shape):
    def before():
        x = relay.var("x", shape=x_shape)
        weight = relay.var("weight", shape=w_shape)
        y = relay.nn.conv2d(
            x,
            weight,
            kernel_size=(3, 3),
            padding=(1, 1),
            data_layout="NCHW4c",
            kernel_layout="OIHW4i4o",
        )
        y = relay.strided_slice(y, begin=[0, 0], end=[1, -1], strides=[1, 8])
        y = relay.Function([x, weight], y)
        return tvm.IRModule.from_expr(y)

    def alter_conv2d(attrs, inputs, tinfos, out_type):
        data, weight = inputs
        new_attrs = dict(attrs)
        new_attrs["data_layout"] = "NCHW"
        new_attrs["kernel_layout"] = "OIHW"
        return relay.nn.conv2d(data, weight, **new_attrs)

    with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d):
        be = transform.InferType()(before())
        print('='*40, 'before', '='*40)
        print(be)
        af = transform.AlterOpLayout()(be)
        print('='*40, 'after', '='*40)
        print(af)
        xnp = np.random.rand(*x_shape).astype(np.float32)
        wnp = np.random.rand(*w_shape).astype(np.float32)
        be_res = relay.create_executor("debug", be).evaluate()(xnp, wnp).numpy()
        af_res = relay.create_executor("debug", af).evaluate()(xnp, wnp).numpy()
        tvm.testing.assert_allclose(be_res, af_res, rtol=1e-3, atol=1e-3)

test1(x_shape=(1, 1, 1, 1, 4), w_shape=(9, 1, 3, 3, 4, 4)) # incorrect numerical result
# test1(x_shape=(1, 1, 1, 1, 4), w_shape=(11, 1, 3, 3, 4, 4)) # crash at InferType

The module before:

def @main(%x: Tensor[(1, 1, 1, 1, 4), float32], %weight: Tensor[(9, 1, 3, 3, 4, 4), float32]) -> Tensor[(1, 1, 1, 1, 4), float32] {
  %0 = nn.conv2d(%x, %weight, padding=[1, 1, 1, 1], kernel_size=[3, 3], data_layout="NCHW4c", kernel_layout="OIHW4i4o") /* ty=Tensor[(1, 9, 1, 1, 4), float32] */;
  strided_slice(%0, begin=[0, 0], end=[1, -1], strides=[1, 8], axes=None) /* ty=Tensor[(1, 1, 1, 1, 4), float32] */
}

and after:

def @main(%x: Tensor[(1, 1, 1, 1, 4), float32], %weight: Tensor[(9, 1, 3, 3, 4, 4), float32]) -> Tensor[(1, 1, 1, 1, 4), float32] {
  %0 = layout_transform(%x, src_layout="NCHW4c", dst_layout="NCHW") /* ty=Tensor[(1, 4, 1, 1), float32] */;
  %1 = layout_transform(%weight, src_layout="OIHW4i4o", dst_layout="OIHW") /* ty=Tensor[(36, 4, 3, 3), float32] */;
  %2 = nn.conv2d(%0, %1, padding=[1, 1, 1, 1], kernel_size=[3, 3]) /* ty=Tensor[(1, 36, 1, 1), float32] */;
  %3 = strided_slice(%2, begin=[0, 0], end=[1, -1], strides=[1, 8], axes=None) /* ty=Tensor[(1, 5, 1, 1), float32] */;
  layout_transform(%3, src_layout="NCHW", dst_layout="NCHW4c") /* ty=Tensor[(1, 1, 1, 1, 4), float32] */
}

Specifically, I am doing conv_NCHW4c_out[;,::8,...] (a 8-stride slice at the primal C dimension of NCHW4c). After altering layout into NCHW, the compiler does not wrap strided_slice with any layout_transformations nor adjust its attributes, so the semantic gets changed to conv_NCHW_out[:,::8,...], which means picking 1 every 8 elements, while what we need is to pick 4 elements every 4*8=32 elements for conv_NCHW_out

It seems that StridedSliceInferCorrectLayout is responsible for this.

BTW, the layout_transform seems weird in the latter IR:

%3 = strided_slice(%2, begin=[0, 0], end=[1, -1], strides=[1, 8], axes=None) /* ty=Tensor[(1, 5, 1, 1), float32] */;
layout_transform(%3, src_layout="NCHW", dst_layout="NCHW4c") /* ty=Tensor[(1, 1, 1, 1, 4), float32] */

The resultant tensor has smaller shape (1,1,1,1,4) than the before-transform one (1,5,1,1), and the reason I think is that (1,5,1,1) is not a valid input to be converted to the layout of NCHW4c, and I thought layout_transform should be able to detect and reject that?

Environment

• TVM: commit e334942
• CUDA version: 10.0
• System: Ubuntu 16.04
• GCC 5.4
• Build options: -DUSE_RELAY_DEBUG=ON -DUSE_CUBLAS=ON -DUSE_LLVM=ON -DUSE_CUDA=ON -DUSE_CUDNN=ON

[masahi]

Thanks, I'll take a look at this.