def add_b_conv(fmap_shape, filter_shape, pad_, stride_, dilation_,
tile_hh=0, tile_coco=0, tile_mm=0, tile_kk=0, tile_nn=0, bypass_l1=False,
use_bias=False, block_size=16, conv_dtype='float16'):
conv, a_value, b_value, bias_value, kernel_name, dim_info = add_b_conv_compute(fmap_shape, filter_shape, pad_, stride_, dilation_,
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, bypass_l1,
use_bias, block_size, conv_dtype)
# schedule
s = akg.tvm.create_schedule(conv.op)
print(conv, a_value, b_value, bias_value)
attrs = {}
attrs["pragma_reschedule"] = True
attrs["pragma_rmselfdep"] = False
attrs['dim'] = dim_info
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
if use_bias:
mod = akg.build(s, [a_value, b_value, bias_value, conv], "cce", name=kernel_name, attrs=attrs, polyhedral=True)
else:
mod = akg.build(s, [a_value, b_value, conv], "cce", name=kernel_name, attrs=attrs, polyhedral=True)
source_code = mod.imported_modules[0].get_source()
cce_path = '.'
utils.create_code(kernel_name, cce_path, source_code)
return mod
def div_mod_issue(data_shape, weight_shape, case_number):
if (case_number == 0):
A = akg.tvm.placeholder(data_shape, dtype='float16', name='input0')
divisor = 2
stage1 = akg.tvm.compute(
data_shape,
lambda n, c, h, w: A[n, c / divisor, h, w] + 1,
name="stage1")
op_vars = [A, stage1]
s = akg.tvm.create_schedule([stage1.op])
akg.lower(s, op_vars, simple_mode=True, polyhedral=True)
with akg.build_config(add_lower_pass=cce.debug_mode(0),
dump_pass_ir=True):
mod = akg.build(s, op_vars, "cce", name="test1", polyhedral=True)
return mod
else:
A = akg.tvm.placeholder(data_shape, dtype='float16', name='input0')
B = akg.tvm.placeholder(weight_shape, dtype='float16', name='input1')
divisor = 3
stage1 = akg.tvm.compute(
data_shape,
lambda n, c, h, w: A[n, c / divisor, h, w] + 1,
name="stage1")
stage2 = akg.tvm.compute(
weight_shape,
lambda n, c, h, w: stage1[0, c, 0, 0] + B[n, c, h, w],
name="stage2")
op_vars = [A, B, stage2]
s = akg.tvm.create_schedule([stage2.op])
akg.lower(s, op_vars, simple_mode=True, polyhedral=True)
with akg.build_config(add_lower_pass=cce.debug_mode(0),
dump_pass_ir=True):
mod_stage2 = akg.build(s,
op_vars,
"cce",
name="test2",
polyhedral=True)
return mod_stage2
def range_run(start, limit, delta, dtype, attrs):
t_range = tvm_range.range_value(start, limit, delta, dtype)
# Create module
sch = akg.tvm.create_schedule(t_range.op)
kernel_name = "range"
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod = akg.build(sch, [t_range], "cce", name=kernel_name, attrs=attrs, polyhedral=True)
print(mod.imported_modules[0].get_source())
# Generate data for testing the op
expect = np.asarray(list(range(start, limit, delta)))
output = np.full((max(0, (limit - start) / delta),), np.nan, dtype)
output = utils.mod_launch(mod, (output, ), expect=expect)
return tuple(), output, expect, compare_tensor(output, expect, rtol=5e-03, equal_nan=True)
def elemwise_sum_manual_schedule(input_shape, polyhedral=False, attrs=None):
"""manually schedule"""
b = akg.tvm.placeholder(input_shape, dtype='float16', name="b")
c = akg.tvm.placeholder(input_shape, dtype='float16', name="c")
a = akg.tvm.compute(input_shape,
lambda *indices: b(*indices) + c(*indices))
ss = akg.tvm.create_schedule([a.op])
ss.cache_read(b, "local.UB", [a])
ss.cache_read(c, "local.UB", [a])
ss.cache_write(a, "local.UB")
ss[a].set_scope("local.UB")
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod = akg.build(ss, [b, c, a],
"cce",
name="test_manual_schedule",
attrs=attrs,
polyhedral=polyhedral)
return mod
def conv_relu(fmap_shape,
filter_shape,
pad_,
stride_,
dilation_,
tile_hh=0,
tile_coco=0,
tile_mm=0,
tile_kk=0,
tile_nn=0,
bypass_l1=False,
use_bias=False,
block_size=16,
conv_dtype='float16'):
conv, a_value, b_value, bias_value, kernel_name, dim_info = add_a_conv_compute(
fmap_shape, filter_shape, pad_, stride_, dilation_, tile_hh, tile_coco,
tile_mm, tile_kk, tile_nn, bypass_l1, use_bias, block_size, conv_dtype)
# leakly relu
negative_slope = 0.0
slope_tmp = akg.tvm.const(negative_slope, dtype=conv_dtype)
# negative_slope*x
out = akg.lang.cce.vmuls(conv, slope_tmp)
# max(x,negative_slope*x)
out = akg.lang.cce.vmax(out, conv)
# schedule
s = akg.tvm.create_schedule(conv.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
if use_bias:
mod = akg.build(s, [a_value, b_value, bias_value, conv],
"cce",
name=kernel_name,
attrs={"dim": dim_info},
polyhedral=True)
else:
mod = akg.build(s, [a_value, b_value, conv],
"cce",
name=kernel_name,
attrs={"dim": dim_info},
polyhedral=True)
return mod
def conv_02(fmap_shape, filter_shape, pad_, stride_, dilation_,
tile_hh=0, tile_coco=0, tile_mm=0, tile_kk=0, tile_nn=0, bypass_l1=False,
use_bias=False, block_size=16, conv_dtype='float16'):
# input shape (NCHW -> NC1HWC0)
in_n, in_c, in_h, in_w = fmap_shape
in_c = (in_c + block_size - 1) // block_size * block_size
# kernel shape (NCHW -> NC1HWC0 -> Fractal)
k_n, k_c, k_h, k_w = filter_shape
k_c = (k_c + block_size - 1) // block_size * block_size
k_n = (k_n + block_size - 1) // block_size * block_size
input_shape_nc1hwc0 = (in_n, in_c // block_size, in_h, in_w, block_size)
in_n, _, in_h, in_w, _ = input_shape_nc1hwc0
kernel_shape_nc1hwc0 = (k_n, k_c // block_size, k_h, k_w, block_size)
k_n, _, k_h, k_w, _ = kernel_shape_nc1hwc0
kernel_shape_fractal = (k_c // block_size * k_h * k_w, k_n // block_size, block_size, block_size)
# A placeholder (NC1HWCO)
A = akg.tvm.placeholder(input_shape_nc1hwc0, dtype=conv_dtype, name="input0")
# B_placeholder (fractal)
B = akg.tvm.placeholder(kernel_shape_fractal, dtype=conv_dtype, name="input1")
if use_bias:
bias_shape_nc1hwc0 = (1, k_n // block_size, 1, 1, block_size)
bias_name = "input2"
bias_value = akg.tvm.placeholder(bias_shape_nc1hwc0, dtype=conv_dtype, name=bias_name)
else:
bias_name = 'None'
bias_value = None
conv_forward = conv_compute_forward(fmap_shape, filter_shape, pad_, stride_, dilation_, A, B, bias_value,
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, bypass_l1,
use_bias, block_size, conv_dtype)
k_hw = k_h * k_w
const_shift = k_hw - 1
# B in Fractal format; result in Fractal format
def flip_weight(B, k_c, k_hw, const_shift):
out_shape = (B.shape[1].value * k_hw, k_c // block_size, block_size, block_size)
B_flip = akg.tvm.compute(out_shape,
lambda i0, i1, i2, i3:
B[i1 * k_hw + const_shift - truncmod(i0, k_hw), floordiv(i0, k_hw), i3, i2],
name=B.name + "_flipped")
return B_flip
# H in 5D format; result in 5D format
def strided_head(H, s_h, s_w):
n, c1, h, w, c0 = H.shape
out_shape = (n, c1, (h - 1) * s_h + 1, (w - 1) * s_w + 1, c0)
H_strided = akg.tvm.compute(out_shape,
lambda i0, i1, i2, i3, i4:
akg.tvm.expr.Select(akg.tvm.any(truncmod(i2, s_h) != 0, truncmod(i3, s_w) != 0),
akg.tvm.const(0.0, dtype="float16"),
H[i0, i1, floordiv(i2, s_h), floordiv(i3, s_w), i4]),
name=H.name + "_strided")
return H_strided
# A in 5D format; result in 5D format
def transpose_data(A):
out_shape = (A.shape[1].value * block_size, A.shape[0].value // block_size,
A.shape[2].value, A.shape[3].value, block_size)
A_transpose = akg.tvm.compute(out_shape,
lambda j0, j1, j2, j3, j4:
A[j1 * block_size + j4, floordiv(j0, block_size), j2, j3, truncmod(j0, block_size)],
name=A.name + "_transposed")
return A_transpose
# Head is in 5D format; result in Fractal format
def transpose_convert_head(Head):
out_shape = ((Head.shape[0].value // block_size) * Head.shape[2].value * Head.shape[3].value,
Head.shape[1].value, block_size, block_size)
tmp_6D_shape = (Head.shape[0].value // block_size, block_size,
Head.shape[1].value, Head.shape[2].value, Head.shape[3].value, block_size)
Head_6D = akg.topi.reshape(Head, tmp_6D_shape)
Head_6D_transpose = akg.topi.transpose(Head_6D, (0, 3, 4, 2, 5, 1))
Head_transpose_convert = akg.topi.reshape(Head_6D_transpose, out_shape)
return Head_transpose_convert
HEAD = akg.tvm.placeholder(conv_forward.shape, name="Head", dtype='float16')
Head_transposed_NCHW = (HEAD.shape[1].value * HEAD.shape[4].value, HEAD.shape[0].value,
HEAD.shape[2].value, HEAD.shape[3].value)
s_h, s_w = stride_
Head_strided_NCHW = (HEAD.shape[0].value, HEAD.shape[1].value * HEAD.shape[4].value,
(HEAD.shape[2].value - 1) * s_h + 1, (HEAD.shape[3].value - 1) * s_w + 1)
A_transposed_NCHW = (in_c, in_n, in_h, in_w)
K_flip_rot_NCHW = (k_c, k_n, k_h, k_w)
Head_transposed_converted = transpose_convert_head(HEAD)
pld_Head_transposed_converted = akg.tvm.placeholder(Head_transposed_converted.shape,
name="Head_trans_fractal", dtype=conv_dtype)
A_transposed = transpose_data(A)
pld_A_transposed = akg.tvm.placeholder(A_transposed.shape, name="A_trans", dtype=conv_dtype)
info = dim.Dim()
info.setdim(index=0, axis=0, tilel1=1, tilel0=1)
info.setdim(index=0, axis=1, tilel1=1, tilel0=1)
info.setdim(index=0, axis=2, tilel1=1, tilel0=1)
info.setdim(index=0, axis=3, tilel1=1, tilel0=1)
B_flip = flip_weight(B, k_c, k_hw, const_shift)
pld_B_flipped = akg.tvm.placeholder(B_flip.shape, name="B_flip", dtype=conv_dtype)
s_flipped = akg.tvm.create_schedule(B_flip.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_weight_flipped = akg.build(s_flipped, [B, B_flip], "cce", name=B.name + "_flipped",
attrs={"dim": str(info)}, polyhedral=True)
s_transposed_converted = akg.tvm.create_schedule(Head_transposed_converted.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_head_transposed_converted = akg.build(s_transposed_converted, [HEAD, Head_transposed_converted],
"cce", name="H_trans_converted",
attrs={"dim": str(info)},
polyhedral=True)
Head_strided = strided_head(HEAD, s_h, s_w)
pld_Head_strided = akg.tvm.placeholder(Head_strided.shape, name="Head_trans_5D", dtype=conv_dtype)
s_strided = akg.tvm.create_schedule(Head_strided.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_head_strided = akg.build(s_strided, [HEAD, Head_strided],
"cce", name="H_strided", attrs={"dim": str(info)}, polyhedral=True)
s_transposed = akg.tvm.create_schedule(A_transposed.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_transposed = akg.build(s_transposed, [A, A_transposed], "cce",
name="A_transposed", attrs={"dim": str(info)}, polyhedral=True)
ad_attrs = {"ad_conv_enable": 1, "ad_conv_reuse_conv": 1}
jacs = list(akg.differentiate(conv_forward, [A], HEAD, ad_attrs, [pld_Head_strided, pld_B_flipped, None]))
info = set_dims(Head_strided_NCHW, (k_c, k_n, k_h, k_w), (k_h - 1, k_w - 1), (1, 1), (1, 1),
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
sjac = akg.tvm.create_schedule([jacs[0].op])
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_AD_data = akg.build(sjac, [pld_Head_strided, pld_B_flipped, jacs[0]], "cce",
name="conv_AD_data", attrs={"dim": str(info)}, polyhedral=True)
conv_data = conv_compute_forward(Head_strided_NCHW, K_flip_rot_NCHW,
(k_h - 1, k_h - 1, k_w - 1, k_w - 1), (1, 1), (1, 1),
pld_Head_strided, pld_B_flipped, None,
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, bypass_l1,
use_bias, block_size, conv_dtype)
info = set_dims(Head_strided_NCHW, (k_c, k_n, k_h, k_w), (k_h - 1, k_w - 1), (1, 1), (1, 1),
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
s_data = akg.tvm.create_schedule(conv_data.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_data = akg.build(s_data, [pld_Head_strided, pld_B_flipped, conv_data], "cce",
name="conv_data", attrs={"dim": str(info)}, polyhedral=True)
ad_attrs = {"ad_conv_enable": 1, "ad_conv_reuse_conv": 1}
jacs = list(akg.differentiate(conv_forward, [B], HEAD, ad_attrs, [pld_A_transposed, pld_Head_transposed_converted, None]))
info = set_dims(A_transposed_NCHW, Head_transposed_NCHW, (0, 0), (1, 1), (s_h, s_w),
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
sjac = akg.tvm.create_schedule([jacs[0].op])
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_AD_weight = akg.build(sjac, [pld_A_transposed, pld_Head_transposed_converted, jacs[0]], "cce",
name="conv_AD_weight", attrs={"dim": str(info)}, polyhedral=True)
conv_weight = conv_compute_forward(A_transposed_NCHW, Head_transposed_NCHW,
(0, 0, 0, 0), (1, 1), (s_h, s_w),
pld_A_transposed, pld_Head_transposed_converted, None,
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, bypass_l1,
use_bias, block_size, conv_dtype)
info = set_dims(A_transposed_NCHW, Head_transposed_NCHW, (0, 0), (1, 1), (s_h, s_w),
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
s_weight = akg.tvm.create_schedule(conv_weight.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_weight = akg.build(s_weight, [pld_A_transposed, pld_Head_transposed_converted, conv_weight], "cce",
name="conv_weight", attrs={"dim": str(info)}, polyhedral=True)
return mod_AD_data, mod_AD_weight, mod_transposed, mod_head_transposed_converted, mod_head_strided, mod_weight_flipped
def conv_01(fmap_shape, filter_shape, pad_, stride_, dilation_,
tile_hh=0, tile_coco=0, tile_mm=0, tile_kk=0, tile_nn=0,
use_bias=False, block_size=16, conv_dtype='float16'):
# input shape (NCHW -> NC1HWC0)
in_n, in_c, in_h, in_w = fmap_shape
in_c = (in_c + block_size - 1) // block_size * block_size
# kernel shape (NCHW -> NC1HWC0 -> Fractal)
k_n, k_c, k_h, k_w = filter_shape
k_c = (k_c + block_size - 1) // block_size * block_size
k_n = (k_n + block_size - 1) // block_size * block_size
input_shape_nc1hwc0 = (in_n, in_c // block_size, in_h, in_w, block_size)
kernel_shape_nc1hwc0 = (k_n, k_c // block_size, k_h, k_w, block_size)
k_n, _, k_h, k_w, _ = kernel_shape_nc1hwc0
kernel_shape_fractal = (k_c // block_size * k_h * k_w, k_n // block_size, block_size, block_size)
# A placeholder (NC1HWCO)
A = akg.tvm.placeholder(input_shape_nc1hwc0, dtype=conv_dtype, name="input0")
# B_placeholder (fractal)
B = akg.tvm.placeholder(kernel_shape_fractal, dtype=conv_dtype, name="input1")
data = [A, B]
if use_bias:
bias_shape_nc1hwc0 = (1, k_n // block_size, 1, 1, block_size)
bias_name = "input2"
bias_value = akg.tvm.placeholder(bias_shape_nc1hwc0, dtype=conv_dtype, name=bias_name)
data.append(bias_value)
else:
bias_name = 'None'
bias_value = None
conv, _ = conv_origin.conv(data, fmap_shape, filter_shape, pad_, stride_, dilation_, use_bias)
kernel_name = 'conv_ad'
k_n, k_c, k_h, k_w = filter_shape
k_c = (k_c + block_size - 1) // block_size * block_size
k_n = (k_n + block_size - 1) // block_size * block_size
k_hw = k_h * k_w
const_shift = k_hw - 1
# B in Fractal format; result in Fractal format
def flip_weight(B, k_c, k_hw, const_shift):
out_shape = (B.shape[1].value * k_hw, k_c // block_size, block_size, block_size)
B_flip = akg.tvm.compute(out_shape,
lambda i0, i1, i2, i3: B[i1 * k_hw + const_shift - truncmod(i0, k_hw),
floordiv(i0, k_hw), i3, i2],
name=B.name + "_flipped")
return B_flip
def strided_head(H, s_h, s_w):
n, c1, h, w, c0 = H.shape
out_shape = (n, c1, (h - 1) * s_h + 1, (w - 1) * s_w + 1, c0)
H_strided = akg.tvm.compute(out_shape, lambda i0, i1, i2, i3, i4:
akg.tvm.expr.Select(akg.tvm.any(truncmod(i2, s_h) != 0,
truncmod(i3, s_w) != 0),
akg.tvm.const(0.0, dtype="float16"),
H[i0, i1, floordiv(i2, s_h), floordiv(i3, s_w), i4]),
name=H.name + "_strided")
return H_strided
B_flip = flip_weight(B, k_c, k_hw, const_shift)
pld_B_flip = akg.tvm.placeholder(B_flip.shape, name="inp1_flipped", dtype='float16')
HEAD = akg.tvm.placeholder(conv.shape, name="Head", dtype='float16')
HEAD_n, HEAD_c1, HEAD_h, HEAD_w, HEAD_c0 = HEAD.shape
info = set_dims((HEAD_n.value, HEAD_c1.value * HEAD_c0.value, HEAD_h.value, HEAD_w.value),
(k_c, k_n, k_h, k_w), (2, 2), (1, 1), (1, 1),
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
s_h, s_w = stride_
if (s_h == 1) and (s_w == 1):
ad_attrs = {"ad_conv_enable": 1, "ad_conv_reuse_conv": 1}
jacs = list(akg.differentiate(conv, [A], HEAD, ad_attrs, [HEAD, pld_B_flip, None]))
sjac = akg.tvm.create_schedule([jacs[0].op])
op_vars = [HEAD, pld_B_flip, jacs[0]]
info = set_dims((HEAD_n.value, HEAD_c1.value * HEAD_c0.value, HEAD_h.value, HEAD_w.value),
(k_c, k_n, k_h, k_w), (k_h - 1, k_w - 1), (1, 1), (1, 1),
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
else:
Head_strided = strided_head(HEAD, s_h, s_w)
pld_Head_strided = akg.tvm.placeholder(Head_strided.shape, name="head_strided", dtype='float16')
ad_attrs = {"ad_conv_enable": 1, "ad_conv_reuse_conv": 1}
jacs = list(akg.differentiate(conv, [A], HEAD, ad_attrs, [pld_Head_strided, pld_B_flip, None]))
sjac = akg.tvm.create_schedule([jacs[0].op])
op_vars = [pld_Head_strided, pld_B_flip, jacs[0]]
h_n, h_c1, h_h, h_w, h_c0 = pld_Head_strided.shape
info = set_dims((h_n.value, h_c1.value * h_c0.value, h_h.value, h_w.value), (k_c, k_n, k_h, k_w),
(k_h - 1, k_w - 1), (1, 1), (1, 1), tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_backward = akg.build(sjac, op_vars, "cce", name=kernel_name, attrs={"dim": str(info)}, polyhedral=True)
def transpose_data(A):
out_shape = (A.shape[1] * block_size, truncdiv(A.shape[0], block_size), A.shape[2], A.shape[3], block_size)
A_transpose = akg.tvm.compute(out_shape,
lambda j0, j1, j2, j3, j4:
A[j1 * block_size + j4, truncdiv(j0, block_size), j2, j3, truncmod(j0, block_size)],
name=A.name + "_transposed")
return A_transpose
# Head is in 5D format
# Output is in Fractal format
def transpose_convert_head(Head):
out_shape = ((floordiv(Head.shape[0].value, block_size)) * Head.shape[2].value * Head.shape[3].value,
Head.shape[1].value, block_size, block_size)
tmp_6D_shape = (floordiv(Head.shape[0].value, block_size),
block_size, Head.shape[1].value, Head.shape[2].value, Head.shape[3].value, block_size)
Head_6D = akg.topi.reshape(Head, tmp_6D_shape)
# Transpose from (N//block_size_N, block_size_N, C//block_size_C, H, W, block_size_C)
# to (N//block_size_N, H, W, C//block_size_C, block_size_C, block_size_N,)
Head_6D_transpose = akg.topi.transpose(Head_6D, (0, 3, 4, 2, 5, 1))
Head_transpose_convert = akg.topi.reshape(Head_6D_transpose, out_shape)
return Head_transpose_convert
X_transposed = transpose_data(A)
pld_X_transposed = akg.tvm.placeholder(X_transposed.shape, name="inp0_transposed", dtype='float16')
if (s_h > 1) or (s_w > 1):
Head_transposed_converted = strided_head(HEAD, s_h, s_w)
else:
Head_transposed_converted = HEAD
strided_head_n, strided_head_c1, strided_head_h, strided_head_w, strided_head_c0 = Head_transposed_converted.shape
Head_transposed_converted = transpose_convert_head(Head_transposed_converted)
s_transposed_converted = akg.tvm.create_schedule(Head_transposed_converted.op)
pld_Head_transposed_converted = akg.tvm.placeholder(Head_transposed_converted.shape,
name="head_transposed",
dtype='float16')
ad_attrs = {"ad_conv_enable": 1, "ad_conv_reuse_conv": 1}
jacs = list(akg.differentiate(conv, [B], HEAD, ad_attrs, [pld_X_transposed, pld_Head_transposed_converted, None]))
sjac = akg.tvm.create_schedule([jacs[0].op])
op_vars = [HEAD, pld_X_transposed, pld_Head_transposed_converted, jacs[0]]
in_n, in_c1, in_h, in_w, in_c0 = A.shape
info = set_dims((in_c1.value * in_c0.value, in_n.value, in_h.value, in_w.value),
(strided_head_c1.value * strided_head_c0.value, strided_head_n.value,
strided_head_h.value, strided_head_w.value),
(0, 0), (1, 1), (1, 1),
tile_hh, tile_coco, tile_mm, tile_kk, tile_nn, block_size)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_backward2 = akg.build(sjac, op_vars, "cce",
name="conv_backward_weight",
attrs={"dim": str(info)},
polyhedral=True)
return mod_backward, mod_backward2
def test_CCE_Conv(fmap_shape,
filter_shape,
pad_,
stride_,
tile_hh=0,
tile_coco=0,
tile_mm=0,
tile_kk=0,
tile_nn=0,
bypass_l1=False,
use_bias=False,
kernel_name="quant_conv",
cce_path='.'):
# input shape (NCHW -> NC1HWC0)
in_n, in_c, in_h, in_w = fmap_shape
input_shape_nc1hwc0 = (in_n, in_c // block_size, in_h, in_w, block_size)
# out_shape_nc1hwc0 = (in_n, in_c // 32, in_h, in_w, 32)
in_n, in_c1, in_h, in_w, in_c0 = input_shape_nc1hwc0
# kernel shape (NCHW -> NC1HWC0 -> Fractal)
k_n, k_c, k_h, k_w = filter_shape
kernel_shape_nc1hwc0 = (k_n, k_c // 32, k_h, k_w, 32)
k_n, k_c1, k_h, k_w, k_c0 = kernel_shape_nc1hwc0
kernel_shape_fractal = (k_c // 32 * k_h * k_w, k_n // 16, 16, 32)
f_ko, f_no, f_ni, f_ki = kernel_shape_fractal
# bias shape
bias_shape_nc1hwc0 = (1, k_n // block_size, 1, 1, block_size)
# padding ((padding_h, padding_w) -> (padding_top, padding_bottom, padding_left, padding_right))
padding = (pad_[0], pad_[0], pad_[1], pad_[1])
p_top, p_bottom, p_left, p_right = padding
# stride (stride_h, stride_w)
s_h, s_w = stride_
# A placeholder (NC1HWCO)
A = akg.tvm.placeholder(input_shape_nc1hwc0, dtype=conv_dtype, name='FMap')
# B_placeholder (fractal)
B = akg.tvm.placeholder(kernel_shape_fractal, dtype='int8', name='Filter')
ScaleQ = akg.tvm.placeholder((16, ), dtype='float16', name='ScaleQ')
OffsetQ = akg.tvm.placeholder((16, ), dtype='float16', name='OffsetQ')
out_shape_nc1hwc0 = (in_n, in_c // 32, in_h, in_w, 32)
q_n, q_c1, q_h, q_w, q_c0 = out_shape_nc1hwc0
# print out_shape_nc1hwc0
Quant = akg.tvm.compute(out_shape_nc1hwc0,
lambda qn, qc1, qh, qw, qc0:
(A[qn, qc1 + qc0 // 16, qh, qw, qc0 % 16] * ScaleQ[
0] + OffsetQ[0]).astype('int8'),
name='QuantOUT',
attrs={'no_inline': 1})
if use_bias:
bias_name = 'bias'
bias_value = akg.tvm.placeholder(bias_shape_nc1hwc0,
dtype=conv_dtype,
name=bias_name)
else:
bias_name = 'None'
# Create reduction variables
kc1 = akg.tvm.reduce_axis((0, k_c1), name='kc1')
kh = akg.tvm.reduce_axis((0, k_h), name='kh')
kw = akg.tvm.reduce_axis((0, k_w), name='kw')
kc0 = akg.tvm.reduce_axis((0, k_c0), name='kc0')
out_h = (in_h + p_top + p_bottom - k_h) // (s_h) + 1
tile_out_h = (tile_hh - k_h) // s_h + 1
out_w = (in_w + p_left + p_right - k_w) // (s_w) + 1
out_shape_nc1hwc0 = (in_n, k_n // block_size, out_h, out_w, block_size)
out_n, out_c1, out_h, out_w, out_c0 = out_shape_nc1hwc0
if (tile_coco > 0):
c1_cut = tile_coco // block_size
else:
c1_cut = out_c1
# set dim
index = 0
info = dim.Dim()
if (q_c1 > 1):
info.setdim(index=index, axis="KO", tilel1=q_c1, tilel0=q_c1) # ko
if (q_h > 1):
info.setdim(index=index,
axis="C1",
tilel1=tile_out_h,
tilel0=tile_out_h) # c1
if (q_w > 1):
info.setdim(index=index, axis="C0", tilel1=q_w, tilel0=q_w) # c0
if (q_c0 > 1):
info.setdim(index=index, axis="KI", tilel1=q_c0, tilel0=q_c0) # ki
index += 1
if (out_c1 > 1):
info.setdim(index=index, axis="C1", tilel1=c1_cut, tilel0=0) # c1
if (out_h > 1):
info.setdim(index=index, axis="H", tilel1=tile_out_h, tilel0=0) # h
if (out_w > 1):
info.setdim(index=index, axis="W", tilel1=out_w, tilel0=0) # w
if (out_c0 > 1):
info.setdim(index=index, axis="C0", tilel1=out_c0, tilel0=0) # c0
if (in_c1 > 1):
info.setdim(index=index, axis="KC1", tilel1=in_c1 / 2, tilel0=0) # kc1
if (k_h > 1):
info.setdim(index=index, axis="KH", tilel1=k_h, tilel0=0) # kh
if (k_w > 1):
info.setdim(index=index, axis="KW", tilel1=k_w, tilel0=0) # kw
info = str(info)
# Compute the convolution
output_name = "output0"
output_bias_name = "output1"
# print out_shape_nc1hwc0
C = akg.tvm.compute(
out_shape_nc1hwc0,
lambda n, c1, h, w, c0: akg.tvm.sum(akg.tvm.if_then_else(
akg.tvm.any((h * s_h + kh) < p_top, (h * s_h + kh) >
(in_h + p_top - 1), (w * s_w + kw) < p_left,
(w * s_w + kw) >
(in_w + p_left - 1)), akg.tvm.const(0.0, 'int8'),
Quant[n, kc1, (h * s_h + kh - p_top),
(w * s_w + kw - p_left), kc0]) * B[
(kc1 * k_h + kh) * k_w + kw, c1, c0, kc0],
axis=[kc1, kh, kw, kc0]),
name=output_name,
attrs={
"pragma_conv_kernel_n": k_n,
"pragma_conv_kernel_h": k_h,
"pragma_conv_kernel_w": k_w,
"pragma_conv_padding_top": p_top,
"pragma_conv_padding_bottom": p_bottom,
"pragma_conv_padding_left": p_left,
"pragma_conv_padding_right": p_right,
"pragma_conv_dilation_h": 1,
"pragma_conv_dilation_w": 1,
"pragma_conv_bypass_l1": 1 if bypass_l1 else 0,
"pragma_conv_stride_h": s_h,
"pragma_conv_stride_w": s_w,
"pragma_conv_fm_n": in_n,
"pragma_conv_fm_c": in_c,
"pragma_conv_fm_h": in_h,
"pragma_conv_fm_w": in_w,
"pragma_conv_h_cut": (h_window_cut - 1) * s_h + k_h,
"pragma_conv_w_cut": (in_w + p_left + p_right),
"pragma_conv_co_cut": c1_cut * k_c0,
"pragma_conv_m_cut": tile_mm,
"pragma_conv_k_cut": tile_kk,
"pragma_conv_n_cut": tile_nn,
"feature": Quant.op.name,
"filter": B.op.name,
"bias": bias_name,
"res": output_name,
"res_bias": output_bias_name
})
if use_bias:
cube = akg.tvm.compute(out_shape_nc1hwc0,
lambda n, c1, h, w, c0: C[n, c1, h, w, c0] +
bias_value[0, c1, 0, 0, c0],
name=output_bias_name)
else:
cube = C
if fusion:
# leakly relu
negative_slope = 0.0
slope_tmp = akg.tvm.const(negative_slope, dtype=conv_dtype)
# negative_slope*x
out = akg.lang.cce.vmuls(cube, slope_tmp)
# max(x,negative_slope*x)
out = akg.lang.cce.vmax(out, cube)
else:
out = cube
# schedule
s = akg.tvm.create_schedule(out.op)
attrs = {}
attrs["pragma_reschedule"] = 1
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
if fusion:
if use_bias:
mod = akg.build(s, [A, B, ScaleQ, OffsetQ, bias_value, out],
"cce",
name=kernel_name,
attrs=attrs,
attrs={"dim": info},
polyhedral=True)
else:
mod = akg.build(s, [A, B, ScaleQ, OffsetQ, out],
"cce",
name=kernel_name,
attrs=attrs,
attrs={"dim": info},
polyhedral=True)
else:
if use_bias:
mod = akg.build(s, [A, B, ScaleQ, OffsetQ, bias_value, out],
"cce",
name=kernel_name,
attrs=attrs,
attrs={"dim": info},
polyhedral=True)
else:
mod = akg.build(s, [A, B, ScaleQ, OffsetQ, out],
"cce",
name=kernel_name,
attrs=attrs,
attrs={"dim": info},
polyhedral=True)
source_code = mod.imported_modules[0].get_source()
# print(source_code)
# utils.create_code(kernel_name, cce_path, source_code)
if run_cce:
run_conv(mod, fmap_shape, filter_shape, pad_[0], stride_[0], use_bias)
def group_conv_forward(_n,
_h,
_w,
_c_i,
_c_o,
group,
_k_h,
_k_w,
_a,
_b,
bias_value,
pad_h,
pad_w,
_s_h,
_s_w,
cut_h,
cut_co,
cut_m,
cut_k,
cut_n,
block_size,
use_bias=False,
kernel_name='group_conv'):
if (not isinstance(_n, int)):
_n, _h, _w, _c_i, _c_o, group, _k_h, _k_w = expr_to_int(
(_n, _h, _w, _c_i, _c_o, group, _k_h, _k_w))
pad_h, pad_w, _s_h, _s_w = expr_to_int((pad_h, pad_w, _s_h, _s_w))
cut_h, cut_co, cut_m, cut_k, cut_n, block_size = expr_to_int(
(cut_h, cut_co, cut_m, cut_k, cut_n, block_size))
conv_dtype = 'float16'
if cut_h == _h:
cut_h += pad_h + pad_h
assert _c_o % group == 0 and _c_i % group == 0
assert _c_o % block_size == 0 and (_c_i // group) % block_size == 0
if (use_bias):
bias = bias_value
_o_h = (_h + 2 * pad_h - _k_h) // _s_h + 1
_o_w = (_w + 2 * pad_w - _k_w) // _s_w + 1
kc1 = akg.tvm.reduce_axis((0, _c_i // block_size // group), name='kc1')
kh = akg.tvm.reduce_axis((0, _k_h), name='kh')
kw = akg.tvm.reduce_axis((0, _k_w), name='kw')
kc0 = akg.tvm.reduce_axis((0, block_size), name='kc0')
p_top, p_bottom, p_left, p_right = pad_h, pad_h, pad_w, pad_w
output_name = 'output'
output_bias_name = 'output_bias'
C = akg.tvm.compute(
(_n, _c_o // block_size, _o_h, _o_w, block_size),
lambda n, c1, h, w, c0: akg.lang.cce.mmad(akg.tvm.if_then_else(
akg.tvm.any((h * _s_h + kh) < p_top, (h * _s_h + kh) >
(_h + p_top - 1), (w * _s_w + kw) < p_left,
(w * _s_w + kw) >
(_w + p_left - 1)), akg.tvm.const(0.0, conv_dtype),
_a[n, c1 // ((_c_o // block_size) // group) *
((_c_i // block_size) // group) + kc1, (h * _s_h + kh - p_top),
(w * _s_w + kw - p_left), kc0]) * _b[
(kc1 * _k_h + kh) * _k_w + kw, c1, c0, kc0],
axis=[kc1, kh, kw, kc0]),
attrs={
"pragma_conv_kernel_n": _c_o,
"pragma_conv_kernel_h": _k_h,
"pragma_conv_kernel_w": _k_w,
"pragma_conv_padding_top": p_top,
"pragma_conv_padding_bottom": p_bottom,
"pragma_conv_padding_left": p_left,
"pragma_conv_padding_right": p_right,
"pragma_conv_bypass_l1": 1,
"pragma_conv_stride_h": _s_h,
"pragma_conv_stride_w": _s_w,
"pragma_conv_fm_n": _n,
"pragma_conv_fm_c": _c_i,
"pragma_conv_fm_h": _h,
"pragma_conv_fm_w": _w,
"pragma_conv_dilation_h": 1,
"pragma_conv_dilation_w": 1,
"pragma_conv_h_cut": cut_h,
"pragma_conv_w_cut": _w + 2 * pad_w,
"pragma_conv_co_cut": cut_co,
"pragma_conv_m_cut": cut_m,
"pragma_conv_k_cut": cut_k,
"pragma_conv_n_cut": cut_n,
"feature": _a.op.name,
"filter": _b.op.name,
"bias": 'bias',
"res": output_name,
"res_bias": output_bias_name
},
name=output_name)
if use_bias:
out = akg.tvm.compute(
C.shape,
lambda n, c1, h, w, c0: C[n, c1, h, w, c0] + bias[0, c1, 0, 0, c0],
name=output_bias_name)
bufs = [_a, _b, bias, out]
else:
out = C
bufs = [_a, _b, out]
# create schedule for cce
s = akg.tvm.create_schedule([out.op])
# set dim
info = set_dims_group(cut_h, cut_co, cut_m, cut_k, cut_n,
expr_to_int(out.shape), _c_i, _c_o, group, _k_h,
_k_w, _s_h, block_size)
# build
with akg.build_config(add_lower_pass=cce.debug_mode(0),
dump_pass_ir=False):
mod = akg.build(s,
bufs,
"cce",
name=kernel_name,
attrs={"dim": info},
polyhedral=True)
return out
def group_conv_ad(_n,
_h,
_w,
_c_i,
_c_o,
group,
_k_h,
_k_w,
pad_h,
pad_w,
_s_h,
_s_w,
cut_h,
cut_co,
cut_m,
cut_k,
cut_n,
block_size,
use_bias=False,
kernel_name='group_conv'):
conv_dtype = 'float16'
_a = akg.tvm.placeholder((_n, _c_i // block_size, _h, _w, block_size),
name="input0",
dtype=conv_dtype)
_b = akg.tvm.placeholder(((_c_i // group) // block_size * _k_h * _k_w,
_c_o // block_size, block_size, block_size),
name="input1",
dtype=conv_dtype)
mod_forward = group_conv_forward(_n, _h, _w, _c_i, _c_o, group, _k_h, _k_w,
_a, _b, None, pad_h, pad_w, _s_h, _s_w,
cut_h, cut_co, cut_m, cut_k, cut_n,
block_size)
_o_h = mod_forward.shape[2].value
_o_w = mod_forward.shape[3].value
head = akg.tvm.placeholder(mod_forward.shape,
name="head",
dtype=conv_dtype)
# (_n,_c_o,_o_h,_o_w)--(stride)-->(_n,_c_o,(_o_h-1)*_s_h+1,
# (_o_w-1)*_s_w+1)--(5d)-->(_n,_c_o/16,(_o_h-1)*_s_h+1,(_o_w-1)*_s_w+1,16)
pld_head_strided = akg.tvm.placeholder(
(_n, _c_o // block_size, (_o_h - 1) * _s_h + 1,
(_o_w - 1) * _s_w + 1, block_size),
name="head_strided_5d",
dtype=conv_dtype)
# (_c_o,_c_i//group,_k_h,_k_w)--(flip)-->
# (_c_i,_c_o//group,_k_h,_k_w)--(Fractal)-->((_c_o//group)/16*_k_h*_k_w, _c_i/16,16,16)
pld_b_flipped = akg.tvm.placeholder(
((_c_o // group) // block_size * _k_h * _k_w, _c_i // block_size,
block_size, block_size),
name="b_flip",
dtype=conv_dtype)
# b in Fractal format; result in Fractal format
b_group_flipped = group_flip_weight(_b, _k_h, _k_w, group,
_c_o // group // block_size,
_c_i // group // block_size,
block_size)
s_gr_fl = akg.tvm.create_schedule([b_group_flipped.op])
info = dim.Dim()
info.setdim(index=0, axis=0, tilel1=1, tilel0=1)
info.setdim(index=0, axis=1, tilel1=1, tilel0=1)
info.setdim(index=0, axis=2, tilel1=1, tilel0=1)
info.setdim(index=0, axis=3, tilel1=1, tilel0=1)
with akg.build_config(add_lower_pass=cce.debug_mode(0),
dump_pass_ir=False):
mod_b_group_flip = akg.build(s_gr_fl, [_b, b_group_flipped],
"cce",
name="b_group_flip",
attrs={"dim": str(info)},
polyhedral=True)
head_strided = strided_head(head, _s_h, _s_w)
s_striding = akg.tvm.create_schedule(head_strided.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0),
dump_pass_ir=False):
mod_head_strided = akg.build(s_striding, [head, head_strided],
"cce",
name="h_strided",
attrs={"dim": str(info)},
polyhedral=True)
a_transposed = transpose_regroup(_a, block_size, group)
s_transposed_nc = akg.tvm.create_schedule(a_transposed.op)
info = dim.Dim()
info.setdim(index=0, axis=0, tilel1=16, tilel0=16)
info.setdim(index=0, axis=1, tilel1=1, tilel0=1)
info.setdim(index=0, axis=2, tilel1=1, tilel0=1)
info.setdim(index=0, axis=3, tilel1=1, tilel0=1)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_transposed_nc = akg.build(s_transposed_nc, [_a, a_transposed],
"cce",
name="a_transposed",
attrs={"dim": str(info)},
polyhedral=True)
head_transposed_convert = transpose_convert_head(head, block_size)
s_transposed_convert = akg.tvm.create_schedule(head_transposed_convert.op)
info = dim.Dim()
info.setdim(index=0, axis=0, tilel1=1, tilel0=1)
info.setdim(index=0, axis=1, tilel1=1, tilel0=1)
info.setdim(index=0, axis=2, tilel1=1, tilel0=1)
info.setdim(index=0, axis=3, tilel1=1, tilel0=1)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_transposed_convert = akg.build(s_transposed_convert,
[head, head_transposed_convert],
"cce",
name="a_transposed",
attrs={"dim": str(info)},
polyhedral=True)
# Begin with the ad kernels
ad_attrs = {"ad_conv_enable": 1}
_jacs_data = list(
akg.differentiate(mod_forward, [_a], head, ad_attrs,
[pld_head_strided, pld_b_flipped, None]))
cut_h_e, cut_co_e, cut_m_e, cut_k_e, cut_n_e = ((_o_h - 1) * _s_h + 1 + 2 *
(_k_h - 1 - pad_h), 16,
_h * _w, 48, 16)
cut_m_e = ((cut_m_e + block_size - 1) // block_size) * block_size
info = set_dims_group(cut_h_e, cut_co_e, cut_m_e, cut_k_e, cut_n_e,
expr_to_int(_a.shape), _c_o, _c_i, group, _k_h, _k_w,
_s_h, block_size)
s_data = akg.tvm.create_schedule([_jacs_data[0].op])
# low_data = akg.lower(s_data, [pld_head_strided, pld_b_flipped, _jacs_data[0]], simple_mode=True)
with akg.build_config(add_lower_pass=cce.debug_mode(0),
dump_pass_ir=False):
mod_ad_data = akg.build(
s_data, [pld_head_strided, pld_b_flipped, _jacs_data[0]],
"cce",
name="conv_ad_data",
attrs={"dim": info},
polyhedral=True)
# (_n,_c_i,_h,_w)--(trans)-->(_c_i,_n,_h,_w)--(regroup)-->
# (_c_i//group,_n*group,_h,_w)--(5d)-->(_c_i//group,(_n*group)/16,_h,_w,16)
pld_x_trans = akg.tvm.placeholder(
(_c_i // group, (_n * group) // block_size, _h, _w, block_size),
name="x_trans_5d",
dtype=conv_dtype)
# (_n,_c_o,_o_h,_o_w)--(trans)-->
# (_c_o,_n,_o_h,_o_w)--(Fractal)-->(_n/16*_o_h*_o_w, _c_o/16,16,16)
pld_head_trans_converted = akg.tvm.placeholder(
(_n // block_size * _o_h * _o_w, _c_o // block_size, block_size,
block_size),
name="head_trans_convert",
dtype=conv_dtype)
# ad_attrs = {"ad_conv_enable": 1}
_jacs_weights = list(
akg.differentiate(mod_forward, [_b], head, ad_attrs,
[pld_x_trans, pld_head_trans_converted, None]))
cut_h_e, cut_co_e, cut_m_e, cut_k_e, cut_n_e = (_h + 2 * pad_h, 16,
_k_h * _k_w, 48, 16)
cut_m_e = ((cut_m_e + block_size - 1) // block_size) * block_size
info = set_dims_group(
cut_h_e, cut_co_e, cut_m_e, cut_k_e, cut_n_e,
(_c_i // group, _c_o // block_size, _k_h, _k_w, block_size),
_n * group, _c_o, group, _o_h, _o_w, 1, block_size)
s_weights = akg.tvm.create_schedule([_jacs_weights[0].op])
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod_ad_weights = akg.build(
s_weights,
[pld_x_trans, pld_head_trans_converted, _jacs_weights[0]],
"cce",
name="conv_ad_weights",
attrs={"dim": info},
polyhedral=True)
print("Forward input data shape: ", _a.shape)
print("Forward input weight shape: ", _b.shape)
print("Forward output shape: ", mod_forward.shape)
print("Backward wrt. DATA input data shape: ", pld_head_strided.shape)
print("Backward wrt. DATA input weight shape: ", pld_b_flipped.shape)
print("Backward wrt. DATA output shape: ", _jacs_data[0].shape)
print("Backward wrt. WEIGHT input data shape: ", pld_x_trans.shape)
print("Backward wrt. WEIGHT input weight shape: ",
pld_head_trans_converted.shape)
print("Backward wrt. WEIGHT output shape: ", _jacs_weights[0].shape)
return mod_ad_data, mod_ad_weights, mod_b_group_flip, mod_head_strided, mod_transposed_nc, mod_transposed_convert
def psroialign_compute(fm_shape, roi_shape, class_num, group_size, sample_h,
sample_w, scale):
'''
:param fm_shape: (n, c_dim, h, w) where: c_dim = group_size * group_size * (class_num + 1)
:param roi_shape: (roi_num, 16, 1, 1). there are 5 value on dim C: score, x1, y1, x2, y2. The other 11 num is pads
:param class_num:
:param group_size:
:param sample_h:
:param sample_w:
:param scale:
:return:
'''
dtype = "float16"
fm_data = akg.tvm.placeholder(fm_shape, name="fm_data", dtype=dtype)
roi_data = akg.tvm.placeholder(roi_shape, name="roi_data", dtype=dtype)
scale_const = akg.tvm.const(scale, dtype=dtype)
sample_h_const = akg.tvm.const(sample_h, "int32")
sample_w_const = akg.tvm.const(sample_w, "int32")
two_const = akg.tvm.const(2, "float16")
one_const = akg.tvm.const(1, "float16")
group_size_const = akg.tvm.const(group_size, "int32")
bin_num = group_size * group_size
# ==============================================================
# step 1: scale coordinates size in original image to size in feature map
# ==============================================================
COSIZE = 16
roi_num = roi_shape[0]
aligned_roi_num = do_align(roi_num, COSIZE)
# 4 means x1, y1, x2, y2
# roi_shape[0] must be equal to COSIZE
scaled_coors = akg.tvm.compute(
(4, aligned_roi_num, 1, 1),
lambda n, c, h, w: roi_data[c, 1 + n, h, w] * scale_const,
name='scaled_coors')
# ==============================================================
# step 2: compute the width and height of roi
# ==============================================================
# 2 stands for width and height
width_height_shape = (2, aligned_roi_num, 1, 1)
width_height_of_rois = akg.tvm.compute(
width_height_shape,
lambda n, c, h, w: scaled_coors[n + 2, c, h, w] - scaled_coors[n, c, h,
w],
name='width_height_of_rois')
width_shape = (aligned_roi_num, )
width_of_rois = akg.tvm.compute(
width_shape,
lambda n: scaled_coors[2, n, 0, 0] - scaled_coors[0, n, 0, 0],
name='width_of_rois')
width_shape = (aligned_roi_num, )
height_of_rois = akg.tvm.compute(
width_shape,
lambda n: scaled_coors[1, n, 0, 0] - scaled_coors[3, n, 0, 0],
name='height_of_rois')
# ==============================================================
# step 3: compute the bias of the coordinates of all samples
# ==============================================================
# samples_shape = (aligned_roi_num, bin_num, sample_h, sample_w)
# unit_nums = akg.tvm.compute((2,), lambda i: two_const * group_size_const \
# * akg.tvm.expr.Select(i == 0, sample_w_const, sample_h_const), name = 'uint_nums')
# width_height_shape(0, x, x, x) indicates the width of a single unit which is separated by samples
# and width_height_shape(1, x, x, x) the height
# unit_lengths = akg.tvm.compute(width_height_shape, lambda n, c, h, w: width_height_of_rois(n, c, h, w) / unit_nums(n), \
# name = 'uint_lengths')
unit_w_lengths = akg.tvm.compute(
width_shape,
lambda n: width_of_rois(n) / sample_w_const * group_size_const,
name='uint_w_lengths')
unit_h_lengths = akg.tvm.compute(
width_shape,
lambda n: height_of_rois(n) / sample_h_const * group_size_const,
name='uint_h_lengths')
# samples_coors_x_shape = (aligned_roi_num, 1, group_size * sample_h, group_size * sample_w)
# samples_x_coors_bias = akg.tvm.compute(samples_coors_x_shape, lambda n, c, h, w: unit_w_lengths[n] * \
# (one_const + w * two_const), name = 'samples_x_coors_bias')
#
# samples_y_coors_bias = akg.tvm.compute(samples_coors_x_shape, lambda n, c, h, w: unit_h_lengths[n] * \
# (one_const + w * two_const), name = 'samples_y_coors_bias')
#
# samples_x_coors = akg.tvm.compute(samples_coors_x_shape, lambda n, c, h, w: \
# samples_x_coors_bias(n, c, h, w) + scaled_coors(1, c, 1, 1), name = 'samples_x_coors')
# samples_y_coors = akg.tvm.compute(samples_coors_x_shape, lambda n, c, h, w: \
# samples_y_coors_bias(n, c, h, w) + scaled_coors(2, c, 1, 1), name = 'samples_y_coors')
sample_w_bias_shape = (1, group_size, sample_w, aligned_roi_num)
# sample_w_bias = akg.tvm.compute(sample_w_bias_shape, lambda n, c, h, w: unit_w_lengths[w] * \
# (one_const + two_const * (c * sample_w_const + h)), name = 'samples_w_bias')
# sample_w_bias = akg.tvm.compute(sample_w_bias_shape, lambda n, c, h, w: unit_w_lengths[w] * \
# (one_const + two_const * (sample_w_const)), name = 'samples_w_bias')
sample_h_bias_shape = (1, group_size, sample_h, aligned_roi_num)
# sample_h_bias = akg.tvm.compute(sample_h_bias_shape, lambda n, c, h, w: unit_h_lengths[w] * \
# (one_const + two_const * (c * sample_h_const + h)), name = 'samples_h_bias')
# sample_h_bias = akg.tvm.compute(sample_h_bias_shape, lambda n, c, h, w: unit_h_lengths[w] * \
# (one_const + two_const * (sample_h_const)), name = 'samples_h_bias')
@akg.tvm.hybrid.script(capture=locals())
def gen_bias(h_value, unit_lengths, ratio):
output = output_tensor((1, group_size, h_value, aligned_roi_num),
'float16')
strides = allocate((aligned_roi_num, ), 'float16', 'local')
for w in range(0, aligned_roi_num):
strides[w] = half(0.0)
for c in range(0, group_size):
for h in range(0, 1):
for w in range(0, aligned_roi_num):
output[0, c, h, w] = unit_lengths[w]
# strides[w] += unit_lengths[w] * ratio * half(h_value)
for h in range(1, h_value):
for w in range(0, aligned_roi_num):
output[0, c, h, w] = output[0, c, h - 1,
w] + ratio * unit_lengths[w]
return output
sample_w_bias = gen_bias(sample_w_const, unit_w_lengths, two_const)
sample_h_bias = gen_bias(sample_h_const, unit_h_lengths, two_const)
samples_x_coors = akg.tvm.compute(
sample_w_bias_shape,
lambda n, c, h, w: sample_w_bias(n, c, h, w) + scaled_coors(
0, w, 0, 0),
name='samples_x_coors')
samples_y_coors = akg.tvm.compute(
sample_h_bias_shape,
lambda n, c, h, w: sample_h_bias(n, c, h, w) + scaled_coors(
1, w, 0, 0),
name='samples_y_coors')
# ==============================================================
# step 4: compute the low and high coordinates of samples for bilinear
# ==============================================================
# samples_x_coors_low = akg.tvm.compute(sample_w_bias_shape, lambda *indices: \
# akg.lang.cce.floor(samples_x_coors(*indices)), name = 'samples_x_coors_low')
# samples_x_coors_high = akg.tvm.compute(sample_w_bias_shape, lambda *indices: \
# akg.lang.cce.ceil(samples_x_coors(*indices)), name = 'samples_x_coors_high')
# samples_y_coors_low = akg.tvm.compute(sample_h_bias_shape, lambda *indices: \
# akg.lang.cce.floor(samples_y_coors(*indices)), name = 'samples_y_coors_low')
# samples_y_coors_high = akg.tvm.compute(sample_h_bias_shape, lambda *indices: \
# akg.lang.cce.ceil(samples_y_coors(*indices)), name = 'samples_y_coors_high')
samples_x_coors_low = akg.lang.cce.floor(samples_x_coors)
samples_x_coors_high = akg.lang.cce.ceil(samples_x_coors)
samples_y_coors_low = akg.lang.cce.floor(samples_y_coors)
samples_y_coors_high = akg.lang.cce.ceil(samples_y_coors)
# samples_x_coors_low = akg.tvm.compute(sample_w_bias_shape, lambda *indices: \
# akg.topi.cast(samples_x_coors(*indices), 'int32'), name = 'samples_x_coors_low')
# samples_x_coors_high = akg.tvm.compute(sample_w_bias_shape, lambda *indices: \
# samples_x_coors_low(*indices) + akg.topi.cast(one_const, 'int32'), name = 'samples_x_coors_high')
# samples_y_coors_low = akg.tvm.compute(sample_h_bias_shape, lambda *indices: \
# akg.topi.cast(samples_y_coors(*indices), 'int32'), name = 'samples_y_coors_low')
# samples_y_coors_high = akg.tvm.compute(sample_h_bias_shape, lambda *indices: \
# samples_y_coors_low(*indices) + akg.topi.cast(one_const, 'int32'), name = 'samples_y_coors_high')
# ==============================================================
# step 5: compute the weight of low and high coordinates for bilinear
# ==============================================================
# wlx = akg.tvm.compute(samples_coors_x_shape, lambda *indices: samples_x_coors_high(*indices) - samples_x_coors(*indices))
# whx = akg.tvm.compute(samples_coors_x_shape, lambda *indices: one_const - wlx(*indices))
#
# wly = akg.tvm.compute(samples_coors_x_shape, lambda *indices: samples_y_coors_high(*indices) - samples_y_coors(*indices))
# why = akg.tvm.compute(samples_coors_x_shape, lambda *indices: one_const - wly(*indices))
#
# wlxXwly = akg.tvm.compute(samples_coors_x_shape, lambda *indices: wlx(*indices) * wly(*indices))
# whxXwly = akg.tvm.compute(samples_coors_x_shape, lambda *indices: whx(*indices) * wly(*indices))
# wlxXwhy = akg.tvm.compute(samples_coors_x_shape, lambda *indices: wlx(*indices) * why(*indices))
# whxXwhy = akg.tvm.compute(samples_coors_x_shape, lambda *indices: whx(*indices) * why(*indices))
wlx = akg.tvm.compute(sample_w_bias_shape,
lambda *indices: samples_x_coors_high(*indices) -
samples_x_coors(*indices),
name='wlx')
whx = akg.tvm.compute(sample_w_bias_shape,
lambda *indices: one_const - wlx(*indices),
name='whx')
wly = akg.tvm.compute(sample_h_bias_shape,
lambda *indices: samples_y_coors_high(*indices) -
samples_y_coors(*indices),
name='wly')
why = akg.tvm.compute(sample_h_bias_shape,
lambda *indices: one_const - wly(*indices),
name='why')
samples_shape = (group_size, group_size, sample_h, sample_w,
aligned_roi_num)
wlxXwly = akg.tvm.compute(
samples_shape,
lambda i, j, m, n, k: wlx(0, j, n, k) * wly(0, i, m, k),
name='wlxXwly')
whxXwly = akg.tvm.compute(
samples_shape,
lambda i, j, m, n, k: whx(0, j, n, k) * wly(0, i, m, k),
name='whxXwly')
wlxXwhy = akg.tvm.compute(
samples_shape,
lambda i, j, m, n, k: wlx(0, j, n, k) * why(0, i, m, k),
name='wlxXwhy')
whxXwhy = akg.tvm.compute(
samples_shape,
lambda i, j, m, n, k: whx(0, j, n, k) * why(0, i, m, k),
name='whxXwhy')
boundaries_values_shape = (4, sample_h, sample_w, aligned_roi_num)
bin_values_shape = (1, class_num + 1, bin_num, aligned_roi_num)
gap_values_shape = (class_num + 1, aligned_roi_num)
@akg.tvm.hybrid.script
def fetch_data(shape, fm_in, c_idx, bin_idx, bin_num, group_size, sample_h,
sample_w, roi_num, x_low, x_high, y_low, y_high, one_value):
boundaries_values = output_tensor(shape, 'float16')
for i in range(0, sample_h):
for j in range(0, sample_w):
for k in range(0, roi_num):
# assume batch is 1
# w_low_idx = x_low[0, bin_idx % group_size, j, k]
# w_high_idx = x_high[0, bin_idx % group_size, j, k]
#
# h_low_idx = y_low[0, bin_idx // group_size, i, k]
# h_high_idx = y_high[0, bin_idx // group_size, i, k]
#x_low, y_low
boundaries_values[0, i, j, k] = one_value
boundaries_values[1, i, j, k] = one_value
boundaries_values[2, i, j, k] = one_value
boundaries_values[3, i, j, k] = one_value
# boundaries_values[0, i, j, k] = fm_in[0, c_idx * bin_num + bin_idx, h_low_idx, w_low_idx]
#
# #x_high, y_low
# boundaries_values[1, i, j, k] = fm_in[0, c_idx * bin_num + bin_idx, h_low_idx, w_high_idx]
#
# #x_low, y_high
# boundaries_values[2, i, j, k] = fm_in[0, c_idx * bin_num + bin_idx, h_high_idx, w_low_idx]
#
# #x_high, y_high
# boundaries_values[3, i, j, k] = fm_in[0, c_idx * bin_num + bin_idx, h_high_idx, w_high_idx]
return boundaries_values
@akg.tvm.hybrid.script(capture=locals())
def compute_bilinear_maxpool_gap(fm_in, x_low, x_high, y_low, y_high,
wlxXwly_, whxXwly_, wlxXwhy_, whxXwhy_,
one_value):
bin_values = allocate(bin_values_shape, 'float16', 'local')
# global average result
gap_values = output_tensor(gap_values_shape, 'float16')
for c in range(0, class_num + 1):
for b in range(0, bin_num):
boundaries_values = fetch_data(boundaries_values_shape, fm_in,
c, b, bin_num, group_size,
sample_h, sample_w, roi_num,
x_low, x_high, y_low, y_high,
one_value)
k_w = b % group_size
k_h = b // group_size
for n in range(0, roi_num):
bin_values[0, c, b, n] = half(0.0)
for h in range(0, sample_h):
for w in range(0, sample_w):
for n in range(0, roi_num):
# bilinear
tmp = boundaries_values[0, h, w, n] * wlxXwly_[k_h, k_w, h, w, n] + \
boundaries_values[1, h, w, n] * whxXwly_[k_h, k_w, h, w, n] + \
boundaries_values[2, h, w, n] * wlxXwhy_[k_h, k_w, h, w, n] + \
boundaries_values[3, h, w, n] * whxXwhy_[k_h, k_w, h, w, n]
# maxpooling
if tmp > bin_values[0, c, b, n]:
bin_values[0, c, b, n] = tmp
# global average pooling
for j in range(0, roi_num):
tmp1 = bin_values[0, c, 0, j]
for k in range(1, bin_num):
tmp1 += bin_values[0, c, k, j]
gap_values[c, j] = tmp1 / bin_num
return gap_values
# ==============================================================
# step 6: compute results of bilinear, maxpooling and global average pooling
# ==============================================================
out = compute_bilinear_maxpool_gap(fm_data, samples_x_coors_low,
samples_x_coors_high,
samples_y_coors_low,
samples_y_coors_high, wlxXwly, whxXwly,
wlxXwhy, whxXwhy, one_const)
# out = wlxXwhy
# info = dim.Dim()
# info.setdim(index=0, head = 0, body = 0, tail = 0, tilel1 = 1, tilel0 = 1)
# info.setdim(index=0, head = 0, body = 0, tail = 0, tilel1 = 1, tilel0 = 1)
s = akg.tvm.create_schedule(out.op)
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
# mod = akg.tvm.build(s, [fm_data, roi_data, out], "cce", name="psroialign", attrs = {"dim" : str(info)}, polyhedral=True)
mod = akg.build(s, [fm_data, roi_data, out],
"cce",
name="psroialign",
polyhedral=True)
return mod
def group_conv(N,
H,
W,
CI,
CO,
group,
KH,
KW,
PAD_H,
PAD_W,
SH,
SW,
cutH,
cutCo,
cutM,
cutK,
cutN,
block_size,
use_bias=False,
kernel_name='conv'):
"""
split channels of FeatureMap to some groups,every group has its filter-kernel
Args:
args1:a list,the size is 3 if use_bias else the size is 2;
data[0] akg.tvm.Tensor of type float16 ,shape 5D(N, CI//C0, C0, H, W)
data[1] akg.tvm.Tensor of type float16 ,shape 6D(CI//(CI//C0)//C0, KH, KW, k_ch*CI//C0, C0, C0)
data[2] akg.tvm.Tensor of type float16 ,shape 5D(N, CI*k_ch//C0, OH, OW, C0)
N:batchsize
H:height of featureMap
W:width of featureMap
CI:channel of featureMap
C0:num of Filters
group:num of spliting channels of FeatureMap
KH:height of Filter
KW:width of Filter
PAD_H:padding pixels in vertical direction
PAD_W:padding pixels in horizontal direction
SH:stride in vertical direction
SW:stride in horizontal direction
block_size:a int var
use_bias:a bool value
Returns:
akg.tvm.Tensor of same type as data, shape is 5D(N, C0//block_size, block_size, OH, OW)
"""
conv_dtype = "float16"
if cutH == H:
cutH += PAD_H + PAD_H
assert CO % group == 0 and CI % group == 0
assert CO % block_size == 0 and (CI // group) % block_size == 0
# (N, CI, H, W) -> (N, C0, H, W, C1)
A = akg.tvm.placeholder((N, CI // block_size, H, W, block_size),
dtype=conv_dtype,
name="A")
# (CO, CI // group, KH, KW) -> (CI // group // block * KH * KW, CO // block, block, block)
B = akg.tvm.placeholder((CI // group // block_size * KH * KW,
CO // block_size, block_size, block_size),
dtype=conv_dtype,
name="B")
bias = akg.tvm.placeholder((1, CO // block_size, 1, 1, block_size),
dtype=conv_dtype,
name="bias")
OH = (H + 2 * PAD_H - KH) // SH + 1
OW = (W + 2 * PAD_W - KW) // SW + 1
kc1 = akg.tvm.reduce_axis((0, CI // block_size // group), name="kc1")
kh = akg.tvm.reduce_axis((0, KH), name="kh")
kw = akg.tvm.reduce_axis((0, KW), name="kw")
kc0 = akg.tvm.reduce_axis((0, block_size), name="kc0")
p_top, p_bottom, p_left, p_right = PAD_H, PAD_H, PAD_W, PAD_W
output_name = "output"
output_bias_name = "output_bias"
C = akg.tvm.compute(
(N, CO // block_size, OH, OW, block_size),
lambda n, c1, h, w, c0: akg.lang.cce.mmad(akg.tvm.if_then_else(
akg.tvm.any((h * SH + kh) < p_top, (h * SH + kh) > (H + p_top - 1),
(w * SW + kw) < p_left, (w * SW + kw) >
(W + p_left - 1)), akg.tvm.const(0.0, conv_dtype),
A[n, c1 // ((CO // block_size) // group) * (
(CI // block_size) // group) + kc1, (h * SH + kh - p_top),
(w * SW + kw - p_left), kc0]) * B[
(kc1 * KH + kh) * KW + kw, c1, c0, kc0],
axis=[kc1, kh, kw, kc0]),
attrs={
"pragma_conv_kernel_n": CO,
"pragma_conv_kernel_h": KH,
"pragma_conv_kernel_w": KW,
"pragma_conv_padding_top": p_top,
"pragma_conv_padding_bottom": p_bottom,
"pragma_conv_padding_left": p_left,
"pragma_conv_padding_right": p_right,
"pragma_conv_bypass_l1": 1,
"pragma_conv_stride_h": SH,
"pragma_conv_stride_w": SW,
"pragma_conv_fm_n": N,
"pragma_conv_fm_c": CI,
"pragma_conv_fm_h": H,
"pragma_conv_fm_w": W,
"pragma_conv_dilation_h": 1,
"pragma_conv_dilation_w": 1,
"pragma_conv_h_cut": cutH,
"pragma_conv_w_cut": W + 2 * PAD_W,
"pragma_conv_co_cut": cutCo,
"pragma_conv_m_cut": cutM,
"pragma_conv_k_cut": cutK,
"pragma_conv_n_cut": cutN,
"feature": A.op.name,
"filter": B.op.name,
"bias": bias.op.name,
"res": output_name,
"res_bias": output_bias_name
},
name=output_name)
if use_bias:
out = akg.tvm.compute(
C.shape,
lambda n, c1, h, w, c0: C[n, c1, h, w, c0] + bias[0, c1, 0, 0, c0],
name=output_bias_name)
bufs = [A, B, bias, out]
else:
out = C
bufs = [A, B, out]
# create schedule for cce
s = akg.tvm.create_schedule([out.op])
# set cut / tiling
out_n, out_c1, out_h, out_w, out_c0 = akg.topi.util.get_const_tuple(
out.shape)
# set dim
tile_out_h = (cutH - KH) // SH + 1
info = dim.Dim()
if (out_n > 1):
info.setdim(index=0, axis=0, tilel1=1, tilel0=0) # n
if (out_c1 > 1):
info.setdim(index=0, axis=0, tilel1=cutCo // block_size,
tilel0=0) # c1
if (out_h > 1):
info.setdim(index=0, axis='H', tilel1=tile_out_h, tilel0=0) # h
if (out_w > 1):
info.setdim(index=0, axis=3, tilel1=out_w, tilel0=0) # w
if (out_c0 > 1):
info.setdim(index=0, axis=4, tilel1=out_c0, tilel0=0) # c0
assert CI // block_size // group == 1
if (CI // block_size // group > 1):
info.setdim(index=0,
axis=5,
tilel1=CI // block_size // group,
tilel0=0) # kc1
if (KH > 1):
info.setdim(index=0, axis=5, tilel1=KH, tilel0=0) # kh
if (KW > 1):
info.setdim(index=0, axis=5, tilel1=KW, tilel0=0) # kw
# build
with akg.build_config(add_lower_pass=cce.debug_mode(0), dump_pass_ir=True):
mod = akg.build(s,
bufs,
"cce",
name=kernel_name,
attrs={"dim": str(info)},
polyhedral=True)
return OH, OW, A, B, C, mod
