From a50e257a0ddbefa89f14609e088c88b92c542d53 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Apr 2019 19:31:58 -0500 Subject: [PATCH 1/8] Add an expression node for calls to external functions This change adds the ExternalCall node and execution support for it. External functions may be passed as part of the execution context. --- grudge/execution.py | 14 ++++++- grudge/symbolic/compiler.py | 40 +++++++++++++----- grudge/symbolic/dofdesc_inference.py | 3 ++ grudge/symbolic/mappers/__init__.py | 31 ++++++++++++++ grudge/symbolic/primitives.py | 63 ++++++++++++++++++++-------- test/test_grudge.py | 29 +++++++++++++ 6 files changed, 150 insertions(+), 30 deletions(-) diff --git a/grudge/execution.py b/grudge/execution.py index c20aa4bc..9084b823 100644 --- a/grudge/execution.py +++ b/grudge/execution.py @@ -94,13 +94,23 @@ class ExecutionMapper(mappers.Evaluator, value = ary return value - def map_call(self, expr): + def map_external_call(self, expr): from pymbolic.primitives import Variable assert isinstance(expr.function, Variable) + args = [self.rec(p) for p in expr.parameters] - # FIXME: Make a way to register functions + return self.context[expr.function.name](*args) + + def map_call(self, expr): + from pymbolic.primitives import Variable + assert isinstance(expr.function, Variable) args = [self.rec(p) for p in expr.parameters] + + # Function lookup precedence: + # * Numpy functions + # * OpenCL functions + from numbers import Number representative_arg = args[0] if ( diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py index b27d5585..2b25f0f0 100644 --- a/grudge/symbolic/compiler.py +++ b/grudge/symbolic/compiler.py @@ -667,13 +667,14 @@ def aggregate_assignments(inf_mapper, instructions, result, for assignee in insn.get_assignees()) from pytools import partition - from grudge.symbolic.primitives import DTAG_SCALAR + from grudge.symbolic.primitives import DTAG_SCALAR, ExternalCall unprocessed_assigns, other_insns = partition( lambda insn: ( isinstance(insn, Assign) and not isinstance(insn, ToDiscretizationScopedAssign) and not isinstance(insn, FromDiscretizationScopedAssign) + and not isinstance(insn.exprs[0], ExternalCall) and not any( inf_mapper.infer_for_name(n).domain_tag == DTAG_SCALAR for n in insn.names)), @@ -886,6 +887,10 @@ class ToLoopyExpressionMapper(mappers.IdentityMapper): expr, "%s_%d" % (expr.aggregate.name, subscript)) + def map_external_call(self, expr): + raise ValueError( + "Cannot map external call '%s' into loopy" % expr.function) + def map_call(self, expr): if isinstance(expr.function, sym.CFunction): from pymbolic import var @@ -970,8 +975,11 @@ class ToLoopyInstructionMapper(object): self.dd_inference_mapper = dd_inference_mapper def map_insn_assign(self, insn): - from grudge.symbolic.primitives import OperatorBinding - if len(insn.exprs) == 1 and isinstance(insn.exprs[0], OperatorBinding): + from grudge.symbolic.primitives import OperatorBinding, ExternalCall + + if ( + len(insn.exprs) == 1 + and isinstance(insn.exprs[0], (OperatorBinding, ExternalCall))): return insn iname = "grdg_i" @@ -1261,6 +1269,17 @@ class OperatorCompiler(mappers.IdentityMapper): prefix=name_hint) return result_var + def map_external_call(self, expr, codegen_state): + return self.assign_to_new_var( + codegen_state, + type(expr)( + expr.function, + [self.assign_to_new_var( + codegen_state, + self.rec(par, codegen_state)) + for par in expr.parameters], + expr.dd)) + def map_call(self, expr, codegen_state): from grudge.symbolic.primitives import CFunction if isinstance(expr.function, CFunction): @@ -1268,15 +1287,14 @@ class OperatorCompiler(mappers.IdentityMapper): 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( - codegen_state, - self.rec(par, codegen_state)) - for par in expr.parameters])) + codegen_state, + type(expr)( + expr.function, + [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, codegen_state): try: diff --git a/grudge/symbolic/dofdesc_inference.py b/grudge/symbolic/dofdesc_inference.py index 92be126f..96ead088 100644 --- a/grudge/symbolic/dofdesc_inference.py +++ b/grudge/symbolic/dofdesc_inference.py @@ -191,6 +191,9 @@ class DOFDescInferenceMapper(RecursiveMapper, CSECachingMapperMixin): # FIXME return arg_dds[0] + def map_external_call(self, expr): + return expr.dd + # }}} # {{{ instruction mappings diff --git a/grudge/symbolic/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py index d52f7ac5..d1b70475 100644 --- a/grudge/symbolic/mappers/__init__.py +++ b/grudge/symbolic/mappers/__init__.py @@ -215,6 +215,12 @@ class IdentityMapperMixin(LocalOpReducerMixin, FluxOpReducerMixin): map_ones = map_grudge_variable map_node_coordinate_component = map_grudge_variable + def map_external_call(self, expr, *args, **kwargs): + return type(expr)( + self.rec(expr.function, *args, **kwargs), + self.rec(expr.parameters, *args, **kwargs), + dd=expr.dd) + # }}} @@ -268,6 +274,15 @@ class DependencyMapper( map_ones = _map_leaf map_node_coordinate_component = _map_leaf + def map_external_call(self, expr): + result = self.map_call(expr) + if self.include_calls == "descend_args": + # Unlike regular calls, we regard the function as an argument, + # because it's user-supplied (and thus we need to pick it up as a + # dependency). + result = self.combine((result, self.rec(expr.function))) + return result + class FlopCounter( CombineMapperMixin, @@ -281,6 +296,9 @@ class FlopCounter( def map_c_function(self, expr): return 1 + def map_external_call(self, expr): + return 1 + def map_ones(self, expr): return 0 @@ -849,6 +867,15 @@ class StringifyMapper(pymbolic.mapper.stringifier.StringifyMapper): def map_interpolation(self, expr, enclosing_prec): return "Interp" + self._format_op_dd(expr) + def map_external_call(self, expr, enclosing_prec): + from pymbolic.mapper.stringifier import PREC_CALL, PREC_NONE + return ( + self.parenthesize_if_needed( + "External:%s:%s" % ( + self.map_call(expr, PREC_NONE), self._format_dd(expr.dd)), + enclosing_prec, + PREC_CALL)) + class PrettyStringifyMapper( pymbolic.mapper.stringifier.CSESplittingStringifyMapperMixin, @@ -1224,6 +1251,7 @@ class CollectorMixin(OperatorReducerMixin, LocalOpReducerMixin, FluxOpReducerMix def map_constant(self, expr, *args, **kwargs): return set() + map_variable = map_constant map_grudge_variable = map_constant map_c_function = map_grudge_variable @@ -1242,6 +1270,9 @@ class BoundOperatorCollector(CSECachingMapperMixin, CollectorMixin, CombineMappe def __init__(self, op_class): self.op_class = op_class + def map_external_call(self, expr): + return self.map_call(expr) + map_common_subexpression_uncached = \ CombineMapper.map_common_subexpression diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index 231a70df..65676118 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -34,12 +34,13 @@ from meshmode.discretization.connection import ( # noqa from pymbolic.primitives import ( # noqa cse_scope as cse_scope_base, - make_common_subexpression as cse, If, Comparison) + make_common_subexpression as cse, If, Comparison, Expression) from pymbolic.geometric_algebra import MultiVector from pytools.obj_array import join_fields, make_obj_array # noqa -class ExpressionBase(pymbolic.primitives.Expression): +class GrudgeStringifiable(object): + def stringifier(self): from grudge.symbolic.mappers import StringifyMapper return StringifyMapper @@ -74,6 +75,7 @@ Symbols .. autoclass:: Variable .. autoclass:: ScalarVariable +.. autoclass:: ExternalCall .. autoclass:: make_sym_array .. autoclass:: make_sym_mv .. autoclass:: CFunction @@ -289,7 +291,16 @@ class HasDOFDesc(object): discretization on which this property is given. """ - def __init__(self, dd): + def __init__(self, *args, **kwargs): + # The remaining arguments are passed to the chained superclass. + + if "dd" in kwargs: + dd = kwargs.pop("dd") + else: + dd = args[-1] + args = args[:-1] + + super(HasDOFDesc, self).__init__(*args, **kwargs) self.dd = dd def __getinitargs__(self): @@ -298,9 +309,7 @@ class HasDOFDesc(object): def with_dd(self, dd): """Return a copy of *self*, modified to the given DOF descriptor. """ - return type(self)( - *self.__getinitargs__()[:-1], - dd=dd or self.dd) + return type(self)(*self.__getinitargs__()) # }}} @@ -311,7 +320,7 @@ class cse_scope(cse_scope_base): # noqa DISCRETIZATION = "grudge_discretization" -class Variable(HasDOFDesc, ExpressionBase, pymbolic.primitives.Variable): +class Variable(HasDOFDesc, GrudgeStringifiable, pymbolic.primitives.Variable): """A user-supplied input variable with a known :class:`DOFDesc`. """ init_arg_names = ("name", "dd") @@ -320,8 +329,7 @@ class Variable(HasDOFDesc, ExpressionBase, pymbolic.primitives.Variable): if dd is None: dd = DD_VOLUME - HasDOFDesc.__init__(self, dd) - pymbolic.primitives.Variable.__init__(self, name) + super(Variable, self).__init__(name, dd) def __getinitargs__(self): return (self.name, self.dd,) @@ -337,6 +345,30 @@ class ScalarVariable(Variable): super(ScalarVariable, self).__init__(name, DD_SCALAR) +class ExternalCall(HasDOFDesc, GrudgeStringifiable, pymbolic.primitives.Call): + + init_arg_names = ("function", "parameters", "dd") + + def __getinitargs__(self): + return (self.function, self.parameters, self.dd) + + def get_hash(self): + # Object arrays are permitted as parameters, but are not hashable. + params = [] + for param in self.parameters: + if isinstance(param, np.ndarray): + assert param.dtype == np.object + param = tuple(param) + params.append(param) + + return hash( + (type(self).__name__, self.function) + + tuple(params) + + (self.dd,)) + + mapper_method = "map_external_call" + + def make_sym_array(name, shape, dd=None): def var_factory(name): return Variable(name, dd) @@ -349,13 +381,10 @@ def make_sym_mv(name, dim, var_factory=None): make_sym_array(name, dim, var_factory)) -class CFunction(pymbolic.primitives.Variable): +class CFunction(GrudgeStringifiable, pymbolic.primitives.Variable): """A symbol representing a C-level function, to be used as the function argument of :class:`pymbolic.primitives.Call`. """ - def stringifier(self): - from grudge.symbolic.mappers import StringifyMapper - return StringifyMapper def __call__(self, *exprs): from pytools.obj_array import with_object_array_or_scalar_n_args @@ -379,7 +408,7 @@ bessel_y = CFunction("bessel_y") # {{{ technical helpers -class OperatorBinding(ExpressionBase): +class OperatorBinding(GrudgeStringifiable, pymbolic.primitives.Expression): init_arg_names = ("op", "field") def __init__(self, op, field): @@ -424,16 +453,16 @@ class PrioritizedSubexpression(pymbolic.primitives.CommonSubexpression): # }}} -class Ones(ExpressionBase, HasDOFDesc): +class Ones(GrudgeStringifiable, HasDOFDesc, Expression): def __getinitargs__(self): - return () + return (self.dd,) mapper_method = intern("map_ones") # {{{ geometry data -class DiscretizationProperty(ExpressionBase, HasDOFDesc): +class DiscretizationProperty(GrudgeStringifiable, HasDOFDesc, Expression): pass diff --git a/test/test_grudge.py b/test/test_grudge.py index cf820f4c..3dfa3984 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -514,6 +514,35 @@ def test_bessel(ctx_factory): assert z < 1e-15 +def test_ExternalCall(ctx_factory): + cl_ctx = ctx_factory() + queue = cl.CommandQueue(cl_ctx) + + def double(x): + return 2 * x + + from meshmode.mesh.generation import generate_regular_rect_mesh + + dims = 2 + + mesh = generate_regular_rect_mesh(a=(0,) * dims, b=(1,) * dims, n=(4,) * dims) + discr = DGDiscretizationWithBoundaries(cl_ctx, mesh, order=1) + + ones = sym.Ones(sym.DD_VOLUME) + from pymbolic.primitives import Variable + op = ( + ones * 3 + + sym.ExternalCall( + Variable("double"), + (ones,), + sym.DD_VOLUME)) + + bound_op = bind(discr, op) + + result = bound_op(queue, double=double) + assert (result == 5).get().all() + + # You can test individual routines by typing # $ python test_grudge.py 'test_routine()' -- GitLab From 56543a47350001241767c0fc9b19adea45d8187b Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Apr 2019 19:41:17 -0500 Subject: [PATCH 2/8] Fix indentation --- grudge/symbolic/compiler.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py index 2b25f0f0..ce896235 100644 --- a/grudge/symbolic/compiler.py +++ b/grudge/symbolic/compiler.py @@ -1271,14 +1271,14 @@ class OperatorCompiler(mappers.IdentityMapper): def map_external_call(self, expr, codegen_state): return self.assign_to_new_var( - codegen_state, - type(expr)( - expr.function, - [self.assign_to_new_var( - codegen_state, - self.rec(par, codegen_state)) - for par in expr.parameters], - expr.dd)) + codegen_state, + type(expr)( + expr.function, + [self.assign_to_new_var( + codegen_state, + self.rec(par, codegen_state)) + for par in expr.parameters], + expr.dd)) def map_call(self, expr, codegen_state): from grudge.symbolic.primitives import CFunction @@ -1288,13 +1288,13 @@ class OperatorCompiler(mappers.IdentityMapper): # 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( - codegen_state, - self.rec(par, codegen_state)) - for par in expr.parameters])) + codegen_state, + type(expr)( + expr.function, + [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, codegen_state): try: -- GitLab From fecf4224691b5aadbe0e2ae163e301332effdb46 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Apr 2019 19:43:37 -0500 Subject: [PATCH 3/8] flake8 fix --- test/test_grudge.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_grudge.py b/test/test_grudge.py index 3dfa3984..115694f0 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -514,7 +514,7 @@ def test_bessel(ctx_factory): assert z < 1e-15 -def test_ExternalCall(ctx_factory): +def test_ExternalCall(ctx_factory): # noqa cl_ctx = ctx_factory() queue = cl.CommandQueue(cl_ctx) -- GitLab From 2f992e19d238739b9edc90fc37d2e86ac074d0cc Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Apr 2019 19:45:01 -0500 Subject: [PATCH 4/8] Remove object array stuff (doesn't work) --- grudge/symbolic/primitives.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index 65676118..e9af1ee9 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -352,20 +352,6 @@ class ExternalCall(HasDOFDesc, GrudgeStringifiable, pymbolic.primitives.Call): def __getinitargs__(self): return (self.function, self.parameters, self.dd) - def get_hash(self): - # Object arrays are permitted as parameters, but are not hashable. - params = [] - for param in self.parameters: - if isinstance(param, np.ndarray): - assert param.dtype == np.object - param = tuple(param) - params.append(param) - - return hash( - (type(self).__name__, self.function) - + tuple(params) - + (self.dd,)) - mapper_method = "map_external_call" -- GitLab From 8fee177f22e5d83b7ca097604474cdeed8524783 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Apr 2019 19:47:14 -0500 Subject: [PATCH 5/8] Document --- grudge/symbolic/primitives.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index e9af1ee9..b8fe6c5a 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -346,6 +346,8 @@ class ScalarVariable(Variable): class ExternalCall(HasDOFDesc, GrudgeStringifiable, pymbolic.primitives.Call): + """A call to a user-supplied function with a :class:`DOFDesc`. + """ init_arg_names = ("function", "parameters", "dd") -- GitLab From 9ad13fd263e9f9759be4f6b7126e7543ef1bf002 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 4 Apr 2019 02:49:17 +0200 Subject: [PATCH 6/8] pymbolic.primitives.Expression -> Expression --- grudge/symbolic/primitives.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index b8fe6c5a..d7a3f585 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -396,7 +396,7 @@ bessel_y = CFunction("bessel_y") # {{{ technical helpers -class OperatorBinding(GrudgeStringifiable, pymbolic.primitives.Expression): +class OperatorBinding(GrudgeStringifiable, Expression): init_arg_names = ("op", "field") def __init__(self, op, field): -- GitLab From d1dee68bd739f7ab824e7bde4b2d1793d8d29aeb Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Apr 2019 19:56:38 -0500 Subject: [PATCH 7/8] Redo the inheritance hierarchy to use ExpressionBase again --- grudge/symbolic/primitives.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index b8fe6c5a..571547fd 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -34,12 +34,12 @@ from meshmode.discretization.connection import ( # noqa from pymbolic.primitives import ( # noqa cse_scope as cse_scope_base, - make_common_subexpression as cse, If, Comparison, Expression) + make_common_subexpression as cse, If, Comparison) from pymbolic.geometric_algebra import MultiVector from pytools.obj_array import join_fields, make_obj_array # noqa -class GrudgeStringifiable(object): +class ExpressionBase(pymbolic.primitives.Expression): def stringifier(self): from grudge.symbolic.mappers import StringifyMapper @@ -320,7 +320,7 @@ class cse_scope(cse_scope_base): # noqa DISCRETIZATION = "grudge_discretization" -class Variable(HasDOFDesc, GrudgeStringifiable, pymbolic.primitives.Variable): +class Variable(HasDOFDesc, ExpressionBase, pymbolic.primitives.Variable): """A user-supplied input variable with a known :class:`DOFDesc`. """ init_arg_names = ("name", "dd") @@ -345,7 +345,7 @@ class ScalarVariable(Variable): super(ScalarVariable, self).__init__(name, DD_SCALAR) -class ExternalCall(HasDOFDesc, GrudgeStringifiable, pymbolic.primitives.Call): +class ExternalCall(HasDOFDesc, ExpressionBase, pymbolic.primitives.Call): """A call to a user-supplied function with a :class:`DOFDesc`. """ @@ -369,7 +369,7 @@ def make_sym_mv(name, dim, var_factory=None): make_sym_array(name, dim, var_factory)) -class CFunction(GrudgeStringifiable, pymbolic.primitives.Variable): +class CFunction(ExpressionBase, pymbolic.primitives.Variable): """A symbol representing a C-level function, to be used as the function argument of :class:`pymbolic.primitives.Call`. """ @@ -396,7 +396,7 @@ bessel_y = CFunction("bessel_y") # {{{ technical helpers -class OperatorBinding(GrudgeStringifiable, pymbolic.primitives.Expression): +class OperatorBinding(ExpressionBase): init_arg_names = ("op", "field") def __init__(self, op, field): @@ -441,16 +441,13 @@ class PrioritizedSubexpression(pymbolic.primitives.CommonSubexpression): # }}} -class Ones(GrudgeStringifiable, HasDOFDesc, Expression): - def __getinitargs__(self): - return (self.dd,) - +class Ones(HasDOFDesc, ExpressionBase): mapper_method = intern("map_ones") # {{{ geometry data -class DiscretizationProperty(GrudgeStringifiable, HasDOFDesc, Expression): +class DiscretizationProperty(HasDOFDesc, ExpressionBase): pass -- GitLab From f5c141a2b51a88c12200cfbf2da6e3988307bce5 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Apr 2019 20:04:15 -0500 Subject: [PATCH 8/8] Fix whitespace changes --- grudge/symbolic/compiler.py | 1 + grudge/symbolic/primitives.py | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py index ce896235..4a8ee89e 100644 --- a/grudge/symbolic/compiler.py +++ b/grudge/symbolic/compiler.py @@ -1287,6 +1287,7 @@ class OperatorCompiler(mappers.IdentityMapper): 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)( diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index 571547fd..d35da7c0 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -40,7 +40,6 @@ from pytools.obj_array import join_fields, make_obj_array # noqa class ExpressionBase(pymbolic.primitives.Expression): - def stringifier(self): from grudge.symbolic.mappers import StringifyMapper return StringifyMapper -- GitLab