Loopy sprint

Jenkinsfile

@@ -40,7 +40,7 @@ pipeline {
               sh 'mkdir tmp'
               dir('tmp') {
                 timestamps {
-                  sh '../scripts/firedrake-install $COMPLEX --tinyasm --disable-ssh --minimal-petsc --slepc --documentation-dependencies --install thetis --install gusto --install icepack --install irksome --install femlium --no-package-manager || (cat firedrake-install.log && /bin/false)'
+                  sh '../scripts/firedrake-install $COMPLEX --tinyasm --disable-ssh --minimal-petsc --slepc --documentation-dependencies --install thetis --install gusto --install icepack --install irksome --no-package-manager|| (cat firedrake-install.log && /bin/false)'
                 }
               }
             }

firedrake/assemble_expressions.py

@@ -286,7 +286,7 @@ def split(self):
     def args(self):
         """Tuple of par_loop arguments for the expression."""
         args = []
-        if self.lvalue in self.rcoefficients:
+        if isinstance(self, AugmentedAssign) or self.lvalue in self.rcoefficients:
             args.append(Arg(weakref.ref(self.lvalue.dat), access=op2.RW))
         else:
             args.append(Arg(weakref.ref(self.lvalue.dat), access=op2.WRITE))
@@ -425,11 +425,14 @@ def pointwise_expression_kernel(exprs, scalar_type):
             except KeyError:
                 continue
             plargs.append(arg)
-            args.append(loopy.GlobalArg(var.name, shape=var.shape, dtype=c.dat.dtype))
+            is_input = arg.access in [op2.INC, op2.MAX, op2.MIN, op2.READ, op2.RW]
+            is_output = arg.access in [op2.INC, op2.MAX, op2.MIN, op2.RW, op2.WRITE]
+            args.append(loopy.GlobalArg(var.name, shape=var.shape, dtype=c.dat.dtype, is_input=is_input, is_output=is_output))
     assert len(coefficients) == 0
-    knl = generate(impero_c, args, scalar_type, kernel_name="expression_kernel",
+    name = "expression_kernel"
+    knl = generate(impero_c, args, scalar_type, kernel_name=name,
                    return_increments=False)
-    return firedrake.op2.Kernel(knl, knl.name), plargs
+    return firedrake.op2.Kernel(knl, name), plargs
 
 
 class dereffed(object):

firedrake/extrusion_utils.py

@@ -102,20 +102,29 @@ def _get_lp_domains(_inames, _extents):
         domains = []
         dd = _get_arity_axis_inames('d')
         domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
-        domains.extend(_get_lp_domains(('c', 'l'), (base_coord_dim, 2)))
+        domains.extend(_get_lp_domains(('c',), (base_coord_dim,)))
+        if layer_heights == 1:
+            domains.extend(_get_lp_domains(('l',), (2,)))
+        else:
+            domains.append("[layer] -> { [l] : 0 <= l <= 1 & 0 <= l + layer < %d}" % layer_heights)
         instructions = """
         ext_coords[{dd}, l, c] = base_coords[{dd}, c]
         ext_coords[{dd}, l, {base_coord_dim}] = ({hv})
         """.format(dd=', '.join(dd),
                    base_coord_dim=base_coord_dim,
                    hv=height_var)
-        ast = lp.make_function(domains, instructions, data, name="pyop2_kernel_uniform_extrusion", target=lp.CTarget(),
+        name = "pyop2_kernel_uniform_extrusion"
+        ast = lp.make_function(domains, instructions, data, name=name, target=lp.CTarget(),
                                seq_dependencies=True, silenced_warnings=["summing_if_branches_ops"])
     elif extrusion_type == 'radial':
         domains = []
         dd = _get_arity_axis_inames('d')
         domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
-        domains.extend(_get_lp_domains(('c', 'k', 'l'), (base_coord_dim, ) * 2 + (2, )))
+        domains.extend(_get_lp_domains(('c', 'k'), (base_coord_dim, ) * 2))
+        if layer_heights == 1:
+            domains.extend(_get_lp_domains(('l',), (2,)))
+        else:
+            domains.append("[layer] -> { [l] : 0 <= l <= 1 & 0 <= l + layer < %d}" % layer_heights)
         instructions = """
         <{RealType}> tt[{dd}] = 0
         <{RealType}> bc[{dd}] = 0
@@ -128,7 +137,8 @@ def _get_lp_domains(_inames, _extents):
         """.format(RealType=RealType,
                    dd=', '.join(dd),
                    hv=height_var)
-        ast = lp.make_function(domains, instructions, data, name="pyop2_kernel_radial_extrusion", target=lp.CTarget(),
+        name = "pyop2_kernel_radial_extrusion"
+        ast = lp.make_function(domains, instructions, data, name=name, target=lp.CTarget(),
                                seq_dependencies=True, silenced_warnings=["summing_if_branches_ops"])
     elif extrusion_type == 'radial_hedgehog':
         # Only implemented for interval in 2D and triangle in 3D.
@@ -203,12 +213,13 @@ def _get_lp_domains(_inames, _extents):
                    _dd=', '.join(_dd),
                    ninst=n_dict[tdim][adim],
                    hv=height_var)
-        ast = lp.make_function(domains, instructions, data, name="pyop2_kernel_radial_hedgehog_extrusion", target=lp.CTarget(),
+        name = "pyop2_kernel_radial_hedgehog_extrusion"
+        ast = lp.make_function(domains, instructions, data, name=name, target=lp.CTarget(),
                                seq_dependencies=True, silenced_warnings=["summing_if_branches_ops"])
     else:
         raise NotImplementedError('Unsupported extrusion type "%s"' % extrusion_type)
 
-    kernel = op2.Kernel(ast, ast.name)
+    kernel = op2.Kernel(ast, name)
     op2.ParLoop(kernel,
                 ext_coords.cell_set,
                 ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),

firedrake/interpolation.py

@@ -236,11 +236,17 @@ def _interpolator(V, tensor, expr, subset, arguments, access):
     if not isinstance(expr, firedrake.Expression):
         if expr.ufl_domain() and expr.ufl_domain() != V.mesh():
             raise NotImplementedError("Interpolation onto another mesh not supported.")
-        ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element,
-                                                                                                                       domain=V.mesh(),
-                                                                                                                       parameters=parameters,
-                                                                                                                       coffee=False)
-        kernel = op2.Kernel(ast, ast.name, requires_zeroed_output_arguments=True)
+        kernel = compile_expression_dual_evaluation(expr, to_element,
+                                                    domain=V.mesh(),
+                                                    parameters=parameters,
+                                                    coffee=False)
+        ast = kernel.ast
+        oriented = kernel.oriented
+        needs_cell_sizes = kernel.needs_cell_sizes
+        coefficients = kernel.coefficients
+        first_coeff_fake_coords = kernel.first_coefficient_fake_coords
+        name = kernel.name
+        kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True)
     elif hasattr(expr, "eval"):
         to_pts = []
         for dual in to_element.fiat_equivalent.dual_basis():

firedrake/parloops.py

@@ -75,13 +75,15 @@ def _form_loopy_kernel(kernel_domains, instructions, measure, args, **kwargs):
     kargs = []
 
     for var, (func, intent) in args.items():
+        is_input = intent in [INC, READ, RW, MAX, MIN]
+        is_output = intent in [INC, RW, WRITE, MAX, MIN]
         if isinstance(func, constant.Constant):
             if intent is not READ:
                 raise RuntimeError("Only READ access is allowed to Constant")
             # Constants modelled as Globals, so no need for double
             # indirection
             ndof = func.dat.cdim
-            kargs.append(loopy.GlobalArg(var, dtype=func.dat.dtype, shape=(ndof,)))
+            kargs.append(loopy.GlobalArg(var, dtype=func.dat.dtype, shape=(ndof,), is_input=is_input, is_output=is_output))
         else:
             # Do we have a component of a mixed function?
             if isinstance(func, Indexed):
@@ -93,7 +95,7 @@ def _form_loopy_kernel(kernel_domains, instructions, measure, args, **kwargs):
             else:
                 if func.function_space().ufl_element().family() == "Real":
                     ndof = func.function_space().dim()  # == 1
-                    kargs.append(loopy.GlobalArg(var, dtype=func.dat.dtype, shape=(ndof,)))
+                    kargs.append(loopy.GlobalArg(var, dtype=func.dat.dtype, shape=(ndof,), is_input=is_input, is_output=is_output))
                     continue
                 else:
                     if len(func.function_space()) > 1:
@@ -104,7 +106,7 @@ def _form_loopy_kernel(kernel_domains, instructions, measure, args, **kwargs):
             if measure.integral_type() == 'interior_facet':
                 ndof *= 2
             # FIXME: shape for facets [2][ndof]?
-            kargs.append(loopy.GlobalArg(var, dtype=dtype, shape=(ndof, cdim)))
+            kargs.append(loopy.GlobalArg(var, dtype=dtype, shape=(ndof, cdim), is_input=is_input, is_output=is_output))
         kernel_domains = kernel_domains.replace(var+".dofs", str(ndof))
 
     if kernel_domains == "":

firedrake/preconditioners/pmg.py

@@ -426,8 +426,8 @@ def prolongation_transfer_kernel_aij(Pk, P1):
     expr = TestFunction(P1)
     to_element = create_base_element(Pk.ufl_element())
 
-    ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element, coffee=False)
-    kernel = op2.Kernel(ast, ast.name)
+    ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _, name = compile_expression_dual_evaluation(expr, to_element, coffee=False)
+    kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True)
     return kernel
 
 
@@ -577,16 +577,16 @@ def make_kernels(self, Vf, Vc):
                    Rc[j] += Afc[i*{dimc} + j] * Rf[i] * w[i];
         }}
         """
-        restrict_kernel = op2.Kernel(restrict_code, "restriction")
+        restrict_kernel = op2.Kernel(restrict_code, "restriction", requires_zeroed_output_arguments=True)
         return prolong_kernel, restrict_kernel
 
     @staticmethod
     def prolongation_transfer_kernel_action(Vf, expr):
         from tsfc import compile_expression_dual_evaluation
         from tsfc.finatinterface import create_base_element
         to_element = create_base_element(Vf.ufl_element())
-        ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element, coffee=False)
-        return op2.Kernel(ast, ast.name)
+        ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _, name = compile_expression_dual_evaluation(expr, to_element, coffee=False)
+        return op2.Kernel(ast, name, requires_zeroed_output_arguments=True)
 
     @staticmethod
     @lru_cache(maxsize=20)
@@ -743,8 +743,10 @@ def make_blas_kernels(Vf, Vc):
         """
 
         from firedrake.slate.slac.compiler import BLASLAPACK_LIB, BLASLAPACK_INCLUDE
-        prolong_kernel = op2.Kernel(prolong_code, "prolongation", include_dirs=BLASLAPACK_INCLUDE.split(), ldargs=BLASLAPACK_LIB.split())
-        restrict_kernel = op2.Kernel(restrict_code, "restriction", include_dirs=BLASLAPACK_INCLUDE.split(), ldargs=BLASLAPACK_LIB.split())
+        prolong_kernel = op2.Kernel(prolong_code, "prolongation", include_dirs=BLASLAPACK_INCLUDE.split(),
+                                    ldargs=BLASLAPACK_LIB.split(), requires_zeroed_output_arguments=True)
+        restrict_kernel = op2.Kernel(restrict_code, "restriction", include_dirs=BLASLAPACK_INCLUDE.split(),
+                                     ldargs=BLASLAPACK_LIB.split(), requires_zeroed_output_arguments=True)
         return prolong_kernel, restrict_kernel
 
     @staticmethod

firedrake/slate/slac/compiler.py

@@ -36,7 +36,7 @@
 
 from pyop2.utils import get_petsc_dir, as_tuple
 from pyop2.mpi import COMM_WORLD
-from pyop2.codegen.rep2loopy import solve_fn_lookup, inv_fn_lookup
+from pyop2.codegen.rep2loopy import SolveCallable, INVCallable
 
 import firedrake.slate.slate as slate
 import numpy as np
@@ -162,16 +162,13 @@ def generate_loopy_kernel(slate_expr, tsfc_parameters=None):
     builder = LocalLoopyKernelBuilder(expression=slate_expr,
                                       tsfc_parameters=tsfc_parameters)
 
-    loopy_merged = merge_loopy(slate_loopy, output_arg, builder, var2terminal)
-    loopy_merged = loopy.register_function_id_to_in_knl_callable_mapper(loopy_merged, inv_fn_lookup)
-    loopy_merged = loopy.register_function_id_to_in_knl_callable_mapper(loopy_merged, solve_fn_lookup)
+    name = "slate_wrapper"
+    loopy_merged = merge_loopy(slate_loopy, output_arg, builder, var2terminal, name)
+    loopy_merged = loopy.register_callable(loopy_merged, INVCallable.name, INVCallable())
+    loopy_merged = loopy.register_callable(loopy_merged, SolveCallable.name, SolveCallable())
 
-    # WORKAROUND: Generate code directly from the loopy kernel here,
-    # then attach code as a c-string to the op2kernel
-    code = loopy.generate_code_v2(loopy_merged).device_code()
-    code = code.replace(f'void {loopy_merged.name}', f'static void {loopy_merged.name}')
-    loopykernel = op2.Kernel(code,
-                             loopy_merged.name,
+    loopykernel = op2.Kernel(loopy_merged,
+                             name,
                              include_dirs=BLASLAPACK_INCLUDE.split(),
                              ldargs=BLASLAPACK_LIB.split())
 
@@ -630,7 +627,7 @@ def gem_to_loopy(gem_expr, var2terminal, scalar_type):
     shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1,)
     idx = make_indices(len(shape))
     indexed_gem_expr = gem.Indexed(gem_expr, idx)
-    args = ([loopy.GlobalArg("output", shape=shape, dtype=scalar_type)]
+    args = ([loopy.GlobalArg("output", shape=shape, dtype=scalar_type, is_output=True, is_input=True)]
             + [loopy.GlobalArg(var.name, shape=var.shape, dtype=scalar_type)
                for var in var2terminal.keys()])
     ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)]

firedrake/slate/slac/kernel_builder.py

@@ -512,9 +512,9 @@ def layer_integral_predicates(self, tensor, integral_type):
 
         # TODO: Variable layers
         nlayer = pym.Variable(self.layer_count)
-        which = {"interior_facet_horiz_top": pym.Comparison(layer, "<", nlayer),
+        which = {"interior_facet_horiz_top": pym.Comparison(layer, "<", nlayer[0]),
                  "interior_facet_horiz_bottom": pym.Comparison(layer, ">", 0),
-                 "exterior_facet_top": pym.Comparison(layer, "==", nlayer),
+                 "exterior_facet_top": pym.Comparison(layer, "==", nlayer[0]),
                  "exterior_facet_bottom": pym.Comparison(layer, "==", 0)}[integral_type]
 
         return [which]
@@ -633,9 +633,12 @@ def initialise_terminals(self, var2terminal, coefficients):
 
         return inits, tensor2temp
 
-    def slate_call(self, kernel, temporaries):
+    def slate_call(self, prg, temporaries):
+        name, = prg.callables_table.keys()
+        kernel = prg.callables_table[name].subkernel
+        output_var = pym.Variable(kernel.args[0].name)
         # Slate kernel call
-        reads = []
+        reads = [output_var]
         for t in temporaries:
             shape = t.shape
             name = t.name
@@ -652,12 +655,17 @@ def generate_wrapper_kernel_args(self, tensor2temp, templated_subkernels):
         args = [loopy.GlobalArg(self.coordinates_arg, shape=coords_extent,
                                 dtype=self.tsfc_parameters["scalar_type"])]
 
-        for loopy_inner in templated_subkernels:
-            for arg in loopy_inner.args[1:]:
-                if arg.name == self.cell_orientations_arg or\
-                   arg.name == self.cell_size_arg:
-                    if arg not in args:
-                        args.append(arg)
+        if self.bag.needs_cell_orientations:
+            ori_extent = self.extent(self.expression.ufl_domain().cell_orientations())
+            args.append(loopy.GlobalArg(self.cell_orientations_arg,
+                                        shape=ori_extent,
+                                        dtype=self.tsfc_parameters["scalar_type"]))
+
+        if self.bag.needs_cell_sizes:
+            siz_extent = self.extent(self.expression.ufl_domain().cell_sizes)
+            args.append(loopy.GlobalArg(self.cell_size_arg,
+                                        shape=siz_extent,
+                                        dtype=self.tsfc_parameters["scalar_type"]))
 
         for coeff in self.bag.coefficients.values():
             if isinstance(coeff, OrderedDict):
@@ -687,8 +695,7 @@ def generate_wrapper_kernel_args(self, tensor2temp, templated_subkernels):
         if self.bag.needs_mesh_layers:
             args.append(loopy.GlobalArg(self.layer_count, shape=(),
                         dtype=np.int32))
-            args.append(loopy.TemporaryVariable(self.layer_arg, shape=(),
-                        dtype=np.int32, address_space=loopy.AddressSpace.GLOBAL))
+            args.append(loopy.ValueArg(self.layer_arg, dtype=np.int32))
 
         for tensor_temp in tensor2temp.values():
             args.append(tensor_temp)
@@ -715,7 +722,9 @@ def generate_tsfc_calls(self, terminal, loopy_tensor):
                     raise ValueError("Integral type '%s' not recognized" % integral_type)
 
                 # Prepare lhs and args for call to tsfc kernel
-                output = self.generate_lhs(slate_tensor, pym.Variable(loopy_tensor.name))
+                output_var = pym.Variable(loopy_tensor.name)
+                reads.append(output_var)
+                output = self.generate_lhs(slate_tensor, output_var)
                 kernel_data = self.collect_tsfc_kernel_data(mesh, cxt_kernel.coefficients, self.bag.coefficients, kinfo)
                 reads.extend(self.loopify_tsfc_kernel_data(kernel_data))
 

firedrake/slate/slac/utils.py

@@ -14,8 +14,7 @@
 from functools import singledispatch
 import firedrake.slate.slate as sl
 import loopy as lp
-from loopy.program import make_program
-from loopy.transform.callable import register_callable_kernel
+from loopy.transform.callable import merge
 import itertools
 
 
@@ -282,7 +281,7 @@ def topological_sort(exprs):
     return schedule
 
 
-def merge_loopy(slate_loopy, output_arg, builder, var2terminal):
+def merge_loopy(slate_loopy, output_arg, builder, var2terminal, name):
     """ Merges tsfc loopy kernels and slate loopy kernel into a wrapper kernel."""
     from firedrake.slate.slac.kernel_builder import SlateWrapperBag
     coeffs = builder.collect_coefficients()
@@ -307,12 +306,24 @@ def merge_loopy(slate_loopy, output_arg, builder, var2terminal):
     domains = builder.bag.index_creator.domains
 
     # Generates the loopy wrapper kernel
-    slate_wrapper = lp.make_function(domains, insns, args, name="slate_wrapper",
-                                     seq_dependencies=True, target=lp.CTarget())
+    slate_wrapper = lp.make_function(domains, insns, args, name=name,
+                                     seq_dependencies=True, target=lp.CTarget(),
+                                     lang_version=(2018, 2))
 
     # Generate program from kernel, so that one can register kernels
-    prg = make_program(slate_wrapper)
+    from pyop2.codegen.loopycompat import _match_caller_callee_argument_dimension_
+    from loopy.kernel.function_interface import CallableKernel
+
     for tsfc_loopy in tsfc_kernels:
-        prg = register_callable_kernel(prg, tsfc_loopy)
-    prg = register_callable_kernel(prg, slate_loopy)
-    return prg
+        slate_wrapper = merge([slate_wrapper, tsfc_loopy])
+        names = tsfc_loopy.callables_table
+        for name in names:
+            if isinstance(slate_wrapper.callables_table[name], CallableKernel):
+                slate_wrapper = _match_caller_callee_argument_dimension_(slate_wrapper, name)
+    slate_wrapper = merge([slate_wrapper, slate_loopy])
+    names = slate_loopy.callables_table
+    for name in names:
+        if isinstance(slate_wrapper.callables_table[name], CallableKernel):
+            slate_wrapper = _match_caller_callee_argument_dimension_(slate_wrapper, name)
+
+    return slate_wrapper