custom_kernels

Full Documentation for hippynn.custom_kernels package. Click here for a summary page.

Custom Kernels for hip-nn interaction sum.

This module provides implementations in pytorch, numba, cupy, and triton.

Pytorch implementations take extra memory, but launch faster than numba kernels. Numba kernels use far less memory, but do come with some launching overhead on GPUs. Cupy kernels only work on the GPU, but are faster than numba. Cupy kernels require numba for CPU operations. Triton custom kernels only work on the GPU, and are generaly faster than CUPY. Triton kernels revert to numba or pytorch as available on module import.

On import, this module attempts to set the custom kernels as specified by the user in hippynn.settings.

See the Custom Kernels section of the documentation for more information.

Depending on your available packages, you may have the following options:

  • “pytorch”: dense pytorch operations.

  • “sparse”: sparse pytorch operations. Can be faster than pure pytorch for large enough systems, and will not require as much memory. May require latest pytorch version. Cannot cover all circumstances; should error if encountering a result that this implementation cannot cover..

  • “numba”: numba implementation of custom kernels, beats pytorch-based kernels.

  • “cupy”: cupy implementation of custom kernels, better than numba

  • “triton”: triton-based custom kernels, uses auto-tuning and the triton compiler. This is usually the best option.

The available items are stored in the variable hippynn.custom_kernels.CUSTOM_KERNELS_AVAILABLE. The active implementation is stored in hippynn.custom_kernels.CUSTOM_KERNELS_ACTIVE.

For more information, see Custom Kernels

exception CustomKernelError[source]

Bases: Exception

envsum(*args, **kwargs)

See hippynn.custom_kernels.env_pytorch.envsum() for more information.

featsum(*args, **kwargs)

See hippynn.custom_kernels.env_pytorch.featsum() for more information.

populate_custom_kernel_availability()[source]

Check available imports and populate the list of available custom kernels.

This function changes the global variable custom_kernels.CUSTOM_KERNELS_AVAILABLE

Returns:

sensesum(*args, **kwargs)

See hippynn.custom_kernels.env_pytorch.sensesum() for more information.

set_custom_kernels(active: bool | str = True) → str[source]

Activate or deactivate custom kernels for interaction.

This function changes the global variables:
Special non-implementation-name values are:

  • True: - Use the best GPU kernel from recommended implementations, error if none are available.

  • False: - equivalent to “pytorch”

  • “auto”: - Equivalently to True if recommended is available, else equivalent to “pytorch”

Parameters:

active – implementation name to activate

Returns:

active, actual implementation selected.

CUSTOM_KERNELS_ACTIVE = 'pytorch'

Which custom kernel implementation is currently active.

CUSTOM_KERNELS_AVAILABLE = ['pytorch', 'sparse']

List of available kernel implementations based on currently installed packages..

Submodules