VTA ALU#
import tvm
from tvm import te
import numpy as np
from tvm import topi
from tvm.contrib.utils import tempdir
import vta
import vta.testing
from vta.testing import simulator
np.random.seed(0xDEADB)
def _run(env, remote):
def check_alu(tvm_op, np_op=None, use_imm=False, test_name=None):
"""Test ALU"""
m = 8
n = 8
imm = np.random.randint(1, 5)
# compute
a = te.placeholder((m, n, env.BATCH, env.BLOCK_OUT), name="a", dtype=env.acc_dtype)
a_buf = te.compute(
(m, n, env.BATCH, env.BLOCK_OUT), lambda *i: a(*i), "a_buf"
) # DRAM->SRAM
if use_imm:
res_buf = te.compute(
(m, n, env.BATCH, env.BLOCK_OUT), lambda *i: tvm_op(a_buf(*i), imm), "res_buf"
) # compute
else:
b = te.placeholder((m, n, env.BATCH, env.BLOCK_OUT), name="b", dtype=env.acc_dtype)
b_buf = te.compute(
(m, n, env.BATCH, env.BLOCK_OUT), lambda *i: b(*i), "b_buf"
) # DRAM->SRAM
res_buf = te.compute(
(m, n, env.BATCH, env.BLOCK_OUT),
lambda *i: tvm_op(a_buf(*i), b_buf(*i)),
"res_buf",
) # compute5B
res = te.compute(
(m, n, env.BATCH, env.BLOCK_OUT),
lambda *i: res_buf(*i).astype(env.inp_dtype),
"res",
) # SRAM->DRAM
# schedule
s = te.create_schedule(res.op)
s[a_buf].set_scope(env.acc_scope) # SRAM
s[a_buf].pragma(a_buf.op.axis[0], env.dma_copy) # DRAM->SRAM
s[res_buf].set_scope(env.acc_scope) # SRAM
s[res_buf].pragma(res_buf.op.axis[0], env.alu) # compute
s[res].pragma(res.op.axis[0], env.dma_copy) # SRAM->DRAM
if not use_imm:
s[b_buf].set_scope(env.acc_scope) # SRAM
s[b_buf].pragma(b_buf.op.axis[0], env.dma_copy) # DRAM->SRAM
if not remote:
return
# build
with vta.build_config():
if use_imm:
mod = vta.build(s, [a, res], tvm.target.Target("ext_dev", host=env.target_host))
else:
mod = vta.build(
s, [a, b, res], tvm.target.Target("ext_dev", host=env.target_host)
)
temp = tempdir()
mod.save(temp.relpath("load_act.o"))
remote.upload(temp.relpath("load_act.o"))
f = remote.load_module("load_act.o")
# verify
dev = remote.ext_dev(0)
a_np = np.random.randint(-16, 16, size=(m, n, env.BATCH, env.BLOCK_OUT)).astype(a.dtype)
if use_imm:
res_np = np_op(a_np, imm) if np_op else tvm_op(a_np, imm)
else:
b_np = np.random.randint(-16, 16, size=(m, n, env.BATCH, env.BLOCK_OUT)).astype(
b.dtype
)
res_np = np_op(a_np, b_np) if np_op else tvm_op(a_np, b_np)
res_np = res_np.astype(res.dtype)
a_nd = tvm.nd.array(a_np, dev)
res_nd = tvm.nd.array(np.zeros((m, n, env.BATCH, env.BLOCK_OUT)).astype(res.dtype), dev)
if env.TARGET in ["sim", "tsim"]:
simulator.clear_stats()
if use_imm:
f(a_nd, res_nd)
else:
b_nd = tvm.nd.array(b_np, dev)
f(a_nd, b_nd, res_nd)
np.testing.assert_equal(res_np, res_nd.numpy())
if env.TARGET in ["sim", "tsim"]:
sim_stats = simulator.stats()
print("ALU {} execution statistics:".format(test_name))
for k, v in sim_stats.items():
print("\t{:<16}: {:>16}".format(k, v))
check_alu(lambda x, y: x << y, np.left_shift, use_imm=True, test_name="SHL")
check_alu(tvm.te.max, np.maximum, use_imm=True, test_name="MAX")
check_alu(tvm.te.max, np.maximum, test_name="MAX")
check_alu(lambda x, y: x + y, use_imm=True, test_name="ADD")
check_alu(lambda x, y: x + y, test_name="ADD")
check_alu(lambda x, y: x >> y, np.right_shift, use_imm=True, test_name="SHR")
vta.testing.run(_run)
ALU SHL execution statistics:
inp_load_nbytes : 0
wgt_load_nbytes : 0
acc_load_nbytes : 4096
uop_load_nbytes : 4
out_store_nbytes: 1024
gemm_counter : 0
alu_counter : 64
ALU MAX execution statistics:
inp_load_nbytes : 0
wgt_load_nbytes : 0
acc_load_nbytes : 4096
uop_load_nbytes : 4
out_store_nbytes: 1024
gemm_counter : 0
alu_counter : 64
ALU MAX execution statistics:
inp_load_nbytes : 0
wgt_load_nbytes : 0
acc_load_nbytes : 8192
uop_load_nbytes : 4
out_store_nbytes: 1024
gemm_counter : 0
alu_counter : 64
ALU ADD execution statistics:
inp_load_nbytes : 0
wgt_load_nbytes : 0
acc_load_nbytes : 4096
uop_load_nbytes : 4
out_store_nbytes: 1024
gemm_counter : 0
alu_counter : 64
ALU ADD execution statistics:
inp_load_nbytes : 0
wgt_load_nbytes : 0
acc_load_nbytes : 8192
uop_load_nbytes : 4
out_store_nbytes: 1024
gemm_counter : 0
alu_counter : 64
ALU SHR execution statistics:
inp_load_nbytes : 0
wgt_load_nbytes : 0
acc_load_nbytes : 4096
uop_load_nbytes : 4
out_store_nbytes: 1024
gemm_counter : 0
alu_counter : 64
2023-09-25 13:12:35.136 INFO load_module /tmp/tmp7oy9i8lt/load_act.o
2023-09-25 13:12:35.287 INFO load_module /tmp/tmp7oy9i8lt/load_act.o
2023-09-25 13:12:35.458 INFO load_module /tmp/tmp7oy9i8lt/load_act.o
2023-09-25 13:12:35.609 INFO load_module /tmp/tmp7oy9i8lt/load_act.o
2023-09-25 13:12:35.780 INFO load_module /tmp/tmp7oy9i8lt/load_act.o
2023-09-25 13:12:35.933 INFO load_module /tmp/tmp7oy9i8lt/load_act.o