10-sum-factorisation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import firedrake\n",
    "except ImportError:\n",
    "    !wget \"https://fem-on-colab.github.io/releases/firedrake-install-real.sh\" -O \"/tmp/firedrake-install.sh\" && bash \"/tmp/firedrake-install.sh\"\n",
    "    import firedrake"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimal complexity finite element assembly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sum factorisation\n",
    "Suppose we are assemblying the 2-form of the Laplace operator on a hexahedral element, e.g. the Q3 element:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<figure>\n",
    "  <center><img src=\"image/Q3_hexahedron.png\" alt=\"Q3\" style=\"width: 400px;\"/></center>\n",
    "  <center><figcaption>$Q_3$ element on hexahedron. Image from <a href=\"http://femtable.org/\">the periodic table.</a></figcaption></center>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since each quadrature point $X$ can be represented as a 3-tuple $q=\\{q_1,q_2,q_3\\}$, and each basis function by a 3-tuple $i = \\{i_1,i_2,i_3\\}$, the naive local assembly kernel for the local tensor $A_{ij}$ contains the loop structure:\n",
    "```\n",
    "for q1, q2, q3\n",
    "  for i1, i2, i3\n",
    "    for j1, j2, j3\n",
    "      A[i1,i2,i3,j1,j2,j3] += ...\n",
    "```\n",
    "This requires $O(N_{q}^3N_{i}^6)$ FLOPs. For polynomial degree $p$, both $N_q$ and $N_i$ are $O(p)$, so this local assembly requires $O(p^9)$ FLOPs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For *tensor product elements* like this, we can rearrange the contraction over quadrature points and hoist invariant sub-expressions out of the innermost loop into temporary variables. This is known as *sum factorisation*:\n",
    "```\n",
    "for q1, i1, j1\n",
    "  t1[i1,j1] += ...\n",
    "for q2, i2, j2\n",
    "  t2[i2,j2] += ...\n",
    "for q3\n",
    "  for i1, i2, i3\n",
    "    for j1, j2, j3\n",
    "      A[i1,i2,i3,j1,j2,j3] += t1*t2*...\n",
    "```\n",
    "This reduces the complexity to $O(p^7)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TSFC \\[1\\], the form compiler of Firedrake, is capable of exploiting this intrinsic structure of the finite element, provided by FInAT \\[2\\], and apply sum factorisation automatically to generate assembly kernels with optimal algorithmic complexity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_format = 'svg'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (11, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from firedrake import *\n",
    "set_log_level(ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can create a hexahedral mesh by extruding a quadrilateral mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = ExtrudedMesh(UnitSquareMesh(10, 10, quadrilateral=True), 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's choose the continuous Lagrange element of degree 5 as our function space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = 5\n",
    "V = FunctionSpace(mesh, \"CG\", p)\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)\n",
    "a = dot(grad(u), grad(v)) *dx  # Laplace operator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Firedrake internalises the process of local assembly. In order to look at the kernel, we need to import the compilation interface from TSFC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tsfc import compile_form"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TSFC organises the optimisation passes into *modes*. Let's first try the *vanilla* mode, which does as little optimisation as possible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kernel_vanilla, = compile_form(a, parameters={\"mode\": \"vanilla\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TSFC also lets us estimate the number of FLOPs performed by the kernel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Local assembly FLOPs with vanilla mode is {0:.3g}\".format(kernel_vanilla.flop_count))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The default optimisation mode in TSFC is *spectral*, which applies sum factorisation to determine the tensor contraction order, and at each level, apply *argument factorisation* \\[3\\] to rearrange the expression using associative and distributive laws. Since *spectral* is the default mode, we do not need to specify it in the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kernel_spectral, = compile_form(a)\n",
    "print(\"Local assembly FLOPs with spectral mode is {0:.3g}\".format(kernel_spectral.flop_count))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a 43x reduction in FLOPs. Not bad, but there's opportunity to do better. For spectral elements, if we use the Gauss–Lobatto–Legendre (GLL) quadrature scheme, which has quadrature points collated with the Lagrange basis function nodes, then we know that the basis function tabulation is an indentity matrix. TSFC and FInAT can further simplify the loop structure of the local assembly kernels. This reduces the complexity to $O(p^5)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to specify the GLL quadrature scheme for hexahedra. We can do this with FIAT, which defines GLL on intervals, and FInAT, which makes the tensor product scheme from the interval scheme."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import FIAT, finat\n",
    "\n",
    "def gauss_lobatto_legendre_line_rule(degree):\n",
    "    fiat_make_rule = FIAT.quadrature.GaussLobattoLegendreQuadratureLineRule\n",
    "    fiat_rule = fiat_make_rule(FIAT.ufc_simplex(1), degree + 1)\n",
    "    finat_ps = finat.point_set.GaussLobattoLegendrePointSet\n",
    "    points = finat_ps(fiat_rule.get_points())\n",
    "    weights = fiat_rule.get_weights()\n",
    "    return finat.quadrature.QuadratureRule(points, weights)\n",
    "\n",
    "def gauss_lobatto_legendre_cube_rule(dimension, degree):\n",
    "    make_tensor_rule = finat.quadrature.TensorProductQuadratureRule\n",
    "    result = gauss_lobatto_legendre_line_rule(degree)\n",
    "    for _ in range(1, dimension):\n",
    "        line_rule = gauss_lobatto_legendre_line_rule(degree)\n",
    "        result = make_tensor_rule([result, line_rule])\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by creating the spectral finite element function space of the same polynomial degree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "element = FiniteElement('CG', mesh.ufl_cell(), degree=p, variant='spectral')\n",
    "V = FunctionSpace(mesh, element)\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to tell Firedrake to use the GLL quadratures for numerical integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gll_quadrature_rule = gauss_lobatto_legendre_cube_rule(dimension=3, degree=p)\n",
    "a_gll = dot(grad(u), grad(v)) *dx(scheme=gll_quadrature_rule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kernel_gll, = compile_form(a_gll)\n",
    "print(\"Local assembly FLOPs with GLL quadrature is {0:.3g}\".format(kernel_gll.flop_count))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a further 10x reduction in FLOPs.\n",
    "\n",
    "Now, let's verify that we achieve the expected asymptotic algorithmic complexity with respect to polynomial degrees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flops = defaultdict(list)\n",
    "ps = range(1, 33)  # polynomial degrees\n",
    "modes = {\n",
    "    'gll': {'mode': 'spectral', 'variant': 'spectral', 'rule': gauss_lobatto_legendre_cube_rule},\n",
    "    'spectral': {'mode': 'spectral', 'variant': None, 'rule': lambda *args: None},\n",
    "    'vanilla': {'mode': 'vanilla', 'variant': None, 'rule': lambda *args: None}\n",
    "}\n",
    "\n",
    "for p in ps:\n",
    "    for mode in modes:\n",
    "        element = FiniteElement('CG', mesh.ufl_cell(), degree=p, variant=modes[mode]['variant'])\n",
    "        V = FunctionSpace(mesh, element)\n",
    "        u = TrialFunction(V)\n",
    "        v = TestFunction(V)\n",
    "        a = dot(grad(u), grad(v))*dx(scheme=modes[mode]['rule'](3, p))\n",
    "        kernel, = compile_form(a, parameters={\"mode\": modes[mode]['mode']})\n",
    "        flops[mode].append(kernel.flop_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "for mode in modes:\n",
    "    ax.plot(ps, flops[mode], label=mode)\n",
    "x = numpy.linspace(1, 32, 100)\n",
    "for p, style, offset in zip([5,7,9], ['-.','--',':'], [10, 3, 5]):\n",
    "    ax.plot(x, numpy.power(x, p)*offset, label=r\"$p^{0}$\".format(p), color='grey', linestyle=style)\n",
    "ax.legend(loc='upper left');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's do the same analysis on a different problem:  $a (u,v) = \\int_\\Omega\\left( \\nabla\\times u \\right) \\cdot \\left( \\nabla\\times v \\right)\\,\\text{d}x$, on the H(curl) conforming NCE element."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This might take some time to run\n",
    "flops_curl = defaultdict(list)\n",
    "ps_curl = range(1, 17)\n",
    "for p in ps_curl:\n",
    "    for mode in modes:\n",
    "        element = FiniteElement('NCE', mesh.ufl_cell(), degree=p, variant=modes[mode]['variant'])\n",
    "        V = FunctionSpace(mesh, element)\n",
    "        u = TrialFunction(V)\n",
    "        v = TestFunction(V)\n",
    "        a = dot(curl(u), curl(v))*dx(scheme=modes[mode]['rule'](3, p))\n",
    "        kernel, = compile_form(a, parameters={\"mode\": modes[mode]['mode']})\n",
    "        flops_curl[mode].append(kernel.flop_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "for mode in modes:\n",
    "    ax.plot(ps_curl, flops_curl[mode], label=mode)\n",
    "x = numpy.linspace(1, 16, 100)\n",
    "for p, style, offset in zip([5,7,9], ['-.','--',':'], [800,40,60]):\n",
    "    ax.plot(x, numpy.power(x, p)*offset, label=r\"$p^{0}$\".format(p), color='grey', linestyle=style)\n",
    "ax.legend(loc='upper left');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\\[1\\] Homolya, M., Mitchell, L., Luporini, F. and Ham, D.A., 2018. TSFC: a structure-preserving form compiler. SIAM Journal on Scientific Computing, 40(3), pp.C401-C428.\n",
    "\n",
    "\\[2\\] Homolya, M., Kirby, R.C. and Ham, D.A., 2017. Exposing and exploiting structure: optimal code generation for high-order finite element methods. arXiv preprint arXiv:1711.02473.\n",
    "\n",
    "\\[3\\] Luporini, F., Ham, D.A. and Kelly, P.H., 2017. An algorithm for the optimization of finite element integration loops. ACM Transactions on Mathematical Software (TOMS), 44(1), p.3."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}