WIP changes...

Author: Diptorup Deb

numba_dppy/compiler.py

@@ -28,7 +28,7 @@
 import dpctl.program as dpctl_prog
 import numpy as np
 
-from . import spirv_generator
+from . import spirv_generator, target
 
 from numba.core.compiler import DefaultPassBuilder, CompilerBase
 from numba_dppy.dppy_parfor_diagnostics import ExtendedParforDiagnostics
@@ -143,15 +143,15 @@ def compile_kernel(sycl_queue, pyfunc, args, access_types, debug=False):
     if not sycl_queue:
         # This will be get_current_queue
         sycl_queue = dpctl.get_current_queue()
-    breakpoint()
     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)
-    # The kernel objet should have a reference to the target context it is compiled for.
-    # This is needed as we intend to shape the behavior of the kernel down the line
-    # depending on the target context. For example, we want to link our kernel object
-    # with implementation containing atomic operations only when atomic operations
-    # are being used in the kernel.
+    kernel = cres.library.get_function(cres.fndesc.llvm_func_name)
+    breakpoint()
+    target.set_dppy_kernel(kernel)
+    # kernel = cres.target_context.prepare_ocl_kernel(func, cres.signature.args)
+    # A reference to the target context is stored in the DPPYKernel to
+    # reference the context later in code generation. For example, we link
+    # the kernel object with a spir_func defining atomic operations only
+    # when atomic operations are used in the kernel.
     oclkern = DPPYKernel(
         context=cres.target_context,
         sycl_queue=sycl_queue,
@@ -170,7 +170,6 @@ def compile_kernel_parfor(sycl_queue, func_ir, args, args_with_addrspaces, debug
             print(a, type(a))
             if isinstance(a, types.npytypes.Array):
                 print("addrspace:", a.addrspace)
-
     cres = compile_with_dppy(func_ir, None, args_with_addrspaces, debug=debug)
     func = cres.library.get_function(cres.fndesc.llvm_func_name)
 
@@ -269,7 +268,7 @@ def __init__(self, cres):
 def _ensure_valid_work_item_grid(val, sycl_queue):
 
     if not isinstance(val, (tuple, list, int)):
-        error_message = "Cannot create work item dimension from " "provided argument"
+        error_message = "Cannot create work item dimension from provided argument"
         raise ValueError(error_message)
 
     if isinstance(val, int):
@@ -290,7 +289,7 @@ def _ensure_valid_work_item_grid(val, sycl_queue):
 def _ensure_valid_work_group_size(val, work_item_grid):
 
     if not isinstance(val, (tuple, list, int)):
-        error_message = "Cannot create work item dimension from " "provided argument"
+        error_message = "Cannot create work item dimension from provided argument"
         raise ValueError(error_message)
 
     if isinstance(val, int):

numba_dppy/dppy_array_type.py

@@ -84,8 +84,8 @@ class DPPYArrayModel(StructModel):
     def __init__(self, dmm, fe_type):
         ndim = fe_type.ndim
         members = [
-            ("meminfo", types.MemInfoPointer(fe_type.dtype)),
-            ("parent", types.pyobject),
+            ("meminfo", types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace)),
+            ("parent", types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace)),
             ("nitems", types.intp),
             ("itemsize", types.intp),
             ("data", types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace)),

numba_dppy/dppy_lowerer.py

@@ -56,7 +56,7 @@
 import dpctl
 from numba_dppy.target import DPPYTargetContext
 from numba_dppy.dppy_array_type import DPPYArray
-from numba_dppy.utils.constants import address_space
+from numba_dppy.utils import address_space
 
 
 def _print_block(block):
@@ -358,7 +358,7 @@ def addrspace_from(params, def_addr):
                 addrspaces.append(None)
         return addrspaces
 
-    addrspaces = addrspace_from(parfor_params, address_space.SPIR_GLOBAL)
+    addrspaces = addrspace_from(parfor_params, address_space.GLOBAL)
 
     if config.DEBUG_ARRAY_OPT >= 1:
         print("parfor_params = ", parfor_params, type(parfor_params))

numba_dppy/ocl/ocldecl.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from numba import types
-from numba.core.typing.npydecl import register_number_classes, parse_dtype, parse_shape
+from numba.core.typing.npydecl import parse_dtype, parse_shape
 from numba.core.typing.templates import (
     AttributeTemplate,
     ConcreteTemplate,
@@ -22,8 +22,7 @@
     signature,
     Registry,
 )
-import numba_dppy, numba_dppy as dppy
-from numba_dppy import target
+import numba_dppy as dppy
 from numba_dppy.dppy_array_type import DPPYArray
 from numba_dppy.utils import address_space
 
@@ -169,7 +168,7 @@ def typer(shape, dtype):
                     dtype=nb_dtype,
                     ndim=ndim,
                     layout="C",
-                    addrspace=address_space.SPIR_LOCAL,
+                    addrspace=address_space.LOCAL,
                 )
 
         return typer

numba_dppy/ocl/oclimpl.py

@@ -29,7 +29,7 @@
 from numba_dppy.codegen import SPIR_DATA_LAYOUT
 from numba_dppy.dppy_array_type import DPPYArray
 from numba_dppy.ocl.atomics import atomic_helper
-from numba_dppy.utils.constants import address_space
+from numba_dppy.utils import address_space
 
 from . import stubs
 
@@ -207,14 +207,14 @@ def insert_and_call_atomic_fn(
     else:
         raise TypeError("Atomic operation is not supported for type %s" % (dtype.name))
 
-    if addrspace == address_space.SPIR_LOCAL:
+    if addrspace == address_space.LOCAL:
         name = name + "_local"
     else:
         name = name + "_global"
 
     assert ll_p != None
     assert name != ""
-    ll_p.addrspace = address_space.SPIR_GENERIC
+    ll_p.addrspace = address_space.GENERIC
 
     mod = builder.module
     if sig.return_type == types.void:
@@ -228,7 +228,7 @@ def insert_and_call_atomic_fn(
     fn = mod.get_or_insert_function(fnty, name)
     fn.calling_convention = target.CC_SPIR_FUNC
 
-    generic_ptr = context.addrspacecast(builder, ptr, address_space.SPIR_GENERIC)
+    generic_ptr = context.addrspacecast(builder, ptr, address_space.GENERIC)
 
     return builder.call(fn, [generic_ptr, val])
 
@@ -272,10 +272,10 @@ def native_atomic_add(context, builder, sig, args):
 
     ptr_type = context.get_value_type(dtype).as_pointer()
     if not hasattr(aryty, "addrspace"):
-        ptr_type.addrspace = address_space.SPIR_GLOBAL
-        ptr = context.addrspacecast(builder, ptr, address_space.SPIR_GLOBAL)
+        ptr_type.addrspace = address_space.GLOBAL
+        ptr = context.addrspacecast(builder, ptr, address_space.GLOBAL)
     else:
-        ptr_type.addrspace = address_space.SPIR_LOCAL
+        ptr_type.addrspace = address_space.LOCAL
     retty = context.get_value_type(sig.return_type)
     spirv_fn_arg_types = [
         ptr_type,
@@ -409,7 +409,7 @@ def atomic_add(context, builder, sig, args, name):
         lary = context.make_array(aryty)(context, builder, ary)
         ptr = cgutils.get_item_pointer(context, builder, aryty, lary, indices)
 
-        if isinstance(aryty, DPPYArray) and aryty.addrspace == address_space.SPIR_LOCAL:
+        if isinstance(aryty, DPPYArray) and aryty.addrspace == address_space.LOCAL:
             return insert_and_call_atomic_fn(
                 context,
                 builder,
@@ -418,7 +418,7 @@ def atomic_add(context, builder, sig, args, name):
                 dtype,
                 ptr,
                 val,
-                address_space.SPIR_LOCAL,
+                address_space.LOCAL,
             )
         else:
             return insert_and_call_atomic_fn(
@@ -429,7 +429,7 @@ def atomic_add(context, builder, sig, args, name):
                 dtype,
                 ptr,
                 val,
-                address_space.SPIR_GLOBAL,
+                address_space.GLOBAL,
             )
     else:
         raise ImportError("Atomic support is not present, can not perform atomic_add")
@@ -445,7 +445,7 @@ def dppy_local_array_integer(context, builder, sig, args):
         shape=(length,),
         dtype=dtype,
         symbol_name="_dppy_lmem",
-        addrspace=address_space.SPIR_LOCAL,
+        addrspace=address_space.LOCAL,
     )
 
 
@@ -460,7 +460,7 @@ def dppy_local_array_tuple(context, builder, sig, args):
         shape=shape,
         dtype=dtype,
         symbol_name="_dppy_lmem",
-        addrspace=address_space.SPIR_LOCAL,
+        addrspace=address_space.LOCAL,
     )
 
 
@@ -473,7 +473,7 @@ def _generic_array(context, builder, shape, dtype, symbol_name, addrspace):
     lldtype = context.get_data_type(dtype)
     laryty = Type.array(lldtype, elemcount)
 
-    if addrspace == address_space.SPIR_LOCAL:
+    if addrspace == address_space.LOCAL:
         lmod = builder.module
 
         # Create global variable in the requested address-space
@@ -506,7 +506,7 @@ def _make_array(
     dtype,
     shape,
     layout="C",
-    addrspace=address_space.SPIR_GENERIC,
+    addrspace=address_space.GENERIC,
 ):
     ndim = len(shape)
     # Create array object

numba_dppy/printimpl.py

@@ -19,14 +19,14 @@
 from numba.core import types, typing, cgutils
 from numba.core.imputils import Registry
 
-from numba_dppy.utils.constants import address_space
+from numba_dppy.utils import address_space
 
 registry = Registry()
 lower = registry.lower
 
 
 def declare_print(lmod):
-    voidptrty = lc.Type.pointer(lc.Type.int(8), addrspace=address_space.SPIR_GENERIC)
+    voidptrty = lc.Type.pointer(lc.Type.int(8), addrspace=address_space.GENERIC)
     printfty = lc.Type.function(lc.Type.int(), [voidptrty], var_arg=True)
     printf = lmod.get_or_insert_function(printfty, "printf")
     return printf