This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git
The following commit(s) were added to refs/heads/master by this push:
new efd35e8 [Relay] Fix an assertion exposed by loop vectorizer (#4916)
efd35e8 is described below
commit efd35e86992815d7f46a77dfced79180964472ac
Author: wpan11nv <[email protected]>
AuthorDate: Thu Feb 20 09:25:35 2020 -0800
[Relay] Fix an assertion exposed by loop vectorizer (#4916)
- Allows uniform conditions for select expressions (the same as halide)
exposed by the loop vectorizer.
Signed-off-by: Wei Pan <[email protected]>
---
src/tir/ir/expr.cc | 3 ++-
tests/python/relay/test_op_level4.py | 1 +
2 files changed, 3 insertions(+), 1 deletion(-)
diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc
index 2284474..07cae5e 100644
--- a/src/tir/ir/expr.cc
+++ b/src/tir/ir/expr.cc
@@ -158,7 +158,8 @@ PrimExpr SelectNode::make(PrimExpr condition, PrimExpr
true_value, PrimExpr fals
CHECK(true_value.defined()) << "ValueError: true_value is undefined";
CHECK(false_value.defined()) << "ValueError: true_value is undefined";
CHECK(condition.dtype().is_bool());
- CHECK_EQ(condition.dtype().lanes(), true_value.dtype().lanes());
+ CHECK(condition.dtype().lanes() == true_value.dtype().lanes() ||
+ condition.dtype().lanes() == 1);
CHECK(false_value.dtype() == true_value.dtype()) << "TypeError: mismatched
types";
ObjectPtr<SelectNode> node = make_object<SelectNode>();
diff --git a/tests/python/relay/test_op_level4.py
b/tests/python/relay/test_op_level4.py
index c5cd708..44b51f2 100644
--- a/tests/python/relay/test_op_level4.py
+++ b/tests/python/relay/test_op_level4.py
@@ -323,6 +323,7 @@ def test_strided_set():
op_res = intrp.evaluate(func)(x_data, v_data)
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
+ verify((3, 4, 16), [0, 0, 0], [4, -5, 4], [1, -1, 2], (3, 1, 2))
verify((3, 4, 3), [0, 0, 0], [4, -5, 4], [1, -1, 2], (3, 1, 2))
verify((3, 4, 3), [1, 1, 0], [4, 4, 3], [2, 1, 1], (1, 3, 3))
verify((3, 4, 3), [1, -1, 0], [4, -5, 3], [2, -1, 1], (1, 4, 3))
