Source code for firedrake.slate.slac.compiler

"""This is Slate's Linear Algebra Compiler. This module is
responsible for generating C kernel functions representing
symbolic linear algebra expressions written in Slate.

This linear algebra compiler uses both Firedrake's form compiler,
the Two-Stage Form Compiler (TSFC) and loopy. TSFC provides this
compiler with appropriate kernel functions (in C) for evaluating integral
expressions (finite element variational forms written in UFL).
import time
from hashlib import md5

from firedrake_citations import Citations
from firedrake.tsfc_interface import SplitKernel, KernelInfo, TSFCKernel

from firedrake.slate.slac.kernel_builder import LocalLoopyKernelBuilder
from firedrake.slate.slac.utils import slate_to_gem, merge_loopy
from firedrake.slate.slac.optimise import optimise

from firedrake import tsfc_interface
from firedrake.logging import logger
from firedrake.parameters import parameters
from firedrake.petsc import get_petsc_variables
from firedrake.utils import complex_mode
from gem import impero_utils
from itertools import chain

from pyop2.utils import get_petsc_dir
from pyop2.mpi import COMM_WORLD
from pyop2.codegen.rep2loopy import SolveCallable, INVCallable

import firedrake.slate.slate as slate
import numpy as np
import loopy
import gem
from gem import indices as make_indices
from tsfc.kernel_args import OutputKernelArg, CoefficientKernelArg
from tsfc.loopy import generate as generate_loopy
import copy

from petsc4py import PETSc

__all__ = ['compile_expression']

GREEN = "\033[1;37;32m%s\033[0m"

    PETSC_DIR, PETSC_ARCH = get_petsc_dir()
except ValueError:
    PETSC_DIR, = get_petsc_dir()
    PETSC_ARCH = None

if COMM_WORLD.rank == 0:
    petsc_variables = get_petsc_variables()
    BLASLAPACK_LIB = petsc_variables.get("BLASLAPACK_LIB", "")
    BLASLAPACK_INCLUDE = petsc_variables.get("BLASLAPACK_INCLUDE", "")
    BLASLAPACK_LIB = COMM_WORLD.bcast(None, root=0)
    BLASLAPACK_INCLUDE = COMM_WORLD.bcast(None, root=0)

cell_to_facets_dtype = np.dtype(np.int8)

class SlateKernel(TSFCKernel):
    def _cache_key(cls, expr, compiler_parameters):
        return md5(
            (expr.expression_hash + str(sorted(compiler_parameters.items()))).encode()).hexdigest(), expr.ufl_domains()[0].comm

    def __init__(self, expr, compiler_parameters):
        if self._initialized:
        self.split_kernel = generate_loopy_kernel(expr, compiler_parameters)
        self._initialized = True

[docs] def compile_expression(slate_expr, compiler_parameters=None): """Takes a Slate expression `slate_expr` and returns the appropriate ``pyop2.op2.Kernel`` object representing the Slate expression. :arg slate_expr: a :class:`~.Tensor` expression. :arg tsfc_parameters: an optional `dict` of form compiler parameters to be passed to TSFC during the compilation of ufl forms. Returns: A ``tuple`` containing a ``SplitKernel(idx, kinfo)`` """ if complex_mode: raise NotImplementedError("SLATE doesn't work in complex mode yet") if not isinstance(slate_expr, slate.TensorBase): raise ValueError("Expecting a `TensorBase` object, not %s" % type(slate_expr)) # Update default parameters with passed parameters # The deepcopy is needed because parameters is a nested dict params = copy.deepcopy(parameters) if compiler_parameters and "slate_compiler" in compiler_parameters.keys(): params["slate_compiler"].update(compiler_parameters.pop("slate_compiler")) if compiler_parameters: params["form_compiler"].update(compiler_parameters) # If the expression has already been symbolically compiled, then # simply reuse the produced kernel. cache = slate_expr._metakernel_cache key = str(sorted(params.items())) try: return cache[key] except KeyError: kernel = SlateKernel(slate_expr, params).split_kernel return cache.setdefault(key, kernel)
def get_temp_info(loopy_kernel): """Get information about temporaries in loopy kernel. Returns memory in bytes and number of temporaries. """ mems = [temp.nbytes for temp in loopy_kernel.temporary_variables.values()] mem_total = sum(mems) num_temps = len(loopy_kernel.temporary_variables) # Get number of temporaries of different shapes shapes = {} for temp in loopy_kernel.temporary_variables.values(): shape = temp.shape if temp.storage_shape is not None: shape = temp.storage_shape shapes[len(shape)] = shapes.get(len(shape), 0) + 1 return mem_total, num_temps, mems, shapes def generate_loopy_kernel(slate_expr, compiler_parameters=None): cpu_time = time.time() if len(slate_expr.ufl_domains()) > 1: raise NotImplementedError("Multiple domains not implemented.") Citations().register("Gibson2018") orig_expr = slate_expr # Optimise slate expr, e.g. push blocks as far inward as possible if compiler_parameters["slate_compiler"]["optimise"]: slate_expr = optimise(slate_expr, compiler_parameters["slate_compiler"]) # Create a loopy builder for the Slate expression, # e.g. contains the loopy kernels coming from TSFC gem_expr, var2terminal = slate_to_gem(slate_expr, compiler_parameters["slate_compiler"]) scalar_type = compiler_parameters["form_compiler"]["scalar_type"] (slate_loopy, slate_loopy_event), output_arg = gem_to_loopy(gem_expr, var2terminal, scalar_type) builder = LocalLoopyKernelBuilder(expression=slate_expr, tsfc_parameters=compiler_parameters["form_compiler"]) name = "slate_wrapper" loopy_merged, arguments, events = merge_loopy(slate_loopy, output_arg, builder, var2terminal, name) loopy_merged = loopy.register_callable(loopy_merged,, INVCallable()) loopy_merged = loopy.register_callable(loopy_merged,, SolveCallable()) loopykernel = tsfc_interface.as_pyop2_local_kernel(loopy_merged, name, len(arguments), include_dirs=BLASLAPACK_INCLUDE.split(), ldargs=BLASLAPACK_LIB.split(), events=events+(slate_loopy_event,)) # map the coefficients in the order that PyOP2 needs orig_coeffs = orig_expr.coefficients() new_coeffs = slate_expr.coefficients() map_new_to_orig = [orig_coeffs.index(c) for c in new_coeffs] coefficient_numbers = tuple((map_new_to_orig[n], split_map) for (n, split_map) in slate_expr.coeff_map) coefficients = list(filter(lambda elm: isinstance(elm, CoefficientKernelArg), arguments)) # do the same for constants orig_constants = orig_expr.constants() new_constants = slate_expr.constants() constant_numbers = tuple(orig_constants.index(c) for c in new_constants) assert len(list(chain(*(map[1] for map in coefficient_numbers)))) == len(coefficients), \ "KernelInfo must be generated with a coefficient map that maps EXACTLY all coefficients that are in its arguments attribute." assert len(loopy_merged.callables_table[name].subkernel.args) - int(builder.bag.needs_mesh_layers) == len(arguments), \ "Outer loopy kernel must have the same amount of args as there are in arguments" kinfo = KernelInfo(kernel=loopykernel, integral_type="cell", # slate can only do things as contributions to the cell integrals oriented=builder.bag.needs_cell_orientations, subdomain_id=("otherwise",), domain_number=0, coefficient_numbers=coefficient_numbers, constant_numbers=constant_numbers, needs_cell_facets=builder.bag.needs_cell_facets, pass_layer_arg=builder.bag.needs_mesh_layers, needs_cell_sizes=builder.bag.needs_cell_sizes, arguments=arguments, events=events) # Cache the resulting kernel # Slate kernels are never split, so indicate that with None in the index slot. idx = tuple([None]*slate_expr.rank) % "compile_slate_expression finished in %g seconds.", time.time() - cpu_time) return (SplitKernel(idx, kinfo),) def gem_to_loopy(gem_expr, var2terminal, scalar_type): """ Method encapsulating stage 2. Converts the gem expression dag into imperoc first, and then further into loopy. :return slate_loopy: 2-tuple of loopy kernel for slate operations and loopy GlobalArg for the output variable. """ # Creation of return variables for outer loopy 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) output_loopy_arg = loopy.GlobalArg("output", shape=shape, dtype=scalar_type, is_input=True, is_output=True) args = [output_loopy_arg] + [loopy.GlobalArg(, shape=var.shape, dtype=scalar_type) for var in var2terminal.keys()] ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)] preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr]) # glue assignments to return variable assignments = list(zip(ret_vars, preprocessed_gem_expr)) # Part A: slate to impero_c impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False) # Part B: impero_c to loopy output_arg = OutputKernelArg(output_loopy_arg) return generate_loopy(impero_c, args, scalar_type, "slate_loopy", [], log=PETSc.Log.isActive()), output_arg