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 __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)