diff --git a/devito/passes/iet/languages/openmp.py b/devito/passes/iet/languages/openmp.py index c5c0fab8247..60b0eb0c413 100644 --- a/devito/passes/iet/languages/openmp.py +++ b/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): diff --git a/devito/passes/iet/parpragma.py b/devito/passes/iet/parpragma.py index 44ee6afd6cc..f38480002ac 100644 --- a/devito/passes/iet/parpragma.py +++ b/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,6 +125,10 @@ def make_simd(self, iet): return iet, {} + @iet_pass + def make_simd(self, iet): + return self._make_simd(iet) + class PragmaIteration(ParallelIteration): @@ -126,7 +136,9 @@ 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) diff --git a/devito/tools/data_structures.py b/devito/tools/data_structures.py index 539f75d5934..9c5710bce24 100644 --- a/devito/tools/data_structures.py +++ b/devito/tools/data_structures.py @@ -11,7 +11,7 @@ __all__ = ['Bunch', 'EnrichedTuple', 'ReducerMap', 'DefaultOrderedDict', 'OrderedSet', 'PartialOrderTuple', 'DAG', 'frozendict', - 'UnboundedMultiTuple'] + 'UnboundTuple', 'UnboundedMultiTuple'] class Bunch(object):