def _to_cuda_object(x, env): if isinstance(x, CudaObject): return x if isinstance(x, Constant): ctype = _typerules.get_ctype_from_scalar(env.mode, x.obj) return CudaObject(str(x.obj).lower(), ctype) assert False
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 _call_ufunc(ufunc, args, dtype, env): if len(args) != ufunc.nin: raise ValueError('invalid number of arguments') in_types = [] for x in args: if is_constants([x]): t = _typerules.get_ctype_from_scalar(env.mode, x.obj).dtype else: t = x.ctype.dtype in_types.append(t) op = _typerules.guess_routine(ufunc, in_types, dtype, env.mode) if op is None: raise TypeError( f'"{ufunc.name}" does not support for the input types: {in_types}') if op.error_func is not None: op.error_func() if ufunc.nout == 1 and op.routine.startswith('out0 = '): out_type = _types.Scalar(op.out_types[0]) expr = op.routine.replace('out0 = ', '') in_params = [] for x, t in zip(args, op.in_types): x = _astype_scalar(x, _types.Scalar(t), 'same_kind', env) x = _to_cuda_object(x, env) in_params.append(x) can_use_inline_expansion = True for i in range(ufunc.nin): if len(list(re.finditer(r'in{}'.format(i), op.routine))) > 1: can_use_inline_expansion = False if can_use_inline_expansion: # Code pass for readable generated code for i, x in enumerate(in_params): expr = expr.replace(f'in{i}', x.code) expr = '(' + expr.replace('out0_type', str(out_type)) + ')' env.preambles.add(ufunc._preamble) else: template_typenames = ', '.join( [f'typename T{i}' for i in range(ufunc.nin)]) ufunc_name = f'{ufunc.name}_{str(numpy.dtype(op.out_types[0]))}' params = ', '.join([f'T{i} in{i}' for i in range(ufunc.nin)]) ufunc_code = f"""template <{template_typenames}> __device__ {out_type} {ufunc_name}({params}) {{ return {expr}; }} """ env.preambles.add(ufunc_code) in_params = ', '.join([a.code for a in in_params]) expr = f'{ufunc_name}({in_params})' return CudaObject(expr, out_type) raise NotImplementedError(f'ufunc `{ufunc.name}` is not supported.')
def _to_cuda_object(x, env): if isinstance(x, CudaObject): return x if isinstance(x, Constant): ctype = _typerules.get_ctype_from_scalar(env.mode, x.obj) code = _typerules.get_cuda_code_from_constant(x.obj, ctype) return CudaObject(code, ctype) if isinstance(x, Range): raise TypeError('range object cannot be interpreted as a cuda object.') assert False
def _emit_cuda_object_from_constants(x, env): ctype = _typerules.get_ctype_from_scalar(env.mode, x) return CudaObject(str(x), ctype)