def get_ref_data(): out_grad_np = np.random.uniform(size=out_grad_shape).astype(dtype) filter_np = np.random.uniform(size=filter_shape).astype(dtype) dilated_out_grad_np = topi.testing.dilate_python( out_grad_np, [1, stride_h, stride_w, 1]) # padding params in forward propagation fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple( [padding_h, padding_w], (filter_h, filter_w)) # padding params in backward propagation bpad_top = filter_h - 1 - fpad_top bpad_bottom = (filter_h - 1 - fpad_bottom) + (stride_h - 1) bpad_left = filter_w - 1 - fpad_left bpad_right = (filter_w - 1 - fpad_right) + (stride_w - 1) padded_out_grad = np.zeros( (batch, dilated_out_grad_np.shape[1] + bpad_top + bpad_bottom, dilated_out_grad_np.shape[2] + bpad_left + bpad_right, out_channel)) padded_out_grad[:, bpad_top:dilated_out_grad_np.shape[1] + bpad_top, bpad_left:dilated_out_grad_np.shape[2] + bpad_left, :] = dilated_out_grad_np in_grad_np = np.zeros((batch, in_h, in_w, in_channel)) for b in range(batch): for c in range(in_channel): for m in range(channel_multiplier): in_grad_np[b, :, :, c] += signal.convolve2d(padded_out_grad[b, :, :, c*channel_multiplier+m], \ filter_np[:, :, c, m], mode='valid')[0:in_h, 0:in_w] return (out_grad_np, filter_np, in_grad_np)
def np_conv(na, nw, padding, stride=1): batch, in_channel, in_height, in_width = na.shape _, num_filter, kernel_h, kernel_w = nw.shape if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple( padding, (kernel_h, kernel_w)) pad_h = pad_top + pad_bottom pad_w = pad_left + pad_right out_channel = num_filter out_height = (in_height - kernel_h + pad_h) // stride_h + 1 out_width = (in_width - kernel_w + pad_w) // stride_w + 1 nb = np.zeros((batch, out_channel, out_height, out_width)) for n in range(batch): for f in range(out_channel): for c in range(in_channel): if pad_h > 0 or pad_w > 0: apad = np.zeros((in_height + pad_h, in_width + pad_w)) apad[pad_top:pad_top + in_height, pad_left:pad_left + in_width] = na[n, c] else: apad = na[n, c] out = scipy.signal.convolve2d(apad, np.rot90(np.rot90(nw[f, c])), mode='valid') nb[n, f] += out[::stride, ::stride] return nb
def conv2d_grad(orig, grad): """Gradient of conv2d""" attrs = orig.attrs data, weight = orig.args data_shape = get_const_tuple(data.checked_type.shape) weight_shape = get_const_tuple(weight.checked_type.shape) _, _, grad_h, grad_w = get_const_tuple(orig.checked_type.shape) batch, in_channel, in_h, in_w = data_shape out_channel, _, filter_h, filter_w = weight_shape # infer output_padding fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple( get_const_tuple(attrs.padding), (filter_h, filter_w)) stride_h, stride_w = get_const_tuple(attrs.strides) dilation_h, dilation_w = get_const_tuple(attrs.dilation) out_h = (grad_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h out_w = (grad_w - 1) * stride_w - fpad_left - fpad_right + filter_w output_padding = (in_h - out_h, in_w - out_w) assert attrs.data_layout == 'NCHW', 'only support NCHW data layout' assert attrs.kernel_layout == 'OIHW', 'only support OIHW kernel layout' assert attrs.out_layout in ['', 'NCHW'], 'only support NCHW output layout' backward_data = _nn.conv2d_transpose(grad, weight, strides=attrs.strides, padding=attrs.padding, dilation=attrs.dilation, groups=attrs.groups, output_padding=output_padding) grad = tile(grad, [1, in_channel // attrs.groups, 1, 1]) grad = reshape(grad, [-1, 1, 0, 0]) # batch * oc * ic // groups, 1, oh, ow data = reshape(data, [1, -1, 0, 0]) # 1, batch * ic, ih, iw backward_weight = _nn.conv2d(data, grad, strides=attrs.dilation, padding=attrs.padding, dilation=attrs.strides, groups=in_channel * batch) # infer shape of backward_weight padded_weight_grad_h = (in_h - (grad_h - 1) * stride_h - 1 + fpad_top + fpad_bottom) \ // dilation_h + 1 padded_weight_grad_w = (in_w - (grad_w - 1) * stride_w - 1 + fpad_left + fpad_right) \ // dilation_w + 1 backward_weight = reshape(backward_weight, [ batch, in_channel // attrs.groups, out_channel, padded_weight_grad_h, padded_weight_grad_w ]) backward_weight = _sum(backward_weight, axis=0) backward_weight = transpose(backward_weight, [1, 0, 2, 3]) assert padded_weight_grad_h >= filter_h assert padded_weight_grad_w >= filter_w if padded_weight_grad_h > filter_h or padded_weight_grad_w > filter_w: backward_weight = strided_slice(backward_weight, begin=[0, 0, 0, 0], end=[None, None, filter_h, filter_w]) return [backward_data, backward_weight]
def compile_conv2d_NHWC_gemm_int8_arm(batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1, add_bias=False, add_relu=False): pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel, kernel)) padding_sum = pad_top + pad_left + pad_bottom + pad_right print("Workload: (%d, %d, %d, %d, %d, %d, %d, %d)" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) in_height = in_width = in_size A = te.placeholder((batch, in_height, in_width, in_channel), name='A', dtype='int8') W = te.placeholder((kernel, kernel, in_channel, num_filter), name='W', dtype='int8') bias = te.placeholder((num_filter,), name='bias', dtype='int8') dtype = 'int32' device = "llvm --device arm_cpu --mtriple aarch64-linux-gnu" ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Compiling on arm AArch64 target: %s" % device) with tvm.target.create(device): assert is_aarch64_arm(), "AArch64 target not recognized" C = topi.arm_cpu.compute_conv2d_NHWC_quantized(A, W, (stride, stride), padding, (dilation, dilation), dtype) if add_bias: C = topi.add(C, bias) if add_relu: C = topi.nn.relu(C) s = topi.arm_cpu.schedule_conv2d_NHWC_quantized([C]) if add_bias: tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func = tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) else: func = tvm.build(s, [A, W, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation))
def _conv2d_nhwc_python(a_np, w_np, stride, padding): """Convolution operator in NHWC layout. Parameters ---------- a_np : numpy.ndarray 4-D with shape [batch, in_height, in_width, in_channel] w_np : numpy.ndarray 4-D with shape [filter_height, filter_width, in_channel, num_filter] stride : int or a list/tuple of two ints Stride size, or [stride_height, stride_width] padding : int or str or a list/tuple of 2 or 4 ints Padding size, or ['VALID', 'SAME'], or [pad_height, pad_width] for 2 ints, or [pad_top, pad_left, pad_bottom, pad_right] for 2 ints Returns ------- b_np : np.ndarray 4-D with shape [batch, out_height, out_width, out_channel] """ batch, in_height, in_width, in_channel = a_np.shape kernel_h, kernel_w, _, num_filter = w_np.shape if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple( padding, (kernel_h, kernel_w)) pad_h = pad_top + pad_bottom pad_w = pad_left + pad_right # compute the output shape out_channel = num_filter out_height = (in_height - kernel_h + pad_h) // stride_h + 1 out_width = (in_width - kernel_w + pad_w) // stride_w + 1 # change the layout from NHWC to NCHW at = a_np.transpose((0, 3, 1, 2)) wt = w_np.transpose((3, 2, 0, 1)) bt = np.zeros((batch, out_channel, out_height, out_width)) # computation for n in range(batch): for f in range(out_channel): for c in range(in_channel): if pad_h > 0 or pad_w > 0: apad = np.zeros((in_height + pad_h, in_width + pad_w)) apad[pad_top:pad_top + in_height, pad_left:pad_left + in_width] = at[n, c] else: apad = at[n, c] out = scipy.signal.convolve2d(apad, np.rot90(np.rot90(wt[f, c])), mode='valid') bt[n, f] += out[::stride_h, ::stride_w] return bt.transpose((0, 2, 3, 1))
def conv2d_transpose_nchw_python(a_np, w_np, stride, padding, output_padding=(0, 0)): """Transposed convolution operator in NCHW layout. Parameters ---------- a_np : numpy.ndarray 4-D with shape [batch, in_channel, in_height, in_width] w_np : numpy.ndarray 4-D with shape [in_channel, num_filter, filter_height, filter_width] stride : int or a list/tuple of two ints Stride size, or [stride_height, stride_width] padding : int or str Padding size, or ['VALID', 'SAME'] Returns ------- b_np : np.ndarray 4-D with shape [batch, out_channel, out_height, out_width] """ batch, in_c, in_h, in_w = a_np.shape _, out_c, filter_h, filter_w = w_np.shape if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride opad_h, opad_w = output_padding # dilate stage dilated_a_np = topi.testing.dilate_python(a_np, [1, 1, stride_h, stride_w]) # padding stage fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple( padding, (filter_h, filter_w)) bpad_top = filter_h - 1 - fpad_top bpad_bottom = filter_h - 1 - fpad_bottom + opad_h bpad_left = filter_w - 1 - fpad_left bpad_right = filter_w - 1 - fpad_right + opad_w padded_a_np = np.zeros((batch, in_c, dilated_a_np.shape[2]+bpad_top+bpad_bottom, \ dilated_a_np.shape[3]+bpad_left+bpad_right)) padded_a_np[:, :, bpad_top:dilated_a_np.shape[2]+bpad_top, \ bpad_left:dilated_a_np.shape[3]+bpad_left] = dilated_a_np # convolution stage out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h + opad_h out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w + opad_w b_np = np.zeros((batch, out_c, out_h, out_w)) for n in range(batch): for f in range(out_c): for c in range(in_c): out = scipy.signal.convolve2d(padded_a_np[n, c], w_np[c, f], mode='valid') b_np[n, f] += out return b_np
def conv2d_transpose_packed(cfg, data, kernel, strides, padding, out_dtype, output_padding=(0, 0)): """Packed conv2d_transpose compute""" ishape = get_const_tuple(data.shape) kshape = get_const_tuple(kernel.shape) b, c_i, i_h, i_w, t_b, t_ci = ishape c_o, _, k_h, k_w, t_co, t_ci = kshape stride_h, stride_w = strides opad_h, opad_w = output_padding # FIXME(tmoreau89): currently IR pass breaks when output padding != (0,0) assert opad_h == 0 and opad_w == 0, "VTA does not support output padding for now" # derive padding parameters fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple( padding, (k_h, k_w)) bpad_top = k_h - 1 - fpad_top bpad_bottom = k_h - 1 - fpad_bottom + opad_h bpad_left = k_w - 1 - fpad_left bpad_right = k_w - 1 - fpad_right + opad_w # padding stage dilated_input = topi.nn.dilate(data, [1, 1, stride_h, stride_w, 1, 1]) data_pad = topi.nn.pad(dilated_input, [0, 0, bpad_top, bpad_left, 0, 0], [0, 0, bpad_bottom, bpad_right, 0, 0]) # convolution transpose stage out_h = (i_h - 1) * stride_h - fpad_top - fpad_bottom + k_h + opad_h out_w = (i_w - 1) * stride_w - fpad_left - fpad_right + k_w + opad_w oshape = (b, c_o, out_h, out_w, t_b, t_co) d_c = te.reduce_axis((0, c_i), name='d_c') d_h = te.reduce_axis((0, k_h), name='d_h') d_w = te.reduce_axis((0, k_w), name='d_w') d_ci = te.reduce_axis((0, t_ci), name='d_ci') out = te.compute(oshape, lambda i_n, i_c, i_h, i_w, j_n, j_c: te. sum(data_pad(i_n, d_c, i_h + d_h, i_w + d_w, j_n, d_ci). astype(out_dtype) * kernel[i_c, d_c, d_h, d_w, j_c, d_ci].astype(out_dtype), axis=[d_c, d_h, d_w, d_ci]), tag="packed_conv2d_transpose", name='res') cfg.add_flop(2 * np.prod(topi.util.get_const_tuple(oshape)) * kshape[2] * kshape[3] * ishape[1] * ishape[-1]) return out
def conv2d_transpose_nchw_python(a_np, w_np, stride, padding): """Transposed convolution operator in NCHW layout. Parameters ---------- a_np : numpy.ndarray 4-D with shape [batch, in_channel, in_height, in_width] w_np : numpy.ndarray 4-D with shape [in_channel, num_filter, filter_height, filter_width] stride : int or a list/tuple of two ints Stride size, or [stride_height, stride_width] padding : int or str Padding size, or ['VALID', 'SAME'] Returns ------- b_np : np.ndarray 4-D with shape [batch, out_channel, out_height, out_width] """ batch, in_c, in_h, in_w = a_np.shape _, out_c, filter_h, filter_w = w_np.shape if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride # dilate stage dilated_a_np = topi.testing.dilate_python(a_np, [1, 1, stride_h, stride_w]) # padding stage fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w)) bpad_top = filter_h - 1 - fpad_top bpad_bottom = filter_h - 1 - fpad_bottom bpad_left = filter_w - 1 - fpad_left bpad_right = filter_w - 1 - fpad_right padded_a_np = np.zeros((batch, in_c, dilated_a_np.shape[2]+bpad_top+bpad_bottom, \ dilated_a_np.shape[3]+bpad_left+bpad_right)) padded_a_np[:, :, bpad_top:dilated_a_np.shape[2]+bpad_top, \ bpad_left:dilated_a_np.shape[3]+bpad_left] = dilated_a_np # convolution stage out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w b_np = np.zeros((batch, out_c, out_h, out_w)) for n in range(batch): for f in range(out_c): for c in range(in_c): out = scipy.signal.convolve2d( padded_a_np[n, c], w_np[c, f], mode='valid') b_np[n, f] += out return b_np
def get_ref_data(): out_grad_np = np.random.uniform(size=out_grad_shape).astype(dtype) input_np = np.random.uniform(size=in_shape).astype(dtype) dilated_out_grad_np = topi.testing.dilate_python(out_grad_np, [1, stride_h, stride_w, 1]) pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple([padding_h, padding_w], (filter_h, filter_w)) padded_input_np = np.zeros((batch, in_h+pad_top+pad_bottom, in_w+pad_left+pad_right, in_channel)) padded_input_np[:, pad_top:in_h+pad_top, pad_left:in_w+pad_left, :] = input_np weight_grad_np = np.zeros((filter_h, filter_w, in_channel, channel_multiplier)) for c in range(in_channel): for m in range(channel_multiplier): for b in range(batch): weight_grad_np[:, :, c, m] += signal.convolve2d(padded_input_np[b, :, :, c], \ np.rot90(dilated_out_grad_np[b, :, :, c*channel_multiplier+m%channel_multiplier], 2), \ mode='valid')[0:filter_h, 0:filter_w] return (out_grad_np, input_np, weight_grad_np)
def get_ref_data(): out_grad_np = np.random.uniform(size=out_grad_shape).astype(dtype) input_np = np.random.uniform(size=in_shape).astype(dtype) dilated_out_grad_np = topi.testing.dilate_python(out_grad_np, [1, stride_h, stride_w, 1]) pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple([padding_h, padding_w], (filter_h, filter_w)) padded_input_np = np.zeros((batch, in_h+pad_top+pad_bottom, in_w+pad_left+pad_right, in_channel)) padded_input_np[:, pad_top:in_h+pad_top, pad_left:in_w+pad_left, :] = input_np weight_grad_np = np.zeros((filter_h, filter_w, in_channel, channel_multiplier)) for c in range(in_channel): for m in range(channel_multiplier): for b in range(batch): weight_grad_np[:, :, c, m] += signal.convolve2d(padded_input_np[b, :, :, c], \ np.rot90(dilated_out_grad_np[b, :, :, c*channel_multiplier+m%channel_multiplier], 2), \ mode='valid')[0:filter_h, 0:filter_w] return (out_grad_np, input_np, weight_grad_np)
def conv2d_transpose_nchw_python(a_np, w_np, stride, padding): """Transposed convolution operator in NCHW layout. Parameters ---------- a_np : numpy.ndarray 4-D with shape [batch, in_channel, in_height, in_width] w_np : numpy.ndarray 4-D with shape [num_filter, in_channel, filter_height, filter_width] stride : int or a list/tuple of two ints Stride size, or [stride_height, stride_width] padding : int or str Padding size, or ['VALID', 'SAME'] Returns ------- b_np : np.ndarray 4-D with shape [batch, out_channel, out_height, out_width] """ batch, in_c, in_h, in_w = a_np.shape out_c, _, filter_h, filter_w = w_np.shape if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride # dilate stage dilated_a_np = topi.testing.dilate_python(a_np, [1, 1, stride_h, stride_w]) # padding stage fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w)) bpad_top = filter_h - 1 - fpad_top bpad_bottom = filter_h - 1 - fpad_bottom bpad_left = filter_w - 1 - fpad_left bpad_right = filter_w - 1 - fpad_right padded_a_np = np.zeros((batch, in_c, dilated_a_np.shape[2]+bpad_top+bpad_bottom, \ dilated_a_np.shape[3]+bpad_left+bpad_right)) padded_a_np[:, :, bpad_top:dilated_a_np.shape[2]+bpad_top, \ bpad_left:dilated_a_np.shape[3]+bpad_left] = dilated_a_np # convolution stage rotated_w_np = np.rot90(w_np, k=2, axes=(2, 3)) b_np = topi.testing.conv2d_nchw_python(padded_a_np, rotated_w_np, stride=1, padding='VALID') return b_np
def _declatation_conv2d_transpose(cfg, data, kernel, strides, padding, out_dtype): ishape = get_const_tuple(data.shape) kshape = get_const_tuple(kernel.shape) b, c_i, i_h, i_w, t_b, t_ci = ishape c_o, _, k_h, k_w, t_co, t_ci = kshape stride_h, stride_w = strides # derive padding parameters fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple( padding, (k_h, k_w)) bpad_top = k_h - 1 - fpad_top bpad_bottom = k_h - 1 - fpad_bottom bpad_left = k_w - 1 - fpad_left bpad_right = k_w - 1 - fpad_right # padding stage dilated_input = topi.nn.dilate(data, [1, 1, stride_h, stride_w, 1, 1]) data_pad = topi.nn.pad(dilated_input, [0, 0, bpad_top, bpad_left, 0, 0], [0, 0, bpad_bottom, bpad_right, 0, 0]) # convolution transpose stage out_h = (i_h - 1) * stride_h - fpad_top - fpad_bottom + k_h out_w = (i_w - 1) * stride_w - fpad_left - fpad_right + k_w oshape = (b, c_o, out_h, out_w, t_b, t_co) d_c = tvm.reduce_axis((0, c_i), name='d_c') d_h = tvm.reduce_axis((0, k_h), name='d_h') d_w = tvm.reduce_axis((0, k_w), name='d_w') d_ci = tvm.reduce_axis((0, t_ci), name='d_ci') out = tvm.compute(oshape, lambda i_n, i_c, i_h, i_w, j_n, j_c: tvm. sum(data_pad(i_n, d_c, i_h + d_h, i_w + d_w, j_n, d_ci). astype(out_dtype) * kernel[i_c, d_c, d_h, d_w, j_c, d_ci].astype(out_dtype), axis=[d_c, d_h, d_w, d_ci]), tag="packed_conv2d_transpose", name='res') cfg.add_flop(2 * np.prod(topi.util.get_const_tuple(oshape)) * kshape[2] * kshape[3] * ishape[1] * ishape[-1]) return out
def get_ref_data(): out_grad_np = np.random.uniform(size=out_grad_shape).astype(dtype) filter_np = np.random.uniform(size=filter_shape).astype(dtype) dilated_out_grad_np = topi.testing.dilate_python(out_grad_np, [1, stride_h, stride_w, 1]) # padding params in forward propagation fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple([padding_h, padding_w], (filter_h, filter_w)) # padding params in backward propagation bpad_top = filter_h - 1 - fpad_top bpad_bottom = (filter_h - 1 - fpad_bottom) + (stride_h - 1) bpad_left = filter_w - 1 - fpad_left bpad_right = (filter_w - 1 - fpad_right) + (stride_w - 1) padded_out_grad = np.zeros((batch, dilated_out_grad_np.shape[1]+bpad_top+bpad_bottom, dilated_out_grad_np.shape[2]+bpad_left+bpad_right, out_channel)) padded_out_grad[:, bpad_top:dilated_out_grad_np.shape[1]+bpad_top, bpad_left:dilated_out_grad_np.shape[2]+bpad_left, :] = dilated_out_grad_np in_grad_np = np.zeros((batch, in_h, in_w, in_channel)) for b in range(batch): for c in range(in_channel): for m in range(channel_multiplier): in_grad_np[b, :, :, c] += signal.convolve2d(padded_out_grad[b, :, :, c*channel_multiplier+m], \ filter_np[:, :, c, m], mode='valid')[0:in_h, 0:in_w] return (out_grad_np, filter_np, in_grad_np)
def verify_conv2d_NCHWc_int8(batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1, add_bias=False, add_relu=False): pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel, kernel)) padding_sum = pad_top + pad_left + pad_bottom + pad_right print("Workload: (%d, %d, %d, %d, %d, %d, %d, %d)" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) in_height = in_width = in_size A = tvm.placeholder((batch, in_channel, in_height, in_width), name='A', dtype='int8') W = tvm.placeholder((num_filter, in_channel, kernel, kernel), name='W', dtype='int8') bias = tvm.placeholder((num_filter // oc_block_factor, 1, 1, oc_block_factor), name='bias', dtype='int8') a_shape = get_const_tuple(A.shape) w_shape = get_const_tuple(W.shape) bias_shape = get_const_tuple(bias.shape) dtype = A.dtype @memoize("topi.tests.test_topi_conv2d_int8.verify_conv2d_nchw") def get_ref_data(): a_np = np.random.randint(low=-128, high=127, size=a_shape).astype(dtype) w_np = np.random.randint(low=-128, high=128, size=w_shape).astype(dtype) b_np = np.random.uniform(size=bias_shape).astype(dtype) dw_np = topi.testing.dilate_python(w_np, (1, 1, dilation, dilation)) c_np = topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding).astype(dtype) # convert to NCHWc _, _, out_height, out_width = c_np.shape c_np = c_np.reshape((batch, num_filter // oc_block_factor, oc_block_factor, \ out_height, out_width)).transpose(0, 1, 3, 4, 2) if add_bias: b_np = np.random.uniform(size=bias_shape).astype(dtype) c_np += b_np if add_relu: c_np = np.maximum(c_np, 0) return a_np, w_np, b_np, c_np a_np, w_np, b_np, c_np = get_ref_data() def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return if device == "cuda" and not tvm.contrib.nvcc.have_int8(ctx.compute_version): print("Skip because int8 intrinsics are not available") return print("Running on target: %s" % device) with tvm.target.create(device): C = topi.nn.conv2d(A, W, (stride, stride), padding, (dilation, dilation), layout='NCHW', out_dtype=dtype) if add_bias: C = topi.add(C, bias) if add_relu: C = topi.nn.relu(C) s = topi.generic.schedule_conv2d_nchw([C]) a = tvm.nd.array(a_np, ctx) w = tvm.nd.array(w_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), ctx) if add_bias: tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func = tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, b, c) else: func = tvm.build(s, [A, W, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, c) tvm.testing.assert_allclose(c.asnumpy(), c_np, rtol=1e-5) for device in ["cuda"]: check_device(device)
def verify_conv2d_nchw(batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1, add_bias=False, add_relu=False,\ use_cudnn=False): pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel, kernel)) padding_sum = pad_top + pad_left + pad_bottom + pad_right print("Workload: (%d, %d, %d, %d, %d, %d, %d, %d)" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) in_height = in_width = in_size A = tvm.placeholder((batch, in_channel, in_height, in_width), name='A') W = tvm.placeholder((num_filter, in_channel, kernel, kernel), name='W') bias = tvm.placeholder((num_filter, 1, 1), name='bias') a_shape = get_const_tuple(A.shape) w_shape = get_const_tuple(W.shape) bias_shape = get_const_tuple(bias.shape) dtype = A.dtype @memoize("topi.tests.test_topi_conv2d_nchw.verify_conv2d_nchw") def get_ref_data(): a_np = np.random.uniform(size=a_shape).astype(dtype) w_np = np.random.uniform(size=w_shape).astype(dtype) b_np = np.random.uniform(size=bias_shape).astype(dtype) dw_np = topi.testing.dilate_python(w_np, (1, 1, dilation, dilation)) c_np = topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding) if add_bias: c_np += b_np if add_relu: c_np = np.maximum(c_np, 0) return a_np, w_np, b_np, c_np a_np, w_np, b_np, c_np = get_ref_data() def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): C = topi.nn.conv2d(A, W, (stride, stride), padding, (dilation, dilation), layout='NCHW', out_dtype=dtype) if add_bias: C = topi.add(C, bias) if add_relu: C = topi.nn.relu(C) s = topi.generic.schedule_conv2d_nchw([C]) a = tvm.nd.array(a_np, ctx) w = tvm.nd.array(w_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), ctx) if add_bias: func = tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, b, c) else: func = tvm.build(s, [A, W, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, c) tvm.testing.assert_allclose(c.asnumpy(), c_np, rtol=1e-4) for device in get_all_backend(): with autotvm.tophub.context(device): # load tophub pre-tuned parameters check_device(device) if use_cudnn: check_device("cuda -model=unknown -libs=cudnn")
def _gen_cfg(cfg, data, kernel, strides, padding, dilation, num_tile): """_gen_cfg""" if len(kernel.shape) == 4: co_, _, kh_, kw_ = get_const_tuple(kernel.shape) else: # kernel tensor is pre packed co_, _, kh_, kw_, vc_ = get_const_tuple(kernel.shape) co_ = co_ * vc_ if isinstance(dilation, int): dilation_h = dilation_w = dilation else: dilation_h, dilation_w = dilation n_, ci_, ih_, iw_ = get_const_tuple(data.shape) dilated_kernel_h = (kh_ - 1) * dilation_h + 1 dilated_kernel_w = (kw_ - 1) * dilation_w + 1 pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple( padding, (dilated_kernel_h, dilated_kernel_w)) hstr, wstr = strides if isinstance(strides, (tuple, list)) else (strides, strides) oh_ = (ih_ + pad_top + pad_bottom - dilated_kernel_h) // hstr + 1 ow_ = (iw_ + pad_left + pad_right - dilated_kernel_w) // wstr + 1 n, co, oh, ow = cfg.axis(n_), cfg.axis(co_), cfg.axis(oh_), cfg.axis(ow_) ci, kh, kw = cfg.reduce_axis(ci_), cfg.reduce_axis(kh_), cfg.reduce_axis( kw_) if num_tile == 2: # for arm cpu candidate_vc = [] for iv in range(3, co_): if co_ % iv == 0: candidate_vc.append([co_ // iv, iv]) candidate_vc.append([1, co_]) co, vc = cfg.define_split("tile_co", co, num_outputs=2, policy="candidate", candidate=candidate_vc) oh, vh = cfg.define_split("tile_oh", oh, num_outputs=2) ow, vw = cfg.define_split("tile_ow", ow, num_outputs=2) elif num_tile == 3: # for mali gpu co, _, vc = cfg.define_split("tile_co", co, num_outputs=3) oh, _, vh = cfg.define_split("tile_oh", oh, num_outputs=3) ow, _, vw = cfg.define_split("tile_ow", ow, num_outputs=3) else: raise RuntimeError("Invalid num_tile") cfg.define_reorder( "reorder_0", [n, co, oh, ow, ci, kh, kw, vh, vw, vc], policy="candidate", candidate=[ [n, co, oh, ow, ci, kh, kw, vh, vw, vc], ], ) vc_ = cfg["tile_co"].size[-1] vh_ = cfg["tile_oh"].size[-1] vw_ = cfg["tile_ow"].size[-1] is_var = False return (is_var, vh_, vw_, vc_)
def depthwise_conv2d_with_workload_NCHWc(batch, in_channel, in_height, channel_multiplier, filter_height, stride, padding, dilation=1): in_width = in_height filter_channel = in_channel filter_width = filter_height stride_h = stride_w = stride assert dilation == 1, "depthwise_conv2d_NCHWc currently does not support dilation." pad_h, pad_w, _, _ = get_pad_tuple(padding, (filter_height, filter_width)) padding_args = (pad_h, pad_w) out_channel = filter_channel * channel_multiplier # for testing functionality, # we choose arbitrary block size that can divide the channel, # regardless of the performance. oc_block = 1 for bn in range(16, 0, -1): if out_channel % bn == 0: oc_block = bn break ic_block = 1 for bn in range(oc_block, 0, -1): if in_channel % bn == 0: ic_block = bn break # placeholder Input = tvm.placeholder((batch, in_channel//ic_block, in_height, in_width, ic_block), name='Input') Filter = tvm.placeholder((out_channel//oc_block, filter_height, filter_width, oc_block), name='Filter') in_layout = "NCHW%dc" % ic_block out_layout = "NCHW%dc" % oc_block dtype = 'float32' def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): # declare DepthwiseConv2d = topi.nn.depthwise_conv2d_NCHWc(Input, Filter, (stride_h, stride_w), padding_args, (dilation, dilation), in_layout, out_layout, dtype) # TODO: add scale_shift implement for NCHWc and add test here Relu = topi.nn.relu(DepthwiseConv2d) # schedule s1 = topi.generic.schedule_depthwise_conv2d_nchw(DepthwiseConv2d) s2 = topi.generic.schedule_depthwise_conv2d_nchw(Relu) # build the kernels f1 = tvm.build(s1, [Input, Filter, DepthwiseConv2d], device) f2 = tvm.build(s2, [Input, Filter, Relu], device) # Prepare pod type for test data closure input_shape = (batch, in_channel, in_height, in_width) filter_shape = (filter_channel, channel_multiplier, filter_height, filter_width) # Use memoize, pickle the test data for next time use. @memoize("topi.tests.test_topi_depthwise_conv2d.NCHWc") def get_ref_data(): input_np = np.random.uniform(size=input_shape).astype(dtype) filter_np = np.random.uniform(size=filter_shape).astype(dtype) # correctness with scipy depthwise_conv2d_scipy = topi.testing.depthwise_conv2d_python_nchw( input_np, filter_np, stride, padding) relu_scipy = np.maximum(depthwise_conv2d_scipy, 0) return (_transform_data(input_np, ic_block), _transform_kernel(filter_np, oc_block), _transform_data(depthwise_conv2d_scipy, oc_block), _transform_data(relu_scipy, oc_block)) # Get the test data (input_np, filter_np, depthwise_conv2d_scipy, relu_scipy) = get_ref_data() input_tvm = tvm.nd.array(input_np, ctx) filter_tvm = tvm.nd.array(filter_np, ctx) depthwise_conv2d_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(DepthwiseConv2d.shape), dtype=DepthwiseConv2d.dtype), ctx) relu_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(Relu.shape), dtype=Relu.dtype), ctx) # launch kernel 1 (depthwise_conv2d) f1(input_tvm, filter_tvm, depthwise_conv2d_tvm) # launch kernel 2 (depthwise_conv2d + relu) f2(input_tvm, filter_tvm, relu_tvm) tvm.testing.assert_allclose(depthwise_conv2d_tvm.asnumpy(), depthwise_conv2d_scipy, rtol=1e-5) tvm.testing.assert_allclose(relu_tvm.asnumpy(), relu_scipy, rtol=1e-5) # test llvm only for now since depthwise_conv2d_NCHWc implement is missing in other backend. for device in ["llvm"]: with autotvm.tophub.context(device): # load tophub pre-tuned parameters check_device(device)
def fused_convs(input_data, filters, resnet_block=False): out_dtype = input_data.dtype Input = None nodes = [input_data] params = [input_data] for f in filters: Input = nodes[-1] Filter = f.placeholder layout = f.layout depthwise = f.depthwise kernel = f.kernel stride = f.stride padding = f.padding dilation = f.dilation assert not (depthwise and kernel == 1) # Don't consider 1by1 depthwise padded_count = 0 conv_count = 0 depthwise_count = 0 if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride if isinstance(dilation, int): dilation_h = dilation_w = dilation else: dilation_h, dilation_w = dilation batch, in_height, in_width, in_channel = Input.shape if f.NHWC_transpose: # HWOI kernel_h, kernel_w, tmp, kernel_channel = Filter.shape else: # HWIO kernel_h, kernel_w, kernel_channel, tmp = Filter.shape if depthwise: channel_multiplier = tmp else: num_filter = tmp # compute the output shape dilated_kernel_h = (kernel_h - 1) * dilation_h + 1 dilated_kernel_w = (kernel_w - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( padding, (dilated_kernel_h, dilated_kernel_w)) out_channel = simplify(in_channel * channel_multiplier) if depthwise else num_filter out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1) out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1) if f.kernel > 1: print("Padding is needed!") pad_before = [0, pad_top, pad_left, 0] pad_after = [0, pad_down, pad_right, 0] PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput_{}".format(padded_count)) padded_count += 1 nodes.append(PaddedInput) # Update Input Input = PaddedInput batch, in_height, in_width, in_channel = Input.shape if not depthwise: rc = tvm.reduce_axis((0, in_channel), name='rc') if kernel > 1: ry = tvm.reduce_axis((0, kernel_h), name='ry') rx = tvm.reduce_axis((0, kernel_w), name='rx') if not depthwise: # Normal convolution if kernel > 1: Output = tvm.compute( (batch, out_height, out_width, out_channel), lambda nn, yy, xx, ff: tvm.sum( Input[nn, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rc].astype(out_dtype) * (Filter[ry, rx, ff, rc] if f.NHWC_transpose else Filter[ry, rx, rc, ff]).astype(out_dtype), axis=[ry, rx, rc]), name="Conv2dOutput_{}".format(conv_count), tag="conv2d_nhwc") else: # Only reduce rc axis Output = tvm.compute( (batch, out_height, out_width, out_channel), lambda nn, yy, xx, ff: tvm.sum( Input[nn, yy * stride_h, xx * stride_w, rc].astype(out_dtype) * (Filter[0, 0, ff, rc] if f.NHWC_transpose else Filter[0, 0, rc, ff]).astype(out_dtype), axis=[rc]), name="Conv2dOutput_{}".format(conv_count), tag="conv2d_nhwc") conv_count += 1 else: # Depthwise convolution (kernel > 1) Output = tvm.compute( (batch, out_height, out_width, out_channel), lambda b, i, j, c: tvm.sum( (Input[b, i*stride_h + ry*dilation_h, j*stride_w + rx*dilation_w, tvm.indexdiv(c, channel_multiplier)].astype(out_dtype) * (Filter[ry, rx, tvm.indexmod(c, channel_multiplier), tvm.indexdiv(c, channel_multiplier)] if f.NHWC_transpose else Filter[ry, rx, tvm.indexdiv(c, channel_multiplier), tvm.indexmod(c, channel_multiplier)]).astype(out_dtype)), axis=[ry, rx]), name='DepthwiseConv2dOutput_{}'.format(depthwise_count), tag="depthwise_nhwc") depthwise_count += 1 nodes.append(Output) params.append(Filter) if resnet_block: First = nodes[0] Last = nodes[-1] assert (first.shape == last.shape) Output = tvm.compute( (batch, out_height, out_width, out_channel), lambda b, i, j, c: tvm.sum( (First[b, i, j, c].astype(out_dtype) + (Last[b, i, j, c]).astype(out_dtype))), name='ElementwiseAddOutput_{}'.format(depthwise_count), tag="elem_nhwc") nodes.append(Output) params.append(nodes[-1]) # Final output return nodes, params
def depthwise_conv2d_with_workload_nchw(batch, in_channel, in_height, channel_multiplier, filter_height, stride, padding, dilation=1): in_width = in_height filter_channel = in_channel filter_width = filter_height stride_h = stride_w = stride if dilation == 1: # here we transform the padding argument from 'str' to 'tuple' , # because we need this to match the "workload" tuple to the records in TopHub pad_h, pad_w, _, _ = get_pad_tuple(padding, (filter_height, filter_width)) padding_args = (pad_h, pad_w) else: padding_args = padding # placeholder Input = tvm.placeholder((batch, in_channel, in_height, in_width), name='Input') Filter = tvm.placeholder((filter_channel, channel_multiplier, filter_height, filter_width), name='Filter') Scale = tvm.placeholder((in_channel * channel_multiplier,), name='Scale') Shift = tvm.placeholder((in_channel * channel_multiplier,), name='Shift') dtype = 'float32' def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): # declare DepthwiseConv2d = topi.nn.depthwise_conv2d_nchw(Input, Filter, (stride_h, stride_w), padding_args, dilation, dtype) ScaleShift = topi.nn.scale_shift_nchw(DepthwiseConv2d, Scale, Shift) Relu = topi.nn.relu(ScaleShift) # schedule s1 = topi.generic.schedule_depthwise_conv2d_nchw(DepthwiseConv2d) s2 = topi.generic.schedule_depthwise_conv2d_nchw(ScaleShift) s3 = topi.generic.schedule_depthwise_conv2d_nchw(Relu) # build the kernels f1 = tvm.build(s1, [Input, Filter, DepthwiseConv2d], device) f2 = tvm.build(s2, [Input, Filter, Scale, Shift, ScaleShift], device) f3 = tvm.build(s3, [Input, Filter, Scale, Shift, Relu], device) # Prepare pod type for test data closure input_shape = get_const_tuple(Input.shape) filter_shape = get_const_tuple(Filter.shape) scale_shape = get_const_tuple(Scale.shape) shift_shape = get_const_tuple(Shift.shape) scale_shift_shape = get_const_tuple(ScaleShift.shape) # Use memoize, pickle the test data for next time use. @memoize("topi.tests.test_topi_depthwise_conv2d.nchw") def get_ref_data(): input_np = np.random.uniform(size=input_shape).astype(dtype) filter_np = np.random.uniform(size=filter_shape).astype(dtype) dilated_filter_np = topi.testing.dilate_python(filter_np, (1, 1, dilation, dilation)) scale_np = np.random.uniform(size=scale_shape).astype(dtype) shift_np = np.random.uniform(size=shift_shape).astype(dtype) # correctness with scipy depthwise_conv2d_scipy = topi.testing.depthwise_conv2d_python_nchw( input_np, dilated_filter_np, stride, padding) scale_shift_scipy = np.zeros(shape=scale_shift_shape) for c in range(in_channel * channel_multiplier): scale_shift_scipy[:,c,:,:] = depthwise_conv2d_scipy[:,c,:,:] * scale_np[c] + shift_np[c] relu_scipy = np.maximum(scale_shift_scipy, 0) return (input_np, filter_np, scale_np, shift_np, depthwise_conv2d_scipy, scale_shift_scipy, relu_scipy) # Get the test data (input_np, filter_np, scale_np, shift_np, depthwise_conv2d_scipy, scale_shift_scipy, relu_scipy) = get_ref_data() input_tvm = tvm.nd.array(input_np, ctx) filter_tvm = tvm.nd.array(filter_np, ctx) scale_tvm = tvm.nd.array(scale_np, ctx) shift_tvm = tvm.nd.array(shift_np, ctx) depthwise_conv2d_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(DepthwiseConv2d.shape), dtype=DepthwiseConv2d.dtype), ctx) scale_shift_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(ScaleShift.shape), dtype=ScaleShift.dtype), ctx) relu_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(Relu.shape), dtype=Relu.dtype), ctx) # launch kernel 1 (depthwise_conv2d) timer_1 = f1.time_evaluator(f1.entry_name, ctx, number=1) tcost_1 = timer_1(input_tvm, filter_tvm, depthwise_conv2d_tvm).mean # launch kernel 2 (depthwise_conv2d + scale_shift) timer_2 = f2.time_evaluator(f2.entry_name, ctx, number=1) tcost_2 = timer_2(input_tvm, filter_tvm, scale_tvm, shift_tvm, scale_shift_tvm).mean # launch kernel 3 (depthwise_conv2d + scale_shift + relu) timer_3 = f3.time_evaluator(f3.entry_name, ctx, number=1) tcost_3 = timer_3(input_tvm, filter_tvm, scale_tvm, shift_tvm, relu_tvm).mean tvm.testing.assert_allclose(depthwise_conv2d_tvm.asnumpy(), depthwise_conv2d_scipy, rtol=1e-5) tvm.testing.assert_allclose(scale_shift_tvm.asnumpy(), scale_shift_scipy, rtol=1e-5) tvm.testing.assert_allclose(relu_tvm.asnumpy(), relu_scipy, rtol=1e-5) for device in get_all_backend(): with autotvm.tophub.context(device): # load tophub pre-tuned parameters check_device(device)
def deformable_conv2d_nchw_python(a_np, offset_np, w_np, stride, padding, dilation, deformable_groups, groups): """Deformable convolution operator in NCHW layout. Parameters ---------- a_np : numpy.ndarray 4-D with shape [batch, in_channel, in_height, in_width] offset_np : numpy.ndarray 4-D with shape [batch, deformable_groups * filter_height * filter_width * 2, out_height, out_width] w_np : numpy.ndarray 4-D with shape [num_filter, in_channel, filter_height, filter_width] stride : int or a list/tuple of two ints Stride size, or [stride_height, stride_width] padding : int or str or a list/tuple of 2 or 4 ints Padding size, or ['VALID', 'SAME'], or [pad_height, pad_width] for 2 ints, or [pad_top, pad_left, pad_bottom, pad_right] for 2 ints dilation : int or a list/tuple of two ints Dilation size, or [dilate_height, dilate_width] deformable_groups : int Number of deformable groups groups : int Number of groups Returns ------- b_np : np.ndarray 4-D with shape [batch, out_channel, out_height, out_width] """ batch, in_channel, in_height, in_width = a_np.shape out_channel, _, kernel_h, kernel_w = w_np.shape out_height, out_width = offset_np.shape[-2:] dtype = a_np.dtype ic_per_dgroup = in_channel // deformable_groups assert groups == 1, "deformable_conv2d_nchw_python does not support groups > 1" if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride pad_top, pad_left, _, _ = get_pad_tuple(padding, (kernel_h, kernel_w)) if isinstance(dilation, int): dilation_h = dilation_w = dilation else: dilation_h, dilation_w = dilation def _bilinear(n, c, h, w): low_h, low_w = int(h), int(w) high_h = min(low_h + 1, in_height - 1) high_w = min(low_w + 1, in_width - 1) y_lerp = h - low_h x_lerp = w - low_w bottom = (1 - x_lerp) * a_np[n, c, low_h, low_w] + x_lerp * a_np[n, c, low_h, high_w] top = (1 - x_lerp) * a_np[n, c, high_h, low_w] + x_lerp * a_np[n, c, high_h, high_w] return (1 - y_lerp) * bottom + y_lerp * top a_deform = np.zeros((batch, in_channel, out_height, out_width, kernel_h, kernel_w), dtype=dtype) for n, h, w in itertools.product(range(batch), range(out_height), range(out_width)): offset = offset_np[n, :, h, w].reshape(deformable_groups, kernel_h, kernel_w, 2) in_h = h * stride_h - pad_top in_w = w * stride_w - pad_left index_h_base, index_w_base = np.meshgrid( np.arange(in_h, in_h + kernel_h * dilation_h, dilation_h, dtype=offset_np.dtype), np.arange(in_w, in_w + kernel_w * dilation_w, dilation_w, dtype=offset_np.dtype), indexing='ij') for c, kh, kw in itertools.product(range(in_channel), range(kernel_h), range(kernel_w)): dg = c // ic_per_dgroup index_h = index_h_base + offset[dg, ..., 0] index_w = index_w_base + offset[dg, ..., 1] y, x = index_h[kh, kw], index_w[kh, kw] if y < 0 or y >= in_height or x < 0 or x >= in_width: continue a_deform[n, c, h, w, kh, kw] = _bilinear(n, c, y, x) b_np = np.zeros((batch, out_channel, out_height, out_width), dtype=dtype) for n, c, f, h, w in itertools.product(range(batch), range(in_channel), range(out_channel), range(out_height), range(out_width)): b_np[n, f, h, w] += np.tensordot(a_deform[n, c, h, w], w_np[f, c]) return b_np
def _depthwise_spatial_pack(args, data, kernel, strides, padding, dilation, out_dtype): """depthwise_conv2d_arm_cpu's inner implement""" is_var, u_vh, u_vw, u_vc = args out_dtype = out_dtype or data.dtype u_n, u_c, ih, iw = data.shape if is_var else get_const_tuple(data.shape) if isinstance(dilation, int): dilation_h = dilation_w = dilation else: dilation_h, dilation_w = dilation if len(kernel.shape) == 4: pre_packed = False u_c, um, ukh, ukw = kernel.shape if is_var else get_const_tuple( kernel.shape) else: # kernel tensor is pre packed pre_packed = True u_c, um, ukh, ukw, u_vc = kernel.shape if is_var else get_const_tuple( kernel.shape) u_c = u_c * u_vc dilated_kernel_h = (ukh - 1) * dilation_h + 1 dilated_kernel_w = (ukw - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( padding, (dilated_kernel_h, dilated_kernel_w)) hstr, wstr = strides if isinstance(strides, (tuple, list)) else (strides, strides) u_oh = (ih + pad_top + pad_down - dilated_kernel_h) // hstr + 1 u_ow = (iw + pad_left + pad_right - dilated_kernel_w) // wstr + 1 # pack data hpad = pad_top + pad_down wpad = pad_left + pad_right dopad = hpad != 0 or wpad != 0 if dopad: data_pad = pad( data, (0, 0, pad_top, pad_left), (0, 0, pad_down, pad_right), name="data_pad", ) else: data_pad = data oh_div = u_oh // u_vh ow_div = u_ow // u_vw kvshape = (u_c // u_vc, um, ukh, ukw, u_vc) ovshape = (u_n, u_c * um // u_vc, oh_div, u_ow // u_vw, u_vh, u_vw, u_vc) oshape = (u_n, u_c * um, oh_div * u_vh, ow_div * u_vw) if dilation_h != 1 or dilation_w != 1: # undilate input data dvshape = (u_n, oh_div, ow_div, u_c, ukh, ukw, u_vh, u_vw) data_vec = tvm.compute( dvshape, lambda n, h, w, c, kh, kw, vh, vw: data_pad[n][c][ (h * u_vh + vh) * hstr + kh * dilation_h][ (w * u_vw + vw) * wstr + kw * dilation_w], name="data_vec_undilated", ) else: dvshape = (u_n, oh_div, ow_div, u_c, u_vh * hstr + ukh - 1, u_vw * wstr + ukw - 1) data_vec = tvm.compute( dvshape, lambda n, h, w, c, vh, vw: data_pad[n][c][h * u_vh * hstr + vh][ w * u_vw * wstr + vw], name="data_vec", ) if pre_packed: kernel_vec = kernel else: kernel_vec = tvm.compute( kvshape, lambda co, m, kh, kw, vc: kernel[co * u_vc + vc][m][kh][kw], name="kernel_vec", ) kh = tvm.reduce_axis((0, ukh), name="kh") kw = tvm.reduce_axis((0, ukw), name="kw") if dilation_h != 1 or dilation_w != 1: conv = tvm.compute( ovshape, lambda n, co, h, w, vh, vw, vc: tvm.sum( data_vec[n, h, w, (co * u_vc + vc) // um, kh, kw, vh, vw]. astype(out_dtype) * kernel_vec[co // um, co % um, kh, kw, vc ].astype(out_dtype), axis=[kh, kw], ), name="depthwise_conv", ) else: conv = tvm.compute( ovshape, lambda n, co, h, w, vh, vw, vc: tvm.sum( data_vec[n, h, w, (co * u_vc + vc) // um, vh * hstr + kh, vw * wstr + kw].astype(out_dtype) * kernel_vec[ co // um, co % um, kh, kw, vc].astype(out_dtype), axis=[kh, kw], ), name="depthwise_conv", ) output = tvm.compute( oshape, lambda n, co, h, w: conv[n][co // u_vc][h // u_vh][w // u_vw][h % u_vh] [w % u_vw][co % u_vc], name="output_unpack", tag="spatial_depthwise_conv_nchw_output", ) return output
def verify_conv2d_NHWC_gemm_int8(batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1, add_bias=False, add_relu=False): pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple( padding, (kernel, kernel)) padding_sum = pad_top + pad_left + pad_bottom + pad_right print("Workload: (%d, %d, %d, %d, %d, %d, %d, %d)" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) in_height = in_width = in_size A = te.placeholder((batch, in_height, in_width, in_channel), name='A', dtype='int8') W = te.placeholder((kernel, kernel, in_channel, num_filter), name='W', dtype='int8') bias = te.placeholder((num_filter, ), name='bias', dtype='int8') a_shape = get_const_tuple(A.shape) w_shape = get_const_tuple(W.shape) bias_shape = get_const_tuple(bias.shape) dtype = A.dtype @memoize("topi.tests.test_topi_conv2d_int8.verify_conv2d_nchw") def get_ref_data(): a_np = np.random.randint(low=-128, high=127, size=a_shape).astype(dtype) w_np = np.random.randint(low=-128, high=128, size=w_shape).astype(dtype) b_np = np.random.uniform(size=bias_shape).astype(dtype) dw_np = topi.testing.dilate_python(w_np, (dilation, dilation, 1, 1)) c_np = topi.testing.conv2d_nhwc_python(a_np, dw_np, stride, padding).astype(dtype) if add_bias: b_np = np.random.uniform(size=bias_shape).astype(dtype) c_np += b_np if add_relu: c_np = np.maximum(c_np, 0) return a_np, w_np, b_np, c_np a_np, w_np, b_np, c_np = get_ref_data() def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): C = topi.arm_cpu.compute_conv2d_NHWC_quantized( A, W, (stride, stride), padding, (dilation, dilation), dtype) if add_bias: C = topi.add(C, bias) if add_relu: C = topi.nn.relu(C) s = topi.arm_cpu.schedule_conv2d_NHWC_quantized([C]) a = tvm.nd.array(a_np, ctx) w = tvm.nd.array(w_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), ctx) if add_bias: tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func = tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, b, c) else: func = tvm.build(s, [A, W, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, c) tvm.testing.assert_allclose(c.asnumpy(), c_np, rtol=1e-5) check_device("llvm")
def verify_conv2d_nchw( batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1, add_bias=False, add_relu=False, devices=['cuda', 'llvm -device=arm_cpu', 'opencl -device=mali']): pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple( padding, (kernel, kernel)) padding_sum = pad_top + pad_left + pad_bottom + pad_right print("Workload: (%d, %d, %d, %d, %d, %d, %d, %d)" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) in_height = in_width = in_size A = te.placeholder((batch, in_channel, in_height, in_width), name='A') W = te.placeholder((num_filter, in_channel, kernel, kernel), name='W') bias = te.placeholder((num_filter, 1, 1), name='bias') a_shape = get_const_tuple(A.shape) w_shape = get_const_tuple(W.shape) bias_shape = get_const_tuple(bias.shape) dtype = A.dtype @memoize("topi.tests.test_topi_conv2d_nchw.verify_conv2d_nchw") def get_ref_data(): a_np = np.random.uniform(size=a_shape).astype(dtype) w_np = np.random.uniform(size=w_shape).astype(dtype) b_np = np.random.uniform(size=bias_shape).astype(dtype) dw_np = topi.testing.dilate_python(w_np, (1, 1, dilation, dilation)) c_np = topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding) if add_bias: b_np = np.random.uniform(size=bias_shape).astype(dtype) c_np += b_np if add_relu: c_np = np.maximum(c_np, 0) return a_np, w_np, b_np, c_np a_np, w_np, b_np, c_np = get_ref_data() def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): fcompute, fschedule = topi.testing.dispatch( device, _conv2d_nchw_winograd_implement) C = fcompute(A, W, stride, padding, dilation, dtype) if add_bias: C = topi.add(C, bias) if add_relu: C = topi.nn.relu(C) s = fschedule([C]) a = tvm.nd.array(a_np, ctx) w = tvm.nd.array(w_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), ctx) if add_bias: func = tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, b, c) else: func = tvm.build(s, [A, W, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, c) rtol = 1e-3 tvm.testing.assert_allclose(c.asnumpy(), c_np, rtol=rtol) for device in devices: check_device(device)
def conv2d_direct_simd_compute(cfg, data, kernel, strides, padding, dilation, out_dtype): """Compute function for Cortex-M7 SIMD implementation of conv2d.""" assert isinstance(strides, int) or len(strides) == 2 assert isinstance(dilation, int) or len(dilation) == 2 if isinstance(strides, int): stride_h = stride_w = strides else: stride_h, stride_w = strides if isinstance(dilation, int): dilation_h = dilation_w = dilation else: dilation_h, dilation_w = dilation batch_size, in_height, in_width, in_channels = data.shape kernel_h, kernel_w, out_channels, _ = kernel.shape # compute the output shape dilated_kernel_h = (kernel_h - 1) * dilation_h + 1 dilated_kernel_w = (kernel_w - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( padding, (dilated_kernel_h, dilated_kernel_w)) out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1) out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1) pad_before = [0, pad_top, pad_left, 0] pad_after = [0, pad_down, pad_right, 0] padded_data = pad(data, pad_before, pad_after, name='padded_data') rc = te.reduce_axis((0, in_channels), name='rc') ry = te.reduce_axis((0, kernel_h), name='ry') rx = te.reduce_axis((0, kernel_w), name='rx') conv = te.compute( (batch_size, out_height, out_width, out_channels), lambda nn, yy, xx, ff: te.sum( padded_data[nn, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rc].astype(out_dtype) * kernel[ry, rx, ff, rc].astype(out_dtype), axis=[ry, rx, rc]), name='conv2d', tag='conv2d_nhwc') ########################### # Config Space Definition # ########################### n, oh, ow, co = (cfg.axis(batch_size.value), cfg.axis(out_height.value), cfg.axis(out_width.value), cfg.axis(out_channels.value)) kh, kw, ci = (cfg.reduce_axis(kernel_h.value), cfg.reduce_axis(kernel_w.value), cfg.reduce_axis(in_channels.value)) assert in_channels.value % 4 == 0 owo, owi = cfg.define_split('tile_ow', ow, policy='factors', num_outputs=2) cio, cii = cfg.define_split('tile_ci', ci, policy='factors', num_outputs=2, filter=lambda x: x.size[-1] % 4 == 0) coo, coi = cfg.define_split('tile_co', co, policy='factors', num_outputs=2) cfg.define_reorder('reorder_0_simd', [n, oh, owo, owi, coo, coi, kh, kw, cio, cii], policy='candidate', candidate=[ [n, oh, kh, kw, owo, coo, cio, owi, coi, cii], [n, oh, kh, kw, coo, owo, cio, owi, coi, cii], [n, kh, kw, oh, owo, coo, cio, owi, coi, cii], [n, kh, kw, oh, coo, owo, cio, owi, coi, cii]]) cfg.define_knob('auto_unroll_max_step', [0, 2, 4, 8, 16, 32]) cfg.define_knob('unroll_explicit', [0, 1]) return conv
def depthwise_conv2d_with_workload_NCHWc(batch, in_channel, in_height, channel_multiplier, filter_height, stride, padding, dilation=1): in_width = in_height filter_channel = in_channel filter_width = filter_height stride_h = stride_w = stride assert dilation == 1, "depthwise_conv2d_NCHWc currently does not support dilation." pad_h, pad_w, _, _ = get_pad_tuple(padding, (filter_height, filter_width)) padding_args = (pad_h, pad_w) out_channel = filter_channel * channel_multiplier # for testing functionality, # we choose arbitrary block size that can divide the channel, # regardless of the performance. oc_block = 1 for bn in range(16, 0, -1): if out_channel % bn == 0: oc_block = bn break ic_block = 1 for bn in range(oc_block, 0, -1): if in_channel % bn == 0: ic_block = bn break # placeholder Input = tvm.placeholder((batch, in_channel//ic_block, in_height, in_width, ic_block), name='Input') Filter = tvm.placeholder((out_channel//oc_block, 1, filter_height, filter_width, 1, oc_block), name='Filter') in_layout = "NCHW%dc" % ic_block out_layout = "NCHW%dc" % oc_block dtype = 'float32' def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): # declare DepthwiseConv2d = topi.nn.depthwise_conv2d_NCHWc(Input, Filter, (stride_h, stride_w), padding_args, (dilation, dilation), in_layout, out_layout, dtype) # TODO: add scale_shift implement for NCHWc and add test here Relu = topi.nn.relu(DepthwiseConv2d) # schedule s1 = topi.generic.schedule_depthwise_conv2d_nchw(DepthwiseConv2d) s2 = topi.generic.schedule_depthwise_conv2d_nchw(Relu) # build the kernels f1 = tvm.build(s1, [Input, Filter, DepthwiseConv2d], device) f2 = tvm.build(s2, [Input, Filter, Relu], device) # Prepare pod type for test data closure input_shape = (batch, in_channel, in_height, in_width) filter_shape = (filter_channel, channel_multiplier, filter_height, filter_width) # Use memoize, pickle the test data for next time use. @memoize("topi.tests.test_topi_depthwise_conv2d.NCHWc") def get_ref_data(): input_np = np.random.uniform(size=input_shape).astype(dtype) filter_np = np.random.uniform(size=filter_shape).astype(dtype) # correctness with scipy depthwise_conv2d_scipy = topi.testing.depthwise_conv2d_python_nchw( input_np, filter_np, stride, padding) relu_scipy = np.maximum(depthwise_conv2d_scipy, 0) return (_transform_data(input_np, ic_block), _transform_kernel(filter_np, oc_block), _transform_data(depthwise_conv2d_scipy, oc_block), _transform_data(relu_scipy, oc_block)) # Get the test data (input_np, filter_np, depthwise_conv2d_scipy, relu_scipy) = get_ref_data() input_tvm = tvm.nd.array(input_np, ctx) filter_tvm = tvm.nd.array(filter_np, ctx) depthwise_conv2d_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(DepthwiseConv2d.shape), dtype=DepthwiseConv2d.dtype), ctx) relu_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(Relu.shape), dtype=Relu.dtype), ctx) # launch kernel 1 (depthwise_conv2d) print(filter_tvm.shape) f1(input_tvm, filter_tvm, depthwise_conv2d_tvm) # launch kernel 2 (depthwise_conv2d + relu) f2(input_tvm, filter_tvm, relu_tvm) tvm.testing.assert_allclose(depthwise_conv2d_tvm.asnumpy(), depthwise_conv2d_scipy, rtol=1e-5) tvm.testing.assert_allclose(relu_tvm.asnumpy(), relu_scipy, rtol=1e-5) # test llvm only for now since depthwise_conv2d_NCHWc implement is missing in other backend. for device in ["llvm"]: with autotvm.tophub.context(device): # load tophub pre-tuned parameters check_device(device)
def _depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=None): """Depthwise convolution nchw forward operator. Parameters ---------- Input : tvm.Tensor 4-D with shape [batch, in_channel, in_height, in_width] Filter : tvm.Tensor 4-D with shape [in_channel, channel_multiplier, filter_height, filter_width] stride : tuple of two ints The spatial stride along height and width padding : int or str Padding size, or ['VALID', 'SAME'] dilation: int or a list/tuple of two ints dilation size, or [dilation_height, dilation_width] out_dtype: str, optional Output data type Returns ------- Output : tvm.Tensor 4-D with shape [batch, out_channel, out_height, out_width] """ out_dtype = Input.dtype if out_dtype is None else out_dtype if isinstance(stride, int): stride_h = stride_w = stride else: stride_h, stride_w = stride if isinstance(dilation, int): dilation_h = dilation_w = dilation else: dilation_h, dilation_w = dilation batch, in_channel, in_height, in_width = Input.shape # shape of dilated kernel filter_channel, channel_multiplier, filter_height, filter_width = Filter.shape dilated_kernel_h = (filter_height - 1) * dilation_h + 1 dilated_kernel_w = (filter_width - 1) * dilation_w + 1 pad_top, pad_left, pad_down, pad_right = get_pad_tuple( padding, (dilated_kernel_h, dilated_kernel_w)) out_channel = simplify(in_channel * channel_multiplier) out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1) out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1) # padding stage pad_before = [0, 0, pad_top, pad_left] pad_after = [0, 0, pad_down, pad_right] PaddedInput = topi.nn.pad(Input, pad_before, pad_after, name="PaddedInput") # depthconv stage di = tvm.te.reduce_axis((0, filter_height), name='di') dj = tvm.te.reduce_axis((0, filter_width), name='dj') Output = tvm.te.compute( (batch, out_channel, out_height, out_width), lambda b, c, i, j: tvm.te.sum( (PaddedInput[b, c/channel_multiplier, i*stride_h+di*dilation_h, j*stride_w+dj*dilation_w].astype(out_dtype) * Filter[c/channel_multiplier, c%channel_multiplier, di, dj].astype(out_dtype)), axis=[di, dj]), name='DepthwiseConv2d', tag="depthwise_conv2d_nchw") return Output
def decl_winograd(cfg, data, kernel, strides, padding, layout, out_dtype, VK=6, VP=8, packed_output=False): # return _baseline_winograd(cfg, data, kernel, strides, padding, layout, out_dtype) N, CI, IH, IW = get_const_tuple(data.shape) CO, _, KH, KW = get_const_tuple(kernel.shape) HSTR, WSTR = strides if isinstance(strides, (tuple, list)) else (strides, strides) HPAD, WPAD, _, _ = get_pad_tuple(padding, kernel) assert layout == 'NCHW' assert KH == 3 and KW == 3 and HPAD == 1 and WPAD == 1 and HSTR == 1 and WSTR == 1 data_pad = pad(data, (0, 0, HPAD, WPAD), name="data_pad") A_data = np.array( [[1, 1, 1, 1, 1, 32, 32, 0], [0, 1, -1, 2, -2, 16, -16, 0], [0, 1, 1, 4, 4, 8, 8, 0], [0, 1, -1, 8, -8, 4, -4, 0], [0, 1, 1, 16, 16, 2, 2, 0], [0, 1, -1, 32, -32, 1, -1, 1]], dtype=np.float32).T G_data = np.array( [[1, 0, 0], [-2 / 9, -2 / 9, -2 / 9], [-2 / 9, 2 / 9, -2 / 9], [1 / 90, 1 / 45, 2 / 45], [1 / 90, -1 / 45, 2 / 45], [1 / 45, 1 / 90, 1 / 180], [1 / 45, -1 / 90, 1 / 180], [0, 0, 1]], dtype=np.float32) B_data = np.array([[1, 0, -21 / 4, 0, 21 / 4, 0, -1, 0], [0, 1, 1, -17 / 4, -17 / 4, 1, 1, 0], [0, -1, 1, 17 / 4, -17 / 4, -1, 1, 0], [0, 1 / 2, 1 / 4, -5 / 2, -5 / 4, 2, 1, 0], [0, -1 / 2, 1 / 4, 5 / 2, -5 / 4, -2, 1, 0], [0, 2, 4, -5 / 2, -5, 1 / 2, 1, 0], [0, -2, 4, 5 / 2, -5, -1 / 2, 1, 0], [0, -1, 0, 21 / 4, 0, -21 / 4, 0, 1]], dtype=np.float32).T m = A_data.shape[1] r = 3 alpha = m + r - 1 C = CI H = (IH + 2 * HPAD - 3) // HSTR + 1 W = (IW + 2 * WPAD - 3) // WSTR + 1 nH, nW = (H + m - 1) // m, (W + m - 1) // m def round_up(a, b): return ((a + b - 1) // b) * b K = round_up(CO, VK) P = round_up(N * nH * nW, VP) assert K % VK == 0 assert P % VP == 0 G = const_matrix(G_data, 'G') r_kh = tvm.reduce_axis((0, KH), 'r_kh') r_kw = tvm.reduce_axis((0, KW), 'r_kw') assert K >= CO if K > CO: kernel_pad = pad(kernel, (0, 0, 0, 0), (K - CO, 0, 0, 0), name="kernel_pad") else: kernel_pad = kernel input_tile = tvm.placeholder(shape=(P // VP, C, alpha, alpha, VP), dtype='float32', name="input_tile") U = tvm.placeholder(shape=(K // VK, alpha, alpha, C, VK), dtype='float32', name="U") #U = tvm.compute( # (K // VK, alpha, alpha, C, VK), lambda k, eps, nu, c, kk: # tvm.sum(kernel_pad[k * VK + kk][c][r_kh][r_kw].astype(out_dtype) * # G[eps][r_kh] * G[nu][r_kw], axis=[r_kh, r_kw]), name='U') ## pack input tile #input_tile = tvm.compute((P // VP, C, alpha, alpha, VP), # lambda b, c, eps, nu, bb: # data_pad[(b*VP+bb) // (nH*nW)][c][(b*VP+bb) // nW % nH * m + eps] # [(b*VP+bb) % nW * m + nu], # name='d') def compute_B_T_dot_X(b, c, eps, nu, bb): temp_expr = {} for j in range(alpha): wd0 = input_tile[b][c][0][j][bb] - input_tile[b][c][6][j][bb] d4_sub_d2 = input_tile[b][c][4][j][bb] - input_tile[b][c][2][j][bb] wd7 = input_tile[b][c][7][j][bb] - input_tile[b][c][1][j][bb] d3_sub_d5 = input_tile[b][c][3][j][bb] - input_tile[b][c][5][j][bb] wd1 = input_tile[b][c][2][j][bb] + input_tile[b][c][6][j][bb] wd2 = input_tile[b][c][1][j][bb] + input_tile[b][c][5][j][bb] wd4 = input_tile[b][c][5][j][bb] + input_tile[b][c][1][j][bb] * 0.25 wd5 = input_tile[b][c][6][j][bb] - input_tile[b][c][4][j][bb] * 5 wd3 = input_tile[b][c][6][j][bb] + input_tile[b][c][2][j][bb] * 0.25 wd6 = input_tile[b][c][1][j][bb] + input_tile[b][c][5][j][bb] * 0.25 wd0 = wd0 + d4_sub_d2 * 5.25 wd7 = wd7 + d3_sub_d5 * 5.25 wd1 = wd1 - input_tile[b][c][4][j][bb] * 4.25 wd2 = wd2 - input_tile[b][c][3][j][bb] * 4.25 wd3 = wd3 - input_tile[b][c][4][j][bb] * 1.25 wd5 = wd5 + input_tile[b][c][2][j][bb] * 4 wd4 = wd4 - input_tile[b][c][3][j][bb] * 1.25 wd6 = wd6 - input_tile[b][c][3][j][bb] * 1.25 temp_expr[(0, j)] = wd0 temp_expr[(1, j)] = wd1 + wd2 temp_expr[(2, j)] = wd1 - wd2 temp_expr[(3, j)] = wd3 + wd4 * 2 temp_expr[(4, j)] = wd3 - wd4 * 2 temp_expr[(5, j)] = wd5 + wd6 * 2 temp_expr[(6, j)] = wd5 - wd6 * 2 temp_expr[(7, j)] = wd7 now = tvm.const(0.0, "float32") for ii in range(alpha): for jj in range(alpha): now = tvm.select(tvm.all(eps == ii, nu == jj), temp_expr[(ii, jj)], now) return now B_T_dot_X = tvm.compute((P // VP, C, alpha, alpha, VP), compute_B_T_dot_X, name="B_T_dot_X") def compute_X_dot_B(b, eps, nu, c, bb): temp_expr = {} for i in range(alpha): wd0 = B_T_dot_X[b][c][i][0][bb] - B_T_dot_X[b][c][i][6][bb] d4_sub_d2 = B_T_dot_X[b][c][i][4][bb] - B_T_dot_X[b][c][i][2][bb] wd7 = B_T_dot_X[b][c][i][7][bb] - B_T_dot_X[b][c][i][1][bb] d3_sub_d5 = B_T_dot_X[b][c][i][3][bb] - B_T_dot_X[b][c][i][5][bb] wd1 = B_T_dot_X[b][c][i][2][bb] + B_T_dot_X[b][c][i][6][bb] wd2 = B_T_dot_X[b][c][i][1][bb] + B_T_dot_X[b][c][i][5][bb] wd4 = B_T_dot_X[b][c][i][5][bb] + B_T_dot_X[b][c][i][1][bb] * 0.25 wd5 = B_T_dot_X[b][c][i][6][bb] - B_T_dot_X[b][c][i][4][bb] * 5 wd3 = B_T_dot_X[b][c][i][6][bb] + B_T_dot_X[b][c][i][2][bb] * 0.25 wd6 = B_T_dot_X[b][c][i][1][bb] + B_T_dot_X[b][c][i][5][bb] * 0.25 wd0 = wd0 + d4_sub_d2 * 5.25 wd7 = wd7 + d3_sub_d5 * 5.25 wd1 = wd1 - B_T_dot_X[b][c][i][4][bb] * 4.25 wd2 = wd2 - B_T_dot_X[b][c][i][3][bb] * 4.25 wd3 = wd3 - B_T_dot_X[b][c][i][4][bb] * 1.25 wd5 = wd5 + B_T_dot_X[b][c][i][2][bb] * 4 wd4 = wd4 - B_T_dot_X[b][c][i][3][bb] * 1.25 wd6 = wd6 - B_T_dot_X[b][c][i][3][bb] * 1.25 temp_expr[(i, 0)] = wd0 temp_expr[(i, 1)] = wd1 + wd2 temp_expr[(i, 2)] = wd1 - wd2 temp_expr[(i, 3)] = wd3 + wd4 * 2 temp_expr[(i, 4)] = wd3 - wd4 * 2 temp_expr[(i, 5)] = wd5 + wd6 * 2 temp_expr[(i, 6)] = wd5 - wd6 * 2 temp_expr[(i, 7)] = wd7 now = tvm.const(0.0, "float32") for ii in range(alpha): for jj in range(alpha): now = tvm.select(tvm.all(eps == ii, nu == jj), temp_expr[(ii, jj)], now) return now V = tvm.compute((P // VP, alpha, alpha, C, VP), compute_X_dot_B, name="V") # batch gemm c = tvm.reduce_axis((0, C), name='c') M = tvm.compute((K // VK, P // VP, alpha, alpha, VK, VP), lambda k, b, eps, nu, kk, bb: tvm.sum( U[k][eps][nu][c][kk] * V[b][eps][nu][c][bb], axis=c), name='M') def compute_A_T_dot_M(k, b, eps, nu, kk, bb): temp_expr = {} for j in range(alpha): m1_add_m2 = M[k][b][1][j][kk][bb] + M[k][b][2][j][kk][bb] m1_sub_m2 = M[k][b][1][j][kk][bb] - M[k][b][2][j][kk][bb] m3_add_m4 = M[k][b][3][j][kk][bb] + M[k][b][4][j][kk][bb] m3_sub_m4 = M[k][b][3][j][kk][bb] - M[k][b][4][j][kk][bb] m5_add_m6 = M[k][b][5][j][kk][bb] + M[k][b][6][j][kk][bb] m5_sub_m6 = M[k][b][5][j][kk][bb] - M[k][b][6][j][kk][bb] s0 = M[k][b][0][j][kk][bb] + m1_add_m2 s5 = M[k][b][7][j][kk][bb] + m1_sub_m2 s1 = m1_sub_m2 + m5_sub_m6 * 16 s4 = m1_add_m2 + m3_add_m4 * 16 s2 = m1_add_m2 + 8 * m5_add_m6 s3 = m1_sub_m2 + 8 * m3_sub_m4 s0 = s0 + m5_add_m6 * 32 s5 = s5 + m3_sub_m4 * 32 s1 = s1 + m3_sub_m4 * 2 s4 = s4 + m5_add_m6 * 2 s0 = s0 + m3_add_m4 s5 = s5 + m5_sub_m6 s2 = s2 + m3_add_m4 * 4 s3 = s3 + m5_sub_m6 * 4 temp_expr[(0, j)] = s0 temp_expr[(1, j)] = s1 temp_expr[(2, j)] = s2 temp_expr[(3, j)] = s3 temp_expr[(4, j)] = s4 temp_expr[(5, j)] = s5 now = tvm.const(0.0, "float32") for ii in range(m): for jj in range(alpha): now = tvm.select(tvm.all(eps == ii, nu == jj), temp_expr[(ii, jj)], now) return now A_T_dot_M = tvm.compute((K // VK, P // VP, m, alpha, VK, VP), compute_A_T_dot_M, name="A_T_dot_M") def compute_X_dot_A(k, b, eps, nu, kk, bb): temp_expr = {} for i in range(m): m1_add_m2 = A_T_dot_M[k][b][i][1][kk][bb] + A_T_dot_M[k][b][i][2][ kk][bb] m1_sub_m2 = A_T_dot_M[k][b][i][1][kk][bb] - A_T_dot_M[k][b][i][2][ kk][bb] m3_add_m4 = A_T_dot_M[k][b][i][3][kk][bb] + A_T_dot_M[k][b][i][4][ kk][bb] m3_sub_m4 = A_T_dot_M[k][b][i][3][kk][bb] - A_T_dot_M[k][b][i][4][ kk][bb] m5_add_m6 = A_T_dot_M[k][b][i][5][kk][bb] + A_T_dot_M[k][b][i][6][ kk][bb] m5_sub_m6 = A_T_dot_M[k][b][i][5][kk][bb] - A_T_dot_M[k][b][i][6][ kk][bb] s0 = A_T_dot_M[k][b][i][0][kk][bb] + m1_add_m2 s5 = A_T_dot_M[k][b][i][7][kk][bb] + m1_sub_m2 s1 = m1_sub_m2 + m5_sub_m6 * 16 s4 = m1_add_m2 + m3_add_m4 * 16 s2 = m1_add_m2 + 8 * m5_add_m6 s3 = m1_sub_m2 + 8 * m3_sub_m4 s0 = s0 + m5_add_m6 * 32 s5 = s5 + m3_sub_m4 * 32 s1 = s1 + m3_sub_m4 * 2 s4 = s4 + m5_add_m6 * 2 s0 = s0 + m3_add_m4 s5 = s5 + m5_sub_m6 s2 = s2 + m3_add_m4 * 4 s3 = s3 + m5_sub_m6 * 4 temp_expr[(i, 0)] = s0 temp_expr[(i, 1)] = s1 temp_expr[(i, 2)] = s2 temp_expr[(i, 3)] = s3 temp_expr[(i, 4)] = s4 temp_expr[(i, 5)] = s5 now = tvm.const(0.0, "float32") for ii in range(m): for jj in range(m): now = tvm.select(tvm.all(eps == ii, nu == jj), temp_expr[(ii, jj)], now) return now Y = tvm.compute((K // VK, P // VP, m, m, VK, VP), compute_X_dot_A, name="Y") # unpack output def _output(n, k_, h, w): b_idx = n * nH * nW + (h // m) * nW + w // m b = b_idx // VP bb = b_idx % VP k = k_ // VK kk = k_ % VK return Y[k][b][h % m][w % m][kk][bb] output = tvm.compute((N, CO, H, W), _output, name='output', tag='winograd_conv_output') if cfg: cfg.add_flop(2 * N * K * H * W * KH * KW * C) return Y, input_tile, U, output
def depth_1by1_fused(Input, Filter_d, Filter_1, stride_d, padding_d='SAME', dilation_d=1, out_dtype=None, layout="NCHW"): """Fused depthwise convolution + 1x1 convolution forward operator (NCHW & NHWC). Parameters ---------- Input : tvm.Tensor 4-D with shape [batch, in_channel, in_height, in_width] (NCHW) or [batch, in_height, in_width, in_channel] (NHWC) Filter_d : tvm.Tensor 4-D with shape [in_channel, in_channel * channel_multiplier, filter_height, filter_width] or [filter_height, filter_width, in_channel, in_channel * channel_multiplier] Filter_1 : tvm.Tensor 4-D with shape [out_channel, in_channel * channel_multiplier, 0, 0] or [0, 0, out_channel, in_channel * channel_multiplier] stride_d : tuple of two ints The spatial stride along height and width padding_d : int or str Padding size, or ['VALID', 'SAME'] dilation_d: int or a list/tuple of two ints dilation size, or [dilation_height, dilation_width] out_dtype: str, optional Output data type Returns ------- output : tvm.Tensor 4-D with shape [batch, out_height, out_width, out_channel] """ assert layout in ["NCHW", "NHWC"] out_dtype = Input.dtype if out_dtype is None else out_dtype if isinstance(stride_d, int): stride_h_d = stride_w_d = stride_d else: stride_h_d, stride_w_d = stride_d if isinstance(dilation_d, int): dilation_h_d = dilation_w_d = dilation_d else: dilation_h_d, dilation_w_d = dilation_d if layout == "NCHW": if dilation_h_d != 1 or dilation_w_d != 1: Filter_d = dilate(Filter_d, (1, 1, dilation_h_d, dilation_w_d)) batch, in_channel_d, in_height_d, in_width_d = Input.shape filter_channel, _, filter_height, filter_width = Filter_d.shape num_filter, channel, _, _ = Filter_1.shape else: # NHWC if dilation_h_d != 1 or dilation_w_d != 1: Filter_d = dilate(Filter_d, (dilation_h_d, dilation_w_d, 1, 1)) batch, in_height_d, in_width_d, in_channel_d = Input.shape filter_height, filter_width, filter_channel, _ = Filter_d.shape _, _, num_filter, channel = Filter_1.shape pad_top, pad_left, pad_down, pad_right = get_pad_tuple( padding_d, (filter_height, filter_width)) out_channel = simplify(in_channel_d) out_height = simplify((in_height_d - filter_height + pad_top + pad_down) // stride_h_d + 1) out_width = simplify((in_width_d - filter_width + pad_left + pad_right) // stride_w_d + 1) out_channel = num_filter # padding stage if layout == "NCHW": pad_before = [0, 0, pad_top, pad_left] pad_after = [0, 0, pad_down, pad_right] else: # NHWC pad_before = [0, pad_top, pad_left, 0] pad_after = [0, pad_down, pad_right, 0] PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput") # depthconv stage di = tvm.reduce_axis((0, filter_height), name='di') dj = tvm.reduce_axis((0, filter_width), name='dj') # 1by1 stage c = tvm.reduce_axis((0, out_channel), name='c') if layout == "NCHW": Output = tvm.compute( (batch, out_channel, out_height, out_width), lambda b, f, i, j: tvm.sum( (PaddedInput[b, c, i * stride_h_d + di, j * stride_w_d + dj]. astype(out_dtype) * Filter_d[c, 0, di, dj].astype( out_dtype) * Filter_1[f, c, 0, 0].astype(out_dtype)), axis=[di, dj, c]), name='Depthwise1by1Fused', tag="depthwise_1by1_fused_nchw") else: # NHWC Output = tvm.compute( (batch, out_height, out_width, out_channel), lambda b, i, j, f: tvm.sum( (PaddedInput[b, i * stride_h_d + di, j * stride_w_d + dj, c]. astype(out_dtype) * Filter_d[di, dj, c, 0].astype( out_dtype) * Filter_1[0, 0, c, f].astype(out_dtype)), axis=[di, dj, c]), name='Depthwise1by1Fused', tag="depthwise_1by1_fused_nhwc") return Output
def _conv_spatial_pack_asm(args, data, kernel, strides, padding, dilation, out_dtype): """_conv_spatial_pack_asm""" is_var, vh_, vw_, vc_ = args # create workload according to raw arguments out_dtype = out_dtype or data.dtype n_, ci_, ih_, iw_ = data.shape if is_var else get_const_tuple(data.shape) if isinstance(dilation, int): dilation_h = dilation_w = dilation else: dilation_h, dilation_w = dilation if len(kernel.shape) == 4: pre_packed = False co_, _, kh_, kw_ = kernel.shape if is_var else get_const_tuple( kernel.shape) else: # kernel tensor is pre packed pre_packed = True co_, _, kh_, kw_, vc_ = kernel.shape if is_var else get_const_tuple( kernel.shape) co_ = co_ * vc_ dilated_kernel_h = (kh_ - 1) * dilation_h + 1 dilated_kernel_w = (kw_ - 1) * dilation_w + 1 pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple( padding, (dilated_kernel_h, dilated_kernel_w)) hstr, wstr = strides if isinstance(strides, (tuple, list)) else (strides, strides) oh_ = (ih_ + pad_top + pad_bottom - dilated_kernel_h) // hstr + 1 ow_ = (iw_ + pad_left + pad_right - dilated_kernel_w) // wstr + 1 data_pad = pad(data, [0, 0, pad_top, pad_left], [0, 0, pad_bottom, pad_right]) oh_div = oh_ // vh_ ow_div = ow_ // vw_ kvshape = (co_ // vc_, ci_, kh_, kw_, vc_) ovshape = (n_, co_ // vc_, oh_div, ow_div, vh_, vw_, vc_) oshape = (n_, co_, oh_div * vh_, ow_div * vw_) if dilation_h != 1 or dilation_w != 1: # undilate input data dvshape = (n_, oh_ // vh_, ow_ // vw_, kh_, kw_, vh_, vw_, ci_) data_vec = tvm.compute( dvshape, lambda n, h, w, kh, kw, vh, vw, ci: data_pad[n][ci][ (h * vh_ + vh) * hstr + kh * dilation_h][ (w * vw_ + vw) * wstr + kw * dilation_w], name="data_vec_undilated", ) else: dvshape = ( n_, oh_ // vh_, ow_ // vw_, (vh_ - 1) * hstr + kh_, (vw_ - 1) * wstr + kw_, ci_, ) data_vec = tvm.compute( dvshape, lambda n, h, w, vh, vw, ci: data_pad[n][ci][h * vh_ * hstr + vh][ w * vw_ * wstr + vw], name="data_vec", ) if pre_packed: kernel_vec = kernel else: kernel_vec = tvm.compute( kvshape, lambda co, ci, kh, kw, vc: kernel[co * vc_ + vc][ci][kh][kw], name="kernel_vec", ) ci = tvm.reduce_axis((0, ci_), name="ci") kh = tvm.reduce_axis((0, kh_), name="kh") kw = tvm.reduce_axis((0, kw_), name="kw") # asm begin---- type_map = { "int8": "int32", "uint8": "uint32", "float32": "float32", "float16": "float16", } acum_dtype = type_map[data.dtype] attrs = { "SH": hstr, "SW": wstr, "PH": pad_top, "PW": pad_left, "DILA_H": dilation_h, "DILA_W": dilation_w, "VH": vh_, "VW": vw_, "VC": vc_, "ACUM_DTYPE": acum_dtype, } # asm end---- if dilation_h != 1 or dilation_w != 1: conv = tvm.compute( ovshape, lambda n, co, h, w, vh, vw, vc: tvm.sum( data_vec[n, h, w, kh, kw, vh, vw, ci].astype(out_dtype) * kernel_vec[co, ci, kh, kw, vc].astype(out_dtype), axis=[ci, kh, kw], ), name="conv", attrs=attrs, ) else: conv = tvm.compute( ovshape, lambda n, co, h, w, vh, vw, vc: tvm.sum( data_vec[n, h, w, vh * hstr + kh, vw * wstr + kw, ci].astype( out_dtype) * kernel_vec[co, ci, kh, kw, vc].astype( out_dtype), axis=[ci, kh, kw], ), name="conv", attrs=attrs, ) output = tvm.compute( oshape, lambda n, co, h, w: conv[n][co // vc_][h // vh_][w // vw_][h % vh_][ w % vw_][co % vc_], name="output_unpack", tag="asm_conv2d_output", ) return output
def depthwise_conv2d_with_workload_nchw(batch, in_channel, in_height, channel_multiplier, filter_height, stride, padding, dilation=1): in_width = in_height filter_channel = in_channel filter_width = filter_height stride_h = stride_w = stride if dilation == 1: # here we transform the padding argument from 'str' to 'tuple' , # because we need this to match the "workload" tuple to the records in TopHub pad_h, pad_w, _, _ = get_pad_tuple(padding, (filter_height, filter_width)) padding_args = (pad_h, pad_w) else: padding_args = padding # placeholder Input = tvm.placeholder((batch, in_channel, in_height, in_width), name='Input') Filter = tvm.placeholder((filter_channel, channel_multiplier, filter_height, filter_width), name='Filter') Scale = tvm.placeholder((in_channel * channel_multiplier,), name='Scale') Shift = tvm.placeholder((in_channel * channel_multiplier,), name='Shift') dtype = 'float32' def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): # declare DepthwiseConv2d = topi.nn.depthwise_conv2d_nchw(Input, Filter, (stride_h, stride_w), padding_args, dilation, dtype) ScaleShift = topi.nn.scale_shift_nchw(DepthwiseConv2d, Scale, Shift) Relu = topi.nn.relu(ScaleShift) # schedule s1 = topi.generic.schedule_depthwise_conv2d_nchw(DepthwiseConv2d) s2 = topi.generic.schedule_depthwise_conv2d_nchw(ScaleShift) s3 = topi.generic.schedule_depthwise_conv2d_nchw(Relu) # build the kernels f1 = tvm.build(s1, [Input, Filter, DepthwiseConv2d], device) f2 = tvm.build(s2, [Input, Filter, Scale, Shift, ScaleShift], device) f3 = tvm.build(s3, [Input, Filter, Scale, Shift, Relu], device) # Prepare pod type for test data closure input_shape = get_const_tuple(Input.shape) filter_shape = get_const_tuple(Filter.shape) scale_shape = get_const_tuple(Scale.shape) shift_shape = get_const_tuple(Shift.shape) scale_shift_shape = get_const_tuple(ScaleShift.shape) # Use memoize, pickle the test data for next time use. @memoize("topi.tests.test_topi_depthwise_conv2d.nchw") def get_ref_data(): input_np = np.random.uniform(size=input_shape).astype(dtype) filter_np = np.random.uniform(size=filter_shape).astype(dtype) dilated_filter_np = topi.testing.dilate_python(filter_np, (1, 1, dilation, dilation)) scale_np = np.random.uniform(size=scale_shape).astype(dtype) shift_np = np.random.uniform(size=shift_shape).astype(dtype) # correctness with scipy depthwise_conv2d_scipy = topi.testing.depthwise_conv2d_python_nchw( input_np, dilated_filter_np, stride, padding) scale_shift_scipy = np.zeros(shape=scale_shift_shape) for c in range(in_channel * channel_multiplier): scale_shift_scipy[:,c,:,:] = depthwise_conv2d_scipy[:,c,:,:] * scale_np[c] + shift_np[c] relu_scipy = np.maximum(scale_shift_scipy, 0) return (input_np, filter_np, scale_np, shift_np, depthwise_conv2d_scipy, scale_shift_scipy, relu_scipy) # Get the test data (input_np, filter_np, scale_np, shift_np, depthwise_conv2d_scipy, scale_shift_scipy, relu_scipy) = get_ref_data() input_tvm = tvm.nd.array(input_np, ctx) filter_tvm = tvm.nd.array(filter_np, ctx) scale_tvm = tvm.nd.array(scale_np, ctx) shift_tvm = tvm.nd.array(shift_np, ctx) depthwise_conv2d_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(DepthwiseConv2d.shape), dtype=DepthwiseConv2d.dtype), ctx) scale_shift_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(ScaleShift.shape), dtype=ScaleShift.dtype), ctx) relu_tvm = tvm.nd.array(np.zeros(shape=get_const_tuple(Relu.shape), dtype=Relu.dtype), ctx) # launch kernel 1 (depthwise_conv2d) timer_1 = f1.time_evaluator(f1.entry_name, ctx, number=1) tcost_1 = timer_1(input_tvm, filter_tvm, depthwise_conv2d_tvm).mean # launch kernel 2 (depthwise_conv2d + scale_shift) timer_2 = f2.time_evaluator(f2.entry_name, ctx, number=1) tcost_2 = timer_2(input_tvm, filter_tvm, scale_tvm, shift_tvm, scale_shift_tvm).mean # launch kernel 3 (depthwise_conv2d + scale_shift + relu) timer_3 = f3.time_evaluator(f3.entry_name, ctx, number=1) tcost_3 = timer_3(input_tvm, filter_tvm, scale_tvm, shift_tvm, relu_tvm).mean tvm.testing.assert_allclose(depthwise_conv2d_tvm.asnumpy(), depthwise_conv2d_scipy, rtol=1e-5) tvm.testing.assert_allclose(scale_shift_tvm.asnumpy(), scale_shift_scipy, rtol=1e-5) tvm.testing.assert_allclose(relu_tvm.asnumpy(), relu_scipy, rtol=1e-5) for device in get_all_backend(): with autotvm.tophub.context(device): # load tophub pre-tuned parameters check_device(device)
def verify_conv2d_nhwc(batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1, add_bias=False, add_relu=False, devices='cuda', bgemm="direct"): """Test the conv2d with winograd for nhwc layout""" pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel, kernel)) padding_sum = pad_top + pad_left + pad_bottom + pad_right print("Workload: (%d, %d, %d, %d, %d, %d, %d, %d)" % ( batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) in_height = in_width = in_size A = te.placeholder((batch, in_height, in_width, in_channel), name='A') W = te.placeholder((kernel, kernel, in_channel, num_filter), name='W') bias = te.placeholder((1, 1, 1, num_filter), name='bias') a_shape = get_const_tuple(A.shape) w_shape = get_const_tuple(W.shape) bias_shape = get_const_tuple(bias.shape) dtype = A.dtype @memoize("topi.tests.test_topi_conv2d_nhwc.verify_conv2d_nhwc") def get_ref_data(): a_np = np.random.uniform(size=a_shape).astype(dtype) w_np = np.random.uniform(size=w_shape).astype(dtype) b_np = np.random.uniform(size=bias_shape).astype(dtype) dw_np = topi.testing.dilate_python(w_np, (dilation, dilation, 1, 1)) c_np = topi.testing.conv2d_nhwc_python(a_np, dw_np, stride, padding) if add_bias: b_np = np.random.uniform(size=bias_shape).astype(dtype) c_np += b_np if add_relu: c_np = np.maximum(c_np, 0) return a_np, w_np, b_np, c_np a_np, w_np, b_np, c_np = get_ref_data() def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): if bgemm == "direct": fcompute, fschedule = topi.testing.dispatch(device, _conv2d_nhwc_winograd_direct) elif bgemm == "tensorcore": fcompute, fschedule = topi.testing.dispatch(device, _conv2d_nhwc_winograd_tensorcore) C = fcompute(A, W, stride, padding, dilation, 'float32') if add_bias: C = topi.add(C, bias) if add_relu: C = topi.nn.relu(C) s = fschedule([C]) a = tvm.nd.array(a_np, ctx) w = tvm.nd.array(w_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), ctx) if add_bias: func = tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % ( batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, b, c) else: func = tvm.build(s, [A, W, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % ( batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, c) tvm.testing.assert_allclose(c.asnumpy(), c_np, rtol=2e-3) check_device(devices)
def verify_conv2d_NCHWc(batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1, add_bias=False, add_relu=False, dtype="float32"): pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple( padding, (kernel, kernel)) padding_sum = pad_top + pad_left + pad_bottom + pad_right in_height = in_width = in_size print( "Workload: (%d, %d, %d, %d, %d, %d, %d)" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum)) # for testing functionality, # we choose arbitrary block size that can divide the channel, # regardless of the performance. oc_block = 1 for bn in range(16, 0, -1): if num_filter % bn == 0: oc_block = bn break ic_block = 1 for bn in range(oc_block, 0, -1): if in_channel % bn == 0: ic_block = bn break A = tvm.placeholder( (batch, in_channel // ic_block, in_height, in_width, ic_block), name='A') W = tvm.placeholder((num_filter // oc_block, in_channel // ic_block, kernel, kernel, ic_block, oc_block), name='W') bias = tvm.placeholder((num_filter // oc_block, 1, 1, oc_block), name='bias') @memoize("topi.tests.test_topi_conv2d_NCHWc.verify_conv2d_NCHWc") def get_ref_data(): a_np = np.random.uniform(size=(batch, in_channel, in_height, in_width)).astype(dtype) w_np = np.random.uniform(size=(num_filter, in_channel, kernel, kernel)).astype(dtype) b_np = np.random.uniform(size=(num_filter, 1, 1)).astype(dtype) dw_np = topi.testing.dilate_python(w_np, (1, 1, dilation, dilation)) c_np = topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding) if add_bias: c_np += b_np if add_relu: c_np = np.maximum(c_np, 0) return _transform_data(a_np, ic_block), _transform_kernel(w_np, ic_block, oc_block), \ _transform_bias(b_np, oc_block), _transform_data(c_np, oc_block) a_np, w_np, b_np, c_np = get_ref_data() def check_device(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) with tvm.target.create(device): C = topi.x86.conv2d_NCHWc(A, W, (stride, stride), padding, (dilation, dilation), 'NCHW%dc' % ic_block, "NCHW%dc" % oc_block, dtype) if add_bias: C = topi.add(C, bias) if add_relu: C = topi.nn.relu(C) s = topi.x86.schedule_conv2d_NCHWc([C]) a = tvm.nd.array(a_np, ctx) w = tvm.nd.array(w_np, ctx) b = tvm.nd.array(b_np, ctx) c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), ctx) if add_bias: func = tvm.build(s, [A, W, bias, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, b, c) else: func = tvm.build(s, [A, W, C], device, name="relu_%d_%d_%d_%d_%d_%d_%d_%d" % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation)) func(a, w, c) tvm.testing.assert_allclose(c.asnumpy(), c_np, rtol=1e-3) # test llvm only for now since conv2d_NCHWc implement is missing in other backend. for device in ["llvm"]: with autotvm.tophub.context( device): # load tophub pre-tuned parameters check_device(device)