assemble_matrix_impl.h

// Copyright (C) 2018-2026 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "DofMap.h"
#include "Form.h"
#include "FunctionSpace.h"
#include "traits.h"
#include <algorithm>
#include <concepts>
#include <dolfinx/la/utils.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <functional>
#include <iterator>
#include <span>
#include <stdexcept>
#include <vector>
 
namespace dolfinx::fem::impl
{
bool has_bc(auto& dofs, auto& bc, auto bs)
{
  for (auto dof : dofs)
    for (int k = 0; k < bs; ++k)
      if (bc[bs * dof + k])
        return true;
  return false;
};

template <bool LiftingMode, typename AB, MDSpan2Int32 XD, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<AB>::value_type>
  requires ScratchBuffer<AB, T>
void assemble_cells_matrix(
    la::MatSet<T> auto mat_set, GeometryPack<XD, U> geometry,
    const IndexList auto& cells, const FormArgumentCells<T> auto& arg0,
    const FormArgumentCells<T> auto& arg1, std::span<const std::int8_t> bc0,
    std::span<const std::int8_t> bc1, const FEkernel<T, U> auto& kernel,
    std::span<const T> constants,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs, AB Ab,
    ScratchBuffer<U> auto cdofs_b)
{
  if (std::ranges::empty(cells))
    return;
 
  // By value: the sizes below fold only if not read through a
  // reference (see fem::DofMapPack). mdspan and span are two-word
  // copies.
  const auto& P0 = arg0.transform;
  const auto& P1T = arg1.transform;
  std::span<const std::uint32_t> cell_info0 = arg0.cell_info;
  std::span<const std::uint32_t> cell_info1 = arg1.cell_info;
  const auto dmap0 = arg0.dofmap.map;
  const auto bs0 = arg0.dofmap.bs;
  // By reference: a generated range (e.g. iota) does not convert to a
  // span, and a caller's std::vector must not be copied.
  const auto& cells0 = arg0.dofmap.entities;
  const auto dmap1 = arg1.dofmap.map;
  const auto bs1 = arg1.dofmap.bs;
  const auto& cells1 = arg1.dofmap.entities;
 
  const auto num_dofs0 = dmap0.extent(1);
  const auto num_dofs1 = dmap1.extent(1);
  constexpr std::size_t nd0 = static_dofs_per_cell<decltype(dmap0)>;
  constexpr std::size_t nd1 = static_dofs_per_cell<decltype(dmap1)>;
  const std::size_t ndim0 = bs0 * num_dofs0;
  const std::size_t ndim1 = bs1 * num_dofs1;
  const auto num_x_dofs_cell = geometry.dofmap.extent(1);
 
  assert(Ab.size() == ndim0 * ndim1);
  assert(cdofs_b.size() == 3 * static_cast<std::size_t>(num_x_dofs_cell));
 
  const bool p0_set = is_transform_set(P0);
  const bool p1t_set = is_transform_set(P1T);
 
  const T* coeffs_data = coeffs.data_handle();
  const std::size_t cstride = coeffs.extent(1);
 
  // Iterate over active cells
  const std::size_t num_cells = std::ranges::size(cells);
  assert(std::ranges::size(cells0) == num_cells);
  assert(std::ranges::size(cells1) == num_cells);
  for (std::size_t c = 0; c < num_cells; ++c)
  {
    // Cell index in integration domain mesh (c), test function mesh
    // (c0) and trial function mesh (c1)
    std::int32_t cell = cells[c];
    std::int32_t cell0 = cells0[c];
    std::int32_t cell1 = cells1[c];
 
    std::span<const std::int32_t, nd0> dofs0(
        dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
    std::span<const std::int32_t, nd1> dofs1(
        dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
 
    // In "LiftingMode" only execute kernel if there are BCs on column space
    if constexpr (LiftingMode)
    {
      if (!has_bc(dofs1, bc1, bs1))
        continue;
    }
 
    gather_cell_coordinates(geometry, cell, cdofs_b.data());
 
    // Tabulate tensor
    std::ranges::fill(Ab, T(0));
    kernel(Ab.data(), coeffs_data + c * cstride, constants.data(),
           cdofs_b.data(), nullptr, nullptr, nullptr);
 
    // Compute A = P_0 \tilde{A} P_1^T (dof transformation)
    if (p0_set)
      P0(Ab, cell_info0, cell0, ndim1); // B = P0 \tilde{A}
    if (p1t_set)
      P1T(Ab, cell_info1, cell1, ndim0); // A =  B P1_T
 
    // In lifting mode only BC dofs are assembled, while in standard mode these
    // row/column dofs are zeroed.
    if constexpr (!LiftingMode)
    {
      // Zero rows and columns for BCs
      if (!bc0.empty())
      {
        for (std::size_t i = 0; i < num_dofs0; ++i)
        {
          for (int k = 0; k < bs0; ++k)
          {
            if (bc0[bs0 * dofs0[i] + k])
            {
              // Zero row bs0 * i + k
              const int row = bs0 * i + k;
              std::fill_n(std::next(Ab.begin(), ndim1 * row), ndim1, T(0));
            }
          }
        }
      }
 
      if (!bc1.empty())
      {
        for (std::size_t j = 0; j < num_dofs1; ++j)
        {
          for (int k = 0; k < bs1; ++k)
          {
            if (bc1[bs1 * dofs1[j] + k])
            {
              // Zero column bs1 * j + k
              int col = bs1 * j + k;
              for (std::size_t row = 0; row < ndim0; ++row)
                Ab[row * ndim1 + col] = 0;
            }
          }
        }
      }
    }
 
    mat_set(dofs0, dofs1, Ab);
  }
}

template <bool LiftingMode, typename AB, MDSpan2Int32 XD, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<AB>::value_type>
  requires ScratchBuffer<AB, T>
void assemble_entities_matrix(
    la::MatSet<T> auto mat_set, GeometryPack<XD, U> geometry,
    md::mdspan<const std::int32_t,
               std::extents<std::size_t, md::dynamic_extent, 2>>
        entities,
    const FormArgumentEntities<T> auto& arg0,
    const FormArgumentEntities<T> auto& arg1, std::span<const std::int8_t> bc0,
    std::span<const std::int8_t> bc1, const FEkernel<T, U> auto& kernel,
    std::span<const T> constants,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms, AB Ab,
    ScratchBuffer<U> auto cdofs_b)
{
  if (entities.empty())
    return;
 
  // By value: the sizes below fold only if not read through a
  // reference (see fem::DofMapPack). mdspan and span are two-word
  // copies.
  const auto& P0 = arg0.transform;
  const auto& P1T = arg1.transform;
  std::span<const std::uint32_t> cell_info0 = arg0.cell_info;
  std::span<const std::uint32_t> cell_info1 = arg1.cell_info;
  const auto dmap0 = arg0.dofmap.map;
  const auto bs0 = arg0.dofmap.bs;
  const auto entities0 = arg0.dofmap.entities;
  const auto dmap1 = arg1.dofmap.map;
  const auto bs1 = arg1.dofmap.bs;
  const auto entities1 = arg1.dofmap.entities;
 
  const auto num_dofs0 = dmap0.extent(1);
  const auto num_dofs1 = dmap1.extent(1);
  constexpr std::size_t nd0 = static_dofs_per_cell<decltype(dmap0)>;
  constexpr std::size_t nd1 = static_dofs_per_cell<decltype(dmap1)>;
  const std::size_t ndim0 = bs0 * num_dofs0;
  const std::size_t ndim1 = bs1 * num_dofs1;
  const auto num_x_dofs_cell = geometry.dofmap.extent(1);
  assert(entities0.size() == entities.size());
  assert(entities1.size() == entities.size());
  assert(Ab.size() == ndim0 * ndim1);
  assert(cdofs_b.size() == 3 * static_cast<std::size_t>(num_x_dofs_cell));
 
  // P0/P1T do not change across entities in this call, so whether each is a
  // set (non-null) transform is loop-invariant -- checked once here rather
  // than on every entity.
  const bool p0_set = is_transform_set(P0);
  const bool p1t_set = is_transform_set(P1T);
 
  const T* coeffs_data = coeffs.data_handle();
  const std::size_t cstride = coeffs.extent(1);
 
  for (std::size_t f = 0; f < entities.extent(0); ++f)
  {
    // Cell in the integration domain, local entity index relative to the
    // integration domain cell, and cells in the test and trial function
    // meshes
    std::int32_t cell = entities(f, 0);
    std::int32_t local_entity = entities(f, 1);
    std::int32_t cell0 = entities0(f, 0);
    std::int32_t cell1 = entities1(f, 0);
 
    std::span<const std::int32_t, nd0> dofs0(
        dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
    std::span<const std::int32_t, nd1> dofs1(
        dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
 
    // Check for BCs on column space
    if constexpr (LiftingMode)
    {
      if (!has_bc(dofs1, bc1, bs1))
        continue;
    }
 
    gather_cell_coordinates(geometry, cell, cdofs_b.data());
 
    // Permutations
    std::uint8_t perm = perms.empty() ? 0 : perms(cell, local_entity);
 
    // Tabulate tensor
    std::ranges::fill(Ab, T(0));
    kernel(Ab.data(), coeffs_data + f * cstride, constants.data(),
           cdofs_b.data(), &local_entity, &perm, nullptr);
    if (p0_set)
      P0(Ab, cell_info0, cell0, ndim1);
    if (p1t_set)
      P1T(Ab, cell_info1, cell1, ndim0);
 
    // Don't clear rows/cols in LiftingMode
    if constexpr (!LiftingMode)
    {
      // Zero rows and columns for BCs
      if (!bc0.empty())
      {
        for (std::size_t i = 0; i < num_dofs0; ++i)
        {
          for (int k = 0; k < bs0; ++k)
          {
            if (bc0[bs0 * dofs0[i] + k])
            {
              // Zero row bs0 * i + k
              const int row = bs0 * i + k;
              std::fill_n(std::next(Ab.begin(), ndim1 * row), ndim1, T(0));
            }
          }
        }
      }
 
      if (!bc1.empty())
      {
        for (std::size_t j = 0; j < num_dofs1; ++j)
        {
          for (int k = 0; k < bs1; ++k)
          {
            if (bc1[bs1 * dofs1[j] + k])
            {
              // Zero column bs1 * j + k
              int col = bs1 * j + k;
              for (std::size_t row = 0; row < ndim0; ++row)
                Ab[row * ndim1 + col] = 0;
            }
          }
        }
      }
    }
 
    mat_set(dofs0, dofs1, Ab);
  }
}

template <bool LiftingMode, typename AB, MDSpan2Int32 XD, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<AB>::value_type>
  requires ScratchBuffer<AB, T>
void assemble_interior_facets_matrix(
    la::MatSet<T> auto mat_set, GeometryPack<XD, U> geometry,
    md::mdspan<const std::int32_t,
               std::extents<std::size_t, md::dynamic_extent, 2, 2>>
        facets,
    const FormArgumentFacets<T> auto& arg0,
    const FormArgumentFacets<T> auto& arg1, std::span<const std::int8_t> bc0,
    std::span<const std::int8_t> bc1, const FEkernel<T, U> auto& kernel,
    std::span<const T> constants,
    md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                    md::dynamic_extent>>
        coeffs,
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms, AB Ab,
    ScratchBuffer<U> auto cdofs_b, ScratchBuffer<std::int32_t> auto dofs_b,
    ScratchBuffer<T> auto Ae_block_b)
{
  if (facets.empty())
    return;
 
  // By value: the sizes below fold only if not read through a
  // reference (see fem::DofMapPack). mdspan and span are two-word
  // copies.
  const auto& P0 = arg0.transform;
  const auto& P1T = arg1.transform;
  std::span<const std::uint32_t> cell_info0 = arg0.cell_info;
  std::span<const std::uint32_t> cell_info1 = arg1.cell_info;
  const auto dmap0 = arg0.dofmap.map;
  const auto bs0 = arg0.dofmap.bs;
  const auto facets0 = arg0.dofmap.entities;
  const auto dmap1 = arg1.dofmap.map;
  const auto bs1 = arg1.dofmap.bs;
  const auto facets1 = arg1.dofmap.entities;
 
  // Data structures used in assembly
  const auto num_x_dofs_cell = geometry.dofmap.extent(1);
  assert(cdofs_b.size() == 2 * 3 * static_cast<std::size_t>(num_x_dofs_cell));
  U* cdofs0 = cdofs_b.data();
  U* cdofs1 = cdofs_b.data() + 3 * num_x_dofs_cell;
 
  const auto dmap0_size = dmap0.extent(1);
  const auto dmap1_size = dmap1.extent(1);
  // The joint (two-cell) dofmap length is a constant exactly when the
  // per-cell one is; dynamic_extent propagates through otherwise.
  constexpr std::size_t nd0
      = std::remove_cvref_t<decltype(dmap0)>::static_extent(1);
  constexpr std::size_t nd1
      = std::remove_cvref_t<decltype(dmap1)>::static_extent(1);
  constexpr std::size_t njoint0
      = nd0 == md::dynamic_extent ? std::dynamic_extent : 2 * nd0;
  constexpr std::size_t njoint1
      = nd1 == md::dynamic_extent ? std::dynamic_extent : 2 * nd1;
  std::size_t num_rows = bs0 * 2 * dmap0_size;
  std::size_t num_cols = bs1 * 2 * dmap1_size;
 
  // Dofmap data structures
  assert(dofs_b.size() == (2 * dmap0_size) + (2 * dmap1_size));
  std::span<std::int32_t> dofs_all(dofs_b);
  std::span<std::int32_t, njoint0> dmapjoint0(dofs_all.data(), 2 * dmap0_size);
  std::span<std::int32_t, njoint1> dmapjoint1(dofs_all.data() + 2 * dmap0_size,
                                              2 * dmap1_size);
 
  assert(facets0.size() == facets.size());
  assert(facets1.size() == facets.size());
  assert(Ab.size() == num_rows * num_cols);
 
  // Buffer used to gather a contiguous (test, trial) block of Ab when
  // one of the two cells attached to the facet does not exist in the
  // test/trial function domain (e.g. an interface between two
  // domains) -- the sparsity pattern only holds entries for blocks
  // where both cells exist, so such blocks must be inserted
  // individually rather than as part of the full joint block.
  assert(Ae_block_b.size() == dmap0_size * bs0 * dmap1_size * bs1);
  std::span<T> Ae_block_all(Ae_block_b);
 
  const T* coeffs_data = coeffs.data_handle();
  const std::size_t cstride = 2 * coeffs.extent(2);
 
  auto insert_block = [&Ae_block_all, &Ab, &bs0, &bs1, &num_cols,
                       &mat_set](std::span<const std::int32_t> rdofs,
                                 std::span<const std::int32_t> cdofs,
                                 std::size_t row_offset, std::size_t col_offset)
  {
    if (rdofs.empty() or cdofs.empty())
      return;
    auto Ae_block = Ae_block_all.first(rdofs.size() * bs0 * cdofs.size() * bs1);
    for (std::size_t i = 0; i < rdofs.size() * bs0; ++i)
    {
      auto row
          = std::next(Ab.begin(), (row_offset + i) * num_cols + col_offset);
      std::copy_n(row, cdofs.size() * bs1,
                  std::next(Ae_block.begin(), i * cdofs.size() * bs1));
    }
    mat_set(rdofs, cdofs, Ae_block);
  };
 
  // P0/P1T do not change across facets in this call, so whether each is a
  // set (non-null) transform is loop-invariant -- checked once here rather
  // than on every facet.
  const bool p0_set = is_transform_set(P0);
  const bool p1t_set = is_transform_set(P1T);
 
  for (std::size_t f = 0; f < facets.extent(0); ++f)
  {
    // Cells in integration domain,  test function domain and trial
    // function domain
    std::array cells{facets(f, 0, 0), facets(f, 1, 0)};
    std::array cells0{facets0(f, 0, 0), facets0(f, 1, 0)};
    std::array cells1{facets1(f, 0, 0), facets1(f, 1, 0)};
 
    // Local facets indices
    std::array local_facet{facets(f, 0, 1), facets(f, 1, 1)};
 
    // Get cell geometry
    gather_cell_coordinates(geometry, cells[0], cdofs0);
    gather_cell_coordinates(geometry, cells[1], cdofs1);
 
    // Get dof maps for cells and pack
    // When integrating over interfaces between two domains, the test function
    // might only be defined on one side, so we check which cells exist in the
    // test function domain
    std::span<const std::int32_t> dmap0_cell0
        = cells0[0] >= 0
              ? std::span(dmap0.data_handle() + cells0[0] * dmap0_size,
                          dmap0_size)
              : std::span<const std::int32_t>();
    std::span<const std::int32_t> dmap0_cell1
        = cells0[1] >= 0
              ? std::span(dmap0.data_handle() + cells0[1] * dmap0_size,
                          dmap0_size)
              : std::span<const std::int32_t>();
 
    std::ranges::copy(dmap0_cell0, dmapjoint0.begin());
    std::ranges::copy(dmap0_cell1, std::next(dmapjoint0.begin(), dmap0_size));
 
    // Check which cells exist in the trial function domain
    std::span<const std::int32_t> dmap1_cell0
        = cells1[0] >= 0
              ? std::span(dmap1.data_handle() + cells1[0] * dmap1_size,
                          dmap1_size)
              : std::span<const std::int32_t>();
    std::span<const std::int32_t> dmap1_cell1
        = cells1[1] >= 0
              ? std::span(dmap1.data_handle() + cells1[1] * dmap1_size,
                          dmap1_size)
              : std::span<const std::int32_t>();
 
    // Check for BCs on column space. The per-cell dofmaps are tested
    // rather than the joint one: a side absent from the trial function
    // domain leaves its half of the joint buffer holding the previous
    // facet's dofs, which would spuriously run the kernel.
    if constexpr (LiftingMode)
    {
      if (!has_bc(dmap1_cell0, bc1, bs1) and !has_bc(dmap1_cell1, bc1, bs1))
        continue;
    }
 
    std::ranges::copy(dmap1_cell0, dmapjoint1.begin());
    std::ranges::copy(dmap1_cell1, std::next(dmapjoint1.begin(), dmap1_size));
 
    // Tabulate tensor
    std::ranges::fill(Ab, T(0));
    std::array perm = perms.empty()
                          ? std::array<std::uint8_t, 2>{0, 0}
                          : std::array{perms(cells[0], local_facet[0]),
                                       perms(cells[1], local_facet[1])};
    kernel(Ab.data(), coeffs_data + f * cstride, constants.data(),
           cdofs_b.data(), local_facet.data(), perm.data(), nullptr);
 
    // Local element layout is a 2x2 block matrix with structure
    //
    //   cell0cell0  |  cell0cell1
    //   cell1cell0  |  cell1cell1
    //
    // where each block is element tensor of size (dmap0, dmap1).
 
    // Only apply transformation when cells exist
    if (p0_set and cells0[0] >= 0)
      P0(Ab, cell_info0, cells0[0], num_cols);
    if (p0_set and cells0[1] >= 0)
    {
      std::span sub_Ae0(Ab.data() + bs0 * dmap0_size * num_cols,
                        bs0 * dmap0_size * num_cols);
      P0(sub_Ae0, cell_info0, cells0[1], num_cols);
    }
    if (p1t_set and cells1[0] >= 0)
      P1T(Ab, cell_info1, cells1[0], num_rows);
 
    if (p1t_set and cells1[1] >= 0)
    {
      for (std::size_t row = 0; row < num_rows; ++row)
      {
        // DOFs for dmap1 and cell1 are not stored contiguously in the
        // block matrix, so each row needs a separate span access
        std::span sub_Ae1(Ab.data() + row * num_cols + bs1 * dmap1_size,
                          bs1 * dmap1_size);
        P1T(sub_Ae1, cell_info1, cells1[1], 1);
      }
    }
 
    // Clear rows/cols if not in LiftingMode
    if constexpr (!LiftingMode)
    {
      // Zero rows and columns for BCs
      if (!bc0.empty())
      {
        for (std::size_t i = 0; i < dmapjoint0.size(); ++i)
        {
          for (int k = 0; k < bs0; ++k)
          {
            if (bc0[bs0 * dmapjoint0[i] + k])
            {
              // Zero row bs0 * i + k
              std::fill_n(std::next(Ab.begin(), num_cols * (bs0 * i + k)),
                          num_cols, T(0));
            }
          }
        }
      }
 
      if (!bc1.empty())
      {
        for (std::size_t j = 0; j < dmapjoint1.size(); ++j)
        {
          for (int k = 0; k < bs1; ++k)
          {
            if (bc1[bs1 * dmapjoint1[j] + k])
            {
              // Zero column bs1 * j + k
              for (std::size_t m = 0; m < num_rows; ++m)
                Ab[m * num_cols + bs1 * j + k] = 0;
            }
          }
        }
      }
    }
 
    // The common case is that a cell exists on both sides of the
    // facet for both the test and trial function domains, in which
    // case the full joint block can be inserted in one go. Otherwise
    // (e.g. an interface between two domains), only the blocks
    // corresponding to existing (test, trial) cell pairs are present
    // in the sparsity pattern, so each must be inserted individually.
    if (cells0[0] >= 0 and cells0[1] >= 0 and cells1[0] >= 0 and cells1[1] >= 0)
      mat_set(dmapjoint0, dmapjoint1, Ab);
    else
    {
      insert_block(dmap0_cell0, dmap1_cell0, 0, 0);
      insert_block(dmap0_cell0, dmap1_cell1, 0, bs1 * dmap1_size);
      insert_block(dmap0_cell1, dmap1_cell0, bs0 * dmap0_size, 0);
      insert_block(dmap0_cell1, dmap1_cell1, bs0 * dmap0_size,
                   bs1 * dmap1_size);
    }
  }
}

template <bool LiftingMode, dolfinx::scalar T, std::floating_point U>
void assemble_matrix(
    la::MatSet<T> auto mat_set, const Form<T, U>& a,
    md::mdspan<const U, md::extents<std::size_t, md::dynamic_extent, 3>> x,
    std::span<const T> constants,
    const std::map<std::pair<IntegralType, int>,
                   std::pair<std::span<const T>, int>>& coefficients,
    std::span<const std::int8_t> bc0, std::span<const std::int8_t> bc1)
{
  // Integration domain mesh
  std::shared_ptr<const mesh::Mesh<U>> mesh = a.mesh();
  assert(mesh);
 
  // Test function mesh
  auto mesh0 = a.function_spaces().at(0)->mesh();
  assert(mesh0);
 
  // Trial function mesh
  auto mesh1 = a.function_spaces().at(1)->mesh();
  assert(mesh1);
 
  // TODO: Mixed topology with exterior and interior facet integrals.
  //
  // NOTE: Can't just loop over cell types for interior facet integrals
  // because we have a kernel per combination of comparable cell types,
  // rather than one per cell type. Also, we need the dofmaps for two
  // different cell types at the same time.
  const int num_cell_types = mesh->topology()->cell_types().size();
  for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
  {
    // Geometry dofmap and data
    mdspan2_t x_dofmap = mesh->geometry().dofmaps().at(cell_type_idx);
    GeometryPack geometry{x_dofmap, x};
 
    // Get dofmap data
    std::shared_ptr<const fem::DofMap> dofmap0
        = a.function_spaces().at(0)->dofmaps().at(cell_type_idx);
    std::shared_ptr<const fem::DofMap> dofmap1
        = a.function_spaces().at(1)->dofmaps().at(cell_type_idx);
    assert(dofmap0);
    assert(dofmap1);
    md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> dofs0
        = dofmap0->map();
    const int bs0 = dofmap0->bs();
    md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> dofs1
        = dofmap1->map();
    const int bs1 = dofmap1->bs();
 
    // Buffers reused across all integral kernels for this cell type,
    // sized for the worst case (interior facets, which touch two
    // cells). The kernels require an exactly-sized buffer, so the
    // one-cell integrals get a leading view.
    std::vector<T> Ab((2 * bs0 * dofs0.extent(1))
                      * (2 * bs1 * dofs1.extent(1)));
    std::vector<U> cdofs_b(2 * 3 * x_dofmap.extent(1));
    std::span Ab1 = std::span(Ab).first((bs0 * dofs0.extent(1))
                                        * (bs1 * dofs1.extent(1)));
    std::span cdofs_b1 = std::span(cdofs_b).first(3 * x_dofmap.extent(1));
    std::size_t dmap0_size = dofmap0->map().extent(1);
    std::size_t dmap1_size = dofmap1->map().extent(1);
    std::vector<std::int32_t> dmap_b((2 * dmap0_size) + (2 * dmap1_size));
    std::vector<T> Ae_block_b(dmap0_size * bs0 * dmap1_size * bs1);
 
    auto element0 = a.function_spaces().at(0)->elements(cell_type_idx);
    assert(element0);
    auto element1 = a.function_spaces().at(1)->elements(cell_type_idx);
    assert(element1);
    const fem::DofTransformKernel<T> auto& P0
        = element0->template dof_transformation_fn<T>(doftransform::standard);
    const fem::DofTransformKernel<T> auto& P1T
        = element1->template dof_transformation_right_fn<T>(
            doftransform::transpose);
 
    std::span<const std::uint32_t> cell_info0;
    std::span<const std::uint32_t> cell_info1;
    if (element0->needs_dof_transformations()
        or element1->needs_dof_transformations())
    {
      mesh0->topology_mutable()->create_cell_permutations();
      mesh1->topology_mutable()->create_cell_permutations();
      cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info());
      cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info());
    }
 
    for (int i = 0; i < a.num_integrals(IntegralType::cell, cell_type_idx); ++i)
    {
      auto fn = a.kernel(IntegralType::cell, i, cell_type_idx);
      assert(fn);
      std::span cells = a.domain(IntegralType::cell, i, cell_type_idx);
      std::span cells0 = a.domain_arg(IntegralType::cell, 0, i, cell_type_idx);
      std::span cells1 = a.domain_arg(IntegralType::cell, 1, i, cell_type_idx);
      auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
      assert(cells.size() * cstride == coeffs.size());
      impl::dispatch_bs(
          bs0, bs1,
          [&mat_set, &geometry, &cells, &dofs0, &cells0, &P0, &cell_info0,
           &dofs1, &cells1, &P1T, &cell_info1, &bc0, &bc1, &fn, &coeffs,
           cstride, &constants, &Ab1, &cdofs_b1](auto bs0, auto bs1)
          {
            impl::assemble_cells_matrix<LiftingMode>(
                mat_set, geometry, cells,
                FormArgument{DofMapPack{dofs0, bs0, cells0}, P0, cell_info0},
                FormArgument{DofMapPack{dofs1, bs1, cells1}, P1T, cell_info1},
                bc0, bc1, fn, constants,
                md::mdspan(coeffs.data(), cells.size(), cstride), Ab1,
                cdofs_b1);
          });
    }
 
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> facet_perms;
    if (a.needs_facet_permutations())
    {
      facet_perms = impl::entity_permutations(
          *mesh->topology_mutable(), IntegralType::interior_facet,
          mesh->topology()->cell_types()[cell_type_idx]);
    }
 
    for (int i = 0;
         i < a.num_integrals(IntegralType::interior_facet, cell_type_idx); ++i)
    {
      if (num_cell_types > 1)
      {
        throw std::invalid_argument("Interior facet integrals with mixed "
                                    "topology aren't supported yet");
      }
 
      using mdspanx22_t
          = md::mdspan<const std::int32_t,
                       md::extents<std::size_t, md::dynamic_extent, 2, 2>>;
      using mdspanx2x_t
          = md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                            md::dynamic_extent>>;
 
      auto fn = a.kernel(IntegralType::interior_facet, i, 0);
      assert(fn);
      auto& [coeffs, cstride]
          = coefficients.at({IntegralType::interior_facet, i});
 
      std::span facets = a.domain(IntegralType::interior_facet, i, 0);
      std::span facets0 = a.domain_arg(IntegralType::interior_facet, 0, i, 0);
      std::span facets1 = a.domain_arg(IntegralType::interior_facet, 1, i, 0);
      assert((facets.size() / 4) * 2 * cstride == coeffs.size());
      impl::dispatch_bs(
          bs0, bs1,
          [&mat_set, &geometry, &facets, &dofs0, &facets0, &P0, &cell_info0,
           &dofs1, &facets1, &P1T, &cell_info1, &bc0, &bc1, &fn, &coeffs,
           cstride, &constants, &facet_perms, &Ab, &cdofs_b, &dmap_b,
           &Ae_block_b](auto bs0, auto bs1)
          {
            impl::assemble_interior_facets_matrix<LiftingMode>(
                mat_set, geometry,
                mdspanx22_t(facets.data(), facets.size() / 4, 2, 2),
                FormArgument{DofMapPack{dofs0, bs0,
                                        mdspanx22_t(facets0.data(),
                                                    facets0.size() / 4, 2, 2)},
                             P0, cell_info0},
                FormArgument{DofMapPack{dofs1, bs1,
                                        mdspanx22_t(facets1.data(),
                                                    facets1.size() / 4, 2, 2)},
                             P1T, cell_info1},
                bc0, bc1, fn, constants,
                mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride),
                facet_perms, std::span(Ab), std::span(cdofs_b),
                std::span(dmap_b), std::span(Ae_block_b));
          });
    }
 
    for (auto itg_type : {fem::IntegralType::exterior_facet,
                          fem::IntegralType::vertex, fem::IntegralType::ridge})
    {
      const int num_itg = a.num_integrals(itg_type, cell_type_idx);
      if (num_itg == 0)
        continue;
 
      // Each integral type is over entities of a different
      // codimension, so only the permutations this form actually
      // integrates over are computed.
      md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms;
      if (a.needs_facet_permutations())
      {
        perms = impl::entity_permutations(
            *mesh->topology_mutable(), itg_type,
            mesh->topology()->cell_types()[cell_type_idx]);
      }
 
      for (int i = 0; i < num_itg; ++i)
      {
        if (num_cell_types > 1)
        {
          throw std::invalid_argument("Exterior facet integrals with mixed "
                                      "topology aren't supported yet");
        }
 
        using mdspanx2_t
            = md::mdspan<const std::int32_t,
                         md::extents<std::size_t, md::dynamic_extent, 2>>;
 
        auto fn = a.kernel(itg_type, i, 0);
        assert(fn);
        auto& [coeffs, cstride] = coefficients.at({itg_type, i});
 
        std::span e = a.domain(itg_type, i, 0);
        mdspanx2_t entities(e.data(), e.size() / 2, 2);
        std::span e0 = a.domain_arg(itg_type, 0, i, 0);
        mdspanx2_t entities0(e0.data(), e0.size() / 2, 2);
        std::span e1 = a.domain_arg(itg_type, 1, i, 0);
        mdspanx2_t entities1(e1.data(), e1.size() / 2, 2);
        assert((entities.size() / 2) * cstride == coeffs.size());
        impl::dispatch_bs(
            bs0, bs1,
            [&mat_set, &geometry, &entities, &dofs0, &entities0, &P0,
             &cell_info0, &dofs1, &entities1, &P1T, &cell_info1, &bc0, &bc1,
             &fn, &coeffs, cstride, &constants, &perms, &Ab1,
             &cdofs_b1](auto bs0, auto bs1)
            {
              impl::assemble_entities_matrix<LiftingMode>(
                  mat_set, geometry, entities,
                  FormArgument{DofMapPack{dofs0, bs0, entities0}, P0,
                               cell_info0},
                  FormArgument{DofMapPack{dofs1, bs1, entities1}, P1T,
                               cell_info1},
                  bc0, bc1, fn, constants,
                  md::mdspan(coeffs.data(), entities.extent(0), cstride), perms,
                  Ab1, cdofs_b1);
            });
      }
    }
  }
}
} // namespace dolfinx::fem::impl