spikingjelly.activation_based.cuda_kernel.cuda_utils package#

spikingjelly.activation_based.cuda_kernel.cuda_utils.env_flag_enabled(var_name: str) → bool[源代码]#
spikingjelly.activation_based.cuda_kernel.cuda_utils.use_cupy_custom_op() → bool[源代码]#
spikingjelly.activation_based.cuda_kernel.cuda_utils.register_python_object(obj: Any, key: str) → int[源代码]#
spikingjelly.activation_based.cuda_kernel.cuda_utils.resolve_python_object(obj_id: int) → Any[源代码]#
spikingjelly.activation_based.cuda_kernel.cuda_utils.python_object_registry_key(obj: Any) → str[源代码]#
spikingjelly.activation_based.cuda_kernel.cuda_utils.cpu_timer(f: Callable, *args, **kwargs)[源代码]#

API Language: 中文 | English

  • 中文

  • 中文

计算在CPU上执行 f(*args, **kwargs) 所需的时间

参数:

f (Callable) -- 函数

返回:

用时，单位是毫秒

返回类型:

float

  • English

  • English

Returns the used time for calling f(*args, **kwargs) in CPU

参数:

f (Callable) -- a function

返回:

used time in milliseconds

返回类型:

float

spikingjelly.activation_based.cuda_kernel.cuda_utils.cuda_timer(device: device | int, f: Callable, *args, **kwargs)[源代码]#

API Language: 中文 | English

  • 中文

  • 中文

计算在CUDA上执行 f(*args, **kwargs) 所需的时间

参数:

  • device (Union[torch.device, int]) -- f 运行的CUDA设备

  • f (Callable) -- 函数

返回:

用时，单位是毫秒

返回类型:

float

  • English

  • English

Returns the used time for calling f(*args, **kwargs) in CUDA

参数:

  • device (Union[torch.device, int]) -- on which cuda device that f is running

  • f (Callable) -- a function

返回:

used time in milliseconds

返回类型:

float

spikingjelly.activation_based.cuda_kernel.cuda_utils.cal_fun_t(n: int, device: str | device | int, f: Callable, *args, **kwargs)[源代码]#

API Language: 中文 | English

  • 中文

测量在 device 上执行 n 次 f(*args, **kwargs) 的平均用时

备注

当 n > 1 时，实际上会执行 2n 次，然后返回后 n 次的平均用时，以减小误差。

参数:

  • n (int) -- 重复的次数

  • device (Union[str, torch.device, int]) -- f 执行的设备，可以为 'cpu' 或CUDA设备

  • f (Callable) -- 函数

返回:

用时，单位是毫秒

返回类型:

float

  • English

Returns the used time averaged by calling f(*args, **kwargs) over n times

Note

If n > 1, this function will call f for 2n times and return the average used time by the last n times to reduce the measure error.

参数:

  • n (int) -- repeat times

  • device (Union[str, torch.device, int]) -- on which cuda device that f is running. It can be 'cpu' or a cuda deivce

  • f (Callable) -- function

返回:

used time in milliseconds

返回类型:

float

spikingjelly.activation_based.cuda_kernel.cuda_utils.cal_blocks(numel: int, threads: int = -1)[源代码]#

API Language: 中文 | English

  • 中文

  • 中文

参数:

  • numel (int) -- 并行执行的CUDA内核的数量

  • threads (int) -- 每个cuda block中threads的数量，默认为-1，表示使用 configure.cuda_threads

返回:

blocks的数量

返回类型:

int

此函数返回 blocks的数量，用来按照 kernel((blocks,), (configure.cuda_threads,), ...) 调用 cupy.RawKernel

  • English

  • English

参数:

  • numel (int) -- the number of parallel CUDA kernels

  • threads (int) -- the number of threads in each cuda block. The defaule value is -1, indicating to use configure.cuda_threads

返回:

the number of blocks

返回类型:

int

Returns the number of blocks to call cupy.RawKernel by kernel((blocks,), (threads,), ...)

spikingjelly.activation_based.cuda_kernel.cuda_utils.get_contiguous(*args)[源代码]#

API Language: 中文 | English

  • 中文

  • 中文

将 *args 中所有的 torch.Tensor 或 cupy.ndarray 进行连续化。

备注

连续化的操作无法in-place，因此本函数返回一个新的list。

返回:

一个元素全部为连续的 torch.Tensor 或 cupy.ndarray 的 list

返回类型:

list

  • English

  • English

返回:

a list that contains the contiguous torch.Tensor or cupy.ndarray

返回类型:

list

Makes torch.Tensor or cupy.ndarray in *args to be contiguous

Note

The making contiguous operation can not be done in-place. Hence, this function will return a new list.

spikingjelly.activation_based.cuda_kernel.cuda_utils.wrap_args_to_raw_kernel(device: int, *args)[源代码]#

API Language: 中文 | English

  • 中文

  • 中文

参数:

device (int) -- raw kernel运行的CUDA设备

返回:

一个包含用来调用 cupy.RawKernel 的 tuple

返回类型:

tuple

此函数可以包装 torch.Tensor 和 cupy.ndarray 并将其作为 cupy.RawKernel.__call__ 的 args

  • English

  • English

参数:

device (int) -- on which CUDA device the raw kernel will run

返回:

a tuple that contains args to call cupy.RawKernel

返回类型:

tuple

This function can wrap torch.Tensor or cupy.ndarray to args in cupy.RawKernel.__call__

class spikingjelly.activation_based.cuda_kernel.cuda_utils.DeviceEnvironment(device: int)[源代码]#

基类：object

API Language: 中文 | English

  • 中文

  • 中文

这个模块可以被用作在指定的 device 上执行CuPy函数的上下文，用来避免 torch.cuda.current_device() 被CuPy意外改变( cupy/cupy#6569 )。

代码示例：

with DeviceEnvironment(device):
    kernel((blocks,), (configure.cuda_threads,), ...)

  • English

  • English

参数:

device (int) -- the CUDA device

This module is used as a context to make CuPy use the specific device, and avoids torch.cuda.current_device() is changed by CuPy ( cupy/cupy#6569 ).

Codes example:

with DeviceEnvironment(device):
    kernel((blocks,), (configure.cuda_threads,), ...)