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/adjoint_utils/assembly.py

@@ -36,8 +36,12 @@ def wrapper(form, *args, **kwargs):
             if not annotate:
                 return output
 
-            if not isinstance(output, (float, Function, Cofunction)):
+            if not isinstance(output, (numbers.Real, Function, Cofunction)):
                 raise NotImplementedError("Taping for complex-valued 0-forms not yet done!")
+            # pyadjoint's AdjFloat requires Python float; coerce numpy scalars
+            # (numpy.float64 is a float subclass so this is a no-op in fp64 mode)
+            if not isinstance(output, (float, Function, Cofunction)):
+                output = float(output)
             output = create_overloaded_object(output)
             block = AssembleBlock(form, ad_block_tag=ad_block_tag)
 

firedrake/adjoint_utils/blocks/assembly.py

@@ -1,3 +1,4 @@
+import numbers
 import ufl
 import firedrake
 from ufl.domain import extract_domains
@@ -136,7 +137,7 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
             else:
                 dform += firedrake.action(firedrake.derivative(form, c_rep),
                                           tlm_value)
-        if not isinstance(dform, float):
+        if not isinstance(dform, numbers.Real):
             dform = ufl.algorithms.expand_derivatives(dform)
             dform = firedrake.assemble(dform)
         return dform
@@ -184,7 +185,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
             else:
                 ddform += firedrake.derivative(dform, c2_rep, tlm_input)
 
-        if not isinstance(ddform, float):
+        if not isinstance(ddform, numbers.Real):
             ddform = ufl.algorithms.expand_derivatives(ddform)
             if not (isinstance(ddform, ufl.ZeroBaseForm)
                     or (isinstance(ddform, ufl.Form) and ddform.empty())):
@@ -200,6 +201,8 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
     def recompute_component(self, inputs, block_variable, idx, prepared):
         form = prepared
         output = firedrake.assemble(form)
+        if isinstance(output, numbers.Real) and not isinstance(output, float):
+            output = float(output)
         output = create_overloaded_object(output)
         if isinstance(output, firedrake.Function):
             return maybe_disk_checkpoint(output)

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/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:
@@ -607,6 +607,8 @@ def _at(self, arg, *args, **kwargs):
             if not np.allclose(arg.imag, 0):
                 raise ValueError("Provided points have non-zero imaginary part")
             arg = arg.real.copy()
+        # Point location (libspatialindex) needs float64 coords, not ScalarType.
+        arg = np.asarray(arg, dtype=np.float64)
 
         dont_raise = kwargs.get('dont_raise', False)
 

firedrake/interpolation.py

@@ -947,7 +947,9 @@ def _create_permutation_mat(self, mat_type: Literal["aij", "baij"]) -> PETSc.Mat
         # Vector and Tensor valued functions are stored in a flattened array, so
         # we need to space out the column indices according to the block size
         cols = (self.target_space.block_size * perm[:, None] + numpy.arange(self.target_space.block_size, dtype=IntType)[None, :]).reshape(-1)
-        mat.setValuesCSR(rows, cols, numpy.ones_like(cols, dtype=IntType))
+        # Matrix values must be PetscScalar, not IntType: petsc4py rejects an unsafe
+        # int -> float32 cast when PETSc is built in single precision.
+        mat.setValuesCSR(rows, cols, numpy.ones_like(cols, dtype=ScalarType))
         mat.assemble()
         if self.forward_reduce and not self.ufl_interpolate.is_adjoint:
             # The mat we have constructed thus far takes us from the input-ordering VOM to the

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