MatIS: Fix parallel LGMap

Author: pbrubeck

firedrake/assemble.py

@@ -1366,6 +1366,10 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions):
         except SparsityFormatError:
             raise ValueError("Monolithic matrix assembly not supported for systems "
                              "with R-space blocks")
+        if mat_type == "is":
+            rlgmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space())
+            clgmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space())
+            sparsity._lgmaps = (rlgmap, clgmap)
         return sparsity
 
     def _make_maps_and_regions(self):
@@ -2177,3 +2181,42 @@ def index_function_spaces(form, indices):
             return tuple(a.ufl_function_space()[i] for i, a in zip(indices, form.arguments()))
         else:
             raise AssertionError
+
+
+def masked_lgmap(lgmap, mask, block=True):
+    if block:
+        indices = lgmap.block_indices.copy()
+        bsize = lgmap.getBlockSize()
+    else:
+        indices = lgmap.indices.copy()
+        bsize = 1
+    indices[mask] = -1
+    return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm)
+
+
+def unghosted_lgmap(lgmap, V, block=True):
+    block_size = lgmap.getBlockSize() if block else 1
+
+    mesh_dm = V.mesh().topology_dm
+    depth = mesh_dm.getDepth()
+    start, end = mesh_dm.getDepthStratum(depth)
+
+    own = []
+    for W in V:
+        section = W.dm.getDefaultSection()
+        for seed in range(start, end):
+            # Do not loop over ghost cells
+            if mesh_dm.getLabelValue("pyop2_ghost", seed) != -1:
+                continue
+            closure, _ = mesh_dm.getTransitiveClosure(seed, useCone=True)
+            for p in closure:
+                dof = section.getDof(p)
+                if dof <= 0:
+                    continue
+                off = section.getOffset(p)
+                # Local indices within W
+                W_indices = range(block_size * off, block_size * (off + dof))
+                own.extend(W_indices)
+
+    mask = numpy.setdiff1d(range(len(lgmap.indices)), own)
+    return masked_lgmap(lgmap, mask, block=block)

pyop2/types/mat.py

@@ -58,6 +58,7 @@ def __init__(self, dsets, maps_and_regions, name=None, nest=None, block_sparse=N
         if self._initialized:
             return
         self._dsets = dsets
+        self._lgmaps = None
         self._maps_and_regions = maps_and_regions
         self._block_sparse = block_sparse
         self._diagonal_block = diagonal_block
@@ -619,15 +620,15 @@ def _init_monolithic(self):
         rlgmap = rset.unblocked_lgmap
         clgmap = cset.unblocked_lgmap
         if self.mat_type == "is":
+            # TODO monolithic lgmaps
+            rlgmap, clgmap = self.sparsity._lgmaps
             create = mat.createIS
-        elif self.mat_type.endswith("aij"):
-            create = mat.createAIJ
         else:
-            raise ValueError(f"Unsupported mat_type {mat_type}")
+            create = mat.createAIJ
         size = ((self.nrows, None), (self.ncols, None))
         create(size, bsize=1, comm=self.comm)
         mat.setLGMap(rmap=rlgmap, cmap=clgmap)
-        mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnzi))
+        mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnz))
         self.handle = mat
         self._blocks = []
         rows, cols = self.sparsity.shape
@@ -691,6 +692,7 @@ def _init_block(self):
         rdim, cdim = self.dims[0][0]
 
         if self.mat_type == "is":
+            row_lg, col_lg = self.sparsity._lgmaps
             block_sparse = False
             create = mat.createIS
         elif rdim == cdim and rdim > 1 and self.sparsity._block_sparse:
@@ -703,10 +705,9 @@ def _init_block(self):
             # the /dof/ sparsity.
             block_sparse = False
             create = mat.createAIJ
-        create(size=((self.nrows, None),
-                     (self.ncols, None)),
-               bsize=(rdim, cdim),
-               comm=self.comm)
+        size = ((self.nrows, None), (self.ncols, None))
+        create(size, bsize=(rdim, cdim), comm=self.comm)
+
         mat.setLGMap(rmap=row_lg, cmap=col_lg)
         mat.setPreallocationNNZ((self.sparsity.nnz, self.sparsity.onnz))
         # Stash entries destined for other processors
@@ -826,6 +827,8 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None):
         self.change_assembly_state(Mat.INSERT_VALUES)
         if len(rows) > 0:
             if self.handle.type == "is":
+                own = self.handle.getLocalSize()[0]
+                rows = rows[rows < own]
                 self.handle.assemble()
                 self.handle.zeroRowsColumnsLocal(rows, diag_val)
             else:

tests/firedrake/regression/test_assemble.py

@@ -142,6 +142,7 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh):
 
     assert A2.M is A1.M
 
+
 @pytest.mark.parametrize("dirichlet_bcs", [False, True])
 def test_assemble_matis(mesh, dirichlet_bcs):
     V = FunctionSpace(mesh, "CG", 1)
@@ -159,9 +160,9 @@ def test_assemble_matis(mesh, dirichlet_bcs):
 
     aij = assemble(a, bcs=bcs, mat_type="aij").petscmat
     aij.axpy(-1, aijnew)
-    aij.view()
     assert np.allclose(aij[:, :], 0)
 
+
 def test_assemble_diagonal(mesh):
     V = FunctionSpace(mesh, "P", 3)
     u = TrialFunction(V)