compiler: Fix multi-field petscsolve in parallel

Author: ZoeLeibowitz

devito/petsc/iet/callbacks.py

@@ -1255,13 +1255,16 @@ def _submat_callback_body(self):
         row_idx = DummyExpr(objs['rowidx'], IntDiv(i, objs['dof']))
         col_idx = DummyExpr(objs['colidx'], Mod(i, objs['dof']))
 
-        # Query global size from each sub-DM via a temporary Vec.
+        # Query global and local sizes from each sub-DM via a temporary Vec.
         get_row_vec = petsc_call(
             'DMGetGlobalVector', [objs['Subdms'].indexed[objs['rowidx']], Byref(tvec)]
         )
         get_row_size = petsc_call(
             'VecGetSize', [tvec, Byref(objs['subblockrows'])]
         )
+        get_row_local_size = petsc_call(
+            'VecGetLocalSize', [tvec, Byref(objs['sublocalrows'])]
+        )
         restore_row_vec = petsc_call(
             'DMRestoreGlobalVector', [objs['Subdms'].indexed[objs['rowidx']], Byref(tvec)]
         )
@@ -1271,15 +1274,19 @@ def _submat_callback_body(self):
         get_col_size = petsc_call(
             'VecGetSize', [tvec, Byref(objs['subblockcols'])]
         )
+        get_col_local_size = petsc_call(
+            'VecGetLocalSize', [tvec, Byref(objs['sublocalcols'])]
+        )
         restore_col_vec = petsc_call(
             'DMRestoreGlobalVector', [objs['Subdms'].indexed[objs['colidx']], Byref(tvec)]
         )
 
-        mat_create = petsc_call('MatCreate', [sobjs['comm'], Byref(objs['block'])])
+        petsc_obj_comm = Call('PetscObjectComm', arguments=[PetscObjectCast(objs['J'])])
+        mat_create = petsc_call('MatCreate', [petsc_obj_comm, Byref(objs['block'])])
 
         mat_set_sizes = petsc_call(
             'MatSetSizes', [
-                objs['block'], 'PETSC_DECIDE', 'PETSC_DECIDE',
+                objs['block'], objs['sublocalrows'], objs['sublocalcols'],
                 objs['subblockrows'], objs['subblockcols']
             ]
         )
@@ -1329,9 +1336,11 @@ def _submat_callback_body(self):
             col_idx,
             get_row_vec,
             get_row_size,
+            get_row_local_size,
             restore_row_vec,
             get_col_vec,
             get_col_size,
+            get_col_local_size,
             restore_col_vec,
             mat_create,
             mat_set_sizes,

devito/petsc/iet/type_builder.py

@@ -290,6 +290,8 @@ def _extend_build(self, base_dict):
     'submat_arr': PointerMat(name='submat_arr'),
     'subblockrows': PetscInt('subblockrows'),
     'subblockcols': PetscInt('subblockcols'),
+    'sublocalrows': PetscInt('sublocalrows'),
+    'sublocalcols': PetscInt('sublocalcols'),
     'rowidx': PetscInt('rowidx'),
     'colidx': PetscInt('colidx'),
     'J': Mat('J'),

devito/petsc/solve.py

@@ -142,7 +142,6 @@ def linear_solve_args(self):
         residual = Residual(target, exprs, arrays, self.time_mapper, jacobian.scdiag)
         initial_guess = InitialGuess(target, exprs, arrays, self.time_mapper)
 
-        # TODO: Extend this to mixed case
         constrain_bc = (
             ConstrainBC(target, exprs, arrays)
             if self.constrain_bcs

devito/petsc/types/object.py

@@ -402,11 +402,6 @@ class NofSubMats(Scalar, LocalType):
     pass
 
 
-# Can this be attached to the consrain bc object in metadata maybe? probs
-# shoulnd't be here
-Counter = PetscInt(name='count')
-
-
 FREE_PRIORITY = {
     JacobianStruct: 0,
     PETScArrayObject: 1,