Skip to content

Commit

Permalink
assembly directly using global dof indices rather than assembling the…
Browse files Browse the repository at this point in the history 
…n rearranging
  • Loading branch information
@mscroggs
mscroggs committed Apr 26, 2024
1 parent 85917f3 commit d618a79
Show file tree
Hide file tree
Showing 4 changed files with 98 additions and 106 deletions.
74 changes: 12 additions & 62 deletions src/assembly/batched.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ use crate::traits::function::FunctionSpace;
#[cfg(feature = "mpi")]
use crate::traits::function::FunctionSpaceInParallel;
use crate::traits::grid::{CellType, GridType, ReferenceMapType, TopologyType};
#[cfg(feature = "mpi")]
use crate::traits::types::Ownership;
use crate::traits::types::ReferenceCellType;
use green_kernels::types::EvalType;
Expand Down Expand Up @@ -211,9 +210,11 @@ fn assemble_batch_singular<
.unwrap();
}
}
output.rows.push(*test_dof);
output.cols.push(*trial_dof);
output.data.push(sum);
if trial_space.ownership(*trial_dof) == Ownership::Owned {
output.rows.push(test_space.global_dof_index(*test_dof));
output.cols.push(trial_space.global_dof_index(*trial_dof));
output.data.push(sum);
}
}
}
}
Expand Down Expand Up @@ -471,9 +472,11 @@ fn assemble_batch_singular_correction<
};
}
}
output.rows.push(*test_dof);
output.cols.push(*trial_dof);
output.data.push(sum);
if trial_space.ownership(*trial_dof) == Ownership::Owned {
output.rows.push(test_space.global_dof_index(*test_dof));
output.cols.push(trial_space.global_dof_index(*trial_dof));
output.data.push(sum);
}
}
}
}
Expand Down Expand Up @@ -1042,34 +1045,7 @@ pub trait BatchedAssembler: Sync + Sized {
trial_space: &(impl FunctionSpaceInParallel<SerialSpace = SerialTrialSpace> + FunctionSpace),
test_space: &(impl FunctionSpaceInParallel<SerialSpace = SerialTestSpace> + FunctionSpace),
) -> CsrMatrix<Self::T> {
let local_matrix =
self.assemble_singular_into_csr(trial_space.local_space(), test_space.local_space());

let ownership = trial_space.ownership();
let test_global_dof_numbers = test_space.global_dof_numbers();
let trial_global_dof_numbers = trial_space.global_dof_numbers();

let mut rows = vec![];
let mut cols = vec![];
let mut data = vec![];
let mut r = 0;
for (i, index) in local_matrix.indices().iter().enumerate() {
while i >= local_matrix.indptr()[r + 1] {
r += 1;
}
if ownership[*index] == Ownership::Owned {
rows.push(test_global_dof_numbers[r]);
cols.push(trial_global_dof_numbers[*index]);
data.push(local_matrix.data()[i]);
}
}
CsrMatrix::from_aij(
[test_space.global_size(), trial_space.global_size()],
&rows,
&cols,
&data,
)
.unwrap()
self.assemble_singular_into_csr(trial_space.local_space(), test_space.local_space())
}

/// Assemble the singular correction into a dense matrix
Expand Down Expand Up @@ -1139,36 +1115,10 @@ pub trait BatchedAssembler: Sync + Sized {
trial_space: &(impl FunctionSpaceInParallel<SerialSpace = SerialTrialSpace> + FunctionSpace),
test_space: &(impl FunctionSpaceInParallel<SerialSpace = SerialTestSpace> + FunctionSpace),
) -> CsrMatrix<Self::T> {
let local_matrix = self.assemble_singular_correction_into_csr(
self.assemble_singular_correction_into_csr(
trial_space.local_space(),
test_space.local_space(),
);

let ownership = trial_space.ownership();
let test_global_dof_numbers = test_space.global_dof_numbers();
let trial_global_dof_numbers = trial_space.global_dof_numbers();

let mut rows = vec![];
let mut cols = vec![];
let mut data = vec![];
let mut r = 0;
for (i, index) in local_matrix.indices().iter().enumerate() {
while i >= local_matrix.indptr()[r + 1] {
r += 1;
}
if ownership[*index] == Ownership::Owned {
rows.push(test_global_dof_numbers[r]);
cols.push(trial_global_dof_numbers[*index]);
data.push(local_matrix.data()[i]);
}
}
CsrMatrix::from_aij(
[test_space.global_size(), trial_space.global_size()],
&rows,
&cols,
&data,
)
.unwrap()
}

/// Assemble into a dense matrix
Expand Down
107 changes: 72 additions & 35 deletions src/function/function_space/parallel.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,17 +16,58 @@ use mpi::{
use rlst::RlstScalar;
use std::collections::HashMap;

/// The local function space on a process
pub struct LocalFunctionSpace<'a, T: RlstScalar, GridImpl: GridType<T = T::Real>> {
serial_space: SerialFunctionSpace<'a, T, GridImpl>,
global_size: usize,
global_dof_numbers: Vec<usize>,
ownership: Vec<Ownership>,
}

impl<'a, T: RlstScalar, GridImpl: GridType<T = T::Real>> FunctionSpace
for LocalFunctionSpace<'a, T, GridImpl>
{
type Grid = GridImpl;
type FiniteElement = CiarletElement<T>;

fn grid(&self) -> &Self::Grid {
self.serial_space.grid
}
fn element(&self, cell_type: ReferenceCellType) -> &CiarletElement<T> {
self.serial_space.element(cell_type)
}
fn get_local_dof_numbers(&self, entity_dim: usize, entity_number: usize) -> &[usize] {
self.serial_space
.get_local_dof_numbers(entity_dim, entity_number)
}
fn local_size(&self) -> usize {
self.serial_space.local_size()
}
fn global_size(&self) -> usize {
self.global_size
}
fn cell_dofs(&self, cell: usize) -> Option<&[usize]> {
self.serial_space.cell_dofs(cell)
}
fn cell_colouring(&self) -> HashMap<ReferenceCellType, Vec<Vec<usize>>> {
self.serial_space.cell_colouring()
}
fn global_dof_index(&self, local_dof_index: usize) -> usize {
self.global_dof_numbers[local_dof_index]
}
fn ownership(&self, local_dof_index: usize) -> Ownership {
self.ownership[local_dof_index]
}
}

/// A parallel function space
pub struct ParallelFunctionSpace<
'a,
T: RlstScalar,
GridImpl: ParallelGridType + GridType<T = T::Real>,
> {
grid: &'a GridImpl,
local_space: SerialFunctionSpace<'a, T, <GridImpl as ParallelGridType>::LocalGridType>,
global_dof_numbers: Vec<usize>,
owners: Vec<Ownership>,
global_size: usize,
local_space: LocalFunctionSpace<'a, T, <GridImpl as ParallelGridType>::LocalGridType>,
}

impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>>
Expand All @@ -50,16 +91,8 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>>
elements.insert(*cell, e_family.element(*cell));
}

let local_space = SerialFunctionSpace {
grid: grid.local_grid(),
elements,
entity_dofs,
cell_dofs,
size: dofmap_size,
};

// Assign global DOF numbers
let mut global_dof_numbers = vec![0; local_space.local_size()];
let mut global_dof_numbers = vec![0; dofmap_size];
let mut ghost_indices = vec![vec![]; size as usize];
let mut ghost_cells = vec![vec![]; size as usize];
let mut ghost_cell_dofs = vec![vec![]; size as usize];
Expand Down Expand Up @@ -123,7 +156,7 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>>
let local_ghost_dofs = gcells
.iter()
.zip(&gdofs)
.map(|(c, d)| local_space.cell_dofs(*c).unwrap()[*d])
.map(|(c, d)| cell_dofs[*c][*d])
.collect::<Vec<_>>();
let global_ghost_dofs = local_ghost_dofs
.iter()
Expand All @@ -136,7 +169,7 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>>
}
}
// receive ghost info
let mut owners = vec![Ownership::Owned; local_space.local_size()];
let mut ownership = vec![Ownership::Owned; dofmap_size];
for p in 0..size {
if p != rank {
let process = comm.process_at_rank(p);
Expand All @@ -147,43 +180,41 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>>
.zip(local_ghost_dofs.iter().zip(&global_ghost_dofs))
{
global_dof_numbers[*i] = *g;
owners[*i] = Ownership::Ghost(p as usize, *l);
ownership[*i] = Ownership::Ghost(p as usize, *l);
}
}
}

Self {
grid,
local_space,
global_dof_numbers,
owners,
let serial_space = SerialFunctionSpace {
grid: grid.local_grid(),
elements,
entity_dofs,
cell_dofs,
size: dofmap_size,
};
let local_space = LocalFunctionSpace {
serial_space,
global_size,
}
global_dof_numbers,
ownership,
};

Self { grid, local_space }
}
}

impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>> FunctionSpaceInParallel
for ParallelFunctionSpace<'a, T, GridImpl>
{
type ParallelGrid = GridImpl;
type SerialSpace = SerialFunctionSpace<'a, T, <GridImpl as ParallelGridType>::LocalGridType>;
type SerialSpace = LocalFunctionSpace<'a, T, <GridImpl as ParallelGridType>::LocalGridType>;

/// Get the local space on the process
fn local_space(
&self,
) -> &SerialFunctionSpace<'a, T, <GridImpl as ParallelGridType>::LocalGridType> {
) -> &LocalFunctionSpace<'a, T, <GridImpl as ParallelGridType>::LocalGridType> {
&self.local_space
}

/// Get the global DOF number associated with each local DOF
fn global_dof_numbers(&self) -> &Vec<usize> {
&self.global_dof_numbers
}

/// Get ownership info
fn ownership(&self) -> &Vec<Ownership> {
&self.owners
}
}

impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>> FunctionSpace
Expand All @@ -206,12 +237,18 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType<T = T::Real>> Func
self.local_space.local_size()
}
fn global_size(&self) -> usize {
self.global_size
self.local_space.global_size()
}
fn cell_dofs(&self, cell: usize) -> Option<&[usize]> {
self.local_space.cell_dofs(cell)
}
fn cell_colouring(&self) -> HashMap<ReferenceCellType, Vec<Vec<usize>>> {
self.local_space.cell_colouring()
}
fn global_dof_index(&self, local_dof_index: usize) -> usize {
self.local_space.global_dof_index(local_dof_index)
}
fn ownership(&self, local_dof_index: usize) -> Ownership {
self.local_space.ownership(local_dof_index)
}
}
10 changes: 8 additions & 2 deletions src/function/function_space/serial.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::traits::{
element::{ElementFamily, FiniteElement},
function::FunctionSpace,
grid::{CellType, GridType, TopologyType},
types::ReferenceCellType,
types::{Ownership, ReferenceCellType},
};
use rlst::RlstScalar;
use std::collections::HashMap;
Expand Down Expand Up @@ -62,7 +62,7 @@ impl<'a, T: RlstScalar, GridImpl: GridType<T = T::Real>> FunctionSpace
self.size
}
fn global_size(&self) -> usize {
self.local_size()
self.size
}
fn cell_dofs(&self, cell: usize) -> Option<&[usize]> {
if cell < self.cell_dofs.len() {
Expand Down Expand Up @@ -146,6 +146,12 @@ impl<'a, T: RlstScalar, GridImpl: GridType<T = T::Real>> FunctionSpace
}
colouring
}
fn global_dof_index(&self, local_dof_index: usize) -> usize {
local_dof_index
}
fn ownership(&self, _local_dof_index: usize) -> Ownership {
Ownership::Owned
}
}

#[cfg(test)]
Expand Down
13 changes: 6 additions & 7 deletions src/traits/function.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ use crate::traits::element::FiniteElement;
use crate::traits::grid::GridType;
#[cfg(feature = "mpi")]
use crate::traits::grid::ParallelGridType;
#[cfg(feature = "mpi")]
use crate::traits::types::Ownership;
use crate::traits::types::ReferenceCellType;
use std::collections::HashMap;
Expand Down Expand Up @@ -40,6 +39,12 @@ pub trait FunctionSpace {

/// Compute a colouring of the cells so that no two cells that share an entity with DOFs associated with it are assigned the same colour
fn cell_colouring(&self) -> HashMap<ReferenceCellType, Vec<Vec<usize>>>;

/// Get the global DOF indes associated with a local DOF indec
fn global_dof_index(&self, local_dof_index: usize) -> usize;

/// Get ownership of a local DOF
fn ownership(&self, local_dof_index: usize) -> Ownership;
}

#[cfg(feature = "mpi")]
Expand All @@ -52,10 +57,4 @@ pub trait FunctionSpaceInParallel {

/// Get the local space on the process
fn local_space(&self) -> &Self::SerialSpace;

/// Get the global DOF number associated with each local DOF
fn global_dof_numbers(&self) -> &Vec<usize>;

/// Get ownership info
fn ownership(&self) -> &Vec<Ownership>;
}

0 comments on commit d618a79

Please sign in to comment.