compiler: Make Parizer more easily overridable

devito/passes/iet/languages/openmp.py

@@ -78,14 +78,6 @@ def _make_clauses(cls, **kwargs):
 
         return clauses
 
-    @classmethod
-    def _process_kwargs(cls, **kwargs):
-        kwargs = super()._process_kwargs(**kwargs)
-
-        kwargs.pop('gpu_fit', None)
-
-        return kwargs
-
 
 class ThreadedProdder(Conditional, Prodder):
 

devito/passes/iet/parpragma.py

@@ -13,8 +13,7 @@
 from devito.passes.iet.langbase import (LangBB, LangTransformer, DeviceAwareMixin,
                                         make_sections_from_imask)
 from devito.symbolics import INT, ccode
-from devito.tools import as_tuple, flatten, is_integer, prod
-from devito.tools.data_structures import UnboundTuple
+from devito.tools import UnboundTuple, as_tuple, flatten, is_integer, prod
 from devito.types import Symbol
 
 __all__ = ['PragmaSimdTransformer', 'PragmaShmTransformer',
@@ -45,8 +44,21 @@ def _support_array_reduction(cls, compiler):
     def simd_reg_size(self):
         return self.platform.simd_reg_size
 
-    @iet_pass
-    def make_simd(self, iet):
+    def _make_simd_pragma(self, iet):
+        indexeds = FindSymbols('indexeds').visit(iet)
+        aligned = {i.name for i in indexeds if i.function.is_DiscreteFunction}
+        if aligned:
+            simd = self.lang['simd-for-aligned']
+            simd = as_tuple(simd(','.join(sorted(aligned)), self.simd_reg_size))
+        else:
+            simd = as_tuple(self.lang['simd-for'])
+
+        return simd
+
+    def _make_simd(self, iet):
+        """
+        Carry out the bulk of `make_simd`.
+        """
         mapper = {}
         for tree in retrieve_iteration_tree(iet):
             candidates = [i for i in tree if i.is_ParallelRelaxed]
@@ -101,13 +113,7 @@ def make_simd(self, iet):
                     continue
 
             # Add SIMD pragma
-            indexeds = FindSymbols('indexeds').visit(candidate)
-            aligned = {i.name for i in indexeds if i.function.is_DiscreteFunction}
-            if aligned:
-                simd = self.lang['simd-for-aligned']
-                simd = as_tuple(simd(','.join(sorted(aligned)), self.simd_reg_size))
-            else:
-                simd = as_tuple(self.lang['simd-for'])
+            simd = self._make_simd_pragma(candidate)
             pragmas = candidate.pragmas + simd
 
             # Add VECTORIZED property
@@ -119,14 +125,20 @@ def make_simd(self, iet):
 
         return iet, {}
 
+    @iet_pass
+    def make_simd(self, iet):
+        return self._make_simd(iet)
+
 
 class PragmaIteration(ParallelIteration):
 
     def __init__(self, *args, parallel=None, schedule=None, chunk_size=None,
                  nthreads=None, ncollapsed=None, reduction=None, tile=None,
                  gpu_fit=None, **kwargs):
 
-        construct = self._make_construct(parallel=parallel)
+        construct = self._make_construct(
+            parallel=parallel, ncollapsed=ncollapsed, tile=tile
+        )
         clauses = self._make_clauses(
             ncollapsed=ncollapsed, chunk_size=chunk_size, nthreads=nthreads,
             reduction=reduction, schedule=schedule, tile=tile, gpu_fit=gpu_fit,
@@ -604,6 +616,7 @@ def _make_partree(self, candidates, nthreads=None, index=None):
         if self._is_offloadable(root):
             body = self.DeviceIteration(gpu_fit=self.gpu_fit,
                                         ncollapsed=len(collapsable) + 1,
+                                        tile=self.par_tile,
                                         **root.args)
             partree = ParallelTree([], body, nthreads=nthreads)
 

devito/tools/data_structures.py

@@ -11,7 +11,7 @@
 
 __all__ = ['Bunch', 'EnrichedTuple', 'ReducerMap', 'DefaultOrderedDict',
            'OrderedSet', 'PartialOrderTuple', 'DAG', 'frozendict',
-           'UnboundedMultiTuple']
+           'UnboundTuple', 'UnboundedMultiTuple']
 
 
 class Bunch(object):