(numba/dppl) Copy ir & funcdesc using same cache. Introduce new function to copy IR. Do not copy non-copyable objects (#78)

Author: AlexanderKalistratov

dppl_lowerer.py

@@ -962,16 +962,189 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
+def relatively_deep_copy(obj, memo):
+    # WARNING: there are some issues with genarators which were not investigated and root cause is not found.
+    # Though copied IR seems to work fine there are some extra references kept on generator objects which may result
+    # in memory "leak"
+
+    obj_id = id(obj)
+    if obj_id in memo:
+        return memo[obj_id]
+
+    from numba.core.dispatcher import _DispatcherBase
+    from numba.core.types.functions import Function, Dispatcher
+    from numba.core.bytecode import FunctionIdentity
+    from numba.core.typing.templates import Signature
+    from numba.dppl.compiler import DPPLFunctionTemplate
+    from numba.core.compiler import CompileResult
+    from numba.np.ufunc.dufunc import DUFunc
+    from ctypes import _CFuncPtr
+    from cffi.api import FFI
+    from types import ModuleType
+    from numba.core.types.abstract import Type
+
+    # objects which shouldn't or can't be copied and it's ok not to copy it.
+    if isinstance(obj, (FunctionIdentity, _DispatcherBase, Function, Type, Dispatcher, ModuleType,
+                        Signature, DPPLFunctionTemplate, CompileResult,
+                        DUFunc, _CFuncPtr, FFI,
+                        type, str, bool, type(None))):
+        return obj
+
+    from numba.core.ir import Global, FreeVar
+    from numba.core.ir import FunctionIR
+    from numba.core.postproc import PostProcessor
+    from numba.core.funcdesc import FunctionDescriptor
+
+    if isinstance(obj, FunctionDescriptor):
+        cpy = FunctionDescriptor(native=obj.native, modname=obj.modname, qualname=obj.qualname,
+                                 unique_name=obj.unique_name, doc=obj.doc,
+                                 typemap=relatively_deep_copy(obj.typemap, memo),
+                                 restype=obj.restype,
+                                 calltypes=relatively_deep_copy(obj.calltypes, memo),
+                                 args=obj.args, kws=obj.kws, mangler=None,
+                                 argtypes=relatively_deep_copy(obj.argtypes, memo),
+                                 inline=obj.inline, noalias=obj.noalias, env_name=obj.env_name,
+                                 global_dict=obj.global_dict)
+        # mangler parameter is not saved in FunctionDescriptor, but used to generated name.
+        # So pass None as mangler parameter and then copy mangled_name by hands
+        cpy.mangled_name = obj.mangled_name
+
+        memo[obj_id] = cpy
+
+        return cpy
+
+    if isinstance(obj, FunctionIR):
+        # PostProcessor do the following:
+        # 1. canonicolize cfg, modifying IR
+        # 2. fills internal generators status
+        # 3. creates and fills VariableLifetime object
+        # We can't copy this objects. So in order to have copy of it we need run PostProcessor on copied IR.
+        # This means, that in case PostProcess wasn't run for original object copied object would defer.
+        # In order to avoid this we are running PostProcess on original object firstly.
+        # This means that copy of IR actually has a side effect on it.
+        pp = PostProcessor(obj)
+        pp.run()
+        cpy = FunctionIR(blocks=relatively_deep_copy(obj.blocks, memo),
+                         is_generator=relatively_deep_copy(obj.is_generator, memo),
+                         func_id=relatively_deep_copy(obj.func_id, memo),
+                         loc=obj.loc,
+                         definitions=relatively_deep_copy(obj._definitions, memo),
+                         arg_count=obj.arg_count,
+                         arg_names=relatively_deep_copy(obj.arg_names, memo))
+        pp = PostProcessor(cpy)
+        pp.run()
+
+        memo[obj_id] = cpy
+
+        return cpy
+
+    if isinstance(obj, Global):
+        cpy = Global(name=obj.name, value=obj.value, loc=obj.loc)
+        memo[obj_id] = cpy
+
+        return cpy
+
+    if isinstance(obj, FreeVar):
+        cpy = FreeVar(index=obj.index, name=obj.name, value=obj.value, loc=obj.loc)
+        memo[obj_id] = cpy
+
+        return cpy
+
+    # for containers we need to copy container itself first. And then fill it with copied items.
+    if isinstance(obj, list):
+        cpy = copy.copy(obj)
+        cpy.clear()
+        for item in obj:
+            cpy.append(relatively_deep_copy(item, memo))
+        memo[obj_id] = cpy
+        return cpy
+    elif isinstance(obj, dict):
+        cpy = copy.copy(obj)
+        cpy.clear()
+        for key, item in obj.items():
+            cpy[relatively_deep_copy(key, memo)] = relatively_deep_copy(item, memo)
+        memo[obj_id] = cpy
+        return cpy
+    elif isinstance(obj, tuple):
+        # subclass constructors could have different parameters than superclass.
+        # e.g. tuple and namedtuple constructors accepts quite different parameters.
+        # it is better to have separate section for namedtuple
+        tpl = tuple([relatively_deep_copy(item, memo) for item in obj])
+        if type(obj) == tuple:
+            cpy = tpl
+        else:
+            cpy = type(obj)(*tpl)
+        memo[obj_id] = cpy
+        return cpy
+    elif isinstance(obj, set):
+        cpy = copy.copy(obj)
+        cpy.clear()
+        for item in obj:
+            cpy.add(relatively_deep_copy(item, memo))
+        memo[obj_id] = cpy
+        return cpy
+
+    # some python objects are not copyable. In such case exception would be raised
+    # it is just a convinient point to find such objects
+    try:
+        cpy = copy.copy(obj)
+    except Exception as e:
+        raise e
+
+    # __slots__ for subclass specify only members declared in subclass. So to get all members we need to go through
+    # all supeclasses
+    def get_slots_members(obj):
+        keys = []
+        typ = obj
+        if not isinstance(typ, type):
+            typ = type(obj)
+
+        try:
+            if len(typ.__slots__):
+                keys.extend(typ.__slots__)
+            if len(typ.__bases__):
+                for base in typ.__bases__:
+                    keys.extend(get_slots_members(base))
+        except:
+            pass
+
+        return keys
+
+    memo[obj_id] = cpy
+    keys = []
+
+    # Objects have either __dict__ or __slots__ or neither.
+    # If object has none of it and it is copyable we already made a copy, just return it
+    # If object is not copyable we shouldn't reach this point.
+    try:
+        keys = obj.__dict__.keys()
+    except:
+        try:
+            obj.__slots__
+            keys = get_slots_members(obj)
+        except:
+            return cpy
+
+    for key in keys:
+        attr = getattr(obj, key)
+        attr_cpy = relatively_deep_copy(attr, memo)
+        setattr(cpy, key, attr_cpy)
+
+    return cpy
+
+
 class DPPLLower(Lower):
     def __init__(self, context, library, fndesc, func_ir, metadata=None):
         Lower.__init__(self, context, library, fndesc, func_ir, metadata)
-        fndesc_cpu = copy.copy(fndesc)
-        fndesc_cpu.calltypes = fndesc.calltypes.copy()
-        fndesc_cpu.typemap = fndesc.typemap.copy()
+        memo = {}
+
+        fndesc_cpu = relatively_deep_copy(fndesc, memo)
+        func_ir_cpu = relatively_deep_copy(func_ir, memo)
+
 
         cpu_context = context.cpu_context if isinstance(context, DPPLTargetContext) else context
-        self.gpu_lower = Lower(context, library, fndesc, func_ir.copy(), metadata)
-        self.cpu_lower = Lower(cpu_context, library, fndesc_cpu, func_ir.copy(), metadata)
+        self.gpu_lower = Lower(context, library, fndesc, func_ir, metadata)
+        self.cpu_lower = Lower(cpu_context, library, fndesc_cpu, func_ir_cpu, metadata)
 
     def lower(self):
         # Basically we are trying to lower on GPU first and if failed - try to lower on CPU.
@@ -991,11 +1164,9 @@ def lower(self):
         # different solution should be used.
 
         try:
-            #lowering.lower_extensions[parfor.Parfor] = lower_parfor_rollback
             lowering.lower_extensions[parfor.Parfor].append(lower_parfor_rollback)
             self.gpu_lower.lower()
             self.base_lower = self.gpu_lower
-            #lowering.lower_extensions[parfor.Parfor] = numba.parfors.parfor_lowering._lower_parfor_parallel
             lowering.lower_extensions[parfor.Parfor].pop()
         except Exception as e:
             if numba.dppl.compiler.DEBUG:
@@ -1015,80 +1186,13 @@ def create_cpython_wrapper(self, release_gil=False):
 
 
 def copy_block(block):
-    def relatively_deep_copy(obj, memo):
-        obj_id = id(obj)
-        if obj_id in memo:
-            return memo[obj_id]
-
-        from numba.core.dispatcher import Dispatcher
-        from numba.core.types.functions import Function
-        from types import ModuleType
-
-        if isinstance(obj, (Dispatcher, Function, ModuleType)):
-            return obj
-
-        if isinstance(obj, list):
-            cpy = copy.copy(obj)
-            cpy.clear()
-            for item in obj:
-                cpy.append(relatively_deep_copy(item, memo))
-            memo[obj_id] = cpy
-            return cpy
-        elif isinstance(obj, dict):
-            cpy = copy.copy(obj)
-            cpy.clear()
-            # do we need to copy keys?
-            for key, item in obj.items():
-                cpy[relatively_deep_copy(key, memo)] = relatively_deep_copy(item, memo)
-            memo[obj_id] = cpy
-            return cpy
-        elif isinstance(obj, tuple):
-            cpy = type(obj)([relatively_deep_copy(item, memo) for item in obj])
-            memo[obj_id] = cpy
-            return cpy
-        elif isinstance(obj, set):
-            cpy = copy.copy(obj)
-            cpy.clear()
-            for item in obj:
-                cpy.add(relatively_deep_copy(item, memo))
-            memo[obj_id] = cpy
-            return cpy
-
-        cpy = copy.copy(obj)
-
-        memo[obj_id] = cpy
-        keys = []
-        try:
-            keys = obj.__dict__.keys()
-        except:
-            try:
-                keys = obj.__slots__
-            except:
-                return cpy
-
-        for key in keys:
-            attr = getattr(obj, key)
-            attr_cpy = relatively_deep_copy(attr, memo)
-            setattr(cpy, key, attr_cpy)
-
-        return cpy
-
     memo = {}
     new_block = ir.Block(block.scope, block.loc)
     new_block.body = [relatively_deep_copy(stmt, memo) for stmt in block.body]
     return new_block
 
 
 def lower_parfor_rollback(lowerer, parfor):
-    try:
-        cache_parfor_races = copy.copy(parfor.races)
-        cache_parfor_params = copy.copy(parfor.params)
-        cache_parfor_loop_body = {key: copy_block(block) for key, block in parfor.loop_body.items()}
-        cache_parfor_init_block = parfor.init_block.copy()
-        cache_parfor_loop_nests = parfor.loop_nests.copy()
-    except Exception as e:
-        raise CopyIRException("Failed to copy IR") from e
-
     try:
         _lower_parfor_gufunc(lowerer, parfor)
         if numba.dppl.compiler.DEBUG:
@@ -1098,12 +1202,6 @@ def lower_parfor_rollback(lowerer, parfor):
         msg = "Failed to lower parfor on DPPL-device.\nTo see details set environment variable NUMBA_DPPL_DEBUG=1"
         warnings.warn(NumbaPerformanceWarning(msg, parfor.loc))
         raise e
-    finally:
-        parfor.params = cache_parfor_params
-        parfor.loop_body = cache_parfor_loop_body
-        parfor.init_block = cache_parfor_init_block
-        parfor.loop_nests = cache_parfor_loop_nests
-        parfor.races = cache_parfor_races
 
 
 def dppl_lower_array_expr(lowerer, expr):

tests/dppl/test_dppl_fallback.py

@@ -6,31 +6,20 @@
 from numba import dppl
 from numba.dppl.testing import unittest
 from numba.dppl.testing import DPPLTestCase
+from numba.tests.support import captured_stderr
 import dpctl
 import sys
 import io
 
 
 @unittest.skipUnless(dpctl.has_gpu_queues(), 'test only on GPU system')
 class TestDPPLFallback(DPPLTestCase):
-
-    def capture_stderr(self, func):
-        backup = sys.stderr
-        sys.stderr = io.StringIO()
-        result = func()
-        out = sys.stderr.getvalue()
-        sys.stderr.close()
-        sys.stderr = backup
-
-        return out, result
-
-    def test_dppl_fallback(self):
-
+    def test_dppl_fallback_inner_call(self):
         @numba.jit
         def fill_value(i):
             return i
 
-        def np_rand_fallback():
+        def inner_call_fallback():
             x = 10
             a = np.empty(shape=x, dtype=np.float32)
 
@@ -39,17 +28,28 @@ def np_rand_fallback():
 
             return a
 
-        def run_dppl():
-            dppl = numba.njit(parallel={'offload':True})(np_rand_fallback)
-            return dppl()
+        with captured_stderr() as msg:
+            dppl = numba.njit(parallel={'offload':True})(inner_call_fallback)
+            dppl_result = dppl()
+
+        ref_result = inner_call_fallback()
+
+        np.testing.assert_array_equal(dppl_result, ref_result)
+        self.assertTrue('Failed to lower parfor on DPPL-device' in msg.getvalue())
+
+    def test_dppl_fallback_reductions(self):
+        def reduction(a):
+            return np.amax(a)
 
-        ref = np_rand_fallback
+        a = np.ones(10)
+        with captured_stderr() as msg:
+            dppl = numba.njit(parallel={'offload':True})(reduction)
+            dppl_result = dppl(a)
 
-        err, dppl_result = self.capture_stderr(run_dppl)
-        ref_result = ref()
+        ref_result = reduction(a)
 
         np.testing.assert_array_equal(dppl_result, ref_result)
-        self.assertTrue('Failed to lower parfor on DPPL-device' in err)
+        self.assertTrue('Failed to lower parfor on DPPL-device' in msg.getvalue())
 
 
 if __name__ == '__main__':