def __init__(self, prefix="_expr", max_vectors_in_batch_expr=None):
IdentityMapper.__init__(self)
self.prefix = prefix
self.max_vectors_in_batch_expr = max_vectors_in_batch_expr
self.code = []
self.expr_to_var = {}
self.assigned_names = set()
def __init__(self, prefix="_expr", max_vectors_in_batch_expr=None):
IdentityMapper.__init__(self)
self.prefix = prefix
self.max_vectors_in_batch_expr = max_vectors_in_batch_expr
self.code = []
self.expr_to_var = {}
self.assigned_names = set()
def map_operator_binding(self, expr):
from hedge.optemplate import \
FluxOperatorBase, \
BoundaryPair, OperatorBinding, \
FluxExchangeOperator
if isinstance(expr, OperatorBinding):
if isinstance(expr.op, FluxOperatorBase):
if isinstance(expr.field, BoundaryPair):
# we're only worried about internal fluxes
return IdentityMapper.map_operator_binding(self, expr)
# by now we've narrowed it down to a bound interior flux
def func_on_scalar_or_vector(func, arg_fields):
# No CSE necessary here--the compiler CSE's these
# automatically.
from hedge.tools import is_obj_array, make_obj_array
if is_obj_array(arg_fields):
# arg_fields (as an object array) isn't hashable
# --make it so by turning it into a tuple
arg_fields = tuple(arg_fields)
return make_obj_array([
func(i, arg_fields) for i in range(len(arg_fields))
])
else:
return func(0, (arg_fields, ))
from hedge.mesh import TAG_RANK_BOUNDARY
def exchange_and_cse(rank):
return func_on_scalar_or_vector(
lambda i, args: FluxExchangeOperator(i, rank, args),
expr.field)
from pymbolic.primitives import flattened_sum
return flattened_sum([expr] + [
OperatorBinding(
expr.op,
BoundaryPair(expr.field, exchange_and_cse(rank),
TAG_RANK_BOUNDARY(rank)))
for rank in self.interacting_ranks
])
else:
return IdentityMapper.map_operator_binding(self, expr)
def map_operator_binding(self, expr):
from hedge.optemplate import \
FluxOperatorBase, \
BoundaryPair, OperatorBinding, \
FluxExchangeOperator
if isinstance(expr, OperatorBinding):
if isinstance(expr.op, FluxOperatorBase):
if isinstance(expr.field, BoundaryPair):
# we're only worried about internal fluxes
return IdentityMapper.map_operator_binding(self, expr)
# by now we've narrowed it down to a bound interior flux
def func_on_scalar_or_vector(func, arg_fields):
# No CSE necessary here--the compiler CSE's these
# automatically.
from hedge.tools import is_obj_array, make_obj_array
if is_obj_array(arg_fields):
# arg_fields (as an object array) isn't hashable
# --make it so by turning it into a tuple
arg_fields = tuple(arg_fields)
return make_obj_array([
func(i, arg_fields)
for i in range(len(arg_fields))])
else:
return func(0, (arg_fields,))
from hedge.mesh import TAG_RANK_BOUNDARY
def exchange_and_cse(rank):
return func_on_scalar_or_vector(
lambda i, args: FluxExchangeOperator(i, rank, args),
expr.field)
from pymbolic.primitives import flattened_sum
return flattened_sum([expr]
+ [OperatorBinding(expr.op, BoundaryPair(
expr.field,
exchange_and_cse(rank),
TAG_RANK_BOUNDARY(rank)))
for rank in self.interacting_ranks])
else:
return IdentityMapper.map_operator_binding(self, expr)
def map_call(self, expr):
from hedge.optemplate.primitives import CFunction
if isinstance(expr.function, CFunction):
return IdentityMapper.map_call(self, expr)
else:
# If it's not a C-level function, it shouldn't get muddled up into
# a vector math expression.
return self.assign_to_new_var(
type(expr)(
expr.function,
[self.assign_to_new_var(self.rec(par))
for par in expr.parameters]))
def map_call(self, expr):
from hedge.optemplate.primitives import CFunction
if isinstance(expr.function, CFunction):
return IdentityMapper.map_call(self, expr)
else:
# If it's not a C-level function, it shouldn't get muddled up into
# a vector math expression.
return self.assign_to_new_var(
type(expr)(expr.function, [
self.assign_to_new_var(self.rec(par))
for par in expr.parameters
]))
def internal_map_flux(self, flux_bind):
from hedge.optemplate import IdentityMapper
return IdentityMapper.map_operator_binding(self, flux_bind)
def internal_map_flux(self, flux_bind):
from hedge.optemplate import IdentityMapper
return IdentityMapper.map_operator_binding(self, flux_bind)
def __call__(self, expr, type_hints={}):
from hedge.optemplate.mappers.type_inference import TypeInferrer
self.typedict = TypeInferrer()(expr, type_hints)
# {{{ flux batching
# Fluxes can be evaluated faster in batches. Here, we find flux
# batches that we can evaluate together.
# For each FluxRecord, find the other fluxes its flux depends on.
flux_queue = self.get_contained_fluxes(expr)
for fr in flux_queue:
fr.dependencies = set(sf.flux_expr
for sf in self.get_contained_fluxes(fr.flux_expr)) \
- set([fr.flux_expr])
# Then figure out batches of fluxes to evaluate
self.flux_batches = []
admissible_deps = set()
while flux_queue:
present_batch = set()
i = 0
while i < len(flux_queue):
fr = flux_queue[i]
if fr.dependencies <= admissible_deps:
present_batch.add(fr)
flux_queue.pop(i)
else:
i += 1
if present_batch:
# bin batched operators by representative operator
batches_by_repr_op = {}
for fr in present_batch:
batches_by_repr_op[fr.repr_op] = \
batches_by_repr_op.get(fr.repr_op, set()) \
| set([fr.flux_expr])
for repr_op, batch in batches_by_repr_op.iteritems():
self.flux_batches.append(
self.FluxBatch(
repr_op=repr_op,
flux_exprs=list(batch)))
admissible_deps |= set(fr.flux_expr for fr in present_batch)
else:
raise RuntimeError("cannot resolve flux evaluation order")
# }}}
# Once flux batching is figured out, we also need to know which
# derivatives are going to be needed, because once the
# rst-derivatives are available, it's best to calculate the
# xyz ones and throw the rst ones out. It's naturally good if
# we can avoid computing (or storing) some of the xyz ones.
# So figure out which XYZ derivatives of what are needed.
self.diff_ops = self.collect_diff_ops(expr)
# Flux exchange also works better when batched.
self.flux_exchange_ops = self.collect_flux_exchange_ops(expr)
# Finally, walk the expression and build the code.
result = IdentityMapper.__call__(self, expr)
# Then, put the toplevel expressions into variables as well.
from hedge.tools import with_object_array_or_scalar
result = with_object_array_or_scalar(self.assign_to_new_var, result)
return Code(self.aggregate_assignments(self.code, result), result)
def __call__(self, expr, type_hints={}):
# Put the result expressions into variables as well.
from hedge.optemplate import make_common_subexpression as cse
expr = cse(expr, "_result")
from hedge.optemplate.mappers.type_inference import TypeInferrer
self.typedict = TypeInferrer()(expr, type_hints)
# {{{ flux batching
# Fluxes can be evaluated faster in batches. Here, we find flux
# batches that we can evaluate together.
# For each FluxRecord, find the other fluxes its flux depends on.
flux_queue = self.get_contained_fluxes(expr)
for fr in flux_queue:
fr.dependencies = set(sf.flux_expr
for sf in self.get_contained_fluxes(fr.flux_expr)) \
- set([fr.flux_expr])
# Then figure out batches of fluxes to evaluate
self.flux_batches = []
admissible_deps = set()
while flux_queue:
present_batch = set()
i = 0
while i < len(flux_queue):
fr = flux_queue[i]
if fr.dependencies <= admissible_deps:
present_batch.add(fr)
flux_queue.pop(i)
else:
i += 1
if present_batch:
# bin batched operators by representative operator
batches_by_repr_op = {}
for fr in present_batch:
batches_by_repr_op[fr.repr_op] = \
batches_by_repr_op.get(fr.repr_op, set()) \
| set([fr.flux_expr])
for repr_op, batch in batches_by_repr_op.iteritems():
self.flux_batches.append(
self.FluxBatch(repr_op=repr_op,
flux_exprs=list(batch)))
admissible_deps |= set(fr.flux_expr for fr in present_batch)
else:
raise RuntimeError("cannot resolve flux evaluation order")
# }}}
# Used for diff batching
self.diff_ops = self.collect_diff_ops(expr)
# Flux exchange also works better when batched.
self.flux_exchange_ops = self.collect_flux_exchange_ops(expr)
# Finally, walk the expression and build the code.
result = IdentityMapper.__call__(self, expr)
return Code(self.aggregate_assignments(self.code, result), result)
def __call__(self, expr, type_hints={}):
# Put the result expressions into variables as well.
from hedge.optemplate import make_common_subexpression as cse
expr = cse(expr, "_result")
from hedge.optemplate.mappers.type_inference import TypeInferrer
self.typedict = TypeInferrer()(expr, type_hints)
# {{{ flux batching
# Fluxes can be evaluated faster in batches. Here, we find flux
# batches that we can evaluate together.
# For each FluxRecord, find the other fluxes its flux depends on.
flux_queue = self.get_contained_fluxes(expr)
for fr in flux_queue:
fr.dependencies = set(sf.flux_expr
for sf in self.get_contained_fluxes(fr.flux_expr)) \
- set([fr.flux_expr])
# Then figure out batches of fluxes to evaluate
self.flux_batches = []
admissible_deps = set()
while flux_queue:
present_batch = set()
i = 0
while i < len(flux_queue):
fr = flux_queue[i]
if fr.dependencies <= admissible_deps:
present_batch.add(fr)
flux_queue.pop(i)
else:
i += 1
if present_batch:
# bin batched operators by representative operator
batches_by_repr_op = {}
for fr in present_batch:
batches_by_repr_op[fr.repr_op] = \
batches_by_repr_op.get(fr.repr_op, set()) \
| set([fr.flux_expr])
for repr_op, batch in batches_by_repr_op.iteritems():
self.flux_batches.append(
self.FluxBatch(
repr_op=repr_op,
flux_exprs=list(batch)))
admissible_deps |= set(fr.flux_expr for fr in present_batch)
else:
raise RuntimeError("cannot resolve flux evaluation order")
# }}}
# Used for diff batching
self.diff_ops = self.collect_diff_ops(expr)
# Flux exchange also works better when batched.
self.flux_exchange_ops = self.collect_flux_exchange_ops(expr)
# Finally, walk the expression and build the code.
result = IdentityMapper.__call__(self, expr)
return Code(self.aggregate_assignments(self.code, result), result)
