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() 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(): 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')
# split output core_extent = 4 xh, xw = s[res_host].op.axis xwo, xwi = s[res_host].split(xw, nparts=core_extent) s[res_host].reorder(xwo, xh, xwi) s[res_host].pragma(xh, env.dma_copy_from_buf) # compute_at s[a_buf].compute_at(s[res_host], xwo) s[b_buf].compute_at(s[res_host], xwo) s[res_acc].compute_at(s[res_host], xwo) s[res_buf].compute_at(s[res_host], xwo) # thread bind s[res_host].bind(xwo, tvm.thread_axis('coreIdx')) print(nnpu.lower(s, [a, b, res_host], simple_mode=True)) func = nnpu.build(s, [a, b, res_host], 'nnpu', 'llvm', 'nnpu_func') # print('------------------- device module 1 asm code: ') # print(func.imported_modules[0].get_source('ll')) ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=shape1_tiled, dtype=a.dtype, low=-16, high=16) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=shape2_tiled, dtype=b.dtype, low=-16, high=16)
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())
# ------ this ends the computation description. ------ #==================================# # ------ begin scheduling ------ #==================================# s = nnpu.create_schedule([tile_host.op]) # since all operations are scratchpad copy, all we need to do is pragma. # this is done by the helper functions, so nothing to do here. #==================================# # ------ this ends the scheduling ------ #==================================# print(tvm.lower(s, [a, tile_host], simple_mode=True)) print(nnpu.lower(s, [a, tile_host], simple_mode=True)) func = nnpu.build(s, [a, tile_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')) # exit() ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=(8, 8), dtype=a.dtype, low = -10000, high = 10000) a_nd = tvm.nd.array(a_np, ctx) #b_np = np.random.randint(size=(4, 32), dtype=b.dtype, low = -10000, high = 10000) #b_nd = tvm.nd.array(b_np, ctx) re_nd = tvm.nd.array(np.zeros((2, 2, 4, 4), dtype=tile_host.dtype), ctx)
s[kernel_buf].compute_at(s[res_host], oco) s[feature_buf].compute_at(s[res_host], pwo) s[conv].compute_at(s[res_host], pwo) s[pooling_buf].compute_at(s[res_host], pwi) # add copy pragma. s[feature_buf].pragma(feature_buf.op.axis[-1], env.dma_copy_to_buf) s[kernel_buf].pragma(kernel_buf.op.axis[-1], env.dma_copy_to_buf) s[res_host].pragma(oci, env.dma_copy_from_buf) s[conv_buf].pragma(conv_buf.op.axis[1], env.copy_acc2buf) s[conv].pragma(s[conv].leaf_iter_vars[-2], env.scratchpad_copy) #==================================# # ------ this ends the scheduling ------ #==================================# print(nnpu.lower(s, [feature, kernel, res_host], simple_mode=True)) # func = tvm.build(s, [feature, kernel, res_host], 'llvm', 'llvm', 'nnpu_conv') func = nnpu.build(s, [feature, kernel, res_host], 'nnpu', 'llvm', 'nnpu_conv') # print('------------------- device module 1 asm code: ') # print(func.imported_modules[0].get_source('asm')) 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')) # exit(0) ctx = tvm.nd.TVMContext(13, 0) fm_np = np.random.randint(size=shape, dtype=feature.dtype, low = -16, high = 16) fm_nd = tvm.nd.array(fm_np, ctx)
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() 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')
# we can move gemm into acc2buffer copy. xo, xi = s[out_buf].split(out_buf.op.axis[0], factor=gemm_shape[0]) s[prod_buf].compute_at(s[out_buf], xo) s[out_buf].pragma(xi, env.copy_acc2buf) # split and tensorize VAddV. nvctr_unit = env.cfg['vector_unit']['size'] xo, xi = s[res_buf].split(res_buf.op.axis[0], factor=nvctr_unit) s[res_buf].tensorize(xi, env.intrins.get('VAddV', mode='w')) #==================================# # ------ this ends the scheduling ------ #==================================# # with nnpu.build_config(dump_pass_ir=True): with nnpu.build_config(): print(nnpu.lower(s, [weight, data, bias, res_host], simple_mode=True)) func = nnpu.build(s, [weight, data, bias, res_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=weight_shape, dtype=weight.dtype, low=-32, high=32)
def test(): with ScheduleProcHelper(): env = nnpu.get_env() # nnpu.set_dump(True) dtype_n, dtype_w = env.cfg['dtype_n'], env.cfg['dtype_w'] shape = (48, ) nvctr_unit = env.cfg['vector_unit']['size'] assert shape[0] % nvctr_unit == 0, 'error' a = tvm.placeholder(shape, dtype_n, 'a') b = tvm.placeholder(shape, dtype_n, 'b') sph = ScheduleProcHelper.current a_buf, a_dram = nnpu.utils.CopyHtoBuf(a, 'a', sph) b_buf, b_dram = nnpu.utils.CopyHtoBuf(b, 'b', sph) 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) plus2 = tvm.compute(shape, lambda i: c_host[i] + tvm.const(2, 'int8'), 'plus2') s = tvm.create_schedule([plus2.op]) sph.Transform(s) xo, xi = s[c_buf].split(c_buf.op.axis[0], factor=nvctr_unit) s[c_buf].tensorize(xi, env.intrins.get('VAddV', mode='n')) print(nnpu.lower(s, [a, b, plus2], simple_mode=True)) # exit() func = nnpu.build(s, [a, b, plus2], 'nnpu', 'llvm', name='nnpu_exp') # exit() ctx = tvm.nd.TVMContext(13, 0) # print('------------------- host module llvm IR: ') # print(func.get_source('ll')) 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')) for i in range(0, 5): 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) plus2_nd = tvm.nd.array(np.zeros(shape).astype(plus2.dtype), ctx) # exit() func(a_nd, b_nd, plus2_nd) print('a = ') print(a_np) print('b = ') print(b_np) print('a + b + 2 =') print(plus2_nd.asnumpy()) print("numpy ground truth is") gt = a_np + b_np + 2 print(gt) np.testing.assert_allclose(plus2_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')
xo, xi = s[sum_buf].split(sum_buf.op.axis[0], factor=insn_shape[0]) yo, yi = s[sum_buf].split(sum_buf.op.axis[1], factor=insn_shape[1]) s[sum_buf].reorder(xo, yo, xi, yi) s[sum_buf].tensorize(xi, env.intrins.get('MAddM', shape=insn_shape, mode='n')) # xo, xi = s[sum_buf].mul_buf(mul_buf.op.axis[0], factor=insn_shape[0]) # yo, yi = s[sum_buf].split(sum_buf.op.axis[1], factor=insn_shape[1]) # s[sum_buf].reorder(xo, yo, xi, yi) s[mul_buf].tile(mul_buf.op.axis[0], mul_buf.op.axis[1], insn_shape[0], insn_shape[1]) s[mul_buf].tensorize( s[mul_buf].leaf_iter_vars[2], env.intrins.get('MMulM', shape=insn_shape, mode='inc')) print(nnpu.lower(s, [a, b, sum_host, mul_host], simple_mode=True)) func = nnpu.build(s, [a, b, sum_host, mul_host], 'nnpu', 'llvm', name='nnpu_func') ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=shape, dtype=a.dtype, low=-32, high=32) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=shape, dtype=b.dtype, low=-32, high=32) b_nd = tvm.nd.array(b_np, ctx) sum_nd = tvm.nd.array(np.zeros(shape, dtype=sum_host.dtype), ctx) mul_nd = tvm.nd.array(np.zeros(shape, dtype=mul_host.dtype), ctx)
def test(): env = nnpu.get_env() shape = (16, 16) a_host = tvm.placeholder(shape, env.cfg['dtype_n'], 'a_host') a = tvm.compute(shape, lambda *i: a_host(*i), name='a') a_buf = tvm.compute(shape, lambda *i: a(*i), name='a_buf') vctr_shape = (1, 16) b_host = tvm.placeholder(vctr_shape, env.cfg['dtype_n'], 'b_host') b = tvm.compute(vctr_shape, lambda *i: b_host(*i), name='b') b_buf = tvm.compute(vctr_shape, lambda *i: b(*i), name='b_buf') dtype_w = env.cfg['dtype_w'] mul_shape = (1, 16) k = tvm.reduce_axis((0, 16), 'k') c_buf = tvm.compute( mul_shape, lambda i, j: tvm.sum( b_buf[i, k].astype(dtype_w) * a_buf[j, k].astype(dtype_w), axis=k)) out_shape = (16, ) bias_host = tvm.placeholder(out_shape, env.cfg['dtype_w'], 'bias_host') bias = tvm.compute(out_shape, lambda *i: bias_host(*i), 'bias') bias_buf = tvm.compute(out_shape, lambda *i: bias(*i), 'bias_buf') #c = tvm.compute(out_shape, lambda *i: c_buf(*i), name='c') #c_host = tvm.compute(out_shape, lambda *i: c(*i), name='c_host') out_buf = tvm.compute(out_shape, lambda i: c_buf[0, i] + bias_buf[i], 'out_buf') out = tvm.compute(out_shape, lambda *i: out_buf(*i), 'out') out_host = tvm.compute(out_shape, lambda *i: out(*i), 'out_host') s = tvm.create_schedule(out_host.op) # mark variable scopes s[a].set_scope(env.dram_scope) s[b].set_scope(env.dram_scope) s[bias].set_scope(env.dram_scope) s[out].set_scope(env.dram_scope) s[a_buf].set_scope(env.uni_scratchpad_scope) s[b_buf].set_scope(env.uni_scratchpad_scope) s[c_buf].set_scope(env.uni_scratchpad_scope) s[bias_buf].set_scope(env.uni_scratchpad_scope) s[out_buf].set_scope(env.uni_scratchpad_scope) #print(dir(s[b].op.body)) # mark compiler pragmas s[a].pragma(s[a].op.axis[0], env.dma_copy_pragma) s[b].pragma(s[b].op.axis[0], env.dma_copy_pragma) s[bias].pragma(s[bias].op.axis[0], env.dma_copy_pragma) s[out_host].pragma(s[out_host].op.axis[0], env.dma_copy_pragma) s[a_buf].pragma(s[a_buf].op.axis[0], env.scratchpad_ls) s[b_buf].pragma(s[b_buf].op.axis[0], env.scratchpad_ls) s[bias_buf].pragma(s[bias_buf].op.axis[0], env.scratchpad_ls) s[out].pragma(s[out].op.axis[0], env.scratchpad_ls) #s[a_buf].compute_at(s[b_buf], b_buf.op.axis[0]) # tensorize #s[b_buf].tensorize(s[b_buf].op.axis[1], env.intrins.get('VEXP', mode='inc')) s[c_buf].tensorize(s[c_buf].op.axis[0], env.intrins.get('GEMM', shape=(1, 16, 16), mode='inc')) #outer, inner = out_buf.op.axis #s[out_buf].reorder(inner, outer) #print(outer) #print(tvm.lower(s, [a_host, b_host, bias_host, out_host], simple_mode=True)) s[out_buf].tensorize(s[out_buf].op.axis[0], env.intrins.get('VAddV', mode='w')) # build print(tvm.lower(s, [a_host, b_host, bias_host, out_host], simple_mode=True)) print( nnpu.lower(s, [a_host, b_host, bias_host, out_host], simple_mode=True)) #exit() func = nnpu.build(s, [a_host, b_host, bias_host, out_host], 'nnpu', 'llvm', name='nnpu_exp') print('function built: ') #print(func.get_source()) # prepare data ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=shape, dtype=a_host.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=vctr_shape, dtype=b_host.dtype, low=0, high=64) #b_np = np.random.random(size=vctr_shape).astype(b_host.dtype) b_nd = tvm.nd.array(b_np, ctx) bias_np = np.random.randint(size=out_shape, dtype=bias_host.dtype, low=0, high=64) #bias_np = np.random.random(size=out_shape).astype(bias_host.dtype) bias_nd = tvm.nd.array(bias_np, ctx) out_nd = tvm.nd.array(np.zeros(out_shape).astype(out_host.dtype), ctx) # run func(a_nd, b_nd, bias_nd, out_nd) print('run finished') print('a=') print(a_np) print('b=') print(b_np) print('bias=') print(bias_np) print('out=') print(out_nd.asnumpy()) print('numpy ground truth is: ') gt = np.dot(b_np.astype(dtype_w), a_np.astype(dtype_w).transpose((1, 0))).reshape( (16, )) + bias_np print(gt) np.testing.assert_allclose(out_nd.asnumpy(), gt)
def test(): pass if (False): print('-----') with ScheduleProcHelper(): env = nnpu.get_env() shape = (16, 64) a_host = tvm.placeholder(shape, env.cfg['dtype_n'], 'a_host') a_buf, _ = nnpu.utils.CopyHtoBuf(a_host, 'a') vctr_shape = (64, ) b_host = tvm.placeholder(vctr_shape, env.cfg['dtype_n'], 'b_host') b_buf, _ = nnpu.utils.CopyHtoBuf(b_host, 'b') dtype_w = env.cfg['dtype_w'] out_shape = (4, 16) k = tvm.reduce_axis((0, 16), 'k') c_buf = tvm.compute( out_shape, lambda j, i: tvm.sum(a_buf[i, j * 16 + k].astype( dtype_w) * b_buf[j * 16 + k].astype(dtype_w), axis=k)) utils.MarkScope(c_buf) c_host, _ = utils.CopyBufToH(c_buf, 'c') s = nnpu.create_schedule(c_host.op) # mark variable scopes # tensorize s[c_buf].tensorize( s[c_buf].op.axis[1], env.intrins.get('GEMM', shape=(16, 16, 1), mode='inc', reduce=True)) # build print(tvm.lower(s, [a_host, b_host, c_host], simple_mode=True)) print(nnpu.lower(s, [a_host, b_host, c_host], simple_mode=True)) #exit() func = nnpu.build(s, [a_host, b_host, c_host], 'nnpu', 'llvm', name='nnpu_exp') print('function built: ') print('------------------- device module 1 asm code: ') print(func.imported_modules[0].get_source('asm')) #print(func.get_source()) # prepare data ctx = tvm.nd.TVMContext(13, 0) a_np = np.random.randint(size=shape, dtype=a_host.dtype, low=-32, high=32) # a_np = np.ones(shape).astype(a_host.dtype) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=vctr_shape, dtype=b_host.dtype, low=-16, high=16) # b_np = np.ones(vctr_shape).astype(b_host.dtype) b_nd = tvm.nd.array(b_np, ctx) out_nd = tvm.nd.array(np.zeros(out_shape).astype(c_host.dtype), ctx) # run func(a_nd, b_nd, out_nd) print('run finished') print('a=') print(a_np) print('b=') print(b_np) print('out=') out_np = out_nd.asnumpy() out_np = np.sum(out_np, axis=0) print(out_np) print('numpy ground truth is: ') gt = np.dot(a_np.astype(dtype_w), b_np.astype(dtype_w)) #gt = np.greater(np.dot(a_np.astype(dtype_w), b_np.astype(dtype_w)), bias_np) print(gt) np.testing.assert_allclose(out_np, gt)
xrow = sum1.op.axis[0] s[sum1].reorder(xblock, xrow, xcol) s[sum1].tensorize( xrow, env.intrins.get('MReduceSumRow', shape=(nRow, factor), scope_out='acc', mode='w')) s[sum2].tensorize(sum2.op.reduce_axis[0], env.intrins.get('VReduceSum', shape=(8, ), mode='w')) xo, xi = s[softmax].split(softmax.op.axis[0], 16) s[softmax].tensorize(xi, env.intrins.get('VDivS', mode='w')) print(nnpu.lower(s, [a, softmax_host], simple_mode=True)) func = nnpu.build(s, [a, softmax_host], 'nnpu', 'llvm', 'nnpu_func') 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.random(shape).astype(a.dtype) * 2 a_nd = tvm.nd.array(a_np, ctx) # sigmoid_nd = tvm.nd.array(np.zeros(shape, dtype=sigmoid_host.dtype), ctx) softmax_nd = tvm.nd.array(np.zeros(shape, dtype=softmax_host.dtype), ctx)
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() 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 print('a = ') print(a_np) print('reduce sum result is:') print(c_nd.asnumpy()) print("numpy ground truth is") gt = np.sum(a_np, axis=(1,), keepdims=True) print(gt) np.testing.assert_allclose(c_nd.asnumpy(), gt) print('reduce max result is:') print(max_nd.asnumpy()) np.testing.assert_allclose(max_nd.asnumpy(), np.max(a_np, axis=(1,), keepdims=True)) print('reduce min result is:') print(min_nd.asnumpy()) np.testing.assert_allclose(min_nd.asnumpy(), np.min(a_np, axis=(1,), keepdims=True))
s[sigmoid].pragma(sigmoid.op.axis[0], env.dma_copy_from_buf) # tensorize vector_unit_size = 32 xo, xi = s[exp_buf].split(exp_buf.op.axis[0], vector_unit_size) s[exp_buf].tensorize(xi, env.intrins.get('VExp', mode='w', size=vector_unit_size)) xo, xi = s[log_buf].split(log_buf.op.axis[0], vector_unit_size) s[log_buf].tensorize(xi, env.intrins.get('VLog', mode='w', size=vector_unit_size)) xo, xi = s[tanh_buf].split(tanh_buf.op.axis[0], vector_unit_size) s[tanh_buf].tensorize(xi, env.intrins.get('VTanh', mode='w', size=vector_unit_size)) xo, xi = s[sigmoid_buf].split(sigmoid_buf.op.axis[0], vector_unit_size) s[sigmoid_buf].tensorize(xi, env.intrins.get('VSigmoid', mode='w', size=vector_unit_size)) print(nnpu.lower(s, [a, exp, log, tanh, sigmoid], simple_mode=True)) func = nnpu.build(s, [a, exp, log, tanh, sigmoid], 'nnpu', 'llvm', 'nnpu_func') 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.random(shape).astype(a.dtype) * 2 a_nd = tvm.nd.array(a_np, ctx) exp_nd = tvm.nd.array(np.zeros(shape, dtype=exp.dtype), ctx) log_nd = tvm.nd.array(np.zeros(shape, dtype=log.dtype), ctx)
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')
x, y, c = s[pooling_host].op.axis xo, yo, xi, yi = s[pooling_host].tile(x, y, dim_x, dim_y) x_vt, xo = s[pooling_host].split(xo, nparts=2) co, ci = s[pooling_host].split(c, dim_c) s[pooling_host].reorder(x_vt, xo, yo, co, xi, yi, ci) s[pooling_buf].compute_at(s[pooling_host], co) s[pooling_host].bind(x_vt, tvm.thread_axis('cthread')) # pragma s[data_buf].pragma(data_buf.op.axis[0], env.dma_copy_to_buf) s[pooling_host].pragma(xi, env.dma_copy_from_buf) #==================================# # ------ this ends the scheduling ------ #==================================# # print(tvm.lower(s, [data, pooling_host], simple_mode=True)) print(nnpu.lower(s, [data, pooling_host], simple_mode=True)) # exit() func = nnpu.build(s, [data, pooling_host], 'nnpu', 'llvm', name='nnpu_func') ctx = tvm.nd.TVMContext(13, 0) data_np = np.random.randint(size=in_shape, dtype=data.dtype, low=-128, high=127) data_nd = tvm.nd.array(data_np, ctx) res_nd = tvm.nd.array(np.zeros(out_shape, dtype=pooling_host.dtype), ctx)
# factor_x, factor_y = 1, 8 # to split outter loop s[out_host].reorder(xo, yo, xi, yi, tx, ty) s[out_host].pragma(yi, env.dma_copy_from_buf) # bind to virtual thread # s[out_host].bind(by, tvm.thread_axis("cthread")) # compute_at s[out_acc].compute_at(s[out_host], xo) # s[out_buf].compute_at(s[out_host], xi) s[res_buf].compute_at(s[out_host], xi) # s[bias_buf].compute_at(s[out_host], by) print(tvm.lower(s, [a, b, bias, out_host], simple_mode=True)) # exit() print(nnpu.lower(s, [a, b, bias, out_host], simple_mode=True)) # exit(0) func = nnpu.build(s, [a, b, bias, out_host], 'nnpu', 'llvm', '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=shape1, dtype=a.dtype, low=-16, high=16) a_nd = tvm.nd.array(a_np, ctx) b_np = np.random.randint(size=shape2, dtype=b.dtype, low=-16, high=16) b_nd = tvm.nd.array(b_np, ctx) bias_np = np.random.randint(size=(shape2[0], ), dtype=bias.dtype, low=-128,
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, 16) a_host = tvm.placeholder(shape, env.cfg['dtype_n'], 'a_host') print('a host ' + str(a_host)) a = tvm.compute(shape, lambda *i: a_host(*i), name='a') a_buf = tvm.compute(shape, lambda *i: a(*i), name='a_buf') b_buf = tvm.compute( shape, lambda i, j: tvm.log(a_buf[i, j].astype(env.cfg['dtype_w'])), name='b_buf') b = tvm.compute(shape, lambda *i: b_buf(*i), name='b') b_host = tvm.compute(shape, lambda *i: b(*i), name='b_host') s = tvm.create_schedule(b_host.op) # mark variable scopes s[a].set_scope(env.dram_scope) s[b].set_scope(env.dram_scope) s[a_buf].set_scope(env.uni_scratchpad_scope) s[b_buf].set_scope(env.uni_scratchpad_scope) #print # (dir(s[b].op.body)) # mark compiler pragmas s[a].pragma(s[a].op.axis[0], env.dma_copy_pragma) s[b_host].pragma(s[b_host].op.axis[0], env.dma_copy_pragma) s[a_buf].pragma(s[a_buf].op.axis[0], env.scratchpad_ls) s[b].pragma(s[b].op.axis[0], env.scratchpad_ls) s[a_buf].compute_at(s[b_buf], b_buf.op.axis[0]) # tensorize s[b_buf].tensorize(s[b_buf].op.axis[1], env.intrins.get('VLOG', mode='inc')) # build print(tvm.lower(s, [a_host, b_host], simple_mode=True)) print(nnpu.lower(s, [a_host, b_host], simple_mode=True)) #exit() func = nnpu.build(s, [a_host, b_host], 'nnpu', 'llvm', name='nnpu_log') print('function built: ') #print(func.get_source()) # prepare data ctx = tvm.nd.TVMContext(13, 0) #??? print('i want to know:') print(ctx.exist) a_np = np.random.randint(size=shape, dtype=a_host.dtype, low=1, high=20) a_nd = tvm.nd.array(a_np, ctx) b_nd = tvm.nd.array(np.zeros(shape).astype(b_host.dtype), ctx) # run func(a_nd, b_nd) print('run finished') b_np = b_nd.asnumpy() print('a=') print(a_np) print('b=') print(b_np) print('ground truth =') gt = np.log(a_np, dtype=b_host.dtype) print(gt) np.testing.assert_allclose(b_np, 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
np.zeros((packed_shape[0] // 8, packed_shape[1] // 8, 8, 8), dtype='int8')) gt_func(a_nd, gt_nd) real_nd = tvm.nd.array( np.zeros((packed_shape[0] // 8, packed_shape[1] // 8, 8, 8), dtype='int8')) func(a_nd, real_nd) gt_np = gt_nd.asnumpy() real_np = real_nd.asnumpy() # print(gt_np) # print(real_np) np.testing.assert_allclose(gt_np, real_np) print('test passed') exit() print(nnpu.lower(s, [feature, tiled], simple_mode=True)) func = nnpu.build(s, [feature, tiled], 'nnpu', 'llvm', 'im2col_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')) a_np = np.random.randint(size=(fh, fw, fc), dtype='int8', low=-128, high=127) a_nd = tvm.nd.array(a_np) gt_nd = tvm.nd.array( np.zeros((packed_shape[0] // 8, packed_shape[1] // 8, 8, 8), dtype='int8'))