diff --git a/examples/gas_dynamics/lbm-simple.py b/examples/gas_dynamics/lbm-simple.py
index f99d5b76994b46f079b41af9f74bf69789ae18f5..2d3496da340fbda01adeb33b6a59290b76b1c698 100644
--- a/examples/gas_dynamics/lbm-simple.py
+++ b/examples/gas_dynamics/lbm-simple.py
@@ -63,12 +63,12 @@ def main(write_output=True, dtype=np.float32):
from grudge.data import CompiledExpressionData
def ic_expr(t, x, fields):
- from grudge.symbolic import CFunction
+ from grudge.symbolic import FunctionSymbol
from pymbolic.primitives import IfPositive
from pytools.obj_array import make_obj_array
- tanh = CFunction("tanh")
- sin = CFunction("sin")
+ tanh = FunctionSymbol("tanh")
+ sin = FunctionSymbol("sin")
rho = 1
u0 = 0.05
diff --git a/examples/wave/wiggly.py b/examples/wave/wiggly.py
index f786074443d2123882bcaecf85972777b6694622..11c0121214982b70fb891b443dfec267ce5defb6 100644
--- a/examples/wave/wiggly.py
+++ b/examples/wave/wiggly.py
@@ -70,8 +70,8 @@ def main(write_output=True,
from grudge.models.wave import StrongWaveOperator
op = StrongWaveOperator(-1, discr.dimensions,
source_f=
- sym.CFunction("sin")(source_omega*sym.ScalarParameter("t"))
- * sym.CFunction("exp")(
+ sym.FunctionSymbol("sin")(source_omega*sym.ScalarParameter("t"))
+ * sym.FunctionSymbol("exp")(
-np.dot(sym_source_center_dist, sym_source_center_dist)
/ source_width**2),
dirichlet_tag="boundary",
diff --git a/grudge/execution.py b/grudge/execution.py
index c20aa4bc6271f1b9a8f5620c8b3021c22563880c..9e48aba39cd6e3aa5481f0bf9f31f3340381e4f6 100644
--- a/grudge/execution.py
+++ b/grudge/execution.py
@@ -31,6 +31,7 @@ from pytools import memoize_in
import grudge.symbolic.mappers as mappers
from grudge import sym
+from grudge.function_registry import base_function_registry
import logging
logger = logging.getLogger(__name__)
@@ -50,6 +51,7 @@ class ExecutionMapper(mappers.Evaluator,
super(ExecutionMapper, self).__init__(context)
self.discrwb = bound_op.discrwb
self.bound_op = bound_op
+ self.function_registry = bound_op.function_registry
self.queue = queue
# {{{ expression mappings -------------------------------------------------
@@ -95,45 +97,8 @@ class ExecutionMapper(mappers.Evaluator,
return value
def map_call(self, expr):
- from pymbolic.primitives import Variable
- assert isinstance(expr.function, Variable)
-
- # FIXME: Make a way to register functions
-
args = [self.rec(p) for p in expr.parameters]
- from numbers import Number
- representative_arg = args[0]
- if (
- isinstance(representative_arg, Number)
- or (isinstance(representative_arg, np.ndarray)
- and representative_arg.shape == ())):
- func = getattr(np, expr.function.name)
- return func(*args)
-
- cached_name = "map_call_knl_"
-
- i = Variable("i")
- func = Variable(expr.function.name)
- if expr.function.name == "fabs": # FIXME
- func = Variable("abs")
- cached_name += "abs"
- else:
- cached_name += expr.function.name
-
- @memoize_in(self.bound_op, cached_name)
- def knl():
- knl = lp.make_kernel(
- "{[i]: 0<=i<n}",
- [
- lp.Assignment(Variable("out")[i],
- func(Variable("a")[i]))
- ], default_offset=lp.auto)
- return lp.split_iname(knl, "i", 128, outer_tag="g.0", inner_tag="l.0")
-
- assert len(args) == 1
- evt, (out,) = knl()(self.queue, a=args[0])
-
- return out
+ return self.function_registry[expr.function.name](self.queue, *args)
def map_nodal_sum(self, op, field_expr):
# FIXME: Could allow array scalars
@@ -473,12 +438,15 @@ class MPISendFuture(object):
# {{{ bound operator
class BoundOperator(object):
- def __init__(self, discrwb, discr_code, eval_code, debug_flags, allocator=None):
+
+ def __init__(self, discrwb, discr_code, eval_code, debug_flags,
+ function_registry, allocator=None):
self.discrwb = discrwb
self.discr_code = discr_code
self.eval_code = eval_code
self.operator_data_cache = {}
self.debug_flags = debug_flags
+ self.function_registry = function_registry
self.allocator = allocator
def __str__(self):
@@ -625,7 +593,8 @@ def process_sym_operator(discrwb, sym_operator, post_bind_mapper=None,
def bind(discr, sym_operator, post_bind_mapper=lambda x: x,
- debug_flags=set(), allocator=None):
+ debug_flags=set(), allocator=None,
+ function_registry=base_function_registry):
# from grudge.symbolic.mappers import QuadratureUpsamplerRemover
# sym_operator = QuadratureUpsamplerRemover(self.quad_min_degrees)(
# sym_operator)
@@ -648,9 +617,10 @@ def bind(discr, sym_operator, post_bind_mapper=lambda x: x,
dumper=dump_sym_operator)
from grudge.symbolic.compiler import OperatorCompiler
- discr_code, eval_code = OperatorCompiler(discr)(sym_operator)
+ discr_code, eval_code = OperatorCompiler(discr, function_registry)(sym_operator)
bound_op = BoundOperator(discr, discr_code, eval_code,
+ function_registry=function_registry,
debug_flags=debug_flags, allocator=allocator)
if "dump_op_code" in debug_flags:
diff --git a/grudge/function_registry.py b/grudge/function_registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..31961b6f48a26183b09c6673e04e860eeb7128d1
--- /dev/null
+++ b/grudge/function_registry.py
@@ -0,0 +1,196 @@
+from __future__ import division, with_statement
+
+__copyright__ = """
+Copyright (C) 2013 Andreas Kloeckner
+Copyright (C) 2019 Matt Wala
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+
+import loopy as lp
+import numpy as np
+
+from pytools import RecordWithoutPickling, memoize_in
+
+
+# {{{ function
+
+class FunctionNotFound(KeyError):
+ pass
+
+
+class Function(RecordWithoutPickling):
+ """
+ .. attribute:: identifier
+ .. attribute:: supports_codegen
+ .. automethod:: __call__
+ .. automethod:: get_result_dofdesc
+ """
+
+ def __init__(self, identifier, **kwargs):
+ super(Function, self).__init__(identifier=identifier, **kwargs)
+
+ def __call__(self, queue, *args, **kwargs):
+ """Call the function implementation, if available."""
+ raise TypeError("function '%s' is not callable" % self.identifier)
+
+ def get_result_dofdesc(self, arg_dds):
+ """Return the :class:`grudge.symbolic.primitives.DOFDesc` for the return value
+ of the function.
+
+ :arg arg_dds: A list of :class:`grudge.symbolic.primitives.DOFDesc` instances
+ for each argument
+ """
+ raise NotImplementedError
+
+
+class CElementwiseUnaryFunction(Function):
+
+ supports_codegen = True
+
+ def get_result_dofdesc(self, arg_dds):
+ assert len(arg_dds) == 1
+ return arg_dds[0]
+
+ def __call__(self, queue, arg):
+ func_name = self.identifier
+
+ from numbers import Number
+ if (
+ isinstance(arg, Number)
+ or (isinstance(arg, np.ndarray)
+ and arg.shape == ())):
+ func = getattr(np, func_name)
+ return func(arg)
+
+ cached_name = "map_call_knl_"
+
+ from pymbolic.primitives import Variable
+ i = Variable("i")
+
+ if self.identifier == "fabs": # FIXME
+ # Loopy has a type-adaptive "abs", but no "fabs".
+ func_name = "abs"
+
+ cached_name += func_name
+
+ @memoize_in(self, cached_name)
+ def knl():
+ knl = lp.make_kernel(
+ "{[i]: 0<=i<n}",
+ [
+ lp.Assignment(Variable("out")[i],
+ Variable(func_name)(Variable("a")[i]))
+ ], default_offset=lp.auto)
+ return lp.split_iname(knl, "i", 128, outer_tag="g.0", inner_tag="l.0")
+
+ evt, (out,) = knl()(queue, a=arg)
+ return out
+
+
+class CBesselFunction(Function):
+
+ supports_codegen = True
+
+ def get_result_dofdesc(self, arg_dds):
+ assert len(arg_dds) == 2
+ return arg_dds[1]
+
+
+class FixedDOFDescExternalFunction(Function):
+
+ supports_codegen = False
+
+ def __init__(self, identifier, implementation, dd):
+ super(FixedDOFDescExternalFunction, self).__init__(
+ identifier,
+ implementation=implementation,
+ dd=dd)
+
+ def __call__(self, queue, *args, **kwargs):
+ return self.implementation(queue, *args, **kwargs)
+
+ def get_result_dofdesc(self, arg_dds):
+ return self.dd
+
+# }}}
+
+
+# {{{ function registry
+
+class FunctionRegistry(RecordWithoutPickling):
+ def __init__(self, id_to_function=None):
+ if id_to_function is None:
+ id_to_function = {}
+
+ super(FunctionRegistry, self).__init__(
+ id_to_function=id_to_function)
+
+ def register(self, function):
+ """Return a copy of *self* with *function* registered."""
+
+ if function.identifier in self.id_to_function:
+ raise ValueError("function '%s' is already registered"
+ % function.identifier)
+
+ new_id_to_function = self.id_to_function.copy()
+ new_id_to_function[function.identifier] = function
+ return self.copy(id_to_function=new_id_to_function)
+
+ def __getitem__(self, function_id):
+ try:
+ return self.id_to_function[function_id]
+ except KeyError:
+ raise FunctionNotFound(
+ "unknown function: '%s'"
+ % function_id)
+
+ def __contains__(self, function_id):
+ return function_id in self.id_to_function
+
+# }}}
+
+
+def _make_bfr():
+ bfr = FunctionRegistry()
+
+ bfr = bfr.register(CElementwiseUnaryFunction("sqrt"))
+ bfr = bfr.register(CElementwiseUnaryFunction("exp"))
+ bfr = bfr.register(CElementwiseUnaryFunction("fabs"))
+ bfr = bfr.register(CElementwiseUnaryFunction("sin"))
+ bfr = bfr.register(CElementwiseUnaryFunction("cos"))
+ bfr = bfr.register(CBesselFunction("bessel_j"))
+ bfr = bfr.register(CBesselFunction("bessel_y"))
+
+ return bfr
+
+
+base_function_registry = _make_bfr()
+
+
+def register_external_function(
+ function_registry, identifier, implementation, dd):
+ return function_registry.register(
+ FixedDOFDescExternalFunction(
+ identifier, implementation, dd))
+
+# vim: foldmethod=marker
diff --git a/grudge/models/em.py b/grudge/models/em.py
index f0e44f903d226e28710709651e948293a534d1b8..bf7495e2c9447b54d37d49a78905e549d9df54d8 100644
--- a/grudge/models/em.py
+++ b/grudge/models/em.py
@@ -355,7 +355,7 @@ class MaxwellOperator(HyperbolicOperator):
return 1/sqrt(self.epsilon*self.mu) # a number
else:
import grudge.symbolic as sym
- return sym.NodalMax()(1/sym.CFunction("sqrt")(self.epsilon*self.mu))
+ return sym.NodalMax()(1/sym.FunctionSymbol("sqrt")(self.epsilon*self.mu))
def max_eigenvalue(self, t, fields=None, discr=None, context={}):
if self.fixed_material:
diff --git a/grudge/models/gas_dynamics/__init__.py b/grudge/models/gas_dynamics/__init__.py
index e5a8ddc9418f1bac4776f252d2951fab3d3c1a2b..4df408d811a826a1f924b855f0ff8a9fce8a29f2 100644
--- a/grudge/models/gas_dynamics/__init__.py
+++ b/grudge/models/gas_dynamics/__init__.py
@@ -326,8 +326,8 @@ class GasDynamicsOperator(TimeDependentOperator):
def characteristic_velocity_optemplate(self, state):
from grudge.symbolic.operators import ElementwiseMaxOperator
- from grudge.symbolic.primitives import CFunction
- sqrt = CFunction("sqrt")
+ from grudge.symbolic.primitives import FunctionSymbol
+ sqrt = FunctionSymbol("sqrt")
sound_speed = cse(sqrt(
self.equation_of_state.gamma*self.cse_p(state)/self.cse_rho(state)),
@@ -743,8 +743,8 @@ class GasDynamicsOperator(TimeDependentOperator):
volq_flux = self.flux(self.volq_state())
faceq_flux = self.flux(self.faceq_state())
- from grudge.symbolic.primitives import CFunction
- sqrt = CFunction("sqrt")
+ from grudge.symbolic.primitives import FunctionSymbol
+ sqrt = FunctionSymbol("sqrt")
speed = self.characteristic_velocity_optemplate(self.state())
diff --git a/grudge/models/wave.py b/grudge/models/wave.py
index 40710eb2363b7e475f1be299f1b69540343309be..0ac60a336a4f72f12ba4f3eeb0bc2d79c9f07509 100644
--- a/grudge/models/wave.py
+++ b/grudge/models/wave.py
@@ -459,7 +459,7 @@ class VariableCoefficientWeakWaveOperator(HyperbolicOperator):
self.radiation_tag])
def max_eigenvalue(self, t, fields=None, discr=None):
- return sym.NodalMax()(sym.CFunction("fabs")(self.c))
+ return sym.NodalMax()(sym.FunctionSymbol("fabs")(self.c))
# }}}
diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py
index f8dc173e771010ff597fd3831f83ffdccfabc0a4..498e5172254ca7acb4a55fb25057e837d024fcfb 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -607,6 +607,8 @@ def aggregate_assignments(inf_mapper, instructions, result,
max_vectors_in_batch_expr):
from pymbolic.primitives import Variable
+ function_registry = inf_mapper.function_registry
+
# {{{ aggregation helpers
def get_complete_origins_set(insn, skip_levels=0):
@@ -674,6 +676,7 @@ def aggregate_assignments(inf_mapper, instructions, result,
isinstance(insn, Assign)
and not isinstance(insn, ToDiscretizationScopedAssign)
and not isinstance(insn, FromDiscretizationScopedAssign)
+ and not is_external_call(insn.exprs[0], function_registry)
and not any(
inf_mapper.infer_for_name(n).domain_tag == DTAG_SCALAR
for n in insn.names)),
@@ -824,9 +827,23 @@ def aggregate_assignments(inf_mapper, instructions, result,
# {{{ to-loopy mapper
+def is_external_call(expr, function_registry):
+ from pymbolic.primitives import Call
+ if not isinstance(expr, Call):
+ return False
+ return not is_function_loopyable(expr.function, function_registry)
+
+
+def is_function_loopyable(function, function_registry):
+ from grudge.symbolic.primitives import FunctionSymbol
+ assert isinstance(function, FunctionSymbol)
+ return function_registry[function.name].supports_codegen
+
+
class ToLoopyExpressionMapper(mappers.IdentityMapper):
def __init__(self, dd_inference_mapper, temp_names, iname):
self.dd_inference_mapper = dd_inference_mapper
+ self.function_registry = dd_inference_mapper.function_registry
self.temp_names = temp_names
self.iname = iname
from pymbolic import var
@@ -887,8 +904,9 @@ class ToLoopyExpressionMapper(mappers.IdentityMapper):
"%s_%d" % (expr.aggregate.name, subscript))
def map_call(self, expr):
- if isinstance(expr.function, sym.CFunction):
+ if is_function_loopyable(expr.function, self.function_registry):
from pymbolic import var
+
func_name = expr.function.name
if func_name == "fabs":
func_name = "abs"
@@ -968,11 +986,18 @@ def bessel_function_mangler(kernel, name, arg_dtypes):
class ToLoopyInstructionMapper(object):
def __init__(self, dd_inference_mapper):
self.dd_inference_mapper = dd_inference_mapper
+ self.function_registry = dd_inference_mapper.function_registry
self.insn_count = 0
def map_insn_assign(self, insn):
from grudge.symbolic.primitives import OperatorBinding
- if len(insn.exprs) == 1 and isinstance(insn.exprs[0], OperatorBinding):
+
+ if (
+ len(insn.exprs) == 1
+ and (
+ isinstance(insn.exprs[0], OperatorBinding)
+ or is_external_call(
+ insn.exprs[0], self.function_registry))):
return insn
iname = "grdg_i"
@@ -1097,7 +1122,8 @@ class CodeGenerationState(Record):
class OperatorCompiler(mappers.IdentityMapper):
- def __init__(self, discr, prefix="_expr", max_vectors_in_batch_expr=None):
+ def __init__(self, discr, function_registry,
+ prefix="_expr", max_vectors_in_batch_expr=None):
super(OperatorCompiler, self).__init__()
self.prefix = prefix
@@ -1113,6 +1139,7 @@ class OperatorCompiler(mappers.IdentityMapper):
self.assigned_names = set()
self.discr = discr
+ self.function_registry = function_registry
from pytools import UniqueNameGenerator
self.name_gen = UniqueNameGenerator()
@@ -1146,7 +1173,8 @@ class OperatorCompiler(mappers.IdentityMapper):
del self.discr_code
from grudge.symbolic.dofdesc_inference import DOFDescInferenceMapper
- inf_mapper = DOFDescInferenceMapper(discr_code + eval_code)
+ inf_mapper = DOFDescInferenceMapper(
+ discr_code + eval_code, self.function_registry)
eval_code = aggregate_assignments(
inf_mapper, eval_code, result, self.max_vectors_in_batch_expr)
@@ -1266,8 +1294,7 @@ class OperatorCompiler(mappers.IdentityMapper):
return result_var
def map_call(self, expr, codegen_state):
- from grudge.symbolic.primitives import CFunction
- if isinstance(expr.function, CFunction):
+ if is_function_loopyable(expr.function, self.function_registry):
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
diff --git a/grudge/symbolic/dofdesc_inference.py b/grudge/symbolic/dofdesc_inference.py
index 92be126f7f3f081b668247e1fe25a73b122a4887..c44a7940b9ec142171a0b0a67885c041ec9d43cb 100644
--- a/grudge/symbolic/dofdesc_inference.py
+++ b/grudge/symbolic/dofdesc_inference.py
@@ -76,7 +76,8 @@ class InferrableMultiAssignment(object):
class DOFDescInferenceMapper(RecursiveMapper, CSECachingMapperMixin):
- def __init__(self, assignments, name_to_dofdesc=None, check=True):
+ def __init__(self, assignments, function_registry,
+ name_to_dofdesc=None, check=True):
"""
:arg assignments: a list of objects adhering to
:class:`InferrableMultiAssignment`.
@@ -98,6 +99,8 @@ class DOFDescInferenceMapper(RecursiveMapper, CSECachingMapperMixin):
self.name_to_dofdesc = name_to_dofdesc
+ self.function_registry = function_registry
+
def infer_for_name(self, name):
try:
return self.name_to_dofdesc[name]
@@ -186,10 +189,9 @@ class DOFDescInferenceMapper(RecursiveMapper, CSECachingMapperMixin):
self.rec(par)
for par in expr.parameters]
- assert arg_dds
-
- # FIXME
- return arg_dds[0]
+ return (
+ self.function_registry[expr.function.name]
+ .get_result_dofdesc(arg_dds))
# }}}
diff --git a/grudge/symbolic/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py
index 8b9bf73b7a7d5080ffe0652a52c0e1926681b749..df53cadf86e69cf3ca97081cf9c3618570120450 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -211,8 +211,7 @@ class IdentityMapperMixin(LocalOpReducerMixin, FluxOpReducerMixin):
# it's a leaf--no changing children
return expr
- map_c_function = map_grudge_variable
-
+ map_function_symbol = map_grudge_variable
map_ones = map_grudge_variable
map_node_coordinate_component = map_grudge_variable
@@ -279,7 +278,7 @@ class FlopCounter(
def map_grudge_variable(self, expr):
return 0
- def map_c_function(self, expr):
+ def map_function_symbol(self, expr):
return 1
def map_ones(self, expr):
@@ -841,12 +840,12 @@ class StringifyMapper(pymbolic.mapper.stringifier.StringifyMapper):
self.rec(expr.op, PREC_NONE),
self.rec(expr.field, PREC_NONE))
- def map_c_function(self, expr, enclosing_prec):
- return expr.name
-
def map_grudge_variable(self, expr, enclosing_prec):
return "%s:%s" % (expr.name, self._format_dd(expr.dd))
+ def map_function_symbol(self, expr, enclosing_prec):
+ return expr.name
+
def map_interpolation(self, expr, enclosing_prec):
return "Interp" + self._format_op_dd(expr)
@@ -1232,7 +1231,7 @@ class CollectorMixin(OperatorReducerMixin, LocalOpReducerMixin, FluxOpReducerMix
return OrderedSet()
map_grudge_variable = map_constant
- map_c_function = map_grudge_variable
+ map_function_symbol = map_constant
map_ones = map_grudge_variable
map_node_coordinate_component = map_grudge_variable
diff --git a/grudge/symbolic/operators.py b/grudge/symbolic/operators.py
index 53fb1422619db271786abb9c2c17df5b3f9a97c0..a1899ed21ef6431f7b57a81fcf143680433b7b53 100644
--- a/grudge/symbolic/operators.py
+++ b/grudge/symbolic/operators.py
@@ -602,16 +602,16 @@ def norm(p, arg, dd=None):
if p == 2:
norm_squared = sym.NodalSum(dd_in=dd)(
- sym.CFunction("fabs")(
+ sym.FunctionSymbol("fabs")(
arg * sym.MassOperator()(arg)))
if isinstance(norm_squared, np.ndarray):
norm_squared = norm_squared.sum()
- return sym.CFunction("sqrt")(norm_squared)
+ return sym.FunctionSymbol("sqrt")(norm_squared)
elif p == np.Inf:
- result = sym.NodalMax(dd_in=dd)(sym.CFunction("fabs")(arg))
+ result = sym.NodalMax(dd_in=dd)(sym.FunctionSymbol("fabs")(arg))
from pymbolic.primitives import Max
if isinstance(result, np.ndarray):
diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py
index 231a70dfa374d8c0e60e2830ff13427f5d33c0dd..c59f1518b5034a9b69c3aca4a22c0c4bab92d7a6 100644
--- a/grudge/symbolic/primitives.py
+++ b/grudge/symbolic/primitives.py
@@ -76,7 +76,6 @@ Symbols
.. autoclass:: ScalarVariable
.. autoclass:: make_sym_array
.. autoclass:: make_sym_mv
-.. autoclass:: CFunction
.. function :: sqrt(arg)
.. function :: exp(arg)
@@ -289,7 +288,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 +306,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 +326,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,)
@@ -349,30 +354,21 @@ def make_sym_mv(name, dim, var_factory=None):
make_sym_array(name, dim, var_factory))
-class CFunction(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
+class FunctionSymbol(ExpressionBase, pymbolic.primitives.Variable):
+ """A symbol to be used as the function argument of
+ :class:`pymbolic.primitives.Call`.
- def __call__(self, *exprs):
- from pytools.obj_array import with_object_array_or_scalar_n_args
- from functools import partial
- return with_object_array_or_scalar_n_args(
- partial(pymbolic.primitives.Expression.__call__, self),
- *exprs)
+ """
- mapper_method = "map_c_function"
+ mapper_method = "map_function_symbol"
-sqrt = CFunction("sqrt")
-exp = CFunction("exp")
-sin = CFunction("sin")
-cos = CFunction("cos")
-bessel_j = CFunction("bessel_j")
-bessel_y = CFunction("bessel_y")
+sqrt = FunctionSymbol("sqrt")
+exp = FunctionSymbol("exp")
+sin = FunctionSymbol("sin")
+cos = FunctionSymbol("cos")
+bessel_j = FunctionSymbol("bessel_j")
+bessel_y = FunctionSymbol("bessel_y")
# }}}
@@ -424,16 +420,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 97c5a012fce82c377373889ac6c053d9189db096..288922e07529aa70bffbc1f812efe582e146ff60 100644
--- a/test/test_grudge.py
+++ b/test/test_grudge.py
@@ -539,6 +539,40 @@ def test_bessel(ctx_factory):
assert z < 1e-15
+def test_external_call(ctx_factory):
+ cl_ctx = ctx_factory()
+ queue = cl.CommandQueue(cl_ctx)
+
+ def double(queue, 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)
+ op = (
+ ones * 3
+ + sym.FunctionSymbol("double")(ones))
+
+ from grudge.function_registry import (
+ base_function_registry, register_external_function)
+
+ freg = register_external_function(
+ base_function_registry,
+ "double",
+ implementation=double,
+ dd=sym.DD_VOLUME)
+
+ bound_op = bind(discr, op, function_registry=freg)
+
+ result = bound_op(queue, double=double)
+ assert (result == 5).get().all()
+
+
# You can test individual routines by typing
# $ python test_grudge.py 'test_routine()'