Implement alignment support for local and shared arrays.

Author: tpn

numba_cuda/numba/cuda/cudadecl.py

@@ -25,7 +25,7 @@
 
 class Cuda_array_decl(CallableTemplate):
     def generic(self):
-        def typer(shape, dtype):
+        def typer(shape, dtype, alignment=None):
 
             # Only integer literals and tuples of integer literals are valid
             # shapes
@@ -39,10 +39,13 @@ def typer(shape, dtype):
             else:
                 return None
 
+            # N.B. alignment validation happens in types.Array().
+
             ndim = parse_shape(shape)
             nb_dtype = parse_dtype(dtype)
             if nb_dtype is not None and ndim is not None:
-                return types.Array(dtype=nb_dtype, ndim=ndim, layout='C')
+                return types.Array(dtype=nb_dtype, ndim=ndim, layout='C',
+                                   alignment=alignment)
 
         return typer
 

numba_cuda/numba/cuda/cudaimpl.py

@@ -72,9 +72,14 @@ def dim3_z(context, builder, sig, args):
 # -----------------------------------------------------------------------------
 
 @lower(cuda.const.array_like, types.Array)
+@lower(cuda.const.array_like, types.Array, types.IntegerLiteral)
+@lower(cuda.const.array_like, types.Array, types.NoneType)
 def cuda_const_array_like(context, builder, sig, args):
     # This is a no-op because CUDATargetContext.make_constant_array already
     # created the constant array.
+    if len(sig.args) > 1:
+        # XXX-140: How do we handle alignment here?
+        pass
     return args[0]
 
 
@@ -92,34 +97,62 @@ def _get_unique_smem_id(name):
 
 
 @lower(cuda.shared.array, types.IntegerLiteral, types.Any)
+@lower(cuda.shared.array, types.IntegerLiteral, types.Any, types.IntegerLiteral)
+@lower(cuda.shared.array, types.IntegerLiteral, types.Any, types.NoneType)
 def cuda_shared_array_integer(context, builder, sig, args):
     length = sig.args[0].literal_value
     dtype = parse_dtype(sig.args[1])
+    alignment = None
+    if len(sig.args) == 3:
+        try:
+            alignment = sig.args[2].literal_value
+        except AttributeError:
+            pass
     return _generic_array(context, builder, shape=(length,), dtype=dtype,
                           symbol_name=_get_unique_smem_id('_cudapy_smem'),
                           addrspace=nvvm.ADDRSPACE_SHARED,
-                          can_dynsized=True)
+                          can_dynsized=True, alignment=alignment)
 
 
+# XXX-140: Should I just use types.Any for the last alignment arg?
+
 @lower(cuda.shared.array, types.Tuple, types.Any)
 @lower(cuda.shared.array, types.UniTuple, types.Any)
+@lower(cuda.shared.array, types.Tuple, types.Any, types.IntegerLiteral)
+@lower(cuda.shared.array, types.UniTuple, types.Any, types.IntegerLiteral)
+@lower(cuda.shared.array, types.Tuple, types.Any, types.NoneType)
+@lower(cuda.shared.array, types.UniTuple, types.Any, types.NoneType)
 def cuda_shared_array_tuple(context, builder, sig, args):
     shape = [ s.literal_value for s in sig.args[0] ]
     dtype = parse_dtype(sig.args[1])
+    alignment = None
+    if len(sig.args) == 3:
+        try:
+            alignment = sig.args[2].literal_value
+        except AttributeError:
+            pass
     return _generic_array(context, builder, shape=shape, dtype=dtype,
                           symbol_name=_get_unique_smem_id('_cudapy_smem'),
                           addrspace=nvvm.ADDRSPACE_SHARED,
-                          can_dynsized=True)
+                          can_dynsized=True, alignment=alignment)
 
 
 @lower(cuda.local.array, types.IntegerLiteral, types.Any)
+@lower(cuda.local.array, types.IntegerLiteral, types.Any, types.IntegerLiteral)
+@lower(cuda.local.array, types.IntegerLiteral, types.Any, types.NoneType)
 def cuda_local_array_integer(context, builder, sig, args):
     length = sig.args[0].literal_value
     dtype = parse_dtype(sig.args[1])
+    alignment = None
+    if len(sig.args) == 3:
+        try:
+            alignment = sig.args[2].literal_value
+        except AttributeError:
+            pass
     return _generic_array(context, builder, shape=(length,), dtype=dtype,
                           symbol_name='_cudapy_lmem',
                           addrspace=nvvm.ADDRSPACE_LOCAL,
-                          can_dynsized=False)
+                          can_dynsized=False, alignment=alignment)
 
 
 @lower(cuda.local.array, types.Tuple, types.Any)
@@ -954,7 +987,7 @@ def ptx_nanosleep(context, builder, sig, args):
 
 
 def _generic_array(context, builder, shape, dtype, symbol_name, addrspace,
-                   can_dynsized=False):
+                   can_dynsized=False, alignment=None):
     elemcount = reduce(operator.mul, shape, 1)
 
     # Check for valid shape for this type of allocation.
@@ -981,17 +1014,37 @@ def _generic_array(context, builder, shape, dtype, symbol_name, addrspace,
         # NVVM is smart enough to only use local memory if no register is
         # available
         dataptr = cgutils.alloca_once(builder, laryty, name=symbol_name)
+
+        # If the user has specified a custom alignment, just set the align
+        # attribute on the alloca IR directly.  We don't do any additional
+        # hand-holding here like checking the underlying data type's alignment
+        # or rounding up to the next power of 2--those checks will have already
+        # been done by the time we see the alignment value.
+        if alignment is not None:
+            dataptr.align = alignment
     else:
         lmod = builder.module
 
         # Create global variable in the requested address space
         gvmem = cgutils.add_global_variable(lmod, laryty, symbol_name,
                                             addrspace)
-        # Specify alignment to avoid misalignment bug
-        align = context.get_abi_sizeof(lldtype)
-        # Alignment is required to be a power of 2 for shared memory. If it is
-        # not a power of 2 (e.g. for a Record array) then round up accordingly.
-        gvmem.align = 1 << (align - 1 ).bit_length()
+
+        # If the caller hasn't specified a custom alignment, obtain the
+        # underlying dtype alignment from the ABI and then round it up to
+        # a power of two.  Otherwise, just use the caller's alignment.
+        #
+        # N.B. The caller *could* provide a valid-but-smaller-than-natural
+        #      alignment here; we'll assume the caller knows what they're
+        #      doing and let that through without error.
+
+        if alignment is None:
+            abi_alignment = context.get_abi_alignment(lldtype)
+            # Ensure a power of two alignment.
+            actual_alignment = 1 << (abi_alignment - 1).bit_length()
+        else:
+            actual_alignment = alignment
+
+        gvmem.align = actual_alignment
 
         if dynamic_smem:
             gvmem.linkage = 'external'
@@ -1041,7 +1094,8 @@ def _generic_array(context, builder, shape, dtype, symbol_name, addrspace,
 
     # Create array object
     ndim = len(shape)
-    aryty = types.Array(dtype=dtype, ndim=ndim, layout='C')
+    aryty = types.Array(dtype=dtype, ndim=ndim, layout='C',
+                        alignment=alignment)
     ary = context.make_array(aryty)(context, builder)
 
     context.populate_array(ary,

numba_cuda/numba/cuda/tests/cudapy/test_array_alignment.py

@@ -0,0 +1,114 @@
+import itertools
+import numpy as np
+from numba import cuda
+from numba.cuda.testing import CUDATestCase
+from numba.core.errors import TypingError
+import unittest
+
+
+# Set to true if you want to see dots printed for each subtest.
+NOISY = True
+
+
+# N.B. We name the test class TestArrayAddressAlignment to avoid name conflict
+#      with the test_alignment.TestArrayAlignment class.
+
+
+class TestArrayAddressAlignment(CUDATestCase):
+    """
+    Test cuda.local.array and cuda.shared.array support for an alignment
+    keyword argument.
+    """
+
+    def test_array_alignment(self):
+        shapes = (1, 3, 4, 8, 9, 50)
+        dtypes = (np.uint8, np.uint16, np.uint32, np.uint64)
+        alignments = (None, 8, 16, 32, 64, 128, 256)
+        array_types = [(0, 'local'), (1, 'shared')]
+
+        items = itertools.product(array_types, shapes, dtypes, alignments)
+
+        for (which, array_type), shape, dtype, alignment in items:
+            with self.subTest(array_type=array_type, shape=shape,
+                              dtype=dtype, alignment=alignment):
+                @cuda.jit
+                def f(loc, shrd, which):
+                    i = cuda.grid(1)
+                    if which == 0:
+                        local_array = cuda.local.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            loc[0] = local_array.ctypes.data
+                    else:
+                        shared_array = cuda.shared.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            shrd[0] = shared_array.ctypes.data
+
+                loc = np.zeros(1, dtype=np.uint64)
+                shrd = np.zeros(1, dtype=np.uint64)
+                f[1, 1](loc, shrd, which)
+
+                if alignment is not None:
+                    address = loc[0] if which == 0 else shrd[0]
+                    alignment_mod = int(address % alignment)
+                    self.assertEqual(alignment_mod, 0)
+
+                if NOISY:
+                    print('.', end='', flush=True)
+
+    def test_invalid_aligments(self):
+        shapes = (1, 3, 4, 8, 9, 50)
+        dtypes = (np.uint8, np.uint16, np.uint32, np.uint64)
+        alignments = (-1, 0, 3, 5, 7, 9, 15, 17, 31, 33, 63, 65)
+        array_types = [(0, 'local'), (1, 'shared')]
+
+        items = itertools.product(array_types, shapes, dtypes, alignments)
+
+        for (which, array_type), shape, dtype, alignment in items:
+            with self.subTest(array_type=array_type, shape=shape,
+                              dtype=dtype, alignment=alignment):
+                @cuda.jit
+                def f(local_array, shared_array, which):
+                    i = cuda.grid(1)
+                    if which == 0:
+                        local_array = cuda.local.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            local_array[0] = local_array.ctypes.data
+                    else:
+                        shared_array = cuda.shared.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            shared_array[0] = shared_array.ctypes.data
+
+                loc = np.zeros(1, dtype=np.uint64)
+                shrd = np.zeros(1, dtype=np.uint64)
+
+                with self.assertRaises(TypingError) as raises:
+                    f[1, 1](loc, shrd, which)
+                exc = str(raises.exception)
+                self.assertIn("Alignment must be", exc)
+
+                if NOISY:
+                    print('.', end='', flush=True)
+
+    def test_array_like(self):
+        # XXX-140: TODO; need to flush out the array_like stuff more.
+        pass
+
+
+if __name__ == '__main__':
+    unittest.main()