Hi, everyone, I got a problem here. In Tesorflow, there is a operator names “split”, it can split a tensor to more than one tensor and with different shape. For example, the input tensor’s shape is [6, 16, 16], the split axis is 0, the size of splits is [1, 2, 3], which means the output tensors’ shape are [1, 16, 16], [2, 16, 16] and [3, 16, 16]. So I need to use compute_at to combine the first axis of input and outputs. But here comes an error:
Could you please tell me how to archive this operator?
Maybe I didn’t make it clear. The operator runs on specific hardware, which cache is not big enough, so I have to use tvm.split() to split axis that make sure I can move the data to cache, as the last example, my cache just can store 16 fp16 data, so I hope I can use tvm to generate the code just like this:
for(i0, 0, 6){
for(i1, 0, 16){
input_data.cache[0] = input_data[(i0 * 16) + i1]
if(i0 < 1){
res_0[i0*16 + i1] = input_data.cache[0]
}
if(i0 >=1){
if(i0 < 3){
res_1[i0*16 + i1 - 16] = input_data.cache[0]
}
}
if(i0 >= 3){
res_2[i0*16 + i1 - 48] = input_data.cache[0]
}
}
}
How do I use tvm to get this code?