def __call__(self,
grid,
block,
args,
shared_mem=0,
stream=None,
enable_cooperative_groups=False):
in_types = []
for x in args:
if isinstance(x, cupy.ndarray):
t = _types.CArray.from_ndarray(x)
elif numpy.isscalar(x):
t = _typerules.get_ctype_from_scalar(self._mode, x)
else:
raise TypeError(f'{type(x)} is not supported for RawKernel')
in_types.append(t)
in_types = tuple(in_types)
kern = self._cache.get(in_types)
if kern is None:
result = _compile.transpile(
self._func,
['extern "C"', '__global__'],
self._mode,
in_types,
_types.Void(),
)
fname = result.func_name
module = cupy._core.core.compile_with_cache(
source=result.code, options=('-D CUPY_JIT_MODE', ))
kern = module.get_function(fname)
self._cache[in_types] = kern
self._cached_codes[in_types] = result.code
kern(grid, block, args, shared_mem, stream, enable_cooperative_groups)
def __call__(self, grid, block, args):
in_types = []
for x in args:
if isinstance(x, cupy.ndarray):
t = _types.Array.from_ndarray(x)
elif numpy.isscalar(x):
t = _typerules.get_ctype_from_scalar(self._mode, x)
else:
raise TypeError(f'{type(x)} is not supported for RawKernel')
in_types.append(t)
in_types = tuple(in_types)
kern = self._cache.get(in_types)
if kern is None:
result = _compile.transpile(
self._func,
['extern "C"', '__global__'],
self._mode,
in_types,
_types.Void(),
)
fname = result.func_name
module = cupy.core.core.compile_with_cache(result.code)
kern = module.get_function(fname)
self._cache[in_types] = kern
kern(grid, block, args)
def __call__(self,
grid,
block,
args,
shared_mem=0,
stream=None,
enable_cooperative_groups=False):
"""Calls the CUDA kernel.
The compilation will be deferred until the first function call.
CuPy's JIT compiler infers the types of arguments at the call
time, and will cache the compiled kernels for speeding up any
subsequent calls.
Args:
grid (tuple of int): Size of grid in blocks.
block (tuple of int): Dimensions of each thread block.
args (tuple):
Arguments of the kernel. The type of all elements must be
``bool``, ``int``, ``float``, ``complex``, NumPy scalar or
``cupy.ndarray``.
shared_mem (int):
Dynamic shared-memory size per thread block in bytes.
stream (cupy.cuda.Stream): CUDA stream.
.. seealso:: :ref:`jit_kernel_definition`
"""
in_types = []
for x in args:
if isinstance(x, cupy.ndarray):
t = _types.CArray.from_ndarray(x)
elif numpy.isscalar(x):
t = _typerules.get_ctype_from_scalar(self._mode, x)
else:
raise TypeError(f'{type(x)} is not supported for RawKernel')
in_types.append(t)
in_types = tuple(in_types)
kern = self._cache.get(in_types)
if kern is None:
result = _compile.transpile(
self._func,
['extern "C"', '__global__'],
self._mode,
in_types,
_types.Void(),
)
fname = result.func_name
# workaround for hipRTC: as of ROCm 4.1.0 hipRTC still does not
# recognize "-D", so we have to compile using hipcc...
backend = 'nvcc' if runtime.is_hip else 'nvrtc'
module = core.compile_with_cache(source=result.code,
options=('-DCUPY_JIT_MODE',
'--std=c++11'),
backend=backend)
kern = module.get_function(fname)
self._cache[in_types] = kern
self._cached_codes[in_types] = result.code
kern(grid, block, args, shared_mem, stream, enable_cooperative_groups)
def _transpile_function(func, attributes, mode, consts, in_types, ret_type):
"""Transpile the function
Args:
func (ast.FunctionDef): Target function.
attributes (str): The attributes of target function.
mode ('numpy' or 'cuda'): The rule for typecast.
consts (dict): The dictionary with keys as variable names and
values as concrete data object.
in_types (list of _types.TypeBase): The types of arguments.
ret_type (_types.TypeBase): The type of return value.
Returns:
code (str): The generated CUDA code.
env (Environment): More details of analysis result of the function,
which includes preambles, estimated return type and more.
"""
consts = dict([(k, Constant(v)) for k, v, in consts.items()])
if not isinstance(func, ast.FunctionDef):
# TODO(asi1024): Support for `ast.ClassDef`.
raise NotImplementedError('Not supported: {}'.format(type(func)))
if len(func.decorator_list) > 0:
if sys.version_info >= (3, 9):
# Code path for Python versions that support `ast.unparse`.
for deco in func.decorator_list:
deco_code = ast.unparse(deco)
if deco_code not in ['rawkernel', 'vectorize']:
warnings.warn(
f'Decorator {deco_code} may not supported in JIT.',
RuntimeWarning)
arguments = func.args
if arguments.vararg is not None:
raise NotImplementedError('`*args` is not supported currently.')
if len(arguments.kwonlyargs) > 0: # same length with `kw_defaults`.
raise NotImplementedError(
'keyword only arguments are not supported currently .')
if arguments.kwarg is not None:
raise NotImplementedError('`**kwargs` is not supported currently.')
if len(arguments.defaults) > 0:
raise NotImplementedError(
'Default values are not supported currently.')
args = [arg.arg for arg in arguments.args]
if len(args) != len(in_types):
raise TypeError(
f'{func.name}() takes {len(args)} positional arguments '
f'but {len(in_types)} were given.')
params = dict([(x, CudaObject(x, t)) for x, t in zip(args, in_types)])
env = Environment(mode, consts, params, ret_type)
body = _transpile_stmts(func.body, True, env)
params = ', '.join([env[a].ctype.declvar(a) for a in args])
local_vars = [v.ctype.declvar(n) + ';' for n, v in env.locals.items()]
if env.ret_type is None:
env.ret_type = _types.Void()
head = f'{attributes} {env.ret_type} {func.name}({params})'
code = CodeBlock(head, local_vars + body)
return str(code), env
def _transpile_function_internal(func, attributes, mode, consts, in_types,
ret_type):
consts = dict([(k, Constant(v)) for k, v, in consts.items()])
if not isinstance(func, ast.FunctionDef):
# TODO(asi1024): Support for `ast.ClassDef`.
raise NotImplementedError('Not supported: {}'.format(type(func)))
if len(func.decorator_list) > 0:
if sys.version_info >= (3, 9):
# Code path for Python versions that support `ast.unparse`.
for deco in func.decorator_list:
deco_code = ast.unparse(deco)
if not any(word in deco_code
for word in ['rawkernel', 'vectorize']):
warnings.warn(
f'Decorator {deco_code} may not supported in JIT.',
RuntimeWarning)
arguments = func.args
if arguments.vararg is not None:
raise NotImplementedError('`*args` is not supported currently.')
if len(arguments.kwonlyargs) > 0: # same length with `kw_defaults`.
raise NotImplementedError(
'keyword only arguments are not supported currently .')
if arguments.kwarg is not None:
raise NotImplementedError('`**kwargs` is not supported currently.')
if len(arguments.defaults) > 0:
raise NotImplementedError(
'Default values are not supported currently.')
args = [arg.arg for arg in arguments.args]
if len(args) != len(in_types):
raise TypeError(
f'{func.name}() takes {len(args)} positional arguments '
f'but {len(in_types)} were given.')
params = dict([(x, CudaObject(x, t)) for x, t in zip(args, in_types)])
env = Environment(mode, consts, params, ret_type)
body = _transpile_stmts(func.body, True, env)
params = ', '.join([env[a].ctype.declvar(a) for a in args])
local_vars = [v.ctype.declvar(n) + ';' for n, v in env.locals.items()]
if env.ret_type is None:
env.ret_type = _types.Void()
head = f'{attributes} {env.ret_type} {func.name}({params})'
code = CodeBlock(head, local_vars + body)
return str(code), env
