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

[mkatanbaf]

Running autotvm.task.extract_from_program on this relay module which is a conv1d with NWC format and WOI and int8 input dtype shows this error:

relay module
#[version = "0.0.5"]
def @main(%input: Tensor[(3, 12, 10), int16]) {
nn.conv1d(%input, meta[relay.Constant][0], dilation=[2], kernel_size=[4], data_layout="NWC", kernel_layout="WOI", out_layout="NWC", 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": [
"P6G0lvBAXt0AAAAAAAAAAAEAAAAAAAAAAwAAAAAQAQAEAAAAAAAAABgAAAAAAAAACgAAAAAAAACABwAAAAAAAAYA+P8AAP7//v/4/wEA+v/7/wYACAD7/wgA/v/7////AAD8/wQAAgAEAP7/BgAGAAEA+f8GAPf//f/5//n/BQD8//b/AAAEAAEA/f8HAPv/CQAHAPz/BgD5/wUAAgD6////BwAEAAIABAAFAAIAAAACAAkA+f8AAAMABQAEAAIAAgD6////AAD4//b/+P/7/wMA/////wMA+v/+//z/+//3/wYAAQACAAcA/P8JAAUA9v8AAPb/BwAFAPj/AQD3/wcAAwABAPv/BwD4//v/CQADAAMACQAIAP7/BAADAAQACAD5/wMA//8HAAYA+v8HAP3/+P8FAP//BQAEAP7/CAAFAPj///8DAAgA+v/9//f/BwAGAAkACQD4/wAA+/8AAPv/AAAIAPb/9/8IAAkA/f8IAAYAAAAFAP7/CQD8/wYABAD8//v/CAD3/wUA+v/5//z/BwAEAP7/+f8FAP7/+//4//b/CQD3/wEAAgD+/wgA///8/wQA+f8IAPv///8GAAEA+P////f//f8CAP7/CQAJAP7/BgD//wEA/f8IAP7//P/+//7/AAADAPf/BQAEAAUABQD4//r/BwAAAAYAAwD7//r/+f8BAP7////7//v/AgD8//v/CQAAAPj/BAD9//f//P8EAAcA/P/4//v/9//7/wcA9/8CAAEAAQD+/wIA/v/4//v/AQAIAAkA+f/4////+f//////BQAGAPr/AAACAAMABAD3/////f/3/wUACAAAAP3/BQD+/wcAAAD3/wgA+/8BAPf/AQD5//n/AAAFAAkA+//4/wEA///5//n/CQD7//n/BgAJAAYAAQABAP7/9v8GAPf//P/+//r//P8AAAEA+//4//r/BgD9//7//v/4/wUA///+//n/9/8DAAQA/P/6//r//v/3/wUA9/8JAP7/AQD6//z//f8CAP7///8AAAcA/f8BAPz//P/2/wIAAAAGAAkAAQABAAIA+f/6/wQAAgAAAAcABAD8//n/AwD9//v/+P/6//7/+f8EAAEACQAIAAEA/v8EAAgA/f/3//n/+/////r//v/+//n/9v/6//j/AAAJAP3/+/8BAAEABQAEAAUA9v8HAAYA/v/3//r//P8BAPz/CQAEAAEA/v8IAPn/BQAGAPf/BgADAPz/BwD+/wAABAABAPz/BwD5/wUA+f8CAPj/9v/8//3/CQD3//b//P8HAAgA9/8IAP7/CAD+/wQAAgACAPn//P/2//b/BAD8/wcA///3/wgA+P/5//7//v/3/wEABAD5//r/AgABAPb/CQAFAAgABgD8//n/CQD6//j/AgAAAP//AQAFAPz/BQAGAAcA/v8JAAUAAAADAP7/+P////b/AwAHAAYA/f/5//7/9/8BAPz/AwD9/wcACQACAPf/9//9//7/CAAFAAMACAD//wcA///6//3//P/+//b/+f8AAP3/CQAAAAEACQD3/wMACAAEAAQA+f/9//n/AQADAPv/9v8GAPb////8//n/+P8IAPr/9/8IAAIA/P8EAAcAAwAEAPr/+v8GAAAA+v/2//r//P8HAPz/+P8GAP//BQAAAAUACQD2//j/CAABAAcA+//3/wUABwABAAUA/v8JAPj/BQAJAPn/9//6/wcAAgD3//v/BgAIAPv//f///wIA9v8GAAQA+f8HAAYACQD///f/9/8JAAYACQD+//3/BQAJAP7/AwD9//f/BwAIAAUA+f/3//7/BQD5/wgA9//3/wIA9//4/wMAAAD7/wYAAAD6//v/+P8GAPn/CAAAAAEAAwAEAPj/AAD9//z/CQAEAAAA///7//7/AQAAAPj//f8CAAYA9v/+//z///8CAAMACAABAPv/BgAFAP//BgD9/wEAAQD9//v/BQAAAAYAAQD7//j///8FAAEACAD+/wMAAAAIAPf/AQD6/wAA///5////BAD4//z/CAD4//b/AQD2/wYA/P8GAAAABQD9/wEA9//+//z/CAAAAPz//f/7/wgA+v/8//r/CAAAAPr/BwAEAPn/CQD9//n/9/8EAPb/BAAJAP3/9v8EAP//BQD2/wIA+P/8//3/+/8IAPr/9v8BAAMA/v8DAP7/AwAJAAIA+f8BAP7//v8JAAcACQACAPv//f/+/wkA/P/7/wcA+f8JAAIA+/8GAPb/AwAAAPz/9//4//n//P8CAPr/BwAFAAcA9/8EAAkA/f8BAPv/CQAAAAMA+v/3//n/9/8CAAAA///2//z/BgD9//r//v/8//j/BAAHAP//9//5//n/9//3/wgAAAAGAAcABgAEAAcA+f8AAPv/BwD5/wAA/P8JAAkAAgAHAAIAAAD9/wIABAAHAAkA/f8JAP7/AQD3/wEA/f8DAAgA9/8BAPz/BQABAP7/CQD+//f/AgAJAPr/+P/+//r/BgD///b/BAAJAAUABAD5//v/9v/3/wEABAD6//v/CQD6//f/AAAGAAIAAgD5/wYAAQD3/wcA//8IAPb/BwD7/wkA9/8GAPz/CQD5/wUACAADAAkAAgD+/wMA+//2/wIA/v/5//f/9v8CAAMAAwD5/w=="
],
"attrs": {"tvm_version": "0.9.dev0"}
}

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

traceback
File "issue.py", line 102, in
tasks = tvm.autotvm.task.extract_from_program(mod["main"], {}, TARGET)
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/relay_integration.py", line 90, in extract_from_program
return extract_from_multiple_program([mod], [params], target, ops=ops)
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/relay_integration.py", line 156, in extract_from_multiple_program
tsk = create(task_name, args, target=target)
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/task.py", line 485, in create
sch, _ = ret.func(*args)
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/task.py", line 240, in call
return self._default_func(*args, **kwargs)
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/task.py", line 246, in _default_func
out = self.fcompute(*args, **kwargs)
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/topi_integration.py", line 165, in wrapper
node = topi_compute(cfg, *args)
File "/home/mohamadk/tvm/tvm/python/tvm/topi/arm_cpu/conv1d.py", line 32, in conv1d_nwc_dsp
return conv1d_nwc_dsp_compute(cfg, data, kernel, strides, padding, dilation, out_dtype)
File "/home/mohamadk/tvm/tvm/python/tvm/topi/arm_cpu/mprofile/dsp/conv1d.py", line 108, in conv1d_nwc_dsp_compute
filter=None if cfg.is_fallback else lambda x: x.size[-1] % 4 == 0,
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/space.py", line 730, in define_split
return self._add_new_transform(SplitSpace, name, axes, policy, **kwargs)
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/space.py", line 837, in _add_new_transform
self._entity_map[name] = space[0]
File "/home/mohamadk/tvm/tvm/python/tvm/autotvm/task/space.py", line 92, in getitem
return self.entities[index]
IndexError: list index out of range

code

import numpy as np
import tvm
from tvm import relay
import tvm.contrib.utils
from tvm.contrib.download import download_testdata
from tvm.relay.backend import Runtime


data_shape = (3, 12, 10)
kernel_size = 4
num_filter = 24
strides = 1
padding = 0
dilation = 2
dtype = "int16"
data_layout = "NWC"
kernel_layout = "WOI"
    
ishape = data_shape
wshape = (kernel_size, data_shape[-1], num_filter)
weight_data = np.random.randint(low=-10, high=10, size=wshape, dtype=dtype)

input1 = relay.var("input", relay.TensorType(ishape, dtype))

if kernel_layout == "WOI":
    weight1 = relay.const(np.moveaxis(weight_data, 1, -1))
else:
    weight1 = relay.const(weight_data)

out1 = relay.op.nn.conv1d(
    input1,
    weight1,
    kernel_size=kernel_size,
    strides=strides,
    padding=padding,
    dilation=dilation,
    data_layout="NWC",
    kernel_layout=kernel_layout,
    out_dtype="int32",
    out_layout="NWC",
)
mod = tvm.IRModule.from_expr(relay.Function([input1], out1))
print(mod.astext(show_meta_data=True))

inputs = {"input": np.random.randint(low=-128, high=127, size=ishape, dtype=dtype)}

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

pass_context = tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True})
with pass_context:
    tasks = tvm.autotvm.task.extract_from_program(mod["main"], {}, TARGET)

cc @elvin-n @Icemist

[mkatanbaf]

similar issue is reported for conv2d: #11216