Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optional numba support #66

Open
jorenham opened this issue Oct 30, 2023 · 0 comments
Open

Optional numba support #66

jorenham opened this issue Oct 30, 2023 · 0 comments
Assignees
Labels
contrib Integration with 3rd party packages performance Runtime performance and benchmarking
Milestone

Comments

@jorenham
Copy link
Owner

jorenham commented Oct 30, 2023

Add contrib.numba, and use a noop @jit if not available on relevant functions.

Additionally, the scipy.integrate.quad integrand can be sped up with cfunc, see https://numba.readthedocs.io/en/stable/user/cfunc.html#example

In places where scipy.special functions are used, some trickery is needed.
For example, this snippet is used to make scipy.special.erfi work within numba jitted functions for np.float64 input:

# lmo/contrib/numba.py
def _overload_scipy_special_erfi():
    _erfi_f8 = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double)(
        numba.extending.get_cython_function_address(
            'scipy.special.cython_special',
            '__pyx_fuse_1erfi',
        ),
    )
    
    @numba.extending.overload(scipy.special.erfi)
    def numba_erfi(*args):
        match args:
            case (numba.types.Float(),):
                def _numba_erfi(*args):
                    return _erfi_f8(*args)
                return _numba_erfi
            case _:
                return None

def overload_scipy_special():
    _overload_scipy_special_erfi()
# lmo/pyproject.toml
[tool.poetry.plugins.numba_extensions]
init = "lmo.contrib.numba:overload_scipy_special"
@jorenham jorenham added enhancement New feature or request contrib Integration with 3rd party packages labels Oct 30, 2023
@jorenham jorenham added this to the v1.0.0 milestone Mar 29, 2024
@jorenham jorenham added performance Runtime performance and benchmarking and removed enhancement New feature or request labels Mar 29, 2024
@jorenham jorenham self-assigned this Mar 29, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
contrib Integration with 3rd party packages performance Runtime performance and benchmarking
Projects
None yet
Development

No branches or pull requests

1 participant