66import warnings
77
88from numba .core import config , serialize , sigutils , types , typing , utils
9+ from numba .core .caching import Cache , CacheImpl
910from numba .core .compiler_lock import global_compiler_lock
1011from numba .core .dispatcher import Dispatcher
1112from numba .core .errors import NumbaPerformanceWarning , NumbaDeprecationWarning
@@ -106,23 +107,50 @@ def __init__(self, py_func, argtypes, link=None, debug=False,
106107 self ._codelibrary = lib
107108 self .call_helper = cres .call_helper
108109
110+ # The following are referred to by the cache implementation. Note:
111+ # - There are no referenced environments in CUDA.
112+ # - Kernels don't have lifted code.
113+ # - reload_init is only for parfors.
114+ self .target_context = tgt_ctx
115+ self .fndesc = cres .fndesc
116+ self .environment = cres .environment
117+ self ._referenced_environments = []
118+ self .lifted = []
119+ self .reload_init = []
120+
121+ @property
122+ def library (self ):
123+ return self ._codelibrary
124+
125+ @property
126+ def type_annotation (self ):
127+ return self ._type_annotation
128+
129+ def _find_referenced_environments (self ):
130+ return self ._referenced_environments
131+
132+ @property
133+ def codegen (self ):
134+ return self .target_context .codegen ()
135+
109136 @property
110137 def argument_types (self ):
111138 return tuple (self .signature .args )
112139
113140 @classmethod
114- def _rebuild (cls , cooperative , name , argtypes , codelibrary , link , debug ,
115- lineinfo , call_helper , extensions ):
141+ def _rebuild (cls , cooperative , name , signature , codelibrary ,
142+ debug , lineinfo , call_helper , extensions ):
116143 """
117144 Rebuild an instance.
118145 """
119146 instance = cls .__new__ (cls )
120147 # invoke parent constructor
121148 super (cls , instance ).__init__ ()
122149 # populate members
150+ instance .entry_point = None
123151 instance .cooperative = cooperative
124152 instance .entry_name = name
125- instance .argument_types = tuple ( argtypes )
153+ instance .signature = signature
126154 instance ._type_annotation = None
127155 instance ._codelibrary = codelibrary
128156 instance .debug = debug
@@ -140,7 +168,7 @@ def _reduce_states(self):
140168 Stream information is discarded.
141169 """
142170 return dict (cooperative = self .cooperative , name = self .entry_name ,
143- argtypes = self .argtypes , codelibrary = self .codelibrary ,
171+ signature = self .signature , codelibrary = self ._codelibrary ,
144172 debug = self .debug , lineinfo = self .lineinfo ,
145173 call_helper = self .call_helper , extensions = self .extensions )
146174
@@ -479,6 +507,31 @@ def __call__(self, *args):
479507 self .stream , self .sharedmem )
480508
481509
510+ class CUDACacheImpl (CacheImpl ):
511+ def reduce (self , kernel ):
512+ return kernel ._reduce_states ()
513+
514+ def rebuild (self , target_context , payload ):
515+ return _Kernel ._rebuild (** payload )
516+
517+ def check_cachable (self , cres ):
518+ # CUDA Kernels are always cachable - the reasons for an entity not to
519+ # be cachable are:
520+ #
521+ # - The presence of lifted loops, or
522+ # - The presence of dynamic globals.
523+ #
524+ # neither of which apply to CUDA kernels.
525+ return True
526+
527+
528+ class CUDACache (Cache ):
529+ """
530+ Implements a cache that saves and loads CUDA kernels and compile results.
531+ """
532+ _impl_class = CUDACacheImpl
533+
534+
482535class CUDADispatcher (Dispatcher , serialize .ReduceMixin ):
483536 '''
484537 CUDA Dispatcher object. When configured and called, the dispatcher will
@@ -517,6 +570,9 @@ def __init__(self, py_func, targetoptions, pipeline_class=CUDACompiler):
517570 def _numba_type_ (self ):
518571 return cuda_types .CUDADispatcher (self )
519572
573+ def enable_caching (self ):
574+ self ._cache = CUDACache (self .py_func )
575+
520576 @functools .lru_cache (maxsize = 128 )
521577 def configure (self , griddim , blockdim , stream = 0 , sharedmem = 0 ):
522578 griddim , blockdim = normalize_kernel_dimensions (griddim , blockdim )
@@ -718,30 +774,45 @@ def compile_device(self, args):
718774
719775 return cres
720776
777+ def add_overload (self , kernel , argtypes ):
778+ c_sig = [a ._code for a in argtypes ]
779+ self ._insert (c_sig , kernel , cuda = True )
780+ self .overloads [argtypes ] = kernel
781+
721782 def compile (self , sig ):
722783 '''
723784 Compile and bind to the current context a version of this kernel
724785 specialized for the given signature.
725786 '''
726787 argtypes , return_type = sigutils .normalize_signature (sig )
727788 assert return_type is None or return_type == types .none
789+
790+ # Do we already have an in-memory compiled kernel?
728791 if self .specialized :
729792 return next (iter (self .overloads .values ()))
730793 else :
731794 kernel = self .overloads .get (argtypes )
732- if kernel is None :
795+ if kernel is not None :
796+ return kernel
797+
798+ # Can we load from the disk cache?
799+ kernel = self ._cache .load_overload (sig , self .targetctx )
800+
801+ if kernel is not None :
802+ self ._cache_hits [sig ] += 1
803+ else :
804+ # We need to compile a new kernel
805+ self ._cache_misses [sig ] += 1
733806 if not self ._can_compile :
734807 raise RuntimeError ("Compilation disabled" )
735- kernel = _Kernel (self .py_func , argtypes ,
736- ** self .targetoptions )
737- # Inspired by _DispatcherBase.add_overload, but differs slightly
738- # because we're inserting a _Kernel object instead of a compiled
739- # function.
740- c_sig = [a ._code for a in argtypes ]
741- self ._insert (c_sig , kernel , cuda = True )
742- self .overloads [argtypes ] = kernel
743808
809+ kernel = _Kernel (self .py_func , argtypes , ** self .targetoptions )
810+ # We call bind to force codegen, so that there is a cubin to cache
744811 kernel .bind ()
812+ self ._cache .save_overload (sig , kernel )
813+
814+ self .add_overload (kernel , argtypes )
815+
745816 return kernel
746817
747818 def inspect_llvm (self , signature = None ):
0 commit comments