490 lines
17 KiB
Python
490 lines
17 KiB
Python
'''
|
|
Use llvmlite to create executable functions from SymPy expressions
|
|
|
|
This module requires llvmlite (https://github.com/numba/llvmlite).
|
|
'''
|
|
|
|
import ctypes
|
|
|
|
from sympy.external import import_module
|
|
from sympy.printing.printer import Printer
|
|
from sympy.core.singleton import S
|
|
from sympy.tensor.indexed import IndexedBase
|
|
from sympy.utilities.decorator import doctest_depends_on
|
|
|
|
llvmlite = import_module('llvmlite')
|
|
if llvmlite:
|
|
ll = import_module('llvmlite.ir').ir
|
|
llvm = import_module('llvmlite.binding').binding
|
|
llvm.initialize()
|
|
llvm.initialize_native_target()
|
|
llvm.initialize_native_asmprinter()
|
|
|
|
|
|
__doctest_requires__ = {('llvm_callable'): ['llvmlite']}
|
|
|
|
|
|
class LLVMJitPrinter(Printer):
|
|
'''Convert expressions to LLVM IR'''
|
|
def __init__(self, module, builder, fn, *args, **kwargs):
|
|
self.func_arg_map = kwargs.pop("func_arg_map", {})
|
|
if not llvmlite:
|
|
raise ImportError("llvmlite is required for LLVMJITPrinter")
|
|
super().__init__(*args, **kwargs)
|
|
self.fp_type = ll.DoubleType()
|
|
self.module = module
|
|
self.builder = builder
|
|
self.fn = fn
|
|
self.ext_fn = {} # keep track of wrappers to external functions
|
|
self.tmp_var = {}
|
|
|
|
def _add_tmp_var(self, name, value):
|
|
self.tmp_var[name] = value
|
|
|
|
def _print_Number(self, n):
|
|
return ll.Constant(self.fp_type, float(n))
|
|
|
|
def _print_Integer(self, expr):
|
|
return ll.Constant(self.fp_type, float(expr.p))
|
|
|
|
def _print_Symbol(self, s):
|
|
val = self.tmp_var.get(s)
|
|
if not val:
|
|
# look up parameter with name s
|
|
val = self.func_arg_map.get(s)
|
|
if not val:
|
|
raise LookupError("Symbol not found: %s" % s)
|
|
return val
|
|
|
|
def _print_Pow(self, expr):
|
|
base0 = self._print(expr.base)
|
|
if expr.exp == S.NegativeOne:
|
|
return self.builder.fdiv(ll.Constant(self.fp_type, 1.0), base0)
|
|
if expr.exp == S.Half:
|
|
fn = self.ext_fn.get("sqrt")
|
|
if not fn:
|
|
fn_type = ll.FunctionType(self.fp_type, [self.fp_type])
|
|
fn = ll.Function(self.module, fn_type, "sqrt")
|
|
self.ext_fn["sqrt"] = fn
|
|
return self.builder.call(fn, [base0], "sqrt")
|
|
if expr.exp == 2:
|
|
return self.builder.fmul(base0, base0)
|
|
|
|
exp0 = self._print(expr.exp)
|
|
fn = self.ext_fn.get("pow")
|
|
if not fn:
|
|
fn_type = ll.FunctionType(self.fp_type, [self.fp_type, self.fp_type])
|
|
fn = ll.Function(self.module, fn_type, "pow")
|
|
self.ext_fn["pow"] = fn
|
|
return self.builder.call(fn, [base0, exp0], "pow")
|
|
|
|
def _print_Mul(self, expr):
|
|
nodes = [self._print(a) for a in expr.args]
|
|
e = nodes[0]
|
|
for node in nodes[1:]:
|
|
e = self.builder.fmul(e, node)
|
|
return e
|
|
|
|
def _print_Add(self, expr):
|
|
nodes = [self._print(a) for a in expr.args]
|
|
e = nodes[0]
|
|
for node in nodes[1:]:
|
|
e = self.builder.fadd(e, node)
|
|
return e
|
|
|
|
# TODO - assumes all called functions take one double precision argument.
|
|
# Should have a list of math library functions to validate this.
|
|
def _print_Function(self, expr):
|
|
name = expr.func.__name__
|
|
e0 = self._print(expr.args[0])
|
|
fn = self.ext_fn.get(name)
|
|
if not fn:
|
|
fn_type = ll.FunctionType(self.fp_type, [self.fp_type])
|
|
fn = ll.Function(self.module, fn_type, name)
|
|
self.ext_fn[name] = fn
|
|
return self.builder.call(fn, [e0], name)
|
|
|
|
def emptyPrinter(self, expr):
|
|
raise TypeError("Unsupported type for LLVM JIT conversion: %s"
|
|
% type(expr))
|
|
|
|
|
|
# Used when parameters are passed by array. Often used in callbacks to
|
|
# handle a variable number of parameters.
|
|
class LLVMJitCallbackPrinter(LLVMJitPrinter):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def _print_Indexed(self, expr):
|
|
array, idx = self.func_arg_map[expr.base]
|
|
offset = int(expr.indices[0].evalf())
|
|
array_ptr = self.builder.gep(array, [ll.Constant(ll.IntType(32), offset)])
|
|
fp_array_ptr = self.builder.bitcast(array_ptr, ll.PointerType(self.fp_type))
|
|
value = self.builder.load(fp_array_ptr)
|
|
return value
|
|
|
|
def _print_Symbol(self, s):
|
|
val = self.tmp_var.get(s)
|
|
if val:
|
|
return val
|
|
|
|
array, idx = self.func_arg_map.get(s, [None, 0])
|
|
if not array:
|
|
raise LookupError("Symbol not found: %s" % s)
|
|
array_ptr = self.builder.gep(array, [ll.Constant(ll.IntType(32), idx)])
|
|
fp_array_ptr = self.builder.bitcast(array_ptr,
|
|
ll.PointerType(self.fp_type))
|
|
value = self.builder.load(fp_array_ptr)
|
|
return value
|
|
|
|
|
|
# ensure lifetime of the execution engine persists (else call to compiled
|
|
# function will seg fault)
|
|
exe_engines = []
|
|
|
|
# ensure names for generated functions are unique
|
|
link_names = set()
|
|
current_link_suffix = 0
|
|
|
|
|
|
class LLVMJitCode:
|
|
def __init__(self, signature):
|
|
self.signature = signature
|
|
self.fp_type = ll.DoubleType()
|
|
self.module = ll.Module('mod1')
|
|
self.fn = None
|
|
self.llvm_arg_types = []
|
|
self.llvm_ret_type = self.fp_type
|
|
self.param_dict = {} # map symbol name to LLVM function argument
|
|
self.link_name = ''
|
|
|
|
def _from_ctype(self, ctype):
|
|
if ctype == ctypes.c_int:
|
|
return ll.IntType(32)
|
|
if ctype == ctypes.c_double:
|
|
return self.fp_type
|
|
if ctype == ctypes.POINTER(ctypes.c_double):
|
|
return ll.PointerType(self.fp_type)
|
|
if ctype == ctypes.c_void_p:
|
|
return ll.PointerType(ll.IntType(32))
|
|
if ctype == ctypes.py_object:
|
|
return ll.PointerType(ll.IntType(32))
|
|
|
|
print("Unhandled ctype = %s" % str(ctype))
|
|
|
|
def _create_args(self, func_args):
|
|
"""Create types for function arguments"""
|
|
self.llvm_ret_type = self._from_ctype(self.signature.ret_type)
|
|
self.llvm_arg_types = \
|
|
[self._from_ctype(a) for a in self.signature.arg_ctypes]
|
|
|
|
def _create_function_base(self):
|
|
"""Create function with name and type signature"""
|
|
global link_names, current_link_suffix
|
|
default_link_name = 'jit_func'
|
|
current_link_suffix += 1
|
|
self.link_name = default_link_name + str(current_link_suffix)
|
|
link_names.add(self.link_name)
|
|
|
|
fn_type = ll.FunctionType(self.llvm_ret_type, self.llvm_arg_types)
|
|
self.fn = ll.Function(self.module, fn_type, name=self.link_name)
|
|
|
|
def _create_param_dict(self, func_args):
|
|
"""Mapping of symbolic values to function arguments"""
|
|
for i, a in enumerate(func_args):
|
|
self.fn.args[i].name = str(a)
|
|
self.param_dict[a] = self.fn.args[i]
|
|
|
|
def _create_function(self, expr):
|
|
"""Create function body and return LLVM IR"""
|
|
bb_entry = self.fn.append_basic_block('entry')
|
|
builder = ll.IRBuilder(bb_entry)
|
|
|
|
lj = LLVMJitPrinter(self.module, builder, self.fn,
|
|
func_arg_map=self.param_dict)
|
|
|
|
ret = self._convert_expr(lj, expr)
|
|
lj.builder.ret(self._wrap_return(lj, ret))
|
|
|
|
strmod = str(self.module)
|
|
return strmod
|
|
|
|
def _wrap_return(self, lj, vals):
|
|
# Return a single double if there is one return value,
|
|
# else return a tuple of doubles.
|
|
|
|
# Don't wrap return value in this case
|
|
if self.signature.ret_type == ctypes.c_double:
|
|
return vals[0]
|
|
|
|
# Use this instead of a real PyObject*
|
|
void_ptr = ll.PointerType(ll.IntType(32))
|
|
|
|
# Create a wrapped double: PyObject* PyFloat_FromDouble(double v)
|
|
wrap_type = ll.FunctionType(void_ptr, [self.fp_type])
|
|
wrap_fn = ll.Function(lj.module, wrap_type, "PyFloat_FromDouble")
|
|
|
|
wrapped_vals = [lj.builder.call(wrap_fn, [v]) for v in vals]
|
|
if len(vals) == 1:
|
|
final_val = wrapped_vals[0]
|
|
else:
|
|
# Create a tuple: PyObject* PyTuple_Pack(Py_ssize_t n, ...)
|
|
|
|
# This should be Py_ssize_t
|
|
tuple_arg_types = [ll.IntType(32)]
|
|
|
|
tuple_arg_types.extend([void_ptr]*len(vals))
|
|
tuple_type = ll.FunctionType(void_ptr, tuple_arg_types)
|
|
tuple_fn = ll.Function(lj.module, tuple_type, "PyTuple_Pack")
|
|
|
|
tuple_args = [ll.Constant(ll.IntType(32), len(wrapped_vals))]
|
|
tuple_args.extend(wrapped_vals)
|
|
|
|
final_val = lj.builder.call(tuple_fn, tuple_args)
|
|
|
|
return final_val
|
|
|
|
def _convert_expr(self, lj, expr):
|
|
try:
|
|
# Match CSE return data structure.
|
|
if len(expr) == 2:
|
|
tmp_exprs = expr[0]
|
|
final_exprs = expr[1]
|
|
if len(final_exprs) != 1 and self.signature.ret_type == ctypes.c_double:
|
|
raise NotImplementedError("Return of multiple expressions not supported for this callback")
|
|
for name, e in tmp_exprs:
|
|
val = lj._print(e)
|
|
lj._add_tmp_var(name, val)
|
|
except TypeError:
|
|
final_exprs = [expr]
|
|
|
|
vals = [lj._print(e) for e in final_exprs]
|
|
|
|
return vals
|
|
|
|
def _compile_function(self, strmod):
|
|
global exe_engines
|
|
llmod = llvm.parse_assembly(strmod)
|
|
|
|
pmb = llvm.create_pass_manager_builder()
|
|
pmb.opt_level = 2
|
|
pass_manager = llvm.create_module_pass_manager()
|
|
pmb.populate(pass_manager)
|
|
|
|
pass_manager.run(llmod)
|
|
|
|
target_machine = \
|
|
llvm.Target.from_default_triple().create_target_machine()
|
|
exe_eng = llvm.create_mcjit_compiler(llmod, target_machine)
|
|
exe_eng.finalize_object()
|
|
exe_engines.append(exe_eng)
|
|
|
|
if False:
|
|
print("Assembly")
|
|
print(target_machine.emit_assembly(llmod))
|
|
|
|
fptr = exe_eng.get_function_address(self.link_name)
|
|
|
|
return fptr
|
|
|
|
|
|
class LLVMJitCodeCallback(LLVMJitCode):
|
|
def __init__(self, signature):
|
|
super().__init__(signature)
|
|
|
|
def _create_param_dict(self, func_args):
|
|
for i, a in enumerate(func_args):
|
|
if isinstance(a, IndexedBase):
|
|
self.param_dict[a] = (self.fn.args[i], i)
|
|
self.fn.args[i].name = str(a)
|
|
else:
|
|
self.param_dict[a] = (self.fn.args[self.signature.input_arg],
|
|
i)
|
|
|
|
def _create_function(self, expr):
|
|
"""Create function body and return LLVM IR"""
|
|
bb_entry = self.fn.append_basic_block('entry')
|
|
builder = ll.IRBuilder(bb_entry)
|
|
|
|
lj = LLVMJitCallbackPrinter(self.module, builder, self.fn,
|
|
func_arg_map=self.param_dict)
|
|
|
|
ret = self._convert_expr(lj, expr)
|
|
|
|
if self.signature.ret_arg:
|
|
output_fp_ptr = builder.bitcast(self.fn.args[self.signature.ret_arg],
|
|
ll.PointerType(self.fp_type))
|
|
for i, val in enumerate(ret):
|
|
index = ll.Constant(ll.IntType(32), i)
|
|
output_array_ptr = builder.gep(output_fp_ptr, [index])
|
|
builder.store(val, output_array_ptr)
|
|
builder.ret(ll.Constant(ll.IntType(32), 0)) # return success
|
|
else:
|
|
lj.builder.ret(self._wrap_return(lj, ret))
|
|
|
|
strmod = str(self.module)
|
|
return strmod
|
|
|
|
|
|
class CodeSignature:
|
|
def __init__(self, ret_type):
|
|
self.ret_type = ret_type
|
|
self.arg_ctypes = []
|
|
|
|
# Input argument array element index
|
|
self.input_arg = 0
|
|
|
|
# For the case output value is referenced through a parameter rather
|
|
# than the return value
|
|
self.ret_arg = None
|
|
|
|
|
|
def _llvm_jit_code(args, expr, signature, callback_type):
|
|
"""Create a native code function from a SymPy expression"""
|
|
if callback_type is None:
|
|
jit = LLVMJitCode(signature)
|
|
else:
|
|
jit = LLVMJitCodeCallback(signature)
|
|
|
|
jit._create_args(args)
|
|
jit._create_function_base()
|
|
jit._create_param_dict(args)
|
|
strmod = jit._create_function(expr)
|
|
if False:
|
|
print("LLVM IR")
|
|
print(strmod)
|
|
fptr = jit._compile_function(strmod)
|
|
return fptr
|
|
|
|
|
|
@doctest_depends_on(modules=('llvmlite', 'scipy'))
|
|
def llvm_callable(args, expr, callback_type=None):
|
|
'''Compile function from a SymPy expression
|
|
|
|
Expressions are evaluated using double precision arithmetic.
|
|
Some single argument math functions (exp, sin, cos, etc.) are supported
|
|
in expressions.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
args : List of Symbol
|
|
Arguments to the generated function. Usually the free symbols in
|
|
the expression. Currently each one is assumed to convert to
|
|
a double precision scalar.
|
|
expr : Expr, or (Replacements, Expr) as returned from 'cse'
|
|
Expression to compile.
|
|
callback_type : string
|
|
Create function with signature appropriate to use as a callback.
|
|
Currently supported:
|
|
'scipy.integrate'
|
|
'scipy.integrate.test'
|
|
'cubature'
|
|
|
|
Returns
|
|
=======
|
|
|
|
Compiled function that can evaluate the expression.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> import sympy.printing.llvmjitcode as jit
|
|
>>> from sympy.abc import a
|
|
>>> e = a*a + a + 1
|
|
>>> e1 = jit.llvm_callable([a], e)
|
|
>>> e.subs(a, 1.1) # Evaluate via substitution
|
|
3.31000000000000
|
|
>>> e1(1.1) # Evaluate using JIT-compiled code
|
|
3.3100000000000005
|
|
|
|
|
|
Callbacks for integration functions can be JIT compiled.
|
|
>>> import sympy.printing.llvmjitcode as jit
|
|
>>> from sympy.abc import a
|
|
>>> from sympy import integrate
|
|
>>> from scipy.integrate import quad
|
|
>>> e = a*a
|
|
>>> e1 = jit.llvm_callable([a], e, callback_type='scipy.integrate')
|
|
>>> integrate(e, (a, 0.0, 2.0))
|
|
2.66666666666667
|
|
>>> quad(e1, 0.0, 2.0)[0]
|
|
2.66666666666667
|
|
|
|
The 'cubature' callback is for the Python wrapper around the
|
|
cubature package ( https://github.com/saullocastro/cubature )
|
|
and ( http://ab-initio.mit.edu/wiki/index.php/Cubature )
|
|
|
|
There are two signatures for the SciPy integration callbacks.
|
|
The first ('scipy.integrate') is the function to be passed to the
|
|
integration routine, and will pass the signature checks.
|
|
The second ('scipy.integrate.test') is only useful for directly calling
|
|
the function using ctypes variables. It will not pass the signature checks
|
|
for scipy.integrate.
|
|
|
|
The return value from the cse module can also be compiled. This
|
|
can improve the performance of the compiled function. If multiple
|
|
expressions are given to cse, the compiled function returns a tuple.
|
|
The 'cubature' callback handles multiple expressions (set `fdim`
|
|
to match in the integration call.)
|
|
>>> import sympy.printing.llvmjitcode as jit
|
|
>>> from sympy import cse
|
|
>>> from sympy.abc import x,y
|
|
>>> e1 = x*x + y*y
|
|
>>> e2 = 4*(x*x + y*y) + 8.0
|
|
>>> after_cse = cse([e1,e2])
|
|
>>> after_cse
|
|
([(x0, x**2), (x1, y**2)], [x0 + x1, 4*x0 + 4*x1 + 8.0])
|
|
>>> j1 = jit.llvm_callable([x,y], after_cse)
|
|
>>> j1(1.0, 2.0)
|
|
(5.0, 28.0)
|
|
'''
|
|
|
|
if not llvmlite:
|
|
raise ImportError("llvmlite is required for llvmjitcode")
|
|
|
|
signature = CodeSignature(ctypes.py_object)
|
|
|
|
arg_ctypes = []
|
|
if callback_type is None:
|
|
for _ in args:
|
|
arg_ctype = ctypes.c_double
|
|
arg_ctypes.append(arg_ctype)
|
|
elif callback_type in ('scipy.integrate', 'scipy.integrate.test'):
|
|
signature.ret_type = ctypes.c_double
|
|
arg_ctypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_double)]
|
|
arg_ctypes_formal = [ctypes.c_int, ctypes.c_double]
|
|
signature.input_arg = 1
|
|
elif callback_type == 'cubature':
|
|
arg_ctypes = [ctypes.c_int,
|
|
ctypes.POINTER(ctypes.c_double),
|
|
ctypes.c_void_p,
|
|
ctypes.c_int,
|
|
ctypes.POINTER(ctypes.c_double)
|
|
]
|
|
signature.ret_type = ctypes.c_int
|
|
signature.input_arg = 1
|
|
signature.ret_arg = 4
|
|
else:
|
|
raise ValueError("Unknown callback type: %s" % callback_type)
|
|
|
|
signature.arg_ctypes = arg_ctypes
|
|
|
|
fptr = _llvm_jit_code(args, expr, signature, callback_type)
|
|
|
|
if callback_type and callback_type == 'scipy.integrate':
|
|
arg_ctypes = arg_ctypes_formal
|
|
|
|
# PYFUNCTYPE holds the GIL which is needed to prevent a segfault when
|
|
# calling PyFloat_FromDouble on Python 3.10. Probably it is better to use
|
|
# ctypes.c_double when returning a float rather than using ctypes.py_object
|
|
# and returning a PyFloat from inside the jitted function (i.e. let ctypes
|
|
# handle the conversion from double to PyFloat).
|
|
if signature.ret_type == ctypes.py_object:
|
|
FUNCTYPE = ctypes.PYFUNCTYPE
|
|
else:
|
|
FUNCTYPE = ctypes.CFUNCTYPE
|
|
|
|
cfunc = FUNCTYPE(signature.ret_type, *arg_ctypes)(fptr)
|
|
return cfunc
|