def compile_ptx(pyfunc, args, debug=False, lineinfo=False, device=False,
fastmath=False, cc=None, opt=True):
"""Compile a Python function to PTX for a given set of argument types.
:param pyfunc: The Python function to compile.
:param args: A tuple of argument types to compile for.
:param debug: Whether to include debug info in the generated PTX.
:type debug: bool
:param lineinfo: Whether to include a line mapping from the generated PTX
to the source code. Usually this is used with optimized
code (since debug mode would automatically include this),
so we want debug info in the LLVM but only the line
mapping in the final PTX.
:type lineinfo: bool
:param device: Whether to compile a device function. Defaults to ``False``,
to compile global kernel functions.
:type device: bool
:param fastmath: Whether to enable fast math flags (ftz=1, prec_sqrt=0,
prec_div=, and fma=1)
:type fastmath: bool
:param cc: Compute capability to compile for, as a tuple ``(MAJOR, MINOR)``.
Defaults to ``(5, 2)``.
:type cc: tuple
:param opt: Enable optimizations. Defaults to ``True``.
:type opt: bool
:return: (ptx, resty): The PTX code and inferred return type
:rtype: tuple
"""
if debug and opt:
msg = ("debug=True with opt=True (the default) "
"is not supported by CUDA. This may result in a crash"
" - set debug=False or opt=False.")
warn(NumbaInvalidConfigWarning(msg))
nvvm_options = {
'debug': debug,
'lineinfo': lineinfo,
'fastmath': fastmath,
'opt': 3 if opt else 0
}
cres = compile_cuda(pyfunc, None, args, debug=debug, lineinfo=lineinfo,
nvvm_options=nvvm_options)
resty = cres.signature.return_type
if device:
lib = cres.library
else:
fname = cres.fndesc.llvm_func_name
tgt = cres.target_context
filename = cres.type_annotation.filename
linenum = int(cres.type_annotation.linenum)
lib, kernel = tgt.prepare_cuda_kernel(cres.library, fname,
cres.signature.args, debug,
nvvm_options, filename, linenum)
cc = cc or config.CUDA_DEFAULT_PTX_CC
ptx = lib.get_asm_str(cc=cc)
return ptx, resty
def _get_ptxes(self, cc=None):
if not cc:
ctx = devices.get_context()
device = ctx.device
cc = device.compute_capability
ptxes = self._ptx_cache.get(cc, None)
if ptxes:
return ptxes
arch = nvvm.get_arch_option(*cc)
options = self._nvvm_options.copy()
options['arch'] = arch
if not nvvm.NVVM().is_nvvm70:
# Avoid enabling debug for NVVM 3.4 as it has various issues. We
# need to warn the user that we're doing this if any of the
# functions that they're compiling have `debug=True` set, which we
# can determine by checking the NVVM options.
for lib in self.linking_libraries:
if lib._nvvm_options.get('debug'):
msg = ("debuginfo is not generated for CUDA versions "
f"< 11.2 (debug=True on function: {lib.name})")
warn(NumbaInvalidConfigWarning(msg))
options['debug'] = False
irs = [str(mod) for mod in self.modules]
if options.get('debug', False):
# If we're compiling with debug, we need to compile modules with
# NVVM one at a time, because it does not support multiple modules
# with debug enabled:
# https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#source-level-debugging-support
ptxes = [nvvm.llvm_to_ptx(ir, **options) for ir in irs]
else:
# Otherwise, we compile all modules with NVVM at once because this
# results in better optimization than separate compilation.
ptxes = [nvvm.llvm_to_ptx(irs, **options)]
# Sometimes the result from NVVM contains trailing whitespace and
# nulls, which we strip so that the assembly dump looks a little
# tidier.
ptxes = [x.decode().strip('\x00').strip() for x in ptxes]
if config.DUMP_ASSEMBLY:
print(("ASSEMBLY %s" % self._name).center(80, '-'))
print(self._join_ptxes(ptxes))
print('=' * 80)
self._ptx_cache[cc] = ptxes
return ptxes
def jit(func_or_sig=None,
device=False,
inline=False,
link=[],
debug=None,
opt=True,
**kws):
"""
JIT compile a python function conforming to the CUDA Python specification.
If a signature is supplied, then a function is returned that takes a
function to compile.
:param func_or_sig: A function to JIT compile, or a signature of a function
to compile. If a function is supplied, then a
:class:`numba.cuda.compiler.AutoJitCUDAKernel` is returned. If a
signature is supplied, then a function is returned. The returned
function accepts another function, which it will compile and then return
a :class:`numba.cuda.compiler.AutoJitCUDAKernel`.
.. note:: A kernel cannot have any return value.
:param device: Indicates whether this is a device function.
:type device: bool
:param link: A list of files containing PTX source to link with the function
:type link: list
:param debug: If True, check for exceptions thrown when executing the
kernel. Since this degrades performance, this should only be used for
debugging purposes. If set to True, then ``opt`` should be set to False.
Defaults to False. (The default value can be overridden by setting
environment variable ``NUMBA_CUDA_DEBUGINFO=1``.)
:param fastmath: When True, enables fastmath optimizations as outlined in
the :ref:`CUDA Fast Math documentation <cuda-fast-math>`.
:param max_registers: Request that the kernel is limited to using at most
this number of registers per thread. The limit may not be respected if
the ABI requires a greater number of registers than that requested.
Useful for increasing occupancy.
:param opt: Whether to compile from LLVM IR to PTX with optimization
enabled. When ``True``, ``-opt=3`` is passed to NVVM. When
``False``, ``-opt=0`` is passed to NVVM. Defaults to ``True``.
:type opt: bool
:param lineinfo: If True, generate a line mapping between source code and
assembly code. This enables inspection of the source code in NVIDIA
profiling tools and correlation with program counter sampling.
:type lineinfo: bool
"""
if link and config.ENABLE_CUDASIM:
raise NotImplementedError('Cannot link PTX in the simulator')
if kws.get('boundscheck'):
raise NotImplementedError("bounds checking is not supported for CUDA")
if kws.get('argtypes') is not None:
msg = _msg_deprecated_signature_arg.format('argtypes')
raise DeprecationError(msg)
if kws.get('restype') is not None:
msg = _msg_deprecated_signature_arg.format('restype')
raise DeprecationError(msg)
if kws.get('bind') is not None:
msg = _msg_deprecated_signature_arg.format('bind')
raise DeprecationError(msg)
debug = config.CUDA_DEBUGINFO_DEFAULT if debug is None else debug
fastmath = kws.get('fastmath', False)
extensions = kws.get('extensions', [])
if debug and opt:
msg = ("debug=True with opt=True (the default) "
"is not supported by CUDA. This may result in a crash"
" - set debug=False or opt=False.")
warn(NumbaInvalidConfigWarning(msg))
if device and kws.get('link'):
raise ValueError("link keyword invalid for device function")
if sigutils.is_signature(func_or_sig):
if config.ENABLE_CUDASIM:
def jitwrapper(func):
return FakeCUDAKernel(func, device=device, fastmath=fastmath)
return jitwrapper
argtypes, restype = sigutils.normalize_signature(func_or_sig)
if restype and not device and restype != types.void:
raise TypeError("CUDA kernel must have void return type.")
def _jit(func):
targetoptions = kws.copy()
targetoptions['debug'] = debug
targetoptions['link'] = link
targetoptions['opt'] = opt
targetoptions['fastmath'] = fastmath
targetoptions['device'] = device
targetoptions['extensions'] = extensions
disp = CUDADispatcher(func, targetoptions=targetoptions)
if device:
disp.compile_device(argtypes)
else:
disp.compile(argtypes)
disp._specialized = True
disp.disable_compile()
return disp
return _jit
else:
if func_or_sig is None:
if config.ENABLE_CUDASIM:
def autojitwrapper(func):
return FakeCUDAKernel(func,
device=device,
fastmath=fastmath)
else:
def autojitwrapper(func):
return jit(func,
device=device,
debug=debug,
opt=opt,
link=link,
**kws)
return autojitwrapper
# func_or_sig is a function
else:
if config.ENABLE_CUDASIM:
return FakeCUDAKernel(func_or_sig,
device=device,
fastmath=fastmath)
else:
targetoptions = kws.copy()
targetoptions['debug'] = debug
targetoptions['opt'] = opt
targetoptions['link'] = link
targetoptions['fastmath'] = fastmath
targetoptions['device'] = device
targetoptions['extensions'] = extensions
return CUDADispatcher(func_or_sig, targetoptions=targetoptions)