How to use compute_at when there are more than one output with different shape?


#1

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?