(numba/dppy) Adds support for local memory allocation inside dppy.kernel, clean ups to barrier, use context manager

* Added local memory allocation and fixed bugs in barrier implementations

* Added support for passing signature in dppy.kernel decorator

* Added bigger example that deals with local memory and barrier

* use barrier and local memory in atomic test

* Change local.alloc to local.static_alloc to indicate this is static allocation

* Added pass to ensure only constants can be given as shape to static local memory allocation

* Change message and fix typo

* Include pass

* Feature/dppy backend context manager (#85)

* Do not use gpu_env instead use the current_env specified in the runtime.

* Update all examples to use dppy context manager.

* Fix Test suites and update API to use dppy context manager

* Uncomment barrier

* Fix caching tests

Author: reazulhoque

compiler.py

@@ -8,6 +8,7 @@
 from numba.typing.templates import AbstractTemplate
 from numba import ctypes_support as ctypes
 from types import FunctionType
+from inspect import signature
 
 import dppy.core as driver
 from . import spirv_generator
@@ -31,6 +32,21 @@ def _raise_invalid_kernel_enqueue_args():
                      "The local size argument is optional.")
     raise ValueError(error_message)
 
+
+def get_ordered_arg_access_types(pyfunc, access_types):
+    # Construct a list of access type of each arg according to their position
+    ordered_arg_access_types = []
+    sig = signature(pyfunc, follow_wrapped=False)
+    for idx, arg_name in enumerate(sig.parameters):
+        if access_types:
+            for key in access_types:
+                if arg_name in access_types[key]:
+                    ordered_arg_access_types.append(key)
+        if len(ordered_arg_access_types) <= idx:
+            ordered_arg_access_types.append(None)
+
+    return ordered_arg_access_types
+
 class DPPyCompiler(CompilerBase):
     """ DPPy Compiler """
 
@@ -62,7 +78,8 @@ def compile_with_dppy(pyfunc, return_type, args, debug):
     # Do not compile (generate native code), just lower (to LLVM)
     flags.set('no_compile')
     flags.set('no_cpython_wrapper')
-    flags.set('nrt')
+    #flags.set('nrt')
+
     # Run compilation pipeline
     if isinstance(pyfunc, FunctionType):
         cres = compiler.compile_extra(typingctx=typingctx,
@@ -95,6 +112,9 @@ def compile_with_dppy(pyfunc, return_type, args, debug):
 def compile_kernel(device, pyfunc, args, access_types, debug=False):
     if DEBUG:
         print("compile_kernel", args)
+    if not device:
+        device = driver.runtime.get_current_device()
+
     cres = compile_with_dppy(pyfunc, None, args, debug=debug)
     func = cres.library.get_function(cres.fndesc.llvm_func_name)
     kernel = cres.target_context.prepare_ocl_kernel(func, cres.signature.args)
@@ -241,7 +261,8 @@ def _ensure_valid_work_group_size(val, work_item_grid):
         val = [val]
 
     if len(val) != len(work_item_grid):
-        error_message = ("Unsupported number of work item dimensions ")
+        error_message = ("Unsupported number of work item dimensions, " +
+                         "dimensions of global and local work items has to be the same ")
         raise ValueError(error_message)
 
     return list(val)
@@ -283,19 +304,19 @@ def forall(self, nelem, local_size=64, queue=None):
     def __getitem__(self, args):
         """Mimick CUDA python's square-bracket notation for configuration.
         This assumes the argument to be:
-            `device_env, global size, local size`
+            `global size, local size`
         """
         ls = None
         nargs = len(args)
         # Check if the kernel enquing arguments are sane
-        if nargs < 2 or nargs > 3:
+        if nargs < 1 or nargs > 2:
             _raise_invalid_kernel_enqueue_args
 
-        device_env = args[0]
-        gs = _ensure_valid_work_item_grid(args[1], device_env)
+        device_env = driver.runtime.get_current_device()
+        gs = _ensure_valid_work_item_grid(args[0], device_env)
         # If the optional local size argument is provided
-        if nargs == 3:
-            ls = _ensure_valid_work_group_size(args[2], gs)
+        if nargs == 2 and args[1] != []:
+            ls = _ensure_valid_work_group_size(args[1], gs)
 
         return self.configure(device_env, gs, ls)
 
@@ -400,7 +421,7 @@ def _unpack_argument(self, ty, val, device_env, retr, kernelargs,
                 dArr = device_env.copy_array_to_device(val)
             elif self.valid_access_types[access_type] == _NUMBA_PVC_WRITE_ONLY:
                 # write_only case, we do not copy the host data
-                dArr = driver.DeviceArray(device_env.get_env_ptr(), val)
+                dArr = device_env.create_device_array(val)
 
             assert (dArr != None), "Problem in allocating device buffer"
             device_arrs[-1] = dArr
@@ -466,7 +487,7 @@ def __call__(self, *args, **kwargs):
         assert not kwargs, "Keyword Arguments are not supported"
         if self.device_env is None:
             try:
-                self.device_env = driver.runtime.get_gpu_device()
+                self.device_env = driver.runtime.get_current_device()
             except:
                 _raise_no_device_found_error()
 

decorators.py

@@ -1,8 +1,7 @@
 from __future__ import print_function, absolute_import, division
 from numba import sigutils, types
-from .compiler import (compile_kernel, JitDPPyKernel,
-                       compile_dppy_func_template, compile_dppy_func)
-from inspect import signature
+from .compiler import (compile_kernel, JitDPPyKernel, compile_dppy_func_template,
+                       compile_dppy_func, get_ordered_arg_access_types)
 
 
 def kernel(signature=None, access_types=None, debug=False):
@@ -17,34 +16,25 @@ def kernel(signature=None, access_types=None, debug=False):
         func = signature
         return autojit(debug=False, access_types=access_types)(func)
     else:
-        return _kernel_jit(signature, debug)
+        return _kernel_jit(signature, debug, access_types)
 
 
 def autojit(debug=False, access_types=None):
     def _kernel_autojit(pyfunc):
-        # Construct a list of access type of each arg according to their position
-        ordered_arg_access_types = []
-        sig = signature(pyfunc, follow_wrapped=False)
-        for idx, arg_name in enumerate(sig.parameters):
-            if access_types:
-                for key in access_types:
-                    if arg_name in access_types[key]:
-                        ordered_arg_access_types.append(key)
-            if len(ordered_arg_access_types) <= idx:
-                ordered_arg_access_types.append(None)
-
+        ordered_arg_access_types = get_ordered_arg_access_types(pyfunc, access_types)
         return JitDPPyKernel(pyfunc, ordered_arg_access_types)
     return _kernel_autojit
 
 
-def _kernel_jit(signature, debug):
+def _kernel_jit(signature, debug, access_types):
     argtypes, restype = sigutils.normalize_signature(signature)
     if restype is not None and restype != types.void:
         msg = ("DPPy kernel must have void return type but got {restype}")
         raise TypeError(msg.format(restype=restype))
 
     def _wrapped(pyfunc):
-        return compile_kernel(pyfunc, argtypes, debug)
+        ordered_arg_access_types = get_ordered_arg_access_types(pyfunc, access_types)
+        return compile_kernel(None, pyfunc, argtypes, ordered_arg_access_types, debug)
 
     return _wrapped
 

device_init.py

@@ -11,11 +11,15 @@
     get_num_groups,
     barrier,
     mem_fence,
-#    shared,
     sub_group_barrier,
     atomic,
+    local,
+    CLK_LOCAL_MEM_FENCE,
+    CLK_GLOBAL_MEM_FENCE,
 )
 
+DEFAULT_LOCAL_SIZE = []
+
 from . import initialize
 
 from .decorators import kernel, func, autojit

dppy_host_fn_call_gen.py

@@ -15,10 +15,10 @@ def __init__(self, lowerer, cres, num_inputs):
         self.context = self.lowerer.context
         self.builder = self.lowerer.builder
 
-        self.gpu_device = driver.runtime.get_gpu_device()
-        self.gpu_device_env = self.gpu_device.get_env_ptr()
-        self.gpu_device_int = int(driver.ffi.cast("uintptr_t",
-                                                  self.gpu_device_env))
+        self.current_device = driver.runtime.get_current_device()
+        self.current_device_env = self.current_device.get_env_ptr()
+        self.current_device_int = int(driver.ffi.cast("uintptr_t",
+                                                  self.current_device_env))
 
         self.kernel_t_obj = cres.kernel._kernel_t_obj[0]
         self.kernel_int = int(driver.ffi.cast("uintptr_t",
@@ -65,8 +65,9 @@ def _init_llvm_types_and_constants(self):
         self.void_ptr_t      = self.context.get_value_type(types.voidptr)
         self.void_ptr_ptr_t  = lc.Type.pointer(self.void_ptr_t)
         self.sizeof_void_ptr = self.context.get_abi_sizeof(self.intp_t)
-        self.gpu_device_int_const = self.context.get_constant(
-                                        types.uintp, self.gpu_device_int)
+        self.current_device_int_const = self.context.get_constant(
+                                            types.uintp,
+                                            self.current_device_int)
 
     def _declare_functions(self):
         create_dppy_kernel_arg_fnty = lc.Type.function(
@@ -173,7 +174,7 @@ def process_kernel_arg(self, var, llvm_arg, arg_type, gu_sig, val_type, index, m
                                              size=self.one, name=buffer_name)
 
             # env, buffer_size, buffer_ptr
-            args = [self.builder.inttoptr(self.gpu_device_int_const, self.void_ptr_t),
+            args = [self.builder.inttoptr(self.current_device_int_const, self.void_ptr_t),
                     self.builder.load(total_size),
                     buffer_ptr]
             self.builder.call(self.create_dppy_rw_mem_buffer, args)
@@ -186,7 +187,7 @@ def process_kernel_arg(self, var, llvm_arg, arg_type, gu_sig, val_type, index, m
 
             # We really need to detect when an array needs to be copied over
             if index < self.num_inputs:
-                args = [self.builder.inttoptr(self.gpu_device_int_const, self.void_ptr_t),
+                args = [self.builder.inttoptr(self.current_device_int_const, self.void_ptr_t),
                         self.builder.load(buffer_ptr),
                         self.one,
                         self.zero,
@@ -263,7 +264,7 @@ def enqueue_kernel_and_read_back(self, loop_ranges):
             self.builder.store(stop,
                                self.builder.gep(dim_stops, [self.context.get_constant(types.uintp, i)]))
 
-        args = [self.builder.inttoptr(self.gpu_device_int_const, self.void_ptr_t),
+        args = [self.builder.inttoptr(self.current_device_int_const, self.void_ptr_t),
                 self.builder.inttoptr(self.context.get_constant(types.uintp, self.kernel_int), self.void_ptr_t),
                 self.context.get_constant(types.uintp, self.total_kernel_args),
                 self.kernel_arg_array,
@@ -276,7 +277,7 @@ def enqueue_kernel_and_read_back(self, loop_ranges):
         # read buffers back to host
         for read_buf in self.read_bufs_after_enqueue:
             buffer_ptr, array_size_member, data_member = read_buf
-            args = [self.builder.inttoptr(self.gpu_device_int_const, self.void_ptr_t),
+            args = [self.builder.inttoptr(self.current_device_int_const, self.void_ptr_t),
                     self.builder.load(buffer_ptr),
                     self.one,
                     self.zero,

dppy_lowerer.py

@@ -618,7 +618,7 @@ def print_arg_with_addrspaces(args):
     # FIXME : We should not always use gpu device, instead select the default
     # device as configured in dppy.
     kernel_func = numba.dppy.compiler.compile_kernel_parfor(
-        driver.runtime.get_gpu_device(),
+        driver.runtime.get_current_device(),
         gufunc_ir,
         gufunc_param_types,
         param_types_addrspaces)

dppy_passbuilder.py

@@ -17,6 +17,7 @@
                                 InlineOverloads)
 
 from .dppy_passes import (
+        DPPyConstantSizeStaticLocalMemoryPass,
         DPPyPreParforPass,
         DPPyParforPass,
         SpirvFriendlyLowering,
@@ -40,6 +41,11 @@ def default_numba_nopython_pipeline(state, pm):
         pm.add_pass(IRProcessing, "processing IR")
         pm.add_pass(WithLifting, "Handle with contexts")
 
+        # Add pass to ensure when users are allocating static
+        # constant memory the size is a constant and can not
+        # come from a closure variable
+        pm.add_pass(DPPyConstantSizeStaticLocalMemoryPass, "dppy constant size for static local memory")
+
         # pre typing
         if not state.flags.no_rewrites:
             pm.add_pass(RewriteSemanticConstants, "rewrite semantic constants")

dppy_passes.py

@@ -2,6 +2,7 @@
 from contextlib import contextmanager
 import warnings
 
+from numba import ir
 import weakref
 from collections import namedtuple, deque
 import operator
@@ -32,7 +33,81 @@
 from numba.parfor import PreParforPass as _parfor_PreParforPass
 from numba.parfor import ParforPass as _parfor_ParforPass
 from numba.parfor import Parfor
-#from numba.npyufunc.dufunc import DUFunc
+
+
+@register_pass(mutates_CFG=True, analysis_only=False)
+class DPPyConstantSizeStaticLocalMemoryPass(FunctionPass):
+
+    _name = "dppy_constant_size_static_local_memory_pass"
+
+    def __init__(self):
+        FunctionPass.__init__(self)
+
+    def run_pass(self, state):
+        """
+        Preprocessing for data-parallel computations.
+        """
+        # Ensure we have an IR and type information.
+        assert state.func_ir
+        func_ir = state.func_ir
+
+        _DEBUG = False
+
+        if _DEBUG:
+            print('Checks if size of OpenCL local address space alloca is a compile-time constant.'.center(80, '-'))
+            print(func_ir.dump())
+
+        work_list = list(func_ir.blocks.items())
+        while work_list:
+            label, block = work_list.pop()
+            for i, instr in enumerate(block.body):
+                if isinstance(instr, ir.Assign):
+                    expr = instr.value
+                    if isinstance(expr, ir.Expr):
+                        if expr.op == 'call':
+                            call_node = block.find_variable_assignment(expr.func.name).value
+                            if isinstance(call_node, ir.Expr) and call_node.attr == "static_alloc":
+                                arg = None
+                                # at first look in keyword arguments to get the shape, which has to be
+                                # constant
+                                if expr.kws:
+                                    for _arg in expr.kws:
+                                        if _arg[0] == "shape":
+                                            arg = _arg[1]
+
+                                if not arg:
+                                    arg = expr.args[0]
+
+                                error = False
+                                # arg can be one constant or a tuple of constant items
+                                arg_type = func_ir.get_definition(arg.name)
+                                if isinstance(arg_type, ir.Expr):
+                                    # we have a tuple
+                                    for item in arg_type.items:
+                                        if not isinstance(func_ir.get_definition(item.name), ir.Const):
+                                            error = True
+                                            break
+
+                                else:
+                                    if not isinstance(func_ir.get_definition(arg.name), ir.Const):
+                                        error = True
+                                        break
+
+                                if error:
+                                    warnings.warn_explicit("The size of the Local memory has to be constant",
+                                                           errors.NumbaError,
+                                                           state.func_id.filename,
+                                                           state.func_id.firstlineno)
+                                    raise
+
+
+
+        if config.DEBUG or config.DUMP_IR:
+            name = state.func_ir.func_id.func_qualname
+            print(("IR DUMP: %s" % name).center(80, "-"))
+            state.func_ir.dump()
+
+        return True
 
 
 @register_pass(mutates_CFG=True, analysis_only=False)