def traverse(op): """Traverse operators from computation graph""" # inline all one-to-one-mapping operators except the last stage (output) if tag.is_broadcast(op.tag): if op not in s.outputs: s[op].compute_inline() for tensor in op.input_tensors: if tensor.op.input_tensors: traverse(tensor.op) if 'conv2d_nchw' in op.tag: # print('Run in x86-rasp schedule') output = op.output(0) conv_out = op.input_tensors[0] kernel_vec = conv_out.op.input_tensors[1] kernel = kernel_vec.op.input_tensors[0] data_vec = conv_out.op.input_tensors[0] data = data_vec.op.input_tensors[0] data_pad = None if isinstance(data.op, tvm.tensor.ComputeOp) and "pad" in data.op.tag: data_pad = data data = data_pad.op.input_tensors[0] padding = infer_pad(data, data_pad) if data_pad is None: stride = infer_stride(data, kernel, output) else: stride = infer_stride(data_pad, kernel, output) wkl = _get_workload(data, kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) return _SCH_TO_SCH_FUNC[type(sch)](s, data, data_pad, data_vec, kernel, kernel_vec, conv_out, output, outs[0])
def traverse(op): """Traverse operators from computation graph""" # inline all one-to-one-mapping operators except the last stage (output) if tag.is_broadcast(op.tag): if op not in s.outputs: s[op].compute_inline() for tensor in op.input_tensors: if tensor.op.input_tensors: traverse(tensor.op) if 'conv2d_nChwc' in op.tag: print('Got conv2d_nChwc tag: ' + str(op.tag)) output = op.output(0) # conv_out = op.input_tensors[0] conv_out = output kernel = conv_out.op.input_tensors[1] # kernel = kernel_vec.op.input_tensors[0] data_vec = conv_out.op.input_tensors[0] data = data_vec.op.input_tensors[0] \ if isinstance(data_vec.op, tvm.tensor.ComputeOp) and len(data_vec.op.input_tensors) > 0 and "pad" not in data_vec.op.tag \ else data_vec data_pad = None if isinstance(data.op, tvm.tensor.ComputeOp) and "pad" in data.op.tag: data_pad = data data = data_pad.op.input_tensors[0] n, ic_chunk, h, w, ic_block = [x.value for x in data.shape] ic = ic_chunk * ic_block original_data = tvm.placeholder((n, ic, h, w), dtype=output.dtype) if data_pad is not None: n, _, pad_h, pad_w, _ = [x.value for x in data_pad.shape] original_data_pad = tvm.placeholder((n, ic, pad_h, pad_w), dtype=output.dtype) padding = infer_pad(original_data, original_data_pad) else: padding = (0, 0) oc, kh, kw = kernel_size original_kernel = tvm.placeholder((oc, ic, kh, kw), dtype=output.dtype) n, oc_chunk, oh, ow, oc_block = [x.value for x in output.shape] original_output = tvm.placeholder((n, oc_chunk * oc_block, oh, ow), dtype=output.dtype) if data_pad is None: stride = infer_stride(original_data, original_kernel, original_output) else: stride = infer_stride(original_data_pad, original_kernel, original_output) wkl = _get_workload(original_data, original_kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) _SCH_TO_SCH_FUNC[type(sch)](s, data, data_pad, data_vec, kernel, conv_out, output, outs[0])
def _schedule_conv(s, data, data_pad, data_vec, kernel, conv_out, output, last): # no stride and padding info here padding = infer_pad(data, data_pad) if data_pad is None: stride = infer_stride(data, kernel, output) else: stride = infer_stride(data_pad, kernel, output) wkl = get_workload(data, kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) # A, W = data, kernel_pack A0, A1 = data_pad, data_vec # schedule data if A0 is not None: s[A0].compute_inline() batch, ic_chunk, ih, ic_block, iw = s[A1].op.axis parallel_axis = s[A1].fuse(ic_chunk, ih) s[A1].parallel(parallel_axis) C, O0, O = conv_out, output, last CC = s.cache_write(C, 'global') batch, oc_chunk, oh, ow, oc_block = s[C].op.axis oh_outer, oh_inner = s[C].split(oh, factor=sch.oh_factor) s[C].vectorize(oc_block) s[CC].compute_at(s[C], oh_outer) _, oc_chunk, oh, ow, oc_block = s[CC].op.axis ic, = s[CC].op.reduce_axis ic_chunk, ic_block = s[CC].split(ic, factor=sch.ic_bn) oh_outer, oh_inner = s[CC].split(oh, factor=sch.oh_factor) ow_outer, ow_inner = s[CC].split(ow, factor=sch.ow_factor) s[CC].reorder(oc_chunk, oh_outer, ow_outer, ic_chunk, ic_block, oh_inner, ow_inner, oc_block) s[CC].vectorize(oc_block) s[CC].unroll(ow_inner) s[CC].unroll(oh_inner) if O0 != O: s[O0].compute_inline() batch, oc, oh, ow = s[O].op.axis oc_chunk, oc_block = s[O].split(oc, factor=sch.oc_bn) oh_outer, oh_inner = s[O].split(oh, factor=sch.oh_factor) ow_outer, ow_inner = s[O].split(ow, factor=sch.ow_factor) s[O].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block) parallel_axis = s[O].fuse(oc_chunk, oh_outer) s[C].compute_at(s[O], parallel_axis) s[O].vectorize(oc_block) s[O].parallel(parallel_axis) return s
def traverse(op): """Traverse operators from computation graph""" # inline all one-to-one-mapping operators except the last stage (output) if tag.is_broadcast(op.tag): if op not in s.outputs: s[op].compute_inline() for tensor in op.input_tensors: if tensor.op.input_tensors: traverse(tensor.op) if 'conv2d_nchw' in op.tag: output = op.output(0) conv_out = op.input_tensors[0] kernel = conv_out.op.input_tensors[1] # kernel = kernel_vec.op.input_tensors[0] data_vec = conv_out.op.input_tensors[0] data = data_vec.op.input_tensors[0] data_pad = None if isinstance(data.op, tvm.tensor.ComputeOp) and "pad" in data.op.tag: data_pad = data data = data_pad.op.input_tensors[0] padding = infer_pad(data, data_pad) _, ic, _, _ = [x.value for x in data.shape] oc = num_filter kh, kw = kernel_size original_kernel = tvm.placeholder((oc, ic, kh, kw)) if data_pad is None: stride = infer_stride(data, original_kernel, output) else: stride = infer_stride(data_pad, original_kernel, output) wkl = _get_workload(data, original_kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) _SCH_TO_SCH_FUNC[type(sch)](s, data, data_pad, data_vec, kernel, conv_out, output, outs[0])
def _schedule_conv(s, data, data_pad, data_vec, kernel, kernel_pack, conv_out, output, last): # print('Run in avx512_conv_common sch') # no stride and padding info here """ C, O0, O = conv_out, output, last batch, oc, oh, ow = s[O].op.axis s[O].parallel(batch) return s """ padding = infer_pad(data, data_pad) if data_pad is None: stride = infer_stride(data, kernel, output) else: stride = infer_stride(data_pad, kernel, output) wkl = _get_workload(data, kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) HPAD, WPAD = wkl.hpad, wkl.wpad DOPAD = (HPAD != 0 and WPAD != 0) A, W = data, kernel_pack A0, A1 = data_pad, data_vec # schedule data if DOPAD: s[A0].compute_inline() batch, ic_chunk, ih, ic_block, iw = s[A1].op.axis parallel_axis = s[A1].fuse(ic_chunk, ih) s[A1].parallel(parallel_axis) s[A1].pragma(batch, "parallel_launch_point") s[A1].pragma(parallel_axis, "parallel_stride_pattern") s[A1].pragma(batch, "parallel_barrier_when_finish") # schedule kernel pack if False: oc_chunk, ic_chunk, oh, ow, ic_block, oc_block = s[W].op.axis s[W].reorder(oc_chunk, oh, ic_chunk, ow, ic_block, oc_block) if sch.oc_bn > 1: s[W].vectorize(oc_block) parallel_axis = s[W].fuse(oc_chunk, oh) s[W].parallel(parallel_axis) s[W].pragma(parallel_axis, "parallel_launch_point") s[W].pragma(parallel_axis, "parallel_stride_pattern") s[W].pragma(parallel_axis, "parallel_barrier_when_finish") # schedule conv C, O0, O = conv_out, output, last CC = s.cache_write(C, 'global') _, oc_chunk, oh, ow, oc_block = s[C].op.axis ow_chunk, ow_block = s[C].split(ow, factor=sch.ur_w) s[C].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block) s[C].fuse(oc_chunk, oh) s[C].vectorize(oc_block) s[CC].compute_at(s[C], ow_chunk) _, oc_chunk, oh, ow, oc_block = s[CC].op.axis ic, kh, kw = s[CC].op.reduce_axis ow_chunk, ow_block = s[CC].split(ow, factor=sch.ur_w) ic_chunk, ic_block = s[CC].split(ic, factor=sch.ic_bn) if sch.unroll_kw: s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kh, ic_block, kw, ow_block, oc_block) s[CC].unroll(kw) else: s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kh, kw, ic_block, ow_block, oc_block) s[CC].fuse(oc_chunk, oh) s[CC].vectorize(oc_block) s[CC].unroll(ow_block) if O0 != O: s[O0].compute_inline() batch, oc, oh, ow = s[O].op.axis ow_chunk, ow_block = s[O].split(ow, factor=sch.ur_w) oc_chunk, oc_block = s[O].split(oc, factor=sch.oc_bn) s[O].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block) parallel_axis = s[O].fuse(oc_chunk, oh) s[C].compute_at(s[O], parallel_axis) s[O].vectorize(oc_block) s[O].parallel(parallel_axis) s[O].pragma(batch, "parallel_launch_point") s[O].pragma(parallel_axis, "parallel_stride_pattern") s[O].pragma(batch, "parallel_barrier_when_finish") return s
def _schedule_conv(s, data, data_pad, data_vec, kernel, kernel_pack, conv_out, output, last): # print('Run in avx512_conv_1x1 sch') # no stride and padding info here padding = infer_pad(data, data_pad) if data_pad is None: stride = infer_stride(data, kernel, output) else: stride = infer_stride(data_pad, kernel, output) wkl = _get_workload(data, kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) A, W = data, kernel_pack A0, A1 = data_pad, data_vec # schedule data if A0 is not None: s[A0].compute_inline() batch, ic_chunk, ih, ic_block, iw = s[A1].op.axis parallel_axis = s[A1].fuse(ic_chunk, ih) s[A1].parallel(parallel_axis) s[A1].pragma(batch, "parallel_launch_point") s[A1].pragma(parallel_axis, "parallel_stride_pattern") s[A1].pragma(batch, "parallel_barrier_when_finish") # schedule kernel pack oc_chunk, ic_chunk, oh, ow, ic_block, oc_block = s[W].op.axis s[W].reorder(oc_chunk, oh, ic_chunk, ow, ic_block, oc_block) if sch.oc_bn > 1: s[W].vectorize(oc_block) parallel_axis = s[W].fuse(oc_chunk, oh) s[W].parallel(parallel_axis) s[W].pragma(parallel_axis, "parallel_launch_point") s[W].pragma(parallel_axis, "parallel_stride_pattern") s[W].pragma(parallel_axis, "parallel_barrier_when_finish") C, O0, O = conv_out, output, last CC = s.cache_write(C, 'global') batch, oc_chunk, oh, ow, oc_block = s[C].op.axis oh_outer, oh_inner = s[C].split(oh, factor=sch.oh_factor) s[C].vectorize(oc_block) s[CC].compute_at(s[C], oh_outer) _, oc_chunk, oh, ow, oc_block = s[CC].op.axis ic, = s[CC].op.reduce_axis ic_chunk, ic_block = s[CC].split(ic, factor=sch.ic_bn) oh_outer, oh_inner = s[CC].split(oh, factor=sch.oh_factor) ow_outer, ow_inner = s[CC].split(ow, factor=sch.ow_factor) s[CC].reorder(oc_chunk, oh_outer, ow_outer, ic_chunk, ic_block, oh_inner, ow_inner, oc_block) s[CC].vectorize(oc_block) s[CC].unroll(ow_inner) s[CC].unroll(oh_inner) if O0 != O: s[O0].compute_inline() batch, oc, oh, ow = s[O].op.axis oc_chunk, oc_block = s[O].split(oc, factor=sch.oc_bn) oh_outer, oh_inner = s[O].split(oh, factor=sch.oh_factor) ow_outer, ow_inner = s[O].split(ow, factor=sch.ow_factor) s[O].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block) parallel_axis = s[O].fuse(oc_chunk, oh_outer) s[C].compute_at(s[O], parallel_axis) s[O].vectorize(oc_block) s[O].parallel(parallel_axis) s[O].pragma(batch, "parallel_launch_point") s[O].pragma(parallel_axis, "parallel_stride_pattern") s[O].pragma(batch, "parallel_barrier_when_finish") return s
def _schedule_conv(s, data, data_pad, data_vec, kernel, conv_out, output, last): print("Run in prepack common sch") # no stride and padding info here padding = infer_pad(data, data_pad) if data_pad is None: stride = infer_stride(data, kernel, output) else: stride = infer_stride(data_pad, kernel, output) wkl = get_workload(data, kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) HPAD, WPAD = wkl.hpad, wkl.wpad DOPAD = (HPAD != 0 and WPAD != 0) # A, W = data, kernel_vec A0, A1 = data_pad, data_vec # schedule data if DOPAD: s[A0].compute_inline() batch, ic_chunk, ih, ic_block, iw = s[A1].op.axis parallel_axis = s[A1].fuse(ic_chunk, ih) s[A1].parallel(parallel_axis) # schedule conv C, O0, O = conv_out, output, last CC = s.cache_write(C, 'global') _, oc_chunk, oh, ow, oc_block = s[C].op.axis ow_chunk, ow_block = s[C].split(ow, factor=sch.reg_n) s[C].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block) s[C].fuse(oc_chunk, oh) s[C].vectorize(oc_block) s[CC].compute_at(s[C], ow_chunk) _, oc_chunk, oh, ow, oc_block = s[CC].op.axis ic, kh, kw = s[CC].op.reduce_axis ow_chunk, ow_block = s[CC].split(ow, factor=sch.reg_n) ic_chunk, ic_block = s[CC].split(ic, factor=sch.ic_bn) if sch.unroll_kw: s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kh, ic_block, kw, ow_block, oc_block) s[CC].unroll(kw) else: s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kh, kw, ic_block, ow_block, oc_block) s[CC].fuse(oc_chunk, oh) s[CC].vectorize(oc_block) s[CC].unroll(ow_block) if O0 != O: s[O0].compute_inline() batch, oc, oh, ow = s[O].op.axis ow_chunk, ow_block = s[O].split(ow, factor=sch.reg_n) oc_chunk, oc_block = s[O].split(oc, factor=sch.oc_bn) s[O].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block) parallel_axis = s[O].fuse(oc_chunk, oh) s[C].compute_at(s[O], parallel_axis) s[O].vectorize(oc_block) s[O].parallel(parallel_axis) return s
def _schedule_im2col_conv2d(wkl, sch, s, data, data_pad, data_col, data_vec, kernel, kernel_vec, conv_out, output, last): # no stride and padding info here padding = infer_pad(data, data_pad) if data_pad is None: stride = infer_stride(data, kernel, output) else: stride = infer_stride(data_pad, kernel, output) H, W = wkl.height, wkl.width CI = wkl.in_filter CO = wkl.out_filter HK, WK = wkl.hkernel, wkl.wkernel HPAD, WPAD = wkl.hpad, wkl.wpad HSTR, WSTR = wkl.hstride, wkl.wstride HCAT, WCAT = HK - 1, WK - 1 DOPAD = (HPAD != 0 and WPAD != 0) P = sch.vp Q = sch.vq UNROLL = sch.unroll A, B, C = data, kernel, last A0, A1, A2 = data_pad, data_col, data_vec B0 = kernel_vec C0, C1 = conv_out, output CC = s.cache_write(C0, "global") AA = s.cache_read(A2, "global", [CC]) BB = s.cache_read(B0, "global", [CC]) ##### Schedule CC _, co, im, vim, vco = s[C0].op.axis s[C0].unroll(vim) s[C0].vectorize(vco) s[CC].compute_at(s[C0], im) _, co, im, vim, vco = s[CC].op.axis ci, hk, wk = s[CC].op.reduce_axis s[CC].reorder(ci, hk, wk, vim, vco) s[CC].unroll(vim) s[CC].vectorize(vco) # s[CC].unroll(ccr) ### Schedule C _, co, h, w = s[C].op.axis im = s[C].fuse(h, w) im, vim = s[C].split(im, P) co, vco = s[C].split(co, Q) s[C].reorder(co, im, vim, vco) if sch.bc == 1: oaxis = co paxis = co else: oco, ico = s[C].split(co, sch.bc) oaxis = oco paxis = ico s[C].parallel(paxis) s[C].pragma(oaxis, "parallel_launch_point") s[C].pragma(paxis, "parallel_stride_pattern") s[C].pragma(oaxis, "parallel_barrier_when_finish") if C1 != C: s[C1].compute_inline() s[C0].compute_at(s[C], paxis) ##### Schedule A if DOPAD: s[A0].compute_inline() s[A1].compute_inline() s[AA].compute_at(s[CC], wk) s[AA].unroll(AA.op.axis[4]) _, im, _, _, _, _ = s[A2].op.axis if sch.ba == 1: oaxis = im paxis = im else: oim, iim = s[A2].split(im, sch.ba) oaxis = oim paxis = iim s[A2].parallel(paxis) s[A2].pragma(oaxis, "parallel_launch_point") s[A2].pragma(paxis, "parallel_stride_pattern") s[A2].pragma(oaxis, "parallel_barrier_when_finish") ##### Schedule B s[BB].compute_at(s[CC], wk) s[BB].vectorize(BB.op.axis[4]) co, _, _, _, _ = s[B0].op.axis if sch.bc == 1: oaxis = co paxis = co else: oco, ico = s[B0].split(co, sch.bc) oaxis = oco paxis = ico s[B0].parallel(paxis) s[B0].pragma(oaxis, "parallel_launch_point") s[B0].pragma(paxis, "parallel_stride_pattern") s[B0].pragma(oaxis, "parallel_barrier_when_finish") return s
def _schedule_spatial_conv2d(s, data, data_pad, data_vec, kernel, kernel_vec, conv_out, output, last): # no stride and padding info here padding = infer_pad(data, data_pad) if data_pad is None: stride = infer_stride(data, kernel, output) else: stride = infer_stride(data_pad, kernel, output) wkl = _get_workload(data, kernel, stride, padding, output.dtype) sch = _get_schedule(wkl) H, W = wkl.height, wkl.width CI, CO = wkl.in_filter, wkl.out_filter HK, WK = wkl.hkernel, wkl.wkernel HPAD, WPAD = wkl.hpad, wkl.wpad HSTR, WSTR = wkl.hstride, wkl.wstride HCAT, WCAT = HK - 1, WK - 1 DOPAD = (HPAD != 0 and WPAD != 0) VH = sch.vh VW = sch.vw VC = sch.vc UNROLL = sch.unroll A, B, C = data, kernel, last A0, A1 = data_pad, data_vec B0 = kernel_vec C0, C1 = conv_out, output CC = s.cache_write(C0, "global") _, co, oh, ow, vh, vw, vc = s[C0].op.axis if UNROLL: s[C0].unroll(vw) s[C0].vectorize(vc) s[CC].compute_at(s[C0], ow) _, co, oh, ow, vh, vw, vc = s[CC].op.axis ci, dh, dw = s[CC].op.reduce_axis s[CC].reorder(ci, dh, vh, dw, vw, vc) if UNROLL: s[CC].unroll(vw) s[CC].vectorize(vc) ##### Schedule A if DOPAD: s[A0].compute_inline() _, h, _, _, _, _ = s[A1].op.axis if sch.ba == 1: oaxis = h paxis = h else: oh, ih = s[A1].split(h, sch.ba) oaxis = oh paxis = ih s[A1].parallel(paxis) s[A1].pragma(oaxis, "parallel_launch_point") s[A1].pragma(paxis, "parallel_stride_pattern") s[A1].pragma(oaxis, "parallel_barrier_when_finish") ##### Schedule B co, _, _, _, _ = s[B0].op.axis if sch.bc == 1: oaxis = co paxis = co else: oco, ico = s[B0].split(co, sch.bc) oaxis = oco paxis = ico s[B0].parallel(paxis) s[B0].pragma(oaxis, "parallel_launch_point") s[B0].pragma(paxis, "parallel_stride_pattern") s[B0].pragma(oaxis, "parallel_barrier_when_finish") ##### Schedule C n, co, h, w = s[C].op.axis co, vc = s[C].split(co, VC) oh, ow, vh, vw = s[C].tile(h, w, VH, VW) s[C].reorder(n, co, oh, ow, vh, vw, vc) if C != C1: s[C1].compute_inline() s[C0].compute_at(s[C], ow) if sch.bc == 1: oaxis = co paxis = co else: oco, ico = s[C].split(co, sch.bc) oaxis = oco paxis = ico s[C].parallel(paxis) s[C].pragma(oaxis, "parallel_launch_point") s[C].pragma(paxis, "parallel_stride_pattern") s[C].pragma(oaxis, "parallel_barrier_when_finish") return s