Traktor/myenv/Lib/site-packages/sympy/printing/llvmjitcode.py

490 lines
17 KiB
Python
Raw Normal View History

2024-05-23 01:57:24 +02:00
'''
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