Read about our upcoming Code of Conduct on this issue

Commit 2555a09a authored by Pierre Augier's avatar Pierre Augier
Browse files

cachedjit for methods!

parent 7e7cadeca2f7
......@@ -40,6 +40,7 @@ from .util import (
modification_date,
is_method,
path_root,
path_cachedjit_classes,
)
from .pythranizer import (
......@@ -53,6 +54,8 @@ from .annotation import (
make_signatures_from_typehinted_func,
)
from .transpiler import produce_code_class
from .log import logger
from . import mpi
from .mpi import Path
......@@ -64,7 +67,7 @@ if mpi.nb_proc == 1:
is_transpiling = False
modules = {}
path_cachedjit_classes = mpi.Path(path_root) / "__cachedjit_classes__"
path_cachedjit_classes = mpi.Path(path_cachedjit_classes)
def _get_fluidpythran_calling_module(index_frame: int = 2):
......@@ -431,17 +434,21 @@ class FluidPythran:
self.classes[cls.__name__] = cls
return cls
jit_methods = {
key: value
for key, value in cls.__dict__.items()
if isinstance(value, FluidPythranTemporaryJITMethod)
}
if jit_methods:
cls = cachedjit_class(cls, jit_methods)
if not self.is_transpiled:
return cls
cls_name = cls.__name__
jit_methods = {}
for key, value in cls.__dict__.items():
if isinstance(value, FluidPythranTemporaryJITMethod):
jit_methods[key] = value
if not isinstance(value, FluidPythranTemporaryMethod):
continue
func = value.func
......@@ -468,10 +475,7 @@ class FluidPythran:
exec(code_new_method, namespace)
setattr(cls, key, functools.wraps(func)(namespace["new_method"]))
if not jit_methods:
return cls
return cachedjit_class(cls, jit_methods)
return cls
def make_signature(self, func, _signature=None, **kwargs):
"""Create signature for a function with values for the template types
......@@ -561,6 +565,8 @@ class FluidPythranTemporaryMethod:
class FluidPythranTemporaryJITMethod:
__fluidpythran__ = "jit_method"
def __init__(self, func, native, xsimd, openmp):
self.func = func
self.native = native
......@@ -581,23 +587,48 @@ def cachedjit_class(cls, jit_methods):
3. replace the methods
"""
# FIXME: to be implemented
cls_name = cls.__name__
mod_name = cls.__module__.__name__
mod_name = cls.__module__
module = sys.modules[mod_name]
# 1. create a Python file with @cachejit functions and methods
python_path_dir = path_cachedjit_classes / mod_name.replace(".", os.path.sep)
python_path = python_path_dir / (cls_name + ".py")
if has_to_build(python_path, cls.__module__.__file__):
python_code = ...
if has_to_build(python_path, module.__file__):
if mpi.rank == 0:
python_path = mpi.PathSeq(python_path)
python_code = produce_code_class(cls, jit=True)
if python_path.exists():
if mpi.rank == 0:
has_to_write = None
if python_path.exists():
with open(python_path, "r") as file:
python_code_file = file.read()
# if mpi.rank == 0:
# path_cachedjit_classes.mkdir(parents=True, exist_ok=True)
# mpi.barrier()
raise NotImplementedError
if python_code_file != python_code:
has_to_write = True
else:
has_to_write = True
if has_to_write:
python_path_dir.mkdir(exist_ok=True, parents=True)
with open(python_path, "w") as file:
file.write(python_code)
python_path = mpi.Path(python_path)
mpi.barrier()
# 2. import the file
python_mod_name = (
path_cachedjit_classes.name + "." + mod_name + "." + cls_name
)
module = import_from_path(python_path, python_mod_name)
# 3. replace the methods
for name_method, method in jit_methods.items():
func = method.func
name_new_method = f"__new_method__{cls.__name__}__{name_method}"
new_method = getattr(module, name_new_method)
setattr(cls, name_method, functools.wraps(func)(new_method))
return cls
......@@ -2,7 +2,7 @@ import numpy as np
# pythran import numpy as np
from fluidpythran import cachedjit, used_by_cachedjit
from fluidpythran import cachedjit, used_by_cachedjit, boost
from .mpi import Path
......@@ -40,3 +40,16 @@ def func2(a):
@cachedjit()
def func_dict(d: "str: float dict"):
return d.popitem()
@boost
class MyClass:
def __init__(self):
self.attr0 = self.attr1 = 1
@cachedjit
def myfunc(self, arg):
return self.attr1 + self.attr0 + arg
def check(self):
assert self.myfunc(1) == 3
import sys
import os
from time import sleep
import numpy as np
......@@ -8,16 +9,24 @@ try:
except ImportError:
pythran = None
from .util import path_cachedjit_classes
from .justintime import path_cachedjit, modules
from .pythranizer import scheduler, wait_for_all_extensions
from . import mpi
from .compat import rmtree
scheduler.nb_cpus = 2
module_name = "fluidpythran.for_test_justintime"
path_pythran_dir = mpi.PathSeq(path_cachedjit) / module_name.replace(
".", os.path.sep
str_relative_path = module_name.replace(".", os.path.sep)
path_cachedjit = mpi.PathSeq(path_cachedjit)
path_pythran_dir = path_cachedjit / str_relative_path
path_classes_dir = path_cachedjit_classes / str_relative_path
path_classes_dir1 = (
path_cachedjit / path_cachedjit_classes.name / str_relative_path
)
......@@ -32,12 +41,15 @@ if mpi.rank == 0:
delete_pythran_files("func2")
delete_pythran_files("func_dict")
if path_classes_dir.exists():
rmtree(path_classes_dir)
rmtree(path_classes_dir1)
mpi.barrier()
def test_cachedjit():
from time import sleep
from .for_test_justintime import func1
a = np.arange(2)
......@@ -50,7 +62,6 @@ def test_cachedjit():
def test_cachedjit_simple():
from time import sleep
from .for_test_justintime import func2
func2(1)
......@@ -78,7 +89,6 @@ def test_cachedjit_simple():
def test_cachedjit_dict():
from time import sleep
from .for_test_justintime import func_dict
d = dict(a=1, b=2)
......@@ -103,3 +113,20 @@ def test_cachedjit_dict():
sleep(0.1)
func_dict(d)
break
def test_cachedjit_method():
from .for_test_justintime import MyClass
obj = MyClass()
obj.check()
if not pythran:
return
obj = MyClass()
obj.check()
wait_for_all_extensions()
obj.check()
......@@ -258,14 +258,14 @@ def make_pythran_code_functions(
return code_pythran
def produce_pythran_code_class(cls):
def produce_code_class(cls, jit=False):
pythran_code = ""
for key, value in cls.__dict__.items():
if (
hasattr(value, "__fluidpythran__")
and value.__fluidpythran__ == "pythran_def_method"
if hasattr(value, "__fluidpythran__") and value.__fluidpythran__ in (
"pythran_def_method",
"jit_method",
):
pythran_code += produce_code_class_func(cls, key)
pythran_code += produce_code_class_func(cls, key, jit)
return pythran_code
......@@ -339,6 +339,9 @@ def produce_code_class_func(cls, func_name, jit=False):
func = cls.__dict__[func_name]
if jit:
func = func.func
new_code, attributes, name_new_func = produce_new_code_method(cls, func)
try:
......@@ -386,7 +389,9 @@ def produce_code_class_func(cls, func_name, jit=False):
)
if jit:
new_code = "@cachejit\n" + new_code
new_code = (
"from fluidpythran import cachedjit\n\n@cachedjit\n" + new_code
)
python_code = pythran_signatures + "\n" + new_code
......@@ -492,7 +497,7 @@ imports: {imports}\n"""
signatures_func_all[name_func].extend(signatures)
for cls in fp.classes.values():
code_pythran += produce_pythran_code_class(cls)
code_pythran += produce_code_class(cls)
else:
if "# FLUIDPYTHRAN_NO_IMPORT" not in code:
......
......@@ -65,6 +65,8 @@ path_root = Path(
os.environ.get("FLUIDPYTHRAN_DIR", Path.home() / ".fluidpythran")
)
path_cachedjit_classes = path_root / "__cachedjit_classes__"
def find_module_name_from_path(path_py: Path):
......@@ -72,6 +74,19 @@ def find_module_name_from_path(path_py: Path):
path = path_py.absolute().parent
module_name = path_py.stem
# special case for cachedjit_classes
try:
path_rel = path.relative_to(path_cachedjit_classes)
except ValueError:
pass
else:
tmp = [path_cachedjit_classes.name]
name_pack = str(path_rel).replace(os.path.sep, ".")
if name_pack:
tmp.append(name_pack)
tmp.append(module_name)
return ".".join(tmp)
while path.parents:
if path == cwd or str(path) in sys.path:
return module_name
......
from fluidpythran.transpiler import produce_new_code_class_func
def boost(obj):
return obj
def cachedjit(method):
return method
@boost
class MyClass:
def __init__(self):
self.attr0 = self.attr1 = 1
@cachedjit
def myfunc(self, arg):
return self.attr1 + self.attr0 + arg
python_code = produce_new_code_class_func(MyClass, "myfunc", jit=True)
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