def test(): env = nnpu.get_env() a = tvm.placeholder((4, 16), env.cfg['dtype_w'], 'a') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) k = tvm.reduce_axis((0, 16), 'k') c_buf = tvm.compute((4, 1), lambda i, j: tvm.sum(a_buf[i,k], axis=k), 'c_buf') sph.MarkScope(c_buf) c_host, c_dram = nnpu.utils.CopyBufToH(c_buf, 'c', sph) k1 = tvm.reduce_axis((0, 16), 'k1') max_buf = tvm.compute((4, 1), lambda i, j: tvm.max(a_buf[i,k1], axis=k1), 'max_buf') sph.MarkScope(max_buf) max_host, max_dram = nnpu.utils.CopyBufToH(max_buf, 'max', sph) k2 = tvm.reduce_axis((0, 16), 'k2') min_buf = tvm.compute((4, 1), lambda i, j: tvm.min(a_buf[i,k2], axis=k2), 'min_buf') sph.MarkScope(min_buf) min_host, min_dram = nnpu.utils.CopyBufToH(min_buf, 'min', sph) # create schedule and tensorize s = tvm.create_schedule([c_host.op, max_host.op, min_host.op]) sph.Transform(s) s[c_buf].tensorize(s[c_buf].op.axis[1], env.intrins.get('VReduceSum', mode='w')) s[max_buf].tensorize(s[max_buf].op.axis[1], env.intrins.get('VReduceMax', mode='w')) s[min_buf].tensorize(s[min_buf].op.axis[1], env.intrins.get('VReduceMin', mode='w')) # build print(nnpu.lower(s, [a, c_host, max_host, min_host], simple_mode=True)) func = nnpu.build(s, [a, c_host, max_host, min_host], 'nnpu', 'llvm', name='nnpu_func') # create data and run ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(4, 16), dtype=a.dtype, low = 0, high = 64) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) c_nd = tvm.nd.array(np.zeros((4, 1)).astype(c_host.dtype), ctx) max_nd = tvm.nd.array(np.zeros((4, 1)).astype(c_host.dtype), ctx) min_nd = tvm.nd.array(np.zeros((4, 1)).astype(c_host.dtype), ctx) func(a_nd, c_nd, max_nd, min_nd) # check results gt = np.sum(a_np, axis=(1,), keepdims=True) np.testing.assert_allclose(c_nd.asnumpy(), gt) np.testing.assert_allclose(max_nd.asnumpy(), np.max(a_np, axis=(1,), keepdims=True)) np.testing.assert_allclose(min_nd.asnumpy(), np.min(a_np, axis=(1,), keepdims=True)) print('test passed')
def test(): env = nnpu.get_env() a = tvm.placeholder((16,16), env.cfg['dtype_n'], 'a') sph = ScheduleProcHelper() Imm = tvm.const(7, env.cfg['dtype_n']) a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) add_buf = tvm.compute((16,16), lambda i,j: Imm+a_buf[i][j], 'add_buf') sph.MarkScope(add_buf) add_host, add_dram = nnpu.utils.CopyBufToH(add_buf, 'add', sph) dtype_w = env.cfg['dtype_w'] mul_buf = tvm.compute((16,16), lambda i,j: a_buf[i][j].astype(dtype_w) * Imm.astype(dtype_w), 'mul_buf') sph.MarkScope(mul_buf) mul_host, mul_dram = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) rsub_buf = tvm.compute((16,16), lambda i,j: Imm-a_buf[i][j], 'rsub_buf') sph.MarkScope(rsub_buf) rsub_host, rsub_dram = nnpu.utils.CopyBufToH(rsub_buf, 'rsub', sph) s = tvm.create_schedule([add_host.op,mul_host.op,rsub_host.op]) sph.Transform(s) s[add_buf].tensorize(s[add_buf].op.axis[0], env.intrins.get('MAddI', shape=(16,16), imm_value=Imm.value, mode='n')) s[mul_buf].tensorize(s[mul_buf].op.axis[0], env.intrins.get('MMulI', shape=(16,16), imm_value=Imm.value, mode='inc')) s[rsub_buf].tensorize(s[rsub_buf].op.axis[0], env.intrins.get('ISubM', shape=(16,16), imm_value=Imm.value, mode='n')) print(nnpu.lower(s, [a,add_host,mul_host,rsub_host], simple_mode=True)) func = nnpu.build(s, [a,add_host,mul_host,rsub_host], 'nnpu', 'llvm', name='nnpu_vmuli') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(16,16), dtype=a.dtype, low = 3, high = 100) a_nd = tvm.nd.array(a_np, ctx) add_nd = tvm.nd.array(np.zeros((16,16)).astype(add_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros((16,16)).astype(mul_host.dtype), ctx) rsub_nd = tvm.nd.array(np.zeros((16,16)).astype(rsub_host.dtype), ctx) func(a_nd, add_nd,mul_nd,rsub_nd) print(a_nd.asnumpy()) print('add result is: ') print(add_nd.asnumpy()) np.testing.assert_allclose(add_nd.asnumpy(), a_np + Imm.value) print('mul result is: ') print(mul_nd.asnumpy()) np.testing.assert_allclose(mul_nd.asnumpy(), a_np.astype(dtype_w) * Imm.value) print('rsub result is: ') print(rsub_nd.asnumpy()) np.testing.assert_allclose(rsub_nd.asnumpy(), Imm.value - a_np ) print('test passed')
def test(): env = nnpu.get_env() shape = (8, 16) a = tvm.placeholder(shape, env.cfg['dtype_n'], 'a') b = tvm.placeholder(shape, env.cfg['dtype_n'], 'b') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) dtype_w = env.cfg['dtype_w'] k = tvm.reduce_axis((0, 16), 'k') dot_buf = tvm.compute((8, ), lambda i: tvm.sum( a_buf[i, k].astype(dtype_w) * b_buf[i, k].astype(dtype_w), k), 'dot_buf') sph.MarkScope(dot_buf) dot_host, dot_dram = nnpu.utils.CopyBufToH(dot_buf, 'sum', sph) s = tvm.create_schedule(dot_host.op) sph.Transform(s) s[dot_buf].tensorize(s[dot_buf].op.axis[0], env.intrins.get('MRowDot', shape=shape, mode='inc')) print(nnpu.lower(s, [a, b, dot_host], simple_mode=True)) func = nnpu.build(s, [a, b, dot_host], 'nnpu', 'llvm', name='nnpu_func') print('------------------- device module 1 llvm IR: ') print(func.imported_modules[0].get_source('ll')) print('------------------- device module 1 asm code: ') print(func.imported_modules[0].get_source('asm')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(8, 16), dtype=a.dtype, low=-32, high=32) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=(8, 16), dtype=b.dtype, low=-32, high=32) b_nd = tvm.nd.array(b_np, ctx) c_nd = tvm.nd.array(np.zeros((8, )).astype(dot_host.dtype), ctx) func(a_nd, b_nd, c_nd) #print('a = ') #print(a_np) #print('b = ') #print(b_np) print(c_nd.asnumpy()) print('ground truth is') gt = np.multiply(a_np, b_np, dtype=dot_host.dtype) gt = np.sum(gt, axis=1) print(gt) np.testing.assert_allclose(c_nd.asnumpy(), gt)
def test(): env = nnpu.get_env() shape = (16, 16) flatten_shape = (shape[0] * shape[1],) a = tvm.placeholder(flatten_shape, env.cfg['dtype_n'], 'a') b = tvm.placeholder(flatten_shape, env.cfg['dtype_n'], 'b') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) sum_buf = tvm.compute(flatten_shape, lambda i: a_buf[i] + b_buf[i], 'sum_buf') sph.MarkScope(sum_buf) sum_host, sum_dram = nnpu.utils.CopyBufToH(sum_buf, 'sum', sph) s = tvm.create_schedule([sum_host.op]) sph.Transform(s) xo, xi = s[sum_buf].split(sum_buf.op.axis[0], 16) s[sum_buf].tensorize(xo, env.intrins.get('MAddM', shape=shape, mode='n')) print(nnpu.lower(s, [a, b, sum_host], simple_mode=True)) func = nnpu.build(s, [a, b, sum_host], 'nnpu', 'llvm', name='nnpu_exp') print('------------------- device module 1 llvm IR: ') print(func.imported_modules[0].get_source('ll')) print('------------------- device module 1 asm code: ') print(func.imported_modules[0].get_source('asm')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=flatten_shape, dtype=a.dtype, low = 0, high = 23) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=flatten_shape, dtype=b.dtype, low = 0, high = 23) b_nd = tvm.nd.array(b_np, ctx) c_nd = tvm.nd.array(np.zeros(flatten_shape).astype(sum_host.dtype), ctx) func(a_nd, b_nd, c_nd) print('a = ') print(a_np) print('b = ') print(b_np) print('a + b = ') print(c_nd.asnumpy()) print("numpy ground truth is") print(a_np + b_np) np.testing.assert_allclose(c_nd.asnumpy(), a_np + b_np)
def test(): env = nnpu.get_env() nnpu.set_device(env) shape = (16, ) bigshape = (4, 64) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] sph = ScheduleProcHelper() a = tvm.placeholder(bigshape, dtype_n, 'a') a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) str_op = 'VAddMerge' k = tvm.reduce_axis((0, 4), 'k') c_buf = tvm.compute((64, ), lambda i: tvm.sum(a_buf[k, i], axis=k), 'c_buf') sph.MarkScope(c_buf) c_host, c_dram = nnpu.utils.CopyBufToH(c_buf, 'c', sph) s = tvm.create_schedule(c_host.op) sph.Transform(s) #tensorize ko, ki = s[c_buf].split(c_buf.op.reduce_axis[0], factor=1) xo, xi = s[c_buf].split(c_buf.op.axis[0], factor=shape[0]) s[c_buf].reorder(xo, ko, ki, xi) #s[c_buf].tensorize(ki, env.intrins.get(str_op, mode='n')) print(nnpu.lower(s, [a, c_host], simple_mode=True)) exit() func = nnpu.build(s, [a, c_host], 'nnpu', 'llvm', name='nnpu_func') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=bigshape, dtype=a.dtype, low=-4, high=4) a_nd = tvm.nd.array(a_np, ctx) c_nd = tvm.nd.array(np.zeros((64, ), dtype=c_host.dtype), ctx) func(a_nd, c_nd) print(str_op) print(c_nd.asnumpy()) gt = np.sum(a_np, axis=0, dtype=dtype_w) print('ground truth=') print(gt) np.testing.assert_allclose(c_nd.asnumpy(), gt)
def test(): env = nnpu.get_env() shape = (16, ) bigshape = (128, ) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] assert bigshape[0] % shape[0] == 0, 'the big vctr size is wrong' n_sheet = bigshape[0] // shape[0] sph = ScheduleProcHelper() a = tvm.placeholder(bigshape, dtype_n, 'a') b = tvm.placeholder(bigshape, dtype_n, 'b') a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) strop = 'VAddV' c_buf = tvm.compute(bigshape, lambda *i: a_buf(*i) + b_buf(*i), 'c_buf') sph.MarkScope(c_buf) c_host, c_dram = nnpu.utils.CopyBufToH(c_buf, 'sum', sph) s = tvm.create_schedule(c_host.op) sph.Transform(s) #tensorize xo, xi = s[c_buf].split(c_buf.op.axis[0], factor=shape[0]) s[c_buf].reorder(xo, xi) s[c_buf].tensorize(xi, env.intrins.get(strop, mode='n')) print(nnpu.lower(s, [a, b, c_host], simple_mode=True)) func = nnpu.build(s, [a, b, c_host], 'nnpu', 'llvm', name='nnpu_func') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=bigshape, dtype=a.dtype, low=-4, high=4) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=bigshape, dtype=b.dtype, low=-4, high=4) b_nd = tvm.nd.array(b_np, ctx) c_nd = tvm.nd.array(np.zeros(bigshape, dtype=c_host.dtype), ctx) func(a_nd, b_nd, c_nd) print(strop) print(c_nd.asnumpy()) gt = a_np + b_np np.testing.assert_allclose(c_nd.asnumpy(), gt)
def test(): env = nnpu.get_env() a = tvm.placeholder((4, 16), 'int16', 'a') b = tvm.placeholder((16, ), 'int16', 'b') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) k = tvm.reduce_axis((0, 16), 'k') c_buf = tvm.compute( (4, 1), lambda i, j: tvm.sum(a_buf[i, k] * b_buf[k], axis=k), 'c_buf') sph.MarkScope(c_buf) c_host, c_dram = nnpu.utils.CopyBufToH(c_buf, 'c', sph) s = tvm.create_schedule(c_host.op) sph.Transform(s) print(s[c_buf]) s[c_buf].tensorize(s[c_buf].op.axis[1], env.intrins.get('VDotV', mode='w')) print(nnpu.lower(s, [a, b, c_host], simple_mode=True)) func = nnpu.build(s, [a, b, c_host], 'nnpu', 'llvm', name='nnpu_func') print('------------------- device module 1 llvm IR: ') print(func.imported_modules[0].get_source('ll')) print('------------------- device module 1 asm code: ') print(func.imported_modules[0].get_source('asm')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(4, 16), dtype=a.dtype, low=0, high=64) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=(16, ), dtype=b.dtype, low=0, high=64) b_nd = tvm.nd.array(b_np, ctx) c_nd = tvm.nd.array(np.zeros((4, 1)).astype(c_host.dtype), ctx) func(a_nd, b_nd, c_nd) print(c_nd.asnumpy()) print("numpy ground truth is") print(np.dot(a_np, b_np))
def test(): env = nnpu.get_env() nnpu.set_device(env) a = tvm.placeholder((4, 4, 16), 'int16', 'a') #b = tvm.placeholder((16, ), 'int16', 'b') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) #b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) k = tvm.reduce_axis((0, 4), 'k0') c_buf = tvm.compute((4, 16), lambda i, j: tvm.sum(a_buf[k, i, j], axis=k), 'c_buf') sph.MarkScope(c_buf) c_host, c_dram = nnpu.utils.CopyBufToH(c_buf, 'c', sph) s = tvm.create_schedule(c_host.op) sph.Transform(s) ko, ki = s[c_buf].split(c_buf.op.reduce_axis[0], factor=1) s[c_buf].reorder(c_buf.op.axis[0], ko, ki, c_buf.op.axis[1]) s[c_buf].tensorize(ki, env.intrins.get('VAddMerge', mode='w', nDim=3)) print(nnpu.lower(s, [a, c_host], simple_mode=True)) func = nnpu.build(s, [a, c_host], 'nnpu', 'llvm', name='nnpu_exp') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(4, 4, 16), dtype=a.dtype, low=-4000, high=4000) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) c_nd = tvm.nd.array(np.zeros((4, 16)).astype(c_host.dtype), ctx) func(a_nd, c_nd) print(c_nd.asnumpy()) print("numpy ground truth is") gt = np.sum(a_np, axis=0) print(gt)
def test(): env = nnpu.get_env() dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] a = tvm.placeholder((64, ), dtype_n, 'a') b = tvm.placeholder((1, ), dtype_n, 'b') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) c_buf = tvm.compute((64, ), lambda i: a_buf[i] + b_buf[0], 'c_buf') sph.MarkScope(c_buf) c_host, _ = nnpu.utils.CopyBufToH(c_buf, 'c', sph) sub_buf = tvm.compute((64, ), lambda i: a_buf[i] - b_buf[0], 'sub_buf') sph.MarkScope(sub_buf) sub_host, _ = nnpu.utils.CopyBufToH(sub_buf, 'sub', sph) rsub_buf = tvm.compute((64, ), lambda i: b_buf[0] - a_buf[i], 'rsub_buf') sph.MarkScope(rsub_buf) rsub_host, _ = nnpu.utils.CopyBufToH(rsub_buf, 'rsub', sph) mul_buf = tvm.compute( (64, ), lambda i: a_buf[i].astype(dtype_w) * b_buf[0].astype(dtype_w), 'mul_buf') sph.MarkScope(mul_buf) mul_host, _ = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) div_buf = tvm.compute((64, ), lambda i: a_buf[i] / b_buf[0], 'div_buf') sph.MarkScope(div_buf) div_host, _ = nnpu.utils.CopyBufToH(div_buf, 'div', sph) rdiv_buf = tvm.compute((64, ), lambda i: b_buf[0] / a_buf[i], 'rdiv_buf') sph.MarkScope(rdiv_buf) rdiv_host, _ = nnpu.utils.CopyBufToH(rdiv_buf, 'rdiv', sph) gtm_buf = tvm.compute((64, ), lambda i: tvm.max(a_buf[i], b_buf[0]), 'gtm_buf') sph.MarkScope(gtm_buf) gtm_host, gtm_dram = nnpu.utils.CopyBufToH(gtm_buf, 'gtm', sph) s = tvm.create_schedule([ c_host.op, sub_host.op, mul_host.op, rsub_host.op, div_host.op, rdiv_host.op, gtm_host.op ]) sph.Transform(s) xo, xi = s[c_buf].split(c_buf.op.axis[0], 16) s[c_buf].tensorize(xi, env.intrins.get('VAddS', mode='n')) xo, xi = s[sub_buf].split(sub_buf.op.axis[0], 16) s[sub_buf].tensorize(xi, env.intrins.get('VSubS', mode='n')) xo, xi = s[rsub_buf].split(rsub_buf.op.axis[0], 16) s[rsub_buf].tensorize(xi, env.intrins.get('SSubV', mode='n')) xo, xi = s[mul_buf].split(mul_buf.op.axis[0], 16) s[mul_buf].tensorize(xi, env.intrins.get('VMulS', mode='inc')) xo, xi = s[div_buf].split(div_buf.op.axis[0], 16) s[div_buf].tensorize(xi, env.intrins.get('VDivS', mode='n')) xo, xi = s[rdiv_buf].split(rdiv_buf.op.axis[0], 16) s[rdiv_buf].tensorize(xi, env.intrins.get('SDivV', mode='n')) xo, xi = s[gtm_buf].split(gtm_buf.op.axis[0], 16) s[gtm_buf].tensorize(xi, env.intrins.get('VGTMS', mode='n')) print( nnpu.lower(s, [ a, b, c_host, sub_host, mul_host, rsub_host, div_host, rdiv_host, gtm_host ], simple_mode=True)) func = nnpu.build(s, [ a, b, c_host, sub_host, mul_host, rsub_host, div_host, rdiv_host, gtm_host ], 'nnpu', 'llvm', name='nnpu_exp') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(64, ), dtype=a.dtype, low=1, high=63) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=(1, ), dtype=b.dtype, low=2, high=31) b_nd = tvm.nd.array(b_np, ctx) c_nd = tvm.nd.array(np.zeros((64, )).astype(c_host.dtype), ctx) sub_nd = tvm.nd.array(np.zeros((64, )).astype(sub_host.dtype), ctx) rsub_nd = tvm.nd.array(np.zeros((64, )).astype(rsub_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros((64, )).astype(mul_host.dtype), ctx) div_nd = tvm.nd.array(np.zeros((64, )).astype(div_host.dtype), ctx) rdiv_nd = tvm.nd.array(np.zeros((64, )).astype(rdiv_host.dtype), ctx) gtm_nd = tvm.nd.array(np.zeros((64, )).astype(gtm_host.dtype), ctx) print('------------------- device module 1 llvm IR: ') print(func.imported_modules[0].get_source('ll')) print('------------------- device module 1 asm code: ') print(func.imported_modules[0].get_source('asm')) func(a_nd, b_nd, c_nd, sub_nd, mul_nd, rsub_nd, div_nd, rdiv_nd, gtm_nd) print('a = ') print(a_np) print('b = ') print(b_np) print('a + b =') print(c_nd.asnumpy()) print('numpy ground truth =') gt = a_np + b_np print(gt) np.testing.assert_allclose(c_nd.asnumpy(), gt) print('a - b =') print(sub_nd.asnumpy()) np.testing.assert_allclose(sub_nd.asnumpy(), a_np - b_np) print('b - a =') print(rsub_nd.asnumpy()) np.testing.assert_allclose(rsub_nd.asnumpy(), b_np - a_np) print('a * b =') print(mul_nd.asnumpy()) np.testing.assert_allclose(mul_nd.asnumpy(), a_np * b_np.astype(dtype_w)) print('a / b =') print(div_nd.asnumpy()) # numpy always round down, while in c, the numerator will be rounded to zero. #np.testing.assert_allclose(div_nd.asnumpy(), a_np / b_np) print('b / a =') print(rdiv_nd.asnumpy()) print('max(a, b)=') print(gtm_nd.asnumpy())
def test(): env = nnpu.get_env() dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] shape = (4, 16) a = tvm.placeholder(shape, dtype_n, 'a') b = tvm.placeholder((16, ), dtype_n, 'b') sph = ScheduleProcHelper() a_buf, _ = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, _ = nnpu.utils.CopyHtoBuf(b, 'b', sph) sum_buf = tvm.compute(shape, lambda i, j: a_buf[i, j] + b_buf[j], 'sum_buf') sph.MarkScope(sum_buf) sum_host, _ = nnpu.utils.CopyBufToH(sum_buf, 'sum', sph) sub_buf = tvm.compute(shape, lambda i, j: a_buf[i, j] - b_buf[j], 'sub_buf') sph.MarkScope(sub_buf) sub_host, _ = nnpu.utils.CopyBufToH(sub_buf, 'sub', sph) mul_buf = tvm.compute( shape, lambda i, j: a_buf[i, j].astype(dtype_w) * b_buf[j].astype(dtype_w), 'sub_buf') sph.MarkScope(mul_buf) mul_host, _ = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) s = tvm.create_schedule([sum_host.op, sub_host.op, mul_host.op]) sph.Transform(s) s[sum_buf].pragma(sum_buf.op.axis[0], 'nnpu.vector', str({ 'code': 'matrix-vector', 'shape': shape })) s[sub_buf].pragma(sub_buf.op.axis[0], 'nnpu.vector', str({ 'code': 'matrix-vector', 'shape': shape })) s[mul_buf].pragma(mul_buf.op.axis[0], 'nnpu.vector', str({ 'code': 'matrix-vector', 'shape': shape })) print(nnpu.lower(s, [a, b, sum_host, sub_host, mul_host], simple_mode=True)) func = nnpu.build(s, [a, b, sum_host, sub_host, mul_host], 'nnpu', 'llvm', name='nnpu_func') print('------------------- device module 1 llvm IR: ') print(func.imported_modules[0].get_source('ir')) print('------------------- device module 1 uop code: ') print(func.imported_modules[0].get_source('uop')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(4, 16), dtype=a.dtype, low=0, high=64) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=(16, ), dtype=b.dtype, low=0, high=64) b_nd = tvm.nd.array(b_np, ctx) sum_nd = tvm.nd.array(np.zeros(shape).astype(sum_host.dtype), ctx) sub_nd = tvm.nd.array(np.zeros(shape).astype(sub_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros(shape).astype(mul_host.dtype), ctx) func(a_nd, b_nd, sum_nd, sub_nd, mul_nd) gt = a_np + b_np np.testing.assert_allclose(sum_nd.asnumpy(), gt) gt = a_np - b_np np.testing.assert_allclose(sub_nd.asnumpy(), gt) gt = a_np.astype(dtype_w) * b_np np.testing.assert_allclose(mul_nd.asnumpy(), gt) print('test passed')
def test(): env = nnpu.get_env() dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] shape = (4, 64) # nvctr_unit = env.cfg['vector_unit']['size'] nvctr_unit = 32 # assert shape[0] % nvctr_unit == 0, 'error' a = tvm.placeholder(shape, dtype_n, 'a') b = tvm.placeholder(shape, dtype_n, 'b') sph = ScheduleProcHelper() b_scope = 'buffer0' a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph, dst_scope=b_scope) c_buf = tvm.compute(shape, lambda *i: a_buf(*i) + b_buf(*i), 'c_buf') sph.MarkScope(c_buf) c_host, c_dram = nnpu.utils.CopyBufToH(c_buf, 'c', sph) mul_buf = tvm.compute( shape, lambda *i: a_buf(*i).astype(dtype_w) * b_buf(*i).astype(dtype_w), 'mul_buf') sph.MarkScope(mul_buf) mul_host, mul_dram = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) gtm_buf = tvm.compute(shape, lambda *i: tvm.max(a_buf(*i), b_buf(*i)), 'gtm_buf') sph.MarkScope(gtm_buf) gtm_host, gtm_dram = nnpu.utils.CopyBufToH(gtm_buf, 'gtm', sph) s = tvm.create_schedule([c_host.op, mul_host.op, gtm_host.op]) sph.Transform(s) # x = s[c_buf].fuse(*c_buf.op.axis) # xo, xi = s[c_buf].split(x, factor=nvctr_unit) params = dict() params['code'] = 'binary' params['size'] = nvctr_unit x = s[c_buf].fuse(*c_buf.op.axis) xo, xi = s[c_buf].split(x, factor=nvctr_unit) s[c_buf].pragma(xi, 'nnpu.vector', str(params)) x = s[mul_buf].fuse(*mul_buf.op.axis) xo, xi = s[mul_buf].split(x, factor=nvctr_unit) s[mul_buf].pragma(xi, 'nnpu.vector', str(params)) x = s[gtm_buf].fuse(*gtm_buf.op.axis) xo, xi = s[gtm_buf].split(x, factor=nvctr_unit) s[gtm_buf].pragma(xi, 'nnpu.vector', str(params)) print(tvm.lower(s, [a, b, c_host, mul_host, gtm_host], simple_mode=True)) print(nnpu.lower(s, [a, b, c_host, mul_host, gtm_host], simple_mode=True)) # exit() func = nnpu.build(s, [a, b, c_host, mul_host, gtm_host], 'nnpu', 'llvm', name='nnpu_exp') print('------------------- device module 1 IR: ') print(func.imported_modules[0].get_source('ir')) print('------------------- device module 1 micro code: ') print(func.imported_modules[0].get_source('uop')) # exit() ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=shape, dtype=a.dtype, low=-64, high=63) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=shape, dtype=b.dtype, low=-64, high=63) b_nd = tvm.nd.array(b_np, ctx) c_nd = tvm.nd.array(np.zeros(shape).astype(c_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros(shape).astype(mul_host.dtype), ctx) gtm_nd = tvm.nd.array(np.zeros(shape).astype(gtm_host.dtype), ctx) # print('------------------- device module 1 llvm IR: ') # print(func.imported_modules[0].get_source('ll')) # print('------------------- device module 1 asm code: ') # print(func.imported_modules[0].get_source('asm')) func(a_nd, b_nd, c_nd, mul_nd, gtm_nd) gt = a_np + b_np np.testing.assert_allclose(c_nd.asnumpy(), gt) gt = np.multiply(a_np, b_np, dtype=mul_host.dtype) np.testing.assert_allclose(mul_nd.asnumpy(), gt) gt = np.maximum(a_np, b_np) np.testing.assert_allclose(gtm_nd.asnumpy(), gt) print('test passed!!')
def test(): env = nnpu.get_env() dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] shape = (4, 16) a = tvm.placeholder(shape, dtype_n, 'a') b = tvm.placeholder((16, ), dtype_n, 'b') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) sum_buf = tvm.compute(shape, lambda i, j: a_buf[i, j] + b_buf[j], 'sum_buf') sph.MarkScope(sum_buf) sum_host, sum_dram = nnpu.utils.CopyBufToH(sum_buf, 'sum', sph) sub_buf = tvm.compute(shape, lambda i, j: a_buf[i, j] - b_buf[j], 'sub_buf') sph.MarkScope(sub_buf) sub_host, sub_dram = nnpu.utils.CopyBufToH(sub_buf, 'sub', sph) mul_buf = tvm.compute( shape, lambda i, j: a_buf[i, j].astype(dtype_w) * b_buf[j].astype(dtype_w), 'sub_buf') sph.MarkScope(mul_buf) mul_host, mul_dram = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) s = tvm.create_schedule([sum_host.op, sub_host.op, mul_host.op]) sph.Transform(s) s[sum_buf].tensorize(s[sum_buf].op.axis[0], env.intrins.get('MAddV', shape=(4, 16), mode='n')) s[sub_buf].tensorize(s[sub_buf].op.axis[0], env.intrins.get('MSubV', shape=(4, 16), mode='n')) s[mul_buf].tensorize(s[mul_buf].op.axis[0], env.intrins.get('MMulV', shape=(4, 16), mode='inc')) print(nnpu.lower(s, [a, b, sum_host, sub_host, mul_host], simple_mode=True)) func = nnpu.build(s, [a, b, sum_host, sub_host, mul_host], 'nnpu', 'llvm', name='nnpu_func') print('------------------- device module 1 llvm IR: ') print(func.imported_modules[0].get_source('ll')) print('------------------- device module 1 asm code: ') print(func.imported_modules[0].get_source('asm')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(4, 16), dtype=a.dtype, low=0, high=64) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=(16, ), dtype=b.dtype, low=0, high=64) b_nd = tvm.nd.array(b_np, ctx) sum_nd = tvm.nd.array(np.zeros(shape).astype(sum_host.dtype), ctx) sub_nd = tvm.nd.array(np.zeros(shape).astype(sub_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros(shape).astype(mul_host.dtype), ctx) func(a_nd, b_nd, sum_nd, sub_nd, mul_nd) print('a = ') print(a_np) print('b = ') print(b_np) print('sum result is ') print(sum_nd.asnumpy()) print("numpy ground truth is") gt = a_np + b_np print(gt) np.testing.assert_allclose(sum_nd.asnumpy(), gt) print('sub result is ') print(sub_nd.asnumpy()) np.testing.assert_allclose(sub_nd.asnumpy(), a_np - b_np) print('mul result is ') print(mul_nd.asnumpy()) np.testing.assert_allclose(mul_nd.asnumpy(), a_np.astype(dtype_w) * b_np) print('test passed')
def test(): env = nnpu.get_env() a = tvm.placeholder((16, ), env.cfg['dtype_w'], 'a') sph = ScheduleProcHelper() Imm = tvm.const(5, env.cfg['dtype_w']) a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) #c_buf = tvm.compute((16, ), lambda i: tvm.select(a_buf[i]>Imm,a_buf[i],Imm), 'c_buf') c_buf = tvm.compute((16, ), lambda i: Imm+a_buf[i], 'c_buf') sph.MarkScope(c_buf) c_host, c_dram = nnpu.utils.CopyBufToH(c_buf, 'c', sph) sub_buf = tvm.compute((16, ), lambda i: a_buf[i] - Imm , 'sub_buf') sph.MarkScope(sub_buf) sub_host, sub_dram = nnpu.utils.CopyBufToH(sub_buf, 'sub', sph) mul_buf = tvm.compute((16, ), lambda i: a_buf[i] * Imm, 'mul_buf') sph.MarkScope(mul_buf) mul_host, mul_dram = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) div_buf = tvm.compute((16, ), lambda i: a_buf[i] / Imm, 'rdiv_buf') sph.MarkScope(div_buf) div_host, div_dram = nnpu.utils.CopyBufToH(div_buf, 'rdiv', sph) gtm_buf = tvm.compute((16, ), lambda i: tvm.max(a_buf[i], Imm), 'gtm_buf') sph.MarkScope(gtm_buf) gtm_host, gtm_dram = nnpu.utils.CopyBufToH(gtm_buf, 'gtm', sph) rsub_buf = tvm.compute((16, ), lambda i: Imm-a_buf[i], 'rsub_buf') sph.MarkScope(rsub_buf) rsub_host, rsub_dram = nnpu.utils.CopyBufToH(rsub_buf, 'rsub', sph) s = tvm.create_schedule([c_host.op, sub_host.op, mul_host.op, div_host.op, gtm_host.op,rsub_host.op]) sph.Transform(s) s[c_buf].tensorize(s[c_buf].op.axis[0], env.intrins.get('VAddI', imm_value=Imm.value,mode='w')) s[sub_buf].tensorize(s[sub_buf].op.axis[0], env.intrins.get('VSubI', imm_value=Imm.value,mode='w')) s[mul_buf].tensorize(s[mul_buf].op.axis[0], env.intrins.get('VMulI', imm_value=Imm.value,mode='w')) s[div_buf].tensorize(s[div_buf].op.axis[0], env.intrins.get('VDivI', imm_value=Imm.value,mode='w')) s[gtm_buf].tensorize(s[gtm_buf].op.axis[0], env.intrins.get('VGTMI', imm_value=Imm.value,mode='w')) s[rsub_buf].tensorize(s[rsub_buf].op.axis[0], env.intrins.get('ISubV', imm_value=Imm.value,mode='w')) print(nnpu.lower(s, [a,c_host,sub_host,mul_host,div_host,gtm_host,rsub_host], simple_mode=True)) func = nnpu.build(s, [a,c_host,sub_host,mul_host,div_host,gtm_host,rsub_host], 'nnpu', 'llvm', name='nnpu_vmuli') print('------------------- device module 1 llvm IR: ') print(func.imported_modules[0].get_source('ll')) print('------------------- device module 1 asm code: ') print(func.imported_modules[0].get_source('asm')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(16, ), dtype=a.dtype, low = 3, high = 122) #a_np = np.random.random(size=shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) c_nd = tvm.nd.array(np.zeros((16, )).astype(c_host.dtype), ctx) sub_nd = tvm.nd.array(np.zeros((16, )).astype(c_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros((16, )).astype(c_host.dtype), ctx) div_nd = tvm.nd.array(np.zeros((16, )).astype(c_host.dtype), ctx) gtm_nd = tvm.nd.array(np.zeros((16, )).astype(c_host.dtype), ctx) rsub_nd = tvm.nd.array(np.zeros((16, )).astype(c_host.dtype), ctx) func(a_nd, c_nd, sub_nd, mul_nd, div_nd, gtm_nd,rsub_nd) print('a = ') print(a_nd.asnumpy()) print('a + {0} = '.format(Imm.value)) print(c_nd.asnumpy()) print('numpy ground truth =') gt = a_np + Imm.value print(gt) np.testing.assert_allclose(c_nd.asnumpy(), gt) print('a - {0} = '.format(Imm.value)) print(sub_nd.asnumpy()) np.testing.assert_allclose(sub_nd.asnumpy(), a_np - Imm.value) print('a * {0} = '.format(Imm.value)) print(mul_nd.asnumpy()) np.testing.assert_allclose(mul_nd.asnumpy(), a_np * Imm.value) print('a / {0} = '.format(Imm.value)) print(div_nd.asnumpy()) np.testing.assert_allclose(div_nd.asnumpy(), a_np / Imm.value) print('a > {0} ? a : {0} = '.format(Imm.value)) print(gtm_nd.asnumpy()) #np.testing.assert_allclose(gtm_nd.asnumpy(), a_np Imm.value) print('{0} - a = '.format(Imm.value)) print(rsub_nd.asnumpy()) np.testing.assert_allclose(rsub_nd.asnumpy(), Imm.value-a_np) print('test passed')
def test(): env = nnpu.get_env() nnpu.set_dump(False) #==================================# # ------ first define shapes ------ #==================================# # input data layout: HWC in_shape = (32, 32, 128) # pooling windows size, height == width. cell_shape = 4 # in this demo we don't do padding, so input data height and width must be divisible to pooling window size. assert in_shape[0] % cell_shape == 0, 'error' assert in_shape[1] % cell_shape == 0, 'error' nvctr_unit = env.cfg['vector_unit']['size'] assert in_shape[2] % nvctr_unit == 0, 'channel not divisible to vector unit size' out_shape = (in_shape[0] // cell_shape,in_shape[1] // cell_shape,in_shape[2]) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] sph = ScheduleProcHelper() str_op = 'VGTMMerge' #=================================================================# # ------ after all shapes defined, begin compute describing. ------ #=================================================================# a = tvm.placeholder(in_shape, dtype_w, 'a') # first copy to scratchpad. a_buf, _1 = nnpu.utils.CopyHtoBuf(a, 'a', sph) # stage 1, find the maximum pixel in every pooling window. # the extent of two reduction axes are sizes of pooling window. k1 = tvm.reduce_axis((0,cell_shape), 'k1') k2 = tvm.reduce_axis((0,cell_shape), 'k2') pooling_buf = tvm.compute(out_shape, lambda i,j,k: tvm.max(a_buf[i * cell_shape + k1, j * cell_shape + k2, k], axis=[k1, k2]), 'pooling_buf') sph.MarkScope(pooling_buf, 'buffer1') # copy back to host. step2_host, step2_dram = nnpu.utils.CopyBufToH(pooling_buf, 'pooling',sph) # ------ this ends the computation description. ------ #==================================# # ------ begin scheduling ------ #==================================# s = tvm.create_schedule(step2_host.op) sph.Transform(s) #tensorize i, j, k = pooling_buf.op.axis k1, k2 = pooling_buf.op.reduce_axis # split the reduce_axis by factor 1, to produce a dummy reduce axis. # this is a trick to enable tensorize, due to limitation of tvm's tensorize pattern matcher. ko, ki = s[pooling_buf].split(k2, factor=1) xo, xi = s[pooling_buf].split(k, factor=16) # reorder axes. # put xo right before ki to eliminate memory dependency between two consecutive VGTMV instruction s[pooling_buf].reorder( i, j, k1, ko, xo, ki, xi) s[pooling_buf].tensorize(ki, env.intrins.get(str_op, scope_out='buffer1', mode='w')) # unroll # s[pooling_buf].unroll(ko) # s[pooling_buf].unroll(xo) #==================================# # ------ this ends the scheduling ------ #==================================# print(nnpu.lower(s, [a, step2_host], simple_mode=True)) # exit() func = nnpu.build(s, [a, step2_host], 'nnpu', 'llvm', name='nnpu_func') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=in_shape, dtype=a.dtype, low = -128, high = 127) a_nd = tvm.nd.array(a_np, ctx) c_nd = tvm.nd.array(np.zeros(out_shape, dtype=step2_host.dtype), ctx) func(a_nd, c_nd) # print("pooling-max") # print(c_nd.asnumpy()) # print("nppooling-max") gt=max_pooling(in_shape,out_shape,cell_shape,a_np,a.dtype) # print(gt) np.testing.assert_allclose(c_nd.asnumpy(), gt) print('test passed')
def test(): env = nnpu.get_env() shape = (4, 16) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] a = tvm.placeholder(shape, dtype_w, 'a') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) k = tvm.reduce_axis((0, 4), 'k') add_buf = tvm.compute((16, ), lambda i: tvm.sum(a_buf[k, i], axis=k), 'add_buf') sph.MarkScope(add_buf) add_host, add_dram = nnpu.utils.CopyBufToH(add_buf, 'add', sph) # k1 = tvm.reduce_axis((0, 4), 'k1') # mul_buf = tvm.compute((16, ), lambda i: tvm.sum(a_buf[k1, i], axis=k1), 'mul_buf') # sph.MarkScope(mul_buf) # mul_host, mul_dram = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) k2 = tvm.reduce_axis((0, 4), 'k2') gtm_buf = tvm.compute((16, ), lambda i: tvm.max(a_buf[k2, i], axis=k2), 'gtm_buf') sph.MarkScope(gtm_buf) gtm_host, gtm_dram = nnpu.utils.CopyBufToH(gtm_buf, 'gtm', sph) s = tvm.create_schedule([add_host.op, gtm_host.op]) sph.Transform(s) ko, ki = s[add_buf].split(add_buf.op.reduce_axis[0], factor=1) s[add_buf].reorder(ko, ki, s[add_buf].op.axis[0]) s[add_buf].tensorize(ki, env.intrins.get('VAddMerge', mode='w')) # ko1, ki1 = s[mul_buf].split(mul_buf.op.reduce_axis[0], factor=1) # s[mul_buf].reorder(ko1, ki1, s[mul_buf].op.axis[0]) # s[mul_buf].tensorize(ki1, env.intrins.get('VMulMerge', mode='w')) ko2, ki2 = s[gtm_buf].split(gtm_buf.op.reduce_axis[0], factor=1) s[gtm_buf].reorder(ko2, ki2, s[gtm_buf].op.axis[0]) s[gtm_buf].tensorize(ki2, env.intrins.get('VGTMMerge', mode='w')) print(nnpu.lower(s, [a, add_host, gtm_host], simple_mode=True)) func = nnpu.build(s, [a, add_host, gtm_host], 'nnpu', 'llvm', name='nnpu_func') #exit() ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(4, 16), dtype=a.dtype, low=-16, high=16) a_nd = tvm.nd.array(a_np, ctx) add_nd = tvm.nd.array(np.zeros((16, )).astype(add_host.dtype), ctx) # mul_nd = tvm.nd.array(np.zeros((16,)).astype(mul_host.dtype), ctx) gtm_nd = tvm.nd.array(np.zeros((16, )).astype(gtm_host.dtype), ctx) print('------------------- device module 1 IR code: ') print(func.imported_modules[0].get_source('ir')) func(a_nd, add_nd, gtm_nd) print('a = ') print(a_np) print('reduce sum row = ') print(add_nd.asnumpy()) print('ground truth is: ') gt = np.sum(a_np, axis=0) print(gt) np.testing.assert_allclose(add_nd.asnumpy(), gt) # print('reduce mul row = ') # print(mul_nd.asnumpy()) # gt = np.multiply.reduce(a_np ,axis=0,dtype = a.dtype) # print(gt) # np.testing.assert_allclose(mul_nd.asnumpy(), gt) print('reduce max row = ') print(gtm_nd.asnumpy()) gt = np.max(a_np, axis=0) print(gt) np.testing.assert_allclose(gtm_nd.asnumpy(), gt)
def test(): env = nnpu.get_env() nnpu.set_device(env) shape = (2, 2, 16) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] a = tvm.placeholder(shape, dtype_w, 'a') sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) k = tvm.reduce_axis((0, 2), 'k') add_buf = tvm.compute( (2, 16), lambda i, j: tvm.sum(a_buf[k, i, j], axis=k), 'add_buf') sph.MarkScope(add_buf) add_host, add_dram = nnpu.utils.CopyBufToH(add_buf, 'add', sph) k1 = tvm.reduce_axis((0, 2), 'k1') mul_buf = tvm.compute( (2, 16), lambda i, j: tvm.sum(a_buf[k1, i, j], axis=k1), 'mul_buf') sph.MarkScope(mul_buf) mul_host, mul_dram = nnpu.utils.CopyBufToH(mul_buf, 'mul', sph) s = tvm.create_schedule([add_host.op, mul_host.op]) sph.Transform(s) ko, ki = s[add_buf].split(add_buf.op.reduce_axis[0], factor=1) s[add_buf].reorder(ko, ki, *(s[add_buf].op.axis)) s[add_buf].tensorize(ki, env.intrins.get('MAddMerge', shape=shape, mode='w')) ko1, ki1 = s[mul_buf].split(mul_buf.op.reduce_axis[0], factor=1) s[mul_buf].reorder(ko1, ki1, *(s[mul_buf].op.axis)) s[mul_buf].tensorize(ki1, env.intrins.get('MMulMerge', shape=shape, mode='w')) print(nnpu.lower(s, [a, add_host, mul_host], simple_mode=True)) func = nnpu.build(s, [a, add_host, mul_host], 'nnpu', 'llvm', name='nnpu_func') #exit() ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(2, 2, 16), dtype=a.dtype, low=-16, high=16) a_nd = tvm.nd.array(a_np, ctx) add_nd = tvm.nd.array(np.zeros((2, 16)).astype(add_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros((2, 16)).astype(mul_host.dtype), ctx) func(a_nd, add_nd, mul_nd) print('a = ') print(a_np) print('reduce sum row = ') print(add_nd.asnumpy()) print('ground truth is: ') gt = np.sum(a_np, axis=0) print(gt) np.testing.assert_allclose(add_nd.asnumpy(), gt) print('reduce mul row = ') print(mul_nd.asnumpy()) gt = np.multiply.reduce(a_np, axis=0, dtype=a.dtype) print(gt) np.testing.assert_allclose(mul_nd.asnumpy(), gt)
def test_ib(): print('aaaa') env = nnpu.get_env() nnpu.set_device(env) shape = (16, ) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] a = tvm.placeholder(shape, dtype_w, name='a') w = shape[0] e = 16 def build_nms_ir(ten_in, ten_out): ib = tvm.ir_builder.create() imm_value = 10 ib.scope_attr(env.nnpu_axis, "coproc_scope", 0) p_in = ib.buffer_ptr(ten_in[0]) p_out = ib.buffer_ptr(ten_out[0]) #with ib.for_range(0,w, name="k") as k: with ib.for_range(0, w / e, name="i") as i: ib.emit( make_intrin_call( "void", 'VAddI', ten_out[0].access_ptr("w", 'uint32') + i * dtype_bytes(dtype_w), ten_in[0].access_ptr("r", 'uint32') + i * dtype_bytes(dtype_w), tvm.const(imm_value, 'float64'), env.cfg['vector_unit']['size'], 3)) stmt = ib.get() return stmt sph = ScheduleProcHelper() a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) sph.MarkScope(a_buf) out = tvm.extern(a_buf.shape, [a_buf], build_nms_ir, in_buffers=[ tvm.decl_buffer(a_buf.shape, dtype_w, data_alignment=dtype_bytes(dtype_w), scope='local.nnpu_scratchpad0') ], out_buffers=[ tvm.decl_buffer(a_buf.shape, dtype_w, data_alignment=dtype_bytes(dtype_w), scope='local.nnpu_scratchpad0') ], dtype=dtype_w, name="test_ir") sph.MarkScope(out) out_host, out_dram = nnpu.utils.CopyBufToH(out, 'out', sph) s = tvm.create_schedule([out_host.op]) sph.Transform(s) print(tvm.lower(s, [a, out_host], simple_mode=True)) print(nnpu.lower(s, [a, out_host], simple_mode=True)) # exit(0) func = nnpu.build(s, [a, out_host], 'nnpu', 'llvm', name='nnpu_test') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(16, ), dtype=a.dtype, low=0, high=127) a_nd = tvm.nd.array(a_np, ctx) b_nd = tvm.nd.array(np.zeros(16, ).astype(out_host.dtype), ctx) func(a_nd, b_nd) print('a = ') print(a_np) print('xjb sum = ') print(b_nd.asnumpy()) return
def test(): env = nnpu.get_env() nnpu.set_dump(False) #==================================# # ------ first define shapes ------ #==================================# # input data layout: HWC in_shape = (32, 32, 128) # pooling windows size, height == width. cell_shape = 4 # in this demo we don't do padding, so input data height and width must be divisible to pooling window size. assert in_shape[0] % cell_shape == 0, 'error' assert in_shape[1] % cell_shape == 0, 'error' nvctr_unit = env.cfg['vector_unit']['size'] assert in_shape[2] % nvctr_unit == 0, 'channel not divisible to vector unit size' out_shape = (in_shape[0] // cell_shape, in_shape[1] // cell_shape, in_shape[2]) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] sph = ScheduleProcHelper() #=================================================================# # ------ after all shapes defined, begin compute describing. ------ #=================================================================# a = tvm.placeholder(in_shape, dtype_w, 'a') # first copy to scratchpad. a_buf, _1 = nnpu.utils.CopyHtoBuf(a, 'a', sph) # stage 1, sum up the pixels in every pooling window. # the extent of two reduction axes are sizes of pooling window. k1 = tvm.reduce_axis((0,cell_shape), 'k1') k2 = tvm.reduce_axis((0,cell_shape), 'k2') pooling_buf = tvm.compute(out_shape, lambda i, j, k: tvm.sum(a_buf[i * cell_shape + k1, j * cell_shape + k2, k], axis=[k1, k2]), 'pooling_buf') sph.MarkScope(pooling_buf) sum_host, _ = nnpu.utils.CopyBufToH(pooling_buf, 'step3', sph) # stage 2, divide by cell_shape^2, to compute average. Imm = tvm.const(cell_shape*cell_shape, env.cfg['dtype_w']) step3_buf = tvm.compute(out_shape, lambda i, j, k: pooling_buf[i,j,k]/Imm, 'step3_buf') sph.MarkScope(step3_buf) # copy back to host. step3_host, step3_dram = nnpu.utils.CopyBufToH(step3_buf, 'step3',sph) # ------ this ends the computation description. ------ #==================================# # ------ begin scheduling ------ #==================================# s = tvm.create_schedule([step3_host.op, sum_host.op]) sph.Transform(s) #tensorize i, j, k = pooling_buf.op.axis k1, k2 = pooling_buf.op.reduce_axis # split the reduce_axis by factor 1, to produce a dummy reduce axis. # this is a trick to enable tensorize, due to limitation of tvm's tensorize pattern matcher. ko, ki = s[pooling_buf].split(k2, factor=1) xo, xi = s[pooling_buf].split(k, factor=nvctr_unit) # reorder axes. # put xo right before ki to eliminate memory dependency between two consecutive VAddV instruction s[pooling_buf].reorder( i, j, k1, ko, xo, ki, xi) s[pooling_buf].tensorize(ki, env.intrins.get('VAddMerge', mode='w')) # unroll # s[pooling_buf].unroll(xo) # s[pooling_buf].unroll(ko) # split and tensorize. xo2, xi2 = s[step3_buf].split(step3_buf.op.axis[2], factor=nvctr_unit) s[step3_buf].reorder( step3_buf.op.axis[0], step3_buf.op.axis[1], xo2, xi2) s[step3_buf].tensorize(xi2, env.intrins.get('VDivI',imm_value=Imm.value, mode='w')) # s[step3_buf].unroll(xo2) #==================================# # ------ this ends the scheduling ------ #==================================# print(nnpu.lower(s, [a, sum_host, step3_host], simple_mode=True)) # exit() func = nnpu.build(s, [a, sum_host, step3_host], 'nnpu', 'llvm', name='nnpu_func') print('------------------- device module 1 TVM IR: ') print(func.imported_modules[0].get_source('ir')) print('------------------- device module 1 uop: ') print(func.imported_modules[0].get_source('uop')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=in_shape, dtype=a.dtype, low = -128, high = 127) a_nd = tvm.nd.array(a_np, ctx) c_nd = tvm.nd.array(np.zeros(out_shape, dtype=step3_host.dtype), ctx) s_nd = tvm.nd.array(np.zeros(out_shape, dtype=step3_host.dtype), ctx) func(a_nd, s_nd, c_nd) # gt = mean_pooling_sum(in_shape, out_shape, cell_shape, a_np, a.dtype) # np.testing.assert_allclose(s_nd.asnumpy(), gt) # print('sum is ok') gt=mean_pooling(in_shape,out_shape,cell_shape,a_np,a.dtype) np.testing.assert_allclose(c_nd.asnumpy(), gt) print('test passed')