Skip to content

Commit 7147ce1

Browse files
authored
Merge pull request #8089 from gmarkall/cuda-caching-2-rebase
Support on-disk caching in the CUDA target
2 parents 2770647 + 944a531 commit 7147ce1

9 files changed

Lines changed: 773 additions & 19 deletions

File tree

‎docs/source/cuda/caching.rst‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
On-disk Kernel Caching
2+
======================
3+
4+
When the ``cache`` keyword argument of the :func:`@cuda.jit <numba.cuda.jit>`
5+
decorator is ``True``, a file-based cache is enabled. This shortens compilation
6+
times when the function was already compiled in a previous invocation.
7+
8+
The cache is maintained in the ``__pycache__`` subdirectory of the directory
9+
containing the source file; if the current user is not allowed to write to it,
10+
the cache implementation falls back to a platform-specific user-wide cache
11+
directory (such as ``$HOME/.cache/numba`` on Unix platforms).
12+
13+
14+
Compute capability considerations
15+
---------------------------------
16+
17+
Separate cache files are maintained for each compute capability. When a cached
18+
kernel is loaded, the compute capability of the device the kernel is first
19+
launched on in the current run is used to determine which version to load.
20+
Therefore, on systems that have multiple GPUs with differing compute
21+
capabilities, the cached versions of kernels are only used for one compute
22+
capability, and recompilation will occur for other compute capabilities.
23+
24+
For example: if a system has two GPUs, one of compute capability 7.5 and one of
25+
8.0, then:
26+
27+
* If a cached kernel is first launched on the CC 7.5 device, then the cached
28+
version for CC 7.5 is used. If it is subsequently launched on the CC 8.0
29+
device, a recompilation will occur.
30+
* If in a subsequent run the cached kernel is first launched on the CC 8.0
31+
device, then the cached version for CC 8.0 is used. A subsequent launch on
32+
the CC 7.5 device will require a recompilation.
33+
34+
This limitation is not expected to present issues in most practical scenarios,
35+
as multi-GPU production systems tend to have identical GPUs within each node.

‎docs/source/cuda/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,5 @@ Numba for CUDA GPUs
2525
external-memory.rst
2626
bindings.rst
2727
cuda_ffi.rst
28+
caching.rst
2829
faq.rst

‎numba/cuda/codegen.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from numba.core import config, serialize
66
from numba.core.codegen import Codegen, CodeLibrary
77
from numba.core.errors import NumbaInvalidConfigWarning
8-
from .cudadrv import devices, driver, nvvm
8+
from .cudadrv import devices, driver, nvvm, runtime
99

1010
import ctypes
1111
import numpy as np
@@ -302,7 +302,7 @@ def _reduce_states(self):
302302
'libraries to link against')
303303
raise RuntimeError(msg)
304304
return dict(
305-
codegen=self._codegen,
305+
codegen=None,
306306
name=self.name,
307307
entry_name=self._entry_name,
308308
module=self._module,
@@ -337,6 +337,8 @@ def _rebuild(cls, codegen, name, entry_name, module, linking_libraries,
337337
instance._max_registers = max_registers
338338
instance._nvvm_options = nvvm_options
339339

340+
return instance
341+
340342

341343
class JITCUDACodegen(Codegen):
342344
"""
@@ -360,3 +362,11 @@ def _create_empty_module(self, name):
360362

361363
def _add_module(self, module):
362364
pass
365+
366+
def magic_tuple(self):
367+
"""
368+
Return a tuple unambiguously describing the codegen behaviour.
369+
"""
370+
ctx = devices.get_context()
371+
cc = ctx.device.compute_capability
372+
return (runtime.runtime.get_version(), cc)

‎numba/cuda/decorators.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313

1414
def jit(func_or_sig=None, device=False, inline=False, link=[], debug=None,
15-
opt=True, **kws):
15+
opt=True, cache=False, **kws):
1616
"""
1717
JIT compile a python function conforming to the CUDA Python specification.
1818
If a signature is supplied, then a function is returned that takes a
@@ -49,6 +49,8 @@ def jit(func_or_sig=None, device=False, inline=False, link=[], debug=None,
4949
assembly code. This enables inspection of the source code in NVIDIA
5050
profiling tools and correlation with program counter sampling.
5151
:type lineinfo: bool
52+
:param cache: If True, enables the file-based cache for this function.
53+
:type cache: bool
5254
"""
5355

5456
if link and config.ENABLE_CUDASIM:
@@ -102,6 +104,9 @@ def _jit(func):
102104

103105
disp = CUDADispatcher(func, targetoptions=targetoptions)
104106

107+
if cache:
108+
disp.enable_caching()
109+
105110
if device:
106111
from numba.core import typeinfer
107112
with typeinfer.register_dispatcher(disp):
@@ -124,7 +129,7 @@ def autojitwrapper(func):
124129
else:
125130
def autojitwrapper(func):
126131
return jit(func, device=device, debug=debug, opt=opt,
127-
link=link, **kws)
132+
link=link, cache=cache, **kws)
128133

129134
return autojitwrapper
130135
# func_or_sig is a function
@@ -140,7 +145,12 @@ def autojitwrapper(func):
140145
targetoptions['fastmath'] = fastmath
141146
targetoptions['device'] = device
142147
targetoptions['extensions'] = extensions
143-
return CUDADispatcher(func_or_sig, targetoptions=targetoptions)
148+
disp = CUDADispatcher(func_or_sig, targetoptions=targetoptions)
149+
150+
if cache:
151+
disp.enable_caching()
152+
153+
return disp
144154

145155

146156
def declare_device(name, sig):

‎numba/cuda/dispatcher.py‎

Lines changed: 84 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import warnings
77

88
from numba.core import config, serialize, sigutils, types, typing, utils
9+
from numba.core.caching import Cache, CacheImpl
910
from numba.core.compiler_lock import global_compiler_lock
1011
from numba.core.dispatcher import Dispatcher
1112
from 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+
482535
class 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):

‎numba/cuda/simulator/api.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ def elapsed_time(self, event):
7777

7878
def jit(func_or_sig=None, device=False, debug=False, argtypes=None,
7979
inline=False, restype=None, fastmath=False, link=None,
80-
boundscheck=None, opt=True
80+
boundscheck=None, opt=True, cache=None
8181
):
8282
# Here for API compatibility
8383
if boundscheck:

0 commit comments

Comments
 (0)