[Bug][AutoTVM] autotvm.task.extract_from_program fails on a conv2d operator

[mehrdadh]

Running autotvm.task.extract_from_program on this relay module which is a conv2d with NHWC format and HWOI and int8 input dtype shows this error:

relay module

#[version = "0.0.5"]
def @main(%input: Tensor[(1, 32, 32, 1), int8]) {
  nn.conv2d(%input, meta[relay.Constant][0], padding=[0, 0, 0, 0], kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWOI", out_layout="NHWC", out_dtype="int32")
}

#[metadata]
{
  "root": 1, 
  "nodes": [
    {
      "type_key": ""
    }, 
    {
      "type_key": "Map", 
      "keys": [
        "relay.Constant"
      ], 
      "data": [2]
    }, 
    {
      "type_key": "Array", 
      "data": [3]
    }, 
    {
      "type_key": "relay.Constant", 
      "attrs": {
        "_checked_type_": "0", 
        "data": "0", 
        "span": "0", 
        "virtual_device_": "4"
      }
    }, 
    {
      "type_key": "VirtualDevice", 
      "attrs": {
        "device_type_int": "-1", 
        "memory_scope": "5", 
        "target": "0", 
        "virtual_device_id": "-1"
      }
    }, 
    {
      "type_key": "runtime.String"
    }
  ], 
  "b64ndarrays": [
    "P6G0lvBAXt0AAAAAAAAAAAEAAAAAAAAABAAAAAAIAQADAAAAAAAAAAMAAAAAAAAADAAAAAAAAAABAAAAAAAAAGwAAAAAAAAA9gX89vcF+wj6AwQC+/gF/gP/APj3/f79CP8GA/oA/gEJ/fn/+AkAAfYABvwF+gcFAf8A+vf/CfsBAQH4BAD2Bfv+9wj6/QH39/cFBgQI9gP8CfwCBwcC+/YIAgQD/f39CQX+BfcCCfoE9v35"
  ], 
  "attrs": {"tvm_version": "0.9.dev0"}
}

target
c -keys=arm_cpu,cpu -link-params=0 -mcpu=cortex-m7

traceback

python/tvm/autotvm/task/relay_integration.py:90: in extract_from_program
    return extract_from_multiple_program([mod], [params], target, ops=ops)
python/tvm/autotvm/task/relay_integration.py:156: in extract_from_multiple_program
    tsk = create(task_name, args, target=target)
python/tvm/autotvm/task/task.py:485: in create
    sch, _ = ret.func(*args)
python/tvm/autotvm/task/task.py:240: in __call__
    return self._default_func(*args, **kwargs)
python/tvm/autotvm/task/task.py:246: in _default_func
    out = self.fcompute(*args, **kwargs)
python/tvm/autotvm/task/topi_integration.py:165: in wrapper
    node = topi_compute(cfg, *args)
python/tvm/topi/arm_cpu/conv2d.py:514: in conv2d_nhwc_dsp
    return conv2d_nhwc_dsp_compute(cfg, data, kernel, strides, padding, dilation, out_dtype)
python/tvm/topi/arm_cpu/mprofile/dsp/conv2d.py:123: in conv2d_nhwc_dsp_compute
    filter=None if cfg.is_fallback else lambda x: x.size[-1] % 4 == 0,
python/tvm/autotvm/task/space.py:730: in define_split
    return self._add_new_transform(SplitSpace, name, axes, policy, **kwargs)
python/tvm/autotvm/task/space.py:837: in _add_new_transform
    self._entity_map[name] = space[0]
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = Split(policy=factors, product=1, num_outputs=2) len=0, index = 0

    def __getitem__(self, index):
        """Get an entity of the space by index
    
        Parameters
        ----------
        index: int
    
        Returns
        -------
        transform entity
        """
>       return self.entities[index]
E       IndexError: list index out of range

cc @elvin-n @Icemist

[mehrdadh]

Error coming from this line which has a TODO:

tvm/python/tvm/topi/arm_cpu/mprofile/dsp/conv2d.py

Line 122 in 588679e

# TODO: check case with in_channels.value % 4 != 0 with AutoTVM

[guberti]

Recently stumbled across this and found the issue. Turns out this issue stems from the fact the kernel in the above case takes up 2.25 words, while our tensor intrinsic only works when the kernel takes up a whole number of words. I'm working on a new schedule in #12969, which should avoid this problem.