def verify(dshape, begin, end, strides, output, slice_mode="end", test_ref=True, dtype="int32"):
x = relay.var("x", relay.TensorType(dshape, "float32"))
ndim = len(dshape)
begin = begin if begin else [0] * ndim
end = end if end else list(dshape)
if strides:
if len(strides) == 1:
strides = strides * ndim
else:
strides = [1] * ndim
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(x_data, begin, end, strides, slice_mode)
data = [x_data, np.array(begin), np.array(end)]
begin = relay.const(begin, dtype=dtype)
end = relay.const(end, dtype=dtype)
if strides:
data.append(np.array(strides))
strides = relay.const(strides, dtype=dtype)
z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode)
else:
z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode)
func = relay.Function([x], z)
func = run_infer_type(func)
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType())
assert isinstance(func2.body, relay.Call)
assert func2.body.op == relay.op.get("strided_slice")
verify_func(func2, [x_data], ref_res)
def expected(x, w1, w2, b1, b2, j, bias_shape1, bias_shape2):
args = [x, w1, w2, b1, b2]
w_stacked = relay.concatenate((w1, w2), axis=0)
y = relay.nn.dense(x, w_stacked, units=3 * j)
n_out_dims = max(len(bias_shape1), 2)
if len(bias_shape1) == 0:
b1 = relay.repeat(relay.expand_dims(b1, -1), j, 0)
elif bias_shape1[-1] == 1:
b1 = relay.repeat(b1, j, len(bias_shape1) - 1)
if len(bias_shape2) == 0:
b2 = relay.repeat(relay.expand_dims(b2, -1), 2 * j, 0)
elif bias_shape2[-1] == 1:
b2 = relay.repeat(b2, 2 * j, len(bias_shape2) - 1)
b = relay.concatenate((b1, b2), axis=max(0, len(bias_shape1) - 1))
y = relay.add(y, b)
begin = [0 for _ in range(n_out_dims - 1)]
end = [-1 for _ in range(n_out_dims - 1)]
strides = [1 for _ in range(n_out_dims)]
y1 = relay.strided_slice(y,
begin=begin + [0],
end=end + [j],
strides=strides,
slice_mode="size")
y2 = relay.strided_slice(y,
begin=begin + [j],
end=end + [2 * j],
strides=strides,
slice_mode="size")
return relay.Function(args, relay.Tuple((y1, y2)))
def expected(x, w1, w2, w3, w4, channels1, channels2, channels3,
channels4):
# use a fixed order of args so alpha equal check can pass
args = [x, w1, w2, w3, w4]
w = relay.concatenate((w1, w2, w4), axis=0)
y = relay.nn.conv2d(x, w, channels=channels1 + channels2 + channels4)
y1 = relay.strided_slice(y,
begin=relay.const([0, 0], "int64"),
end=relay.const([-1, channels1], "int64"),
strides=relay.const([1, 1], 'int64'),
slice_mode="size")
y2 = relay.strided_slice(y,
begin=relay.const([0, channels1], "int64"),
end=relay.const([-1, channels2], "int64"),
strides=relay.const([1, 1], 'int64'),
slice_mode="size")
y3 = relay.nn.conv2d(x, w3)
y4 = relay.strided_slice(y,
begin=relay.const([0, channels1 + channels2],
"int64"),
end=relay.const([-1, channels4], "int64"),
strides=relay.const([1, 1], 'int64'),
slice_mode="size")
y5 = relay.nn.max_pool2d(x)
y = relay.Tuple((y1, y2, y3, y4, y5))
return relay.Function(args, y)
def expected(x, w1, w2, b1, b2, scale1, scale2, newshape1, newshape2, j):
args = [x, w1, w2, b1, b2, scale1, scale2]
w_stacked = relay.concatenate((w1, w2), axis=0)
y = relay.nn.dense(x, w_stacked, units=3 * j)
b = relay.concatenate((b1, b2), axis=0)
y = relay.add(y, b)
scale1 = relay.repeat(scale1, j, 0)
scale2 = relay.repeat(scale2, 2 * j, 0)
scale = relay.concatenate((scale1, scale2), axis=0)
y = relay.multiply(y, scale)
strides = [1, 1]
y1 = relay.strided_slice(y,
begin=[0, 0],
end=[-1, j],
strides=strides,
slice_mode="size")
y2 = relay.strided_slice(y,
begin=[0, j],
end=[-1, 2 * j],
strides=strides,
slice_mode="size")
y1 = relay.reshape(y1, newshape=newshape1)
y2 = relay.reshape(y2, newshape=newshape2)
y = relay.Tuple((y1, y2))
return relay.Function(args, y)
def verify(dshape,
begin,
end,
strides,
output,
slice_mode="end",
test_ref=True,
dtype="int32"):
x = relay.var("x", relay.TensorType(dshape, "float32"))
ndim = len(dshape)
begin = begin if begin else [0] * ndim
end = end if end else list(dshape)
if strides:
if len(strides) == 1:
strides = strides * ndim
else:
strides = [1] * ndim
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(
x_data, begin, end, strides, slice_mode)
data = [
x_data,
np.array(begin, dtype=dtype),
np.array(end, dtype=dtype)
]
begin = relay.var("begin", shape=[len(begin)], dtype=dtype)
end = relay.var("end", shape=[len(end)], dtype=dtype)
inputs = [x, begin, end]
if strides:
data.append(np.array(strides, dtype=dtype))
strides = relay.var("strides", shape=[len(strides)], dtype=dtype)
inputs.append(strides)
z = relay.strided_slice(x,
begin=begin,
end=end,
strides=strides,
slice_mode=slice_mode)
else:
z = relay.strided_slice(x,
begin=begin,
end=end,
slice_mode=slice_mode)
func = relay.Function(inputs, z)
func = run_infer_type(func)
text = func.astext()
if not test_ref:
return
for target, dev in tvm.testing.enabled_targets():
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor("vm",
mod=mod,
device=dev,
target=target)
op_res = intrp.evaluate()(*data)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
def verify(dshape, begin, end, strides, output, slice_mode="end", test_ref=True, dtype="int32"):
ndim = len(dshape)
begin = begin if begin else [0] * ndim
end = end if end else list(dshape)
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(x_data, begin, end, strides, slice_mode)
x = relay.var("x", relay.TensorType((relay.Any(),) * ndim, "float32"))
if strides:
z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode)
else:
z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode)
func = relay.Function([x], z)
func = run_infer_type(func)
text = func.astext()
assert "begin=" in text
assert "end=" in text
if not test_ref:
return
for target, ctx in tvm.testing.enabled_targets():
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor("vm", mod=mod, ctx=ctx, target=target)
op_res = intrp.evaluate()(x_data)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
def expected(x, w1, w2, w3, b1, b2, b3):
# use a fixed order of args so alpha equal check can pass
s1 = w1.type_annotation.shape[1]
s2 = w2.type_annotation.shape[1]
s3 = w3.type_annotation.shape[1]
args = [x, w1, w2, w3, b1, b2, b3]
w = relay.concatenate((w1, w2, w3), axis=1)
b = relay.concatenate((b1, b2, b3), axis=-1)
y = relay.nn.batch_matmul(x, w)
y = relay.add(y, b)
y1 = relay.strided_slice(y,
begin=relay.const([0, 0, 0], "int64"),
end=relay.const([-1, -1, s1], "int64"),
strides=relay.const([1, 1, 1], 'int64'),
slice_mode="size")
y2 = relay.strided_slice(y,
begin=relay.const([0, 0, s1], "int64"),
end=relay.const([-1, -1, s2], "int64"),
strides=relay.const([1, 1, 1], 'int64'),
slice_mode="size")
y3 = relay.strided_slice(y,
begin=relay.const([0, 0, s1+s2], "int64"),
end=relay.const([-1, -1, s3], "int64"),
strides=relay.const([1, 1, 1], 'int64'),
slice_mode="size")
y = relay.Tuple((y1, y2, y3))
return relay.Function(args, y)
def verify(
dshape,
begin,
end,
strides,
output,
axes=None,
slice_mode="end",
test_ref=True,
dtype="int32",
):
x = relay.var("x", relay.TensorType(dshape, "float32"))
ndim = len(dshape)
begin = begin if begin else [0] * ndim
end = end if end else list(dshape)
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(
x_data,
begin,
end,
strides,
slice_mode,
axes=axes,
)
if strides:
z = relay.strided_slice(x,
begin=begin,
end=end,
strides=strides,
axes=axes,
slice_mode=slice_mode)
else:
z = relay.strided_slice(x,
begin=begin,
end=end,
axes=axes,
slice_mode=slice_mode)
func = relay.Function([x], z)
func = run_infer_type(func)
text = func.astext()
assert "begin=" in text
assert "end=" in text
if output:
assert func.body.checked_type == relay.ty.TensorType(
output, "float32")
if not test_ref:
return
for target, dev in tvm.testing.enabled_targets():
op_res = relay.create_executor(
"graph", device=dev, target=target).evaluate(func)(x_data)
tvm.testing.assert_allclose(op_res.numpy(), ref_res)
def expected(x, w1, w2, scale1, scale2, channels1, channels2):
args = [x, w1, w2, scale1, scale2]
w = relay.concatenate((w1, w2), axis=0)
y = relay.nn.conv2d(x, w, channels=channels1 + channels2)
y1 = relay.strided_slice(y, [0, 0], [None, channels1])
y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2])
y1 = relay.multiply(y1, scale1)
y2 = relay.multiply(y2, scale2)
y = relay.Tuple((y1, y2))
return relay.Function(args, y)
def expected(x, w1, w2, scale1, scale2, channels1, channels2):
args = [x, w1, w2, scale1, scale2]
w = relay.concatenate((w1, w2), axis=0)
y = relay.nn.conv2d(x, w, channels=channels1 + channels2)
y1 = relay.strided_slice(y, [0, 0], [None, channels1])
y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2])
y1 = relay.multiply(y1, scale1)
y2 = relay.multiply(y2, scale2)
y = relay.Tuple((y1, y2))
return relay.Function(args, y)
def expected(x, w, channels, repeat):
args = [x, w]
y = x
for i in range(repeat):
w_concat = relay.concatenate((w, w), axis=0)
y = relay.nn.conv2d(y, w_concat, channels=channels*2)
y1 = relay.strided_slice(y, [0, 0], [None, channels])
y2 = relay.strided_slice(y, [0, channels], [None, channels * 2])
y = relay.concatenate((y1, y2), axis=1)
return relay.Function(args, y)
def expected(x, w, channels, repeat):
args = [x, w]
y = x
for i in range(repeat):
w_concat = relay.concatenate((w, w), axis=0)
y = relay.nn.conv2d(y, w_concat, channels=channels * 2)
y1 = relay.strided_slice(y, [0, 0], [None, channels])
y2 = relay.strided_slice(y, [0, channels], [None, channels * 2])
y = relay.concatenate((y1, y2), axis=1)
return relay.Function(args, y)
def expected(dshape):
x = relay.var("x", shape=dshape)
slice1 = relay.strided_slice(x, begin=[0, 0], end=[dshape[1]//2, dshape[1]], strides=[1,1])
slice2 = relay.strided_slice(x, begin=[dshape[1]//2, 0], end=[dshape[0], dshape[1]], strides=[1,1])
out = relay.Tuple((slice1, slice2))
f0 = relay.Function([x], out)
x = relay.var("x", shape=dshape)
y = relay.Call(f0, [x])
return relay.Function([x], y)
def verify(
dshape,
begin_val,
end_val,
strides_val,
output,
slice_mode="end",
test_ref=True,
dtype="int32",
):
x = relay.var("x", relay.TensorType(dshape, "float32"))
ndim = len(dshape)
begin_val = begin_val if begin_val else [0] * ndim
end_val = end_val if end_val else list(dshape)
if strides_val:
if len(strides_val) == 1:
strides_val = strides_val * ndim
else:
strides_val = [1] * ndim
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(
x_data, begin_val, end_val, strides_val, slice_mode)
data = [x_data, np.array(begin_val), np.array(end_val)]
begin = relay.var("begin", relay.TensorType((len(begin_val), ), dtype))
end = relay.var("end", relay.TensorType((len(end_val), ), dtype))
func_params = [x, begin, end]
if strides_val:
data.append(np.array(strides_val))
strides = relay.var("strides",
relay.TensorType((len(strides_val), ), dtype))
z = relay.strided_slice(x,
begin=begin,
end=end,
strides=strides,
slice_mode=slice_mode)
func_params.append(strides)
else:
z = relay.strided_slice(x,
begin=begin,
end=end,
slice_mode=slice_mode)
func = relay.Function(func_params, z)
params = {"begin": begin_val, "end": end_val, "strides": strides_val}
func = run_infer_type(func)
func2 = run_opt_pass(
run_opt_pass(func, transform.DynamicToStatic(), params),
transform.InferType())
assert isinstance(func2.body, relay.Call)
assert func2.body.op == relay.op.get("strided_slice")
verify_func(func2, [x_data], ref_res)
def verify(dshape,
begin,
end,
strides,
output,
slice_mode="end",
attr_const=True,
test_ref=True,
dtype="int32"):
x = relay.var("x", relay.TensorType(dshape, "float32"))
ndim = len(dshape)
begin = begin if begin else [0] * ndim
end = end if end else list(dshape)
# target numpy result
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = tvm.topi.testing.strided_slice_python(
x_data, begin, end, strides, slice_mode)
if attr_const:
begin = relay.const(begin, dtype=dtype)
end = relay.const(end, dtype=dtype)
if strides:
if attr_const:
strides = relay.const(strides, dtype=dtype)
z = relay.strided_slice(x,
begin=begin,
end=end,
strides=strides,
slice_mode=slice_mode)
else:
z = relay.strided_slice(x,
begin=begin,
end=end,
slice_mode=slice_mode)
func = relay.Function([x], z)
func = run_infer_type(func)
text = func.astext()
assert "begin=" in text
assert "end=" in text
if output:
assert func.body.checked_type == relay.ty.TensorType(
output, "float32")
if not test_ref:
return
for target, ctx in ctx_list():
intrp = relay.create_executor("graph", ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
def before(dshape):
x = relay.var("x", shape=dshape)
slice1 = relay.strided_slice(x,
begin=[0, 0],
end=[dshape[1] // 2, dshape[1]],
strides=[1, 1])
slice2 = relay.strided_slice(x,
begin=[dshape[1] // 2, 0],
end=[dshape[0], dshape[1]],
strides=[1, 1])
out = relay.Tuple((slice1, slice2))
return relay.Function([x], out)
def expected(x, w1, w2, w3, w4, channels1, channels2, channels3, channels4):
# use a fixed order of args so alpha equal check can pass
args = [x, w1, w2, w3, w4]
w = relay.concatenate((w1, w2, w4), axis=0)
y = relay.nn.conv2d(x, w, channels=channels1 + channels2 + channels4)
y1 = relay.strided_slice(y, [0, 0], [None, channels1])
y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2])
y3 = relay.nn.conv2d(x, w3)
y4 = relay.strided_slice(y, [0, channels1 + channels2],
[None, channels1 + channels2 + channels4])
y5 = relay.nn.max_pool2d(x)
y = relay.Tuple((y1, y2, y3, y4, y5))
return relay.Function(args, y)
def expected(x, w1, w2, w3, w4, channels1, channels2, channels3, channels4):
# use a fixed order of args so alpha equal check can pass
args = [x, w1, w2, w3, w4]
w = relay.concatenate((w1, w2, w4), axis=0)
y = relay.nn.conv2d(x, w, channels=channels1 + channels2 + channels4)
y1 = relay.strided_slice(y, [0, 0], [None, channels1])
y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2])
y3 = relay.nn.conv2d(x, w3)
y4 = relay.strided_slice(y, [0, channels1 + channels2],
[None, channels1 + channels2 + channels4])
y5 = relay.nn.max_pool2d(x)
y = relay.Tuple((y1, y2, y3, y4, y5))
return relay.Function(args, y)
def expected(x, w1, w2, scale1, scale2, channels1, channels2):
args = [x, w1, w2, scale1, scale2]
w = relay.concatenate((w1, w2), axis=0)
y = relay.nn.conv2d(x, w, channels=channels1 + channels2)
y1 = relay.strided_slice(
y, begin=[0, 0], end=[-1, channels1], strides=[1, 1], slice_mode="size"
)
y2 = relay.strided_slice(
y, begin=[0, channels1], end=[-1, channels2], strides=[1, 1], slice_mode="size"
)
y1 = relay.multiply(y1, scale1)
y2 = relay.multiply(y2, scale2)
y = relay.Tuple((y1, y2))
return relay.Function(args, y)
def expected(x, w, channels, repeat):
args = [x, w]
y = x
for i in range(repeat):
w_concat = relay.concatenate((w, w), axis=0)
y = relay.nn.conv2d(y, w_concat, channels=channels * 2)
y1 = relay.strided_slice(
y, begin=[0, 0], end=[-1, channels], strides=[1, 1], slice_mode="size"
)
y2 = relay.strided_slice(
y, begin=[0, channels], end=[-1, channels], strides=[1, 1], slice_mode="size"
)
y = relay.concatenate((y1, y2), axis=1)
return relay.Function(args, y)
def verify_any_strided_slice(data_shape, begin_shape, end_shape, strides_shape,
data_np_shape, slice_mode="end", const_attrs=False):
# Generate random numpy input data
np_data = np.random.uniform(size=data_np_shape).astype('float32')
np_begin = np.random.randint(2, size=begin_shape, dtype="int32")
np_end = np.random.randint(5, 10, size=end_shape, dtype="int32")
np_strides = np.random.randint(1, 2 if slice_mode == "size" else 3, size=strides_shape, dtype="int32")
# target numpy result
ref_res = topi.testing.strided_slice_python(np_data, np_begin, np_end, np_strides, slice_mode)
# Relay Module
mod = tvm.IRModule()
data = relay.var('data', shape=data_shape, dtype='float32')
if const_attrs:
data = relay.var('data', shape=data_np_shape, dtype='float32')
begin = relay.const(np_begin)
end = relay.const(np_end)
strides = relay.const(np_strides)
args = [data]
np_inputs = [np_data]
else:
begin = relay.var('begin', shape=begin_shape, dtype="int32")
end = relay.var('end', shape=end_shape, dtype="int32")
strides = relay.var('strides', shape=strides_shape, dtype="int32")
args = [data, begin, end, strides]
np_inputs = [np_data, np_begin, np_end, np_strides]
y = relay.strided_slice(data, begin=begin, end=end,
strides=strides, slice_mode=slice_mode)
mod["main"] = relay.Function(args, y)
for kind in ["debug", "vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(*np_inputs)
tvm.testing.assert_allclose(result.asnumpy(), ref_res)
def get_graph():
x = relay.var("x", shape=(1, 2, 2, 4), dtype="int8")
x = relay.strided_slice(x, begin=(0, 0, 0, 0), end=(1, 2, 2, 2))
x = infra.make_ethosu_identity(x)
x = relay.reshape(x, newshape=(1, 2, 4))
x = infra.make_ethosu_identity(x)
return relay.Function(relay.analysis.free_vars(x), x)
def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr,
node_map: tvm.ir.container.Map) -> tvm.relay.Expr:
split_input = post.args[0]
split_begins = list()
split_ends = list()
section_begins_in_split_axis = self.get_section_begin_coords(post)
for split_cord in section_begins_in_split_axis:
# first begin is [0, 0, ... , 0]
begin_shape = [
0 for i in range(len(split_input.checked_type.shape))
]
begin_shape[post.attrs.axis] = split_cord
split_begins.append(begin_shape)
end_shape = list(split_input.checked_type.shape)
# Only the split axis coordinate changes
end_shape[post.attrs.axis] = split_cord
split_ends.append(end_shape)
# Coordinates needs to be shifted left because beginning
# of the next section is the end of the previous
split_ends = split_ends[1:]
# Last section end is the shape of the tensor itself.
split_ends.append(list(split_input.checked_type.shape))
strided_slices = list()
for sb, se in zip(split_begins, split_ends):
strided_slices.append(relay.strided_slice(split_input, sb, se))
return relay.Tuple(strided_slices)
def before():
x = relay.var("x", shape=(1, 32, 28, 28))
weight = relay.var('weight', shape=(32, 32, 3, 3))
y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1))
y = relay.strided_slice(y, begin=[0, 16], end=[None, None])
y = relay.Function(free_vars(y), y)
return y
def conv2relay(node: Conv, batch_size):
node2var = {}
params = {}
for term in node.inputs:
for value in term:
x = relay.var(name_hint=value.node.name, shape=(batch_size, *value.node.output_shape))
node2var[value.node] = x
term_vars = []
for term in node.inputs:
value_vars = []
for value in term:
if value.begin == 0 and value.end == value.node.output_shape[0]:
var = node2var[value.node]
else:
var = relay.strided_slice(node2var[value.node], begin=[0, value.begin, 0, 0],
end=[batch_size, value.end, *value.node.output_shape[1:]])
value_vars.append(var)
term_var = value_vars[0]
for value_var in value_vars[1:]:
term_var = term_var + value_var
term_vars.append(term_var)
if len(term_vars) > 1:
x = relay.concatenate(term_vars, axis=1)
else:
x = term_vars[0]
x = conv2d(x, node, params)
node2var[node] = x
fn = relay.Function(relay.analysis.free_vars(x), x)
if tvm_minor_version() <= 6:
return relay.Module.from_expr(fn), params
else:
return tvm.ir.IRModule.from_expr(fn), params
def before():
x = relay.var("x", shape=(1, 32, 28, 28))
weight = relay.var('weight', shape=(32, 32, 3, 3))
y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1))
y = relay.strided_slice(y, begin=[0, 16], end=[None, None])
y = relay.Function(analysis.free_vars(y), y)
return y
def get_graph():
in_1 = relay.var("x", shape=(1, 16, 16, 8), dtype="int8")
conv_1 = infra.make_ethosu_conv2d(
ifm=in_1,
ifm_channels=8,
ofm_channels=8,
kernel_shape=(1, 1),
padding=(0, 0),
strides=(1, 1),
dilation=(1, 1),
ifm_layout="NHWC",
ofm_layout="NHWC",
)
reshape = relay.reshape(conv_1, (1, 16, 16, 8))
strided_slice = relay.strided_slice(reshape, (0, 0, 0, 0),
(1, 16, 16, 8))
conv_2 = infra.make_ethosu_conv2d(
ifm=strided_slice,
ifm_channels=8,
ofm_channels=8,
kernel_shape=(1, 1),
padding=(0, 0),
strides=(1, 1),
dilation=(1, 1),
ifm_layout="NHWC",
ofm_layout="NHWC",
)
return relay.Function(relay.analysis.free_vars(conv_2), conv_2)
def expected():
if not do_pad:
return before()
x = relay.var("x", shape=data_shape, dtype=dtype)
if db or di:
x_pad = relay.nn.pad(x,
pad_width=((0, db), (0, 0), (0, 0), (0,
di)))
else:
x_pad = x
weight = relay.var("weight", shape=(kernel_shape), dtype=dtype)
if di or do:
weight_pad = relay.nn.pad(weight,
pad_width=((0, 0), (0, 0), (0, di),
(0, do)))
else:
weight_pad = weight
y_pad = relay.nn.conv2d(
x_pad,
weight=weight_pad,
channels=out_channel + do,
kernel_size=(3, 3),
padding=(1, 1),
data_layout="NHWC",
kernel_layout="HWIO",
)
if db or do:
y = relay.strided_slice(y_pad,
begin=[0, 0, 0, 0],
end=out_shape)
else:
y = y_pad
y = relay.Function([x, weight], y)
return y
def _pad_conv2d_NHWC(db, di, do, data, kernel, out_channel, new_attrs,
output_tensor):
# Pad batch size
if db != 0:
data = relay.nn.pad(data, pad_width=((0, db), (0, 0), (0, 0), (0, 0)))
# Pad input channel
if di != 0:
data = relay.nn.pad(data, pad_width=((0, 0), (0, 0), (0, 0), (0, di)))
kernel = relay.nn.pad(kernel,
pad_width=((0, 0), (0, 0), (0, di), (0, 0)))
# Pad output channel
if do != 0:
kernel = relay.nn.pad(kernel,
pad_width=((0, 0), (0, 0), (0, 0), (0, do)))
if do != 0:
new_out_channel = out_channel + do
new_attrs["channels"] = new_out_channel
out = relay.nn.conv2d(data, kernel, **new_attrs)
if db != 0 or do != 0:
original_out_shape = [x.value for x in output_tensor.shape]
out = relay.strided_slice(out,
begin=[0, 0, 0, 0],
end=original_out_shape)
return out
def _get_func():
ifm = relay.var("ifm", shape=(1, 12, 12, 8), dtype="int8")
conv1 = make_ethosu_conv2d(
ifm=ifm,
ifm_channels=8,
ofm_channels=64,
kernel_shape=(1, 1),
padding=(0, 0),
strides=(1, 1),
dilation=(1, 1),
activation="NONE",
ifm_layout="NHWC",
ofm_layout="NHWC",
)
strided_slice = relay.strided_slice(conv1, [0, 0, 0, 0], [1, 6, 6, 128])
conv2 = make_ethosu_conv2d(
ifm=strided_slice,
ifm_channels=64,
ofm_channels=16,
kernel_shape=(3, 3),
padding=(1, 1),
strides=(1, 1),
dilation=(1, 1),
activation="NONE",
ifm_layout="NHWC",
ofm_layout="NHCWB16",
)
func = relay.Function(relay.analysis.free_vars(conv2), conv2)
func = run_opt_pass(func, relay.transform.InferType())
return func
def verify_strided_slice(dshape, begin, end, strides, mode):
x = relay.var("x", relay.TensorType(dshape, "float32"))
if mode == "size":
strides = None
z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=mode)
func = relay.Function([x], z)
x_data = np.random.uniform(size=dshape).astype("float32")
verify_results(func, [x_data], "test_strided_slice", rtol=1e-5, atol=1e-5)
def before():
shape = (tvm.tir.const(10, "int64"),
tvm.tir.const(1, "int64"))
x = relay.var("x", shape=shape)
concat = relay.concatenate([x,x], axis=-1)
out = relay.strided_slice(concat, begin=[np.int64(0)], end=[np.int64(3)])
t = relay.Function(relay.analysis.free_vars(out), out)
return relay.Function(relay.analysis.free_vars(out), out)
def get_graph(x_shape, begin, end, strides=None):
x = relay.var("x", shape=(x_shape), dtype="float32")
if strides:
out = relay.strided_slice(
x,
relay.expr.const(begin, dtype="int32"),
relay.expr.const(end, dtype="int32"),
relay.expr.const(strides, dtype="int32"),
)
else:
out = relay.strided_slice(
x,
relay.expr.const(begin, dtype="int32"),
relay.expr.const(end, dtype="int32"),
)
f = relay.Function([x], out)
return f, {"x": x_shape}, []
def get_graph(get_expected=False):
x = relay.var("x", shape=(1, 2, 2, 4), dtype="int8")
x = infra.make_ethosu_pooling(x, "MAX", (1, 1), 4, (1, 1), (0, 0))
x = relay.strided_slice(x, begin=[0, 0, 0, 0], end=[1, 2, 2, 2])
if not get_expected:
x = infra.make_ethosu_identity(x)
x = infra.make_ethosu_pooling(x, "MAX", (1, 1), 2, (1, 1), (0, 0))
return relay.Function(relay.analysis.free_vars(x), x)
def expected():
x = relay.var("x", shape=(1, 32, 28, 28))
weight = relay.var("weight")
x = relay.layout_transform(x, "NCHW", "NCHW4c")
y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1),
data_layout="NCHW4c")
y = relay.strided_slice(y, begin=[0, 4], end=[None, 8])
y = relay.layout_transform(y, "NCHW4c", "NCHW")
y = relay.Function(free_vars(y), y)
return y
def verify(dshape, begin, end, strides, output, test_ref=True):
x = relay.var("x", relay.TensorType(dshape, "float32"))
z = relay.strided_slice(x, begin=begin, end=end, strides=strides)
func = relay.Function([x], z)
func = relay.ir_pass.infer_type(func)
text = func.astext()
assert "begin=" in text
assert "end=" in text
if output:
assert func.body.checked_type == relay.ty.TensorType(output, "float32")
if not test_ref:
return
x_data = np.random.uniform(size=dshape).astype("float32")
ref_res = topi.testing.strided_slice_python(
x_data, begin, end, strides)
for target, ctx in ctx_list():
intrp = relay.create_executor("graph", ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x_data)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
