Read about our upcoming Code of Conduct on this issue

This instance will be upgraded to Heptapod 0.28.1 on 2022-01-26 at 16:00 UTC+1 (a few minutes of down time)

Commit 8488db24 authored by Pierre Blanc-fatin's avatar Pierre Blanc-fatin
Browse files

Jitted func : fix different names issue & many imports in a single line

parent d2657f8430bd
......@@ -238,13 +238,13 @@ def analyse_aot(code, pathfile):
code_dependance = capturex.make_code_external()
if boosted_dicts["functions"]:
func = next(iter(boosted_dicts["functions"]))
code_dependance, code_ext = get_exterior_code(
code_dependance, code_ext, jitted_dicts, spe = get_exterior_code(
{func: code_dependance}, pathfile, classes="function", relative=True
)
code_dependance = code_dependance[func]
if boosted_dicts["classes"]:
cls = next(iter(boosted_dicts["classes"]))
code_dependance, code_ext = get_exterior_code(
code_dependance, code_ext, jitted_dicts, spe = get_exterior_code(
{cls: code_dependance}, pathfile, classes="classe", relative=True
)
code_dependance = code_dependance[cls]
......
......@@ -29,8 +29,6 @@ def analysis_jit(code, pathfile):
debug("compute ancestors and chains")
ancestors, duc, udc = compute_ancestors_chains(module)
# boosted_dicts = get_boosted_dicts(module, ancestors, duc)
jitted_dicts = get_decorated_dicts(
module, ancestors, duc, pathfile, decorator="jit"
)
......@@ -97,19 +95,32 @@ def analysis_jit(code, pathfile):
codes_dependance_classes[key] = capturex.make_code_external()
special = []
spe = []
if jitted_dicts["methods"]:
codes_dependance_classes, code_ext = get_exterior_code(
codes_dependance_classes, code_ext, jitted_dicts, spe = get_exterior_code(
codes_dependance_classes,
pathfile,
previous_file_name=None,
classes="classe",
relative=False,
jitted_dicts=jitted_dicts,
)
codes_dependance, code_ext = get_exterior_code(
special = special + spe
codes_dependance, code_ext, jitted_dicts, spe = get_exterior_code(
codes_dependance,
pathfile,
previous_file_name=None,
classes="function",
relative=False,
jitted_dicts=jitted_dicts,
)
special = special + spe
return (
jitted_dicts,
codes_dependance,
codes_dependance_classes,
code_ext,
special,
)
return (jitted_dicts, codes_dependance, codes_dependance_classes, code_ext)
......@@ -2,6 +2,7 @@
=============================
"""
import re
import beniget
from pathlib import Path
from textwrap import dedent
......@@ -254,7 +255,7 @@ def filter_external_code(module: object, names: list):
return code_dependance_annotations + "\n".join(lines_code)
def adapt_code_dependance(func: str, codes_dependance: str):
def adapt_code_dependance(func: str, codes_dependance: str, jitted_dicts: dict):
"""
Adapt code_dependance to the call of a jitted function in a jitted function:
- Remove the import transonic
......@@ -262,6 +263,7 @@ def adapt_code_dependance(func: str, codes_dependance: str):
- 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
"""
special = []
module = extast.parse(codes_dependance)
module_body = module.body.copy()
jitted_functions = []
......@@ -273,20 +275,39 @@ def adapt_code_dependance(func: str, codes_dependance: str):
# 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)
module.body.remove(node)
module.body.insert(
0,
[
extast.parse(
"from "
+ node.value.args[0].id
+ " import "
+ node.value.args[0].id
)
],
)
if isinstance(node.value, ast.Call):
if node.value.func.id == "jit":
# if assigned name different from jitted function name
if node.targets[0].id != node.value.args[0].id:
# change function names in jitted dict
# The list "special" contains functions that has to be write from jitted_dicts
# see transonic.justintime:287
def_func = extast.unparse(jitted_dicts["functions"][func])
spl = re.split("(\W+)", def_func)
spl = [
node.value.args[0].id
if x == node.targets[0].id
else x
for x in spl
]
st = "".join(str(e) for e in spl)
jitted_dicts["functions"][func] = extast.parse(st)
special.append(func)
jitted_functions.append(node.value.args[0].id)
else:
jitted_functions.append(node.targets[0].id)
module.body.remove(node)
module.body.insert(
0,
[
extast.parse(
"from "
+ node.value.args[0].id
+ " import "
+ node.value.args[0].id
)
],
)
# remove the definition:
for node in module_body:
# of the jitted function in the file
......@@ -301,7 +322,7 @@ def adapt_code_dependance(func: str, codes_dependance: str):
# remove the importFrom if no function is imported
if not node.names:
module.body.remove(node)
return extast.unparse(module)
return extast.unparse(module), jitted_dicts, special, jitted_functions
code_ext = {"function": {}, "classe": {}}
......@@ -313,10 +334,13 @@ def get_exterior_code(
previous_file_name=None,
classes: str = None,
relative: bool = None,
jitted_dicts: dict = None,
):
""" 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
"""
special = []
treated = []
for func, dep in codes_dependance.items():
if dep:
module_ext = extast.parse(dep)
......@@ -326,9 +350,15 @@ def get_exterior_code(
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]
codes_dependance[
func
], jitted_dicts, spe, treat = adapt_code_dependance(
func, codes_dependance[func], jitted_dicts
)
# the "special" list signals that jitted functions has to be written
# from jitted_dicts (see transonic/justintime.py:287)
special = special + spe
treated = treated + treat
for func, dep in codes_dependance.items():
if dep:
......@@ -338,14 +368,16 @@ def get_exterior_code(
# 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:
if file_name and not file_name in treated:
new_file_name = "__" + func + "__" + file_name
# get the content of the file
try:
with open(str(file_path), "r") as file:
content = file.read()
except:
break
raise NotImplementedError(
file_name + " can no be found"
)
mod = extast.parse(content)
# filter the code and add it to code_ext dict
code_ext[classes][new_file_name] = str(
......@@ -372,4 +404,4 @@ def get_exterior_code(
func,
relative,
)
return codes_dependance, code_ext
return codes_dependance, code_ext, jitted_dicts, special
......@@ -15,5 +15,9 @@ def func_import():
return const + np.pi - np.pi
def func_import2():
return const + np.pi - np.pi
def use_black():
print(black)
......@@ -71,28 +71,27 @@ class MyClass:
assert self.myfunc(1) == 3
# FIXME Support import in a single line
from .for_test_exterior_import_jit import func_import_multilevel, foo
from .for_test_exterior_import_jit import func_import
# FIXME support multilevel imported jitted function call in a jitted function
from .for_test_exterior_import_jit import foo, func_import, func_import2
from numpy import pi
const = 1
# FIXME support multilevel with jit(func_import)
# FIXME support a different name
jitted_func_import = jit(func_import)
jitted_func_import2 = jit(func_import2)
@jit
def main(add: int):
return (
foo
+ add
+ func_import_multilevel()
add
+ foo
+ pi
- pi
- jitted_func_import()
+ func_import2()
+ func0_jitted(1 / 2)
+ const
)
......@@ -103,7 +102,7 @@ class MyClass2:
@jit()
def myfunc(self, arg):
return self.attr1 + self.attr0 + np.abs(arg) + func_import_multilevel()
return self.attr1 + self.attr0 + np.abs(arg) + func_import()
def check(self):
assert self.myfunc(1) == 5
assert self.myfunc(1) == 4
......@@ -47,6 +47,7 @@ Note: During the compilation (the "warmup" of the JIT), the Python function is
used.
"""
from transonic.analyses import extast
import inspect
import itertools
......@@ -123,7 +124,7 @@ class ModuleJIT:
self.jit_functions = {}
source = self.get_source()
self.jitted_dicts, self.codes_dependance, self.codes_dependance_classes, self.code_ext = analysis_jit(
self.jitted_dicts, self.codes_dependance, self.codes_dependance_classes, self.code_ext, self.special = analysis_jit(
source, self.filename
)
# Write exterior code for functions
......@@ -283,8 +284,15 @@ class JIT:
if has_to_write:
src = mod.codes_dependance[func_name]
src += get_source_without_decorator(func)
if func_name in mod.special:
if func_name in mod.jitted_dicts["functions"]:
src += extast.unparse(
mod.jitted_dicts["functions"][func_name]
)
elif func_name in mod.jitted_dicts["methods"]:
src += extast.unparse(mod.jitted_dicts["methods"][func_name])
else:
src += get_source_without_decorator(func)
if path_pythran.exists() and mpi.rank == 0:
with open(path_pythran) as file:
src_old = file.read()
......
......@@ -170,7 +170,8 @@ def test_func0():
assert result == func0(2)
# jitted function that uses a local function
# jitted function that uses a local function, a jitted local function
# and a jitted imported function (with a different name)
def test_main():
from .for_test_justintime import main
......
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