Merge pull request #357 from opesci/yask-integration-phase12

yask: Small fixes to make SIMD work, and tests

devito/yask/__init__.py

@@ -81,21 +81,32 @@ def __init__(self, *args, **kwargs):
 yask_configuration = Parameters('YASK-Configuration')
 yask_configuration.add('compiler', YaskCompiler())
 yask_configuration.add('python-exec', False, [False, True])
-yask_configuration.add('develop-mode', True, [False, True])
-# Sniff the highest Instruction Set Architecture level that we can YASK to use
+# Set the Instruction Set Architecture used by the YASK code generator
 isa, ISAs = 'cpp', ['cpp', 'avx', 'avx2', 'avx512', 'knc']
-if yask_configuration['develop-mode'] is False:
-    cpu_flags = cpuinfo.get_cpu_info()['flags']
-    for i in reversed(ISAs):
-        if i in cpu_flags:
-            isa = i
-            break
 yask_configuration.add('isa', isa, ISAs)
 # Currently YASK also require the CPU architecture (e.g., snb for sandy bridge,
 # hsw for haswell, etc.). At the moment, we simply infer it from the ISA
 arch_mapper = {'cpp': 'intel64', 'avx': 'snb', 'avx2': 'hsw', 'avx512': 'knl'}
 yask_configuration.add('arch', arch_mapper[isa], arch_mapper.values())
 
+
+# In develop-mode, no optimizations are applied to the generated code (e.g., SIMD)
+# When switching to non-develop-mode, optimizations are automatically switched on,
+# sniffing the highest Instruction Set Architecture level available on the current
+# machine and providing it to YASK
+def reset_yask_isa(develop_mode):
+    if develop_mode is True:
+        return
+    cpu_flags = cpuinfo.get_cpu_info()['flags']
+    isa = 'cpp'
+    for i in reversed(ISAs):
+        if i in cpu_flags:
+            isa = i
+            break
+    yask_configuration['isa'] = isa
+    yask_configuration['arch'] = arch_mapper[isa]
+yask_configuration.add('develop-mode', True, [False, True], reset_yask_isa)  # noqa
+
 configuration.add('yask', yask_configuration)
 
 log("Backend successfully initialized!")

devito/yask/wrappers.py

@@ -37,23 +37,40 @@ class YaskGrid(object):
     def __init__(self, grid, shape, radius, dtype):
         """
         Initialize a new :class:`YaskGrid`.
+
+        The storage layout adopted by YASK is as follows: ::
+
+            --------------------------------------------------------------
+            | extra_padding | halo |              | halo | extra_padding |
+            ------------------------    domain    ------------------------
+            |       padding        |              |       padding        |
+            --------------------------------------------------------------
+            |                         allocation                         |
+            --------------------------------------------------------------
+
+        :param grid: The YASK yk::grid that will be wrapped. Data storage will be
+                     allocated if not yet available.
+        :param shape: The "visibility region" of the YaskGrid. The shape should be
+                      at least as big as the domain (in each dimension). If larger,
+                      then users will be allowed to access more data entries,
+                      such as those lying on the halo region.
+        :param radius: The extent of the halo region.
+        :param dtype: The type of the raw data.
         """
         self.grid = grid
         self.shape = shape
         self.dtype = dtype
 
-        # Set up halo sizes
-        self.halo = [0 if i == namespace['time-dim'] else radius
-                     for i in self.dimensions]
-
-        # Allocate memory in YASK-land and initialize it to 0
-        # The following is_storage_allocated check is needed because of self.with_halo
         if not self.is_storage_allocated():
+            # Allocate memory in YASK-land and initialize it to 0
             for i, j in zip(self.dimensions, shape):
                 if i == namespace['time-dim']:
                     assert self.grid.is_dim_used(i)
                     self.grid.set_alloc_size(i, j)
                 else:
+                    # Note, from the YASK docs:
+                    # "If the halo is set to a value larger than the padding size,
+                    # the padding size will be automatically increase to accomodate it."
                     self.grid.set_halo_size(i, radius)
             self.grid.alloc_storage()
             self._reset()
@@ -136,13 +153,22 @@ def _reset(self):
         """
         self[:] = 0.0
 
+    @property
+    def _halo(self):
+        return [0 if i == namespace['time-dim'] else self.get_halo_size(i)
+                for i in self.dimensions]
+
+    @property
+    def _padding(self):
+        return [0 if i == namespace['time-dim'] else self.get_pad_size(i)
+                for i in self.dimensions]
+
     @property
     def _offsets(self):
         offsets = []
-        for i, j in zip(self.dimensions, self.halo):
+        for i, j in zip(self.dimensions, self._padding):
             ofs = 0 if i == namespace['time-dim'] else self.get_first_rank_alloc_index(i)
-            ofs += j
-            offsets.append(ofs)
+            offsets.append(ofs + j)
         return offsets
 
     @property
@@ -152,8 +178,7 @@ def with_halo(self):
         unmasked. This allows the caller to write/read the halo region as well as
         the domain.
         """
-        shape = [i + 2*j for i, j in zip(self.shape, self.halo)]
-        return YaskGrid(self.grid, shape, 0, self.dtype)
+        return YaskGridWithHalo(self.grid, self.shape, 0, self.dtype)
 
     @property
     def name(self):
@@ -196,6 +221,20 @@ def view(self):
         return self[:]
 
 
+class YaskGridWithHalo(YaskGrid):
+
+    """A helper class for YaskGrid wrappers providing access to the halo region."""
+
+    def __init__(self, grid, shape, radius, dtype):
+        super(YaskGridWithHalo, self).__init__(grid, shape, radius, dtype)
+        self.shape = [i + 2*j for i, j in zip(self.shape, self._halo)]
+
+    @property
+    def _offsets(self):
+        offsets = super(YaskGridWithHalo, self)._offsets
+        return [i - j for i, j in zip(offsets, self._halo)]
+
+
 class YaskKernel(object):
 
     """
@@ -423,7 +462,7 @@ def fetch(self, dimensions, shape, dtype):
                               if i != namespace['time-dim']])
 
         # A unique key for this context.
-        key = tuple([dtype] + domain.items())
+        key = tuple([yask_configuration['isa'], dtype] + domain.items())
 
         # Fetch or create a YaskContext
         if key in self:

tests/test_yask.py

@@ -8,6 +8,7 @@
 from devito import (Operator, DenseData, TimeData, PointData,
                     time, t, x, y, z, configuration, clear_cache)  # noqa
 from devito.dle import retrieve_iteration_tree  # noqa
+from devito.yask import arch_mapper, yask_configuration  # noqa
 from devito.yask.wrappers import YaskGrid, contexts  # noqa
 
 # For the acoustic wave test
@@ -130,8 +131,15 @@ class TestOperatorSimple(object):
     def setup_class(cls):
         clear_cache()
 
+    @pytest.fixture(scope='class', autouse=True)
+    def reset_isa(self):
+        """Force back to NO-SIMD after each test, as some tests may optionally
+        switch on SIMD."""
+        yask_configuration['develop-mode'] = True
+
     @pytest.mark.parametrize("space_order", [0, 1, 2])
-    def test_increasing_halo_wo_ofs(self, space_order):
+    @pytest.mark.parametrize("nosimd", [True, False])
+    def test_increasing_halo_wo_ofs(self, space_order, nosimd):
         """
         Apply the trivial equation ``u[t+1,x,y,z] = u[t,x,y,z] + 1`` and check
         that increasing space orders lead to proportionately larger halo regions,
@@ -156,6 +164,9 @@ def test_increasing_halo_wo_ofs(self, space_order):
 
         And so on and so forth.
         """
+        # SIMD on/off
+        yask_configuration['develop-mode'] = nosimd
+
         u = TimeData(name='yu4D', shape=(16, 16, 16), dimensions=(x, y, z),
                      space_order=space_order)
         u.data.with_halo[:] = 0.
@@ -319,7 +330,7 @@ def model(self):
         shape = (60, 70, 80)
         nbpml = 10
         return demo_model(spacing=[15, 15, 15], shape=shape, nbpml=nbpml,
-                          preset='layers', ratio=3)
+                          preset='layers-isotropic', ratio=3)
 
     @pytest.fixture
     def time_params(self, model):