def visit_ClosureNode(self, node):
"""
Compile the inner function.
"""
# Compile closure, skip CFA and type inference
node.func_env.qualified_name = self.get_qname(node)
numba.pipeline.run_env(self.env,
node.func_env,
pipeline_name='compile')
translator = node.func_env.translator
# translator.link()
node.lfunc = translator.lfunc
node.lfunc_pointer = translator.lfunc_pointer
if node.need_numba_func:
return self.create_numba_function(node, node.func_env)
else:
func_name = node.func_def.name
self.symtab[func_name] = Variable(name=func_name,
type=node.type,
is_local=True)
# return nodes.LLVMValueRefNode(node.type, node.lfunc)
# TODO: Remove assignment altogether!
# return nodes.NoneNode()
return nodes.ObjectInjectNode(None, type=object_)
def visit_ClosureNode(self, node):
"""
Compile the inner function.
"""
# Compile inner function, skip type inference
order = numba.pipeline.Pipeline.order
order = order[order.index('type_infer') + 1:]
p, result = numba.pipeline.run_pipeline(
self.context,
node.py_func,
node.type_inferred_ast,
node.type.signature,
symtab=node.symtab,
order=order, # skip type inference
)
node.lfunc = p.translator.lfunc
node.lfunc_pointer = p.translator.lfunc_pointer
if node.need_numba_func:
return self.create_numba_function(node, p)
else:
func_name = node.func_def.name
self.symtab[func_name] = Variable(name=func_name,
type=node.type,
is_local=True)
return ast.Pass()
def create_deferred(type_inferer, node, deferred_cls):
"Create a deferred type for an AST node"
variable = Variable(None)
deferred_type = deferred_cls(variable, type_inferer, node)
variable.type = deferred_type
node.variable = variable
return deferred_type
def _int(typesystem, node, x, base, dst_type=int_):
# Resolve int(x) and float(x) to an equivalent cast
if len(node.args) < 2:
return cast(node, dst_type)
node.variable = Variable(dst_type)
return node
def build_wrapper_translation(env, llvm_module=None):
"""
Generate a wrapper function in the given llvm module.
"""
from numba import pipeline
if llvm_module:
wrapper_module = llvm_module
else:
wrapper_module = env.llvm_context.module
# Create wrapper code generator and wrapper AST
func_name = '__numba_wrapper_%s' % env.crnt.func_name
signature = object_(void.pointer(), object_)
symtab = dict(self=Variable(object_, is_local=True),
args=Variable(object_, is_local=True))
func_env = env.crnt.inherit(
func=fake_pyfunc,
name=func_name,
mangled_name=None, # Force FunctionEnvironment.init()
# to generate a new mangled name.
func_signature=signature,
locals={},
symtab=symtab,
refcount_args=False,
llvm_module=wrapper_module)
# Create wrapper LLVM function
func_env.lfunc = pipeline.get_lfunc(env, func_env)
# Build wrapper ast
wrapper_node = build_wrapper_function_ast(env,
wrapper_lfunc=func_env.lfunc,
llvm_module=wrapper_module)
func_env.ast = wrapper_node
# Specialize and compile wrapper
pipeline.run_env(env, func_env, pipeline_name='late_translate')
keep_alive(fake_pyfunc, func_env.lfunc)
return func_env.translator # TODO: Amend callers to eat func_env
def abs_(context, node, x):
import builtinmodule
argtype = get_type(x)
if argtype.is_array and argtype.is_numeric:
# Handle np.abs() on arrays
dtype = builtinmodule.abstype(argtype.dtype)
result_type = argtype.copy(dtype=dtype)
node.variable = Variable(result_type)
return node
return builtinmodule.abs_(context, node, x)
def visit_FunctionDef(self, node):
if self.function_level == 0:
return self._visit_func_children(node)
signature = self._process_decorators(node)
type = numba_types.ClosureType(signature)
self.symtab[node.name] = Variable(type, is_local=True)
closure = nodes.ClosureNode(node, type, self.func)
type.closure = closure
self.ast.closures.append(closure)
return closure
def infer_unary_math_call(context, call_node, arg, default_result_type=double):
"Resolve calls to math functions to llvm.log.f32() etc"
# signature is a generic signature, build a correct one
type = get_type(call_node.args[0])
if type.is_numeric and type.kind < default_result_type.kind:
type = default_result_type
elif type.is_array and type.dtype.is_int:
type = type.copy(dtype=double)
# signature = minitypes.FunctionType(return_type=type, args=[type])
# result = nodes.MathNode(py_func, signature, call_node.args[0])
nodes.annotate(context.env, call_node, is_math=True)
call_node.variable = Variable(type)
return call_node
def abs_(typesystem, node, x):
from . import builtinmodule
argtype = get_type(x)
nodes.annotate(typesystem.env, node, is_math=True)
if argtype.is_array and argtype.dtype.is_numeric:
# Handle np.abs() on arrays
dtype = builtinmodule.abstype(argtype.dtype)
result_type = argtype.add('dtype', dtype)
node.variable = Variable(result_type)
else:
node = builtinmodule.abs_(typesystem, node, x)
return node
def round_(typesystem, node, number, ndigits):
argtype = get_type(number)
if len(node.args) == 1 and argtype.is_int:
# round(myint) -> float(myint)
return nodes.CoercionNode(node.args[0], double)
if argtype.is_float or argtype.is_int:
dst_type = double
else:
dst_type = object_
node.args[0] = nodes.CoercionNode(node.args[0], object_)
node.variable = Variable(dst_type)
return node # nodes.CoercionNode(node, double)
def infer(typesystem, call_node, *args):
"Resolve calls to llvmmath math calls"
# signature is a generic signature, build a correct one
type = reduce(typesystem.promote, map(get_type, call_node.args))
if type.is_numeric and rank(type) < rank(default_result_type):
type = default_result_type
elif type.is_array and type.dtype.is_int:
type = typesystem.array(double, type.ndim)
call_node.args[:] = nodes.CoercionNode.coerce(call_node.args, type)
# TODO: Remove the abuse below
nodes.annotate(typesystem.env, call_node, is_math=True)
call_node.variable = Variable(type)
return call_node
def round_(context, node, number, ndigits):
# is_math = is_math_function(node.args, round)
argtype = get_type(number)
if len(node.args) == 1 and argtype.is_int:
# round(myint) -> float(myint)
return nodes.CoercionNode(node.args[0], double)
if argtype.is_float or argtype.is_int:
dst_type = double
else:
dst_type = object_
node.args[0] = nodes.CoercionNode(node.args[0], object_)
node.variable = Variable(dst_type)
return node # nodes.CoercionNode(node, double)
def update_closures(self, func_def, scope_type, ext_type):
# Patch closures to get the closure scope as the first argument
for closure in func_def.closures:
# closure.scope_type = scope_type
closure.func_def.scope_type = scope_type
closure.ext_type = ext_type
# patch function parameters
param = ast.Name(id=CLOSURE_SCOPE_ARG_NAME, ctx=ast.Param())
param.variable = Variable(scope_type, is_local=True)
param.type = param.variable.type
closure.symtab[CLOSURE_SCOPE_ARG_NAME] = param.variable
closure.func_def.args.args.insert(0, param)
closure.need_closure_scope = True
# patch closure signature
closure.type.signature.args = (
scope_type, ) + closure.type.signature.args
def get_vars_symtab(self):
return dict(
(var.temp_name,
Variable(name=var.temp_name, is_local=True, type=var.type))
for var in self.variables if not var.code)
def unellipsify(node, slices, subscript_node):
"""
Given an array node `node`, process all AST slices and create the
final type:
- process newaxes (None or numpy.newaxis)
- replace Ellipsis with a bunch of ast.Slice objects
- process integer indices
- append any missing slices in trailing dimensions
"""
type = node.variable.type
if not type.is_array:
assert type.is_object
return object_, node
if (len(slices) == 1 and nodes.is_constant_index(slices[0]) and
slices[0].value.pyval is Ellipsis):
# A[...]
return type, node
result = []
seen_ellipsis = False
# Filter out newaxes
newaxes = [newaxis for newaxis in slices if nodes.is_newaxis(newaxis)]
n_indices = len(slices) - len(newaxes)
full_slice = ast.Slice(lower=None, upper=None, step=None)
full_slice.variable = Variable(typesystem.slice_)
ast.copy_location(full_slice, slices[0])
# process ellipses and count integer indices
indices_seen = 0
for slice_node in slices[::-1]:
slice_type = slice_node.variable.type
if slice_type.is_ellipsis:
if seen_ellipsis:
result.append(full_slice)
else:
nslices = type.ndim - n_indices + 1
result.extend([full_slice] * nslices)
seen_ellipsis = True
elif (slice_type.is_slice or slice_type.is_int or
nodes.is_newaxis(slice_node)):
indices_seen += slice_type.is_int
result.append(slice_node)
else:
# TODO: Coerce all object operands to integer indices?
# TODO: (This will break indexing with the Ellipsis object or
# TODO: with slice objects that we couldn't infer)
return object_, nodes.CoercionNode(node, object_)
# Reverse our reversed processed list of slices
result.reverse()
# append any missing slices (e.g. a2d[:]
result_length = len(result) - len(newaxes)
if result_length < type.ndim:
nslices = type.ndim - result_length
result.extend([full_slice] * nslices)
subscript_node.slice = ast.ExtSlice(result)
ast.copy_location(subscript_node.slice, slices[0])
# create the final array type and set it in value.variable
result_dtype = node.variable.type.dtype
result_ndim = node.variable.type.ndim + len(newaxes) - indices_seen
if result_ndim > 0:
result_type = result_dtype[(slice(None),) * result_ndim]
elif result_ndim == 0:
result_type = result_dtype
else:
result_type = object_
return result_type, node
def typednode(node, type):
"Set a type and simple typed variable on a node"
node.variable = Variable(type)
node.type = type
return node
def range_(typesystem, node, start, stop, step):
node.variable = Variable(typesystem.range_)
node.args = nodes.CoercionNode.coerce(node.args, dst_type=Py_ssize_t)
return node
def abs_(typesystem, node, x):
node.variable = Variable(abstype(get_type(x)))
return node
def abs_(context, node, x):
node.variable = Variable(abstype(get_type(x)))
return node
def abs_(typesystem, node, x):
argtype = typesystem.promote(long_, get_type(x))
dst_type = abstype(argtype)
node.variable = Variable(dst_type)
node.args = [nodes.CoercionNode(x, argtype)]
return node
def pow_(context, call_node, node, power, mod=None):
dst_type = binop_type(context, node, power)
call_node.variable = Variable(dst_type)
return call_node