Modify the typecast for vector dtypes.

Author: nkoskelo

loopy/target/c/codegen/expression.py

@@ -130,8 +130,8 @@ def wrap_in_typecast(self, actual_type: LoopyType, needed_type: LoopyType, s):
         if actual_type != needed_type:
             registry = self.codegen_state.ast_builder.target.get_dtype_registry()
             cast = var("(%s) " % registry.dtype_to_ctype(needed_type))
-
             return cast(s)
+
         return s
 
     def rec(self, expr, type_context=None, needed_type: LoopyType | None = None):  # type: ignore[override]
@@ -283,10 +283,12 @@ def make_var(name):
                 if (
                         isinstance(ary, (ConstantArg, ArrayArg)) or
                         (isinstance(ary, TemporaryVariable) and ary.base_storage)):
-                    # unsubscripted global args are pointers if they are inputs
-                    result = self.make_subscript(ary,
-                                                 make_var(access_info.array_name),
-                                                 0)
+                    # unsubscripted global args are pointers
+                    result = self.make_subscript(
+                            ary,
+                            make_var(access_info.array_name),
+                            (0,))
+
                 else:
                     # unsubscripted temp vars are scalars
                     # (unless they use base_storage)
@@ -412,37 +414,9 @@ def map_remainder(self, expr, type_context):
     def map_if(self, expr, type_context):
         from loopy.types import to_loopy_type
         result_type = self.infer_type(expr)
-        conditional_needed_loopy_type = to_loopy_type(np.bool_)
-        if self.codegen_state.vectorization_info:
-            from loopy.codegen import UnvectorizableError
-            from loopy.expression import VectorizabilityChecker
-            checker = VectorizabilityChecker(self.codegen_state.kernel,
-                                     self.codegen_state.vectorization_info.iname,
-                                     self.codegen_state.vectorization_info.length)
-
-            try:
-                is_vector = checker(expr)
-
-                if is_vector:
-                    """
-                    We could have a vector literal here.
-                    So we may need to type cast the condition.
-                    OpenCL specification states that for ( c ? a : b)
-                    to be vectorized appropriately c must have the same
-                    number of elements in the vector as that of a and b.
-                    Also each element must have the same number of bits,
-                    and c must be an integral type.
-                    """
-                    index_type = to_loopy_type(np.int64)
-                    if type_context == "f":
-                        index_type = to_loopy_type(np.int32)
-                    conditional_needed_loopy_type = to_loopy_type(self.codegen_state.target.vector_dtype(index_type,
-                                                    self.codegen_state.vectorization_info.length))
-            except UnvectorizableError:
-                pass
         return type(expr)(
                 self.rec(expr.condition, type_context,
-                         conditional_needed_loopy_type),
+                         to_loopy_type(np.bool_)),
                 self.rec(expr.then, type_context, result_type),
                 self.rec(expr.else_, type_context, result_type),
                 )

loopy/target/opencl.py

@@ -577,6 +577,47 @@ def map_variable(self, expr, type_context):
                         ",".join([f"{i}" for i in range(vector_length)]) + "))"
                 return Literal(vector_literal)
         return super().map_variable(expr, type_context)
+
+    def map_if(self, expr, type_context):
+        from loopy.types import to_loopy_type
+        result_type = self.infer_type(expr)
+        conditional_needed_loopy_type = to_loopy_type(np.bool_)
+        if self.codegen_state.vectorization_info:
+            from loopy.codegen import UnvectorizableError
+            from loopy.expression import VectorizabilityChecker
+            checker = VectorizabilityChecker(self.codegen_state.kernel,
+                                     self.codegen_state.vectorization_info.iname,
+                                     self.codegen_state.vectorization_info.length)
+
+            try:
+                is_vector = checker(expr)
+
+                if is_vector:
+                    """
+                    We could have a vector literal here.
+                    So we may need to type cast the condition.
+                    OpenCL specification states that for ( c ? a : b)
+                    to be vectorized appropriately c must have the same
+                    number of elements in the vector as that of a and b.
+                    Also each element must have the same number of bits,
+                    and c must be an integral type.
+                    """
+                    index_type = to_loopy_type(np.int64)
+                    if type_context == "f":
+                        index_type = to_loopy_type(np.int32)
+                    length = self.codegen_state.vectorization_info.length
+                    vector_type = self.codegen_state.target.vector_dtype(index_type,
+                                                                         length)
+                    conditional_needed_loopy_type = to_loopy_type(vector_type)
+            except UnvectorizableError:
+                pass
+
+        return type(expr)(
+                self.rec(expr.condition, type_context,
+                         conditional_needed_loopy_type),
+                self.rec(expr.then, type_context, result_type),
+                self.rec(expr.else_, type_context, result_type),
+                )
 # }}}