Add numba backend

[schnellerhase]

Based on @chrisrichardson's work at chris/languages-3

Introduces an optional numba backend for FFCx. In particular the code passes through the same C-function interface of ufcx.h, and therefore the standard assemblers of DOLFINx are compatible.

Important

This does not introduce a (non-optional) dependency on numba - code generation is independent of compile.

Tested functionally on a simple poisson problem, and generation across all demo files.

Generalises some naming schemes to separate language specific implementation from functional concept.

[coveralls]

coverage: 84.44% (+1.4%) from 83.044%
when pulling af5078d on schnellerhase/copy-chris/languages-3
into 818267a on main.

[garth-wells]

Numba should remain an optional dependency. We have consciously not made it a required dependency across FEniCS packages.

Numba can constrain versions of other dependencies heavily, e.g. Numpy, and it often takes a long time to move to new a Python version.

[garth-wells]

@schnellerhase could you post a sample of the generated code?

My interest is that Numba code using the UFCx interface is normally poor from a performance point-of-view because Numba doesn't support static memory allocation. Performance can be (very) good by generating a whole assembler, or possibly good with the tricks to capture (and re-use) tables.

I like Numba as a prototype for supporting multiple languages, but I'm less keen on backends that we wouldn't recommend for efficient computation.

[schnellerhase]

Generated with ffcx --language numba PoissonQuad.py

Details

# This code conforms with the UFC specification version 2018.2.0.dev0
# and was automatically generated by FFCx version 0.11.0.dev0.
#
# This code was generated with the following options:
#
#  {'epsilon': 1e-14,
#   'language': 'numba',
#   'output_directory': '.',
#   'part': 'full',
#   'profile': False,
#   'scalar_type': 'float64',
#   'sum_factorization': False,
#   'table_atol': 1e-09,
#   'table_rtol': 1e-06,
#   'ufl_file': ['./PoissonQuad.py'],
#   'verbosity': 30,
#   'visualise': False}

import numba
import numpy as np
import math

# ufcx enums
interval = 10
triangle = 20
quadrilateral = 30
tetrahedron = 40
hexahedron = 50
vertex = 60
prism = 70
pyramid = 80

cell = 0
exterior_facet = 1
interior_facet = 2

ufcx_basix_element = 0
ufcx_mixed_element = 1
ufcx_quadrature_element = 2
ufcx_basix_custom_element = 3


# Code for integral integral_63fb73da9bd363692a1516abb30921c29f18b039_triangle

def tabulate_tensor_integral_63fb73da9bd363692a1516abb30921c29f18b039_triangle(_A, _w, _c, _coordinate_dofs,
                                   _entity_local_index, _quadrature_permutation, custom_data):

    A = numba.carray(_A, (36))
    w = numba.carray(_w, (0))
    c = numba.carray(_c, (0))
    coordinate_dofs = numba.carray(_coordinate_dofs, (36))
    entity_local_index = numba.carray(_entity_local_index, (2))
    quadrature_permutation = numba.carray(_quadrature_permutation, (0))
        # Quadrature rules 
    weights_39d = np.array([0.054975871827661, 0.054975871827661, 0.054975871827661, 0.1116907948390055, 0.1116907948390055, 0.1116907948390055], dtype=np.float64)
    # Precomputed values of basis functions and precomputations 
    # FE* dimensions: [permutation][entities][points][dofs] 
    FE0_C0_D01_Q39d = np.array([[[[0.6336951459609204, 0.0, -0.6336951459609156, 3.267390291921836, 0.0, -3.2673902919218363],
    [0.6336951459609199, 0.0, 2.2673902919218354, 0.36630485403908425, -2.901085437882756, -0.36630485403908397],
    [-2.267390291921832, 0.0, -0.6336951459609155, 0.3663048540390838, 2.9010854378827484, -0.3663048540390845],
    [-0.7837939636638601, 0.0, 0.78379396366386, 0.43241207267228005, 0.0, -0.43241207267228016],
    [-0.7837939636638598, 0.0, -0.5675879273277195, 1.7837939636638596, 1.35138189099158, -1.7837939636638604],
    [0.5675879273277201, 0.0, 0.7837939636638599, 1.78379396366386, -1.3513818909915805, -1.7837939636638602]]]], dtype=np.float64)
    FE1_C0_D10_Q39d = np.array([[[[0.6336951459609204, 2.267390291921836, 0.0, 0.36630485403908347, -0.36630485403908486, -2.9010854378827564],
    [0.63369514596092, -0.6336951459609156, 0.0, 3.2673902919218363, -3.2673902919218363, 0.0],
    [-2.2673902919218323, -0.6336951459609159, 0.0, 0.3663048540390847, -0.36630485403908375, 2.9010854378827475],
    [-0.78379396366386, -0.5675879273277199, 0.0, 1.7837939636638604, -1.7837939636638602, 1.3513818909915798],
    [-0.78379396366386, 0.7837939636638601, 0.0, 0.43241207267227993, -0.43241207267228016, 0.0],
    [0.5675879273277202, 0.7837939636638601, 0.0, 1.78379396366386, -1.7837939636638607, -1.3513818909915802]]]], dtype=np.float64)
    for iq in range(0, 6):
        # ------------------------ 
        # Section: Jacobian 
        # Inputs: FE0_C0_D01_Q39d, coordinate_dofs, FE1_C0_D10_Q39d 
        # Outputs: J0_c3, J0_c0, J0_c2, J0_c1 
        J0_c0 = 0.0
        J0_c3 = 0.0
        J0_c1 = 0.0
        J0_c2 = 0.0
        for ic in range(0, 6):
            J0_c0 += coordinate_dofs[(ic) * 3] * FE1_C0_D10_Q39d[0, 0, iq, ic]
            J0_c3 += coordinate_dofs[(ic) * 3 + 1] * FE0_C0_D01_Q39d[0, 0, iq, ic]
            J0_c1 += coordinate_dofs[(ic) * 3] * FE0_C0_D01_Q39d[0, 0, iq, ic]
            J0_c2 += coordinate_dofs[(ic) * 3 + 1] * FE1_C0_D10_Q39d[0, 0, iq, ic]
            
        # ------------------------ 
        # ------------------------ 
        # Section: Intermediates 
        # Inputs: J0_c3, J0_c0, J0_c2, J0_c1 
        # Outputs: fw0, fw1, fw2 
        fw0 = 0
        fw1 = 0
        fw2 = 0
        sv_39d_0 = J0_c0 * J0_c3
        sv_39d_1 = J0_c1 * J0_c2
        sv_39d_2 = -sv_39d_1
        sv_39d_3 = sv_39d_0 + sv_39d_2
        sv_39d_4 = J0_c0 / sv_39d_3
        sv_39d_5 = -J0_c1
        sv_39d_6 = sv_39d_5 / sv_39d_3
        sv_39d_7 = sv_39d_4 * sv_39d_4
        sv_39d_8 = sv_39d_4 * sv_39d_6
        sv_39d_9 = sv_39d_6 * sv_39d_6
        sv_39d_10 = J0_c3 / sv_39d_3
        sv_39d_11 = -J0_c2
        sv_39d_12 = sv_39d_11 / sv_39d_3
        sv_39d_13 = sv_39d_12 * sv_39d_12
        sv_39d_14 = sv_39d_10 * sv_39d_12
        sv_39d_15 = sv_39d_10 * sv_39d_10
        sv_39d_16 = sv_39d_7 + sv_39d_13
        sv_39d_17 = sv_39d_8 + sv_39d_14
        sv_39d_18 = sv_39d_15 + sv_39d_9
        sv_39d_19 = np.abs(sv_39d_3)
        sv_39d_20 = sv_39d_16 * sv_39d_19
        sv_39d_21 = sv_39d_17 * sv_39d_19
        sv_39d_22 = sv_39d_18 * sv_39d_19
        fw0 = sv_39d_22 * weights_39d[iq]
        fw1 = sv_39d_21 * weights_39d[iq]
        fw2 = sv_39d_20 * weights_39d[iq]
        # ------------------------ 
        # ------------------------ 
        # Section: Tensor Computation 
        # Inputs: fw0, FE1_C0_D10_Q39d, FE0_C0_D01_Q39d, fw1, fw2 
        # Outputs: A 
        temp_0 = np.full((6,), 0, dtype=np.float64)
        for j in range(0, 6):
            temp_0[j] = fw0 * FE1_C0_D10_Q39d[0, 0, iq, j]
            
        temp_1 = np.full((6,), 0, dtype=np.float64)
        for j in range(0, 6):
            temp_1[j] = fw1 * FE0_C0_D01_Q39d[0, 0, iq, j]
            
        temp_2 = np.full((6,), 0, dtype=np.float64)
        for j in range(0, 6):
            temp_2[j] = fw1 * FE1_C0_D10_Q39d[0, 0, iq, j]
            
        temp_3 = np.full((6,), 0, dtype=np.float64)
        for j in range(0, 6):
            temp_3[j] = fw2 * FE0_C0_D01_Q39d[0, 0, iq, j]
            
        for j in range(0, 6):
            for i in range(0, 6):
                A[6 * (i) + (j)] += FE1_C0_D10_Q39d[0, 0, iq, i] * temp_0[j]
                A[6 * (i) + (j)] += FE1_C0_D10_Q39d[0, 0, iq, i] * temp_1[j]
                A[6 * (i) + (j)] += FE0_C0_D01_Q39d[0, 0, iq, i] * temp_2[j]
                A[6 * (i) + (j)] += FE0_C0_D01_Q39d[0, 0, iq, i] * temp_3[j]
                
            
        # ------------------------ 
        
    

class integral_63fb73da9bd363692a1516abb30921c29f18b039_triangle(object):
    enabled_coefficients = []
    tabulate_tensor = tabulate_tensor_integral_63fb73da9bd363692a1516abb30921c29f18b039_triangle
    needs_facet_permutations = False
    coordinate_element_hash = 16933917890882727401
    domain = 2

# End of code for integral integral_63fb73da9bd363692a1516abb30921c29f18b039_triangle

# Code for integral integral_778fc7abc3ea118a6b2ffae9b84390638c355c49_triangle

def tabulate_tensor_integral_778fc7abc3ea118a6b2ffae9b84390638c355c49_triangle(_A, _w, _c, _coordinate_dofs,
                                   _entity_local_index, _quadrature_permutation, custom_data):

    A = numba.carray(_A, (6))
    w = numba.carray(_w, (6))
    c = numba.carray(_c, (0))
    coordinate_dofs = numba.carray(_coordinate_dofs, (36))
    entity_local_index = numba.carray(_entity_local_index, (2))
    quadrature_permutation = numba.carray(_quadrature_permutation, (0))
        # Quadrature rules 
    weights_1ec = np.array([0.0254224531851035, 0.0254224531851035, 0.0254224531851035, 0.0583931378631895, 0.0583931378631895, 0.0583931378631895, 0.041425537809187, 0.041425537809187, 0.041425537809187, 0.041425537809187, 0.041425537809187, 0.041425537809187], dtype=np.float64)
    # Precomputed values of basis functions and precomputations 
    # FE* dimensions: [permutation][entities][points][dofs] 
    FE0_C0_Q1ec = np.array([[[[-0.05512856699248414, 0.6533077030470595, -0.05512856699248412, 0.2205142679699363, 0.015920894998035756, 0.22051426796993656],
    [-0.05512856699248414, -0.05512856699248414, 0.6533077030470594, 0.22051426796993656, 0.22051426796993662, 0.01592089499803573],
    [0.6533077030470595, -0.05512856699248414, -0.05512856699248409, 0.015920894998035784, 0.22051426796993656, 0.22051426796993667],
    [-0.12499898253509756, 0.00143057951778882, -0.12499898253509749, 0.4999959301403893, 0.24857552527162594, 0.49999593014039123],
    [-0.12499898253509756, -0.12499898253509759, 0.00143057951778888, 0.49999593014038934, 0.49999593014039134, 0.2485755252716258],
    [0.0014305795177897915, -0.12499898253509756, -0.12499898253509747, 0.24857552527162488, 0.4999959301403904, 0.4999959301403902],
    [-0.047496257198800046, 0.17376836365417409, -0.11771516330842909, 0.7901604427658231, 0.06597478591860523, 0.1353078281686268],
    [-0.11771516330842921, 0.17376836365417414, -0.04749625719880005, 0.13530782816862683, 0.06597478591860531, 0.7901604427658231],
    [-0.04749625719880002, -0.11771516330842915, 0.1737683636541741, 0.7901604427658231, 0.13530782816862685, 0.06597478591860523],
    [0.1737683636541741, -0.11771516330842913, -0.04749625719880004, 0.06597478591860526, 0.13530782816862694, 0.7901604427658231],
    [-0.11771516330842917, -0.047496257198800095, 0.1737683636541741, 0.13530782816862683, 0.7901604427658232, 0.06597478591860523],
    [0.17376836365417406, -0.047496257198800115, -0.11771516330842907, 0.06597478591860526, 0.7901604427658232, 0.1353078281686269]]]], dtype=np.float64)
    FE1_C0_D10_Q1ec = np.array([[[[0.7476439420339921, 2.4952878840679835, 0.0, 0.2523560579660072, -0.25235605796600874, -3.2429318261019753],
    [0.7476439420339923, -0.7476439420339918, 0.0, 3.4952878840679844, -3.495287884067985, 0.0],
    [-2.495287884067984, -0.7476439420339915, 0.0, 0.25235605796600863, -0.2523560579660075, 3.2429318261019753],
    [0.0028530193163561757, 1.0057060386327161, 0.0, 0.9971469806836398, -0.9971469806836405, -1.0085590579490724],
    [0.00285301931635629, -0.0028530193163599986, 0.0, 2.0057060386327166, -2.005706038632717, 0.0],
    [-1.0057060386327201, -0.0028530193163599504, 0.0, 0.9971469806836403, -0.99714698068364, 1.00855905794908],
    [0.7874198006207365, 1.546009996485596, 0.0, 1.2414098041351396, -1.2414098041351407, -2.333429797106332],
    [-0.2414098041351399, 1.546009996485596, 0.0, 0.2125801993792637, -0.21258019937926445, -1.304600192350456],
    [0.7874198006207362, 0.24140980413514018, 0.0, 2.5460099964855964, -2.5460099964855964, -1.0288296047558763],
    [-1.5460099964855964, 0.2414098041351399, 0.0, 0.21258019937926426, -0.21258019937926395, 1.304600192350456],
    [-0.24140980413513996, -0.7874198006207358, 0.0, 2.5460099964855964, -2.5460099964855964, 1.0288296047558758],
    [-1.5460099964855962, -0.7874198006207358, 0.0, 1.2414098041351405, -1.2414098041351398, 2.333429797106332]]]], dtype=np.float64)
    FE1_C1_D01_Q1ec = np.array([[[[0.7476439420339924, 0.0, -0.7476439420339914, 3.4952878840679835, 0.0, -3.4952878840679844],
    [0.747643942033992, 0.0, 2.4952878840679835, 0.25235605796600824, -3.2429318261019766, -0.2523560579660078],
    [-2.495287884067984, 0.0, -0.7476439420339913, 0.252356057966008, 3.2429318261019766, -0.25235605796600863],
    [0.0028530193163562824, 0.0, -0.0028530193163599283, 2.0057060386327166, 0.0, -2.0057060386327166],
    [0.0028530193163561774, 0.0, 1.0057060386327161, 0.9971469806836402, -1.0085590579490724, -0.99714698068364],
    [-1.0057060386327201, 0.0, -0.002853019316359882, 0.99714698068364, 1.0085590579490804, -0.9971469806836404],
    [0.7874198006207362, 0.0, 0.24140980413514013, 2.546009996485596, -1.0288296047558765, -2.546009996485596],
    [-0.24140980413513968, 0.0, -0.7874198006207356, 2.546009996485596, 1.0288296047558756, -2.5460099964855964],
    [0.7874198006207361, 0.0, 1.5460099964855958, 1.2414098041351402, -2.3334297971063327, -1.24140980413514],
    [-1.546009996485596, 0.0, -0.7874198006207355, 1.2414098041351398, 2.333429797106332, -1.2414098041351405],
    [-0.24140980413514007, 0.0, 1.5460099964855958, 0.2125801993792643, -1.3046001923504562, -0.212580199379264],
    [-1.546009996485596, 0.0, 0.2414098041351402, 0.2125801993792641, 1.3046001923504562, -0.21258019937926448]]]], dtype=np.float64)
    for iq in range(0, 12):
        # ------------------------ 
        # Section: Coefficient 
        # Inputs: w, FE0_C0_Q1ec 
        # Outputs: w0 
        w0 = 0.0
        for ic in range(0, 6):
            w0 += w[ic] * FE0_C0_Q1ec[0, 0, iq, ic]
            
        # ------------------------ 
        # ------------------------ 
        # Section: Jacobian 
        # Inputs: coordinate_dofs, FE1_C0_D10_Q1ec, FE1_C1_D01_Q1ec 
        # Outputs: J0_c3, J0_c0, J0_c2, J0_c1 
        J0_c0 = 0.0
        J0_c3 = 0.0
        J0_c1 = 0.0
        J0_c2 = 0.0
        for ic in range(0, 6):
            J0_c0 += coordinate_dofs[(ic) * 3] * FE1_C0_D10_Q1ec[0, 0, iq, ic]
            J0_c3 += coordinate_dofs[(ic) * 3 + 1] * FE1_C1_D01_Q1ec[0, 0, iq, ic]
            J0_c1 += coordinate_dofs[(ic) * 3] * FE1_C1_D01_Q1ec[0, 0, iq, ic]
            J0_c2 += coordinate_dofs[(ic) * 3 + 1] * FE1_C0_D10_Q1ec[0, 0, iq, ic]
            
        # ------------------------ 
        # ------------------------ 
        # Section: Intermediates 
        # Inputs: w0, J0_c3, J0_c0, J0_c2, J0_c1 
        # Outputs: fw0 
        fw0 = 0
        sv_1ec_0 = J0_c0 * J0_c3
        sv_1ec_1 = J0_c1 * J0_c2
        sv_1ec_2 = -sv_1ec_1
        sv_1ec_3 = sv_1ec_0 + sv_1ec_2
        sv_1ec_4 = np.abs(sv_1ec_3)
        sv_1ec_5 = sv_1ec_4 * w0
        fw0 = sv_1ec_5 * weights_1ec[iq]
        # ------------------------ 
        # ------------------------ 
        # Section: Tensor Computation 
        # Inputs: fw0, FE0_C0_Q1ec 
        # Outputs: A 
        for i in range(0, 6):
            A[(i)] += fw0 * FE0_C0_Q1ec[0, 0, iq, i]
            
        # ------------------------ 
        
    

class integral_778fc7abc3ea118a6b2ffae9b84390638c355c49_triangle(object):
    enabled_coefficients = [1]
    tabulate_tensor = tabulate_tensor_integral_778fc7abc3ea118a6b2ffae9b84390638c355c49_triangle
    needs_facet_permutations = False
    coordinate_element_hash = 16933917890882727401
    domain = 2

# End of code for integral integral_778fc7abc3ea118a6b2ffae9b84390638c355c49_triangle

# Code for form form_d761aec89406377fcbb58a999b7323d30d8bea53


finite_element_hashes_form_d761aec89406377fcbb58a999b7323d30d8bea53 = [16933917890882727401, 16933917890882727401]
form_integral_offsets_form_d761aec89406377fcbb58a999b7323d30d8bea53 = [0, 1, 1, 1, 1, 1]
form_integrals_form_d761aec89406377fcbb58a999b7323d30d8bea53 = [integral_63fb73da9bd363692a1516abb30921c29f18b039_triangle]
form_integral_ids_form_d761aec89406377fcbb58a999b7323d30d8bea53 = [-1]






class form_d761aec89406377fcbb58a999b7323d30d8bea53(object):

  signature = "9c0f4f100853821a50937b9c69cae563cfece5fc9f0aa8c5f48523cecd75bb4f64e4b0e0d765e482914e1d38e7da285445d3fb618171d980b731a6dccc39063d"
  rank = 2

  num_coefficients = 0
  original_coefficient_positions = None
  coefficient_name_map = None

  num_constants = 0
  constant_ranks = None
  constant_shapes = None
  constant_name_map = None

  finite_element_hashes = finite_element_hashes_form_d761aec89406377fcbb58a999b7323d30d8bea53

  form_integrals = form_integrals_form_d761aec89406377fcbb58a999b7323d30d8bea53
  form_integral_ids = form_integral_ids_form_d761aec89406377fcbb58a999b7323d30d8bea53
  form_integral_offsets = form_integral_offsets_form_d761aec89406377fcbb58a999b7323d30d8bea53


# Alias name
form_PoissonQuad_a = form_d761aec89406377fcbb58a999b7323d30d8bea53

# End of code for form form_d761aec89406377fcbb58a999b7323d30d8bea53

# Code for form form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672

original_coefficient_position_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672 = [0]
finite_element_hashes_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672 = [16933917890882727401, 16933917890882727401]
form_integral_offsets_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672 = [0, 1, 1, 1, 1, 1]
form_integrals_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672 = [integral_778fc7abc3ea118a6b2ffae9b84390638c355c49_triangle]
form_integral_ids_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672 = [-1]

coefficient_names_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672 = ["f"]




class form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672(object):

  signature = "dccbdc7bda573c2c069c489b1644c8ccec0aee82c8241444a95c0be7f953228ee36efc3d489e561dd14558559f08b87b4c345ea95563dd6c91c3fdf07645ce47"
  rank = 1

  num_coefficients = 1
  original_coefficient_positions = original_coefficient_position_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672
  coefficient_name_map = coefficient_names_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672

  num_constants = 0
  constant_ranks = None
  constant_shapes = None
  constant_name_map = None

  finite_element_hashes = finite_element_hashes_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672

  form_integrals = form_integrals_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672
  form_integral_ids = form_integral_ids_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672
  form_integral_offsets = form_integral_offsets_form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672


# Alias name
form_PoissonQuad_L = form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672

# End of code for form form_fb2b1df1128f70fdeeddd69cd46d86aa2a83d672

[jorgensd]

What if we rather make this a jitclass, which initializes the tables once: https://numba.readthedocs.io/en/stable/user/jitclass.html?highlight=jitclass ?

[schnellerhase]

I like Numba as a prototype for supporting multiple languages, but I'm less keen on backends that we wouldn't recommend for efficient computation.

This is our take on this as well. We see the numba backend not as a performant or equal replacement of the C backend, but as an example for extensibility of FFCx and it ensures the interface for another backend is sound before experimenting (externally) in this direction. Polishing the interface to general backends first with numba is the easier step to see where this interface is at.

Currently the numba kernels do not support any kind of JITing or easy calling from a user, therefore anyone calling it assumably knows about the performance impact, but prefers the backend for some other reason. If we can make this even clearer somehow happy to do so. The numba module is self contained and does not influence the C module.

While working with the numba backend I saw its capabilities for development and debugging. The Python code is easier, shorter and a single file. Therefore, I think it also makes a good addition as a development tool.

[chrisrichardson]

Maybe we can output a message (e.g. a comment in the generated code) or on the console, reminding users that the kernel is not suitable for production use.

[michalhabera]

Happy to have this merged as an exploratory language backend.