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)