cupy.ElementwiseKernel#

class cupy.ElementwiseKernel(in_params, out_params, operation, name='kernel', reduce_dims=True, preamble='', no_return=False, return_tuple=False, **kwargs)[source]#

用户定义逐元素核函数。

此类可用于定义支持或不支持广播的逐元素核函数。

核函数在调用 __call__() 方法时进行编译，并为每个设备缓存。编译后的二进制文件也会缓存到 $HOME/.cupy/kernel_cache/ 目录下的一个带哈希文件名的文件中。缓存的二进制文件可由其他进程重用。

参数:

  • in_params (str) – 输入参数列表。

  • out_params (str) – 输出参数列表。

  • operation (str) – 使用 CUDA-C/C++ 编写的循环体。

  • name (str) – 核函数的名称。应设置此参数以提高性能分析的可读性。

  • reduce_dims (bool) – 如果为 False，则数组参数的形状在核函数调用中保持不变。默认情况下，形状会被缩减（即数组会被重塑到最小维度而无需复制）。这可以通过减少索引计算来加快核函数的执行速度。

  • options (tuple) – 传递给 NVRTC 的编译选项。详情请参阅 https://docs.nvda.net.cn/cuda/nvrtc/index.html#group__options。

  • preamble (str) – 插入到 .cu 文件顶部的 CUDA-C/C++ 代码片段。

  • no_return (bool) – 如果为 True，则 __call__ 返回 None。

  • return_tuple (bool) – 如果为 True，则即使只返回单个值，__call__ 也总是返回数组元组。

  • loop_prep (str) – 插入到核函数定义顶部、for 循环之上的 CUDA-C/C++ 代码片段。

  • after_loop (str) – 插入到核函数定义底部的 CUDA-C/C++ 代码片段。

方法

__call__()#

编译并调用逐元素核函数。

只有当核函数未缓存时，才会进行编译。请注意，参数数据类型或维度不同的核函数不兼容。这意味着单个 ElementwiseKernel 对象可能会被编译成多个核函数二进制文件。

参数:

  • args – 核函数的参数。

  • size (int) – 索引的范围大小。默认情况下，范围大小会根据广播结果自动确定。仅当所有 ndarrays 都是 raw 且无法自动确定范围大小时，才必须指定此参数。

  • block_size (int) – 每个块的线程数。默认情况下，该值设置为 128。

返回值:

如果未设置 no_return，则根据 __init__ 方法的 out_params 参数返回数组。如果设置了 no_return，则返回 None。

__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#

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

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

in_params#
kwargs#
name#
nargs#
nin#
no_return#
nout#
operation#
out_params#
params#
preamble#
reduce_dims#
return_tuple#