Merge pull request #2266 from devitocodes/subdim-args

misc: Process args for subdimensions

devito/core/autotuning.py

@@ -51,7 +51,8 @@ def autotune(operator, args, level, mode):
         # WARNING: `copies` keeps references to numpy arrays, which is required
         # to avoid garbage collection to kick in during autotuning and prematurely
         # free the shadow copies handed over to C-land
-        at_args.update({k: writes[k]._C_make_dataobj(v) for k, v in copies.items()})
+        at_args.update({k: writes[k]._C_make_dataobj(alias=writes[k], **copies)
+                        for k in copies})
 
     # Disable halo exchanges through MPI_PROC_NULL
     if mode in ['preemptive', 'destructive']:

devito/ir/support/space.py

@@ -993,7 +993,8 @@ def itintervals(self):
     @cached_property
     def dimensions(self):
         sub_dims = flatten(i._defines for v in self.sub_iterators.values() for i in v)
-        return tuple(filter_ordered(self.itdims + tuple(sub_dims)))
+        r_dims = flatten(i._defines for v in self.relations for i in v)
+        return tuple(filter_ordered(self.itdims + tuple(sub_dims) + tuple(r_dims)))
 
     @cached_property
     def nonderived_directions(self):

devito/types/dense.py

@@ -659,11 +659,13 @@ def initializer(self):
                                           (_C_field_owned_ofs, POINTER(c_int)),
                                           (_C_field_dmap, c_void_p)]}))
 
-    def _C_make_dataobj(self, data):
+    def _C_make_dataobj(self, alias=None, **args):
         """
         A ctypes object representing the DiscreteFunction that can be passed to
         an Operator.
         """
+        key = alias or self
+        data = args[key.name]
         dataobj = byref(self._C_ctype._type_())
         dataobj._obj.data = data.ctypes.data_as(c_restrict_void_p)
         dataobj._obj.size = (c_ulong*self.ndim)(*data.shape)
@@ -863,7 +865,7 @@ def _arg_check(self, args, intervals, **kwargs):
 
     def _arg_finalize(self, args, alias=None):
         key = alias or self
-        return {key.name: self._C_make_dataobj(args[key.name])}
+        return {key.name: self._C_make_dataobj(alias=key, **args)}
 
     # Pickling support
 

devito/types/dimension.py

@@ -1,4 +1,5 @@
 from collections import namedtuple
+import math
 
 import sympy
 from sympy.core.decorators import call_highest_priority
@@ -889,6 +890,23 @@ def free_symbols(self):
             pass
         return retval
 
+    def _arg_values(self, interval, grid=None, **kwargs):
+        # Parent dimension define the interval
+        fact = self._factor.data if self._factor is not None else 1
+        toint = lambda x: math.ceil(x / fact)
+        vals = {}
+        try:
+            vals[self.min_name] = toint(kwargs.get(self.parent.min_name))
+        except (KeyError, TypeError):
+            pass
+
+        try:
+            vals[self.max_name] = toint(kwargs.get(self.parent.max_name))
+        except (KeyError, TypeError):
+            pass
+
+        return vals
+
     def _arg_defaults(self, _min=None, size=None, alias=None):
         defaults = super()._arg_defaults(_min=_min, size=size, alias=alias)