Commit 81067576 authored by Pierre Augier's avatar Pierre Augier
Browse files

Prepare proper backend type formatter

parent 111e4c2043f4
from runpy import run_module
import numpy as np
inp = np.ones(2)
def check(module_name, should_return=None):
mod = run_module(module_name)
trans = mod["Transmitter"](1.)
assert trans.__call__.__doc__ == "My docstring", trans.__call__.__doc__
assert trans.__call__.__name__ == "__call__"
result = trans(inp)
if should_return is not None:
assert np.allclose(result, should_return)
return result
original = check("transmitter")
check("try0", original)
check("try1", original)
check("try2", original)
check("try3", original)
check("try4", original)
import numpy as np
# pythran export __for_method__Transmitter____call__(float, float[:])
# pythran export __for_method__Transmitter____call__(float, int[:])
# pythran export __for_method__Transmitter____call__(float, float[:, :])
# pythran export __for_method__Transmitter____call__(float, int[:, :])
def __for_method__Transmitter____call__(self_freq, inp):
"My docstring"
print("__for_method...")
return inp * np.exp(((np.arange(len(inp)) * self_freq) * 1j))
# pythran export __code_new_method__Transmitter____call__
__code_new_method__Transmitter____call__ = """
def new_method(self, inp):
print("new_method")
return backend_func(self.freq, inp)
"""
import numpy as np
class Transmitter():
freq: float
def __init__(self, freq):
self.freq = float(freq)
def __call__(self, inp: "float[]"):
"""My docstring"""
return inp * np.exp(np.arange(len(inp))*self.freq*1j)
from functools import wraps
import numpy as np
# pythran export __method__Transmitter____call__(float, float[])
def __for_method__Transmitter____call__(self_freq, inp):
return inp * np.exp(np.arange(len(inp))*self_freq*1j)
def __method__Transmitter____call__(self, inp):
self_freq = self.freq
return __for_method__Transmitter____call__(self_freq, inp)
def __call__(self, inp: "float[]"):
"""My docstring"""
return inp * np.exp(np.arange(len(inp))*self.freq*1j)
class Transmitter():
freq: float
def __init__(self, freq):
self.freq = float(freq)
__call__ = wraps(__call__)(__method__Transmitter____call__)
\ No newline at end of file
from functools import wraps
import numpy as np
# pythran export __method__Transmitter____call__(float, float[])
def __for_method__Transmitter____call__(self_freq, inp):
return inp * np.exp(np.arange(len(inp))*self_freq*1j)
def __method__Transmitter____call__(self, inp):
self_freq = self.freq
return __for_method__Transmitter____call__(self_freq, inp)
def my_decor(func):
return wraps(func)(__method__Transmitter____call__)
class Transmitter():
freq: float
def __init__(self, freq):
self.freq = float(freq)
@my_decor
def __call__(self, inp: "float[]"):
"""My docstring"""
return inp * np.exp(np.arange(len(inp))*self.freq*1j)
\ No newline at end of file
from functools import wraps
import numpy as np
import pythran_manual
code = """
def new_method(self, inp):
return backend_func(self.freq, inp)
"""
def my_decor(func):
func_name = func.__name__
name_backend_func = "__for_method__Transmitter__" + func_name
backend_func = pythran_manual.__dict__[name_backend_func]
namespace = {"backend_func": backend_func}
exec(code, namespace)
return wraps(func)(namespace["new_method"])
class Transmitter:
freq: float
def __init__(self, freq):
self.freq = float(freq)
@my_decor
def __call__(self, inp: "float[]"):
"""My docstring"""
return inp * np.exp(np.arange(len(inp)) * self.freq * 1j)
from functools import wraps
import numpy as np
import pythran_manual
class FluidPythranTemporaryMethod:
def __init__(self, func):
self.func = func
def __call__(self, self_bis, *args, **kwargs):
raise RuntimeError("Do not call this function!")
def trans_def_method(func):
return FluidPythranTemporaryMethod(func)
def boost(cls):
cls_name = cls.__name__
for key, value in cls.__dict__.items():
if not isinstance(value, FluidPythranTemporaryMethod):
continue
func = value.func
func_name = func.__name__
name_backend_func = f"__for_method__{cls_name}__{func_name}"
backend_func = pythran_manual.__dict__[name_backend_func]
name_var_code_new_method = f"__code_new_method__{cls_name}__{func_name}"
code_new_method = pythran_manual.__dict__[name_var_code_new_method]
namespace = {"backend_func": backend_func}
exec(code_new_method, namespace)
setattr(cls, key, wraps(func)(namespace["new_method"]))
return cls
@boost
class Transmitter:
freq: float
def __init__(self, freq):
self.freq = float(freq)
@trans_def_method
def __call__(self, inp: "float[]"):
"""My docstring"""
return inp * np.exp(np.arange(len(inp)) * self.freq * 1j)
if __name__ == "__main__":
inp = np.ones(2)
trans = Transmitter(1.)
trans(inp)
import numpy as np
from transonic import boost
@boost
class Transmitter:
freq: float
def __init__(self, freq):
self.freq = float(freq)
@boost
def __call__(self, inp: "float[]"):
"""My docstring"""
return inp * np.exp(np.arange(len(inp)) * self.freq * 1j)
if __name__ == "__main__":
inp = np.ones(2)
freq = 1.0
trans = Transmitter(freq)
def for_check(freq, inp):
return inp * np.exp(np.arange(len(inp)) * freq * 1j)
assert np.allclose(trans(inp), for_check(freq, inp))
from tokenize import tokenize, untokenize, NAME, OP
from token import tok_name
from io import BytesIO
import inspect
import numpy as np
from transonic.log import logger
from transonic.util import get_source_without_decorator, format_str
from transonic.typing import compute_signatures_from_typeobjects
def trans_def_method(func):
func.__transonic__ = "trans_def_method"
return func
# logger.set_level("debug")
from transonic import Array, Type, NDim
A = Array[Type(float, int), NDim(1, 2)]
class Transmitter:
freq: float
def __init__(self):
pass
@trans_def_method
def __call__(self, inp: A):
"""My docstring"""
return inp * np.exp(np.arange(len(inp)) * self.freq * 1j)
@trans_def_method
def call_with_print(self, inp: A):
"""call + print"""
print("call_with_print")
return inp * np.exp(np.arange(len(inp)) * self.freq * 1j)
def produce_pythran_code_class(cls):
pythran_code = ""
for key, value in cls.__dict__.items():
if (
hasattr(value, "__transonic__")
and value.__transonic__ == "trans_def_method"
):
pythran_code += produce_pythran_code_class_func(cls, key)
return pythran_code
def produce_pythran_code_class_func(cls, func_name):
cls_name = cls.__name__
cls_annotations = cls.__annotations__
func = cls.__dict__[func_name]
signature = inspect.signature(func)
types_func = [param.annotation for param in signature.parameters.values()][1:]
args_func = list(signature.parameters.keys())[1:]
src = get_source_without_decorator(func)
tokens = []
attributes = set()
using_self = False
g = tokenize(BytesIO(src.encode("utf-8")).readline)
for toknum, tokval, a, b, c in g:
logger.debug((tok_name[toknum], tokval))
if using_self == "self":
if toknum == OP and tokval == ".":
using_self = tokval
continue
elif toknum == OP and tokval == ",":
tokens.append((NAME, "self"))
using_self = False
else:
raise RuntimeError
if using_self == ".":
if toknum == NAME and tokval in cls_annotations:
using_self = False
tokens.append((NAME, "self_" + tokval))
attributes.add(tokval)
continue
else:
raise RuntimeError
if toknum == NAME and tokval == "self":
using_self = tokval
continue
tokens.append((toknum, tokval))
attributes = sorted(attributes)
types_attrs = [cls_annotations[attr] for attr in attributes]
attributes_self = ["self_" + attr for attr in attributes]
index_self = tokens.index((NAME, "self"))
tokens_attr = []
for ind, attr in enumerate(attributes_self):
tokens_attr.append((NAME, attr))
tokens_attr.append((OP, ","))
tokens = tokens[:index_self] + tokens_attr + tokens[index_self + 2 :]
index_func_name = tokens.index((NAME, func_name))
name_backend_func = f"__for_method__{cls_name}__{func_name}"
tokens[index_func_name] = (NAME, name_backend_func)
new_code = untokenize(tokens).decode("utf-8")
new_code = format_str(new_code)
# args_pythran = attributes + args_func
types_pythran = types_attrs + types_func
pythran_signatures = "\n"
for types_string_signature in compute_signatures_from_typeobjects(
types_pythran
):
pythran_signatures += (
"# pythran export "
+ name_backend_func
+ "("
+ ", ".join(types_string_signature)
+ ")\n"
)
pythran_code = pythran_signatures + "\n" + new_code
name_var_code_new_method = f"__code_new_method__{cls_name}__{func_name}"
str_self_dot_attributes = ", ".join("self." + attr for attr in attributes)
str_args_func = ", ".join(args_func)
pythran_code += (
f"\n# pythran export {name_var_code_new_method}\n"
f'\n{name_var_code_new_method} = """\n\n'
f"def new_method(self, {str_args_func}):\n"
f" return backend_func({str_self_dot_attributes}, {str_args_func})"
'\n\n"""\n'
)
return pythran_code
pythran_code = produce_pythran_code_class(Transmitter)
print(pythran_code)
......@@ -13,7 +13,8 @@ Internal API
from pathlib import Path
from textwrap import indent
from typing import Iterable, Optional
from pprint import pprint
# from pprint import pprint
import transonic
......@@ -34,6 +35,7 @@ from transonic.util import (
from .base_jit import SubBackendJIT
from .for_classes import make_new_code_method_from_nodes
from .typing import TypeFormatter
class Backend:
......@@ -46,11 +48,13 @@ class Backend:
keyword_export = "export"
_SubBackendJIT = SubBackendJIT
needs_compilation = True
_TypeFormatter = TypeFormatter
def __init__(self):
self.name = self.backend_name
self.name_capitalized = self.name.capitalize()
self.jit = self._SubBackendJIT(self.name)
self.type_formatter = TypeFormatter()
self.jit = self._SubBackendJIT(self.name, self.type_formatter)
def _make_code_from_fdef_node(self, fdef):
transformed = TypeHintRemover().visit(fdef)
......@@ -466,14 +470,14 @@ class BackendAOT(Backend):
):
signatures_as_lists_strings = []
for annot in annotations:
print("DEBUG, annot")
pprint(annot)
# print("DEBUG, annot")
# pprint(annot)
signatures_as_lists_strings.extend(
compute_signatures_from_typeobjects(annot)
compute_signatures_from_typeobjects(annot, self.type_formatter)
)
print("DEBUG, signatures_as_lists_strings")
pprint(signatures_as_lists_strings)
# print("DEBUG, signatures_as_lists_strings")
# pprint(signatures_as_lists_strings)
return self._make_header_from_fdef_signatures(
fdef,
......
......@@ -32,9 +32,10 @@ from .for_classes import produce_code_class
class SubBackendJIT:
"""Sub-class for Transonic JIT compilation"""
def __init__(self, name):
def __init__(self, name, backend_type_formatter):
self.name = name
self.name_capitalized = name.capitalize()
self.backend_type_formatter = backend_type_formatter
def make_backend_source(self, info_analysis, func, path_backend):
func_name = func.__name__
......@@ -62,7 +63,9 @@ class SubBackendJIT:
def make_new_header(self, func, arg_types):
# Include signature comming from type hints
signatures = make_signatures_from_typehinted_func(func)
signatures = make_signatures_from_typehinted_func(
func, self.backend_type_formatter
)
exports = set(f"export {signature}" for signature in signatures)
if arg_types != "no types":
......
......@@ -25,12 +25,13 @@ import itertools
from transonic.analyses.extast import unparse, ast, FunctionDef, Name
from transonic.typing import (
compute_pythran_type_from_type,
format_type_as_backend_type,
make_signatures_from_typehinted_func,
)
from .base import BackendAOT, TypeHintRemover, format_str
from .base_jit import SubBackendJIT
from .typing import base_type_formatter
def analyze_array_type(type_):
......@@ -66,7 +67,7 @@ def np_ndarray_type(type_) -> str:
def compute_cython_type_from_pythran_type(type_, memoryview=False):
if isinstance(type_, type):
type_ = compute_pythran_type_from_type(type_)
type_ = format_type_as_backend_type(type_, base_type_formatter)
if type_.endswith("]"):
if memoryview:
......@@ -171,7 +172,9 @@ class SubBackendJITCython(SubBackendJIT):
imports="import cython\n\nimport numpy as np\ncimport numpy as np\n",
)
signatures = make_signatures_from_typehinted_func(func, as_list_str=True)
signatures = make_signatures_from_typehinted_func(
func, self.backend_type_formatter, as_list_str=True
)
for signature in signatures:
header.add_signature(
......@@ -271,8 +274,8 @@ class CythonBackend(BackendAOT):
template_variables = dict(zip(names, set_types))
cython_fused_types[name_cython_type].append(
compute_cython_type_from_pythran_type(
compute_pythran_type_from_type(
ttype, **template_variables
format_type_as_backend_type(
ttype, self.type_formatter, **template_variables
)
)
)
......
......@@ -18,6 +18,8 @@ from transonic.util import (
make_code_from_fdef_node,
)
from .typing import base_type_formatter
def produce_code_class(cls):
pythran_code = ""
......@@ -56,7 +58,7 @@ def make_code_method_jit(cls, func_name):
try:
signatures_as_lists_strings = compute_signatures_from_typeobjects(
types_pythran
types_pythran, base_type_formatter
)
except ValueError:
signatures_as_lists_strings = []
......
class TypeFormatter:
pass
base_type_formatter = TypeFormatter()
......@@ -37,9 +37,9 @@ Internal API
.. autofunction:: compute_signatures_from_typeobjects
.. autofunction:: compute_pythran_types_from_types
.. autofunction:: format_types_as_backend_types
.. autofunction:: make_signature_from_template_variables
.. autofunction:: _make_signature_from_template_variables
.. autofunction:: make_signatures_from_typehinted_func
......@@ -126,7 +126,7 @@ class Type(TemplateVar):
def __repr__(self):
return self.__name__
def format_as_backend_type(self, **kwargs):
def format_as_backend_type(self, backend_type_formatter, **kwargs):
dtype = None
......@@ -309,7 +309,7 @@ class ArrayMeta(type):
return "Array[" + ", ".join(strings) + "]"
def format_as_backend_type(self, **kwargs):
def format_as_backend_type(self, backend_type_formatter, **kwargs):
dtype = ndim = None
......@@ -386,9 +386,11 @@ class UnionMeta(type):
strings.append(string)
return "Union[" + ", ".join(strings) + "]"
def format_as_backend_type(self, **kwargs):
def format_as_backend_type(self, backend_type_formatter, **kwargs):
type_ = kwargs.pop(self.template_var.__name__)
return compute_pythran_type_from_type(type_, **kwargs)
return format_type_as_backend_type(
type_, backend_type_formatter, **kwargs
)
class Union(metaclass=UnionMeta):
......@@ -414,9 +416,11 @@ def normalize_type_name(name):
return name
def compute_pythran_type_from_type(type_, **kwargs):
def format_type_as_backend_type(type_, backend_type_formatter, **kwargs):