Generalise la::Vector to support GPUs (#3855)

* Updates

* Scatter updates

* Generalise

* Update tests

* Generalise

* Generalise

* Updates

* Update assemblers

* Small updates

* Format

* Simplify

* Remove resize

* Move code outside loops

* Fix size

* Lint

* Fix

* Reinstate demo

* Add size check

* Simplify

* Tidy up

* Remove check

* Work on GPU Vector

* Tidy up Vector

* Lint

* Reduce number of functions

* Lint

* Tidy up

* Enable scatter fwd test

* Enable test

* Test simplification

* Use iterators

* Doc fix

* Change pack fn

* Fix concept

* Lint

* Comment out concept

* Small updates

* Small fix and revert

* Lint

* Simplify

* Doc updates

* Add concepts

* Re-enable demo

* Re-enable more code

* Doc fix

* Fix

* Remove commented code

* PR review updates

* Add comments

* Update cpp/dolfinx/common/Scatterer.h

Co-authored-by: Paul T. Kühner <56360279+schnellerhase@users.noreply.github.com>

---------

Co-authored-by: Paul T. Kühner <56360279+schnellerhase@users.noreply.github.com>

Author: garth-wells

cpp/demo/biharmonic/main.cpp

@@ -202,7 +202,7 @@ int main(int argc, char* argv[])
     //  {cpp:class}`DirichletBC` takes two arguments: the value of the
     //  boundary condition, and the part of the boundary on which the
     //  condition applies. In our example, the value of the boundary
-    //  condition (0.0) can represented using a {cpp:class}`Function`,
+    //  condition (0) can represented using a {cpp:class}`Function`,
     //  and the Dirichlet boundary is defined by the indices of degrees
     //  of freedom to which the boundary condition applies. The
     //  definition of the Dirichlet boundary condition then looks as
@@ -212,7 +212,7 @@ int main(int argc, char* argv[])
     auto facets = mesh::exterior_facet_indices(*mesh->topology());
     const auto bdofs = fem::locate_dofs_topological(
         *V->mesh()->topology_mutable(), *V->dofmap(), 1, facets);
-    fem::DirichletBC<T> bc(0.0, bdofs, V);
+    fem::DirichletBC<T> bc(0, bdofs, V);
 
     //  Now, we have specified the variational forms and can consider
     //  the solution of the variational problem. First, we need to
@@ -237,11 +237,11 @@ int main(int argc, char* argv[])
     MatAssemblyBegin(A.mat(), MAT_FINAL_ASSEMBLY);
     MatAssemblyEnd(A.mat(), MAT_FINAL_ASSEMBLY);
 
-    b.set(0.0);
-    fem::assemble_vector(b.mutable_array(), L);
-    fem::apply_lifting<T, U>(b.mutable_array(), {a}, {{bc}}, {}, T(1.0));
+    std::ranges::fill(b.array(), 0);
+    fem::assemble_vector(b.array(), L);
+    fem::apply_lifting(b.array(), {a}, {{bc}}, {}, T(1));
     b.scatter_rev(std::plus<T>());
-    bc.set(b.mutable_array(), std::nullopt);
+    bc.set(b.array(), std::nullopt);
 
     la::petsc::KrylovSolver lu(MPI_COMM_WORLD);
     la::petsc::options::set("ksp_type", "preonly");
@@ -263,7 +263,7 @@ int main(int argc, char* argv[])
 
     //  Save solution in VTK format
     io::VTKFile file(MPI_COMM_WORLD, "u.pvd", "w");
-    file.write<T>({u}, 0.0);
+    file.write<T>({u}, 0);
   }
 
   PetscFinalize();

cpp/demo/custom_kernel/main.cpp

@@ -111,7 +111,7 @@ double assemble_vector0(std::shared_ptr<const fem::FunctionSpace<T>> V,
   auto dofmap = V->dofmap();
   la::Vector<T> b(dofmap->index_map, 1);
   common::Timer timer("Assembler0 std::function (vector)");
-  fem::assemble_vector(b.mutable_array(), L);
+  fem::assemble_vector(b.array(), L);
   b.scatter_rev(std::plus<T>());
   return la::squared_norm(b);
 }
@@ -141,7 +141,7 @@ double assemble_matrix1(const mesh::Geometry<T>& g, const fem::DofMap& dofmap,
   common::Timer timer("Assembler1 lambda (matrix)");
   md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 3>> x(
       g.x().data(), g.x().size() / 3, 3);
-  fem::impl::assemble_cells<T>(
+  fem::impl::assemble_cells_matrix<T>(
       A.mat_add_values(), g.dofmap(), x, cells, {dofmap.map(), 1, cells}, ident,
       {dofmap.map(), 1, cells}, ident, {}, {}, kernel, {}, {}, {}, {});
   A.scatter_rev();
@@ -167,9 +167,9 @@ double assemble_vector1(const mesh::Geometry<T>& g, const fem::DofMap& dofmap,
   md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 3>> x(
       g.x().data(), g.x().size() / 3, 3);
   common::Timer timer("Assembler1 lambda (vector)");
-  fem::impl::assemble_cells<T, 1>([](auto, auto, auto, auto) {},
-                                  b.mutable_array(), g.dofmap(), x, cells,
-                                  {dofmap.map(), 1, cells}, kernel, {}, {}, {});
+  fem::impl::assemble_cells<1>([](auto, auto, auto, auto) {}, b.array(),
+                               g.dofmap(), x, cells, {dofmap.map(), 1, cells},
+                               kernel, {}, {}, {});
   b.scatter_rev(std::plus<T>());
   return la::squared_norm(b);
 }

cpp/demo/hyperelasticity/main.cpp

@@ -86,9 +86,9 @@ class HyperElasticProblem
     return [&](const Vec x, Vec)
     {
       // Assemble b and update ghosts
-      std::span b(_b.mutable_array());
+      std::span b(_b.array());
       std::ranges::fill(b, 0);
-      fem::assemble_vector<T>(b, _l);
+      fem::assemble_vector(b, _l);
       VecGhostUpdateBegin(_b_petsc, ADD_VALUES, SCATTER_REVERSE);
       VecGhostUpdateEnd(_b_petsc, ADD_VALUES, SCATTER_REVERSE);
 
@@ -197,7 +197,7 @@ int main(int argc, char* argv[])
           constexpr U theta = std::numbers::pi / 3;
 
           // New coordinates
-          std::vector<U> fdata(3 * x.extent(1), 0.0);
+          std::vector<U> fdata(3 * x.extent(1), 0);
           md::mdspan<U, md::extents<std::size_t, 3, md::dynamic_extent>> f(
               fdata.data(), 3, x.extent(1));
           for (std::size_t p = 0; p < x.extent(1); ++p)
@@ -282,12 +282,12 @@ int main(int argc, char* argv[])
 
     // Save solution in VTK format
     io::VTKFile file_u(mesh->comm(), "u.pvd", "w");
-    file_u.write<T>({*u}, 0.0);
+    file_u.write<T>({*u}, 0);
 
     // Save Cauchy stress in XDMF format
     io::XDMFFile file_sigma(mesh->comm(), "sigma.xdmf", "w");
     file_sigma.write_mesh(*mesh);
-    file_sigma.write_function(sigma, 0.0);
+    file_sigma.write_function(sigma, 0);
   }
 
   PetscFinalize();

cpp/demo/interpolation-io/main.cpp

@@ -68,7 +68,7 @@ void interpolate_scalar(std::shared_ptr<mesh::Mesh<U>> mesh,
   // ParaView
   io::VTXWriter<U> outfile(mesh->comm(), filename.replace_extension("bp"), {u},
                            "BP4");
-  outfile.write(0.0);
+  outfile.write(0);
   outfile.close();
 #endif
 }
@@ -137,15 +137,15 @@ void interpolate_nedelec(std::shared_ptr<mesh::Mesh<U>> mesh,
   u->interpolate(
       [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
       {
-        std::vector<T> f(2 * x.extent(1), 0.0);
+        std::vector<T> f(2 * x.extent(1), 0);
         std::copy_n(x.data_handle(), f.size(), f.begin());
         return {f, {2, x.extent(1)}};
       },
       cells0);
   u->interpolate(
       [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
       {
-        std::vector<T> f(2 * x.extent(1), 0.0);
+        std::vector<T> f(2 * x.extent(1), 0);
         std::copy_n(x.data_handle(), f.size(), f.begin());
         std::ranges::transform(f, f.begin(), [](auto x) { return x + T(1); });
         return {f, {2, x.extent(1)}};
@@ -185,7 +185,7 @@ void interpolate_nedelec(std::shared_ptr<mesh::Mesh<U>> mesh,
 #ifdef HAS_ADIOS2
   io::VTXWriter<U> outfile(mesh->comm(), filename.replace_extension("bp"),
                            {u_l}, "BP4");
-  outfile.write(0.0);
+  outfile.write(0);
   outfile.close();
 #endif
 }

cpp/demo/interpolation_different_meshes/main.cpp

@@ -53,7 +53,7 @@ int main(int argc, char* argv[])
 
     auto fun = [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
     {
-      std::vector<T> fdata(3 * x.extent(1), 0.0);
+      std::vector<T> fdata(3 * x.extent(1), 0);
       using dextent = md::dextents<std::size_t, 2>;
       md::mdspan<double, dextent> f(fdata.data(), 3, x.extent(1));
       for (std::size_t i = 0; i < x.extent(1); ++i)
@@ -84,9 +84,9 @@ int main(int argc, char* argv[])
 
 #ifdef HAS_ADIOS2
     io::VTXWriter<double> write_tet(mesh_tet->comm(), "u_tet.bp", {u_tet});
-    write_tet.write(0.0);
+    write_tet.write(0);
     io::VTXWriter<double> write_hex(mesh_hex->comm(), "u_hex.bp", {u_hex});
-    write_hex.write(0.0);
+    write_hex.write(0);
 #endif
   }
   MPI_Finalize();

cpp/demo/mixed_poisson/main.cpp

@@ -327,16 +327,16 @@ int main(int argc, char* argv[])
     MatAssemblyEnd(A.mat(), MAT_FINAL_ASSEMBLY);
 
     // Assemble the linear form `L` into RHS vector
-    b.set(0);
-    fem::assemble_vector(b.mutable_array(), L);
+    std::ranges::fill(b.array(), 0);
+    fem::assemble_vector(b.array(), L);
 
     // Modify unconstrained dofs on RHS to account for Dirichlet BC dofs
     // (constrained dofs), and perform parallel update on the vector.
-    fem::apply_lifting<T, U>(b.mutable_array(), {a}, {{bc}}, {}, T(1));
+    fem::apply_lifting(b.array(), {a}, {{bc}}, {}, T(1));
     b.scatter_rev(std::plus<T>());
 
     // Set value for constrained dofs
-    bc.set(b.mutable_array(), std::nullopt);
+    bc.set(b.array(), std::nullopt);
 
     // Create PETSc linear solver
     la::petsc::KrylovSolver lu(MPI_COMM_WORLD);

cpp/demo/poisson/main.cpp

@@ -150,7 +150,7 @@ int main(int argc, char* argv[])
     //  {cpp:class}`DirichletBC` takes two arguments: the value of the
     //  boundary condition, and the part of the boundary on which the
     //  condition applies. In our example, the value of the boundary
-    //  condition (0.0) can represented using a {cpp:class}`Function`,
+    //  condition (0) can represented using a {cpp:class}`Function`,
     //  and the Dirichlet boundary is defined by the indices of degrees
     //  of freedom to which the boundary condition applies. The
     //  definition of the Dirichlet boundary condition then looks as
@@ -175,7 +175,7 @@ int main(int argc, char* argv[])
         });
     std::vector bdofs = fem::locate_dofs_topological(
         *V->mesh()->topology_mutable(), *V->dofmap(), 1, facets);
-    fem::DirichletBC<T> bc(0.0, bdofs, V);
+    fem::DirichletBC<T> bc(0, bdofs, V);
 
     f->interpolate(
         [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
@@ -222,11 +222,11 @@ int main(int argc, char* argv[])
     MatAssemblyBegin(A.mat(), MAT_FINAL_ASSEMBLY);
     MatAssemblyEnd(A.mat(), MAT_FINAL_ASSEMBLY);
 
-    b.set(0.0);
-    fem::assemble_vector(b.mutable_array(), L);
-    fem::apply_lifting<T, U>(b.mutable_array(), {a}, {{bc}}, {}, T(1));
+    std::ranges::fill(b.array(), 0);
+    fem::assemble_vector(b.array(), L);
+    fem::apply_lifting(b.array(), {a}, {{bc}}, {}, T(1));
     b.scatter_rev(std::plus<T>());
-    bc.set(b.mutable_array(), std::nullopt);
+    bc.set(b.array(), std::nullopt);
 
     la::petsc::KrylovSolver lu(MPI_COMM_WORLD);
     la::petsc::options::set("ksp_type", "preonly");
@@ -248,7 +248,7 @@ int main(int argc, char* argv[])
 
     // Save solution in VTK format
     io::VTKFile file(MPI_COMM_WORLD, "u.pvd", "w");
-    file.write<T>({*u}, 0.0);
+    file.write<T>({*u}, 0);
 
 #ifdef HAS_ADIOS2
     // Save solution in VTX format

cpp/demo/poisson_matrix_free/main.cpp

@@ -64,7 +64,7 @@ namespace linalg
 /// @param[in] y
 void axpy(auto&& r, auto alpha, auto&& x, auto&& y)
 {
-  std::ranges::transform(x.array(), y.array(), r.mutable_array().begin(),
+  std::ranges::transform(x.array(), y.array(), r.array().begin(),
                          [alpha](auto x, auto y) { return alpha * x + y; });
 }
 
@@ -178,18 +178,18 @@ void solver(MPI_Comm comm)
 
   // Assemble RHS vector
   la::Vector<T> b(V->dofmap()->index_map, V->dofmap()->index_map_bs());
-  fem::assemble_vector(b.mutable_array(), L);
+  fem::assemble_vector(b.array(), L);
 
   // Apply lifting to account for Dirichlet boundary condition
   // b <- b - A * x_bc
-  bc->set(ui->x()->mutable_array(), std::nullopt, T(-1));
-  fem::assemble_vector(b.mutable_array(), M);
+  bc->set(ui->x()->array(), std::nullopt, T(-1));
+  fem::assemble_vector(b.array(), M);
 
   // Communicate ghost values
   b.scatter_rev(std::plus<T>());
 
   // Set BC dofs to zero (effectively zeroes columns of A)
-  bc->set(b.mutable_array(), std::nullopt, T(0));
+  bc->set(b.array(), std::nullopt, T(0));
 
   b.scatter_fwd();
 
@@ -201,18 +201,18 @@ void solver(MPI_Comm comm)
   auto action = [&M, &ui, &bc, &coeff, &constants](auto& x, auto& y)
   {
     // Zero y
-    y.set(0.0);
+    std::ranges::fill(y.array(), 0);
 
     // Update coefficient ui (just copy data from x to ui)
-    std::ranges::copy(x.array(), ui->x()->mutable_array().begin());
+    std::ranges::copy(x.array(), ui->x()->array().begin());
 
     // Compute action of A on x
     fem::pack_coefficients(M, coeff);
-    fem::assemble_vector(y.mutable_array(), M, std::span<const T>(constants),
+    fem::assemble_vector(y.array(), M, std::span<const T>(constants),
                          fem::make_coefficients_span(coeff));
 
     // Set BC dofs to zero (effectively zeroes rows of A)
-    bc->set(y.mutable_array(), std::nullopt, T(0));
+    bc->set(y.array(), std::nullopt, T(0));
 
     // Accumulate ghost values
     y.scatter_rev(std::plus<T>());
@@ -226,7 +226,7 @@ void solver(MPI_Comm comm)
   int num_it = linalg::cg(*u->x(), b, action, 200, 1e-6);
 
   // Set BC values in the solution vectors
-  bc->set(u->x()->mutable_array(), std::nullopt, T(1));
+  bc->set(u->x()->array(), std::nullopt, T(1));
 
   // Compute L2 error (squared) of the solution vector e = (u - u_d, u
   // - u_d)*dx

cpp/dolfinx/common/IndexMap.cpp

@@ -538,7 +538,7 @@ std::vector<int32_t>
 common::compute_owned_indices(std::span<const std::int32_t> indices,
                               const IndexMap& map)
 {
-  // Assume that indices are sorted and unique
+  // Require that indices are sorted and unique
   assert(std::ranges::is_sorted(indices));
 
   std::span ghosts = map.ghosts();

cpp/dolfinx/common/IndexMap.h

@@ -29,13 +29,14 @@ enum class IndexMapOrder : bool
   any = false      ///< Allow arbitrary ordering of ghost indices in sub-maps
 };
 
-/// @brief Given a sorted vector of indices (local numbering, owned or
+/// @brief Given a sorted list of indices (local indexing, owned or
 /// ghost) and an index map, this function returns the indices owned by
 /// this process, including indices that might have been in the list of
 /// indices on another processes.
-/// @param[in] indices List of indices
-/// @param[in] map The index map
-/// @return Indices owned by the calling process
+///
+/// @param[in] indices List of indices.
+/// @param[in] map The index map.
+/// @return Indices owned by the calling process.
 std::vector<int32_t>
 compute_owned_indices(std::span<const std::int32_t> indices,
                       const IndexMap& map);