Skip to content

Commit 2c66987

Browse files
committed
MatIS: Fix parallel LGMap
1 parent da194e3 commit 2c66987

File tree

3 files changed

+56
-9
lines changed

3 files changed

+56
-9
lines changed

firedrake/assemble.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1366,6 +1366,10 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions):
13661366
except SparsityFormatError:
13671367
raise ValueError("Monolithic matrix assembly not supported for systems "
13681368
"with R-space blocks")
1369+
if mat_type == "is":
1370+
rlgmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space())
1371+
clgmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space())
1372+
sparsity._lgmaps = (rlgmap, clgmap)
13691373
return sparsity
13701374

13711375
def _make_maps_and_regions(self):
@@ -2177,3 +2181,42 @@ def index_function_spaces(form, indices):
21772181
return tuple(a.ufl_function_space()[i] for i, a in zip(indices, form.arguments()))
21782182
else:
21792183
raise AssertionError
2184+
2185+
2186+
def masked_lgmap(lgmap, mask, block=True):
2187+
if block:
2188+
indices = lgmap.block_indices.copy()
2189+
bsize = lgmap.getBlockSize()
2190+
else:
2191+
indices = lgmap.indices.copy()
2192+
bsize = 1
2193+
indices[mask] = -1
2194+
return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm)
2195+
2196+
2197+
def unghosted_lgmap(lgmap, V, block=True):
2198+
block_size = lgmap.getBlockSize() if block else 1
2199+
2200+
mesh_dm = V.mesh().topology_dm
2201+
depth = mesh_dm.getDepth()
2202+
start, end = mesh_dm.getDepthStratum(depth)
2203+
2204+
own = []
2205+
for W in V:
2206+
section = W.dm.getDefaultSection()
2207+
for seed in range(start, end):
2208+
# Do not loop over ghost cells
2209+
if mesh_dm.getLabelValue("pyop2_ghost", seed) != -1:
2210+
continue
2211+
closure, _ = mesh_dm.getTransitiveClosure(seed, useCone=True)
2212+
for p in closure:
2213+
dof = section.getDof(p)
2214+
if dof <= 0:
2215+
continue
2216+
off = section.getOffset(p)
2217+
# Local indices within W
2218+
W_indices = range(block_size * off, block_size * (off + dof))
2219+
own.extend(W_indices)
2220+
2221+
mask = numpy.setdiff1d(range(len(lgmap.indices)), own)
2222+
return masked_lgmap(lgmap, mask, block=block)

pyop2/types/mat.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ def __init__(self, dsets, maps_and_regions, name=None, nest=None, block_sparse=N
5858
if self._initialized:
5959
return
6060
self._dsets = dsets
61+
self._lgmaps = None
6162
self._maps_and_regions = maps_and_regions
6263
self._block_sparse = block_sparse
6364
self._diagonal_block = diagonal_block
@@ -619,15 +620,15 @@ def _init_monolithic(self):
619620
rlgmap = rset.unblocked_lgmap
620621
clgmap = cset.unblocked_lgmap
621622
if self.mat_type == "is":
623+
# TODO monolithic lgmaps
624+
rlgmap, clgmap = self.sparsity._lgmaps
622625
create = mat.createIS
623-
elif self.mat_type.endswith("aij"):
624-
create = mat.createAIJ
625626
else:
626-
raise ValueError(f"Unsupported mat_type {mat_type}")
627+
create = mat.createAIJ
627628
size = ((self.nrows, None), (self.ncols, None))
628629
create(size, bsize=1, comm=self.comm)
629630
mat.setLGMap(rmap=rlgmap, cmap=clgmap)
630-
mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnzi))
631+
mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnz))
631632
self.handle = mat
632633
self._blocks = []
633634
rows, cols = self.sparsity.shape
@@ -691,6 +692,7 @@ def _init_block(self):
691692
rdim, cdim = self.dims[0][0]
692693

693694
if self.mat_type == "is":
695+
row_lg, col_lg = self.sparsity._lgmaps
694696
block_sparse = False
695697
create = mat.createIS
696698
elif rdim == cdim and rdim > 1 and self.sparsity._block_sparse:
@@ -703,10 +705,9 @@ def _init_block(self):
703705
# the /dof/ sparsity.
704706
block_sparse = False
705707
create = mat.createAIJ
706-
create(size=((self.nrows, None),
707-
(self.ncols, None)),
708-
bsize=(rdim, cdim),
709-
comm=self.comm)
708+
size = ((self.nrows, None), (self.ncols, None))
709+
create(size, bsize=(rdim, cdim), comm=self.comm)
710+
710711
mat.setLGMap(rmap=row_lg, cmap=col_lg)
711712
mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnz))
712713
# Stash entries destined for other processors
@@ -826,6 +827,8 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None):
826827
self.change_assembly_state(Mat.INSERT_VALUES)
827828
if len(rows) > 0:
828829
if self.handle.type == "is":
830+
own = self.handle.getLocalSize()[0]
831+
rows = rows[rows < own]
829832
self.handle.assemble()
830833
self.handle.zeroRowsColumnsLocal(rows, diag_val)
831834
else:

tests/firedrake/regression/test_assemble.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh):
142142

143143
assert A2.M is A1.M
144144

145+
145146
@pytest.mark.parametrize("dirichlet_bcs", [False, True])
146147
def test_assemble_matis(mesh, dirichlet_bcs):
147148
V = FunctionSpace(mesh, "CG", 1)
@@ -159,9 +160,9 @@ def test_assemble_matis(mesh, dirichlet_bcs):
159160

160161
aij = assemble(a, bcs=bcs, mat_type="aij").petscmat
161162
aij.axpy(-1, aijnew)
162-
aij.view()
163163
assert np.allclose(aij[:, :], 0)
164164

165+
165166
def test_assemble_diagonal(mesh):
166167
V = FunctionSpace(mesh, "P", 3)
167168
u = TrialFunction(V)

0 commit comments

Comments
 (0)