Read about our upcoming Code of Conduct on this issue

Commit 4e7efda4 authored by Pierre Blanc-fatin's avatar Pierre Blanc-fatin
Browse files

Draft : start numba integration

parent 5bec007471e7
......@@ -537,7 +537,7 @@ def jit_class(cls, jit_methods):
if mpi.has_to_build(python_path, module.__file__):
from transonic.justintime import _get_module_jit
mod = _get_module_jit(5)
mod = _get_module_jit(backend="pythran", index_frame=5)
if mpi.rank == 0:
python_path = mpi.PathSeq(python_path)
python_code = mod.codes_dependance_classes[cls_name] + "\n"
......
......@@ -3,9 +3,11 @@
"""
import re
import inspect
import beniget
from pathlib import Path
from textwrap import dedent
from typing import Callable
import gast as ast
from transonic.analyses import extast
from transonic.analyses.capturex import CaptureX
......
"""numba Backend
==================
"""
from pathlib import Path
from transonic.analyses import analyse_aot
from transonic.analyses import extast
from transonic.util import (
has_to_build,
get_source_without_decorator,
format_str,
write_if_has_to_write,
)
from transonic.log import logger
from .backend import Backend
class NumbaBackend(Backend):
def __init__(self, paths_py):
super.__init__(paths_py)
def make_backend_files(self):
"""Create numba files from a list of Python files"""
paths_out = []
for path in self.paths_py:
path_out = self.make_numba_file(path)
if path_out:
paths_out.append(path_out)
if paths_out:
nb_files = len(paths_out)
if nb_files == 1:
conjug = "s"
else:
conjug = ""
logger.warning(
f"{nb_files} files created or updated need{conjug}"
" to be pythranized"
)
return paths_out
def make_numba_file(self, path_py: Path, force="force", log_level=None):
if log_level is not None:
logger.set_level(log_level)
path_py = Path(path_py)
if not path_py.exists():
raise FileNotFoundError(f"Input file {path_py} not found")
if path_py.absolute().parent.name == "__numba__":
logger.debug(f"skip file {path_py}")
return
if not path_py.name.endswith(".py"):
raise ValueError(
"transonic only processes Python file. Cannot process {path_py}"
)
path_dir = path_py.parent / "__pythran__"
path_numba = path_dir / path_py.name
if not has_to_build(path_numba, path_py) and not force:
logger.warning(f"File {path_numba} already up-to-date.")
return
code_pythran, code_ext = self.make_numba_code(path_py)
if not code_pythran:
return
for file_name, code in code_ext["function"].items():
path_ext_file = path_dir / (file_name + ".py")
write_if_has_to_write(path_ext_file, code, logger.info)
for file_name, code in code_ext["classe"].items():
path_ext_file = path_dir.parent / "__pythran__" / (file_name + ".py")
write_if_has_to_write(path_ext_file, code, logger.info)
if path_numba.exists() and not force:
with open(path_numba) as file:
code_pythran_old = file.read()
if code_pythran_old == code_pythran:
logger.warning(f"Code in file {path_numba} already up-to-date.")
return
logger.debug(f"code_pythran:\n{code_pythran}")
path_dir.mkdir(exist_ok=True)
with open(path_numba, "w") as file:
file.write(code_pythran)
logger.info(f"File {path_numba} written")
return path_numba
def make_numba_code(self, path_py):
"""Create a pythran code from a Python file"""
with open(path_py) as file:
code_module = file.read()
boosted_dicts, code_dependance, annotations, blocks, code_ext = analyse_aot(
code_module, path_py
)
print(boosted_dicts)
import numpy as np
from transonic import jit
@jit(backend="numba")
def fib(n):
if n < 2:
return n
return np.add(fib(n - 1), fib(n - 2))
@jit
def fib2(n):
if n < 2:
return n
return np.add(fib2(n - 1), fib2(n - 2))
......@@ -86,6 +86,7 @@ from .log import logger
from .config import has_to_replace
from transonic.analyses.justintime import analysis_jit
from transonic.analyses.util import get_source_with_numba
modules = {}
......@@ -106,8 +107,9 @@ def set_compile_jit(value):
class ModuleJIT:
"""Representation of a module using jit"""
def __init__(self, frame=None):
def __init__(self, backend, frame=None):
self.backend = backend
if frame is None:
frame = inspect.stack()[1]
......@@ -155,7 +157,7 @@ class ModuleJIT:
return inspect.getsource(mod)
def _get_module_jit(index_frame: int = 2):
def _get_module_jit(backend="pythran", index_frame: int = 2):
"""Get the ModuleJIT instance corresponding to the calling module
Parameters
......@@ -180,7 +182,7 @@ def _get_module_jit(index_frame: int = 2):
if module_name in modules:
return modules[module_name]
else:
return ModuleJIT(frame=frame)
return ModuleJIT(backend=backend, frame=frame)
def make_pythran_type_name(obj: object):
......@@ -213,11 +215,11 @@ def make_pythran_type_name(obj: object):
return name
def jit(func=None, native=False, xsimd=False, openmp=False):
def jit(func=None, backend="pythran", native=False, xsimd=False, openmp=False):
"""Decorator to record that the function has to be jit compiled
"""
decor = JIT(native=native, xsimd=xsimd, openmp=openmp)
decor = JIT(backend=backend, native=native, xsimd=xsimd, openmp=openmp)
if callable(func):
decor._decorator_no_arg = True
return decor(func)
......@@ -229,7 +231,8 @@ class JIT:
"""Decorator used internally by the public jit decorator
"""
def __init__(self, native=False, xsimd=False, openmp=False):
def __init__(self, backend, native=False, xsimd=False, openmp=False):
self.backend = backend
self.native = native
self.xsimd = xsimd
self.openmp = openmp
......@@ -239,11 +242,14 @@ class JIT:
self.compiling = False
self.process = None
def __call__(self, func):
def __call__(self, func, backend=None):
if not has_to_replace:
return func
if self.backend == "numba":
return self.numba_jit(func)
if is_method(func):
return TransonicTemporaryJITMethod(
func, self.native, self.xsimd, self.openmp
......@@ -259,7 +265,7 @@ class JIT:
else:
index_frame = 2
mod = _get_module_jit(index_frame)
mod = _get_module_jit(self.backend, index_frame)
mod.jit_functions[func_name] = self
module_name = mod.module_name
......@@ -468,3 +474,8 @@ class JIT:
return func(*args, **kwargs)
return type_collector
def numba_jit(self, func):
src = get_source_with_numba(func)
print(src)
return func
......@@ -61,7 +61,7 @@ def test_jit():
sleep(0.1)
def fib():
def test_fib():
from .for_test_justintime import fib, use_fib
fib2 = fib(2)
......
from .for_test_numba import fib2, fib
def test_simple():
fib(2)
fib2(2)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment