tvm-vectorization.py

import tvm

def main(n):
    a = tvm.placeholder((n,), dtype='int8')
    c = tvm.compute((n,), lambda i: a[i])

    s = tvm.create_schedule(c.op)

    bx, tx = s[c].split(s[c].op.axis[0], factor=4)

    s[c].bind(bx, tvm.thread_axis('blockIdx.x'))

    s[c].vectorize(tx)

    with tvm.build_config(data_alignment=16):
        f = tvm.build(s, [a,c], 'cuda')

if __name__=='__main__':
    main(1024)
    

'''
Generated CUDA code
extern "C" __global__ void default_function_kernel0( signed char* __restrict__ compute,  signed char* __restrict__ placeholder) {
  (( int*)(compute + (((int)blockIdx.x) * 4)))[0] = (( int*)(placeholder + (((int)blockIdx.x) * 4)))[0];
}
'''