Performance Portable Finite Element Computations in Fluidity with UFL, FFC and PyOP2

Florian Rathgeber¹, Graham Markall¹, Lawrence Mitchell³, Nicolas Loriant¹, Gheorghe-teodor Bercea¹, David Ham^1,2, Paul Kelly¹

¹ Department of Computing, Imperial College London

² Grantham Institute for Climate Change, Imperial College London

³ EPCC, University of Edinburgh

I will wrap up the PyOP2 trio today and show you how to do finite-element computations in the PyOP2 parallel universe and benefit from the performance portability Graham demonstrated earlier.

The PyOP2-Fluidity Tool Chain

• PyOP2: clever way of efficiently executing kernels in parallel on unstructured meshes
• could write FEM kernels by hand and feed in mesh data in PyOP2 Dats -> of course we won't
• need for FEM: tool to generate kernels, tool to provide the mesh and data over the mesh
• this diagram shows the tool chain we provide for this: the PyOP2 part you've already seen in Graham's talk
• UFL and FFC provide the one part and we use the Fluidity library I'll introduce later for the meshes

Separation of Concerns

• This allows us to deliver separation of concerns or rather expand on what FEniCS already provides by drawing in PyOP2 as another layer of abstraction below the UFL layer
• Separation of concerns (CS buzzword): allow scientist to work at the level of their expertise whithout having to worry about the rest:
  • FEM model developer
  • Numerical analyst
  • Parallel programming expert
• freely mix and match UFL and PyOP2 (demonstrated later)

Comparison to the FEniCS Tool Chain

How do we compare to FEniCS:

• UFL/FFC exactly the same
• DOLFIN provides geometry and meshes etc.
• UFC interface not quite comparable to PyOP2 (tailored to CPU computation etc.)
• PyOP2 supports many more hardware architectures -> transparently exploitable for FEM

Interfacing PyOP2 to Fluidity

Fluidity

• open source, general purpose, multi-phase computational fluid dynamics code
• used internationally for complex fluid tasks
• developed at AMCG at Imperial College
• XML-based configuration files with GUI editor

Interfacing PyOP2

• existing interface to access fields from Python
• additional equation type UFL alongside Fluidity's built-in equations
• user provides custom UFL code
• call PyOP2 instead of Fluidity's built-in advection-diffusion solver
• create PyOP2 data structures for accessed fields on the fly

Why we chose Fluidity:

• established, existing user group
• separation of model users and model developers
• design: python data structures "all the way down"

UFL equations in Diamond

Driving PyOP2 from UFL source

Solving the advection-diffusion equation $$\frac{\partial c}{\partial t} + \nabla \cdot (\vec{u} c) = \nabla \cdot (\kappa \nabla c) + F$$

t=state.scalar_fields["Tracer"]   # Coefficient(FiniteElement("CG", "triangle", 1))
u=state.vector_fields["Velocity"] # Coefficient(VectorElement("CG", "triangle", 1))

p=TrialFunction(t)
q=TestFunction(t)

diffusivity = 0.1

M = p * q * dx

d = dt * (diffusivity * dot(grad(q), grad(p)) - dot(grad(q), u) * p) * dx

a = M + 0.5 * d
L = action(M - 0.5 * d, t)

solve(a == L, t)

How does a user drive the Fluidity-PyOP2 tool chain? They write their model in UFL! This slide shows all the code the Fluidity user needs to write to solve an advection-diffusion equation.

Small change compared to DOLFIN: the fields come from Fluidity exposed via its internal state: u and t are both UFL coefficient and Fluidity scalar/vector fields:

• can be used as coefficients in a form
• wrap a PyOP2 Dat

Memorize the solve call on the last line...

Finite element assembly and solve in PyOP2

What goes on behind the scenes of the solve call (simplified example!):

from pyop2 import op2, ffc_interface

def solve(equation, x):
    # Generate kernels for matrix and rhs assembly
    lhs = ffc_interface.compile_form(equation.lhs, "lhs")[0]
    rhs = ffc_interface.compile_form(equation.rhs, "rhs")[0]

    # Omitted: extract coordinates (coords), connectivity (elem_node)
    # and coefficients (tracer t, velocity u)

    # Construct OP2 matrix to assemble into
    sparsity = op2.Sparsity((elem_node, elem_node), sparsity_dim) 
    mat = op2.Mat(sparsity, numpy.float64)
    b = op2.Dat(nodes, 1, np.zeros(nodes.size))

    # Assemble lhs, rhs and solve linear system
    op2.par_loop(lhs, elements(3,3),
             mat((elem_node[op2.i[0]], elem_node[op2.i[1]]), op2.INC),
             coords(elem_node, op2.READ))
    op2.par_loop(rhs, elements(3),
             b(elem_node[op2.i[0]], op2.INC),
             coords(elem_node, op2.READ),
             t(elem_node, op2.READ),
             u(elem_node, op2.READ))

    op2.solve(mat, x, b)

What goes on "under the hood" when you call the solve method from the previous slide?

This is a (strongly) simpliefied implementation specific to the example PDE from the previous slide, showing what needs to happen in PyOP2.

Using PyOP2 for non-FEM kernels

PyOP2: performance portability for any unstructured mesh computations, not limited to FEM!

Use PyOP2 kernel for re-normalising a vector field

vec_norm_code="""
void vec_norm(double *u)
{
  const double n = sqrt(u[0]*u[0]+u[1]*u[1]);
  u[0] /= n;
  u[1] /= n;
}
"""

vec_norm = op2.Kernel(vec_norm_code, "vec_norm")

op2.par_loop(vec_norm, nodes,
             u(op2.IdentityMap, op2.RW))

Speedup relative to single core Fluidity (single 12-core node)

Speedup relative to single core Fluidity (4 12-core nodes)

Conclusions & future work

Conclusions

• Two-layer abstraction for FEM computation from UFL sources
• Decoupling of UFL (FEM) and PyOP2 (parallelisation) layers
• Performance portability for unstructured grid applications: FEM, non-FEM or combinations

Future Work

• Auto-tuning of optimisation parameters (e.g. iteration space)
• Support for curved elements
• Kernel fusion

Resources

• All the code is open source on GitHub and Launchpad
• Contact us: email mapdes@imperial.ac.uk

PyOP2

https://github.com/OP2/PyOP2

FFC

https://code.launchpad.net/~mapdes/ffc/pyop2

Fluidity

https://code.launchpad.net/~fluidity-core/fluidity/floppy_gn

Benchmarks

https://github.com/OP2/PyOP2_benchmarks

This talk

http://kynan.github.com/fenics13

Summary: UFL equations in Fluidity

For each UFL equation in each time step:

• Shell out to Python, execute the user's UFL equation
• FFC generates local assembly kernels for FE forms
• Backend-specific JIT-compilation of kernels and calling code
  • Instant for the sequential and OpenMP (including MPI)
  • PyCUDA for CUDA
  • PyOpenCL for OpenCL
• Agressive caching of forms, generated code and operators