diff --git a/grudge/execution.py b/grudge/execution.py index c20aa4bc6271f1b9a8f5620c8b3021c22563880c..9084b823dcdb2cda502b85891da91bc4b137d510 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 b27d55855237b5e0200b094c2aa164e192022f0b..4a8ee89e87571694da9e4487e3c11b5be0e4900e 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): diff --git a/grudge/symbolic/dofdesc_inference.py b/grudge/symbolic/dofdesc_inference.py index 92be126f7f3f081b668247e1fe25a73b122a4887..96ead0885d2e300e7e87fed4395ec74ff771ccb2 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 d52f7ac5bf99adbb16a618a8c3cb44fb8c65c515..d1b70475491709d8ef75780f38a2e520b13f007f 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 231a70dfa374d8c0e60e2830ff13427f5d33c0dd..d35da7c00648c63fdfa649544439114e5d184dfd 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -74,6 +74,7 @@ Symbols .. autoclass:: Variable .. autoclass:: ScalarVariable +.. autoclass:: ExternalCall .. autoclass:: make_sym_array .. autoclass:: make_sym_mv .. autoclass:: CFunction @@ -289,7 +290,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 +308,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__()) # }}} @@ -320,8 +328,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 +344,18 @@ class ScalarVariable(Variable): super(ScalarVariable, self).__init__(name, DD_SCALAR) +class ExternalCall(HasDOFDesc, ExpressionBase, pymbolic.primitives.Call): + """A call to a user-supplied function with a :class:`DOFDesc`. + """ + + init_arg_names = ("function", "parameters", "dd") + + def __getinitargs__(self): + return (self.function, self.parameters, 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 +368,10 @@ def make_sym_mv(name, dim, var_factory=None): make_sym_array(name, dim, var_factory)) -class CFunction(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`. """ - 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 @@ -424,16 +440,13 @@ class PrioritizedSubexpression(pymbolic.primitives.CommonSubexpression): # }}} -class Ones(ExpressionBase, HasDOFDesc): - def __getinitargs__(self): - return () - +class Ones(HasDOFDesc, ExpressionBase): mapper_method = intern("map_ones") # {{{ geometry data -class DiscretizationProperty(ExpressionBase, HasDOFDesc): +class DiscretizationProperty(HasDOFDesc, ExpressionBase): pass diff --git a/test/test_grudge.py b/test/test_grudge.py index cf820f4c89f259a03e74b1ef37436f4c45833afc..115694f0710934e0d9a20b56e87a85c1955a21d7 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): # noqa + 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()'