tvm-trivial.py

import tvm                      
import numpy  
import time   
              
# The size of the square matrix
N = 100000000 
# The default tensor type in tvm
dtype = "float32"
target = "llvm"
#target = "llvm -mcpu=skylake-avx512"
#target = "llvm -mattr=+avx2"
# target = "llvm -mcpu=core-avx2"                                                                                                                           
# Random generated tensor for testing
a = tvm.nd.array(numpy.random.rand(N,).astype(dtype), tvm.cpu(0))
b = tvm.nd.array(numpy.random.rand(N,).astype(dtype), tvm.cpu(0))
# The expected answer         
answer = a.asnumpy() + b.asnumpy()
              
# Algorithm   
A = tvm.placeholder((N,), name = 'A')
B = tvm.placeholder((N,), name = 'B')
C = tvm.compute(
           A.shape,           
           lambda i: A[i] + B[i],
           name = 'C')        
              
# Default schedule            
s = tvm.create_schedule(C.op) 
i, = C.op.axis
io, ii = s[C].split(i, factor=16)
s[C].vectorize(ii)            
              
print(tvm.lower(s, [A, B, C], simple_mode=True))
              
func = tvm.build(s, [A, B, C], target=target, name = 'add')
assert func   
evaluator = func.time_evaluator(func.entry_name, tvm.cpu(0), number = 100)
c = tvm.nd.array(numpy.zeros((N,), dtype = dtype), tvm.cpu(0))
print('Baseline: %f' % evaluator(a, b, c).mean)
              
numpy.testing.assert_allclose(c.asnumpy(), answer, rtol=1e-5)