def make_intrinsic(intrinsic):
"""
Create an intrinsic function given an Intrinsic.
"""
if intrinsic in cache:
return cache[intrinsic]
# NOTE: don't use numba.jit() and 'exec', it will make inspect.getsource()
# NOTE: fail, and hence won't work in python 2.6 (since meta doesn't work
# NOTE: there)
# Build argument names
nargs = len(intrinsic.func_signature.args)
assert nargs < len(string.ascii_letters)
argnames = ", ".join(string.ascii_letters[:nargs])
# Build source code and environment
args = (intrinsic.name, argnames, argnames)
source = ("def %s(%s): return intrinsic(%s)\n" % args)
func_globals = {'intrinsic': intrinsic}
mod_ast = ast.parse(source)
func_ast = mod_ast.body[0]
# Compile
func_env, _ = pipeline.run_pipeline2(
env, func=None, func_ast=func_ast,
func_signature=intrinsic.func_signature,
function_globals=func_globals)
jitted_func = func_env.numba_wrapper_func
# Populate cache
cache[intrinsic] = jitted_func
return jitted_func
def cf_from_source(source, func_globals):
"Render the SSA graph given python source code"
from numba import pipeline
from numba import environment
mod = ast.parse(source)
func_ast = mod.body[0]
env = environment.NumbaEnvironment.get_environment()
func_env, _ = pipeline.run_pipeline2(
env, None, func_ast, void(), pipeline_name="cf", function_globals=dict(func_globals)
)
return func_env.symtab, func_env.flow # func_env.cfg
def cf_from_source(source, func_globals):
"Render the SSA graph given python source code"
from numba import pipeline
from numba import environment
mod = ast.parse(source)
func_ast = mod.body[0]
env = environment.NumbaEnvironment.get_environment()
func_env, _ = pipeline.run_pipeline2(env,
None,
func_ast,
void(),
pipeline_name='cf',
function_globals=dict(func_globals))
return func_env.symtab, func_env.flow #func_env.cfg
def infer(func, signature=functype(), warn=True, **kwargs):
func_ast = functions._get_ast(func)
env = environment.NumbaEnvironment.get_environment(kwargs.get('env', None))
infer_pipe = env.get_or_add_pipeline('infer', construct_infer_pipeline)
kwargs.update(warn=warn, pipeline_name='infer')
pipe, (signature, symtab, func_ast) = pipeline.run_pipeline2(
env, func, func_ast, signature, **kwargs)
last_block = func_ast.flow.blocks[-2]
symbols = {}
#for block in ast.flow.blocks: print block.symtab
for var_name, var in symtab.iteritems():
if not var.parent_var and not var.is_constant:
var = lookup(last_block, var_name)
if var:
symbols[var_name] = var
return signature, symbols
def infer(func, signature=functype(), warn=True, **kwargs):
func_ast = functions._get_ast(func)
env = environment.NumbaEnvironment.get_environment(kwargs.get('env', None))
infer_pipe = env.get_or_add_pipeline('infer', construct_infer_pipeline)
kwargs.update(warn=warn, pipeline_name='infer')
pipe, (signature, symtab, func_ast) = pipeline.run_pipeline2(
env, func, func_ast, signature, **kwargs)
last_block = func_ast.flow.blocks[-2]
symbols = {}
#for block in ast.flow.blocks: print block.symtab
for var_name, var in symtab.iteritems():
if not var.parent_var and not var.is_constant:
var = lookup(last_block, var_name)
if var:
symbols[var_name] = var
return signature, symbols
def register_array_expression(self, node, lhs=None):
super(ArrayExpressionRewriteNative, self).register_array_expression(
node, lhs)
# llvm_module = llvm.core.Module.new(temp_name("array_expression_module"))
# llvm_module = self.env.llvm_context.module
lhs_type = lhs.type if lhs else node.type
is_expr = lhs is None
if node.type.is_array and lhs_type.ndim < node.type.ndim:
# TODO: this is valid in NumPy if the leading dimensions of the
# TODO: RHS have extent 1
raise error.NumbaError(
node, "Right hand side must have a "
"dimensionality <= %d" % lhs_type.ndim)
# Create ufunc scalar kernel
ufunc_ast, signature, ufunc_builder = get_py_ufunc_ast(self.env, lhs, node)
# Compile ufunc scalar kernel with numba
ast.fix_missing_locations(ufunc_ast)
# func_env = self.env.crnt.inherit(
# func=None, ast=ufunc_ast, func_signature=signature,
# wrap=False, #link=False, #llvm_module=llvm_module,
# )
# pipeline.run_env(self.env, func_env) #, pipeline_name='codegen')
func_env, (_, _, _) = pipeline.run_pipeline2(
self.env, None, ufunc_ast, signature,
function_globals=self.env.crnt.function_globals,
wrap=False, link=False, nopython=True,
#llvm_module=llvm_module, # pipeline_name='codegen',
)
llvm_module = func_env.llvm_module
operands = ufunc_builder.operands
operands = [nodes.CloneableNode(operand) for operand in operands]
if lhs is not None:
lhs = nodes.CloneableNode(lhs)
broadcast_operands = [lhs] + operands
lhs = lhs.clone
else:
broadcast_operands = operands[:]
shape = slicenodes.BroadcastNode(lhs_type, broadcast_operands)
operands = [op.clone for op in operands]
if lhs is None and self.nopython:
raise error.NumbaError(
node, "Cannot allocate new memory in nopython context")
elif lhs is None:
# TODO: determine best output order at runtime
shape = shape.cloneable
lhs = nodes.ArrayNewEmptyNode(lhs_type, shape.clone,
lhs_type.is_f_contig).cloneable
# Build minivect wrapper kernel
context = NumbaStaticArgsContext()
context.llvm_module = llvm_module
# context.llvm_ee = self.env.llvm_context.execution_engine
b = context.astbuilder
variables = [b.variable(name_node.type, "op%d" % i)
for i, name_node in enumerate([lhs] + operands)]
miniargs = [b.funcarg(variable) for variable in variables]
body = miniutils.build_kernel_call(func_env.lfunc.name, signature,
miniargs, b)
minikernel = b.function_from_numpy(
temp_name("array_expression"), body, miniargs)
lminikernel, = context.run_simple(minikernel,
specializers.StridedSpecializer)
# lminikernel.linkage = llvm.core.LINKAGE_LINKONCE_ODR
# pipeline.run_env(self.env, func_env, pipeline_name='post_codegen')
# llvm_module.verify()
del func_env
assert lminikernel.module is llvm_module
# print("---------")
# print(llvm_module)
# print("~~~~~~~~~~~~")
lminikernel = self.env.llvm_context.link(lminikernel)
# Build call to minivect kernel
operands.insert(0, lhs)
args = [shape]
scalar_args = []
for operand in operands:
if operand.type.is_array:
data_p = self.array_attr(operand, 'data')
data_p = nodes.CoercionNode(data_p,
operand.type.dtype.pointer())
if not isinstance(operand, nodes.CloneNode):
operand = nodes.CloneNode(operand)
strides_p = self.array_attr(operand, 'strides')
args.extend((data_p, strides_p))
else:
scalar_args.append(operand)
args.extend(scalar_args)
result = nodes.NativeCallNode(minikernel.type, args, lminikernel)
# Use native slicing in array expressions
slicenodes.mark_nopython(ast.Suite(body=result.args))
if not is_expr:
# a[:] = b[:] * c[:]
return result
# b[:] * c[:], return new array as expression
return nodes.ExpressionNode(stmts=[result], expr=lhs.clone)
def register_array_expression(self, node, lhs=None):
super(ArrayExpressionRewriteNative, self).register_array_expression(
node, lhs)
lhs_type = lhs.type if lhs else node.type
is_expr = lhs is None
if node.type.is_array and lhs_type.ndim < node.type.ndim:
# TODO: this is valid in NumPy if the leading dimensions of the
# TODO: RHS have extent 1
raise error.NumbaError(
node, "Right hand side must have a "
"dimensionality <= %d" % lhs_type.ndim)
# Create ufunc scalar kernel
ufunc_ast, signature, ufunc_builder = self.get_py_ufunc_ast(lhs, node)
signature.struct_by_reference = True
# Compile ufunc scalar kernel with numba
ast.fix_missing_locations(ufunc_ast)
func_env, (_, _, _) = pipeline.run_pipeline2(
self.env, None, ufunc_ast, signature,
function_globals={},
)
# Manual linking
lfunc = func_env.lfunc
# print lfunc
operands = ufunc_builder.operands
functions.keep_alive(self.func, lfunc)
operands = [nodes.CloneableNode(operand) for operand in operands]
if lhs is not None:
lhs = nodes.CloneableNode(lhs)
broadcast_operands = [lhs] + operands
lhs = lhs.clone
else:
broadcast_operands = operands[:]
shape = slicenodes.BroadcastNode(lhs_type, broadcast_operands)
operands = [op.clone for op in operands]
if lhs is None and self.nopython:
raise error.NumbaError(
node, "Cannot allocate new memory in nopython context")
elif lhs is None:
# TODO: determine best output order at runtime
shape = shape.cloneable
lhs = nodes.ArrayNewEmptyNode(lhs_type, shape.clone,
lhs_type.is_f_contig).cloneable
# Build minivect wrapper kernel
context = NumbaproStaticArgsContext()
context.llvm_module = self.env.llvm_context.module
# context.debug = True
context.optimize_broadcasting = False
b = context.astbuilder
variables = [b.variable(name_node.type, "op%d" % i)
for i, name_node in enumerate([lhs] + operands)]
miniargs = [b.funcarg(variable) for variable in variables]
body = miniutils.build_kernel_call(lfunc.name, signature, miniargs, b)
minikernel = b.function_from_numpy(
templating.temp_name("array_expression"), body, miniargs)
lminikernel, ctypes_kernel = context.run_simple(
minikernel, specializers.StridedSpecializer)
# Build call to minivect kernel
operands.insert(0, lhs)
args = [shape]
scalar_args = []
for operand in operands:
if operand.type.is_array:
data_p = self.array_attr(operand, 'data')
data_p = nodes.CoercionNode(data_p,
operand.type.dtype.pointer())
if not isinstance(operand, nodes.CloneNode):
operand = nodes.CloneNode(operand)
strides_p = self.array_attr(operand, 'strides')
args.extend((data_p, strides_p))
else:
scalar_args.append(operand)
args.extend(scalar_args)
result = nodes.NativeCallNode(minikernel.type, args, lminikernel)
# Use native slicing in array expressions
slicenodes.mark_nopython(ast.Suite(body=result.args))
if not is_expr:
# a[:] = b[:] * c[:]
return result
# b[:] * c[:], return new array as expression
return nodes.ExpressionNode(stmts=[result], expr=lhs.clone)
