def convnet():
"""Alternating layout of simple convnet (from image super-resolution).
"""
bias1 = relay.var('bias1', shape=(64,))
bias2 = relay.var('bias2', shape=(64,))
bias3 = relay.var('bias3', shape=(64,))
bias4 = relay.var('bias4', shape=(64,))
weight1 = relay.var('weight1', shape=(64, 1, 5, 5))
weight2 = relay.var('weight2', shape=(64, 64, 3, 3))
weight3 = relay.var('weight3', shape=(64, 64, 3, 3))
weight4 = relay.var('weight4', shape=(64, 64, 3, 3))
data = relay.var("x", shape=(1, 1, 224, 224))
n00 = relay.nn.conv2d(data, weight1, padding=[2, 2], kernel_size=[5, 5])
n01 = relay.expand_dims(bias1, axis=1, num_newaxis=2)
n02 = relay.add(n00, n01)
n03 = relay.nn.relu(n02)
n04 = relay.nn.conv2d(n03, weight2, padding=[1, 1], kernel_size=[3, 3])
n05 = relay.expand_dims(bias2, axis=1, num_newaxis=2)
n06 = relay.add(n04, n05)
n07 = relay.nn.relu(n06)
n08 = relay.nn.conv2d(n07, weight3, padding=[1, 1], kernel_size=[3, 3])
n09 = relay.expand_dims(bias3, axis=1, num_newaxis=2)
n10 = relay.add(n08, n09)
n11 = relay.nn.relu(n10)
n12 = relay.nn.conv2d(n11, weight4, padding=[1, 1], kernel_size=[3, 3])
n13 = relay.expand_dims(bias4, axis=1, num_newaxis=2)
n14 = relay.add(n12, n13)
n15 = relay.reshape(n14, newshape=[1, 1, 3, 3, 224, 224])
n16 = relay.transpose(n15, axes=[0, 1, 4, 2, 5, 3])
net = relay.reshape(n16, newshape=[1, 1, 672, 672])
args = relay.ir_pass.free_vars(net)
return relay.Function(args, net)
def get_net(batch_size, random_len=100, oshape=(3, 64, 64), ngf=128, code=None, dtype="float32"):
"""get net of dcgan generator"""
assert oshape[-1] == 64, "Only support 64x64 image"
assert oshape[-2] == 64, "Only support 64x64 image"
code = relay.var("data", dtype=dtype, shape=(batch_size, random_len)) if code is None else code
dense_weight = relay.var("dense_weight")
dense = relay.nn.dense(code, weight=dense_weight, units=4*4*ngf*8)
relu = relay.nn.relu(dense)
# 4 x 4
reshape = relay.reshape(relu, newshape=(-1, ngf * 8, 4, 4))
# 8 x 8
dc8 = deconv2d_bn_relu(
reshape, ishape=(ngf * 8, 4, 4), oshape=(ngf * 4, 8, 8), kshape=(4, 4), prefix="g2")
# 16x16
dc16 = deconv2d_bn_relu(
dc8, ishape=(ngf * 4, 8, 8), oshape=(ngf * 2, 16, 16), kshape=(4, 4), prefix="g3")
# 32x32
dc32 = deconv2d_bn_relu(
dc16, ishape=(ngf * 2, 16, 16), oshape=(ngf, 32, 32), kshape=(4, 4), prefix="g4")
# 64x64
dc64 = deconv2d(
dc32, ishape=(ngf, 32, 32), oshape=oshape[-3:], kshape=(4, 4), name="g5_deconv")
tanh = relay.tanh(dc64)
args = relay.ir_pass.free_vars(tanh)
return relay.Function(args, tanh)
def test_reshape_infer_type():
n, t, d1, d2 = 10, 20, 100, 20
x = relay.var("x", relay.TensorType((n, t, d1, d2), "float32"))
y = relay.reshape(x, newshape=(n, t, 2000))
assert "newshape=" in y.astext()
yy = relay.ir_pass.infer_type(y)
assert yy.checked_type == relay.TensorType(
(n, t, 2000), "float32")
def verify_reshape(shape, newshape, oshape):
x = relay.var("x", relay.TensorType(shape, "float32"))
z = relay.reshape(x, newshape=newshape)
zz = relay.ir_pass.infer_type(z)
assert "newshape=" in z.astext()
assert zz.checked_type == relay.ty.TensorType(oshape, "float32")
func = relay.Function([x], z)
x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
ref_res = np.reshape(x_data, oshape)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def verify_reshape(shape, newshape, oshape):
x = relay.var("x", relay.TensorType(shape, "float32"))
y = relay.var("y", relay.TensorType((len(newshape), ), "int64"))
z = relay.reshape(x, y)
func = relay.Function([x, y], z)
x_data = np.random.uniform(low=-1, high=1,
size=shape).astype("float32")
x_data = np.ones(shape).astype("float32")
ref_res = np.reshape(x_data, oshape)
check_grad(run_infer_type(func),
inputs=[x_data, np.array(newshape).astype("int64")],
test_inputs=[x_data],
eps=1e-3)
verify_func(func, [x_data, np.array(newshape).astype("int64")],
ref_res)
def test_vm_reshape_tuple(x_shape=(1, 4, 2), y_shape=(1, 2, 10)):
tup = relay.var(
"tup",
type_annotation=relay.TupleType(
[relay.TensorType(x_shape),
relay.TensorType(y_shape)]),
)
out = relay.reshape(relay.TupleGetItem(tup, 0), (1, -1))
f = relay.Function([tup], out)
x_data = np.random.uniform(size=x_shape).astype("float32")
y_data = np.random.uniform(size=y_shape).astype("float32")
for tgt, dev in tvm.testing.enabled_targets():
res = veval(f, (x_data, y_data), device=dev, target=tgt)
tvm.testing.assert_allclose(res.asnumpy(), np.reshape(x_data, (1, -1)))
def verify_reshape(shape, newshape, oshape):
x = relay.var("x", relay.TensorType(shape, "float32"))
z = relay.reshape(x, newshape=newshape)
zz = run_infer_type(z)
assert "newshape=" in z.astext()
assert zz.checked_type == relay.ty.TensorType(oshape, "float32")
func = relay.Function([x], z)
check_grad(func)
x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
ref_res = np.reshape(x_data, oshape)
for target, ctx in ctx_list():
for kind in ["graph", "debug"]:
intrp = relay.create_executor(kind, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def test_inline_no_ops():
input = relay.var("input", shape=(12, 12), dtype="uint8")
slice = relay.strided_slice(input, [0, 0], [6, 6])
relu1 = relay.nn.relu(slice)
reshape = relay.reshape(relu1, (36, ))
relu2 = relay.nn.relu(reshape)
func = relay.Function(relay.analysis.free_vars(relu2), relu2)
func = run_opt_pass(func, relay.transform.InferType())
cached_func = lower_to_te(func)
sch = te.create_schedule([cached_func.outputs[0].op])
inline_no_ops(cached_func, sch)
reshape_tensor = cached_func.outputs[0].op.input_tensors[0]
slice_tensor = reshape_tensor.op.input_tensors[0].op.input_tensors[0]
assert sch[reshape_tensor].attach_type == AttachType.kInline
assert sch[slice_tensor].attach_type == AttachType.kInline
def verify_reshape(shape, newshape, oshape):
x = relay.var("x", relay.TensorType(shape, "float32"))
y = relay.var("y", relay.TensorType(newshape, "float32"))
z = relay.reshape(x, relay.shape_of(y))
func = run_infer_type(relay.Function([x, y], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("reshape")
assert "newshape=" in zz.astext()
assert zz.checked_type == relay.ty.TensorType(oshape, "float32")
x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
y_data = np.random.uniform(low=-1, high=1, size=newshape).astype("float32")
ref_res = np.reshape(x_data, oshape)
verify_func(func2, [x_data, y_data], ref_res)
def make_net(network_arch, nfeat, nlabel, max_inp_len):
out = relay.var('data', shape=(1, nfeat, 1, max_inp_len),
dtype='float32') # NCHW
needs_flat = True
itm_ctr = defaultdict(int)
for i, line in enumerate(network_arch):
if line.startswith('#'):
continue
layer = line.rstrip().split(' ')
itm_ctr[layer[0]] += 1
if layer[0] == 'V': # view
continue
elif layer[0] == 'C':
if layer[1] == 'NFEAT':
layer[1] = nfeat
_ch_in, ch_out, k, s, p = map(int, layer[1:])
assert s == 1
out = _conv(out, ch_out, k, s, p, name=f'conv_{itm_ctr[layer[0]]}')
elif layer[0] == 'GLU':
# axis = int(layer[1])
out = relay.nn.glu(out, 1)
elif layer[0] == 'DO':
out = relay.nn.dropout(out, float(layer[1]))
elif layer[0] == 'RO': # reorder
pass
# out = relay.transpose(out, tuple(map(int, layer[1:])))
elif layer[0] == 'L': # linear
if layer[2] == 'NLABEL':
layer[2] = nlabel
if needs_flat:
out = relay.reshape(out, (1, 908))
out = _dense(out, int(layer[2]), name=f'dense_{itm_ctr[layer[0]]}')
needs_flat = False
else:
raise RuntimeError('Unknown layer type: ' + layer[0])
args = relay.ir_pass.free_vars(out)
return relay.Function(args, out)
def verify_any_reshape(x_shape, newshape, x_np_shape, out_shape, variable_newshape=False):
x = relay.var("x", shape=x_shape, dtype="float32")
relu_x = relay.nn.relu(x)
data = np.random.uniform(size=x_np_shape).astype("float32")
params = [x]
args = [data]
if variable_newshape:
newshape_var = relay.var("newshape", shape=(len(newshape),), dtype="int64")
params.append(newshape_var)
args.append(np.array(newshape, dtype="int64"))
newshape = newshape_var
y = relay.reshape(relu_x, newshape=newshape)
mod = tvm.IRModule()
mod["main"] = relay.Function(params, y)
check_result(args, mod, data, flatten=True)
def _get_func(ifm_shape, reshaped, ifm_layout):
ifm = relay.var("ifm", shape=ifm_shape, dtype="int8")
ifm_reshaped = relay.reshape(ifm, reshaped)
conv = make_ethosu_conv2d(
ifm_reshaped,
reshaped[3],
16,
(3, 3),
(1, 1),
(1, 1),
(1, 1),
activation="NONE",
ifm_layout=ifm_layout,
)
func = relay.Function(relay.analysis.free_vars(conv), conv)
func = run_opt_pass(func, relay.transform.InferType())
return func
def test_full(self):
fill_val = relay.expr.const(1.0)
fill_shape = [10, 1]
net = relay.full(fill_val, fill_shape, "float32")
net = relay.reshape(net, [1, -1])
mod = tvm.IRModule.from_expr(net)
params = {}
xgraph = xf_relay.from_relay(mod, params)
layers = xgraph.get_layers()
assert layers[0].type[0] == "Constant"
assert layers[0].shapes == [1]
assert layers[1].type[0] == "AnyOp"
assert layers[1].shapes == [10, 1]
assert layers[2].type[0] == "Reshape"
assert layers[2].shapes == [1, 10]
def legalize_group_conv(attrs, inputs, types):
"""Legalize group conv / conv_transpose calculation.
Alter weight layout from OIHW to GOIHW / IOHW to GIOHW"""
groups = attrs.groups
data, weight = inputs
if groups == 1:
if "Transpose" not in type(attrs).__name__:
return relay.nn.conv2d(data, weight, **attrs)
return relay.nn.conv2d_transpose(data, weight, **attrs)
OC, IC, H, W = get_shape(weight)
new_attrs = dict(attrs)
weight = relay.reshape(weight, (groups, OC // groups, IC, H, W))
if "Transpose" not in type(attrs).__name__:
new_attrs["kernel_layout"] = "GOIHW"
return relay.nn.conv2d(data, weight, **new_attrs)
new_attrs["kernel_layout"] = "GIOHW"
return relay.nn.conv2d_transpose(data, weight, **new_attrs)
def get_net(batch_size,
random_len=100,
oshape=(3, 64, 64),
ngf=128,
code=None,
dtype="float32"):
"""get net of dcgan generator"""
assert oshape[-1] == 64, "Only support 64x64 image"
assert oshape[-2] == 64, "Only support 64x64 image"
code = relay.var("data", dtype=dtype,
shape=(batch_size, random_len)) if code is None else code
dense_weight = relay.var("dense_weight")
dense = relay.nn.dense(code, weight=dense_weight, units=4 * 4 * ngf * 8)
relu = relay.nn.relu(dense)
# 4 x 4
reshape = relay.reshape(relu, newshape=(-1, ngf * 8, 4, 4))
# 8 x 8
dc8 = deconv2d_bn_relu(reshape,
ishape=(ngf * 8, 4, 4),
oshape=(ngf * 4, 8, 8),
kshape=(4, 4),
prefix="g2")
# 16x16
dc16 = deconv2d_bn_relu(dc8,
ishape=(ngf * 4, 8, 8),
oshape=(ngf * 2, 16, 16),
kshape=(4, 4),
prefix="g3")
# 32x32
dc32 = deconv2d_bn_relu(dc16,
ishape=(ngf * 2, 16, 16),
oshape=(ngf, 32, 32),
kshape=(4, 4),
prefix="g4")
# 64x64
dc64 = deconv2d(dc32,
ishape=(ngf, 32, 32),
oshape=oshape[-3:],
kshape=(4, 4),
name="g5_deconv")
tanh = relay.tanh(dc64)
args = relay.analysis.free_vars(tanh)
return relay.Function(args, tanh)
def tvm_lenet(num_classes=10,
data_shape=(1, 1, 32, 32),
dtype='float32',
alpha=1.0,
is_shallow=False):
from tvm import relay
from tvm.relay.testing import layers
"""Function to construct a Lenet"""
data = relay.var("data", shape=data_shape, dtype=dtype)
conv1 = layers.conv2d(data=data,
channels=6,
kernel_size=(5, 5),
name='conv1')
conv1 = relay.tanh(conv1)
pool2 = relay.nn.avg_pool2d(conv1, pool_size=(2, 2), strides=(2, 2))
conv3 = layers.conv2d(data=pool2,
channels=16,
kernel_size=(5, 5),
name='conv3')
conv3 = relay.tanh(conv3)
pool4 = relay.nn.avg_pool2d(conv3, pool_size=(2, 2), strides=(2, 2))
conv5 = layers.conv2d(data=pool4,
channels=120,
kernel_size=(5, 5),
name='conv5')
conv5 = relay.tanh(conv5)
# Temp
flattened6 = relay.reshape(conv5, (1, 120))
# flattened6 = relay.nn.batch_flatten(conv5)
fcw7 = relay.var('fc7_weight', shape=(120, 84))
fcw7 = relay.transpose(fcw7)
fc7 = relay.nn.dense(data=flattened6, weight=fcw7, units=84)
fc7 = relay.tanh(fc7)
fcw8 = relay.var('fc6_weight', shape=(84, 10))
fcw8 = relay.transpose(fcw8)
fc8 = relay.nn.dense(data=fc7, weight=fcw8, units=10)
softmax = relay.nn.softmax(data=fc8)
fn = relay.Function(relay.analysis.free_vars(softmax), softmax)
return fn
def test_batch_matmul_rewrite():
data = relay.var("data", shape=(1, 4, 16, 16))
data2 = relay.sigmoid(relay.var("data", shape=(4, 16, 64)))
out = relay.nn.conv2d(data, relay.var("weight"), kernel_size=(3, 3), padding=(1, 1), channels=8)
out = relay.nn.batch_flatten(out)
out = relay.reshape(out, [1, 32, 64])
out = relay.nn.batch_matmul(out, data2)
qmod = quantize_and_build(out)
def _check_batch_matmul(node):
if isinstance(node, Call):
if node.op.name in ["nn.batch_matmul", "nn.conv2d"]:
assert node.checked_type.dtype == "int32"
elif node.op.name == "nn.batch_flatten":
assert node.checked_type.dtype == "int8"
# check if batch_matmul is quantized
relay.analysis.post_order_visit(qmod["main"], _check_batch_matmul)
def verify_any_reshape(x_shape, newshape, x_np_shape, out_shape, variable_newshape=False):
x = relay.var('x', shape=x_shape, dtype="float32")
relu_x = relay.nn.relu(x)
data = np.random.uniform(size=x_np_shape).astype('float32')
params = [x]
args = [data]
if variable_newshape:
newshape_var = relay.var('newshape', shape=(len(newshape),), dtype='int64')
params.append(newshape_var)
args.append(np.array(newshape, dtype='int64'))
newshape = newshape_var
y = relay.reshape(relu_x, newshape=newshape)
mod = tvm.IRModule()
mod["main"] = relay.Function(params, y)
for kind in ["debug", "vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(*args).asnumpy()
assert result.shape == out_shape
tvm.testing.assert_allclose(result.flatten(), data.flatten())
def test_vm_reshape_tensor():
x_np = np.random.uniform(size=(8, 16)).astype("float32")
x = relay.var("x", shape=(8, 16), dtype="float32")
y = relay.reshape(x, [-1, 4, 8])
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y)
with tvm.transform.PassContext(opt_level=3):
exec = relay.vm.compile(mod, "llvm")
assert "reshape_tensor" in exec.bytecode
check_result([x_np], x_np.reshape([4, 4, 8]), mod)
x = relay.var("x", shape=(8, 16), dtype="float32")
y = relay.reshape(x, [16, -1])
y = relay.reverse_reshape(y, [-1, 4, 0])
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y)
with tvm.transform.PassContext(opt_level=3):
exec = relay.vm.compile(mod, "llvm")
assert exec.bytecode.count("reshape_tensor") == 1
check_result([x_np], x_np.reshape([4, 4, 8]), mod)
# reshape with symbolic/any shape
for n in [tvm.tir.Any(), tvm.te.size_var("n")]:
x = relay.var("x", shape=(n, 16), dtype="float32")
y = relay.reshape(x, [-1, 4])
y = relay.reshape(y, [0, 2, -1])
mod = tvm.IRModule()
mod["main"] = relay.Function([x], y)
with tvm.transform.PassContext(opt_level=3):
exec = relay.vm.compile(mod, "llvm")
assert exec.bytecode.count("reshape_tensor") == 1
check_result([x_np], x_np.reshape([32, 2, 2]), mod)
# dyn.reshape
x = relay.var("x", shape=(8, 16), dtype="float32")
y = relay.var("y", shape=(3, ), dtype="int32")
z = relay.reshape(x, [-1, 4, 8])
z = relay.reshape(z, y)
mod = tvm.IRModule()
mod["main"] = relay.Function([x, y], z)
with tvm.transform.PassContext(opt_level=3):
exec = relay.vm.compile(mod, "llvm")
assert exec.bytecode.count("reshape_tensor") == 2
assert "reshape_tensor" in exec.bytecode
y_np = np.array([8, 2, 8]).astype("int32")
check_result([x_np, y_np], x_np.reshape([8, 2, 8]), mod)
def reshape_output(output: tvm.relay.Expr, ifm_input_shape: List[int]) -> tvm.relay.Expr:
"""Reshape the output back to the original dimensionality.
Since the NPU must have the brodcastable tensor as the
second operand, the original shape of the first ifm must
be the output shape.
Parameters
----------
output: tvm.relay.Expr
The output to reshape.
ifm_input_shape: List[int]
The shape of the non-reshaped ifm tensor.
Returns
-------
reshaped_output: tvm.relay.Expr
The reshaped output expression.
"""
if len(ifm_input_shape) == 4:
return output
reshaped_output = relay.reshape(output, ifm_input_shape)
return reshaped_output
def test_run(x_data_list, x_shape, new_shape, should_offload_to_trt):
result_arr = [{} for _ in range(len(x_data_list))]
for use_trt in [True, False]:
x = relay.var("x", shape=x_shape, dtype="float32")
out = relay.reshape(x, new_shape)
f = relay.Function([x], out)
mod = tvm.IRModule()
mod["main"] = f
if use_trt:
mod, _ = tensorrt.partition_for_tensorrt(
mod, params={}, remove_no_mac_subgraphs=False
)
assert are_ops_on_trt(mod, op_list=["reshape"]) == should_offload_to_trt
if not skip_runtime_test():
with relay.build_config(opt_level=3):
relay_exec = relay.create_executor("vm", mod=mod, ctx=tvm.cpu(0), target="llvm")
for i, x_data in enumerate(x_data_list):
result_arr[i][use_trt] = relay_exec.evaluate()(x_data)
if not skip_runtime_test():
for i in range(len(x_data_list)):
assert_result_dict_holds(result_arr[i])
def verify_more_dynamic_broadcast_to(x_shape, out_shape):
rank = len(out_shape)
dtype = "float32"
shape_type = "int64"
reshape_shape = relay.Var(
"shape", relay.ty.TensorType((len(x_shape), ), shape_type))
broadcast_shape = relay.Var("shape",
relay.ty.TensorType((rank, ), shape_type))
x = relay.Var("x", relay.ty.TensorType((np.prod(x_shape), ), dtype))
r = relay.reshape(x, reshape_shape)
z = relay.broadcast_to(r, broadcast_shape)
func = relay.Function([x, reshape_shape, broadcast_shape], z)
x = np.random.uniform(size=np.prod(x_shape)).astype(dtype)
ref_res = np.broadcast_to(np.reshape(x, x_shape), out_shape)
for target, dev in tvm.testing.enabled_targets():
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
op_res = relay.create_executor(
kind, mod=mod, device=dev, target=target).evaluate(func)(
x, np.array(x_shape).astype(shape_type),
np.array(out_shape).astype(shape_type))
tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-5)
def expected(x, conv_weight, out_scale, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight]
if blocking:
squeezed_scale = relay.squeeze(out_scale, axis=[0, 2, 3])
conv_weight = relay.multiply(
conv_weight,
relay.reshape(squeezed_scale, (channels // blocking[1], 1, 1, 1, 1, blocking[1])),
)
else:
squeezed_scale = relay.squeeze(out_scale, axis=[1, 2])
conv_weight = relay.multiply(
conv_weight, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3)
)
y = relay.nn.conv2d(
x,
conv_weight,
channels=channels,
kernel_size=(3, 3),
padding=(1, 1),
data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW",
)
return relay.Function(args, y)
def _execute(self):
self.node_dict = {}
# self.node_dict['1'] = relay.const(np.zeros((1, 128)), dtype='int32')
gelu_a = relay.var('gelu_a', shape=())
gelu_b = relay.var('gelu_b', shape=())
gelu_c = relay.var('gelu_c', shape=())
gelu_d = relay.var('gelu_d', shape=())
gelu_e = relay.var('gelu_e', shape=())
self.node_dict['1'] = relay.var('input.1', shape=(1,128), dtype='int32')
self.node_dict['2'] = relay.var('input.2', shape=(1,128), dtype='int32')
for gnode in self.graph:
name = gnode['name']
op_type = gnode['op_type']
attrs = gnode['attrs']
del attrs['A_shape']
del attrs['O_shape']
inputs = gnode['inputs']
if op_type == 'Const':
arr = np.zeros(attrs['shape'], dtype=np.int32)
y = relay.const(arr, dtype='int32')
elif op_type == 'expand_dims':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.expand_dims(x, attrs['axis'], attrs['num_newaxis'])
elif op_type == 'reshape':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.reshape(x, attrs['newshape'])
elif op_type == 'take':
data = get_input(self.node_dict, self.params, inputs[0])
indices = get_input(self.node_dict, self.params, inputs[1])
y = relay.take(data, indices, axis=attrs['axis'][0], mode=attrs['mode'])
elif op_type == 'one_hot':
x = get_input(self.node_dict, self.params, inputs[0])
cc1 = get_input(self.node_dict, self.params, inputs[1])
cc2 = get_input(self.node_dict, self.params, inputs[2])
y = relay.one_hot(x, cc1, cc2, **attrs)
elif op_type == 'strided_slice':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.strided_slice(x, **attrs)
elif op_type == 'mean':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.mean(x, axis=attrs['axis'],
exclude=attrs['exclude'],
keepdims=attrs['keepdims'])
elif op_type == 'nn.dense':
x = get_input(self.node_dict, self.params, inputs[0])
weight = get_input(self.node_dict, self.params, inputs[1])
y = relay.nn.dense(x, weight, units=attrs['units'][0])
elif op_type == 'add':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.add(x1, x2)
elif op_type == 'subtract':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.subtract(x1, x2)
elif op_type == 'multiply':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.multiply(x1, x2)
elif op_type == 'power':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.power(x1, x2)
elif op_type == 'transpose':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.transpose(x, **attrs)
elif op_type == 'tanh':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.tanh(x)
elif op_type == 'squeeze':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.squeeze(x, **attrs)
elif op_type == 'nn.batch_matmul':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.nn.batch_matmul(x1, x2)
elif op_type == 'nn.softmax':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.nn.softmax(x, **attrs)
elif op_type == 'gelu':
x = get_input(self.node_dict, self.params, inputs[0])
y = x * gelu_a * (gelu_b + relay.tanh(
( gelu_c * (x + gelu_d *
relay.power(x, gelu_e)))))
else:
import pdb; pdb.set_trace()
print( 'not supported op %s ' % op_type)
self.node_dict[name] = y
output_name = self.output_node_ids[0]
output = self.node_dict[output_name]
inputs = relay.analysis.free_vars(output)
# inputs = [self.node_dict['1'], self.node_dict['2']]
func = relay.Function(inputs, output)
mod = tvm.IRModule()
mod['main'] = func
with relay.build_config(opt_level=0):
graph, lib, params = relay.build(mod, 'llvm', params={})
self.m = graph_runtime.create(graph, lib, tvm.cpu())
def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
target = tvm.target.Target.current(allow_none=False)
dispatch_ctx = autotvm.task.DispatchContext.current
new_attrs = {k: attrs[k] for k in attrs.keys()}
# Parse the attributes.
padding = attrs.get_int_tuple("padding")
strides = attrs.get_int_tuple("strides")
dilation = attrs.get_int_tuple("dilation")
data_layout = attrs["data_layout"]
kernel_layout = attrs["kernel_layout"]
data_tensor, kernel_tensor = tinfos
data_dtype = data_tensor.dtype
kernel_dtype = kernel_tensor.dtype
out_dtype = out_type.dtype
if isinstance(dispatch_ctx, autotvm.task.ApplyGraphBest):
cfg = dispatch_ctx.query(target, None)
workload = cfg.workload
else:
impl, outs = relay.backend.compile_engine.select_implementation(
relay.op.get("nn.conv2d"), attrs, tinfos, out_type, target)
workload = autotvm.task.get_workload(outs)
if workload is None:
# The best implementation is not an AutoTVM template.
# It may be from the auto-scheduler
if impl.name.find("winograd") != -1:
if dilation != (1, 1):
logger.warning(
"Does not support weight pre-transform for dilated convolution."
)
return None
assert data_layout == "NHWC" and kernel_layout == "HWIO"
N, H, W, CI = get_const_tuple(data_tensor.shape)
KH, KW, _, CO = get_const_tuple(kernel_tensor.shape)
# Pre-compute weight transformation in winograd
tile_size = 4
# HWIO -> OIHW
kernel_transform = relay.transpose(inputs[1],
axes=[3, 2, 0, 1])
# alpha, alpha, CO, CI
weight = relay.nn.contrib_conv2d_winograd_weight_transform(
kernel_transform, tile_size=tile_size)
new_attrs["tile_size"] = tile_size
new_attrs["channels"] = CO
return relay.nn.contrib_conv2d_winograd_without_weight_transform(
inputs[0], weight, **new_attrs)
return None
cfg = dispatch_ctx.query(target, workload)
topi_tmpl = workload[0]
if topi_tmpl == "conv2d_NCHWc.x86":
# we only convert conv2d_NCHW to conv2d_NCHWc for x86
if data_layout == "NCHW" and kernel_layout == "OIHW":
if cfg.is_fallback:
_get_default_config(
cfg,
data_tensor,
kernel_tensor,
strides,
padding,
dilation,
out_dtype,
False,
data_layout,
)
batch_size, in_channel, height, width = get_const_tuple(
data_tensor.shape)
out_channel, _, kh, kw = get_const_tuple(kernel_tensor.shape)
ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1]
# update new attrs
new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
# (oc, ic, h, w) -> (OC, IC, h, w, ic, oc)
new_attrs["kernel_layout"] = "OIHW%di%do" % (ic_bn, oc_bn)
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
# Store altered operator's config
new_data = te.placeholder(
(batch_size, in_channel // ic_bn, height, width, ic_bn),
dtype=data_dtype)
new_kernel = te.placeholder(
(out_channel // oc_bn, in_channel // ic_bn, kh, kw, ic_bn,
oc_bn),
dtype=kernel_tensor.dtype,
)
new_workload = autotvm.task.args_to_workload(
[
new_data,
new_kernel,
strides,
padding,
dilation,
new_attrs["data_layout"],
new_attrs["out_layout"],
out_dtype,
],
topi_tmpl,
)
dispatch_ctx.update(target, new_workload, cfg)
else:
assert _NCHWc_matcher.match(data_layout)
assert _OIHWio_matcher.match(kernel_layout)
return relay.nn.contrib_conv2d_nchwc(*inputs, **new_attrs)
if topi_tmpl == "conv2d_NCHWc_int8.x86":
# TODO(@icemelon9, @anijain2305): Need to support data layout NHWC with kernel layout HWIO
assert data_layout == "NCHW" and kernel_layout == "OIHW"
if cfg.is_fallback:
_get_default_config_int8(
cfg,
data_tensor,
kernel_tensor,
strides,
padding,
dilation,
out_dtype,
False,
data_layout,
)
batch_size, in_channel, height, width = get_const_tuple(
data_tensor.shape)
out_channel, channel_multiplier, kh, kw = get_const_tuple(
kernel_tensor.shape)
ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1]
n_elems = 4
# convert kernel data layout from 4D to 7D
data_expr, kernel_expr = inputs
kernel_IHWO = relay.transpose(kernel_expr, axes=(1, 2, 3, 0))
kernel_IHWOo = relay.reshape(
kernel_IHWO, (in_channel, kh, kw, out_channel // oc_bn, oc_bn))
kernel_OHWoI = relay.transpose(kernel_IHWOo, axes=(3, 1, 2, 4, 0))
kernel_OHWoIi = relay.reshape(
kernel_OHWoI,
(out_channel // oc_bn, kh, kw, oc_bn, in_channel // ic_bn, ic_bn))
kernel_OHWoIie = relay.reshape(
kernel_OHWoIi,
(out_channel // oc_bn, kh, kw, oc_bn, in_channel // ic_bn,
ic_bn // n_elems, n_elems),
)
kernel_OIHWioe = relay.transpose(kernel_OHWoIie,
axes=(0, 4, 1, 2, 5, 3, 6))
# update new attrs
new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
# Store altered operator's config.
new_data = te.placeholder(
(batch_size, in_channel // ic_bn, height, width, ic_bn),
dtype=data_dtype)
new_kernel = te.placeholder(
(out_channel // oc_bn, in_channel // ic_bn, kh, kw,
ic_bn // n_elems, oc_bn, n_elems),
dtype=kernel_dtype,
)
new_workload = autotvm.task.args_to_workload(
[
new_data,
new_kernel,
strides,
padding,
dilation,
new_attrs["data_layout"],
new_attrs["out_layout"],
out_dtype,
],
topi_tmpl,
)
dispatch_ctx.update(target, new_workload, cfg)
return relay.nn.contrib_conv2d_nchwc(data_expr, kernel_OIHWioe,
**new_attrs)
if topi_tmpl == "depthwise_conv2d_NCHWc.x86":
if data_layout == "NCHW" and kernel_layout == "OIHW":
if cfg.is_fallback:
_get_default_config(
cfg,
data_tensor,
kernel_tensor,
strides,
padding,
dilation,
out_dtype,
True,
data_layout,
)
batch_size, in_channel, height, width = get_const_tuple(
data_tensor.shape)
out_channel, channel_multiplier, kh, kw = get_const_tuple(
kernel_tensor.shape)
ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1]
assert channel_multiplier == 1
# update new attrs
new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
new_attrs["kernel_layout"] = "OIHW1i%do" % oc_bn
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
# Store altered operator's config.
new_data = te.placeholder(
(batch_size, in_channel // ic_bn, height, width, ic_bn),
dtype=data_dtype)
new_kernel = te.placeholder(
(out_channel // oc_bn, 1, kh, kw, 1, oc_bn),
dtype=kernel_dtype)
new_workload = autotvm.task.args_to_workload(
[
new_data,
new_kernel,
strides,
padding,
dilation,
new_attrs["data_layout"],
new_attrs["out_layout"],
out_dtype,
],
topi_tmpl,
)
dispatch_ctx.update(target, new_workload, cfg)
else:
assert _NCHWc_matcher.match(data_layout)
assert _OIHWio_matcher.match(kernel_layout)
return relay.nn.contrib_depthwise_conv2d_nchwc(*inputs, **new_attrs)
return None
def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
target = tvm.target.Target.current(allow_none=False)
dispatch_ctx = autotvm.task.DispatchContext.current
new_attrs = {k: attrs[k] for k in attrs.keys()}
strides = attrs.get_int_tuple("strides")
padding = attrs.get_int_tuple("padding")
dilation = attrs.get_int_tuple("dilation")
data_layout = attrs["data_layout"]
kernel_layout = attrs["kernel_layout"]
data, kernel = tinfos
out_dtype = out_type.dtype
impl, outs = relay.backend.compile_engine.select_implementation(
relay.op.get("nn.conv2d"), attrs, tinfos, out_type, target
)
workload = autotvm.task.get_workload(outs)
if workload is None:
# The best implementation is not an AutoTVM template.
# It may be from the auto-scheduler
if impl.name.find("winograd") != -1:
if dilation != (1, 1):
logger.warning("Does not support weight pre-transform for dilated convolution.")
return None
assert data_layout == "NHWC" and kernel_layout == "HWIO"
N, H, W, CI = get_const_tuple(data.shape)
KH, KW, _, CO = get_const_tuple(kernel.shape)
# Pre-compute weight transformation in winograd
tile_size = _pick_tile_size(tinfos[0], tinfos[1], layout="NHWC")
# HWIO -> OIHW
kernel_transform = relay.transpose(inputs[1], axes=[3, 2, 0, 1])
# alpha, alpha, CO, CI
weight = relay.nn.contrib_conv2d_winograd_weight_transform(
kernel_transform, tile_size=tile_size
)
new_attrs["tile_size"] = tile_size
new_attrs["channels"] = CO
return relay.nn.contrib_conv2d_winograd_without_weight_transform(
inputs[0], weight, **new_attrs
)
return None
cfg = dispatch_ctx.query(target, workload)
if cfg.is_fallback: # if is fallback, clear query cache and return None
autotvm.task.clear_fallback_cache(target, workload)
return None
topi_tmpl = workload[0]
idxd = tvm.tir.indexdiv
if topi_tmpl == "conv2d_nchw_spatial_pack.mali":
assert data_layout == "NCHW" and kernel_layout == "OIHW"
N, CI, H, W = get_const_tuple(data.shape)
CO, _, KH, KW = get_const_tuple(kernel.shape)
VC = cfg["tile_co"].size[-1]
new_attrs["kernel_layout"] = "OIHW%do" % VC
new_data = data
new_kernel = te.placeholder((idxd(CO, VC), CI, KH, KW, VC), dtype=kernel.dtype)
new_workload = autotvm.task.args_to_workload(
[new_data, new_kernel, strides, padding, dilation, out_dtype],
"conv2d_nchw_spatial_pack.mali",
)
dispatch_ctx.update(target, new_workload, cfg)
return relay.nn.conv2d(*inputs, **new_attrs)
elif topi_tmpl == "conv2d_nchw_winograd.mali":
assert data_layout == "NCHW" and kernel_layout == "OIHW"
N, CI, H, W = get_const_tuple(data.shape)
CO, _, KH, KW = get_const_tuple(kernel.shape)
tile_size = _pick_tile_size(data, kernel)
VC = cfg["tile_bna"].val
weight_expr = inputs[1]
weight_expr = relay.nn.contrib_conv2d_winograd_weight_transform(
weight_expr, tile_size=tile_size
)
weight_expr = relay.reshape(
weight_expr, newshape=(KH + tile_size - 1, KW + tile_size - 1, idxd(CO, VC), VC, CI)
)
weight_expr = relay.transpose(weight_expr, axes=[0, 1, 2, 4, 3])
new_attrs["tile_size"] = tile_size
new_data = data
new_kernel = te.placeholder(
(KH + tile_size - 1, KW + tile_size - 1, idxd(CO, VC), CI, VC), kernel.dtype
)
new_workload = autotvm.task.args_to_workload(
[new_data, new_kernel, strides, padding, dilation, out_dtype],
"conv2d_nchw_winograd.mali",
)
dispatch_ctx.update(target, new_workload, cfg)
return relay.nn.contrib_conv2d_winograd_without_weight_transform(
inputs[0], weight_expr, **new_attrs
)
else:
return None
def relay_conv2d_weight_grad(c, data, wsize, dout, stride, pad, dil, groups):
# This implementation should match the one in pytorch backend
# (myia.compile.backends.pytorch_conv_grad.conv2d_weight)
assert wsize.is_constant(tuple)
assert stride.is_constant(tuple)
assert pad.is_constant(tuple)
assert dil.is_constant(tuple)
assert groups.is_constant(int)
batch, in_channel, in_h, in_w = data.abstract.xshape()
out_channel, _, filter_h, filter_w = wsize.value
grad_sh0, grad_sh1, grad_h, grad_w = dout.abstract.xshape()
pad_h, pad_w = pad.value
data = c.ref(data)
dout = c.ref(dout)
fpad_h = pad_h * 2
fpad_w = pad_w * 2
fpad_top = (pad_h + 1) // 2
fpad_left = (pad_w + 1) // 2
fpad_bottom = fpad_h - fpad_top
fpad_right = fpad_w - fpad_left
padded_weight_grad_h = (in_h - (grad_h - 1) * stride.value[0] - 1 +
fpad_top + fpad_bottom) // dil.value[0] + 1
padded_weight_grad_w = (in_w - (grad_w - 1) * stride.value[1] - 1 +
fpad_left + fpad_right) // dil.value[1] + 1
dout = relay.tile(dout, [1, in_channel // groups.value, 1, 1])
dout = relay.reshape(dout, [-1, 1, 0, 0])
data = relay.reshape(data, [1, -1, 0, 0])
d = relay.nn.conv2d(
data,
dout,
strides=dil.value,
padding=pad.value,
dilation=stride.value,
groups=batch * in_channel,
)
conv_sh1 = grad_sh0 * grad_sh1 * (in_channel // groups.value)
d = relay.reshape(
d,
[batch, conv_sh1 // batch, padded_weight_grad_h, padded_weight_grad_w],
)
d = relay.sum(d, axis=0)
if groups.value > 1:
d = relay.reshape(
d,
[
grad_sh1,
in_channel // groups.value,
padded_weight_grad_h,
padded_weight_grad_w,
],
)
else:
d = relay.reshape(
d,
[
in_channel // groups.value,
grad_sh1,
padded_weight_grad_h,
padded_weight_grad_w,
],
)
d = relay.transpose(d, [1, 0, 2, 3])
if padded_weight_grad_h > filter_h or padded_weight_grad_w > filter_w:
d = relay.strided_slice(d,
begin=[0, 0, 0, 0],
end=[None, None, filter_h, filter_w])
return d
def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr,
node_map: tvm.ir.container.Map) -> tvm.relay.Expr:
params = ethosu_patterns.MeanParams(post.op.body)
params.ifm.tensor = post.args[0]
ifm_shape = params.ifm.shape
ofm_shape = params.ofm.shape
lut = relay.const([], "int8")
axis = params.axis
reduced_op = params.ifm.tensor
# Enforce 4d input
if len(ifm_shape) < 4:
axis = [x + 1 for x in axis]
if len(ifm_shape) == 3:
ifm_shape = [1, params.height, params.width, ifm_shape[2]]
else:
ifm_shape = [1, params.height, params.width, 1]
reduced_op = relay.reshape(reduced_op, ifm_shape)
filter_height = ifm_shape[1] if 1 in axis else 1
filter_width = ifm_shape[2] if 2 in axis else 1
in_channels = out_channels = ifm_shape[-1]
# If the height is greater than max kernel height, reshape the input
# from [filter_height, filter_width] to [1, (filter_height*filter_width)]
# only in the case the axis is [1, 2].
if axis == [1, 2] and filter_height > 64:
ifm_shape = (ifm_shape[0], 1, filter_height * filter_width,
in_channels)
filter_width = filter_height * filter_width
filter_height = 1
reduced_op = relay.reshape(reduced_op, ifm_shape)
if axis == [1, 2] and params.keepdims:
weight_scale = 1
weight_values = np.ones(
[out_channels, filter_height, filter_width, in_channels])
scale_bias = vela_api.pack_biases(
biases=np.zeros(ifm_shape[-1]),
ifm_scale=params.ifm.q_params.scale_f32,
ifm_dtype=np.dtype(params.ifm.dtype),
weight_scales=np.array([weight_scale], dtype=np.float),
ofm_scale=params.ofm.q_params.scale_f32,
is_activation_tanh_or_sigmoid=False,
)
reduced_op = ethosu_ops.ethosu_depthwise_conv2d(
ifm=reduced_op,
weight=relay.const(weight_values, params.ifm.dtype),
scale_bias=relay.const(scale_bias, "uint8"),
lut=lut,
ifm_scale=float(params.ifm.q_params.scale_f32),
ifm_zero_point=int(params.ifm.q_params.zero_point),
weight_zero_point=0,
ofm_scale=float(params.ofm.q_params.scale_f32),
ofm_zero_point=int(params.ofm.q_params.zero_point),
kernel_shape=(filter_height, filter_width),
ofm_channels=out_channels,
ofm_dtype="int16",
)
n = int(filter_height * filter_width)
eps = 1 / (256 * (n + 1)) if n % 2 == 0 else 0
scalar_tensor = relay.const(np.ones([1, 1, 1, 1], dtype="int16"),
dtype="int16")
reduced_op = ethosu_ops.ethosu_binary_elementwise(
ifm=reduced_op,
ifm2=scalar_tensor,
lut=lut,
operator_type="MUL",
ifm_scale=float(params.ofm.q_params.scale_f32),
ifm_zero_point=int(params.ofm.q_params.zero_point),
ifm2_scale=1 / (n - eps),
ifm2_zero_point=0,
ofm_scale=float(params.ofm.q_params.scale_f32),
ofm_zero_point=int(params.ofm.q_params.zero_point),
ifm_channels=out_channels,
ifm2_channels=out_channels,
reversed_operands=False,
ofm_dtype="int8",
rounding_mode="NATURAL",
)
elif (params.ifm.q_params.scale_f32 == params.ofm.q_params.scale_f32
and params.ifm.q_params.zero_point
== params.ofm.q_params.zero_point):
reduced_op = ethosu_ops.ethosu_pooling(
ifm=reduced_op,
lut=lut,
pooling_type="AVG",
ifm_scale=float(params.ifm.q_params.scale_f32),
ifm_zero_point=0,
ofm_scale=float(params.ofm.q_params.scale_f32),
ofm_zero_point=0,
pool_shape=(filter_height, filter_width),
ofm_channels=out_channels,
rounding_mode="TRUNCATE",
)
else:
weight_scale = 1 / (filter_height * filter_width)
weight_values = np.ones(
[out_channels, filter_height, filter_width, in_channels])
bias = -1 * int(
params.ifm.q_params.zero_point) * filter_height * filter_width
scale_bias = vela_api.pack_biases(
biases=np.ones([ifm_shape[-1]]) * bias,
ifm_scale=params.ifm.q_params.scale_f32,
ifm_dtype=np.dtype(params.ifm.dtype),
weight_scales=np.array([weight_scale], dtype=np.float),
ofm_scale=params.ofm.q_params.scale_f32,
is_activation_tanh_or_sigmoid=False,
)
reduced_op = ethosu_ops.ethosu_depthwise_conv2d(
ifm=reduced_op,
weight=relay.const(weight_values, params.ifm.dtype),
scale_bias=relay.const(scale_bias, "uint8"),
lut=lut,
ifm_scale=float(params.ifm.q_params.scale_f32),
ifm_zero_point=0,
weight_zero_point=0,
ofm_scale=float(params.ofm.q_params.scale_f32),
ofm_zero_point=int(params.ofm.q_params.zero_point),
kernel_shape=(filter_height, filter_width),
ofm_channels=out_channels,
rounding_mode="NATURAL",
)
# Reshape to original ofm shape
if len(ofm_shape) < 4:
reduced_op = relay.reshape(reduced_op, ofm_shape)
return reduced_op
def expected():
x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32")
w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32")
y = relay.nn.conv2d(x, w, padding=(1, 1))
y = relay.reshape(y, newshape=(32, 16, 16))
return relay.Function([x, w], y)
def get_model(input_shape, var_names):
"""Return a model and any parameters it may have."""
a = relay.var(next(var_names), shape=input_shape, dtype="float32")
out = relay.reshape(a, (1, 1, 1000))
out = relay.reshape(out, (1, 1000))
return out
def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
target = tvm.target.Target.current(allow_none=False)
dispatch_ctx = autotvm.task.DispatchContext.current
_, outs = relay.backend.compile_engine.select_implementation(
relay.op.get("nn.conv2d"), attrs, tinfos, out_type, target)
workload = autotvm.task.get_workload(outs)
if workload is None:
# The best implementation is not an AutoTVM template,
# we then assume it's not necessary to alter this op.
return None
cfg = dispatch_ctx.query(target, workload)
if cfg.is_fallback: # if is fallback, clear query cache and return None
autotvm.task.clear_fallback_cache(target, workload)
return None
topi_tmpl = workload[0]
new_attrs = {k: attrs[k] for k in attrs.keys()}
strides = attrs.get_int_tuple("strides")
padding = attrs.get_int_tuple("padding")
dilation = attrs.get_int_tuple("dilation")
data_layout = attrs["data_layout"]
kernel_layout = attrs["kernel_layout"]
data, kernel = tinfos
out_dtype = out_type.dtype
idxd = tvm.tir.indexdiv
if topi_tmpl == "conv2d_nchw_spatial_pack.bifrost":
assert data_layout == "NCHW" and kernel_layout == "OIHW"
N, CI, H, W = get_const_tuple(data.shape)
CO, _, KH, KW = get_const_tuple(kernel.shape)
VC = cfg["tile_co"].size[-1]
new_attrs["kernel_layout"] = "OIHW%do" % VC
new_data = data
new_kernel = te.placeholder((idxd(CO, VC), CI, KH, KW, VC),
dtype=kernel.dtype)
new_workload = autotvm.task.args_to_workload(
[new_data, new_kernel, strides, padding, dilation, out_dtype],
"conv2d_nchw_spatial_pack.bifrost",
)
dispatch_ctx.update(target, new_workload, cfg)
return relay.nn.conv2d(*inputs, **new_attrs)
if topi_tmpl == "conv2d_nchw_winograd.bifrost":
assert data_layout == "NCHW" and kernel_layout == "OIHW"
N, CI, H, W = get_const_tuple(data.shape)
CO, _, KH, KW = get_const_tuple(kernel.shape)
tile_size = 2
weight_expr = inputs[1]
weight_expr = relay.nn.contrib_conv2d_winograd_weight_transform(
weight_expr, tile_size=tile_size)
weight_expr = relay.reshape(weight_expr,
newshape=(KH + tile_size - 1,
KW + tile_size - 1, CO, CI))
new_attrs["tile_size"] = tile_size
new_data = data
new_kernel = te.placeholder(
(KH + tile_size - 1, KW + tile_size - 1, CO, CI), kernel.dtype)
new_workload = autotvm.task.args_to_workload(
[new_data, new_kernel, strides, padding, dilation, out_dtype],
"conv2d_nchw_winograd.bifrost",
)
dispatch_ctx.update(target, new_workload, cfg)
return relay.nn.contrib_conv2d_winograd_without_weight_transform(
inputs[0], weight_expr, **new_attrs)
return None
conv7_weight = relay.var("conv7_weight")
conv8_weight = relay.var("conv8_weight")
bias1 = relay.var("bias", relay.TensorType((96,), "float32"))
bias2 = relay.var("2bias")
bias3 = relay.var("3bias")
bias4 = relay.var("4bias")
bias5 = relay.var("5bias")
bias6 = relay.var("6bias")
bias7 = relay.var("7bias")
bias8 = relay.var("8bias")
simple_net = relay.nn.conv2d(data=data, weight=conv1_weight, kernel_size=(11,11), channels=96, strides=(4,4),padding=(0, 0))
simple_net = relay.nn.bias_add(data=simple_net, bias=bias1)
simple_net = relay.reshape(simple_net, (batch_size,96,55,55))
simple_net = relay.nn.relu(data=simple_net)
simple_net = relay.reshape(simple_net, (batch_size,96,55,55))
simple_net = relay.nn.lrn(data=simple_net,size=5, axis=1, bias=2, alpha=.00001, beta=0.75)
simple_net = relay.reshape(simple_net, (batch_size,96,55,55))
simple_net = relay.nn.max_pool2d(data=simple_net, pool_size=(3, 3), strides=(2, 2), padding=(0, 0), layout="NCHW")
simple_net = relay.nn.conv2d(data=simple_net, weight=conv2_weight, kernel_size=(5,5), channels=256, strides=(1,1),padding=(2, 2))
simple_net = relay.nn.bias_add(data=simple_net, bias=bias2)
simple_net = relay.reshape(simple_net, (batch_size,256,27,27))
simple_net = relay.nn.relu(data=simple_net)
simple_net = relay.reshape(simple_net, (batch_size,256,27,27))
simple_net = relay.nn.lrn(data=simple_net,size=5, axis=1, bias=2, alpha=.00001, beta=0.75)
simple_net = relay.reshape(simple_net, (batch_size,256,27,27))
simple_net = relay.nn.max_pool2d(data=simple_net, pool_size=(3, 3), strides=(2, 2), padding=(0, 0), layout="NCHW")
def relay_conv_transpose2d(c, input, weight, stride, padding, output_padding,
groups, dilation):
"""Implement conv2d_transpose using 10 relay calls including conv2d.
Support all values for groups, dilation, strides, padding and
output padding.
Based on Theano implementation (2020/04/14):
https://github.com/Theano/Theano/blob/master/theano/tensor/nnet/abstract_conv.py#L2927
Need implementation of operation relay.nn.dilate
in TVM relay backend
"""
assert stride.is_constant(tuple)
assert padding.is_constant(tuple)
assert output_padding.is_constant(tuple)
assert dilation.is_constant(tuple)
assert groups.is_constant(int)
data_shape = input.abstract.xshape()
kern_shape = weight.abstract.xshape()
h_in, w_in = data_shape[2:]
filter_h, filter_w = kern_shape[2:]
strides = stride.value
padding = padding.value
dilation = dilation.value
output_padding = output_padding.value
groups = groups.value
data = c.ref(input)
weight = c.ref(weight)
h_out = ((h_in - 1) * strides[0] - 2 * padding[0] + dilation[0] *
(filter_h - 1) + output_padding[0] + 1)
w_out = ((w_in - 1) * strides[1] - 2 * padding[1] + dilation[1] *
(filter_w - 1) + output_padding[1] + 1)
data_dilated = relay.nn.dilate(data, (1, 1) + strides)
data_padded = relay.nn.pad(
data_dilated,
(
(0, 0),
(0, 0),
(0, output_padding[0]),
(0, output_padding[1]),
),
)
# Pre-process kernel,
# from (m0, m1, m2, m3) to (m1 * g, m0 // g, m2, m3).
mshp0 = kern_shape[0] // groups
c_out = kern_shape[1] * groups
kern = relay.reshape(weight, (groups, mshp0) + kern_shape[1:])
# => (g, m0 // g, m1, m2, m3)
kern = relay.op.transpose(kern, axes=(1, 0, 2, 3, 4))
# => (m0 // g, g, m1, m2, m3)
kern = relay.reshape(kern, (mshp0, c_out, kern_shape[-2], kern_shape[-1]))
# => (m0 // g, m1 * g, m2, m3)
kern = relay.op.transpose(kern, (1, 0, 2, 3))
# => (m1 * g, m0 // g, m2, m3)
# Kernel 2 latest dimensions must be flipped
kern = relay.op.transform.reverse(kern, 2)
kern = relay.op.transform.reverse(kern, 3)
# End pre-processing kernel.
img = relay.nn.conv2d(
data_padded,
kern,
groups=groups,
channels=c_out,
padding=[(kern_shape[2 + i] - 1) * dilation[i] for i in range(2)],
dilation=dilation,
)
if any(p != 0 for p in padding):
img = relay.op.transform.strided_slice(
data=img,
begin=[0, 0, padding[0], padding[1]],
end=[None, None, h_out + padding[0], w_out + padding[1]],
)
return img
def _get_model(input_shape, output_shape, dtype, var_names):
"""Return a model and any parameters it may have."""
a = relay.var(next(var_names), shape=input_shape, dtype=dtype)
reshape = relay.reshape(a, output_shape)
return reshape
