dsl: Further work on making MultiSubDomains overrideable

Author: EdCaunt

devito/operator/operator.py

@@ -559,6 +559,7 @@ def _prepare_arguments(self, autotune=None, **kwargs):
         Process runtime arguments passed to ``.apply()` and derive
         default values for any remaining arguments.
         """
+        from IPython import embed; embed()
         # Sanity check -- all user-provided keywords must be known to the Operator
         if not configuration['ignore-unknowns']:
             for k, v in kwargs.items():

devito/types/grid.py

@@ -685,8 +685,7 @@ def __setstate__(self, state):
             self._distributor = SubDistributor(self)
 
 
-class MultiSubDomain(AbstractSubDomain):
-    # FIXME: Should subclass ArgProvider
+class MultiSubDomain(AbstractSubDomain, ArgProvider):
 
     """
     Abstract base class for types representing groups of SubDomains.
@@ -745,6 +744,26 @@ def _bounds_glb_to_loc(cls, dec, m, M):
         return bounds_m, bounds_M
 
 
+class MultiSubDomainFunction(Function):
+    def _arg_values(self, **kwargs):
+        # TODO: Will want estimate-memory utility in due course
+        new = kwargs.get(self.name, self)
+
+        # We support `op.apply(..., msd=msd1, ...)`
+        if isinstance(new, MultiSubDomain):
+            # TODO: Homogenise this nomenclature throughout all MultiSubDomains
+            new = new._subfunction
+
+        return new._arg_defaults(alias=self, **kwargs).reduce_all()
+
+    def _arg_defaults(self, alias=None, **kwargs):
+        # Either apply the overide here, or apply the default values
+        # Will need to replace the function, but also adjust the implicit
+        # dimension bounds -> this bit is trickier, since they only get numbered
+        # during compilation
+        pass
+
+
 class SubDomainSet(MultiSubDomain):
     """
     Class to define a set of N (a positive integer) subdomains.
@@ -875,9 +894,9 @@ def __subdomain_finalize_core__(self, grid):
         # of replacements without risking overwriting.
         i_dim = Dimension(f'n_{str(id(self))}')
         d_dim = DefaultDimension(name='d', default_value=2*grid.dim)
-        sd_func = Function(name=self.name, grid=self._grid,
-                           shape=(self._n_domains, 2*grid.dim),
-                           dimensions=(i_dim, d_dim), dtype=np.int32)
+        sd_func = MultiSubDomainFunction(name=self.name, grid=self._grid,
+                                         shape=(self._n_domains, 2*grid.dim),
+                                         dimensions=(i_dim, d_dim), dtype=np.int32)
 
         dimensions = []
         for i, d in enumerate(grid.dimensions):
@@ -892,6 +911,7 @@ def __subdomain_finalize_core__(self, grid):
             ))
 
         self._dimensions = tuple(dimensions)
+        self._subfunction = sd_func
 
     def __subdomain_finalize__(self):
         self.__subdomain_finalize_core__(self.grid)

tests/test_subdomains.py

@@ -816,8 +816,9 @@ def test_override(self):
         grid = Grid(shape=(6, 8))
         x, y = grid.dimensions
 
-        border0 = Border(grid, 1)
-        border1 = Border(grid, 1, dims={x: 'left', y: 'right'}, corners='nocorners')
+        border0 = Border(grid, 1, name='border0')
+        border1 = Border(grid, 1, dims={x: 'left', y: 'right'},
+                         corners='nocorners', name='border1')
 
         f0 = Function(name='f0', grid=grid, dtype=np.int32)
         f1 = Function(name='f1', grid=grid, dtype=np.int32)