tests/compiler: Add unique options prefixes to tests and modify CoupledCallbackBuilder to generalise to support constraining bcs for mixed problems

Author: ZoeLeibowitz

devito/petsc/iet/callbacks.py

@@ -1128,17 +1128,43 @@ def _submat_callback_body(self):
 
         dm_get_info = petsc_call(
             'DMDAGetInfo', [
-                sobjs['callbackdm'], Null, Byref(sobjs['M']), Byref(sobjs['N']),
+                sobjs['callbackdm'], Null, Null, Null,
                 Null, Null, Null, Null, Byref(objs['dof']), Null, Null, Null, Null, Null
             ]
         )
-        subblock_rows = DummyExpr(objs['subblockrows'], Mul(sobjs['M'], sobjs['N']))
-        subblock_cols = DummyExpr(objs['subblockcols'], Mul(sobjs['M'], sobjs['N']))
 
         ptr = DummyExpr(
             objs['submat_arr']._C_symbol, Deref(objs['Submats']._C_symbol), init=True
         )
 
+        i = Dimension(name='i')
+        tmpvec = sobjs['tmpvec']
+
+        row_idx = DummyExpr(objs['rowidx'], IntDiv(i, objs['dof']))
+        col_idx = DummyExpr(objs['colidx'], Mod(i, objs['dof']))
+
+        # Query constrained global size from each sub-DM via a temporary Vec.
+        # For unconstrained sub-DMs this is equivalent to M*N; for constrained
+        # (BC-excluded) sub-DMs it returns the reduced size automatically.
+        get_row_vec = petsc_call(
+            'DMGetGlobalVector', [objs['Subdms'].indexed[objs['rowidx']], Byref(tmpvec)]
+        )
+        get_row_size = petsc_call(
+            'VecGetSize', [tmpvec, Byref(objs['subblockrows'])]
+        )
+        restore_row_vec = petsc_call(
+            'DMRestoreGlobalVector', [objs['Subdms'].indexed[objs['rowidx']], Byref(tmpvec)]
+        )
+        get_col_vec = petsc_call(
+            'DMGetGlobalVector', [objs['Subdms'].indexed[objs['colidx']], Byref(tmpvec)]
+        )
+        get_col_size = petsc_call(
+            'VecGetSize', [tmpvec, Byref(objs['subblockcols'])]
+        )
+        restore_col_vec = petsc_call(
+            'DMRestoreGlobalVector', [objs['Subdms'].indexed[objs['colidx']], Byref(tmpvec)]
+        )
+
         mat_create = petsc_call('MatCreate', [sobjs['comm'], Byref(objs['block'])])
 
         mat_set_sizes = petsc_call(
@@ -1151,10 +1177,6 @@ def _submat_callback_body(self):
         mat_set_type = petsc_call('MatSetType', [objs['block'], 'MATSHELL'])
 
         malloc = petsc_call('PetscMalloc1', [1, Byref(objs['subctx'])])
-        i = Dimension(name='i')
-
-        row_idx = DummyExpr(objs['rowidx'], IntDiv(i, objs['dof']))
-        col_idx = DummyExpr(objs['colidx'], Mod(i, objs['dof']))
 
         deref_subdm = Dereference(objs['Subdms'], objs['ljacctx'])
 
@@ -1193,12 +1215,18 @@ def _submat_callback_body(self):
         assign_block = DummyExpr(objs['submat_arr'].indexed[i], objs['block'])
 
         iter_body = (
+            row_idx,
+            col_idx,
+            get_row_vec,
+            get_row_size,
+            restore_row_vec,
+            get_col_vec,
+            get_col_size,
+            restore_col_vec,
             mat_create,
             mat_set_sizes,
             mat_set_type,
             malloc,
-            row_idx,
-            col_idx,
             set_rows,
             set_cols,
             dm_set_ctx,
@@ -1233,8 +1261,6 @@ def _submat_callback_body(self):
             mat_get_dm,
             dm_get_app,
             dm_get_info,
-            subblock_rows,
-            subblock_cols,
             ptr,
             BlankLine,
             iteration,

devito/petsc/iet/type_builder.py

@@ -108,13 +108,6 @@ def _extend_build(self, base_dict):
         )
         base_dict['nfields'] = PetscInt(sreg.make_name(prefix='nfields'))
 
-        space_dims = len(self.field_data.grid.dimensions)
-
-        dim_labels = ["M", "N", "P"]
-        base_dict.update({
-            dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims)
-        })
-
         submatrices = self.field_data.jacobian.nonzero_submatrices
 
         base_dict['jacctx'] = JacobianStruct(
@@ -133,6 +126,10 @@ def _extend_build(self, base_dict):
             base_dict[f'{name}Y'] = Vec(f'{name}Y', destroy=False)
             base_dict[f'{name}F'] = Vec(f'{name}F', destroy=False)
 
+        # Temporary Vec used in MatCreateSubMatrices to query constrained
+        # global sizes from sub-DMs (works for both constrained and unconstrained).
+        base_dict['tmpvec'] = Vec('tmpvec', destroy=False)
+
         # Bundle objects/metadata required by the coupled residual callback
         f_components, x_components = [], []
         bundle_mapper = {}

tests/test_petsc.py

@@ -229,9 +229,9 @@ def test_petsc_cast():
     eqn2 = Eq(f2.laplace, 10)
     eqn3 = Eq(f3.laplace, 10)
 
-    petsc1 = petscsolve(eqn1, f1)
-    petsc2 = petscsolve(eqn2, f2)
-    petsc3 = petscsolve(eqn3, f3)
+    petsc1 = petscsolve(eqn1, f1, options_prefix='cast1d')
+    petsc2 = petscsolve(eqn2, f2, options_prefix='cast2d')
+    petsc3 = petscsolve(eqn3, f3, options_prefix='cast3d')
 
     with switchconfig(language='petsc'):
         op1 = Operator(petsc1)
@@ -261,9 +261,9 @@ def test_dmda_create():
     eqn2 = Eq(f2.laplace, 10)
     eqn3 = Eq(f3.laplace, 10)
 
-    petsc1 = petscsolve(eqn1, f1)
-    petsc2 = petscsolve(eqn2, f2)
-    petsc3 = petscsolve(eqn3, f3)
+    petsc1 = petscsolve(eqn1, f1, options_prefix='dmda1d')
+    petsc2 = petscsolve(eqn2, f2, options_prefix='dmda2d')
+    petsc3 = petscsolve(eqn3, f3, options_prefix='dmda3d')
 
     with switchconfig(language='petsc'):
         op1 = Operator(petsc1, opt='noop')