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.
- envsum(sensitivities: Tensor, features: Tensor, pair_first: Tensor, pair_second: Tensor) Tensor[source]
See
hippynn.custom_kernels.env_pytorch.envsum()for more information.
- featsum(env, sense, pair_first, pair_second)[source]
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(env, features, pair_first, pair_second)[source]
See
hippynn.custom_kernels.env_pytorch.sensesum()for more information.
- set_custom_kernels(active: bool | str = True)[source]
Activate or deactivate custom kernels for interaction.
- This function changes the global variables:
hippynn.custom_kernels.sensum()
- Parameters:
active – If true, set custom kernels to the best available. If False, turn them off and default to pytorch. If “triton”, “numba” or “cupy”, use those implementations explicitly. If “auto”, use best available.
- Returns:
None
- CUSTOM_KERNELS_ACTIVE = False
Which custom kernel implementation is currently active.
- CUSTOM_KERNELS_AVAILABLE = []
List of available kernel implementations based on currently installed packages..
Submodules