Read about our upcoming Code of Conduct on this issue

Commit 2bdc6542 authored by Pierre Blanc-fatin's avatar Pierre Blanc-fatin
Browse files

Get_code_ext refactoring + same name not mandatory for a jit call (see...

Get_code_ext refactoring + same name not mandatory for a jit call (see description) + associated test

With a jitted function called in another jitted function: same name like "func = jit(func)" not mandatory
parent 1b1b07951205
from transonic import jit, boost
from for_test_exterior_import_jit import fib
import numpy as np
def func0(a):
return 2 * a
func0 = jit(func0)
fib = jit(fib)
@jit
def jitted_use_func0():
return func0(1)
@jit
def jitted_use_fib():
return fib(1)
# @boost
# class MyClass2:
# def __init__(self):
# self.attr0 = self.attr1 = 1
# @jit()
# def myfunc(self, arg):
# return self.attr1 + self.attr0 + np.abs(arg) + func0(1)
# def check(self):
# assert self.myfunc(1) == 5
if __name__ == "__main__":
jitted_use_func0()
jitted_use_fib()
from transonic import jit
from for_test_exterior_import_jit import fib
fib_jitted = jit(fib)
if __name__ == "__main__":
fib_jitted(1)
\ No newline at end of file
......@@ -93,20 +93,24 @@ def analysis_jit(code, pathfile):
codes_dependance_classes[key] = capturex.make_code_external()
code_ext = {}
code_ext_cls = {}
def adapt_code_dependance(func, codes_dependance: str):
def adapt_code_dependance(func: str, codes_dependance: str):
"""
Adapt code_dependance to the call of a jitted function in a jitted function:
- Remove the import transonic
- Remove the jitted function statement (i.e func = jit(func))
- Add a import statement to the jitted function
- remove the definition of the jitted function if its on the file, or remove the import statement
"""
module = extast.parse(codes_dependance)
module_body = module.body.copy()
jitted_functions = []
for node in module_body:
print(node)
# remove the transonic import
if isinstance(node, ast.ImportFrom):
if node.module == "transonic":
module.body.remove(node)
# remove the jitted function by jit() (i.e. func = jit(func))
# and add the import statement for the jitted function
elif isinstance(node, ast.Assign):
if node.value.func.id == "jit":
jitted_functions.append(node.targets[0].id)
......@@ -122,101 +126,82 @@ def analysis_jit(code, pathfile):
)
],
)
# if in the module, remove the definition of the jitted function by jit()
# remove the definition:
for node in module_body:
# of the jitted function in the file
if isinstance(node, ast.FunctionDef):
if node.name in jitted_functions:
module.body.remove(node)
# of the jitted imported function
if isinstance(node, ast.ImportFrom):
for name in node.names:
if name.name in jitted_functions:
node.names.remove(name)
# remove the importFrom if no function is imported
if not node.names:
module.body.remove(node)
return extast.unparse(module)
code_ext = {"function": {}, "classe": {}}
def get_exterior_code(
codes_dependance: dict, previous_file_name=None, cls=None
codes_dependance: dict,
previous_file_name=None,
classes: str = None,
relative: bool = None,
):
""" Get all imported functions needed by jitted functions add multiple levels,
i.e get functions needed by functions needed by jitted function.
""" Get all imported functions needed by boosted functions and methods at multiple levels,
(i.e get functions needed by functions imported by boosted function) and add them into code_ext
"""
for func, dep in codes_dependance.items():
if not dep:
continue
module_ext = extast.parse(dep)
for node in module_ext.body:
if isinstance(node, (ast.ImportFrom, ast.Import)):
# get the path of the imported module
file_name, file_path = find_path(node, pathfile)
# a jitted function or method needs another jitted function
if file_name == "transonic":
if cls:
codes_dependance_classes[
func
] = adapt_code_dependance(
func, codes_dependance_classes[func]
)
else:
if dep:
module_ext = extast.parse(dep)
for node in module_ext.body:
if isinstance(node, (ast.ImportFrom, ast.Import)):
# get the path of the imported module
file_name, file_path = find_path(node, pathfile)
# a jitted function or method needs another jitted function
if file_name == "transonic":
codes_dependance[func] = adapt_code_dependance(
func, codes_dependance[func]
)
elif file_name:
file_name = "__" + func + "__" + file_name
# get the content of the file
with open(str(file_path), "r") as file:
content = file.read()
mod = extast.parse(content)
if cls:
elif file_name:
file_name = "__" + func + "__" + file_name
# get the content of the file
with open(str(file_path), "r") as file:
content = file.read()
mod = extast.parse(content)
# filter the code and add it to code_ext dict
code_ext_cls[file_name] = str(
code_ext[classes][file_name] = str(
filter_external_code(mod, node.names)
)
# change imported module names
codes_dependance_classes[func] = change_import_name(
codes_dependance_classes[func],
node,
func,
cls="yes",
codes_dependance[func] = change_import_name(
codes_dependance[func], node, func, relative
)
# recursively get the exteriors codes
if code_ext_cls[file_name]:
# recursively get the exterior codes
if code_ext[classes][file_name]:
get_exterior_code(
{func: code_ext_cls[file_name]},
{func: code_ext[classes][file_name]},
file_name,
cls="yes",
classes,
)
if previous_file_name:
code_ext_cls[
code_ext[classes][
previous_file_name
] = change_import_name(
code_ext_cls[previous_file_name],
code_ext[classes][previous_file_name],
node,
func,
cls="yes",
)
else:
# filter the code and add it to code_ext dict
code_ext[file_name] = str(
filter_external_code(mod, node.names)
)
# change imported module names
codes_dependance[func] = change_import_name(
codes_dependance[func], node, func
)
# recursively get the exteriors codes
if code_ext[file_name]:
get_exterior_code(
{func: code_ext[file_name]}, file_name
)
if previous_file_name:
code_ext[
previous_file_name
] = change_import_name(
code_ext[previous_file_name], node, func
relative,
)
return codes_dependance
if jitted_dicts["methods"]:
get_exterior_code(codes_dependance_classes, cls="yes")
get_exterior_code(codes_dependance)
return (
jitted_dicts,
codes_dependance,
codes_dependance_classes,
code_ext,
code_ext_cls,
codes_dependance_classes = get_exterior_code(
codes_dependance_classes, classes="classe", relative=False
)
codes_dependance = get_exterior_code(
codes_dependance, classes="function", relative=False
)
return (jitted_dicts, codes_dependance, codes_dependance_classes, code_ext)
......@@ -201,7 +201,7 @@ def find_path(node: object, pathfile: str):
def change_import_name(
code_dep: str, changed_node: object, func_name: str, cls: str = None
code_dep: str, changed_node: object, func_name: str, relative: str = None
):
""" Change the name of changed_node in code_dep by adding "__" + func + "__"
at the beginning of the imported module, and return the modified code
......@@ -213,7 +213,7 @@ def change_import_name(
node.module = "__" + func_name + "__" + node.module
elif isinstance(node, ast.Import):
node.names[0].name = "__" + func_name + "__" + node.names[0].name
if cls:
if not relative:
node.level = 0
return extast.unparse(mod)
......
......@@ -84,7 +84,15 @@ jitted_func_import = jit(func_import)
@jit
def main(add: int):
return foo + add + func_import_multilevel() + pi - pi
return (
foo
+ add
+ func_import_multilevel()
+ pi
- pi
- jitted_func_import()
+ func0_jitted(1 / 2)
)
@boost
......
......@@ -123,17 +123,17 @@ class ModuleJIT:
self.jit_functions = {}
source = self.get_source()
self.jitted_dicts, self.codes_dependance, self.codes_dependance_classes, self.code_ext, self.code_ext_cls = analysis_jit(
self.jitted_dicts, self.codes_dependance, self.codes_dependance_classes, self.code_ext = analysis_jit(
source, self.filename
)
# Write exterior code for functions
for file_name, code in self.code_ext.items():
for file_name, code in self.code_ext["function"].items():
path_ext = path_jit / self.module_name.replace(".", os.path.sep)
path_ext_file = path_ext / (file_name + ".py")
write_if_has_to_write(path_ext_file, code, logger.info)
# Write exterior code for classes
for file_name, code in self.code_ext_cls.items():
for file_name, code in self.code_ext["classe"].items():
path_ext = (
path_jit.parent
/ "__jit_classes__"
......
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