Aluode/PerceptionLabPortable

1from collections import namedtuple
2
3from numba.core import types, ir
4from numba.core.typing import signature
5
6
7_CallableNode = namedtuple("BoundFunc", ["func", "sig"])
8
9
10class ParforLoweringBuilder:
11    """Helper class for building Numba-IR and lowering for Parfor.
12    """
13    def __init__(self, lowerer, scope, loc):
14        self._lowerer = lowerer
15        self._scope = scope
16        self._loc = loc
17
18    @property
19    def _context(self):
20        return self._lowerer.context
21
22    @property
23    def _typingctx(self):
24        return self._context.typing_context
25
26    @property
27    def _typemap(self):
28        return self._lowerer.fndesc.typemap
29
30    @property
31    def _calltypes(self):
32        return self._lowerer.fndesc.calltypes
33
34    def bind_global_function(self, fobj, ftype, args, kws=None):
35        """Binds a global function to a variable.
36
37        Parameters
38        ----------
39        fobj : object
40            The function to be bound.
41        ftype : types.Type
42        args : Sequence[types.Type]
43        kws : Mapping[str, types.Type]
44
45        Returns
46        -------
47        callable: _CallableNode
48        """
49        if kws is None:
50            kws = {}
51        loc = self._loc
52        varname = f"{fobj.__name__}_func"
53        gvname = f"{fobj.__name__}"
54
55        func_sig = self._typingctx.resolve_function_type(ftype, args, kws)
56        func_var = self.assign(
57            rhs=ir.Global(gvname, fobj, loc=loc), typ=ftype, name=varname
58        )
59        return _CallableNode(func=func_var, sig=func_sig)
60
61    def make_const_variable(self, cval, typ, name="pf_const") -> ir.Var:
62        """Makes a constant variable
63
64        Parameters
65        ----------
66        cval : object
67            The constant value
68        typ : types.Type
69            type of the value
70        name : str
71            variable name to store to
72
73        Returns
74        -------
75        res : ir.Var
76        """
77        return self.assign(
78            rhs=ir.Const(cval, loc=self._loc), typ=typ, name=name
79        )
80
81    def make_tuple_variable(self, varlist, name="pf_tuple") -> ir.Var:
82        """Makes a tuple variable
83
84        Parameters
85        ----------
86        varlist : Sequence[ir.Var]
87            Variables containing the values to be stored.
88        name : str
89            variable name to store to
90
91        Returns
92        -------
93        res : ir.Var
94        """
95        loc = self._loc
96        vartys = [self._typemap[x.name] for x in varlist]
97        tupty = types.Tuple.from_types(vartys)
98        return self.assign(
99            rhs=ir.Expr.build_tuple(varlist, loc), typ=tupty, name=name
100        )
101
102    def assign(self, rhs, typ, name="pf_assign") -> ir.Var:
103        """Assign a value to a new variable
104
105        Parameters
106        ----------
107        rhs : object
108            The value
109        typ : types.Type
110            type of the value
111        name : str
112            variable name to store to
113
114        Returns
115        -------
116        res : ir.Var
117        """
118        loc = self._loc
119        var = self._scope.redefine(name, loc)
120        self._typemap[var.name] = typ
121        assign = ir.Assign(rhs, var, loc)
122        self._lowerer.lower_inst(assign)
123        return var
124
125    def assign_inplace(self, rhs, typ, name) -> ir.Var:
126        """Assign a value to a new variable or inplace if it already exist
127
128        Parameters
129        ----------
130        rhs : object
131            The value
132        typ : types.Type
133            type of the value
134        name : str
135            variable name to store to
136
137        Returns
138        -------
139        res : ir.Var
140        """
141        loc = self._loc
142        var = ir.Var(self._scope, name, loc)
143        assign = ir.Assign(rhs, var, loc)
144        self._typemap.setdefault(var.name, typ)
145        self._lowerer.lower_inst(assign)
146        return var
147
148    def call(self, callable_node, args, kws=None) -> ir.Expr:
149        """Call a bound callable
150
151        Parameters
152        ----------
153        callable_node : _CallableNode
154            The callee
155        args : Sequence[ir.Var]
156        kws : Mapping[str, ir.Var]
157
158        Returns
159        -------
160        res : ir.Expr
161            The expression node for the return value of the call
162        """
163        if kws is None:
164            kws = {}
165        call = ir.Expr.call(callable_node.func, args, kws, loc=self._loc)
166        self._calltypes[call] = callable_node.sig
167        return call
168
169    def setitem(self, obj, index, val) -> ir.SetItem:
170        """Makes a setitem call
171
172        Parameters
173        ----------
174        obj : ir.Var
175            the object being indexed
176        index : ir.Var
177            the index
178        val : ir.Var
179            the value to be stored
180
181        Returns
182        -------
183        res : ir.SetItem
184        """
185        loc = self._loc
186        tm = self._typemap
187        setitem = ir.SetItem(obj, index, val, loc=loc)
188        self._lowerer.fndesc.calltypes[setitem] = signature(
189            types.none, tm[obj.name], tm[index.name], tm[val.name]
190        )
191        self._lowerer.lower_inst(setitem)
192        return setitem
193
194    def getitem(self, obj, index, typ) -> ir.Expr:
195        """Makes a getitem call
196
197        Parameters
198        ----------
199        obj : ir.Var
200            the object being indexed
201        index : ir.Var
202            the index
203        val : ir.Var
204            the ty
205
206        Returns
207        -------
208        res : ir.Expr
209            the retrieved value
210        """
211        tm = self._typemap
212        getitem = ir.Expr.getitem(obj, index, loc=self._loc)
213        self._lowerer.fndesc.calltypes[getitem] = signature(
214            typ, tm[obj.name], tm[index.name],
215        )
216        return getitem
217