From aa6d17f7ab3b7423d012c600c28d9a1899b585cc Mon Sep 17 00:00:00 2001 From: wpan11nv <60017475+wpan11nv@users.noreply.github.com> Date: Thu, 20 Feb 2020 09:25:35 -0800 Subject: [PATCH] [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 --- 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 22844745982f..07cae5e2c746 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 node = make_object(); diff --git a/tests/python/relay/test_op_level4.py b/tests/python/relay/test_op_level4.py index c5cd70818795..44b51f2c2367 100644 --- a/tests/python/relay/test_op_level4.py +++ b/tests/python/relay/test_op_level4.py @@ -323,6 +323,7 @@ def verify(dshape, begin, end, strides, vshape, test_ref=True): 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))