How to implement rgb2yuv in tensor expression

Troubleshooting
#1

Hi, I’m a newbie to TVM, and found TVM could be a general solution to performance tuning not only for NN model. I tried to start from color space transformation and need your help.

With Numpy, rgb2yuv could be implemented like below, how can it be done with TVM? Numpy version:

def rgb2yuv(rgb):
    m = np.array([
        [0.29900, -0.147108,  0.614777],
        [0.58700, -0.288804, -0.514799],
        [0.11400,  0.435912, -0.099978]
    ])
    yuv = np.dot(rgb, m)
    yuv[:,:,1:] += 0.5
    return yuv

My not working version(wsn’t able to find a way to update [i,i,:] at once ):

w=h=32
c=3
rgb = te.placeholder((w,h,c), name='rgb')
def rgb2yuv(i,j,k):
    return rgb[i,j,0],rgb[i,j,1],rgb[i,j,2] // caculation formula ignored
yuv = te.compute((w,h,c), lambda i,j,k:rgb2yuv(i,j,k), name='yuv')
#2

OK, I think I found a workaround, not elegant, but work as I expected

w=640
h=480
c=3
rgb = te.placeholder((w,h,c), name='rgb')
def rgb2y(i,j,k):
    return rgb[i,j,0]*0.257+rgb[i,j,1]*0.504+rgb[i,j,2]*0.098
def rgb2u(i,j,k):
    return rgb[i,j,0]*-0.148+rgb[i,j,1]*-0.291+rgb[i,j,2]*0.439
def rgb2v(i,j,k):
    return rgb[i,j,0]*0.439+rgb[i,j,1]*-0.368+rgb[i,j,2]*-0.071
rgb2yuv=lambda i,j,k:te.if_then_else(k==0,rgb2y(i,j,k), te.if_then_else(k==1,rgb2u(i,j,k),rgb2v(i,j,k)))
yuv = te.compute((w,h,c), rgb2yuv, name='yuv')
s = te.create_schedule(yuv.op)
s[yuv].unroll(yuv.op.axis[2])
print(tvm.lower(s, [rgb, yuv], simple_mode=True))