def numba_funcify_Mul(op, node, **kwargs):
signature = create_numba_signature(node, force_scalar=True)
nary_mul_fn = binary_to_nary_func(node.inputs, "mul", "*")
return numba.njit(signature, inline="always")(nary_mul_fn)
def jit_compile_reducer(node, fn, **kwds):
"""Compile Python source for reduction loops using additional optimizations.
Parameters
==========
node
An node from which the signature can be derived.
fn
The Python function object to compile.
kwds
Extra keywords to be added to the :func:`numba.njit` function.
Returns
=======
A :func:`numba.njit`-compiled function.
"""
signature = create_numba_signature(node, reduce_to_scalar=True)
# Eagerly compile the function using increased optimizations. This should
# help improve nested loop reductions.
with use_optimized_cheap_pass():
res = numba_basic.numba_njit(
signature,
boundscheck=False,
fastmath=config.numba__fastmath,
**kwds,
)(fn)
return res
def create_vectorize_func(
scalar_op_fn: Callable,
node: Apply,
use_signature: bool = False,
identity: Optional[Any] = None,
**kwargs,
) -> Callable:
r"""Create a vectorized Numba function from a `Apply`\s Python function."""
if len(node.outputs) > 1:
raise NotImplementedError(
"Multi-output Elemwise Ops are not supported by the Numba backend")
if use_signature:
signature = [create_numba_signature(node, force_scalar=True)]
else:
signature = []
target = (getattr(node.tag, "numba__vectorize_target", None)
or config.numba__vectorize_target)
numba_vectorized_fn = numba_basic.numba_vectorize(
signature,
identity=identity,
target=target,
fastmath=config.numba__fastmath)
py_scalar_func = getattr(scalar_op_fn, "py_func", scalar_op_fn)
elemwise_fn = numba_vectorized_fn(scalar_op_fn)
elemwise_fn.py_scalar_func = py_scalar_func
return elemwise_fn
def create_vectorize_func(op, node, use_signature=False, identity=None, **kwargs):
scalar_op_fn = numba_funcify(op.scalar_op, node=node, inline="always", **kwargs)
if len(node.outputs) > 1:
raise NotImplementedError(
"Multi-output Elemwise Ops are not supported by the Numba backend"
)
if use_signature:
signature = [create_numba_signature(node, force_scalar=True)]
else:
signature = []
target = (
getattr(node.tag, "numba__vectorize_target", None)
or config.numba__vectorize_target
)
numba_vectorized_fn = numba_basic.numba_vectorize(
signature, identity=identity, target=target, fastmath=config.numba__fastmath
)
py_scalar_func = getattr(scalar_op_fn, "py_func", scalar_op_fn)
elemwise_fn = numba_vectorized_fn(scalar_op_fn)
elemwise_fn.py_scalar_func = py_scalar_func
return elemwise_fn
def numba_funcify_Composite(op, node, **kwargs):
signature = create_numba_signature(node, force_scalar=True)
composite_fn = numba_basic.numba_njit(
signature,
fastmath=config.numba__fastmath)(numba_funcify(op.fgraph,
squeeze_output=True,
**kwargs))
return composite_fn
def numba_funcify_Mul(op, node, **kwargs):
signature = create_numba_signature(node, force_scalar=True)
nary_mul_fn = binary_to_nary_func(node.inputs, "mul", "*")
return numba_basic.numba_njit(signature,
inline="always",
fastmath=config.numba__fastmath)(nary_mul_fn)
def numba_funcify_ScalarOp(op, node, **kwargs):
# TODO: Do we need to cache these functions so that we don't end up
# compiling the same Numba function over and over again?
scalar_func_name = op.nfunc_spec[0]
if scalar_func_name.startswith("scipy."):
func_package = scipy
scalar_func_name = scalar_func_name.split(".", 1)[-1]
else:
func_package = np
if "." in scalar_func_name:
scalar_func = reduce(getattr, [scipy] + scalar_func_name.split("."))
else:
scalar_func = getattr(func_package, scalar_func_name)
scalar_op_fn_name = get_name_for_object(scalar_func)
unique_names = unique_name_generator([scalar_op_fn_name, "scalar_func"],
suffix_sep="_")
input_names = ", ".join(
[unique_names(v, force_unique=True) for v in node.inputs])
global_env = {"scalar_func": scalar_func}
scalar_op_src = f"""
def {scalar_op_fn_name}({input_names}):
return scalar_func({input_names})
"""
scalar_op_fn = compile_function_src(scalar_op_src, scalar_op_fn_name, {
**globals(),
**global_env
})
signature = create_numba_signature(node, force_scalar=True)
return numba_basic.numba_njit(
signature, inline="always",
fastmath=config.numba__fastmath)(scalar_op_fn)
def create_vectorize_func(op,
node,
use_signature=False,
identity=None,
**kwargs):
scalar_op_fn = numba_funcify(op.scalar_op,
node=node,
inline="always",
**kwargs)
if len(node.outputs) > 1:
raise NotImplementedError(
"Multi-output Elemwise Ops are not supported by the Numba backend")
if use_signature:
signature = [create_numba_signature(node, force_scalar=True)]
else:
signature = []
numba_vectorize = numba.vectorize(signature, identity=identity)
elemwise_fn = numba_vectorize(scalar_op_fn)
elemwise_fn.py_scalar_func = scalar_op_fn
return elemwise_fn
def numba_funcify_Composite(op, node, **kwargs):
signature = create_numba_signature(node, force_scalar=True)
composite_fn = numba.njit(signature)(numba_funcify(op.fgraph,
squeeze_output=True,
**kwargs))
return composite_fn
def numba_funcify_ScalarOp(op, node, **kwargs):
# TODO: Do we need to cache these functions so that we don't end up
# compiling the same Numba function over and over again?
scalar_func_name = op.nfunc_spec[0]
if scalar_func_name.startswith("scipy."):
func_package = scipy
scalar_func_name = scalar_func_name.split(".", 1)[-1]
else:
func_package = np
if "." in scalar_func_name:
scalar_func = reduce(getattr, [scipy] + scalar_func_name.split("."))
else:
scalar_func = getattr(func_package, scalar_func_name)
scalar_op_fn_name = get_name_for_object(scalar_func)
unique_names = unique_name_generator([scalar_op_fn_name, "scalar_func"],
suffix_sep="_")
global_env = {"scalar_func": scalar_func}
input_tmp_dtypes = None
if func_package == scipy and hasattr(scalar_func, "types"):
# The `numba-scipy` bindings don't provide implementations for all
# inputs types, so we need to convert the inputs to floats and back.
inp_dtype_kinds = tuple(
np.dtype(inp.type.dtype).kind for inp in node.inputs)
accepted_inp_kinds = tuple(
sig_type.split("->")[0] for sig_type in scalar_func.types)
if not any(
all(dk == ik for dk, ik in zip(inp_dtype_kinds, ok_kinds))
for ok_kinds in accepted_inp_kinds):
# They're usually ordered from lower-to-higher precision, so
# we pick the last acceptable input types
#
# XXX: We should pick the first acceptable float/int types in
# reverse, excluding all the incompatible ones (e.g. `"0"`).
# The assumption is that this is only used by `numba-scipy`-exposed
# functions, although it's possible for this to be triggered by
# something else from the `scipy` package
input_tmp_dtypes = tuple(
np.dtype(k) for k in accepted_inp_kinds[-1])
if input_tmp_dtypes is None:
unique_names = unique_name_generator(
[scalar_op_fn_name, "scalar_func"], suffix_sep="_")
input_names = ", ".join(
[unique_names(v, force_unique=True) for v in node.inputs])
scalar_op_src = f"""
def {scalar_op_fn_name}({input_names}):
return scalar_func({input_names})
"""
else:
global_env["direct_cast"] = numba_basic.direct_cast
global_env["output_dtype"] = np.dtype(node.outputs[0].type.dtype)
input_tmp_dtype_names = {
f"inp_tmp_dtype_{i}": i_dtype
for i, i_dtype in enumerate(input_tmp_dtypes)
}
global_env.update(input_tmp_dtype_names)
unique_names = unique_name_generator(
[scalar_op_fn_name, "scalar_func"] + list(global_env.keys()),
suffix_sep="_")
input_names = [unique_names(v, force_unique=True) for v in node.inputs]
converted_call_args = ", ".join([
f"direct_cast({i_name}, {i_tmp_dtype_name})"
for i_name, i_tmp_dtype_name in zip(input_names,
input_tmp_dtype_names.keys())
])
scalar_op_src = f"""
def {scalar_op_fn_name}({', '.join(input_names)}):
return direct_cast(scalar_func({converted_call_args}), output_dtype)
"""
scalar_op_fn = compile_function_src(scalar_op_src, scalar_op_fn_name, {
**globals(),
**global_env
})
signature = create_numba_signature(node, force_scalar=True)
return numba_basic.numba_njit(
signature, inline="always",
fastmath=config.numba__fastmath)(scalar_op_fn)