Single-precision (fp32) build support

.github/workflows/core.yml

@@ -23,6 +23,14 @@ on:
         description: Whether to test a complex build
         type: boolean
         required: true
+      test_single:
+        description: Whether to test a single-precision (fp32) build
+        type: boolean
+        required: true
+      test_single_complex:
+        description: Whether to test a single-precision complex build
+        type: boolean
+        required: true
       test_cuda:
         description: Whether to test using CUDA-enabled PETSc
         type: boolean
@@ -54,6 +62,14 @@ on:
         description: Whether to test a complex build
         type: boolean
         required: true
+      test_single:
+        description: Whether to test a single-precision (fp32) build
+        type: boolean
+        required: true
+      test_single_complex:
+        description: Whether to test a single-precision complex build
+        type: boolean
+        required: true
       test_cuda:
         description: Whether to test using CUDA-enabled PETSc
         type: boolean
@@ -264,6 +280,22 @@ jobs:
           if [ ${{ inputs.test_complex }} == 'true' ]; then
             matrix='{"scalar_type": "complex", "gpu": "none"}'
           fi
+          if [ ${{ inputs.test_single }} == 'true' ]; then
+            arch='{"scalar_type": "single", "gpu": "none"}'
+            if [ "$matrix" ]; then
+              matrix="$matrix, $arch"
+            else
+              matrix="$arch"
+            fi
+          fi
+          if [ ${{ inputs.test_single_complex }} == 'true' ]; then
+            arch='{"scalar_type": "single-complex", "gpu": "none"}'
+            if [ "$matrix" ]; then
+              matrix="$matrix, $arch"
+            else
+              matrix="$arch"
+            fi
+          fi
           if [ ${{ inputs.test_cuda }} == 'true' ]; then
             arch='{"scalar_type": "default", "gpu": "cuda"}'
             if [ "$matrix" ]; then

.github/workflows/pr.yml

@@ -23,6 +23,8 @@ jobs:
       # Run all tests of a specific configuration (e.g. complex) if the right label
       # is on the PR, otherwise do nothing.
       test_complex: ${{ contains(github.event.pull_request.labels.*.name, 'ci:complex') }}
+      test_single: ${{ contains(github.event.pull_request.labels.*.name, 'ci:single') }}
+      test_single_complex: ${{ contains(github.event.pull_request.labels.*.name, 'ci:single-complex') }}
       test_cuda: ${{ contains(github.event.pull_request.labels.*.name, 'ci:cuda') }}
       test_macos: ${{ contains(github.event.pull_request.labels.*.name, 'ci:macos') }}
     secrets: inherit

.github/workflows/push.yml

@@ -14,6 +14,8 @@ jobs:
       source_ref: ${{ github.ref_name }}
       base_ref: ${{ github.ref_name }}
       test_complex: true
+      test_single: true
+      test_single_complex: true
       test_cuda: true
       test_macos: true
       deploy_website: true

firedrake/cython/petschdr.pxi

@@ -8,6 +8,7 @@ cdef extern from "mpi-compat.h":
 
 cdef extern from "petsc.h":
     ctypedef long PetscInt
+    # Nominal hints only: C compiler resolves PetscReal/PetscScalar from petsc.h, so fp32 builds are correct.
     ctypedef double PetscReal
     ctypedef double PetscScalar
     ctypedef enum PetscBool:

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,27 @@ 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). */
 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,7 @@ 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)
+        coord = np.asarray(coord)
         evaluator = PointEvaluator(self.function_space().mesh(), coord)
         result = evaluator.evaluate(self)
         if len(coord.shape) == 1:

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]) {

firedrake/mesh.py

@@ -37,7 +37,7 @@
 import firedrake.extrusion_utils as eutils
 import firedrake.cython.rtree as rtree
 import firedrake.utils as utils
-from firedrake.utils import as_cstr, IntType, RealType
+from firedrake.utils import as_cstr, as_ctypes, IntType, RealType, RealType_c
 from firedrake.logging import logger
 from firedrake.parameters import parameters
 from firedrake.petsc import PETSc, DEFAULT_PARTITIONER
@@ -422,7 +422,7 @@ def _from_triangle(filename, dim, comm):
                 nodecount = header[0]
                 nodedim = header[1]
                 assert nodedim == dim
-                coordinates = np.loadtxt(nodefile, usecols=list(range(1, dim+1)), skiprows=1, dtype=np.double)
+                coordinates = np.loadtxt(nodefile, usecols=list(range(1, dim+1)), skiprows=1, dtype=PETSc.RealType)
                 assert nodecount == coordinates.shape[0]
 
             with open(basename+".ele") as elefile:
@@ -472,12 +472,10 @@ def plex_from_cell_list(dim, cells, coords, comm, name=None):
     :arg comm: communicator to build the mesh on. Must be a PyOP2 internal communicator
     :kwarg name: name of the plex
     """
-    # These types are /correct/, DMPlexCreateFromCellList wants int
-    # and double (not PetscInt, PetscReal).
     with temp_internal_comm(comm) as icomm:
         if comm.rank == 0:
             cells = np.asarray(cells, dtype=np.int32)
-            coords = np.asarray(coords, dtype=np.double)
+            coords = np.asarray(coords, dtype=PETSc.RealType)
             icomm.bcast(cells.shape, root=0)
             icomm.bcast(coords.shape, root=0)
             # Provide the actual data on rank 0.
@@ -491,7 +489,7 @@ def plex_from_cell_list(dim, cells, coords, comm, name=None):
             # A subsequent call to plex.distribute() takes care of parallel partitioning
             plex = PETSc.DMPlex().createFromCellList(dim,
                                                      np.zeros(cell_shape, dtype=np.int32),
-                                                     np.zeros(coord_shape, dtype=np.double),
+                                                     np.zeros(coord_shape, dtype=PETSc.RealType),
                                                      comm=comm)
     if name is not None:
         plex.setName(name)
@@ -2608,7 +2606,8 @@ def rtree(self):
         coords_max = coords_mid + (tolerance + 0.5)*d
 
         with PETSc.Log.Event("rtree_build"):
-            self._rtree = rtree.build_from_aabb(coords_min, coords_max)
+            # rtree C API requires float64 regardless of PETSc scalar precision.
+            self._rtree = rtree.build_from_aabb(np.asarray(coords_min, dtype=np.float64), np.asarray(coords_max, dtype=np.float64))
         self._saved_coordinate_dat_version = self.coordinates.dat.dat_version
         return self._rtree
 
@@ -2656,7 +2655,7 @@ def locate_cell_and_reference_coordinate(self, x, tolerance=None, cell_ignore=No
             (cell number, reference coordinates) of type (int, numpy array),
             or, when point is not in the domain, (None, None).
         """
-        x = np.asarray(x)
+        x = np.asarray(x, dtype=np.float64)
         if x.size != self.geometric_dimension:
             raise ValueError("Point must have the same geometric dimension as the mesh")
         x = x.reshape((1, self.geometric_dimension))
@@ -2694,11 +2693,12 @@ def locate_cells_ref_coords_and_dists(self, xs, tolerance=None, cells_ignore=Non
             tolerance = self.tolerance
         else:
             self.tolerance = tolerance
-        xs = np.asarray(xs, dtype=utils.ScalarType)
+        # Physical coordinates: always float64 (libspatialindex requires double).
+        xs = np.asarray(xs, dtype=np.float64)
         xs = xs.real.copy()
         if xs.shape[1] != self.geometric_dimension:
             raise ValueError("Point coordinate dimension does not match mesh geometric dimension")
-        Xs = np.empty_like(xs)
+        Xs = np.empty((len(xs), self.geometric_dimension), dtype=RealType)
         npoints = len(xs)
         if cells_ignore is None or cells_ignore[0][0] is None:
             cells_ignore = np.full((npoints, 1), -1, dtype=IntType, order="C")
@@ -2707,14 +2707,14 @@ def locate_cells_ref_coords_and_dists(self, xs, tolerance=None, cells_ignore=Non
         if cells_ignore.shape[0] != npoints:
             raise ValueError("Number of cells to ignore does not match number of points")
         assert cells_ignore.shape == (npoints, cells_ignore.shape[1])
-        ref_cell_dists_l1 = np.empty(npoints, dtype=utils.RealType)
+        ref_cell_dists_l1 = np.empty(npoints, dtype=RealType)
         cells = np.empty(npoints, dtype=IntType)
         assert xs.size == npoints * self.geometric_dimension
         run_c = self._c_locator(tolerance=tolerance)
-        cells_data = cells.ctypes.data_as(ctypes.POINTER(ctypes.c_int))
-        ref_cells_dists = ref_cell_dists_l1.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
+        cells_data = cells.ctypes.data_as(ctypes.POINTER(as_ctypes(IntType)))
+        ref_cells_dists = ref_cell_dists_l1.ctypes.data_as(ctypes.POINTER(as_ctypes(RealType)))
         xs_data = xs.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
-        Xs_data = Xs.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
+        Xs_data = Xs.ctypes.data_as(ctypes.POINTER(as_ctypes(RealType)))
         with PETSc.Log.Event("c_locator_run"):
             run_c(self.coordinates._ctypes, xs_data, Xs_data, ref_cells_dists, cells_data, npoints, cells_ignore.shape[1], cells_ignore)
         return cells, Xs, ref_cell_dists_l1
@@ -2732,7 +2732,9 @@ def _c_locator(self, tolerance=None):
             IntTypeC = as_cstr(IntType)
             src = pq_utils.src_locate_cell(self, tolerance=tolerance)
             src += dedent(f"""
-                int locator(struct Function *f, double *x, double *X, double *ref_cell_dists_l1, {IntTypeC} *cells, {IntTypeC} npoints, size_t ncells_ignore, int* cells_ignore)
+                /* x is double (not {RealType_c}) because libspatialindex requires double for the
+                   spatial index lookup; xs is always cast to float64 before this call. */
+                {IntTypeC} locator(struct Function *f, double *x, {RealType_c} *X, {RealType_c} *ref_cell_dists_l1, {IntTypeC} *cells, {IntTypeC} npoints, size_t ncells_ignore, {IntTypeC}* cells_ignore)
                 {{
                     {IntTypeC} j = 0;  /* index into x and X */
                     for({IntTypeC} i=0; i<npoints; i++) {{
@@ -2746,7 +2748,7 @@ def _c_locator(self, tolerance=None):
                         pointquery_utils.py. If they contain python calls, this loop will
                         not run at c-loop speed. */
                         /* cells_ignore has shape (npoints, ncells_ignore) - find the ith row */
-                        int *cells_ignore_i = cells_ignore + i*ncells_ignore;
+                        {IntTypeC} *cells_ignore_i = cells_ignore + i*ncells_ignore;
                         cells[i] = locate_cell(f, &x[j], {self.geometric_dimension}, &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &ref_cell_dists_l1[i], ncells_ignore, cells_ignore_i);
 
                         for (int k = 0; k < {self.geometric_dimension}; k++) {{
@@ -2777,13 +2779,13 @@ def _c_locator(self, tolerance=None):
             locator = getattr(dll, "locator")
             locator.argtypes = [ctypes.POINTER(function._CFunction),
                                 ctypes.POINTER(ctypes.c_double),
-                                ctypes.POINTER(ctypes.c_double),
-                                ctypes.POINTER(ctypes.c_double),
-                                ctypes.POINTER(ctypes.c_int),
+                                ctypes.POINTER(as_ctypes(RealType)),
+                                ctypes.POINTER(as_ctypes(RealType)),
+                                ctypes.POINTER(as_ctypes(IntType)),
                                 ctypes.c_size_t,
                                 ctypes.c_size_t,
-                                np.ctypeslib.ndpointer(ctypes.c_int, flags="C_CONTIGUOUS")]
-            locator.restype = ctypes.c_int
+                                np.ctypeslib.ndpointer(as_ctypes(IntType), flags="C_CONTIGUOUS")]
+            locator.restype = as_ctypes(IntType)
             return cache.setdefault(tolerance, locator)
 
     @cached_property  # TODO: Recalculate if mesh moves. Extend this for regular meshes.