From 5a962fb66923899364e71715231c485d514f8e0f Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Sun, 4 Nov 2018 23:55:55 -0600
Subject: [PATCH 01/11] Actually implement Evaluator mapper

---
 grudge/symbolic/mappers/__init__.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/grudge/symbolic/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py
index 2b0b2242..73762157 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -1263,7 +1263,32 @@ class FluxExchangeCollector(CSECachingMapperMixin, CollectorMixin, CombineMapper
 # {{{ evaluation
 
 class Evaluator(pymbolic.mapper.evaluator.EvaluationMapper):
-    pass
+    def map_operator_binding(self, expr):
+        return expr.op(self.rec(expr.field))
+
+    def map_node_coordinate_component(self, expr):
+        return expr
+
+    def map_call(self, expr):
+        return type(expr)(
+                expr.function,
+                tuple(self.rec(child, *args, **kwargs)
+                    for child in expr.parameters))
+
+    def map_call_with_kwargs(self, expr):
+        return type(expr)(
+                expr.function,
+                tuple(self.rec(child, *args, **kwargs)
+                    for child in expr.parameters),
+                dict(
+                    (key, self.rec(val, *args, **kwargs))
+                    for key, val in six.iteritems(expr.kw_parameters))
+                    )
+
+    def map_common_subexpression(self, expr):
+        return type(expr)(expr.child, expr.prefix, expr.scope)
+
+
 
 # }}}
 
-- 
GitLab


From 22732e9a0c169de0efea0b67abfff973a2f1af89 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Sun, 4 Nov 2018 23:59:38 -0600
Subject: [PATCH 02/11] Implement initial dagrt->grudge rewriter

---
 examples/dagrt-fusion.py | 141 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 141 insertions(+)
 create mode 100644 examples/dagrt-fusion.py

diff --git a/examples/dagrt-fusion.py b/examples/dagrt-fusion.py
new file mode 100644
index 00000000..2ddf1332
--- /dev/null
+++ b/examples/dagrt-fusion.py
@@ -0,0 +1,141 @@
+from grudge.models.wave import StrongWaveOperator
+from leap.rk import LSRK4Method
+from dagrt.codegen import PythonCodeGenerator
+import numpy as np
+
+dt_method = LSRK4Method(component_id="w")
+dt_code = dt_method.generate()
+from grudge import sym
+
+import dagrt.language as lang
+import pymbolic.mapper as pmap
+import pymbolic.primitives as p
+import grudge.symbolic.mappers as gmap
+
+
+def topological_sort(stmts, root_deps):
+    id_to_stmt = {stmt.id: stmt for stmt in stmts}
+
+    ordered_stmts = []
+    satisfied = set()
+
+    def satisfy_dep(name):
+        if name in satisfied:
+            return
+
+        stmt = id_to_stmt[name]
+        for dep in stmt.depends_on:
+            satisfy_dep(dep)
+        ordered_stmts.append(stmt)
+        satisfied.add(name)
+
+    for d in root_deps:
+        satisfy_dep(d)
+
+    return ordered_stmts
+
+
+class DagrtToGrudgeRewriter(pmap.IdentityMapper):
+    def __init__(self, context):
+        self.context = context
+
+    def map_variable(self, expr):
+        return self.context[expr.name]
+
+    def map_call(self, expr):
+        raise ValueError("function call not expected")
+
+
+class GrudgeArgSubstitutor(gmap.IdentityMapper):
+    def __init__(self, args):
+        self.args = args
+
+    def map_variable(self, expr):
+        if expr.name in self.args:
+            return self.args[expr.name]
+        else:
+            return super().map_variable(expr)
+
+
+def transcribe_phase(dag, phase_name, rhs_name, op_sym_expr):
+    phase = dag.phases[phase_name]
+
+    ctx = {"<t>": sym.var("t", sym.DD_SCALAR),
+           "<dt>": sym.var("dt", sym.DD_SCALAR),
+           "<state>w": sym.make_sym_array("w", 3),
+           "<p>residual": sym.var("residual")
+          }
+    output_vars = [v for v in ctx.keys()]
+    yielded_states = []
+
+    from dagrt.codegen.transform import isolate_function_calls_in_phase
+    ordered_stmts = topological_sort(
+        isolate_function_calls_in_phase(
+            phase,
+            dag.get_stmt_id_generator(),
+            dag.get_var_name_generator()).statements,
+        phase.depends_on)
+
+    for stmt in ordered_stmts:
+        if stmt.condition is not True:
+            raise NotImplementedError(
+                "non-True condition (in statement '%s') not supported"
+                % stmt.id)
+
+        if isinstance(stmt, lang.Nop):
+            pass
+        elif isinstance(stmt, lang.AssignExpression):
+            if not isinstance(stmt.lhs, p.Variable):
+                raise NotImplementedError("lhs of statement %s is not a variable: %s"
+                        % (stmt.id, stmt.lhs))
+            print(stmt.lhs.name)
+            ctx[stmt.lhs.name] = sym.cse(
+                DagrtToGrudgeRewriter(ctx)(stmt.rhs),
+                stmt.lhs.name)
+        elif isinstance(stmt, lang.AssignFunctionCall):
+            if stmt.function_id != rhs_name:
+                raise NotImplementedError("statement '%s' calls unsupported function '%s'"
+                        % (stmt.id, stmt.function_id))
+            if stmt.parameters:
+                raise NotImplementedError(
+                    "statement '%s' calls function '%s' with positional arguments"
+                    % (stmt.id, stmt.function_id))
+            kwargs = {name: sym.cse(DagrtToGrudgeRewriter(ctx)(arg))
+                      for name, arg in stmt.kw_parameters.items()}
+
+            if len(stmt.assignees) != 1:
+                raise NotImplementedError(
+                    "statement '%s' calls function '%s' "
+                    "with more than one LHS"
+                    % (stmt.id, stmt.function_id))
+
+            assignee, = stmt.assignees
+            ctx[assignee] = GrudgeArgSubstitutor(
+                kwargs)(op_sym_expr)
+        elif isinstance(stmt, lang.YieldState):
+            d2g = DagrtToGrudgeRewriter(ctx)
+            yielded_states.append(
+                (stmt.time_id, d2g(stmt.time), stmt.component_id, d2g(stmt.expression)))
+        else:
+            raise NotImplementedError("statement %s is of unsupported type ''%s'"
+                        % (stmt.id, type(stmt).__name__))
+
+    return output_vars, [ctx[ov] for ov in output_vars], yielded_states
+
+
+def main():
+    pass
+    from meshmode.mesh import BTAG_ALL, BTAG_NONE
+    op = StrongWaveOperator(-0.1, 2,
+            dirichlet_tag=BTAG_NONE,
+            neumann_tag=BTAG_NONE,
+            radiation_tag=BTAG_ALL,
+            flux_type="upwind")
+    sym_op = op.sym_operator()
+
+    ov, results, ys = transcribe_phase(dt_code, "primary", "<func>w", sym_op)
+    for ov_i, res_i in zip(ov, results):
+        print(ov_i, sym.pretty(res_i))
+
+if __name__ == "__main__":
+    main()
-- 
GitLab


From 0eea5b75949d315d7d4d47a511fe1dcdab5a5a0a Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Mon, 5 Nov 2018 00:24:25 -0600
Subject: [PATCH 03/11] Un-confuse symbolic and non-symbolic evaluator

---
 examples/dagrt-fusion.py            | 48 ++++++++++++++++++-----------
 grudge/symbolic/mappers/__init__.py |  4 +++
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/examples/dagrt-fusion.py b/examples/dagrt-fusion.py
index 2ddf1332..2fbcc12f 100644
--- a/examples/dagrt-fusion.py
+++ b/examples/dagrt-fusion.py
@@ -1,16 +1,16 @@
 from grudge.models.wave import StrongWaveOperator
 from leap.rk import LSRK4Method
-from dagrt.codegen import PythonCodeGenerator
-import numpy as np
+import numpy as np  # noqa
 
 dt_method = LSRK4Method(component_id="w")
 dt_code = dt_method.generate()
 from grudge import sym
 
 import dagrt.language as lang
-import pymbolic.mapper as pmap
 import pymbolic.primitives as p
 import grudge.symbolic.mappers as gmap
+from pymbolic.mapper.evaluator import EvaluationMapper \
+        as PymbolicEvaluationMapper
 
 
 def topological_sort(stmts, root_deps):
@@ -35,7 +35,10 @@ def topological_sort(stmts, root_deps):
     return ordered_stmts
 
 
-class DagrtToGrudgeRewriter(pmap.IdentityMapper):
+# Use evaluation, not identity mappers to propagate symbolic vectors to
+# outermost level.
+
+class DagrtToGrudgeRewriter(PymbolicEvaluationMapper):
     def __init__(self, context):
         self.context = context
 
@@ -46,24 +49,27 @@ class DagrtToGrudgeRewriter(pmap.IdentityMapper):
         raise ValueError("function call not expected")
 
 
-class GrudgeArgSubstitutor(gmap.IdentityMapper):
+class GrudgeArgSubstitutor(gmap.SymbolicEvaluator):
     def __init__(self, args):
         self.args = args
 
-    def map_variable(self, expr):
+    def map_grudge_variable(self, expr):
         if expr.name in self.args:
             return self.args[expr.name]
         else:
             return super().map_variable(expr)
 
 
-def transcribe_phase(dag, phase_name, rhs_name, op_sym_expr):
+def transcribe_phase(dag, phase_name, rhs_name, sym_operator):
+    sym_operator = gmap.OperatorBinder()(sym_operator)
+
     phase = dag.phases[phase_name]
 
-    ctx = {"<t>": sym.var("t", sym.DD_SCALAR),
-           "<dt>": sym.var("dt", sym.DD_SCALAR),
-           "<state>w": sym.make_sym_array("w", 3),
-           "<p>residual": sym.var("residual")
+    ctx = {
+            "<t>": sym.var("t", sym.DD_SCALAR),
+            "<dt>": sym.var("dt", sym.DD_SCALAR),
+            "<state>w": sym.make_sym_array("w", 3),
+            "<p>residual_w": sym.var("residual")
           }
     output_vars = [v for v in ctx.keys()]
     yielded_states = []
@@ -88,13 +94,16 @@ def transcribe_phase(dag, phase_name, rhs_name, op_sym_expr):
             if not isinstance(stmt.lhs, p.Variable):
                 raise NotImplementedError("lhs of statement %s is not a variable: %s"
                         % (stmt.id, stmt.lhs))
-            print(stmt.lhs.name)
             ctx[stmt.lhs.name] = sym.cse(
                 DagrtToGrudgeRewriter(ctx)(stmt.rhs),
-                stmt.lhs.name)
+                (
+                    stmt.lhs.name
+                    .replace("<", "")
+                    .replace(">", "")))
         elif isinstance(stmt, lang.AssignFunctionCall):
             if stmt.function_id != rhs_name:
-                raise NotImplementedError("statement '%s' calls unsupported function '%s'"
+                raise NotImplementedError(
+                        "statement '%s' calls unsupported function '%s'"
                         % (stmt.id, stmt.function_id))
             if stmt.parameters:
                 raise NotImplementedError(
@@ -111,11 +120,12 @@ def transcribe_phase(dag, phase_name, rhs_name, op_sym_expr):
 
             assignee, = stmt.assignees
             ctx[assignee] = GrudgeArgSubstitutor(
-                kwargs)(op_sym_expr)
+                kwargs)(sym_operator)
         elif isinstance(stmt, lang.YieldState):
             d2g = DagrtToGrudgeRewriter(ctx)
             yielded_states.append(
-                (stmt.time_id, d2g(stmt.time), stmt.component_id, d2g(stmt.expression)))
+                (stmt.time_id, d2g(stmt.time), stmt.component_id,
+                    d2g(stmt.expression)))
         else:
             raise NotImplementedError("statement %s is of unsupported type ''%s'"
                         % (stmt.id, type(stmt).__name__))
@@ -124,7 +134,6 @@ def transcribe_phase(dag, phase_name, rhs_name, op_sym_expr):
 
 
 def main():
-    pass
     from meshmode.mesh import BTAG_ALL, BTAG_NONE
     op = StrongWaveOperator(-0.1, 2,
             dirichlet_tag=BTAG_NONE,
@@ -135,7 +144,10 @@ def main():
 
     ov, results, ys = transcribe_phase(dt_code, "primary", "<func>w", sym_op)
     for ov_i, res_i in zip(ov, results):
-        print(ov_i, sym.pretty(res_i))
+        print(75*"#")
+        print(sym.pretty(res_i))
+        print(f"{ov_i} IS THE ABOVE ^")
+
 
 if __name__ == "__main__":
     main()
diff --git a/grudge/symbolic/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py
index 73762157..4c9ad5d4 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -1263,6 +1263,10 @@ class FluxExchangeCollector(CSECachingMapperMixin, CollectorMixin, CombineMapper
 # {{{ evaluation
 
 class Evaluator(pymbolic.mapper.evaluator.EvaluationMapper):
+    pass
+
+
+class SymbolicEvaluator(pymbolic.mapper.evaluator.EvaluationMapper):
     def map_operator_binding(self, expr):
         return expr.op(self.rec(expr.field))
 
-- 
GitLab


From 0c0a3197bc46ade2fb88d26e20fe5650821155e6 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Mon, 5 Nov 2018 00:34:11 -0600
Subject: [PATCH 04/11] SymbolicEvaluator: placate flake8

---
 grudge/symbolic/mappers/__init__.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/grudge/symbolic/mappers/__init__.py b/grudge/symbolic/mappers/__init__.py
index 4c9ad5d4..ca6d976f 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -1267,19 +1267,19 @@ class Evaluator(pymbolic.mapper.evaluator.EvaluationMapper):
 
 
 class SymbolicEvaluator(pymbolic.mapper.evaluator.EvaluationMapper):
-    def map_operator_binding(self, expr):
-        return expr.op(self.rec(expr.field))
+    def map_operator_binding(self, expr, *args, **kwargs):
+        return expr.op(self.rec(expr.field, *args, **kwargs))
 
-    def map_node_coordinate_component(self, expr):
+    def map_node_coordinate_component(self, expr, *args, **kwargs):
         return expr
 
-    def map_call(self, expr):
+    def map_call(self, expr, *args, **kwargs):
         return type(expr)(
                 expr.function,
                 tuple(self.rec(child, *args, **kwargs)
                     for child in expr.parameters))
 
-    def map_call_with_kwargs(self, expr):
+    def map_call_with_kwargs(self, expr, *args, **kwargs):
         return type(expr)(
                 expr.function,
                 tuple(self.rec(child, *args, **kwargs)
@@ -1292,8 +1292,6 @@ class SymbolicEvaluator(pymbolic.mapper.evaluator.EvaluationMapper):
     def map_common_subexpression(self, expr):
         return type(expr)(expr.child, expr.prefix, expr.scope)
 
-
-
 # }}}
 
 
-- 
GitLab


From c7de48709cfbf1e86d9dae2fe8d411d846daf5e3 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Tue, 27 Nov 2018 14:03:10 -0600
Subject: [PATCH 05/11] Tweak dagrt fusion example to test successful bind

---
 examples/dagrt-fusion.py | 33 +++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/examples/dagrt-fusion.py b/examples/dagrt-fusion.py
index 2fbcc12f..1099be11 100644
--- a/examples/dagrt-fusion.py
+++ b/examples/dagrt-fusion.py
@@ -1,17 +1,16 @@
+import pyopencl as cl
 from grudge.models.wave import StrongWaveOperator
 from leap.rk import LSRK4Method
 import numpy as np  # noqa
 
-dt_method = LSRK4Method(component_id="w")
-dt_code = dt_method.generate()
-from grudge import sym
-
 import dagrt.language as lang
 import pymbolic.primitives as p
 import grudge.symbolic.mappers as gmap
 from pymbolic.mapper.evaluator import EvaluationMapper \
         as PymbolicEvaluationMapper
 
+from grudge import sym, bind, DGDiscretizationWithBoundaries
+
 
 def topological_sort(stmts, root_deps):
     id_to_stmt = {stmt.id: stmt for stmt in stmts}
@@ -66,10 +65,10 @@ def transcribe_phase(dag, phase_name, rhs_name, sym_operator):
     phase = dag.phases[phase_name]
 
     ctx = {
-            "<t>": sym.var("t", sym.DD_SCALAR),
-            "<dt>": sym.var("dt", sym.DD_SCALAR),
-            "<state>w": sym.make_sym_array("w", 3),
-            "<p>residual_w": sym.var("residual")
+            "<t>": sym.var("input_t", sym.DD_SCALAR),
+            "<dt>": sym.var("input_dt", sym.DD_SCALAR),
+            "<state>w": sym.make_sym_array("input_w", 3),
+            "<p>residual_w": sym.var("input_residual")
           }
     output_vars = [v for v in ctx.keys()]
     yielded_states = []
@@ -134,6 +133,9 @@ def transcribe_phase(dag, phase_name, rhs_name, sym_operator):
 
 
 def main():
+    dt_method = LSRK4Method(component_id="w")
+    dt_code = dt_method.generate()
+
     from meshmode.mesh import BTAG_ALL, BTAG_NONE
     op = StrongWaveOperator(-0.1, 2,
             dirichlet_tag=BTAG_NONE,
@@ -148,6 +150,21 @@ def main():
         print(sym.pretty(res_i))
         print(f"{ov_i} IS THE ABOVE ^")
 
+    from meshmode.mesh.generation import generate_regular_rect_mesh
+    mesh = generate_regular_rect_mesh(
+            a=(-0.5,)*2,
+            b=(0.5,)*2,
+            n=(16,)*2)
+
+    cl_ctx = cl.create_some_context()
+    discr = DGDiscretizationWithBoundaries(cl_ctx, mesh, order=3)
+
+    from pytools.obj_array import join_fields
+    results = join_fields(results[0], results[1], *results[2], *results[3])
+    bound_op = bind(discr, results)
+
+    print(bound_op.eval_code)
+
 
 if __name__ == "__main__":
     main()
-- 
GitLab


From 6bd7e7ba5374aa03d2403294932291d94b31e1ae Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Wed, 5 Dec 2018 15:49:22 -0600
Subject: [PATCH 06/11] Fix recursion in CSE in EvalMapper

---
 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 64d26c38..76864e7d 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -1294,7 +1294,7 @@ class SymbolicEvaluator(pymbolic.mapper.evaluator.EvaluationMapper):
                     )
 
     def map_common_subexpression(self, expr):
-        return type(expr)(expr.child, expr.prefix, expr.scope)
+        return type(expr)(self.rec(expr.child), expr.prefix, expr.scope)
 
 # }}}
 
-- 
GitLab


From a55e373194a506b0ee1679e122431c2612974867 Mon Sep 17 00:00:00 2001
From: Fan Kiat Chan <fchan5@illinois.edu>
Date: Sat, 8 Dec 2018 17:42:01 -0600
Subject: [PATCH 07/11] Fix Code class in compiler.py to be aware of
 subscripted variables

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

diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py
index 53db42f1..84a55eea 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -446,7 +446,7 @@ 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:
@@ -454,7 +454,8 @@ class Code(object):
 
         from pytools import flatten
         discardable_vars = set(available_names) - set(flatten(
-            [dep.name for dep in insn.get_dependencies()]
+            # [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))
 
@@ -578,6 +579,10 @@ class Code(object):
                                 future_sub_timer.stop().submit()
                             break
 
+        # print("Reachables:")
+        # for insn in done_insns:
+        #     input(insn)
+
         if len(done_insns) < len(self.instructions):
             print("Unreachable instructions:")
             for insn in set(self.instructions) - done_insns:
-- 
GitLab


From 7c12f0cfe9e5fd6cf340614dd6f09876eaf18072 Mon Sep 17 00:00:00 2001
From: Fan Kiat Chan <fchan5@illinois.edu>
Date: Sat, 8 Dec 2018 22:15:12 -0600
Subject: [PATCH 08/11] Fix linter failures in compiler.py

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

diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py
index 84a55eea..fc94e5eb 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -446,7 +446,8 @@ 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.aggregate.name if isinstance(dep, Subscript)
+                    else dep.name) in available_names
                     for dep in insn.get_dependencies())]
 
         if not available_insns:
@@ -454,8 +455,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()]
+            [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))
 
@@ -579,10 +580,6 @@ class Code(object):
                                 future_sub_timer.stop().submit()
                             break
 
-        # print("Reachables:")
-        # for insn in done_insns:
-        #     input(insn)
-
         if len(done_insns) < len(self.instructions):
             print("Unreachable instructions:")
             for insn in set(self.instructions) - done_insns:
-- 
GitLab


From 9b81cb7e7387412f08b7b91c9c1644ea3804ccb1 Mon Sep 17 00:00:00 2001
From: Fan Kiat Chan <fchan5@illinois.edu>
Date: Tue, 11 Dec 2018 14:53:51 -0600
Subject: [PATCH 09/11] Rename dagrt-fusion.py to dagrt_fusion.py for module
 import

---
 examples/{dagrt-fusion.py => dagrt_fusion.py} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename examples/{dagrt-fusion.py => dagrt_fusion.py} (100%)

diff --git a/examples/dagrt-fusion.py b/examples/dagrt_fusion.py
similarity index 100%
rename from examples/dagrt-fusion.py
rename to examples/dagrt_fusion.py
-- 
GitLab


From bd634d5b7848bfdb12f4b09155bea7c0b448cd66 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Wed, 3 Apr 2019 19:31:58 -0500
Subject: [PATCH 10/11] WIP

---
 grudge/execution.py                  | 14 ++++++-
 grudge/symbolic/compiler.py          | 40 +++++++++++++-----
 grudge/symbolic/dofdesc_inference.py |  3 ++
 grudge/symbolic/mappers/__init__.py  | 58 +++++++++++++------------
 grudge/symbolic/primitives.py        | 63 ++++++++++++++++++++--------
 test/test_grudge.py                  | 29 +++++++++++++
 6 files changed, 150 insertions(+), 57 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 7bb51201..87d993ec 100644
--- a/grudge/symbolic/compiler.py
+++ b/grudge/symbolic/compiler.py
@@ -669,13 +669,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)),
@@ -888,6 +889,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
@@ -972,8 +977,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"
@@ -1263,6 +1271,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):
@@ -1270,15 +1289,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 0cea10df..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
 
@@ -1269,33 +1300,6 @@ class FluxExchangeCollector(CSECachingMapperMixin, CollectorMixin, CombineMapper
 class Evaluator(pymbolic.mapper.evaluator.EvaluationMapper):
     pass
 
-
-class SymbolicEvaluator(pymbolic.mapper.evaluator.EvaluationMapper):
-    def map_operator_binding(self, expr, *args, **kwargs):
-        return expr.op(self.rec(expr.field, *args, **kwargs))
-
-    def map_node_coordinate_component(self, expr, *args, **kwargs):
-        return expr
-
-    def map_call(self, expr, *args, **kwargs):
-        return type(expr)(
-                expr.function,
-                tuple(self.rec(child, *args, **kwargs)
-                    for child in expr.parameters))
-
-    def map_call_with_kwargs(self, expr, *args, **kwargs):
-        return type(expr)(
-                expr.function,
-                tuple(self.rec(child, *args, **kwargs)
-                    for child in expr.parameters),
-                dict(
-                    (key, self.rec(val, *args, **kwargs))
-                    for key, val in six.iteritems(expr.kw_parameters))
-                    )
-
-    def map_common_subexpression(self, expr):
-        return type(expr)(self.rec(expr.child), expr.prefix, expr.scope)
-
 # }}}
 
 
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 ffc1cb8cad7f9dbbeedcdc38a3b079f89f6db0e4 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Sat, 11 May 2019 01:03:01 -0500
Subject: [PATCH 11/11] Implement a function registry; remove the ExternalCall
 node

---
 grudge/execution.py                  |  62 ++-------
 grudge/function_registry.py          | 196 +++++++++++++++++++++++++++
 grudge/symbolic/compiler.py          |  55 ++++----
 grudge/symbolic/dofdesc_inference.py |  15 +-
 grudge/symbolic/mappers/__init__.py  |  38 +-----
 grudge/symbolic/operators.py         |   6 +-
 grudge/symbolic/primitives.py        |  53 +-------
 test/test_grudge.py                  |  18 ++-
 8 files changed, 267 insertions(+), 176 deletions(-)
 create mode 100644 grudge/function_registry.py

diff --git a/grudge/execution.py b/grudge/execution.py
index 9084b823..89db0a1a 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 -------------------------------------------------
@@ -94,56 +96,9 @@ class ExecutionMapper(mappers.Evaluator,
                 value = ary
         return value
 
-    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]
-
-        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 (
-                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
@@ -483,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):
@@ -635,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)
@@ -658,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..bac92cc3
--- /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, *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):
+
+    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 87d993ec..76724445 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
@@ -609,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):
@@ -669,14 +671,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, ExternalCall
+    from grudge.symbolic.primitives import DTAG_SCALAR
 
     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 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)),
@@ -827,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
@@ -889,13 +903,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):
+        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"
@@ -975,13 +986,17 @@ 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
 
     def map_insn_assign(self, insn):
-        from grudge.symbolic.primitives import OperatorBinding, ExternalCall
+        from grudge.symbolic.primitives import OperatorBinding
 
         if (
                 len(insn.exprs) == 1
-                and isinstance(insn.exprs[0], (OperatorBinding, ExternalCall))):
+                and (
+                    isinstance(insn.exprs[0], OperatorBinding)
+                    or is_external_call(
+                        insn.exprs[0], self.function_registry))):
             return insn
 
         iname = "grdg_i"
@@ -1103,7 +1118,7 @@ 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
 
@@ -1119,6 +1134,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()
@@ -1152,7 +1168,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)
@@ -1271,20 +1288,8 @@ 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):
+        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 96ead088..c44a7940 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,13 +189,9 @@ class DOFDescInferenceMapper(RecursiveMapper, CSECachingMapperMixin):
                 self.rec(par)
                 for par in expr.parameters]
 
-        assert arg_dds
-
-        # FIXME
-        return arg_dds[0]
-
-    def map_external_call(self, expr):
-        return expr.dd
+        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 d1b70475..82d5cc28 100644
--- a/grudge/symbolic/mappers/__init__.py
+++ b/grudge/symbolic/mappers/__init__.py
@@ -210,17 +210,9 @@ 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
 
-    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)
-
     # }}}
 
 
@@ -274,15 +266,6 @@ 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,
@@ -293,10 +276,7 @@ class FlopCounter(
     def map_grudge_variable(self, expr):
         return 0
 
-    def map_c_function(self, expr):
-        return 1
-
-    def map_external_call(self, expr):
+    def map_call(self, expr):
         return 1
 
     def map_ones(self, expr):
@@ -858,24 +838,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_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,
@@ -1253,7 +1221,6 @@ class CollectorMixin(OperatorReducerMixin, LocalOpReducerMixin, FluxOpReducerMix
 
     map_variable = map_constant
     map_grudge_variable = map_constant
-    map_c_function = map_grudge_variable
 
     map_ones = map_grudge_variable
     map_node_coordinate_component = map_grudge_variable
@@ -1270,9 +1237,6 @@ 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/operators.py b/grudge/symbolic/operators.py
index 53fb1422..f0126689 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.Variable("fabs")(
                     arg * sym.MassOperator()(arg)))
 
         if isinstance(norm_squared, np.ndarray):
             norm_squared = norm_squared.sum()
 
-        return sym.CFunction("sqrt")(norm_squared)
+        return sym.sqrt(norm_squared)
 
     elif p == np.Inf:
-        result = sym.NodalMax(dd_in=dd)(sym.CFunction("fabs")(arg))
+        result = sym.NodalMax(dd_in=dd)(sym.Variable("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 65676118..98b13fa2 100644
--- a/grudge/symbolic/primitives.py
+++ b/grudge/symbolic/primitives.py
@@ -78,7 +78,6 @@ Symbols
 .. autoclass:: ExternalCall
 .. autoclass:: make_sym_array
 .. autoclass:: make_sym_mv
-.. autoclass:: CFunction
 
 .. function :: sqrt(arg)
 .. function :: exp(arg)
@@ -345,30 +344,6 @@ 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)
@@ -381,27 +356,13 @@ def make_sym_mv(name, dim, var_factory=None):
             make_sym_array(name, dim, var_factory))
 
 
-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 __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
+sqrt = Variable("sqrt")
+exp = Variable("exp")
+sin = Variable("sin")
+cos = Variable("cos")
+bessel_j = Variable("bessel_j")
+bessel_y = Variable("bessel_y")
 
 # }}}
 
diff --git a/test/test_grudge.py b/test/test_grudge.py
index 3dfa3984..47ed3efc 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_external_call(ctx_factory):
     cl_ctx = ctx_factory()
     queue = cl.CommandQueue(cl_ctx)
 
@@ -532,12 +532,18 @@ def test_ExternalCall(ctx_factory):
     from pymbolic.primitives import Variable
     op = (
             ones * 3
-            + sym.ExternalCall(
-                Variable("double"),
-                (ones,),
-                sym.DD_VOLUME))
+            + Variable("double")(ones))
 
-    bound_op = bind(discr, op)
+    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()
-- 
GitLab