How to optimize the realization of indexing?

shyouku · April 23, 2018, 9:43am

Indexing is very common in some ops. When indices are passed with python list, list can not use tvm.var as index in lambda expression. For example, So tvm.select is used:

If input shape is [5,10,10] and permut is [1,3,4,1,2], lower IR can be

However we expect select can be just at the i0 axis, and the best expression may be more efficient without i0 loop like this: #i0 = 0, copy first 10 * 10 for (i1,0,10) { for (i2,0,10) { … } } #i0 = 1, copy 3rd 10 * 10 for (i1,0,10) { for (i2,0,10) { … } } … So how to realize to print IR like this?

merrymercy · April 24, 2018, 2:40pm

Looks like simply unroll i0 can give you the desired results.

I have used unroll + select + inline to pass a constant matrix to tvm. This matrix has many zeros and tvm can simplify the expression to simplest form.

github.com

dmlc/tvm/blob/035b6254e681027ca12d4ee66a2502a3e23c54a5/topi/python/topi/mali/conv2d.py#L68-L81


def const_array(data, name):
""" convert an const array to tvm tensor"""
row, col = data.shape
dtype = str(data.dtype)


def select_array(i, j):
    now = tvm.const(0.0, dtype)
    for ii in range(row):
        for jj in range(col):
            now = tvm.select(tvm.all(i % row == ii, j % col == jj),
                             tvm.const(data[ii][jj], dtype),
                             now)
    return now
return tvm.compute(data.shape, select_array, name=name)