sklam · Jun 13, 2022 · Mar 29, 2022 · Apr 12, 2022 · Apr 12, 2022 · Apr 12, 2022
diff --git a/docs/source/cuda-reference/kernel.rst b/docs/source/cuda-reference/kernel.rst
@@ -537,3 +537,35 @@ Timer Intrinsics
 
     Suspends the thread for a sleep duration approximately close to the delay
     ``ns``, specified in nanoseconds.
+
+
+Vector Types
+~~~~~~~~~~~~
+
+`CUDA Vector Types <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#built-in-vector-types>`_
+are available in the kernel. Unlike native CUDA, vector types are constructed
+directly with their constructor. For example, to construct a ``float3``:
+
+.. code-block:: python3
+
+    from numba.cuda import float3
+
+    # In kernel
+    f3 = float3(float32(0.0), float32(-1.0), float32(1.0))
+
+Besides, vector type can be constructed from a combination of vector types and
+primitive types, as long as the total number of components matches with the result
+vector type. For example, all of the following constructions are valid:
+
+.. code-block:: python3
+
+    zero = uint32(0)
+    u2 = uint2(1, 2)
+    # Construct a 3-component vector with primitive type and a 2-component vector
+    u3 = uint3(zero, u2)
+    # Construct a 4-component vector with 2 2-component vectors
+    u4 = uint4(u2, u2)
+
+The 1st, 2nd, 3rd and 4th component of the vector type can be accessed through fields 
+``x``, ``y``, ``z``, ``w``, respectively. The components are currently immutable
+after construction.
@@ -1,12 +1,14 @@
 # Re export
+import sys
 from .stubs import (threadIdx, blockIdx, blockDim, gridDim, laneid,
                     warpsize, syncthreads, syncthreads_count, syncwarp,
                     syncthreads_and, syncthreads_or, shared, local,
                     const, grid, gridsize, atomic, shfl_sync_intrinsic,
                     vote_sync_intrinsic, match_any_sync, match_all_sync,
                     threadfence_block, threadfence_system,
                     threadfence, selp, popc, brev, clz, ffs, fma, cbrt,
-                    cg, activemask, lanemask_lt, nanosleep, fp16)
+                    cg, activemask, lanemask_lt, nanosleep, fp16,
+                    _vector_type_stubs)
 from .cudadrv.error import CudaSupportError
 from numba.cuda.cudadrv.driver import (BaseCUDAMemoryManager,
                                        HostOnlyCUDAMemoryManager,
@@ -31,6 +33,10 @@
 
 reduce = Reduce = reduction.Reduce
 
+# Expose vector types constructors as module level attributes.
+for vector_type_stub in _vector_type_stubs:
+    setattr(sys.modules[__name__], vector_type_stub.__name__, vector_type_stub)
+
 
 def is_available():
     """Returns a boolean to indicate the availability of a CUDA GPU.

@@ -2,6 +2,7 @@
 This scripts specifies all PTX special objects.
 """
 import functools
+import itertools
 
 
 class Stub(object):
@@ -656,3 +657,40 @@ class habs(Stub):
         Returns the fp16 result of the absolute value.
 
         """
+
+
+#-------------------------------------------------------------------------------
+# vector types
+
+_vector_type_prefix = (
+    "char",
+    "short",
+    "int",
+    "long",
+    "longlong",
+    "uchar",
+    "ushort",
+    "uint",
+    "ulong",
+    "ulonglong",
+    "float",
+    "double"
+)
+_vector_type_element_counts = (1, 2, 3, 4)
+_vector_type_attribute_names = ("x", "y", "z", "w")
+
+_vector_type_stubs = []
+for name, nelem in itertools.product(
+    _vector_type_prefix, _vector_type_element_counts
+):
+    type_name = f"{name}{nelem}"
+    attr_names = _vector_type_attribute_names[:nelem]
+
+    vector_type_stub = type(
+        type_name, (Stub,),
+        {
+            **{attr: lambda self: None for attr in attr_names},
+            **{"_description_": f"<{type_name}>"}
+        }
+    )
+    _vector_type_stubs.append(vector_type_stub)
@@ -19,14 +19,16 @@
 
 class CUDATypingContext(typing.BaseContext):
     def load_additional_registries(self):
-        from . import cudadecl, cudamath, libdevicedecl
+        from . import cudadecl, cudamath, libdevicedecl, vector_types
         from numba.core.typing import enumdecl
 
+        vector_types.initialize_once()
         self.install_registry(cudadecl.registry)
         self.install_registry(cudamath.registry)
         self.install_registry(cmathdecl.registry)
         self.install_registry(libdevicedecl.registry)
         self.install_registry(enumdecl.registry)
+        self.install_registry(vector_types.typing_registry)
 
     def resolve_value_type(self, val):
         # treat other dispatcher object as another device function
@@ -96,12 +98,17 @@ def load_additional_registries(self):
         from numba.cpython import cmathimpl
         from numba.np import arrayobj # noqa: F401
         from numba.np import npdatetime # noqa: F401
-        from . import cudaimpl, printimpl, libdeviceimpl, mathimpl
+        from . import (
+            cudaimpl, printimpl, libdeviceimpl, mathimpl, vector_types
+        )
+
+        vector_types.initialize_once()
         self.install_registry(cudaimpl.registry)
         self.install_registry(printimpl.registry)
         self.install_registry(libdeviceimpl.registry)
         self.install_registry(cmathimpl.registry)
         self.install_registry(mathimpl.registry)
+        self.install_registry(vector_types.impl_registry)
 
     def codegen(self):
         return self._internal_codegen

@@ -0,0 +1,155 @@
+import numpy as np
+
+from numba.cuda.testing import CUDATestCase
+
+from numba import cuda
+from numba.cuda.vector_types import vector_types
+
+
+def make_kernel(vtype):
+    vobj = vtype.user_facing_object
+    base_type = vtype.base_type
+
+    def kernel_1elem(res):
+        v = vobj(base_type(0))
+        res[0] = v.x
+
+    def kernel_2elem(res):
+        v = vobj(base_type(0), base_type(1))
+        res[0] = v.x
+        res[1] = v.y
+
+    def kernel_3elem(res):
+        v = vobj(base_type(0), base_type(1), base_type(2))
+        res[0] = v.x
+        res[1] = v.y
+        res[2] = v.z
+
+    def kernel_4elem(res):
+        v = vobj(
+            base_type(0),
+            base_type(1),
+            base_type(2),
+            base_type(3)
+        )
+        res[0] = v.x
+        res[1] = v.y
+        res[2] = v.z
+        res[3] = v.w
+
+    host_function = {
+        1: kernel_1elem,
+        2: kernel_2elem,
+        3: kernel_3elem,
+        4: kernel_4elem
+    }[vtype.num_elements]
+    return cuda.jit(host_function)
+
+
+def make_fancy_creation_kernel(vtype):
+    base_type = vtype.base_type
+    v2 = getattr(cuda, f"{vtype.name[:-1]}2")
+    v3 = getattr(cuda, f"{vtype.name[:-1]}3")
+    v4 = getattr(cuda, f"{vtype.name[:-1]}4")
+
+    def kernel(res):
+        one = base_type(1.0)
+        two = base_type(2.0)
+        three = base_type(3.0)
+        four = base_type(4.0)
+
+        # 2 3
+        f2 = v2(two, three)
+        # 2 3
+        f2_2 = v2(f2)
+
+        # 2 3 1
+        f3 = v3(f2, one)
+        # 1 2 3
+        f3_2 = v3(one, f2)
+        # 1 2 3
+        f3_3 = v3(one, two, three)
+        # 2 3 1
+        f3_4 = v3(f3)
+
+        # 1 2 3 4
+        f4 = v4(one, two, three, four)
+        # 2 3 1 4
+        f4_2 = v4(f2, one, four)
+        # 1 2 3 4
+        f4_3 = v4(one, f2, four)
+        # 1 4 2 3
+        f4_4 = v4(one, four, f2)
+        # 2 3 2 3
+        f4_5 = v4(f2, f2)
+        # 2 3 1 4
+        f4_6 = v4(f3, four)
+        # 4 2 3 1
+        f4_7 = v4(four, f3)
+        # 1 2 3 4
+        f4_8 = v4(f4)
+
+        res[0] = f2.x
+        res[1] = f2.y
+        res[2] = f2_2.x
+        res[3] = f2_2.y
+
+        j = 4
+        for i, f in enumerate((f3, f3_2, f3_3, f3_4)):
+            res[i * 3 + j] = f.x
+            res[i * 3 + j + 1] = f.y
+            res[i * 3 + j + 2] = f.z
+
+        j += 12
+        for i, f in enumerate((f4, f4_2, f4_3, f4_4, f4_5, f4_6, f4_7, f4_8)):
+            res[i * 4 + j] = f.x
+            res[i * 4 + j + 1] = f.y
+            res[i * 4 + j + 2] = f.z
+            res[i * 4 + j + 3] = f.w
+
+    return cuda.jit(kernel)
+
+
+class TestCudaVectorType(CUDATestCase):
+
+    def setUp(self):
+        """Compile a empty kernel to initialize vector types."""
+        super().setUp()
+
+        @cuda.jit("()")
+        def k():
+            pass
+
+    def test_creation_readout(self):
+        for vty in vector_types.values():
+            with self.subTest(vty=vty):
+                arr = np.zeros((vty.num_elements,))
+                kernel = make_kernel(vty)
+                kernel[1, 1](arr)
+                np.testing.assert_almost_equal(
+                    arr, np.array(range(vty.num_elements))
+                )
+
+    def test_fancy_creation_readout(self):
+        for vty in vector_types.values():
+            kernel = make_fancy_creation_kernel(vty)
+
+            expected = np.array([
+                2, 3,
+                2, 3,
+                2, 3, 1,
+                1, 2, 3,
+                1, 2, 3,
+                2, 3, 1,
+                1, 2, 3, 4,
+                2, 3, 1, 4,
+                1, 2, 3, 4,
+                1, 4, 2, 3,
+                2, 3, 2, 3,
+                2, 3, 1, 4,
+                4, 2, 3, 1,
+                1, 2, 3, 4
+            ])
+            arr = np.zeros(expected.shape)
+            kernel[1, 1](arr)
+            np.testing.assert_almost_equal(arr, expected)