def not_supported(self, expr): # noqa
from symengine.lib.symengine_wrapper import \
PyFunction # pylint: disable=E0611
if isinstance(expr, PyFunction) and \
self.function_name(expr) == "CSE": # pylint: disable=E0611
sympy_expr = expr._sympy_()
return prim.CommonSubexpression(self.rec(expr.args[0]),
sympy_expr.prefix,
sympy_expr.scope)
elif isinstance(expr, symengine.Function) and \
self.function_name(expr) == "CSE":
return prim.CommonSubexpression(self.rec(expr.args[0]))
return SympyLikeToPymbolicMapper.not_supported(self, expr)
def map_substitution(self, expr):
assert isinstance(expr.child, prim.Derivative)
call = expr.child.child
if (isinstance(call.function, prim.Variable)
and call.function.name in ["hankel_1", "bessel_j"]):
function = call.function
order, _ = call.parameters
arg, = expr.values
n_derivs = len(expr.child.variables)
import sympy as sym
# AS (9.1.31)
# http://dlmf.nist.gov/10.6.7
if order >= 0:
order_str = str(order)
else:
order_str = "m"+str(-order)
k = n_derivs
return prim.CommonSubexpression(
2**(-k)*sum(
(-1)**idx*int(sym.binomial(k, idx)) * function(i, arg)
for idx, i in enumerate(range(order-k, order+k+1, 2))),
"d%d_%s_%s" % (n_derivs, function.name, order_str))
else:
return IdentityMapper.map_substitution(self, expr)
def map_call(self, expr):
from loopy.library.reduction import parse_reduction_op
if not isinstance(expr.function, p.Variable):
return IdentityMapper.map_call(self, expr)
name = expr.function.name
if name == "cse":
if len(expr.parameters) in [1, 2]:
if len(expr.parameters) == 2:
if not isinstance(expr.parameters[1], p.Variable):
raise TypeError("second argument to cse() must be a symbol")
tag = expr.parameters[1].name
else:
tag = None
return p.CommonSubexpression(
self.rec(expr.parameters[0]), tag)
else:
raise TypeError("cse takes two arguments")
elif name in ["reduce", "simul_reduce"]:
if len(expr.parameters) >= 3:
operation, inames = expr.parameters[:2]
red_exprs = expr.parameters[2:]
operation = parse_reduction_op(str(operation))
return self._parse_reduction(operation, inames,
tuple(self.rec(red_expr) for red_expr in red_exprs),
allow_simultaneous=(name == "simul_reduce"))
else:
raise TypeError("invalid 'reduce' calling sequence")
elif name == "if":
if len(expr.parameters) == 3:
from pymbolic.primitives import If
return If(*tuple(self.rec(p) for p in expr.parameters))
else:
raise TypeError("if takes three arguments")
else:
# see if 'name' is an existing reduction op
operation = parse_reduction_op(name)
if operation:
# arg_count counts arguments but not inames
if len(expr.parameters) != 1 + operation.arg_count:
raise RuntimeError("invalid invocation of "
"reduction operation '%s': expected %d arguments, "
"got %d instead" % (expr.function.name,
1 + operation.arg_count,
len(expr.parameters)))
inames = expr.parameters[0]
red_exprs = tuple(self.rec(param) for param in expr.parameters[1:])
return self._parse_reduction(operation, inames, red_exprs)
else:
return IdentityMapper.map_call(self, expr)
def map_call(self, expr):
from loopy.library.reduction import parse_reduction_op
if not isinstance(expr.function, p.Variable):
return IdentityMapper.map_call(self, expr)
name = expr.function.name
if name == "cse":
if len(expr.parameters) in [1, 2]:
if len(expr.parameters) == 2:
if not isinstance(expr.parameters[1], p.Variable):
raise TypeError("second argument to cse() must be a symbol")
tag = expr.parameters[1].name
else:
tag = None
return p.CommonSubexpression(
self.rec(expr.parameters[0]), tag)
else:
raise TypeError("cse takes two arguments")
elif name in ["reduce", "simul_reduce"]:
if len(expr.parameters) == 3:
operation, inames, red_expr = expr.parameters
if not isinstance(operation, p.Variable):
raise TypeError("operation argument to reduce() "
"must be a symbol")
operation = parse_reduction_op(operation.name)
return self._parse_reduction(operation, inames, self.rec(red_expr),
allow_simultaneous=(name == "simul_reduce"))
else:
raise TypeError("invalid 'reduce' calling sequence")
elif name == "if":
if len(expr.parameters) == 3:
from pymbolic.primitives import If
return If(*tuple(self.rec(p) for p in expr.parameters))
else:
raise TypeError("if takes three arguments")
else:
# see if 'name' is an existing reduction op
operation = parse_reduction_op(name)
if operation:
if len(expr.parameters) != 2:
raise RuntimeError("invalid invocation of "
"reduction operation '%s'" % expr.function.name)
inames, red_expr = expr.parameters
return self._parse_reduction(operation, inames, self.rec(red_expr))
else:
return IdentityMapper.map_call(self, expr)
def test_flop_counter():
x = prim.Variable("x")
y = prim.Variable("y")
z = prim.Variable("z")
subexpr = prim.CommonSubexpression(3 * (x**2 + y + z))
expr = 3 * subexpr + subexpr
from pymbolic.mapper.flop_counter import FlopCounter, CSEAwareFlopCounter
assert FlopCounter()(expr) == 4 * 2 + 2
assert CSEAwareFlopCounter()(expr) == 4 + 2
def map_CSE(self, expr): # noqa
return prim.CommonSubexpression(
self.rec(expr.args[0]), expr.prefix, expr.scope)
def map_CSE(self, expr):
return prim.CommonSubexpression(self.rec(expr.args[0]), expr.prefix)