diff --git a/pyproject.toml b/pyproject.toml index 845c4d6..318186a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,11 +91,8 @@ format = {composite = ["black", "isort"]} -[tool.pytest] -addopts = "--doctest-modules" - [tool.pytest.ini_options] -addopts = "--pdbcls=IPython.terminal.debugger:TerminalPdb" +addopts = "--doctest-modules --pdbcls=IPython.terminal.debugger:TerminalPdb" [tool.black] diff --git a/src/transonic/aheadoftime.py b/src/transonic/aheadoftime.py index 3fb2642..b146f6e 100644 --- a/src/transonic/aheadoftime.py +++ b/src/transonic/aheadoftime.py @@ -28,6 +28,7 @@ import sys import time from importlib import import_module +from typing import Optional from transonic import mpi from transonic.backends import backends, get_backend_name_module @@ -107,13 +108,14 @@ def _get_transonic_calling_module(backend_name: str = None): def boost( obj=None, - backend: str = None, + backend: str | None = None, inline=False, boundscheck=True, wraparound=True, cdivision=False, nonecheck=True, nogil=False, + fastmath=False ): """Decorator to declare that an object can be accelerated @@ -135,6 +137,7 @@ def boost( wraparound=wraparound, cdivision=cdivision, nonecheck=nonecheck, + fastmath=fastmath ) if callable(obj) or isinstance(obj, type): return decor(obj) diff --git a/src/transonic/backends/numba.py b/src/transonic/backends/numba.py index 6192648..e952a36 100644 --- a/src/transonic/backends/numba.py +++ b/src/transonic/backends/numba.py @@ -22,11 +22,10 @@ from .py import PythonBackend, SubBackendJITPython -def add_numba_comments(code): +def add_numba_comments(code, fastmath_metadata): """Add Numba code in Python comments""" mod = parse(code) new_body = [CommentLine("# __protected__ from numba import njit")] - for node in mod.body: if ( isinstance(node, gast.FunctionDef) @@ -37,8 +36,13 @@ def add_numba_comments(code): ) and not node.name.startswith("__code_new_method__") ): + + is_fastmath = True + if (node.name in fastmath_metadata): + is_fastmath = fastmath_metadata[node.name] + new_body.append( - CommentLine("# __protected__ @njit(cache=True, fastmath=True)") + CommentLine(f"# __protected__ @njit(cache=True, fastmath={is_fastmath})") ) new_body.append(node) @@ -48,14 +52,20 @@ def add_numba_comments(code): class SubBackendJITNumba(SubBackendJITPython): def make_backend_source(self, info_analysis, func, path_backend): - src, has_to_write = super().make_backend_source( - info_analysis, func, path_backend - ) + src, has_to_write = super().make_backend_source(info_analysis, func, path_backend) if not src: return src, has_to_write - return add_numba_comments(src), has_to_write + fastmath_metadata = {} + numba_funcs = info_analysis[0].get('functions', {}).get('numba', {}) + for name, node in numba_funcs.items(): + is_fastmath = False + if hasattr(node, '_transonic_keywords') and 'fastmath' in node._transonic_keywords: + is_fastmath = node._transonic_keywords['fastmath'] + fastmath_metadata[name] = is_fastmath + + return add_numba_comments(src, fastmath_metadata), has_to_write class NumbaBackend(PythonBackend): @@ -92,12 +102,24 @@ def compile_extension( def _make_backend_code(self, path_py, analysis, **kwargs): """Create a backend code from a Python file""" + fastmath_metadata = {} + numba_funcs = analysis[0].get('functions', {}).get('numba', {}) + code, codes_ext, header = super()._make_backend_code(path_py, analysis) + + for name, node in numba_funcs.items(): + is_fastmath = False + + if hasattr(node, '_transonic_keywords'): + if 'fastmath' in node._transonic_keywords: + is_fastmath = node._transonic_keywords['fastmath'] + + fastmath_metadata[name] = is_fastmath if not code: return code, codes_ext, header - code = add_numba_comments(code) + code = add_numba_comments(code , fastmath_metadata) for_meson = kwargs.get("for_meson", False) if for_meson: diff --git a/src/transonic/justintime.py b/src/transonic/justintime.py index 079ae43..5bcde21 100644 --- a/src/transonic/justintime.py +++ b/src/transonic/justintime.py @@ -53,6 +53,7 @@ import time from functools import wraps from pathlib import Path +from typing import Optional from transonic import mpi from transonic.aheadoftime import TransonicTemporaryJITMethod @@ -184,7 +185,7 @@ def _get_module_jit(backend_name: str = None, depth_frame: int = 2, frame=None): return ModuleJIT(backend_name=backend_name, frame=frame) -def jit(func=None, backend: str = None, native=False, xsimd=False, openmp=False): +def jit(func=None, backend: str | None = None, native=False, xsimd=False, openmp=False, fastmath=False): """Decorator to record that the function has to be jit compiled""" frame = get_frame(1) decor = JIT(frame, backend=backend, native=native, xsimd=xsimd, openmp=openmp)