TVM Hexagon test-matmul.py

test-matmul.py

#!/usr/bin/env python3
import tvm
import tvm.contrib.hexagon
import numpy as np
from tvm import te, tir

# Size of the matrices.
N = 32

# Construct the TVM computation.
A = te.placeholder((N, N), name='A', dtype='int16')
B = te.placeholder((N, N), name='B', dtype='int16')
k = te.reduce_axis((0, N), name='k')

# Provide the function that calculates element at position i,j of the output:
# the dot product of the i-th row of A and j-th column of B.
C = te.compute((N,N), lambda i, j: te.sum(A[i][k] * B[k][j], axis=k), name='C')

# Create the schedule.
s = te.create_schedule(C.op);
px, x = s[C].split(s[C].op.axis[0], nparts=1)
s[C].bind(px, te.thread_axis("pipeline"))

target = tvm.target.hexagon('v66', hvx=0)

f = tvm.build(s, [A, B, C], target=target, target_host='llvm', name='mmult')

# Prepare inputs as numpy arrays, and placeholders for outputs.
ctx = tvm.hexagon(0)
a = np.random.randint(0, 16, (N, N), dtype=np.int16)
b = np.random.randint(0, 16, (N, N), dtype=np.int16)
c = tvm.runtime.ndarray.empty((N, N), dtype='int16', ctx=ctx)

print(a)
print(b)

# Invoke the matrix multiplication via function f and through numpy.
f(tvm.runtime.ndarray.array(a, ctx=ctx), tvm.runtime.ndarray.array(b, ctx=ctx), c)
npc = np.matmul(a, b)

print('tvm\n', c.asnumpy())
print('numpy\n', npc)

# Report error if there is any difference in the two outputs.
if not (c.asnumpy() - npc).any():
    print('correct')
else:
    print('wrong')