diff --git a/examples/wave/wave-min.py b/examples/wave/wave-min.py index 59cdbdad5f070691b2d84ea00fa65765a7c32c63..6a4ae710a081c6d5bea2ab17fb054376ed356cdc 100644 --- a/examples/wave/wave-min.py +++ b/examples/wave/wave-min.py @@ -74,7 +74,8 @@ def main(write_output=True, order=4): # print(sym.pretty(op.sym_operator())) bound_op = bind(discr, op.sym_operator()) - # print(bound_op.code) + print(bound_op) + # 1/0 def rhs(t, w): return bound_op(queue, t=t, w=w) @@ -93,18 +94,23 @@ def main(write_output=True, order=4): norm = bind(discr, sym.norm(2, sym.var("u"))) + from time import time + t_last_step = time() + for event in dt_stepper.run(t_end=final_t): if isinstance(event, dt_stepper.StateComputed): assert event.component_id == "w" step += 1 - print(step, event.t, norm(queue, u=event.state_component[0])) + print(step, event.t, norm(queue, u=event.state_component[0]), + time()-t_last_step) vis.write_vtk_file("fld-%04d.vtu" % step, [ ("u", event.state_component[0]), ("v", event.state_component[1:]), ]) + t_last_step = time() if __name__ == "__main__": diff --git a/grudge/discretization.py b/grudge/discretization.py index e9e123e2c1ad017c9aaa988f097632018dded9a2..3231676e7716d7181012ca6465e3aea769382532 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -68,6 +68,16 @@ class Discretization(object): self.order = order self.quad_min_degrees = quad_min_degrees + # {{{ management of discretization-scoped common subexpressions + + from pytools import UniqueNameGenerator + self._discr_scoped_name_gen = UniqueNameGenerator() + + self._discr_scoped_subexpr_to_name = {} + self._discr_scoped_subexpr_name_to_value = {} + + # }}} + # {{{ boundary @memoize_method diff --git a/grudge/execution.py b/grudge/execution.py index 0b5b21e545a6ceec57f353ab57792e7233cdab11..c1280a2904f59fc2ed9c22cf7639abfd3a73498e 100644 --- a/grudge/execution.py +++ b/grudge/execution.py @@ -310,6 +310,19 @@ class ExecutionMapper(mappers.Evaluator, return [(name, self.rec(expr)) for name, expr in zip(insn.names, insn.exprs)], [] + def exec_assign_to_discr_scoped(self, insn): + assignments = [] + for name, expr in zip(insn.names, insn.exprs): + value = self.rec(expr) + self.discr._discr_scoped_subexpr_name_to_value[name] = value + assignments.append((name, value)) + + return assignments, [] + + def exec_assign_from_discr_scoped(self, insn): + return [(insn.name, + self.discr._discr_scoped_subexpr_name_to_value[insn.name])], [] + def exec_vector_expr_assign(self, insn): if self.bound_op.instrumented: def stats_callback(n, vec_expr): @@ -355,13 +368,26 @@ class ExecutionMapper(mappers.Evaluator, # {{{ bound operator class BoundOperator(object): - def __init__(self, discr, code, debug_flags, allocator=None): + def __init__(self, discr, discr_code, eval_code, debug_flags, allocator=None): self.discr = discr - self.code = code + self.discr_code = discr_code + self.eval_code = eval_code self.operator_data_cache = {} self.debug_flags = debug_flags self.allocator = allocator + def __str__(self): + sep = 75 * "=" + "\n" + return ( + sep + + "DISCRETIZATION-SCOPE CODE\n" + + sep + + str(self.discr_code) + "\n" + + sep + + "PER-EVALUATION CODE\n" + + sep + + str(self.eval_code)) + def __call__(self, queue, **context): import pyopencl.array as cl_array @@ -373,11 +399,21 @@ class BoundOperator(object): from pytools.obj_array import with_object_array_or_scalar + # {{{ discr-scope evaluation + + if any(result_var.name not in self.discr._discr_scoped_subexpr_name_to_value + for result_var in self.discr_code.result): + # need to do discr-scope evaluation + discr_eval_context = {} + self.discr_code.execute(ExecutionMapper(queue, discr_eval_context, self)) + + # }}} + new_context = {} for name, var in six.iteritems(context): new_context[name] = with_object_array_or_scalar(replace_queue, var) - return self.code.execute(ExecutionMapper(queue, new_context, self)) + return self.eval_code.execute(ExecutionMapper(queue, new_context, self)) # }}} @@ -474,15 +510,15 @@ def bind(discr, sym_operator, post_bind_mapper=lambda x: x, mesh=discr.mesh) from grudge.symbolic.compiler import OperatorCompiler - code = OperatorCompiler()(sym_operator) + discr_code, eval_code = OperatorCompiler(discr)(sym_operator) - bound_op = BoundOperator(discr, code, + bound_op = BoundOperator(discr, discr_code, eval_code, debug_flags=debug_flags, allocator=allocator) if "dump_op_code" in debug_flags: from grudge.tools import open_unique_debug_file open_unique_debug_file("op-code", ".txt").write( - str(code)) + str(bound_op)) if "dump_dataflow_graph" in debug_flags: bound_op.code.dump_dataflow_graph() diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py index 37dba74606cc7efbf441aa766af8ff59bd39a83a..47822baa495040c9d1ce99798770ea7c83b45908 100644 --- a/grudge/symbolic/compiler.py +++ b/grudge/symbolic/compiler.py @@ -30,6 +30,7 @@ from six.moves import zip, reduce from pytools import Record, memoize_method, memoize from grudge import sym import grudge.symbolic.mappers as mappers +from pymbolic.primitives import Variable, Subscript # {{{ instructions @@ -82,6 +83,7 @@ class Assign(Instruction): """ comment = "" + scope_indicator = "" def __init__(self, names, exprs, **kwargs): Instruction.__init__(self, names=names, exprs=exprs, **kwargs) @@ -119,7 +121,9 @@ class Assign(Instruction): if comment: comment = "/* %s */ " % comment - return "%s <- %s%s" % (self.names[0], comment, self.exprs[0]) + return "%s <-%s %s%s" % ( + self.names[0], self.scope_indicator, comment, + self.exprs[0]) else: if comment: comment = " /* %s */" % comment @@ -132,7 +136,8 @@ class Assign(Instruction): else: dnr_indicator = "" - lines.append(" %s <%s- %s" % (n, dnr_indicator, e)) + lines.append(" %s <%s-%s %s" % ( + n, dnr_indicator, self.scope_indicator, e)) lines.append("}") return "\n".join(lines) @@ -140,6 +145,32 @@ class Assign(Instruction): return exec_mapper.exec_assign +class ToDiscretizationScopedAssign(Assign): + scope_indicator = "(to discr)-" + + def get_execution_method(self, exec_mapper): + return exec_mapper.exec_assign_to_discr_scoped + + +class FromDiscretizationScopedAssign(Assign): + scope_indicator = "(discr)-" + + def __init__(self, name, **kwargs): + Instruction.__init__(self, name=name, **kwargs) + + def get_assignees(self): + return frozenset([self.name]) + + def get_dependencies(self): + return frozenset() + + def __str__(self): + return "%s <-(from discr)" % self.name + + def get_execution_method(self, exec_mapper): + return exec_mapper.exec_assign_from_discr_scoped + + class DiffBatchAssign(Instruction): """ :ivar names: @@ -155,7 +186,7 @@ class DiffBatchAssign(Instruction): """ def get_assignees(self): - return set(self.names) + return frozenset(self.names) @memoize_method def get_dependencies(self): @@ -334,6 +365,41 @@ class Code(object): .write(dot_dataflow_graph(self, max_node_label_length=None)) def __str__(self): + var_to_writer = dict( + (var_name, insn) + for insn in self.instructions + for var_name in insn.get_assignees()) + + if 1: + # {{{ topological sort + + added_insns = set() + ordered_insns = [] + + def insert_insn(insn): + if insn in added_insns: + return + + for dep in insn.get_dependencies(): + try: + writer = var_to_writer[dep.name] + except KeyError: + # input variables won't be found + pass + else: + insert_insn(writer) + + ordered_insns.append(insn) + added_insns.add(insn) + + for insn in self.instructions: + insert_insn(insn) + + assert len(ordered_insns) == len(self.instructions) + assert len(added_insns) == len(self.instructions) + + # }}} + lines = [] for insn in self.instructions: lines.extend(str(insn).split("\n")) @@ -377,7 +443,6 @@ class Code(object): # not consist of just variables. for var in dm(result_expr): - from pymbolic.primitives import Variable assert isinstance(var, Variable) discardable_vars.discard(var.name) @@ -538,18 +603,39 @@ class Code(object): # {{{ compiler +class CodeGenerationState(Record): + """ + .. attribute:: generating_discr_code + """ + + def get_code_list(self, compiler): + if self.generating_discr_code: + return compiler.discr_code + else: + return compiler.eval_code + + class OperatorCompiler(mappers.IdentityMapper): - def __init__(self, prefix="_expr", max_vectors_in_batch_expr=None): + def __init__(self, discr, prefix="_expr", max_vectors_in_batch_expr=None): super(OperatorCompiler, self).__init__() self.prefix = prefix self.max_vectors_in_batch_expr = max_vectors_in_batch_expr - self.code = [] + self.discr_code = [] + self.discr_scope_names_created = set() + self.discr_scope_names_copied_to_eval = set() + + self.eval_code = [] self.expr_to_var = {} self.assigned_names = set() + self.discr = discr + + from pytools import UniqueNameGenerator + self.name_gen = UniqueNameGenerator() + # {{{ collect various optemplate components def collect_diff_ops(self, expr): @@ -569,56 +655,35 @@ class OperatorCompiler(mappers.IdentityMapper): # Used for diff batching self.diff_ops = self.collect_diff_ops(expr) + codegen_state = CodeGenerationState(generating_discr_code=False) # Finally, walk the expression and build the code. - result = super(OperatorCompiler, self).__call__(expr) + result = super(OperatorCompiler, self).__call__(expr, codegen_state) # FIXME: Reenable assignment aggregation - #return Code(self.aggregate_assignments(self.code, result), result) - return Code(self.code, result) + #return Code(self.aggregate_assignments(self.eval_code, result), result) + + from pytools.obj_array import make_obj_array + return ( + Code(self.discr_code, + make_obj_array( + [Variable(name) + for name in self.discr_scope_names_copied_to_eval])), + Code(self.eval_code, result)) # }}} # {{{ variables and names - def get_var_name(self, prefix=None): - def generate_suffixes(): - yield "" - i = 2 - while True: - yield "_%d" % i - i += 1 - - def generate_plain_names(): - i = 0 - while True: - yield self.prefix + str(i) - i += 1 - - if prefix is None: - for name in generate_plain_names(): - if name not in self.assigned_names: - break - else: - for suffix in generate_suffixes(): - name = prefix + suffix - if name not in self.assigned_names: - break - - self.assigned_names.add(name) - return name - - def assign_to_new_var(self, expr, priority=0, prefix=None): - from pymbolic.primitives import Variable, Subscript - + def assign_to_new_var(self, codegen_state, expr, priority=0, prefix=None): # Observe that the only things that can be legally subscripted in # grudge are variables. All other expressions are broken down into # their scalar components. if isinstance(expr, (Variable, Subscript)): return expr - new_name = self.get_var_name(prefix) - self.code.append(self.make_assign( - new_name, expr, priority)) + new_name = self.name_gen(prefix if prefix is not None else "expr") + codegen_state.get_code_list(self).append(Assign( + (new_name,), (expr,), priority=priority)) return Variable(new_name) @@ -626,72 +691,104 @@ class OperatorCompiler(mappers.IdentityMapper): # {{{ map_xxx routines - def map_common_subexpression(self, expr): - try: - return self.expr_to_var[expr.child] - except KeyError: - priority = getattr(expr, "priority", 0) - + def map_common_subexpression(self, expr, codegen_state): + def get_rec_child(codegen_state): if isinstance(expr.child, sym.OperatorBinding): # We need to catch operator bindings here and # treat them specially. They get assigned to their # own variable by default, which would mean the # CSE prefix would be omitted, making the resulting # code less readable. - rec_child = self.map_operator_binding( - expr.child, name_hint=expr.prefix) + return self.map_operator_binding( + expr.child, codegen_state, name_hint=expr.prefix) else: - rec_child = self.rec(expr.child) + return self.rec(expr.child, codegen_state) - cse_var = self.assign_to_new_var(rec_child, - priority=priority, prefix=expr.prefix) + if expr.scope == sym.cse_scope.DISCRETIZATION: + from pymbolic import var + try: + expr_name = self.discr._discr_scoped_subexpr_to_name[expr.child] + except KeyError: + expr_name = "discr." + self.discr._discr_scoped_name_gen( + expr.prefix if expr.prefix is not None else "expr") + self.discr._discr_scoped_subexpr_to_name[expr.child] = expr_name + + assert expr_name.startswith("discr.") + + priority = getattr(expr, "priority", 0) - self.expr_to_var[expr.child] = cse_var - return cse_var + if expr_name not in self.discr_scope_names_created: + new_codegen_state = codegen_state.copy(generating_discr_code=True) + rec_child = get_rec_child(new_codegen_state) + + new_codegen_state.get_code_list(self).append( + ToDiscretizationScopedAssign( + (expr_name,), (rec_child,), priority=priority)) + + self.discr_scope_names_created.add(expr_name) + + if codegen_state.generating_discr_code: + return var(expr_name) + else: + if expr_name in self.discr_scope_names_copied_to_eval: + return var(expr_name) - def map_operator_binding(self, expr, name_hint=None): + self.eval_code.append( + FromDiscretizationScopedAssign( + expr_name, priority=priority)) + + self.discr_scope_names_copied_to_eval.add(expr_name) + + return var(expr_name) + + else: + try: + return self.expr_to_var[expr.child] + except KeyError: + priority = getattr(expr, "priority", 0) + + rec_child = get_rec_child(codegen_state) + + cse_var = self.assign_to_new_var( + codegen_state, rec_child, + priority=priority, prefix=expr.prefix) + + self.expr_to_var[expr.child] = cse_var + return cse_var + + def map_operator_binding(self, expr, codegen_state, name_hint=None): if isinstance(expr.op, sym.RefDiffOperatorBase): - return self.map_ref_diff_op_binding(expr) + return self.map_ref_diff_op_binding(expr, codegen_state) else: # make sure operator assignments stand alone and don't get muddled # up in vector math field_var = self.assign_to_new_var( - self.rec(expr.field)) + codegen_state, + self.rec(expr.field, codegen_state)) result_var = self.assign_to_new_var( + codegen_state, expr.op(field_var), prefix=name_hint) return result_var - def map_ones(self, expr): - # make sure expression stands alone and doesn't get - # muddled up in vector math - return self.assign_to_new_var(expr, prefix="ones") - - def map_node_coordinate_component(self, expr): - # make sure expression stands alone and doesn't get - # muddled up in vector math - return self.assign_to_new_var(expr, prefix="nodes%d" % expr.axis) - - def map_normal_component(self, expr): - # make sure expression stands alone and doesn't get - # muddled up in vector math - return self.assign_to_new_var(expr, prefix="normal%d" % expr.axis) - - def map_call(self, expr): + def map_call(self, expr, codegen_state): from grudge.symbolic.primitives import CFunction if isinstance(expr.function, CFunction): - return super(OperatorCompiler, self).map_call(expr) + return super(OperatorCompiler, self).map_call(expr, codegen_state) 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( + codegen_state, type(expr)( expr.function, - [self.assign_to_new_var(self.rec(par)) + [self.assign_to_new_var( + codegen_state, + self.rec(par, codegen_state)) for par in expr.parameters])) - def map_ref_diff_op_binding(self, expr): + def map_ref_diff_op_binding(self, expr, codegen_state): try: return self.expr_to_var[expr] except KeyError: @@ -700,18 +797,19 @@ class OperatorCompiler(mappers.IdentityMapper): if diff.op.equal_except_for_axis(expr.op) and diff.field == expr.field] - names = [self.get_var_name() for d in all_diffs] + names = [self.name_gen("expr") for d in all_diffs] from pytools import single_valued op_class = single_valued(type(d.op) for d in all_diffs) - self.code.append( + codegen_state.get_code_list(self).append( DiffBatchAssign( names=names, op_class=op_class, operators=[d.op for d in all_diffs], field=self.rec( - single_valued(d.field for d in all_diffs)))) + single_valued(d.field for d in all_diffs), + codegen_state))) from pymbolic import var for n, d in zip(names, all_diffs): @@ -721,14 +819,6 @@ class OperatorCompiler(mappers.IdentityMapper): # }}} - # {{{ instruction producers - - def make_assign(self, name, expr, priority): - return Assign(names=[name], exprs=[expr], - priority=priority) - - # }}} - # {{{ assignment aggregration pass def aggregate_assignments(self, instructions, result): diff --git a/requirements.txt b/requirements.txt index d78f6a5a8fc6d80875aea61cdb076657a3234e5d..acd0138e7520c0973f0cec23992a9823df51e2ec 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ numpy +git+https://github.com/inducer/pymbolic.git git+https://github.com/inducer/islpy.git git+https://github.com/pyopencl/pyopencl.git git+https://github.com/inducer/loopy.git