diff --git a/numba/core/consts.py b/numba/core/consts.py index d062b3208..fe4cbc1d4 100644 --- a/numba/core/consts.py +++ b/numba/core/consts.py @@ -21,6 +21,9 @@ def __init__(self, func_ir): self._func_ir = weakref.proxy(func_ir) self._cache = {} + def __repr__(self): + return str(self._cache) + def infer_constant(self, name, loc=None): """ Infer a constant value for the given variable *name*. diff --git a/numba/core/dispatcher.py b/numba/core/dispatcher.py index 69414a7c5..56e3c5cab 100644 --- a/numba/core/dispatcher.py +++ b/numba/core/dispatcher.py @@ -1101,6 +1101,16 @@ def __init__(self, func_ir, typingctx, targetctx, flags, locals): can_fallback=True, exact_match_required=False) + class LCFuncIR: + """ This is a minimalistic reimplementation of Numba's regular FuncIR + class. We only need these two parts for caching. We either don't + have the information to fill in the regular FuncIR or there are + problems if we do. + """ + def __init__(self, arg_count, func_id): + self.arg_count = arg_count + self.func_id = func_id + def _reduce_states(self): """ Reduce the instance for pickling. This will serialize @@ -1110,8 +1120,8 @@ def _reduce_states(self): NOTE: part of ReduceMixin protocol """ return dict( - uuid=self._uuid, func_ir=self.func_ir, flags=self.flags, - locals=self.locals, extras=self._reduce_extras(), + uuid=self._uuid, func_ir=self.LCFuncIR(self.func_ir.arg_count, self.func_ir.func_id), + flags=self.flags, locals=self.locals, extras=self._reduce_extras(), ) def _reduce_extras(self): @@ -1227,7 +1237,10 @@ class LiftedWith(LiftedCode): can_cache = True def _reduce_extras(self): - return dict(output_types=self.output_types) + if hasattr(self, "output_types"): + return dict(output_types=self.output_types) + else: + return dict() @property def _numba_type_(self): diff --git a/numba/openmp.py b/numba/openmp.py index f9d165dce..1d86e0026 100644 --- a/numba/openmp.py +++ b/numba/openmp.py @@ -52,6 +52,7 @@ llvm_binpath=None llvm_libpath=None + def _init(): global llvm_binpath global llvm_libpath @@ -80,6 +81,7 @@ def _init(): _init() + #---------------------------------------------------------------------------------------------- class NameSlice: @@ -166,6 +168,17 @@ def __init__(self, name, arg=None, load=False, non_arg=False, omp_slice=None): self.non_arg = non_arg self.omp_slice = omp_slice + def __getstate__(self): + state = self.__dict__.copy() + if isinstance(self.arg, lir.instructions.AllocaInstr): + del state['arg'] + return state + + def __setstate__(self, state): + self.__dict__.update(state) + if not hasattr(self, "arg"): + self.arg = None + def var_in(self, var): assert isinstance(var, str) @@ -1120,6 +1133,13 @@ def __init__(self, tags, region_number, loc, firstprivate_dead_after=None): self.alloca_queue = [] self.end_region = None + def __getstate__(self): + state = self.__dict__.copy() + return state + + def __setstate__(self, state): + self.__dict__.update(state) + def replace_var_names(self, namedict): for i in range(len(self.tags)): if isinstance(self.tags[i].arg, ir.Var):