[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 <weip@nvidia.com>

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>();

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))