This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 82b6ea9 Add missing Slice layout fallback check of `stride=1`.
(#10690)
82b6ea9 is described below
commit 82b6ea9fc4e9f513a55b70e31b46ce529e339cb6
Author: Jinkun Lin <[email protected]>
AuthorDate: Thu Mar 24 16:19:40 2022 -0400
Add missing Slice layout fallback check of `stride=1`. (#10690)
* Add missing Slice layout fallback check.
* Fix lint
* jostle ci
---
src/relay/op/tensor/transform.cc | 7 +++++++
tests/python/relay/test_pass_alter_op_layout.py | 21 +++++++++++++++++++++
2 files changed, 28 insertions(+)
diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc
index 3f7da4e..d225d93 100644
--- a/src/relay/op/tensor/transform.cc
+++ b/src/relay/op/tensor/transform.cc
@@ -2711,6 +2711,13 @@ InferCorrectLayoutOutput StridedSliceInferCorrectLayout(
new_begin.push_back(begin[i]);
new_end.push_back(end[i]);
} else {
+ if (strides.defined() && i < strides.size()) {
+ auto stride = strides[i];
+ // arbitrary stride is not supported
+ if (stride.defined() && stride->value != 1) {
+ return out_default;
+ }
+ }
int64_t bg = begin[i];
int64_t ed = end[i];
if (bg % factor || ed % factor) {
diff --git a/tests/python/relay/test_pass_alter_op_layout.py
b/tests/python/relay/test_pass_alter_op_layout.py
index 250c177..cffc33b 100644
--- a/tests/python/relay/test_pass_alter_op_layout.py
+++ b/tests/python/relay/test_pass_alter_op_layout.py
@@ -1499,6 +1499,27 @@ def test_conv2d_strided_slice_packed_to_unpacked():
assert tvm.ir.structural_equal(a, b)
+def test_conv2d_strided_slice_arbitrary_stride():
+ """Test rewriting strided_slice with arbitrary stride"""
+
+ def before():
+ x = relay.var("x", shape=(4, 12, 1, 1))
+ weight = relay.var("weight", shape=(9, 12, 1, 1))
+ y = relay.nn.conv2d(x, weight, channels=9, kernel_size=(1, 1),
padding=(0, 0))
+ y = relay.strided_slice(y, begin=[3], end=[6], strides=[3], axes=[1])
+ y = relay.Function(analysis.free_vars(y), y)
+ return y
+
+ def alter_conv2d(attrs, inputs, tinfos, out_type):
+ data, weight = inputs
+ new_attrs = dict(attrs)
+ new_attrs["data_layout"] = "NCHW3c"
+ return relay.nn.conv2d(data, weight, **new_attrs)
+
+ with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d):
+ run_opt_pass(before(), transform.AlterOpLayout())
+
+
def test_conv2d_reduce_channels():
x = relay.var("data", shape=(1, 8, 48, 48))
y = relay.nn.conv2d(
