Single-precision (fp32) build support

.github/workflows/core.yml

@@ -83,10 +83,10 @@ jobs:
     name: Build and test Firedrake (Linux)
     strategy:
       # We want to know all of the tests which fail, so don't kill real if
-      # complex fails and vice-versa
+      # complex or single fails and vice-versa
       fail-fast: false
       matrix:
-        arch: [default, complex]
+        arch: [default, complex, single]
     runs-on: [self-hosted, Linux]
     container:
       # TODO: set to 'ubuntu:latest'

firedrake/assemble.py

@@ -28,7 +28,7 @@
 from firedrake.petsc import PETSc
 from firedrake.slate import slac, slate
 from firedrake.slate.slac.kernel_builder import CellFacetKernelArg, LayerCountKernelArg
-from firedrake.utils import ScalarType, assert_empty, tuplify
+from firedrake.utils import IntType, ScalarType, assert_empty, tuplify
 from pyop2 import op2
 from pyop2.exceptions import MapValueError, SparsityFormatError
 from functools import cached_property
@@ -1783,8 +1783,8 @@ def _make_mat_global_kernel_arg(self, Vrow, Vcol):
         else:
             rmap_arg, cmap_arg = (V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)._global_kernel_arg for V in [Vrow, Vcol])
             # PyOP2 matrix objects have scalar dims so we flatten them here
-            rdim = numpy.prod(self._get_dim(relem), dtype=int)
-            cdim = numpy.prod(self._get_dim(celem), dtype=int)
+            rdim = numpy.prod(self._get_dim(relem), dtype=IntType)
+            cdim = numpy.prod(self._get_dim(celem), dtype=IntType)
             return op2.MatKernelArg((((rdim, cdim),),), (rmap_arg, cmap_arg), unroll=self._unroll)
 
     @staticmethod
@@ -1879,7 +1879,7 @@ def _as_global_kernel_arg_coefficient(_, self):
 @_as_global_kernel_arg.register(kernel_args.ConstantKernelArg)
 def _as_global_kernel_arg_constant(_, self):
     const = next(self._constants)
-    value_size = numpy.prod(const.ufl_shape, dtype=int)
+    value_size = numpy.prod(const.ufl_shape, dtype=IntType)
     return op2.GlobalKernelArg((value_size,))
 
 

firedrake/cython/supermeshimpl.pyx

@@ -160,8 +160,8 @@ def intersection_finder(mesh_A, mesh_B):
 
     libsupermesh_tree_intersection_finder_query_output(&nindices)
 
-    indices = numpy.empty((nindices,), dtype=int)
-    indptr  = numpy.empty((mesh_A.num_cells() + 1,), dtype=int)
+    indices = numpy.empty((nindices,), dtype=IntType)
+    indptr  = numpy.empty((mesh_A.num_cells() + 1,), dtype=IntType)
 
     libsupermesh_tree_intersection_finder_get_output(&ncells_A, &nindices, <long*>indices.data, <long*>indptr.data)
 

firedrake/evaluate.h

@@ -9,13 +9,13 @@ extern "C" {
 
 struct Function {
 	/* Number of cells in the base mesh */
-	int n_cols;
+	PetscInt n_cols;
 
 	/* 1 if extruded, 0 if not */
 	int extruded;
 
 	/* number of layers for extruded, otherwise 1 */
-	int n_layers;
+	PetscInt n_layers;
 
 	/* Coordinate values and node mapping */
 	PetscScalar *coords;
@@ -36,26 +36,28 @@ struct Function {
 
 typedef PetscReal (*ref_cell_l1_dist)(void *data_,
 				struct Function *f,
-				int cell,
+				PetscInt cell,
 				double *x);
 
 typedef PetscReal (*ref_cell_l1_dist_xtr)(void *data_,
 				struct Function *f,
-				int cell,
-				int layer,
+				PetscInt cell,
+				PetscInt layer,
 				double *x);
 
-extern int locate_cell(struct Function *f,
+extern PetscInt locate_cell(struct Function *f,
 		       double *x,
 		       int dim,
 		       ref_cell_l1_dist try_candidate,
 		       ref_cell_l1_dist_xtr try_candidate_xtr,
 		       void *temp_ref_coords,
 		       void *found_ref_coords,
-		       double *found_ref_cell_dist_l1,
-			   size_t ncells_ignore,
-			   int* cells_ignore);
+		       PetscReal *found_ref_cell_dist_l1,
+		       size_t ncells_ignore,
+		       PetscInt* cells_ignore);
 
+/* x is physical coordinates: always double (libspatialindex requires float64).
+ * Return value is a status code (-1 = not found, 0 = success), not a cell index. */
 extern int evaluate(struct Function *f,
 		    double *x,
 		    PetscScalar *result);

firedrake/function.py

@@ -37,9 +37,9 @@
 
 class _CFunction(ctypes.Structure):
     r"""C struct collecting data from a :class:`Function`"""
-    _fields_ = [("n_cols", c_int),
+    _fields_ = [("n_cols", as_ctypes(IntType)),
                 ("extruded", c_int),
-                ("n_layers", c_int),
+                ("n_layers", as_ctypes(IntType)),
                 ("coords", c_void_p),
                 ("coords_map", POINTER(as_ctypes(IntType))),
                 ("f", c_void_p),
@@ -564,7 +564,8 @@ def evaluate(self, coord, mapping, component, index_values):
         # Called by UFL when evaluating expressions at coordinates
         if component or index_values:
             raise NotImplementedError("Unsupported arguments when attempting to evaluate Function.")
-        coord = np.asarray(coord, dtype=utils.ScalarType)
+        # Point evaluation always uses float64 for geometric robustness
+        coord = np.asarray(coord, dtype=np.float64)
         evaluator = PointEvaluator(self.function_space().mesh(), coord)
         result = evaluator.evaluate(self)
         if len(coord.shape) == 1:

firedrake/functionspaceimpl.py

@@ -19,6 +19,7 @@
 from pyop2.utils import as_tuple
 
 from firedrake import dmhooks
+from firedrake.utils import IntType
 from firedrake.mesh import MeshGeometry, MeshSequenceTopology, MeshSequenceGeometry
 from firedrake.functionspacedata import get_shared_data, create_element
 from firedrake.petsc import PETSc
@@ -569,7 +570,7 @@ def __init__(self, mesh, element, name=None):
         :class:`finat.ufl.mixedelement.TensorElement` have rank 1 and 2
         respectively."""
 
-        self.block_size = int(numpy.prod(self.shape, dtype=int))
+        self.block_size = int(numpy.prod(self.shape, dtype=IntType))
         r"""The total number of degrees of freedom at each function
         space node."""
         self.name = name

firedrake/locate.c

@@ -1,22 +1,21 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <rtree-capi.h>
-#include <float.h>
 #include <evaluate.h>
 
-int locate_cell(struct Function *f,
+PetscInt locate_cell(struct Function *f,
         double *x,
         int dim,
         ref_cell_l1_dist try_candidate,
         ref_cell_l1_dist_xtr try_candidate_xtr,
         void *temp_ref_coords,
         void *found_ref_coords,
-        double *found_ref_cell_dist_l1,
+        PetscReal *found_ref_cell_dist_l1,
         size_t ncells_ignore,
-        int* cells_ignore)
+        PetscInt* cells_ignore)
 {
     RTreeError err;
-    int cell = -1;
+    PetscInt cell = -1;
     int cell_ignore_found = 0;
     /* NOTE: temp_ref_coords and found_ref_coords are actually of type
     struct ReferenceCoords but can't be declared as such in the function
@@ -25,8 +24,8 @@ int locate_cell(struct Function *f,
     surrounds this is declared in pointquery_utils.py. We cast when we use the
     ref_coords_copy function and trust that the underlying memory which the
     pointers refer to is updated as necessary. */
-    double ref_cell_dist_l1 = DBL_MAX;
-    double current_ref_cell_dist_l1 =  -0.5;
+    PetscReal ref_cell_dist_l1 = PETSC_MAX_REAL;
+    PetscReal current_ref_cell_dist_l1 =  -0.5;
     /* NOTE: `tolerance`, which is used throughout this funciton, is a static
        variable defined outside this function when putting together all the C
        code that needs to be compiled - see pointquery_utils.py */
@@ -73,9 +72,9 @@ int locate_cell(struct Function *f,
     }
     else {
         for (size_t i = 0; i < nids; i++) {
-            int nlayers = f->n_layers;
-            int c = ids[i] / nlayers;
-            int l = ids[i] % nlayers;
+            PetscInt nlayers = f->n_layers;
+            PetscInt c = ids[i] / nlayers;
+            PetscInt l = ids[i] % nlayers;
             current_ref_cell_dist_l1 = (*try_candidate_xtr)(temp_ref_coords, f, c, l, x);
             for (size_t j = 0; j < ncells_ignore; j++) {
                 if (ids[i] == cells_ignore[j]) {