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

• spikingjelly.activation_based.triton_kernel.neuron_kernel package
  • torch.compile Compatibility
  • IF
    • multistep_if()
  • LIF
    • multistep_lif()
  • PLIF
    • multistep_plif()

FlexSN Implementation

• spikingjelly.activation_based.triton_kernel.flexsn package
  • 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

• spikingjelly.activation_based.triton_kernel.torch2triton package
  • PyTorch to Graph
    • generate_inference_graph()
    • generate_forward_and_backward_graph()
  • Graph to Triton
    • generate_triton_code_str()
    • compile_triton_code_str()

Spike Compressors

• spikingjelly.activation_based.triton_kernel.compress module
  • bit_spike_compress()
  • bit_spike_decompress()

Utilities

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

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

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

spikingjelly.activation_based.triton_kernel.triton_utils.contiguous_and_device_guard(f: Callable) → Callable

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

spikingjelly.activation_based.triton_kernel.triton_utils.ensure_cleanup_tmp_python_files(f: Callable) → Callable

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

class spikingjelly.activation_based.triton_kernel.dummy.DummyImport

基类：object

Dummy class as an import placeholder.