cython.py 6.3 KB
Newer Older
1
2
3
4
5
"""Cython Backend
==================


"""
6
import copy
Pierre Augier's avatar
Pierre Augier committed
7
import inspect
8

Pierre Augier's avatar
Pierre Augier committed
9
from transonic.analyses import extast
10
11
12
13
from transonic.annotation import (
    compute_pythran_type_from_type,
    make_signatures_from_typehinted_func,
)
14

Pierre Augier's avatar
Pierre Augier committed
15
16
from .base import BackendAOT
from .base_jit import SubBackendJIT
17
18


Pierre Augier's avatar
Pierre Augier committed
19
20
21
22
23
24
def compute_cython_type_from_pythran_type(type_):

    if isinstance(type_, type):
        type_ = compute_pythran_type_from_type(type_)

    if type_.endswith("]"):
25
        start, end = type_.split("[", 1)
Pierre Augier's avatar
Pierre Augier committed
26
27
        if not start.startswith("np."):
            start = "np." + start
28
29
30
31
32
33
34
35

        dim = end.count("[") + 1
        if dim > 1:
            end = ",".join(":" * dim) + "]"

        if end == "]":
            end = ":]"

Pierre Augier's avatar
Pierre Augier committed
36
37
38
        return start + "_t[" + end

    if any(type_.endswith(str(number)) for number in (32, 64, 128)):
39
        return "np." + type_ + "_t"
Pierre Augier's avatar
Pierre Augier committed
40

Pierre Augier's avatar
Pierre Augier committed
41
42
43
    if "dict" in type_:
        return "dict"

Pierre Augier's avatar
Pierre Augier committed
44
45
46
    return "cython." + type_


Pierre Augier's avatar
Pierre Augier committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class HeaderFunction:
    def __init__(
        self,
        path=None,
        name=None,
        arguments=None,
        types: dict = None,
        imports=None,
    ):

        if path is not None:
            self.path = path
            with open(path) as file:
                lines = file.readlines()

            last_line = lines[-1]
            assert last_line.startswith("cpdef ")
            name = last_line.split(" ", 1)[1].split("(", 1)[0]

            parts = [
                part.strip()
                for part in "".join(lines[:-1]).split("ctypedef fused ")
            ]
            imports = parts[0]

            types = {}

            for part in parts[1:]:
                assert part.startswith(f"__{name}_")
                lines = part.split("\n")
                arg_name = lines[0].split(f"__{name}_", 1)[1].split(":", 1)[0]
                types[arg_name] = set(line.strip() for line in lines[1:])

        if types is None:
            if arguments is None:
                raise ValueError
            types = {key: set() for key in arguments}

        if arguments is None:
            arguments = types.keys()

        self.arguments = arguments
        self.name = name
        self.types = types
        self.imports = imports

    def make_code(self):

        bits = [self.imports + "\n\n"]

        for arg, types in self.types.items():
            bits.append(f"ctypedef fused __{self.name}_{arg}:\n")
            for type_ in sorted(types):
                bits.append(f"    {type_}\n")
            bits.append("\n")

        tmp = ", ".join(f"__{self.name}_{arg} {arg}" for arg in self.types)
        bits.append(f"cpdef {self.name}({tmp})")
        code = "".join(bits)
        return code

    def add_signature(self, new_types):
        for new_type, set_types in zip(new_types, self.types.values()):
            set_types.add(new_type)

    def update_with_other_header(self, other):
        if self.name != other.name:
            raise ValueError
        if self.types.keys() != other.types.keys():
            raise ValueError
        for key, value in other.types.items():
            self.types[key].update(value)


Pierre Augier's avatar
Pierre Augier committed
121
class SubCythonJIT(SubBackendJIT):
122
123
124
125
126
127
128
129
    def compute_typename_from_object(self, obj: object):
        """return the Pythran type name"""
        return compute_cython_type_from_pythran_type(
            super().compute_typename_from_object(obj)
        )

    def make_new_header(self, func, arg_types):
        # Include signature comming from type hints
Pierre Augier's avatar
Pierre Augier committed
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
        header = HeaderFunction(
            name=func.__name__,
            arguments=list(inspect.signature(func).parameters.keys()),
            imports="import cython\n\nimport numpy as np\ncimport numpy as np\n",
        )

        signatures = make_signatures_from_typehinted_func(func, as_list_str=True)

        for signature in signatures:
            header.add_signature(
                compute_cython_type_from_pythran_type(type_)
                for type_ in signature
            )

        if arg_types != "no types":
            header.add_signature(arg_types)

        return header
148
149

    def _load_old_header(self, path_backend_header):
Pierre Augier's avatar
Pierre Augier committed
150
        return HeaderFunction(path=path_backend_header)
151
152

    def _merge_header_objects(self, header, header_old):
Pierre Augier's avatar
Pierre Augier committed
153
154
        header.update_with_other_header(header_old)
        return header
155
156

    def _make_header_code(self, header):
Pierre Augier's avatar
Pierre Augier committed
157
        return header.make_code()
158
159


Pierre Augier's avatar
Pierre Augier committed
160
class CythonBackend(BackendAOT):
161
    backend_name = "cython"
162
    suffix_header = ".pxd"
163
    keyword_export = "cpdef"
Pierre Augier's avatar
Pierre Augier committed
164
    _SubBackendJIT = SubCythonJIT
165

166
167
168
    def _make_first_lines_header(self):
        return ["import cython\n\nimport numpy as np\ncimport numpy as np\n"]

169
170
171
172
173
174
175
176
177
178
179
    def _make_header_from_fdef_signatures(
        self, fdef, signatures_as_lists_strings, locals_types=None
    ):

        fdef = extast.ast.FunctionDef(
            name=fdef.name,
            args=copy.deepcopy(fdef.args),
            body=[],
            decorator_list=[],
            returns=None,
        )
180
181
        signatures_func = []
        if signatures_as_lists_strings:
182
183
            # produce ctypedef
            index = 0
Pierre Augier's avatar
Pierre Augier committed
184
            name_type_args = []
185
            for arg in [name.id for name in fdef.args.args]:
186
                ctypedef = []
Pierre Augier's avatar
Pierre Augier committed
187
                name_type_arg = f"__{fdef.name}_{arg}"
Pierre Augier's avatar
Pierre Augier committed
188
189
                name_type_args.append(name_type_arg)
                ctypedef.append(f"ctypedef fused {name_type_arg}:\n")
190
                possible_types = [x[index] for x in signatures_as_lists_strings]
191
                for possible_type in sorted(set(possible_types)):
Pierre Augier's avatar
Pierre Augier committed
192
193
194
                    ctypedef.append(
                        f"   {compute_cython_type_from_pythran_type(possible_type)}\n"
                    )
195
                index += 1
196
                signatures_func.append("".join(ctypedef))
197

198
            # change function parameters
Pierre Augier's avatar
Pierre Augier committed
199
200
201
202
203
            for name in fdef.args.args:
                name.annotation = None
                name.id = name_type_args[0] + " " + name.id
                del name_type_args[0]

204
205
        if locals_types is not None and locals_types:
            # TODO: fused types not supported here
Pierre Augier's avatar
Pierre Augier committed
206
207
208
209
210
211
212
            locals_types = ", ".join(
                f"{k}={compute_cython_type_from_pythran_type(v)}"
                for k, v in locals_types.items()
            )
            signatures_func.append(f"@cython.locals({locals_types})\n")

        signatures_func.append("cp" + extast.unparse(fdef).strip()[:-1] + "\n")
213
        return signatures_func