[incubator-tvm] branch master updated: [Relay] Fix an assertion exposed by loop vectorizer (#4916)

[tqchen]

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 <60017475+wpan1...@users.noreply.github.com>
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 <w...@nvidia.com>
---
 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))