spikingjelly.activation_based.triton_kernel package

备注

Developers should decide whether Triton backend is available by:

try:
    import triton
except ImportError:
    triton = None

if triton is not None:
    # Triton backend is available
    ...
else:
    # Triton backend is not available
    ...

MemoryModule encapsulates this logic.

Predefined Neuron Kernels

• neuron_kernel
  • torch.compile Compatibility
  • IF
    • multistep_if()
  • LIF
    • multistep_lif()
  • PLIF
    • multistep_plif()

FlexSN Implementation

• flexsn
  • Function Info
    • FlexSNInfo
    • extract_info()
  • Template
    • get_flexsn_inference_kernel()
    • get_flexsn_inference_final_state_kernel()
    • get_flexsn_forward_kernel()
    • get_flexsn_backward_kernel()
  • Kernel Builders
    • build_inference_kernel()
    • build_inference_final_state_kernel()
    • build_training_kernels()
  • Wrapper
    • flexsn_backward()
    • flexsn_backward_ncl_bucket()
    • flexsn_forward()
    • flexsn_inference()
    • flexsn_inference_final_state()
  • Custom Ops
    • FlexSNKernelHandle
    • register_flexsn_kernel_handle()
    • retain_flexsn_kernel_handle()
    • retain_owner_flexsn_kernel_handle()
    • release_flexsn_kernel_handle()
    • release_active_flexsn_kernel_handle()
    • attach_flexsn_handle_finalizer()
  • HigherOrderOp & Eager Scan
    • dynamo_hop_available()
    • eager_scan()
    • eager_scan_final_state()
    • FlexSNScan
    • lowerable_scan()
    • lowerable_scan_available()
    • lowerable_scan_final_state()
    • lowerable_while_loop_scan()
    • lowerable_while_loop_available()
    • lowerable_while_loop_scan_final_state()

Torch-to-Triton Transpiler

• torch2triton
  • PyTorch to Graph
    • generate_inference_graph()
    • generate_forward_and_backward_graph()
  • Graph to Triton
    • generate_triton_code_str()
    • compile_triton_code_str()

Spike Compressors

• compress
  • bit_spike_compress()
  • bit_spike_decompress()

Utilities

Borrowed from: AllenYolk/flash-snn fla-org/flash-linear-attention

spikingjelly.activation_based.triton_kernel.triton_utils.normalize_triton_compute_dtype_name(compute_dtype)

参数:

compute_dtype (str | dtype)

返回类型:

str

spikingjelly.activation_based.triton_kernel.triton_utils.normalize_triton_storage_dtype(storage_dtype)

参数:

storage_dtype (str | dtype)

返回类型:

dtype

spikingjelly.activation_based.triton_kernel.triton_utils.is_fp8_dtype(dtype)

参数:

dtype (dtype)

返回类型:

bool

spikingjelly.activation_based.triton_kernel.triton_utils.torch_dtype_to_triton_neuron_dtype_id(dtype)

参数:

dtype (dtype)

返回类型:

int

spikingjelly.activation_based.triton_kernel.triton_utils.triton_neuron_dtype_id_to_torch_dtype(dtype_id)

参数:

dtype_id (int)

返回类型:

dtype

spikingjelly.activation_based.triton_kernel.triton_utils.triton_compute_dtype_name_to_neuron_dtype_id(compute_dtype_name, storage_dtype)

参数:

• compute_dtype_name (str)

• storage_dtype (dtype)

返回类型:

int

spikingjelly.activation_based.triton_kernel.triton_utils.triton_neuron_compute_dtype_id_to_tl_dtype(dtype_id, storage_dtype_id)

参数:

• dtype_id (int)

• storage_dtype_id (int)

spikingjelly.activation_based.triton_kernel.triton_utils.resolve_triton_compute_dtype(compute_dtype, storage_dtype=None)

参数:

• compute_dtype (str | dtype)

• storage_dtype (str | dtype | None)

spikingjelly.activation_based.triton_kernel.triton_utils.torch_dtype_for_triton_compute_dtype(compute_dtype)

参数:

compute_dtype (str | dtype)

返回类型:

dtype

spikingjelly.activation_based.triton_kernel.triton_utils.torch_dtype_for_triton_neuron_compute_dtype_id(dtype_id)

参数:

dtype_id (int)

返回类型:

dtype

spikingjelly.activation_based.triton_kernel.triton_utils.convert_and_store(pointer, value, boundary_check)

spikingjelly.activation_based.triton_kernel.triton_utils.register_op(opname, mutates_args=())

参数:

opname (str)

spikingjelly.activation_based.triton_kernel.triton_utils.wrap_triton(kernel)

spikingjelly.activation_based.triton_kernel.triton_utils.contiguous_and_device_guard(f)

Make sure all input tensors are contiguous and set to the same device.

参数:

f (Callable)

返回类型:

Callable

spikingjelly.activation_based.triton_kernel.triton_utils.use_static_range_for_triton_neuron_kernel(T)

参数:

T (int)

返回类型:

bool

spikingjelly.activation_based.triton_kernel.triton_utils.ensure_cleanup_tmp_python_files(f)

Remove temporary python files returned or created by a wrapped function.

参数:

f (Callable)

返回类型:

Callable

class spikingjelly.activation_based.triton_kernel.dummy.DummyImport

基类：object

Dummy class as an import placeholder.