[OpenCL] Fix type casting error

src/target/source/codegen_opencl.cc

@@ -327,6 +327,10 @@ void CodeGenOpenCL::PrintRestrict(const Var& v, std::ostream& os) {
 
 std::string CodeGenOpenCL::CastFromTo(std::string value, DataType from, DataType target) {
   if (from == target) return value;
+  return CastTo(value, target);
+}
+
+std::string CodeGenOpenCL::CastTo(std::string value, DataType target) {
   std::ostringstream os;
   if (target.lanes() == 1) {
     os << "((";
@@ -512,6 +516,30 @@ void CodeGenOpenCL::VisitExpr_(const MaxNode* op, std::ostream& os) {
   PrintBinaryExpr(op, "max", os, this);
 }
 
+void CodeGenOpenCL::PrintVecBinaryOp(const std::string& op, DataType t, PrimExpr lhs, PrimExpr rhs,
+                                     std::ostream& os) {
+  std::ostringstream oss;
+  if (isalpha(op[0])) {
+    os << op << "(";
+    this->PrintExpr(lhs, oss);
+    os << CastTo(oss.str(), t);
+    oss.str("");
+    os << ", ";
+    this->PrintExpr(rhs, oss);
+    os << CastTo(oss.str(), t);
+    os << ")";
+  } else {
+    os << "(";
+    this->PrintExpr(lhs, oss);
+    os << CastTo(oss.str(), t);
+    oss.str("");
+    os << ' ' << op << ' ';
+    this->PrintExpr(rhs, oss);
+    os << CastTo(oss.str(), t);
+    os << ")";
+  }
+}
+
 void CodeGenOpenCL::SetTextureScope(
     const std::unordered_map<const VarNode*, std::string>& scope) {  // NOLINT(*)
   for (auto& texture : scope) {

src/target/source/codegen_opencl.h

@@ -55,6 +55,7 @@ class CodeGenOpenCL final : public CodeGenC {
                     std::ostream& os);                                           // NOLINT(*)
   void PrintRestrict(const Var& v, std::ostream& os) final;                      // NOLINT(*)
   std::string CastFromTo(std::string value, DataType from, DataType target);     // NOLINT(*)
+  std::string CastTo(std::string value, DataType target);                        // NOLINT(*)
   void SetTextureScope(const std::unordered_map<const VarNode*, std::string>&);  // NOLINT(*)
 
   // overload visitor
@@ -70,6 +71,10 @@ class CodeGenOpenCL final : public CodeGenC {
   void VisitExpr_(const MinNode* op, std::ostream& os) final;
   void VisitExpr_(const MaxNode* op, std::ostream& os) final;
 
+  // Binary vector op.
+  void PrintVecBinaryOp(const std::string& op, DataType op_type, PrimExpr lhs, PrimExpr rhs,
+                        std::ostream& os) final;
+
  private:
   // whether enable fp16 and fp64 extension
   bool enable_fp16_{false};

tests/python/unittest/test_target_codegen_opencl.py

@@ -139,8 +139,45 @@ def check_erf(dev, n, dtype):
     check_erf(dev, 1, "float64")
 
 
+@tvm.testing.requires_gpu
+@tvm.testing.requires_opencl
+def test_opencl_type_casting():
+    def check_type_casting(ctx, n, dtype):
+        block_size = 4
+        C = te.compute(
+            (n,),
+            lambda i: tvm.tir.Select(
+                tvm.tir.all(
+                    *[
+                        i // block_size == tvm.tir.const(3, "int32"),
+                        i % block_size == tvm.tir.const(3, "int32"),
+                    ]
+                ),
+                tvm.tir.const(1, dtype),
+                tvm.tir.const(0, dtype),
+            ),
+            name="C",
+        )
+        s = te.create_schedule(C.op)
+        (tx, vx) = s[C].split(s[C].op.axis[0], factor=block_size)
+        s[C].vectorize(vx)
+        thrx = te.thread_axis("threadIdx.x")
+
+        s[C].bind(tx, thrx)
+        fun = tvm.build(s, [C], target)
+
+        c = tvm.nd.empty((n,), dtype, ctx)
+        # Only need to test compiling here
+        fun(c)
+
+    dev = tvm.device(target, 0)
+
+    check_type_casting(dev, 16, "float32")
+
+
 if __name__ == "__main__":
     test_opencl_ternary_expression()
     test_opencl_inf_nan()
     test_opencl_max()
     test_opencl_erf()
+    test_opencl_type_casting()