def map_subscript(self, expr):
from pymbolic.primitives import CommonSubexpression
if expr.aggregate.name == self.vec_name \
and isinstance(expr.index, int):
return CommonSubexpression(
expr.aggregate.index((expr.index, ) + self.additional_indices))
else:
return IdentityMapper.map_subscript(self, expr)
def map_call(self, expr):
from loopy.library.reduction import parse_reduction_op
from pymbolic.primitives import Variable
if not isinstance(expr.function, Variable):
return IdentityMapper.map_call(self, expr)
name = expr.function.name
if name == "cse":
from pymbolic.primitives import CommonSubexpression
if len(expr.parameters) in [1, 2]:
if len(expr.parameters) == 2:
if not isinstance(expr.parameters[1], Variable):
raise TypeError("second argument to cse() must be a symbol")
tag = expr.parameters[1].name
else:
tag = None
return 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, 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(*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 wrap_intermediate(x):
if len(x) > 1:
from pymbolic.primitives import CommonSubexpression
result = numpy.empty(len(x), dtype=object)
for i, x_i in enumerate(x):
result[i] = CommonSubexpression(x_i)
return result
else:
return x
def __init__(self, child, prefix=None, dtype=None):
CommonSubexpression.__init__(self, child, prefix)
self.dtype = dtype
def __init__(self, child, prefix=None, dtype=None):
CommonSubexpression.__init__(self, child, prefix)
self.dtype = dtype
def map_call(self, expr):
if self.subexpr_histogram.get(expr, 0) > 1:
return CommonSubexpression(expr)
else:
return getattr(IdentityMapper, expr.mapper_method)(self, expr)
def make_penalty_term(power=1):
from pymbolic.primitives import CommonSubexpression
return CommonSubexpression((ElementOrder()**2 / LocalMeshSize())**power,
"penalty")
