phynexis-deps/mfem
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
c577ee0f43500a2fade34e62cc55cacb44be0840
mfem/fem/hybridization_ext.cpp
T

1334 lines
41 KiB
C++
Raw Blame History

// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
// LICENSE and NOTICE for details. LLNL-CODE-806117.
//
// This file is part of the MFEM library. For more information and source code
// availability visit https://mfem.org.
//
// MFEM is free software; you can redistribute it and/or modify it under the
// terms of the BSD-3 license. We welcome feedback and contributions, see file
// CONTRIBUTING.md for details.
#include "hybridization_ext.hpp"
#include "hybridization.hpp"
#include "pfespace.hpp"
#include "../general/forall.hpp"
#include "../linalg/batched/batched.hpp"
#include "../linalg/kernels.hpp"
namespace mfem
{
HybridizationExtension::HybridizationExtension(Hybridization &hybridization_)
: h(hybridization_)
{ }
static int GetNFacesPerElement(const Mesh &mesh)
{
const int dim = mesh.Dimension();
switch (dim)
{
case 2: return mesh.GetElement(0)->GetNEdges();
case 3: return mesh.GetElement(0)->GetNFaces();
default: MFEM_ABORT("Invalid dimension.");
}
}
static bool IsParFESpace(const FiniteElementSpace &fes)
{
#ifdef MFEM_USE_MPI
return dynamic_cast<const ParFiniteElementSpace *>(&fes) != nullptr;
#else
return false;
#endif
}
const Operator &HybridizationExtension::GetProlongation() const
{
return P_pc ? *P_pc : *h.c_fes.GetProlongationMatrix();
}
void HybridizationExtension::ConstructC()
{
Mesh &mesh = *h.fes.GetMesh();
const int nf = mesh.GetNFbyType(FaceType::Interior);
const int m = h.fes.GetFE(0)->GetDof(); // num hat dofs per el
const int n = h.c_fes.GetFaceElement(0)->GetDof(); // num c dofs per face
// Assemble Ct_mat using EA
Vector emat(m * n * 2 * nf);
h.c_bfi->AssembleEAInteriorFaces(h.c_fes, h.fes, emat, false);
const auto *tbe = dynamic_cast<const TensorBasisElement*>(h.fes.GetFE(0));
MFEM_VERIFY(tbe, "");
// Note: copying the DOF map here (instead of using a reference) because
// reading it on GPU can cause issues in other parts of the code when using
// the debug device. The DOF map is accessed in other places without
// explicitly calling HostRead, which fails on non-const access if the device
// pointer is valid.
Array<int> dof_map = tbe->GetDofMap();
Ct_mat.SetSize(m * n * n_el_face);
const auto d_emat = Reshape(emat.Read(), m, n, 2, nf);
const int *d_dof_map = dof_map.Read();
const auto d_face_to_el = Reshape(face_to_el.Read(), 2, 2, nf);
auto d_Ct_mat = Reshape(Ct_mat.Write(), m, n, n_el_face);
Ct_mat = 0.0; // On device, since previous call to Write()
mfem::forall(m*n*2*nf, [=] MFEM_HOST_DEVICE (int idx)
{
const int i_lex = idx % m;
const int j = (idx / m) % n;
const int ie = (idx / m / n) % 2;
const int f = idx / m / n / 2;
const int fi = d_face_to_el(1, ie, f);
// Skip elements belonging to face neighbors of shared faces
if (fi >= 0)
{
// Convert to back to native MFEM ordering in the volume
const int i_s = d_dof_map[i_lex];
const int i = (i_s >= 0) ? i_s : -1 - i_s;
d_Ct_mat(i, j, fi) = d_emat(i_lex, j, ie, f);
}
});
#ifdef MFEM_USE_MPI
if (auto pc_fes = dynamic_cast<ParFiniteElementSpace*>(&h.c_fes))
{
if (pc_fes->Nonconforming())
{
P_pc.reset(pc_fes->GetPartialConformingInterpolation());
}
}
#endif
}
namespace internal
{
template <typename T, int SIZE>
struct LocalMemory
{
T data[SIZE];
MFEM_HOST_DEVICE inline operator T *() const { return (T*)data; }
};
template <typename T>
struct LocalMemory<T,0>
{
MFEM_HOST_DEVICE inline operator T *() const { return (T*)nullptr; }
};
};
template <int MID, int MBD>
void HybridizationExtension::FactorElementMatrices(Vector &AhatInvCt_mat)
{
const Mesh &mesh = *h.fes.GetMesh();
const int ne = mesh.GetNE();
const int m = h.fes.GetFE(0)->GetDof();
const int n = h.c_fes.GetFaceElement(0)->GetDof();
AhatInvCt_mat.SetSize(Ct_mat.Size());
auto d_AhatInvCt = Reshape(AhatInvCt_mat.Write(), m, n, n_el_face);
const int nidofs = idofs.Size();
const int nbdofs = bdofs.Size();
MFEM_VERIFY(nidofs <= MID, "");
MFEM_VERIFY(nbdofs <= MBD, "");
Ahat_ii.SetSize(nidofs*nidofs*ne);
Ahat_ib.SetSize(nidofs*nbdofs*ne);
Ahat_bi.SetSize(nbdofs*nidofs*ne);
Ahat_bb.SetSize(nbdofs*nbdofs*ne);
Ahat_ii_piv.SetSize(nidofs*ne);
Ahat_bb_piv.SetSize(nbdofs*ne);
const auto *d_idofs = idofs.Read();
const auto *d_bdofs = bdofs.Read();
const auto d_hat_dof_marker = Reshape(hat_dof_marker.Read(), m, ne);
auto d_Ahat = Reshape(Ahat.Read(), m, m, ne);
auto d_A_ii = Reshape(Ahat_ii.Write(), nidofs, nidofs, ne);
auto d_A_ib_all = Reshape(Ahat_ib.Write(), nidofs*nbdofs, ne);
auto d_A_bi_all = Reshape(Ahat_bi.Write(), nbdofs*nidofs, ne);
auto d_A_bb_all = Reshape(Ahat_bb.Write(), nbdofs*nbdofs, ne);
auto d_ipiv_ii = Reshape(Ahat_ii_piv.Write(), nidofs, ne);
auto d_ipiv_bb = Reshape(Ahat_bb_piv.Write(), nbdofs, ne);
const auto d_Ct_mat = Reshape(Ct_mat.Read(), m, n, n_el_face);
const auto d_el_face_offsets = el_face_offsets.Read();
static constexpr bool GLOBAL = (MID == 0 && MBD == 0);
using internal::LocalMemory;
mfem::forall(ne, [=] MFEM_HOST_DEVICE (int e)
{
constexpr int MD1D = DofQuadLimits::HDIV_MAX_D1D;
constexpr int MAX_DOFS = 3*MD1D*(MD1D-1)*(MD1D-1);
constexpr int MAX_IDOFS = (MID == 0 && MBD == 0) ? MAX_DOFS : MID;
constexpr int MAX_BDOFS = (MID == 0 && MBD == 0) ? MAX_DOFS : MBD;
LocalMemory<int,MAX_IDOFS> idofs_loc;
LocalMemory<int,MAX_BDOFS> bdofs_loc;
for (int i = 0; i < nidofs; i++) { idofs_loc[i] = d_idofs[i]; }
for (int i = 0; i < nbdofs; i++) { bdofs_loc[i] = d_bdofs[i]; }
LocalMemory<int,MAX_BDOFS> essdofs_loc;
int nbfdofs = 0;
int nessdofs = 0;
for (int i = 0; i < nbdofs; i++)
{
const int dof_idx = bdofs_loc[i];
if (d_hat_dof_marker(dof_idx, e) == ESSENTIAL)
{
essdofs_loc[nessdofs] = dof_idx;
nessdofs += 1;
}
else
{
bdofs_loc[nbfdofs] = dof_idx;
nbfdofs += 1;
}
}
LocalMemory<real_t, MID*MID> A_ii_loc;
LocalMemory<real_t, MBD*MID> A_bi_loc;
LocalMemory<real_t, MID*MBD> A_ib_loc;
LocalMemory<real_t, MBD*MBD> A_bb_loc;
DeviceMatrix A_ii(GLOBAL ? &d_A_ii(0,0,e) : A_ii_loc, nidofs, nidofs);
DeviceMatrix A_ib(GLOBAL ? &d_A_ib_all(0,e) : A_ib_loc, nidofs, nbfdofs);
DeviceMatrix A_bi(GLOBAL ? &d_A_bi_all(0,e) : A_bi_loc, nbfdofs, nidofs);
DeviceMatrix A_bb(GLOBAL ? &d_A_bb_all(0,e) : A_bb_loc, nbfdofs, nbfdofs);
for (int j = 0; j < nidofs; j++)
{
const int jj = idofs_loc[j];
for (int i = 0; i < nidofs; i++)
{
A_ii(i,j) = d_Ahat(idofs_loc[i], jj, e);
}
for (int i = 0; i < nbfdofs; i++)
{
A_bi(i,j) = d_Ahat(bdofs_loc[i], jj, e);
}
}
for (int j = 0; j < nbfdofs; j++)
{
const int jj = bdofs_loc[j];
for (int i = 0; i < nidofs; i++)
{
A_ib(i,j) = d_Ahat(idofs_loc[i], jj, e);
}
for (int i = 0; i < nbfdofs; i++)
{
A_bb(i,j) = d_Ahat(bdofs_loc[i], jj, e);
}
}
LocalMemory<int,MID> ipiv_ii_loc;
LocalMemory<int,MBD> ipiv_bb_loc;
auto ipiv_ii = GLOBAL ? &d_ipiv_ii(0,e) : ipiv_ii_loc;
auto ipiv_bb = GLOBAL ? &d_ipiv_ii(0,e) : ipiv_bb_loc;
kernels::LUFactor(A_ii, nidofs, ipiv_ii);
kernels::BlockFactor(A_ii, nidofs, ipiv_ii, nbfdofs, A_ib, A_bi, A_bb);
kernels::LUFactor(A_bb, nbfdofs, ipiv_bb);
const int begin = d_el_face_offsets[e];
const int end = d_el_face_offsets[e + 1];
for (int f = begin; f < end; ++f)
{
for (int j = 0; j < n; ++j)
{
LocalMemory<real_t,MAX_BDOFS> Sb_inv_Cb_t;
for (int i = 0; i < nbfdofs; ++i)
{
Sb_inv_Cb_t[i] = d_Ct_mat(bdofs_loc[i], j, f);
}
kernels::LUSolve(A_bb, nbfdofs, ipiv_bb, Sb_inv_Cb_t);
for (int i = 0; i < nbfdofs; ++i)
{
const int b_i = bdofs_loc[i];
d_AhatInvCt(b_i, j, f) = Sb_inv_Cb_t[i];
}
for (int i = 0; i < nidofs; ++i)
{
d_AhatInvCt(idofs_loc[i], j, f) = 0.0;
}
for (int i = 0; i < nessdofs; ++i)
{
d_AhatInvCt(essdofs_loc[i], j, f) = 0.0;
}
}
}
// Write out to global memory
if (!GLOBAL)
{
// Note: in the following constructors, avoid using index 0 in
// d_A_{bi,ib,bb}_all when their size is 0.
DeviceMatrix d_A_bi((nbfdofs && nidofs) ?
&d_A_bi_all(0,e) : nullptr,
nbfdofs, nidofs);
DeviceMatrix d_A_ib((nbfdofs && nidofs) ?
&d_A_ib_all(0,e) : nullptr,
nidofs, nbfdofs);
DeviceMatrix d_A_bb((nbfdofs) ? &d_A_bb_all(0,e) : nullptr,
nbfdofs, nbfdofs);
for (int j = 0; j < nidofs; j++)
{
d_ipiv_ii(j,e) = ipiv_ii[j];
for (int i = 0; i < nidofs; i++)
{
d_A_ii(i,j,e) = A_ii(i,j);
}
for (int i = 0; i < nbfdofs; i++)
{
d_A_bi(i,j) = A_bi(i,j);
}
}
for (int j = 0; j < nbfdofs; j++)
{
d_ipiv_bb(j,e) = ipiv_bb[j];
for (int i = 0; i < nidofs; i++)
{
d_A_ib(i,j) = A_ib(i,j);
}
for (int i = 0; i < nbfdofs; i++)
{
d_A_bb(i,j) = A_bb(i,j);
}
}
}
});
}
void HybridizationExtension::ConstructH()
{
const Mesh &mesh = *h.fes.GetMesh();
const int ne = mesh.GetNE();
const int m = h.fes.GetFE(0)->GetDof();
const int n = h.c_fes.GetFaceElement(0)->GetDof();
Vector AhatInvCt_mat;
{
// The dispatch below is based on the following sizes, sorted
// appropriately.
//
// RT(k) in 2D (quads): (interior,boundary) dofs:
// - arbitrary k: 2*(k+1)*(k+2)-4*(k+1), 4*(k+1)
// - k=0: (0,4)
// - k=1: (4,8)
// - k=2: (12,12)
// - k=3: (24,16)
// RT(k) in 3D (hexes): (interior,boundary) dofs:
// - arbitrary k: 3*(k+1)^2*(k+2)-6*(k+1)^2, 6*(k+1)^2
// - k=0: (0,6)
// - k=1: (12,24)
// - k=2: (54,54)
const int NI = idofs.Size();
const int NB = bdofs.Size();
if (NI == 0 && NB <= 4) { FactorElementMatrices<0,4>(AhatInvCt_mat); }
else if (NI == 0 && NB <= 6) { FactorElementMatrices<0,6>(AhatInvCt_mat); }
else if (NI <= 4 && NB <= 8) { FactorElementMatrices<4,8>(AhatInvCt_mat); }
else if (NI <= 12 && NB <= 12) { FactorElementMatrices<12,12>(AhatInvCt_mat); }
else if (NI <= 12 && NB <= 24) { FactorElementMatrices<12,24>(AhatInvCt_mat); }
else if (NI <= 24 && NB <= 16) { FactorElementMatrices<24,16>(AhatInvCt_mat); }
else if (NI <= 54 && NB <= 54) { FactorElementMatrices<54,54>(AhatInvCt_mat); }
// Fallback
else { FactorElementMatrices<0,0>(AhatInvCt_mat); }
}
const auto d_AhatInvCt =
Reshape(AhatInvCt_mat.Read(), m, n, n_el_face);
const int nf = h.fes.GetNFbyType(FaceType::Interior);
Array<int> face_to_face(n_face_face);
Vector CAhatInvCt(n_face_face*n*n);
const auto d_Ct = Reshape(Ct_mat.Read(), m, n, n_el_face);
const auto d_face_to_el = Reshape(face_to_el.Read(), 2, 2, nf);
const auto d_el_to_face = el_to_face.Read();
const auto d_el_face_offsets = el_face_offsets.Read();
auto d_CAhatInvCt = Reshape(CAhatInvCt.Write(), n, n, n_face_face);
auto d_face_to_face = Reshape(face_to_face.Write(), n_face_face);
auto d_face_face_offsets = face_face_offsets.Read();
CAhatInvCt = 0.0;
mfem::forall(nf, [=] MFEM_HOST_DEVICE (int fi)
{
const int begin_f = d_face_face_offsets[fi];
int idx = 0;
for (int ei = 0; ei < 2; ++ei)
{
const int e = d_face_to_el(0, ei, fi);
if (e < 0 || e >= ne) { continue; }
const int begin_i = d_el_face_offsets[e];
const int end_i = d_el_face_offsets[e + 1];
for (int fj_i = begin_i; fj_i < end_i; ++fj_i)
{
const int fj = d_el_to_face[fj_i];
// Allow fi == fj (self-connections)
// Have we seen this face before? It is possible in some
// configurations to encounter the same neighboring face twice
int idx_j = idx;
for (int i = 0; i < idx; ++i)
{
if (d_face_to_face[begin_f + i] == fj)
{
idx_j = i;
break;
}
}
// This is a new face, record it and increment the counter
if (idx_j == idx)
{
d_face_to_face[begin_f + idx] = fj;
idx++;
}
}
}
});
mfem::forall(nf, [=] MFEM_HOST_DEVICE (int fi)
{
const int begin = d_face_face_offsets[fi];
const int end = d_face_face_offsets[fi + 1];
for (int idx_j = begin; idx_j < end; ++idx_j)
{
const int fj = d_face_to_face[idx_j];
for (int ei = 0; ei < 2; ++ei)
{
const int e = d_face_to_el(0, ei, fi);
if (e < 0 || e >= ne) { continue; }
const int fi_i = d_face_to_el(1, ei, fi);
int fj_i = -1;
for (int ej = 0; ej < 2; ++ej)
{
if (d_face_to_el(0, ej, fj) == e)
{
fj_i = d_face_to_el(1, ej, fj);
break;
}
}
if (fj_i >= 0)
{
const real_t *Ct_i = &d_Ct(0, 0, fi_i);
const real_t *AhatInvCt_i = &d_AhatInvCt(0, 0, fj_i);
real_t *CAhatInvCt_i = &d_CAhatInvCt(0, 0, idx_j);
kernels::AddMultAtB(m, n, n, Ct_i, AhatInvCt_i, CAhatInvCt_i);
}
}
}
});
#ifdef MFEM_USE_MPI
auto *c_pfes = dynamic_cast<ParFiniteElementSpace*>(&h.c_fes);
#endif
const int ncdofs_face_nbr = [&]()
{
#ifdef MFEM_USE_MPI
// Only need to handle face neighbor DOFs when there are nonconforming
// (ghost) faces.
if (c_pfes && c_pfes->Nonconforming())
{
c_pfes->ExchangeFaceNbrData();
return c_pfes->GetFaceNbrVSize();
}
#endif
return 0;
}();
const int ncdofs_local = h.c_fes.GetVSize();
const int ncdofs = ncdofs_local + ncdofs_face_nbr;
const ElementDofOrdering ordering = ElementDofOrdering::NATIVE;
const FaceRestriction *face_restr =
h.c_fes.GetFaceRestriction(ordering, FaceType::Interior);
const auto *l2_face_restr =
dynamic_cast<const L2InterfaceFaceRestriction*>(face_restr);
MFEM_ASSERT(l2_face_restr, "");
const auto c_scatter_map = Reshape(l2_face_restr->ScatterMap().Read(), n, nf);
h.H.reset(new SparseMatrix);
h.H->OverrideSize(ncdofs, ncdofs);
h.H->GetMemoryI().New(ncdofs + 1, h.H->GetMemoryI().GetMemoryType());
{
int *I = h.H->WriteI();
mfem::forall(ncdofs, [=] MFEM_HOST_DEVICE (int i) { I[i] = 0; });
mfem::forall(nf*n, [=] MFEM_HOST_DEVICE (int idx_i)
{
const int i = idx_i % n;
const int fi = idx_i / n;
const int ii = c_scatter_map(i, fi);
const int begin = d_face_face_offsets[fi];
const int end = d_face_face_offsets[fi + 1];
for (int idx = begin; idx < end; ++idx)
{
for (int j = 0; j < n; ++j)
{
if (d_CAhatInvCt(i, j, idx) != 0)
{
I[ii]++;
}
}
}
});
}
// At this point, I[i] contains the number of nonzeros in row I. Perform a
// partial sum to get I in CSR format. This is serial, so perform on host.
//
// At the same time, we find any empty rows (corresponding to non-ghost DOFs)
// and add a single nonzero (we will put 1 on the diagonal) and record the
// row index.
Array<int> empty_rows;
{
int *I = h.H->HostReadWriteI();
int empty_row_count = 0;
for (int i = 0; i < ncdofs_local; i++)
{
if (I[i] == 0) { empty_row_count++; }
}
empty_rows.SetSize(empty_row_count);
int empty_row_idx = 0;
int sum = 0;
for (int i = 0; i < ncdofs; i++)
{
int nnz = I[i];
if (nnz == 0 && i < ncdofs_local)
{
empty_rows[empty_row_idx] = i;
empty_row_idx++;
nnz = 1;
}
I[i] = sum;
sum += nnz;
}
I[ncdofs] = sum;
}
const int nnz = h.H->HostReadI()[ncdofs];
h.H->GetMemoryJ().New(nnz, h.H->GetMemoryJ().GetMemoryType());
h.H->GetMemoryData().New(nnz, h.H->GetMemoryData().GetMemoryType());
{
int *I = h.H->ReadWriteI();
int *J = h.H->WriteJ();
real_t *V = h.H->WriteData();
mfem::forall(nf*n, [=] MFEM_HOST_DEVICE (int idx_i)
{
const int i = idx_i % n;
const int fi = idx_i / n;
const int ii = c_scatter_map[i + fi*n];
const int begin = d_face_face_offsets[fi];
const int end = d_face_face_offsets[fi + 1];
for (int idx = begin; idx < end; ++idx)
{
const int fj = d_face_to_face[idx];
for (int j = 0; j < n; ++j)
{
const real_t val = d_CAhatInvCt(i, j, idx);
if (val != 0)
{
const int k = I[ii];
const int jj = c_scatter_map(j, fj);
I[ii]++;
J[k] = jj;
V[k] = val;
}
}
}
});
const int *d_empty_rows = empty_rows.Read();
mfem::forall(empty_rows.Size(), [=] MFEM_HOST_DEVICE (int idx)
{
const int i = d_empty_rows[idx];
const int k = I[i];
I[i]++;
J[k] = i;
V[k] = 1.0;
});
}
// Shift back down (serial, done on host)
{
int *I = h.H->HostReadWriteI();
for (int i = ncdofs - 1; i > 0; --i)
{
I[i] = I[i-1];
}
I[0] = 0;
}
#ifdef MFEM_USE_MPI
if (c_pfes)
{
OperatorHandle dH(h.pH.Type());
if (ncdofs_face_nbr > 0)
{
// Build the "face neighbor prolongation matrix" P_nbr, which maps from
// VDOFs (i.e. L-vector) to L-vectors with face neighbor DOFs. The
// action of P_nbr is equivalent to calling ExchangeFaceNbrData on a
// ParGridFunction. We compute P^t A P with P = P_nbr to assemble the
// face neighbor contributions into a parallel matrix.
ParMesh &pmesh = *c_pfes->GetParMesh();
HYPRE_BigInt ncdofs_bigint = ncdofs;
const HYPRE_BigInt global_ncdofs = pmesh.ReduceInt(ncdofs);
Array<HYPRE_BigInt> rows;
Array<HYPRE_BigInt> *offsets[1] = { &rows };
pmesh.GenerateOffsets(1, &ncdofs_bigint, offsets);
Array<int> I(ncdofs + 1);
auto d_I = I.Write();
mfem::forall(ncdofs + 1, [=] MFEM_HOST_DEVICE (int i) { d_I[i] = i; });
HYPRE_BigInt offset = c_pfes->GetMyDofOffset();
Array<HYPRE_BigInt> J(ncdofs);
auto d_J = J.Write();
mfem::forall(ncdofs_local, [=] MFEM_HOST_DEVICE (int i)
{
d_J[i] = offset + i;
});
const HYPRE_BigInt *map = c_pfes->GetFaceNbrGlobalDofMapArray().Read();
mfem::forall(ncdofs_face_nbr, [=] MFEM_HOST_DEVICE (int i)
{
d_J[ncdofs_local + i] = map[i];
});
Vector V(ncdofs);
V.UseDevice();
V = 1.0;
auto P_face_nbr =
std::make_unique<HypreParMatrix>(
c_pfes->GetComm(), ncdofs, global_ncdofs, c_pfes->GlobalVSize(),
I.HostReadWrite(), J.HostReadWrite(), V.HostReadWrite(), rows,
c_pfes->GetDofOffsets());
HypreParMatrix H_diag(c_pfes->GetComm(), global_ncdofs, rows, h.H.get());
dH.Reset(RAP(&H_diag, P_face_nbr.get()));
P_nbr = std::move(P_face_nbr);
}
else
{
dH.MakeSquareBlockDiag(c_pfes->GetComm(),c_pfes->GlobalVSize(),
c_pfes->GetDofOffsets(), h.H.get());
}
OperatorHandle pP(h.pH.Type());
auto P_hyp = static_cast<HypreParMatrix*>(
P_pc ? P_pc.get() : c_pfes->Dof_TrueDof_Matrix());
pP.ConvertFrom(P_hyp);
h.pH.MakePtAP(dH, pP);
h.H.reset();
}
#endif
}
void HybridizationExtension::MultCt(const Vector &x, Vector &y) const
{
Mesh &mesh = *h.fes.GetMesh();
const int ne = mesh.GetNE();
const int nf = mesh.GetNFbyType(FaceType::Interior);
const int n_hat_dof_per_el = h.fes.GetFE(0)->GetDof();
const int n_c_dof_per_face = h.c_fes.GetFaceElement(0)->GetDof();
const ElementDofOrdering ordering = ElementDofOrdering::NATIVE;
const FaceRestriction *face_restr =
h.c_fes.GetFaceRestriction(ordering, FaceType::Interior);
Vector x_evec(face_restr->Height());
face_restr->Mult(x, x_evec);
const int *d_el_to_face = el_to_face.Read();
const int *d_el_face_offsets = el_face_offsets.Read();
const auto d_Ct = Reshape(Ct_mat.Read(), n_hat_dof_per_el, n_c_dof_per_face,
n_el_face);
const auto d_x_evec = Reshape(x_evec.Read(), n_c_dof_per_face, nf);
auto d_y = Reshape(y.Write(), n_hat_dof_per_el, ne);
mfem::forall(ne * n_hat_dof_per_el, [=] MFEM_HOST_DEVICE (int idx)
{
const int e = idx / n_hat_dof_per_el;
const int i = idx % n_hat_dof_per_el;
d_y(i, e) = 0.0;
const int begin = d_el_face_offsets[e];
const int end = d_el_face_offsets[e+1];
for (int fi = begin; fi < end; ++fi)
{
const int f = d_el_to_face[fi];
for (int j = 0; j < n_c_dof_per_face; ++j)
{
d_y(i, e) += d_Ct(i, j, fi)*d_x_evec(j, f);
}
}
});
}
void HybridizationExtension::MultC(const Vector &x, Vector &y) const
{
Mesh &mesh = *h.fes.GetMesh();
const int ne = mesh.GetNE();
const int nf = mesh.GetNFbyType(FaceType::Interior);
const int n_hat_dof_per_el = h.fes.GetTypicalFE()->GetDof();
const int n_c_dof_per_face = h.c_fes.GetTypicalTraceElement()->GetDof();
const ElementDofOrdering ordering = ElementDofOrdering::NATIVE;
const FaceRestriction *face_restr = h.c_fes.GetFaceRestriction(
ordering, FaceType::Interior);
Vector y_evec(face_restr->Height());
const auto d_face_to_el = Reshape(face_to_el.Read(), 2, 2, nf);
const auto d_Ct = Reshape(Ct_mat.Read(), n_hat_dof_per_el, n_c_dof_per_face,
n_el_face);
auto d_x = Reshape(x.Read(), n_hat_dof_per_el, ne);
auto d_y_evec = Reshape(y_evec.Write(), n_c_dof_per_face, nf);
mfem::forall(nf * n_c_dof_per_face, [=] MFEM_HOST_DEVICE (int idx)
{
const int f = idx / n_c_dof_per_face;
const int j = idx % n_c_dof_per_face;
d_y_evec(j, f) = 0.0;
for (int el_i = 0; el_i < 2; ++el_i)
{
const int e = d_face_to_el(0, el_i, f);
const int fi = d_face_to_el(1, el_i, f);
// Skip face neighbor elements of shared faces
if (e >= ne) { continue; }
for (int i = 0; i < n_hat_dof_per_el; ++i)
{
d_y_evec(j, f) += d_Ct(i, j, fi)*d_x(i, e);
}
}
});
y.SetSize(face_restr->Width());
if (P_nbr)
{
auto l2_face_restr =
dynamic_cast<const L2InterfaceFaceRestriction*>(face_restr);
MFEM_ASSERT(l2_face_restr != nullptr, "");
Vector y_s(P_nbr->Height());
l2_face_restr->MultTransposeShared(y_evec, y_s);
P_nbr->MultTranspose(y_s, y);
}
else
{
face_restr->MultTranspose(y_evec, y);
}
}
void HybridizationExtension::AssembleMatrix(int el, const DenseMatrix &elmat)
{
const int n = elmat.Width();
const real_t *d_elmat = elmat.Read();
real_t *d_Ahat = Ahat.ReadWrite();
const int offset = el*n*n;
mfem::forall(n*n, [=] MFEM_HOST_DEVICE (int i)
{
d_Ahat[offset + i] += d_elmat[i];
});
}
void HybridizationExtension::AssembleBdrMatrix(int bdr_el,
const DenseMatrix &elmat)
{
DenseMatrix B = elmat; // deep copy
const int n = h.fes.GetFE(0)->GetDof();
// Create mapping e2f from element DOF indices to face DOF indices
Array<int> e2f(n);
e2f = -1;
int el;
{
Mesh &mesh = *h.fes.GetMesh();
int info;
mesh.GetBdrElementAdjacentElement(bdr_el, el, info);
Array<int> lvdofs;
lvdofs.Reserve(elmat.Height());
h.fes.FEColl()->SubDofOrder(mesh.GetElementGeometry(el),
mesh.Dimension() - 1, info, lvdofs);
// Convert local element dofs to local element vdofs.
const int vdim = h.fes.GetVDim();
Ordering::DofsToVDofs<Ordering::byNODES>(n/vdim, vdim, lvdofs);
MFEM_ASSERT(lvdofs.Size() == elmat.Height(), "internal error");
B.AdjustDofDirection(lvdofs);
FiniteElementSpace::AdjustVDofs(lvdofs);
// Create a map from local element vdofs to local boundary (face) vdofs.
for (int i = 0; i < lvdofs.Size(); i++)
{
e2f[lvdofs[i]] = i;
}
}
const int offset = el*n*n;
Ahat.HostReadWrite();
for (int j = 0; j < n; ++j)
{
const int j_f = e2f[j];
if (j_f < 0) { continue; }
for (int i = 0; i < n; ++i)
{
const int i_f = e2f[i];
if (i_f < 0) { continue; }
Ahat[offset + i + j*n] += B(i_f, j_f);
}
}
}
void HybridizationExtension::AssembleElementMatrices(const DenseTensor &elmats)
{
const real_t *d_elmats = elmats.Read();
real_t *d_Ahat = Ahat.ReadWrite();
mfem::forall(elmats.TotalSize(), [=] MFEM_HOST_DEVICE (int i)
{
d_Ahat[i] += d_elmats[i];
});
}
void HybridizationExtension::Init(const Array<int> &ess_tdof_list)
{
// Verify that preconditions for the extension are met
const Mesh &mesh = *h.fes.GetMesh();
const int dim = mesh.Dimension();
const int ne = h.fes.GetNE();
const int nf = mesh.GetNFbyType(FaceType::Interior);
const int ndof_per_el = h.fes.GetFE(0)->GetDof();
const int ndof_per_face = h.c_fes.GetFaceElement(0)->GetDof();
MFEM_VERIFY(!h.fes.IsVariableOrder(), "");
MFEM_VERIFY(dim == 2 || dim == 3, "");
MFEM_VERIFY(UsesTensorBasis(h.fes), "");
// Set up array for idofs and bdofs
{
const TensorBasisElement* tbe =
dynamic_cast<const TensorBasisElement*>(h.fes.GetFE(0));
MFEM_VERIFY(tbe != nullptr, "");
const Array<int> &dof_map = tbe->GetDofMap();
const int n_faces_per_el = GetNFacesPerElement(mesh);
Array<int> all_face_dofs;
for (int f = 0; f < n_faces_per_el; ++f)
{
Array<int> face_map(ndof_per_face);
h.fes.GetFE(0)->GetFaceMap(f, face_map);
all_face_dofs.Append(face_map);
}
Array<bool> b_marker(ndof_per_el);
b_marker = false;
for (int i = 0; i < all_face_dofs.Size(); ++i)
{
const int j_s = all_face_dofs[i];
const int j = (j_s >= 0) ? j_s : -1 - j_s;
const int j_nat_s = dof_map[j];
const int j_nat = (j_nat_s >= 0) ? j_nat_s : -1 - j_nat_s;
b_marker[j_nat] = true;
}
for (int i = 0; i < ndof_per_el; ++i)
{
if (b_marker[i]) { bdofs.Append(i); }
else { idofs.Append(i); }
}
}
// Set up face info arrays
el_face_offsets.SetSize(ne + 1);
el_face_offsets = 0;
// Count faces per element
for (int f = 0; f < mesh.GetNumFacesWithGhost(); ++f)
{
const Mesh::FaceInformation info = mesh.GetFaceInformation(f);
if (!info.IsInterior() || info.IsNonconformingCoarse()) { continue; }
el_face_offsets[info.element[0].index + 1] += 1;
if (!info.IsShared())
{
el_face_offsets[info.element[1].index + 1] += 1;
}
}
el_face_offsets.PartialSum();
// Set up element-to-face and face-to-element arrays
n_el_face = el_face_offsets.Last();
el_to_face.SetSize(n_el_face);
face_to_el.SetSize(4 * nf);
{
Array<int> el_face_counter(ne);
el_face_counter = 0;
int face_idx = 0;
for (int f = 0; f < mesh.GetNumFacesWithGhost(); ++f)
{
const Mesh::FaceInformation info = mesh.GetFaceInformation(f);
if (!info.IsInterior() || info.IsNonconformingCoarse()) { continue; }
const int el1 = info.element[0].index;
int &offset1 = el_face_offsets[el1];
el_to_face[offset1] = face_idx;
face_to_el[0 + 4*face_idx] = el1;
face_to_el[1 + 4*face_idx] = offset1;
offset1 += 1;
const int el2 = info.element[1].index;
if (!info.IsShared())
{
int &offset2 = el_face_offsets[el2];
el_to_face[offset2] = face_idx;
face_to_el[2 + 4*face_idx] = el2;
face_to_el[3 + 4*face_idx] = offset2;
offset2 += 1;
}
else
{
face_to_el[2 + 4*face_idx] = ne + el2;
face_to_el[3 + 4*face_idx] = -1;
}
++face_idx;
}
for (int i = ne; i > 0; i--)
{
el_face_offsets[i] = el_face_offsets[i-1];
}
el_face_offsets[0] = 0;
}
// Create the face-to-face connectivity
{
face_face_offsets.SetSize(nf + 1);
const auto d_face_to_el = Reshape(face_to_el.Read(), 2, 2, nf);
const auto d_el_face_offsets = el_face_offsets.Read();
auto d_face_face_offsets = face_face_offsets.Write();
mfem::forall(nf + 1, [=] MFEM_HOST_DEVICE (int i) { d_face_face_offsets[i] = 0; });
mfem::forall(nf, [=] MFEM_HOST_DEVICE (int f)
{
int n_connections = 0;
for (int ie = 0; ie < 2; ++ie)
{
// Number of faces adjacent to e
const int e = d_face_to_el(0, ie, f);
if (e < ne)
{
n_connections += d_el_face_offsets[e + 1] - d_el_face_offsets[e];
// Subtract 1 since we are double-counting the face 'f' (it
// belongs to both adjacent elements).
if (ie > 0) { n_connections -= 1; }
}
}
d_face_face_offsets[f + 1] = n_connections;
});
// TODO: parallel scan on device?
face_face_offsets.HostReadWrite();
face_face_offsets.PartialSum();
n_face_face = face_face_offsets.Last();
}
// Count the number of dofs in the discontinuous version of fes:
num_hat_dofs = ne*ndof_per_el;
{
h.hat_offsets.SetSize(ne + 1);
int *d_hat_offsets = h.hat_offsets.Write();
mfem::forall(ne + 1, [=] MFEM_HOST_DEVICE (int i)
{
d_hat_offsets[i] = i*ndof_per_el;
});
}
Ahat.SetSize(ne*ndof_per_el*ndof_per_el);
Ahat.UseDevice(true);
Ahat = 0.0;
ConstructC();
const ElementDofOrdering ordering = ElementDofOrdering::NATIVE;
const Operator *R_op = h.fes.GetElementRestriction(ordering);
const auto *R = dynamic_cast<const ElementRestriction*>(R_op);
MFEM_VERIFY(R, "");
// Find out which "hat DOFs" are essential (depend only on essential Lagrange
// multiplier DOFs).
{
const int ntdofs = h.fes.GetTrueVSize();
// free_tdof_marker is 1 if the DOF is free, 0 if the DOF is essential
Array<int> free_tdof_marker(ntdofs);
{
int *d_free_tdof_marker = free_tdof_marker.Write();
mfem::forall(ntdofs, [=] MFEM_HOST_DEVICE (int i)
{
d_free_tdof_marker[i] = 1;
});
const int n_ess_dofs = ess_tdof_list.Size();
const int *d_ess_tdof_list = ess_tdof_list.Read();
mfem::forall(n_ess_dofs, [=] MFEM_HOST_DEVICE (int i)
{
d_free_tdof_marker[d_ess_tdof_list[i]] = 0;
});
}
Array<int> free_vdofs_marker;
#ifdef MFEM_USE_MPI
auto *pfes = dynamic_cast<ParFiniteElementSpace*>(&h.fes);
if (pfes)
{
HypreParMatrix *P = pfes->Dof_TrueDof_Matrix();
free_vdofs_marker.SetSize(h.fes.GetVSize());
// TODO: would be nice to do this on device
P->BooleanMult(1, free_tdof_marker.HostRead(),
0, free_vdofs_marker.HostWrite());
}
else
#endif
{
const SparseMatrix *cP = h.fes.GetConformingProlongation();
if (cP)
{
free_vdofs_marker.SetSize(cP->Height());
cP->BooleanMult(free_tdof_marker, free_vdofs_marker);
}
else
{
free_vdofs_marker.MakeRef(free_tdof_marker);
}
}
hat_dof_marker.SetSize(num_hat_dofs);
{
// The gather map from the ElementRestriction operator gives us the
// index of the L-dof corresponding to a given (element, local DOF)
// index pair.
const int *gather_map = R->GatherMap().Read();
const int *d_free_vdofs_marker = free_vdofs_marker.Read();
const auto d_Ct_mat = Reshape(Ct_mat.Read(), ndof_per_el,
ndof_per_face, n_el_face);
const int *d_el_face_offsets = el_face_offsets.Read();
DofType *d_hat_dof_marker = hat_dof_marker.Write();
// Set the hat_dofs_marker to 1 or 0 according to whether the DOF is
// "free" or "essential". (For now, we mark all free DOFs as free
// interior as a placeholder). Then, as a later step, the "free" DOFs
// will be further classified as "interior free" or "boundary free".
mfem::forall(num_hat_dofs, [=] MFEM_HOST_DEVICE (int i)
{
const int j_s = gather_map[i];
const int j = (j_s >= 0) ? j_s : -1 - j_s;
if (d_free_vdofs_marker[j])
{
const int i_loc = i % ndof_per_el;
const int e = i / ndof_per_el;
d_hat_dof_marker[i] = INTERIOR;
const int begin = d_el_face_offsets[e];
const int end = d_el_face_offsets[e + 1];
for (int f = begin; f < end; ++f)
{
for (int k = 0; k < ndof_per_face; ++k)
{
if (d_Ct_mat(i_loc, k, f) != 0.0)
{
d_hat_dof_marker[i] = BOUNDARY;
break;
}
}
}
}
else
{
d_hat_dof_marker[i] = ESSENTIAL;
}
});
}
}
// Create the hat DOF gather map. This is used to apply the action of R and
// R^T
{
const int vsize = h.fes.GetVSize();
hat_dof_gather_map.SetSize(num_hat_dofs);
const int *d_offsets = R->Offsets().Read();
const int *d_indices = R->Indices().Read();
int *d_hat_dof_gather_map = hat_dof_gather_map.Write();
mfem::forall(num_hat_dofs, [=] MFEM_HOST_DEVICE (int i)
{
d_hat_dof_gather_map[i] = -1;
});
mfem::forall(vsize, [=] MFEM_HOST_DEVICE (int i)
{
const int offset = d_offsets[i];
const int j_s = d_indices[offset];
const int hat_dof_index = (j_s >= 0) ? j_s : -1 - j_s;
// Note: -1 is used as a special value (invalid), so the negative
// DOF indices start at -2.
d_hat_dof_gather_map[hat_dof_index] = (j_s >= 0) ? i : (-2 - i);
});
}
}
void HybridizationExtension::MultR(const Vector &x_hat, Vector &x) const
{
const Operator *R = h.fes.GetRestrictionOperator();
// If R is null, then L-vector and T-vector are the same, and we don't need
// an intermediate temporary variable.
//
// If R is not null, we first convert to intermediate L-vector (with the
// correct BCs), and then from L-vector to T-vector.
if (!R)
{
MFEM_ASSERT(x.Size() == h.fes.GetVSize(), "");
tmp2.MakeRef(x, 0);
}
else
{
tmp2.SetSize(R->Width());
R->MultTranspose(x, tmp2);
}
const ElementDofOrdering ordering = ElementDofOrdering::NATIVE;
const auto *restr = static_cast<const ElementRestriction*>(
h.fes.GetElementRestriction(ordering));
const int *gather_map = restr->GatherMap().Read();
const DofType *d_hat_dof_marker = hat_dof_marker.Read();
const real_t *d_evec = x_hat.Read();
real_t *d_lvec = tmp2.ReadWrite();
mfem::forall(num_hat_dofs, [=] MFEM_HOST_DEVICE (int i)
{
// Skip essential DOFs
if (d_hat_dof_marker[i] == ESSENTIAL) { return; }
const int j_s = gather_map[i];
const int sgn = (j_s >= 0) ? 1 : -1;
const int j = (j_s >= 0) ? j_s : -1 - j_s;
d_lvec[j] = sgn*d_evec[i];
});
// Convert from L-vector to T-vector.
if (R) { R->Mult(tmp2, x); }
}
void HybridizationExtension::MultRt(const Vector &b, Vector &b_hat) const
{
Vector b_lvec;
const Operator *R = h.fes.GetRestrictionOperator();
if (!R)
{
b_lvec.MakeRef(const_cast<Vector&>(b), 0, b.Size());
}
else
{
tmp1.SetSize(h.fes.GetVSize());
b_lvec.MakeRef(tmp1, 0, tmp1.Size());
R->MultTranspose(b, b_lvec);
}
b_hat.SetSize(num_hat_dofs);
const int *d_hat_dof_gather_map = hat_dof_gather_map.Read();
const real_t *d_b_lvec = b_lvec.Read();
real_t *d_b_hat = b_hat.Write();
mfem::forall(num_hat_dofs, [=] MFEM_HOST_DEVICE (int i)
{
const int j_s = d_hat_dof_gather_map[i];
if (j_s == -1) // invalid
{
d_b_hat[i] = 0.0;
}
else
{
const int sgn = (j_s >= 0) ? 1 : -1;
const int j = (j_s >= 0) ? j_s : -2 - j_s;
d_b_hat[i] = sgn*d_b_lvec[j];
}
});
}
void HybridizationExtension::MultAhatInv(Vector &x) const
{
const int ne = h.fes.GetMesh()->GetNE();
const int n = h.fes.GetFE(0)->GetDof();
const int nidofs = idofs.Size();
const int nbdofs = bdofs.Size();
const auto d_hat_dof_marker = Reshape(hat_dof_marker.Read(), n, ne);
const auto d_A_ii = Reshape(Ahat_ii.Read(), nidofs, nidofs, ne);
const auto d_A_ib = Reshape(Ahat_ib.Read(), nidofs*nbdofs, ne);
const auto d_A_bi = Reshape(Ahat_bi.Read(), nbdofs*nidofs, ne);
const auto d_A_bb = Reshape(Ahat_bb.Read(), nbdofs*nbdofs, ne);
const auto d_ipiv_ii = Reshape(Ahat_ii_piv.Read(), nidofs, ne);
const auto d_ipiv_bb = Reshape(Ahat_bb_piv.Read(), nbdofs, ne);
const auto *d_idofs = idofs.Read();
const auto *d_bdofs = bdofs.Read();
Vector ivals(nidofs*ne);
Vector bvals(nbdofs*ne);
auto d_ivals = Reshape(ivals.Write(), nidofs, ne);
auto d_bvals = Reshape(bvals.Write(), nbdofs, ne);
auto d_x = Reshape(x.ReadWrite(), n, ne);
mfem::forall(ne, [=] MFEM_HOST_DEVICE (int e)
{
constexpr int MD1D = DofQuadLimits::HDIV_MAX_D1D;
constexpr int MAX_DOFS = 3*MD1D*(MD1D-1)*(MD1D-1);
internal::LocalMemory<int,MAX_DOFS> bdofs_loc;
int nbfdofs = 0;
for (int i = 0; i < nbdofs; i++)
{
const int dof_idx = d_bdofs[i];
if (d_hat_dof_marker(dof_idx, e) != ESSENTIAL)
{
bdofs_loc[nbfdofs] = dof_idx;
nbfdofs += 1;
}
}
for (int i = 0; i < nidofs; ++i)
{
d_ivals(i, e) = d_x(d_idofs[i], e);
}
for (int i = 0; i < nbfdofs; ++i)
{
d_bvals(i, e) = d_x(bdofs_loc[i], e);
}
if (nidofs > 0)
{
// Block forward substitution:
// B1 <- L^{-1} P B1
kernels::LSolve(&d_A_ii(0,0,e), nidofs, &d_ipiv_ii(0,e), &d_ivals(0,e));
// B2 <- B2 - L21 B1
kernels::SubMult(
nidofs, nbfdofs, 1, &d_A_bi(0,e), &d_ivals(0,e), &d_bvals(0, e));
}
// Schur complement solve
kernels::LUSolve(&d_A_bb(0,e), nbfdofs, &d_ipiv_bb(0,e), &d_bvals(0,e));
if (nidofs > 0)
{
// Block backward substitution
// Y1 <- Y1 - U12 X2
kernels::SubMult(
nbfdofs, nidofs, 1, &d_A_ib(0,e), &d_bvals(0,e), &d_ivals(0, e));
// Y1 <- U^{-1} Y1
kernels::USolve(&d_A_ii(0,0,e), nidofs, &d_ivals(0,e));
}
for (int i = 0; i < nidofs; ++i)
{
d_x(d_idofs[i], e) = d_ivals(i, e);
}
for (int i = 0; i < nbfdofs; ++i)
{
d_x(bdofs_loc[i], e) = d_bvals(i, e);
}
});
}
void HybridizationExtension::ReduceRHS(const Vector &b, Vector &b_r) const
{
Vector b_hat(num_hat_dofs);
MultRt(b, b_hat);
{
const auto *d_hat_dof_marker = hat_dof_marker.Read();
auto *d_b_hat = b_hat.ReadWrite();
mfem::forall(num_hat_dofs, [=] MFEM_HOST_DEVICE (int i)
{
if (d_hat_dof_marker[i] == ESSENTIAL) { d_b_hat[i] = 0.0; }
});
}
MultAhatInv(b_hat);
if (IsParFESpace(h.c_fes))
{
const Operator &P = GetProlongation();
Vector bl(P.Height());
b_r.SetSize(P.Width());
MultC(b_hat, bl);
P.MultTranspose(bl, b_r);
}
else
{
MultC(b_hat, b_r);
}
}
void HybridizationExtension::ComputeSolution(
const Vector &b, const Vector &sol_r, Vector &sol) const
{
// tmp1 = A_hat^{-1} ( R^T b - C^T lambda )
Vector b_hat(num_hat_dofs);
MultRt(b, b_hat);
tmp1.SetSize(num_hat_dofs);
if (IsParFESpace(h.c_fes))
{
const Operator &P = GetProlongation();
Vector sol_l(P.Height());
P.Mult(sol_r, sol_l);
MultCt(sol_l, tmp1);
}
else
{
MultCt(sol_r, tmp1);
}
add(b_hat, -1.0, tmp1, tmp1);
// Eliminate essential DOFs
const auto *d_hat_dof_marker = hat_dof_marker.Read();
real_t *d_tmp1 = tmp1.ReadWrite();
mfem::forall(num_hat_dofs, [=] MFEM_HOST_DEVICE (int i)
{
if (d_hat_dof_marker[i] == ESSENTIAL) { d_tmp1[i] = 0.0; }
});
MultAhatInv(tmp1);
MultR(tmp1, sol);
}
}
Reference in New Issue View Git Blame Copy Permalink