From 3d8276922000aee6964d1bfb44e948e2ee50b25a Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sat, 11 May 2019 02:19:18 -0500
Subject: [PATCH 1/9] WIP: Implement a function registry with support for
 external calls

---
 grudge/execution.py                  |  52 ++-----
 grudge/function_registry.py          | 198 +++++++++++++++++++++++++++
 grudge/symbolic/compiler.py          |  60 +++++---
 grudge/symbolic/dofdesc_inference.py |   9 +-
 grudge/symbolic/mappers/__init__.py  |  10 +-
 grudge/symbolic/primitives.py        |  58 +++-----
 test/test_grudge.py                  |  35 +++++
 7 files changed, 319 insertions(+), 103 deletions(-)
 create mode 100644 grudge/function_registry.py

diff --git a/grudge/execution.py b/grudge/execution.py
index c20aa4bc..9e48aba3 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 00000000..6efa39da
--- /dev/null
+++ b/grudge/function_registry.py
@@ -0,0 +1,198 @@
+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, *args):
+        assert len(args) == 1
+
+        func_name = self.identifier
+
+        arg, = args
+        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
+            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/symbolic/compiler.py b/grudge/symbolic/compiler.py
index f8dc173e..c2750e33 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -31,7 +31,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
+from pymbolic.primitives import Variable, Subscript, Call
 from six.moves import intern
 
 from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_1  # noqa: F401
@@ -447,7 +447,8 @@ class Code(object):
         available_insns = [
                 (insn, insn.priority) for insn in self.instructions
                 if insn not in done_insns
-                and all(dep.name in available_names
+                and all((dep.aggregate.name if isinstance(dep, Subscript)
+                    else dep.name) in available_names
                     for dep in insn.get_dependencies())]
 
         if not available_insns:
@@ -455,7 +456,8 @@ class Code(object):
 
         from pytools import flatten
         discardable_vars = set(available_names) - set(flatten(
-            [dep.name for dep in insn.get_dependencies()]
+            [dep.aggregate.name if isinstance(dep, Subscript) else dep.name
+            for dep in insn.get_dependencies()]
             for insn in self.instructions
             if insn not in done_insns))
 
@@ -607,6 +609,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 +678,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 +829,21 @@ def aggregate_assignments(inf_mapper, instructions, result,
 
 # {{{ to-loopy mapper
 
+def is_external_call(expr, function_registry):
+    if not isinstance(expr, Call):
+        return False
+    return not is_function_loopyable(expr.function, function_registry)
+
+
+def is_function_loopyable(function, function_registry):
+    assert isinstance(function, Variable)
+    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,21 +1294,19 @@ 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
             # 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..26e5e19a 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,7 +189,9 @@ class DOFDescInferenceMapper(RecursiveMapper, CSECachingMapperMixin):
                 self.rec(par)
                 for par in expr.parameters]
 
-        assert arg_dds
+        return (
+                self.function_registry[expr.function.name]
+                .get_result_dofdesc(arg_dds))
 
         # FIXME
         return arg_dds[0]
diff --git a/grudge/symbolic/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py
index 8b9bf73b..bfb60879 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -211,8 +211,6 @@ class IdentityMapperMixin(LocalOpReducerMixin, FluxOpReducerMixin):
         # it's a leaf--no changing children
         return expr
 
-    map_c_function = map_grudge_variable
-
     map_ones = map_grudge_variable
     map_node_coordinate_component = map_grudge_variable
 
@@ -279,7 +277,7 @@ class FlopCounter(
     def map_grudge_variable(self, expr):
         return 0
 
-    def map_c_function(self, expr):
+    def map_call(self, expr):
         return 1
 
     def map_ones(self, expr):
@@ -841,9 +839,6 @@ 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))
 
@@ -1232,7 +1227,8 @@ class CollectorMixin(OperatorReducerMixin, LocalOpReducerMixin, FluxOpReducerMix
         return OrderedSet()
 
     map_grudge_variable = map_constant
-    map_c_function = map_grudge_variable
+    # Found in function call nodes
+    map_variable = map_grudge_variable
 
     map_ones = map_grudge_variable
     map_node_coordinate_component = map_grudge_variable
diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py
index 231a70df..3c2dc96e 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,14 @@ 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
-
-    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"
-
-
-sqrt = CFunction("sqrt")
-exp = CFunction("exp")
-sin = CFunction("sin")
-cos = CFunction("cos")
-bessel_j = CFunction("bessel_j")
-bessel_y = CFunction("bessel_y")
+# function symbols
+CFunction = Variable
+sqrt = Variable("sqrt")
+exp = Variable("exp")
+sin = Variable("sin")
+cos = Variable("cos")
+bessel_j = Variable("bessel_j")
+bessel_y = Variable("bessel_y")
 
 # }}}
 
@@ -424,16 +413,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 97c5a012..f646dff4 100644
--- a/test/test_grudge.py
+++ b/test/test_grudge.py
@@ -539,6 +539,41 @@ 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)
+    from pymbolic.primitives import Variable
+    op = (
+            ones * 3
+            + Variable("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()'
 
-- 
GitLab


From fe186d3d2b96f5368397fe1ad46906bdedba02e9 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sat, 11 May 2019 08:08:29 -0500
Subject: [PATCH 2/9] Remove dead code

---
 grudge/symbolic/dofdesc_inference.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/grudge/symbolic/dofdesc_inference.py b/grudge/symbolic/dofdesc_inference.py
index 26e5e19a..c44a7940 100644
--- a/grudge/symbolic/dofdesc_inference.py
+++ b/grudge/symbolic/dofdesc_inference.py
@@ -193,9 +193,6 @@ class DOFDescInferenceMapper(RecursiveMapper, CSECachingMapperMixin):
                 self.function_registry[expr.function.name]
                 .get_result_dofdesc(arg_dds))
 
-        # FIXME
-        return arg_dds[0]
-
     # }}}
 
     # {{{ instruction mappings
-- 
GitLab


From 9055bda3472356fbe83f9266be10a02098e1d5bb Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sun, 12 May 2019 20:00:37 -0500
Subject: [PATCH 3/9] Remove some unrelated changes

---
 grudge/symbolic/compiler.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py
index c2750e33..8c1c5925 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -447,8 +447,7 @@ class Code(object):
         available_insns = [
                 (insn, insn.priority) for insn in self.instructions
                 if insn not in done_insns
-                and all((dep.aggregate.name if isinstance(dep, Subscript)
-                    else dep.name) in available_names
+                and all(dep.name in available_names
                     for dep in insn.get_dependencies())]
 
         if not available_insns:
@@ -456,8 +455,7 @@ class Code(object):
 
         from pytools import flatten
         discardable_vars = set(available_names) - set(flatten(
-            [dep.aggregate.name if isinstance(dep, Subscript) else dep.name
-            for dep in insn.get_dependencies()]
+            [dep.name for dep in insn.get_dependencies()]
             for insn in self.instructions
             if insn not in done_insns))
 
@@ -1299,14 +1297,15 @@ 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:
-- 
GitLab


From 0e42886c32db681dddc1787f669db3e5bb27baa5 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sun, 12 May 2019 20:01:05 -0500
Subject: [PATCH 4/9] Introduce a FunctionSymbol class

---
 examples/gas_dynamics/lbm-simple.py    |  6 +++---
 examples/wave/wiggly.py                |  4 ++--
 grudge/models/em.py                    |  2 +-
 grudge/models/gas_dynamics/__init__.py |  8 ++++----
 grudge/models/wave.py                  |  2 +-
 grudge/symbolic/mappers/__init__.py    |  9 ++++++---
 grudge/symbolic/operators.py           |  6 +++---
 grudge/symbolic/primitives.py          | 23 +++++++++++++++--------
 test/test_grudge.py                    |  3 +--
 9 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/examples/gas_dynamics/lbm-simple.py b/examples/gas_dynamics/lbm-simple.py
index f99d5b76..2d3496da 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 f7860744..11c01212 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/models/em.py b/grudge/models/em.py
index f0e44f90..bf7495e2 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 e5a8ddc9..4df408d8 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 40710eb2..0ac60a33 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/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py
index bfb60879..304bfb4c 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -211,6 +211,7 @@ class IdentityMapperMixin(LocalOpReducerMixin, FluxOpReducerMixin):
         # it's a leaf--no changing children
         return expr
 
+    map_function_symbol = map_grudge_variable
     map_ones = map_grudge_variable
     map_node_coordinate_component = map_grudge_variable
 
@@ -277,7 +278,7 @@ class FlopCounter(
     def map_grudge_variable(self, expr):
         return 0
 
-    def map_call(self, expr):
+    def map_function_symbol(self, expr):
         return 1
 
     def map_ones(self, expr):
@@ -842,6 +843,9 @@ class StringifyMapper(pymbolic.mapper.stringifier.StringifyMapper):
     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
+
     def map_interpolation(self, expr, enclosing_prec):
         return "Interp" + self._format_op_dd(expr)
 
@@ -1227,8 +1231,7 @@ class CollectorMixin(OperatorReducerMixin, LocalOpReducerMixin, FluxOpReducerMix
         return OrderedSet()
 
     map_grudge_variable = map_constant
-    # Found in function call nodes
-    map_variable = 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 53fb1422..a1899ed2 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 3c2dc96e..c59f1518 100644
--- a/grudge/symbolic/primitives.py
+++ b/grudge/symbolic/primitives.py
@@ -354,14 +354,21 @@ def make_sym_mv(name, dim, var_factory=None):
             make_sym_array(name, dim, var_factory))
 
 
-# function symbols
-CFunction = Variable
-sqrt = Variable("sqrt")
-exp = Variable("exp")
-sin = Variable("sin")
-cos = Variable("cos")
-bessel_j = Variable("bessel_j")
-bessel_y = Variable("bessel_y")
+class FunctionSymbol(ExpressionBase, pymbolic.primitives.Variable):
+    """A symbol to be used as the function argument of
+    :class:`pymbolic.primitives.Call`.
+
+    """
+
+    mapper_method = "map_function_symbol"
+
+
+sqrt = FunctionSymbol("sqrt")
+exp = FunctionSymbol("exp")
+sin = FunctionSymbol("sin")
+cos = FunctionSymbol("cos")
+bessel_j = FunctionSymbol("bessel_j")
+bessel_y = FunctionSymbol("bessel_y")
 
 # }}}
 
diff --git a/test/test_grudge.py b/test/test_grudge.py
index f646dff4..288922e0 100644
--- a/test/test_grudge.py
+++ b/test/test_grudge.py
@@ -554,10 +554,9 @@ def test_external_call(ctx_factory):
     discr = DGDiscretizationWithBoundaries(cl_ctx, mesh, order=1)
 
     ones = sym.Ones(sym.DD_VOLUME)
-    from pymbolic.primitives import Variable
     op = (
             ones * 3
-            + Variable("double")(ones))
+            + sym.FunctionSymbol("double")(ones))
 
     from grudge.function_registry import (
             base_function_registry, register_external_function)
-- 
GitLab


From 1873e91fb86506468dd8b0e461017eadf5294b18 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sun, 12 May 2019 20:04:46 -0500
Subject: [PATCH 5/9] Make assertions more strict

---
 grudge/symbolic/compiler.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py
index 8c1c5925..864bbed5 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -31,7 +31,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, Call
+from pymbolic.primitives import Variable, Subscript
 from six.moves import intern
 
 from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_1  # noqa: F401
@@ -828,13 +828,15 @@ 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):
-    assert isinstance(function, Variable)
+    from grudge.primitives import FunctionSymbol
+    assert isinstance(function, FunctionSymbol)
     return function_registry[function.name].supports_codegen
 
 
-- 
GitLab


From c40c36212d126da8d2627096b5a12bbbe39d9420 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sun, 12 May 2019 20:53:18 -0500
Subject: [PATCH 6/9] Fix function printing

---
 grudge/symbolic/mappers/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/grudge/symbolic/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py
index 304bfb4c..df53cadf 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -844,7 +844,7 @@ class StringifyMapper(pymbolic.mapper.stringifier.StringifyMapper):
         return "%s:%s" % (expr.name, self._format_dd(expr.dd))
 
     def map_function_symbol(self, expr, enclosing_prec):
-        return expr
+        return expr.name
 
     def map_interpolation(self, expr, enclosing_prec):
         return "Interp" + self._format_op_dd(expr)
-- 
GitLab


From bea73dfb0ca165de932c1385e3359871f2322a72 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sun, 12 May 2019 20:58:00 -0500
Subject: [PATCH 7/9] Import fix

---
 grudge/symbolic/compiler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py
index 864bbed5..498e5172 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -835,7 +835,7 @@ def is_external_call(expr, function_registry):
 
 
 def is_function_loopyable(function, function_registry):
-    from grudge.primitives import FunctionSymbol
+    from grudge.symbolic.primitives import FunctionSymbol
     assert isinstance(function, FunctionSymbol)
     return function_registry[function.name].supports_codegen
 
-- 
GitLab


From 12401570bf4a3d391507d652acc6c851aec5faf3 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sun, 12 May 2019 21:14:18 -0500
Subject: [PATCH 8/9] Change prototype to accept only a single arg

---
 grudge/function_registry.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/grudge/function_registry.py b/grudge/function_registry.py
index 6efa39da..df709f93 100644
--- a/grudge/function_registry.py
+++ b/grudge/function_registry.py
@@ -71,12 +71,9 @@ class CElementwiseUnaryFunction(Function):
         assert len(arg_dds) == 1
         return arg_dds[0]
 
-    def __call__(self, queue, *args):
-        assert len(args) == 1
-
+    def __call__(self, queue, arg):
         func_name = self.identifier
 
-        arg, = args
         from numbers import Number
         if (
                 isinstance(arg, Number)
-- 
GitLab


From 10d1391f0a85705e68c80f101a0b37d82e9b2e8b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andreas=20Kl=C3=B6ckner?= <inform@tiker.net>
Date: Sat, 1 Jun 2019 22:12:49 +0200
Subject: [PATCH 9/9] Apply suggestion to grudge/function_registry.py

---
 grudge/function_registry.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/grudge/function_registry.py b/grudge/function_registry.py
index df709f93..31961b6f 100644
--- a/grudge/function_registry.py
+++ b/grudge/function_registry.py
@@ -88,6 +88,7 @@ class CElementwiseUnaryFunction(Function):
         i = Variable("i")
 
         if self.identifier == "fabs":  # FIXME
+            # Loopy has a type-adaptive "abs", but no "fabs".
             func_name = "abs"
 
         cached_name += func_name
-- 
GitLab