cupyx.jit._interface._JitRawKernel#

class cupyx.jit._interface._JitRawKernel(func, mode, device)[source]#

JIT CUDA 内核对象。

装饰器 :func:cupyx.jit.rawkernel 将目标函数转换为此类的一个对象。用户不应直接实例化此类。

方法

__call__(grid, block, args, shared_mem=0, stream=None)[source]#

调用 CUDA 内核。

编译将延迟到第一次函数调用时进行。CuPy 的 JIT 编译器在调用时推断参数类型，并将编译好的内核缓存起来以加快后续调用。

参数:

  • grid (整数元组) – 网格的块大小。

  • block (整数元组) – 每个线程块的维度。

  • args (元组) – 内核的参数。所有元素的类型必须是 bool, int, float, complex, NumPy 标量或 cupy.ndarray。

  • shared_mem (int) – 每个线程块的动态共享内存大小（字节）。

  • stream (cupy.cuda.Stream) – CUDA 流。

另见

JIT 内核定义

__getitem__(grid_and_block)[source]#

Numba 风格的内核调用。

另见

JIT 内核定义

__eq__(value, /)#

返回 self==value。

__ne__(value, /)#

返回 self!=value。

__lt__(value, /)#

返回 self<value。

__le__(value, /)#

返回 self<=value。

__gt__(value, /)#

返回 self>value。

__ge__(value, /)#

返回 self>=value。

属性

cached_code#

返回 next(iter(self.cached_codes.values()))。

此属性方法用于调试。返回值不保证保持向后兼容性。

cached_codes#

返回一个字典，其键为输入类型，值为代码。

此属性方法用于调试。返回值不保证保持向后兼容性。