firedrake.slate.slac package

Submodules

firedrake.slate.slac.compiler module

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 COFFEE’s kernel abstract syntax tree (AST) optimizer. TSFC provides this compiler with appropriate kernel functions (in C) for evaluating integral expressions (finite element variational forms written in UFL). COFFEE’s AST base helps with the construction of code blocks throughout the kernel returned by: compile_expression.

The Eigen C++ library (http://eigen.tuxfamily.org/) is required, as all low-level numerical linear algebra operations are performed using this templated function library.

firedrake.slate.slac.compiler.compile_expression(slate_expr, compiler_parameters=None, coffee=False)[source]

Takes a Slate expression slate_expr and returns the appropriate firedrake.op2.Kernel object representing the Slate expression.

Parameters:
  • slate_expr – a :class:’TensorBase’ expression.

  • 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)

firedrake.slate.slac.kernel_builder module

class firedrake.slate.slac.kernel_builder.CellFacetKernelArg(ast_arg)[source]

Bases: KernelArg

class firedrake.slate.slac.kernel_builder.CoefficientInfo(space_index, offset_index, shape, vector, local_temp)

Bases: tuple

Context information for creating coefficient temporaries.

Parameters:
  • space_index – An integer denoting the function space index.

  • offset_index – An integer denoting the starting position in the vector temporary for assignment.

  • shape – A singleton with an integer describing the shape of the coefficient temporary.

  • vector – The slate.AssembledVector containing the relevant data to be placed into the temporary.

  • local_temp – The local temporary for the coefficient vector.

Create new instance of CoefficientInfo(space_index, offset_index, shape, vector, local_temp)

local_temp

Alias for field number 4

offset_index

Alias for field number 1

shape

Alias for field number 2

space_index

Alias for field number 0

vector

Alias for field number 3

class firedrake.slate.slac.kernel_builder.IndexCreator[source]

Bases: object

property domains

ISL domains for the currently known indices.

class firedrake.slate.slac.kernel_builder.LayerCountKernelArg(ast_arg)[source]

Bases: KernelArg

class firedrake.slate.slac.kernel_builder.LocalKernelBuilder(expression, tsfc_parameters=None)[source]

Bases: object

The primary helper class for constructing cell-local linear algebra kernels from Slate expressions.

This class provides access to all temporaries and subkernels associated with a Slate expression. If the Slate expression contains nodes that require operations on already assembled data (such as the action of a slate tensor on a ufl.Coefficient), this class provides access to the expression which needs special handling.

Instructions for assembling the full kernel AST of a Slate expression is provided by the method construct_ast.

Constructor for the LocalKernelBuilder class.

Parameters:
  • expression – a TensorBase object.

  • tsfc_parameters – an optional dict of parameters to provide to TSFC when constructing subkernels associated with the expression.

cell_facet_sym = <coffee.base.Symbol object>
cell_orientations_sym = <coffee.base.Symbol object>
cell_size_sym = <coffee.base.Symbol object>
coefficient(coefficient)[source]

Extracts the kernel arguments corresponding to a particular coefficient. This handles both the case when the coefficient is defined on a mixed or non-mixed function space.

coefficient_map[source]

Generates a mapping from a coefficient to its kernel argument symbol. If the coefficient is mixed, all of its split components will be returned.

context_kernels[source]

Gathers all ContextKernels containing all TSFC kernels, and integral type information.

coord_sym = <coffee.base.Symbol object>
expression_flops[source]
property integral_type

Returns the integral type associated with a Slate kernel.

it_sym = <coffee.base.Symbol object>
mesh_layer_count_sym = <coffee.base.Symbol object>
mesh_layer_sym = <coffee.base.Symbol object>
needs_cell_facets[source]

Searches for any embedded forms (by inspecting the ContextKernels) which require looping over cell facets. If any are found, this function returns True and False otherwise.

needs_mesh_layers[source]

Searches for any embedded forms (by inspecting the ContextKernels) which require mesh level information (extrusion measures). If any are found, this function returns True and False otherwise.

supported_integral_types = ['cell', 'interior_facet', 'exterior_facet', 'interior_facet_horiz_top', 'interior_facet_horiz_bottom', 'interior_facet_vert', 'exterior_facet_top', 'exterior_facet_bottom', 'exterior_facet_vert']
supported_subdomain_types = ['subdomains_exterior_facet', 'subdomains_interior_facet']
terminal_flops[source]
class firedrake.slate.slac.kernel_builder.LocalLoopyKernelBuilder(expression, tsfc_parameters=None)[source]

Bases: object

Constructor for the LocalGEMKernelBuilder class.

Parameters:
  • expression – a TensorBase object.

  • tsfc_parameters – an optional dict of parameters to provide to TSFC when constructing subkernels associated with the expression.

cell_facets_arg_name = 'cell_facets'
cell_orientations_arg_name = 'cell_orientations'
cell_sizes_arg_name = 'cell_sizes'
collect_coefficients()[source]

Saves all coefficients of self.expression where non-mixed coefficients are dicts of form {coeff: (name, extent)} and mixed coefficients are double dicts of form {mixed_coeff: {coeff_per_space: (name, extent)}}.

collect_tsfc_kernel_data(mesh, tsfc_coefficients, wrapper_coefficients, kinfo)[source]

Collect the kernel data aka the parameters fed into the subkernel, that are coordinates, orientations, cell sizes and cofficients.

coordinates_arg_name = 'coords'
extent(coefficient)[source]

Calculation of the range of a coefficient.

facet_integral_predicates(mesh, integral_type, kinfo)[source]
generate_lhs(tensor, temp)[source]

Generation of an lhs for the loopy kernel, which contains the TSFC assembly of the tensor.

generate_tsfc_calls(terminal, loopy_tensor)[source]

A setup method to initialize all the local assembly kernels generated by TSFC. This function also collects any information regarding orientations and extra include directories.

generate_wrapper_kernel_args(tensor2temp)[source]
initialise_terminals(var2terminal, coefficients)[source]

Initilisation of the variables in which coefficients and the Tensors coming from TSFC are saved.

Parameters:

var2terminal – dictionary that maps Slate Tensors to gem Variables

is_integral_type(integral_type, type)[source]
layer_arg_name = 'layer'
layer_count_name = 'layer_count'
layer_integral_predicates(tensor, integral_type)[source]
local_facet_array_arg_name = 'facet_array'
loopify_tsfc_kernel_data(kernel_data)[source]

This method generates loopy arguments from the kernel data, which are then fed to the TSFC loopy kernel. The arguments are arrays and have to be fed element by element to loopy aka they have to be subarrayrefed.

shape(tensor)[source]

A helper method to retrieve tensor shape information. In particular needed for the right shape of scalar tensors.

slate_call(prg, temporaries)[source]
supported_integral_types = ['cell', 'interior_facet', 'exterior_facet', 'interior_facet_horiz_top', 'interior_facet_horiz_bottom', 'interior_facet_vert', 'exterior_facet_top', 'exterior_facet_bottom', 'exterior_facet_vert']
supported_subdomain_types = ['subdomains_exterior_facet', 'subdomains_interior_facet']
tsfc_cxt_kernels(terminal)[source]

Gathers all ContextKernels containing all TSFC kernels, and integral type information.

class firedrake.slate.slac.kernel_builder.SlateWrapperBag(coeffs)[source]

Bases: object

firedrake.slate.slac.optimise module

class firedrake.slate.slac.optimise.ActionBag(coeff, pick_op)

Bases: tuple

Create new instance of ActionBag(coeff, pick_op)

coeff

Alias for field number 0

pick_op

Alias for field number 1

firedrake.slate.slac.optimise.drop_double_transpose(expr)[source]

Remove double transposes from optimised Slate expression.

firedrake.slate.slac.optimise.flip(pick_op)[source]

Flip an index. Using this function essentially reverses the order of multiplication.

firedrake.slate.slac.optimise.optimise(expression, parameters)[source]

Optimises a Slate expression, by pushing blocks and multiplications inside the expression and by removing double transposes.

Parameters:
  • expression – A (potentially unoptimised) Slate expression.

  • parameters – A dict of compiler parameters.

Returns: An optimised Slate expression

firedrake.slate.slac.optimise.push_block(expression)[source]

Executes a Slate compiler optimisation pass. The optimisation is achieved by pushing blocks from the outside to the inside of an expression. Without the optimisation the local TSFC kernels are assembled first and then the result of the assembly kernel gets indexed in the Slate kernel (and further linear algebra operations maybe done on it). The optimisation pass essentially changes the order of assembly and indexing.

Parameters:

expression – A (potentially unoptimised) Slate expression.

Returns: An optimised Slate expression, where Blocks are terminal whereever possible.

firedrake.slate.slac.optimise.push_diag(expression)[source]

Executes a Slate compiler optimisation pass. The optimisation is achieved by pushing DiagonalTensor from the outside to the inside of an expression.

Parameters:

expression – A (potentially unoptimised) Slate expression.

Returns: An optimised Slate expression, where DiagonalTensors are sitting on terminal tensors whereever possible.

firedrake.slate.slac.optimise.push_mul(tensor, options)[source]

Executes a Slate compiler optimisation pass. The optimisation is achieved by pushing coefficients from the outside to the inside of an expression. The optimisation pass essentially changes the order of operations in the expressions so that only matrix-vector products are executed.

Parameters:
  • tensor – A (potentially unoptimised) Slate expression.

  • options – Optimisation pass options, e.g. if the multiplication should be replaced by an action.

Returns: An optimised Slate expression,

where only matrix-vector products are executed whereever possible.

firedrake.slate.slac.tsfc_driver module

class firedrake.slate.slac.tsfc_driver.ContextKernel(tensor, coefficients, original_integral_type, tsfc_kernels)

Bases: tuple

A bundled object containing TSFC subkernels corresponding to a particular integral type.

Parameters:
  • tensor – The terminal Slate tensor corresponding to the list of TSFC assembly kernels.

  • coefficients – The local coefficients of the tensor contained in the integrands (arguments for TSFC subkernels).

  • original_integral_type – The unmodified measure type of the form integrals.

  • tsfc_kernels – A list of local tensor assembly kernels provided by TSFC.

Create new instance of ContextKernel(tensor, coefficients, original_integral_type, tsfc_kernels)

coefficients

Alias for field number 1

original_integral_type

Alias for field number 2

tensor

Alias for field number 0

tsfc_kernels

Alias for field number 3

firedrake.slate.slac.tsfc_driver.compile_terminal_form(tensor, prefix, *, tsfc_parameters=None, coffee=True)[source]

Compiles the TSFC form associated with a Slate Tensor object. This function will return a ContextKernel which stores information about the original tensor, integral types and the corresponding TSFC kernels.

Parameters:
  • tensor – A Slate Tensor.

  • prefix – An optional string indicating the prefix for the subkernel.

  • tsfc_parameters – An optional dict of parameters to provide TSFC.

Returns: A ContextKernel containing all relevant information.

firedrake.slate.slac.tsfc_driver.transform_integrals(integrals)[source]

Generates a mapping of the form:

{original_integral_type: transformed_integrals}

where the original_integral_type is the pre-transformed integral type. The transformed_integrals are an iterable of ufl.Integral`s with the appropriately modified type. For example, an `interior_facet integral will become an exterior_facet integral.

firedrake.slate.slac.utils module

class firedrake.slate.slac.utils.RemoveRestrictions[source]

Bases: MultiFunction

UFL MultiFunction for removing any restrictions on the integrals of forms.

expr(o, *ops)

Reuse object if operands are the same objects.

Use in your own subclass by setting e.g.

expr = MultiFunction.reuse_if_untouched

as a default rule.

positive_restricted(o)[source]
class firedrake.slate.slac.utils.SymbolWithFuncallIndexing(symbol, rank=None, offset=None)[source]

Bases: Symbol

A functionally equivalent representation of a coffee.Symbol, with modified output for rank calls. This is syntactically necessary when referring to symbols of Eigen::MatrixBase objects.

class firedrake.slate.slac.utils.Transformer[source]

Bases: Visitor

Replaces all out-put tensor references with a specified name of :type: Eigen::Matrix with appropriate shape. This class is primarily for COFFEE acrobatics, jumping through nodes and redefining where appropriate.

The default name of "A" is assigned, otherwise a specified name may be passed as the name keyword argument when calling the visitor.

visit_Decl(o, *args, **kwargs)[source]

Visits a declared tensor and changes its type to :template: result Eigen::MatrixBase<Derived>.

i.e. double A[n][m] —> const Eigen::MatrixBase<Derived> &A_

visit_FunDecl(o, *args, **kwargs)[source]

Visits a COFFEE FunDecl object and reconstructs the FunDecl body and header to generate Eigen::MatrixBase C++ template functions.

Creates a template function for each subkernel form.

template <typename Derived>
static inline void foo(Eigen::MatrixBase<Derived> const & A, ...)
{
  [Body...]
}
visit_Node(o, *args, **kwargs)

A visit method that reconstructs nodes if their children have changed.

visit_Symbol(o, *args, **kwargs)[source]

Visits a COFFEE symbol and redefines it as a Symbol with FunCall indexing.

i.e. A[j][k] —> A(j, k)

visit_list(o, *args, **kwargs)[source]

Visits an input of COFFEE objects and returns the complete list of said objects.

visit_object(o, *args, **kwargs)[source]

Visits an object and returns it.

e.g. string —> string

A recursive depth-first search (DFS) algorithm for traversing a DAG consisting of Slate expressions.

Parameters:
  • graph – A DAG whose nodes (vertices) are Slate expressions with edges connected to dependent expressions.

  • node – A starting vertex.

  • visited – A set keeping track of visited nodes.

  • schedule – A list of reverse-postordered nodes. This list is used to produce a topologically sorted list of Slate nodes.

firedrake.slate.slac.utils.merge_loopy(slate_loopy, output_arg, builder, var2terminal, name)[source]

Merges tsfc loopy kernels and slate loopy kernel into a wrapper kernel.

firedrake.slate.slac.utils.slate2gem(expression, options)[source]
firedrake.slate.slac.utils.slate_to_gem(expression, options)[source]

Convert a slate expression to gem.

Parameters:

expression – A slate expression.

Returns:

A singleton list of gem expressions and a mapping from gem variables to UFL “terminal” forms.

firedrake.slate.slac.utils.topological_sort(exprs)[source]

Topologically sorts a list of Slate expressions. The expression graph is constructed by relating each Slate node with a list of dependent Slate nodes.

Parameters:

exprs – A list of Slate expressions.

Module contents