Merge pull request #355 from opesci/yask-integration-phase11

yask: Fixes to make the acoustic wave equation test take off

devito/cgen_utils.py

@@ -145,7 +145,7 @@ def _print_FrozenExpr(self, expr):
 
     def _print_FunctionFromPointer(self, expr):
         indices = [self._print(i) for i in expr.params]
-        return "%s->%s(%s)" % (expr.pointer, expr.name, ', '.join(indices))
+        return "%s->%s(%s)" % (expr.pointer, expr.function, ', '.join(indices))
 
 
 def ccode(expr, **settings):

devito/dse/extended_sympy.py

@@ -54,14 +54,15 @@ class Add(sympy.Add, FrozenExpr):
 
 class FunctionFromPointer(sympy.Symbol):
 
-    def __new__(cls, name, pointer, params=None):
-        obj = sympy.Symbol.__new__(cls, name)
+    def __new__(cls, function, pointer, params=None):
+        obj = sympy.Symbol.__new__(cls, '%s->%s' % (pointer, function))
+        obj.function = function
         obj.pointer = pointer
         obj.params = params or ()
         return obj
 
     def __str__(self):
-        return "%s->%s(%s)" % (self.pointer, self.name,
+        return "%s->%s(%s)" % (self.pointer, self.function,
                                ', '.join([str(i) for i in self.params]))
 
     __repr__ = __str__

devito/operator.py

@@ -453,13 +453,16 @@ def apply(self, **kwargs):
         self.cfunction(*list(arguments.values()))
 
         # Output summary of performance achieved
+        return self._profile_output(dim_sizes)
+
+    def _profile_output(self, dim_sizes):
+        """Return a performance summary of the profiled sections."""
         summary = self.profiler.summary(dim_sizes, self.dtype)
         with bar():
             for k, v in summary.items():
                 name = '%s<%s>' % (k, ','.join('%d' % i for i in v.itershape))
                 info("Section %s with OI=%.2f computed in %.3f s [Perf: %.2f GFlops/s]" %
                      (name, v.oi, v.time, v.gflopss))
-
         return summary
 
     def _profile_sections(self, nodes, parameters):

devito/yask/__init__.py

@@ -6,6 +6,8 @@
 from collections import OrderedDict
 import os
 
+import cpuinfo
+
 from devito import configuration
 from devito.exceptions import InvalidOperator
 from devito.logger import yask as log
@@ -79,9 +81,21 @@ def __init__(self, *args, **kwargs):
 yask_configuration = Parameters('YASK-Configuration')
 yask_configuration.add('compiler', YaskCompiler())
 yask_configuration.add('python-exec', False, [False, True])
-# TODO: this should be somewhat sniffed
-yask_configuration.add('arch', 'snb', ['snb'])
-yask_configuration.add('isa', 'cpp', ['cpp'])
+yask_configuration.add('develop-mode', True, [False, True])
+# Sniff the highest Instruction Set Architecture level that we can YASK to use
+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())
+
 configuration.add('yask', yask_configuration)
 
 log("Backend successfully initialized!")

devito/yask/operator.py

@@ -196,6 +196,9 @@ def apply(self, **kwargs):
             self.cfunction(*list(arguments.values()))
         log("YASK Operator successfully run!")
 
+        # Output summary of performance achieved
+        return self._profile_output(dim_sizes)
+
     @property
     def compile(self):
         """

devito/yask/utils.py

@@ -21,14 +21,14 @@ def make_grid_accesses(node):
     """
 
     def make_grid_gets(expr):
+        mapper = {}
         indexeds = retrieve_indexed(expr)
         data_carriers = [i for i in indexeds if i.base.function.from_YASK]
         for i in data_carriers:
             name = namespace['code-grid-name'](i.base.function.name)
             args = [INT(make_grid_gets(j)) for j in i.indices]
-            handle = make_sharedptr_funcall(namespace['code-grid-get'], args, name)
-            expr = expr.xreplace({i: handle})
-        return expr
+            mapper[i] = make_sharedptr_funcall(namespace['code-grid-get'], args, name)
+        return expr.xreplace(mapper)
 
     mapper = {}
     for i, e in enumerate(FindNodes(Expression).visit(node)):
@@ -45,7 +45,7 @@ def make_grid_gets(expr):
             processed = Element(c.Statement(ccode(handle)))
         else:
             # Writing to a scalar temporary
-            processed = Expression(e.expr.func(lhs, rhs))
+            processed = Expression(e.expr.func(lhs, rhs), dtype=e.dtype)
 
         mapper.update({e: processed})
 

devito/yask/wrappers.py

@@ -229,10 +229,12 @@ def __init__(self, name, yc_soln, domain):
 
         # JIT-compile it
         try:
-            make(os.environ['YASK_HOME'], ['-j', 'YK_CXXOPT=-O0',
-                                           "EXTRA_MACROS=TRACE",
+            opt_level = 1 if yask_configuration['develop-mode'] else 3
+            make(os.environ['YASK_HOME'], ['-j', 'YK_CXXOPT=-O%d' % opt_level,
+                                           # "EXTRA_MACROS=TRACE",
                                            'YK_BASE=%s' % str(name),
                                            'stencil=%s' % yc_soln.get_name(),
+                                           'arch=%s' % yask_configuration['arch'],
                                            '-C', namespace['kernel-path'], 'api'])
         except CompilationError:
             exit("Kernel solution compilation")

tests/test_yask.py

@@ -6,18 +6,25 @@
 pexpect = pytest.importorskip('yask_compiler')  # Run only if YASK is available
 
 from devito import (Operator, DenseData, TimeData, PointData,
-                    t, x, y, z, configuration, clear_cache, info)  # noqa
+                    time, t, x, y, z, configuration, clear_cache)  # noqa
 from devito.dle import retrieve_iteration_tree  # noqa
 from devito.yask.wrappers import YaskGrid, contexts  # noqa
 
 # For the acoustic wave test
 from examples.seismic.acoustic import AcousticWaveSolver  # noqa
-from examples.seismic import demo_model, RickerSource, Receiver  # noqa
+from examples.seismic import demo_model, PointSource, RickerSource, Receiver  # noqa
 
 pytestmark = pytest.mark.skipif(configuration['backend'] != 'yask',
                                 reason="'yask' wasn't selected as backend on startup")
 
 
+def setup_module(module):
+    """Get rid of any YASK modules generated and JIT-compiled in previous runs.
+    This is not strictly necessary for the tests, but it helps in keeping the
+    lib directory clean, which may be helpful for offline analysis."""
+    contexts.dump()
+
+
 @pytest.fixture(scope="module")
 def u(dims):
     u = DenseData(name='yu3D', shape=(16, 16, 16), dimensions=(x, y, z), space_order=0)
@@ -113,7 +120,7 @@ def test_data_arithmetic_nD(self, u):
         assert np.all(arr - u.data == -1.)
 
 
-class TestOperatorExecution(object):
+class TestOperatorSimple(object):
 
     """
     Test execution of "toy" Operators through YASK.
@@ -162,6 +169,20 @@ def test_increasing_halo_wo_ofs(self, space_order):
         assert all(np.all(u.data.with_halo[1, :, i, :] == 0) for i in range(space_order))
         assert all(np.all(u.data.with_halo[1, :, :, i] == 0) for i in range(space_order))
 
+    def test_increasing_multi_steps(self):
+        """
+        Apply the trivial equation ``u[t+1,x,y,z] = u[t,x,y,z] + 1`` for 11
+        timesteps and check that all grid domain values are equal to 11 within
+        ``u[1]`` and equal to 10 within ``u[0]``.
+        """
+        u = TimeData(name='yu4D', shape=(8, 8, 8), dimensions=(x, y, z), space_order=0)
+        u.data.with_halo[:] = 0.
+        op = Operator(Eq(u.forward, u + 1.), subs={t.spacing: 1})
+        op(yu4D=u, t=12)
+        assert 'run_solution' in str(op)
+        assert np.all(u.data[0] == 10.)
+        assert np.all(u.data[1] == 11.)
+
     @pytest.mark.parametrize("space_order", [2])
     def test_fixed_halo_w_ofs(self, space_order):
         """
@@ -280,7 +301,7 @@ def test_irregular_write(self):
         assert all(np.all(u.data[1, :, :, i] == 3 - i) for i in range(4))
 
 
-class TestOperatorRealAcoustic(object):
+class TestOperatorAcoustic(object):
 
     """
     Test the acoustic wave model through YASK.
@@ -293,51 +314,110 @@ class TestOperatorRealAcoustic(object):
         'layers': {'preset': 'layers', 'ratio': 3},
     }
 
-    @pytest.mark.parametrize('mkey, dimensions, time_order, space_order, nbpml', [
-        # 3D tests with varying space orders
-        pytest.mark.xfail(('layers', (60, 70, 80), 2, 4, 10)),
-        pytest.mark.xfail(('layers', (60, 70, 80), 2, 8, 10)),
-    ])
-    def test_acoustic(self, mkey, dimensions, time_order, space_order, nbpml):
-        t0 = 0.0  # Start time
-        tn = 500.  # Final time
-        nrec = 130  # Number of receivers
-
-        # Create model from preset
-        model = demo_model(spacing=[15. for _ in dimensions],
-                           shape=dimensions, nbpml=nbpml, **(self.presets[mkey]))
+    @pytest.fixture
+    def model(self):
+        shape = (60, 70, 80)
+        nbpml = 10
+        return demo_model(spacing=[15, 15, 15], shape=shape, nbpml=nbpml,
+                          preset='layers', ratio=3)
 
+    @pytest.fixture
+    def time_params(self, model):
         # Derive timestepping from model spacing
-        dt = model.critical_dt * (1.73 if time_order == 4 else 1.0)
+        t0 = 0.0  # Start time
+        tn = 500.  # Final time
+        dt = model.critical_dt
         nt = int(1 + (tn-t0) / dt)  # Number of timesteps
-        time_values = np.linspace(t0, tn, nt)  # Discretized time axis
-
+        return t0, tn, nt
+
+    @pytest.fixture
+    def m(self, model):
+        return model.m
+
+    @pytest.fixture
+    def damp(self, model):
+        return model.damp
+
+    @pytest.fixture
+    def u(self, model):
+        space_order = 4
+        time_order = 2
+        return TimeData(name='u', shape=model.shape_domain, dimensions=(x, y, z),
+                        space_order=space_order, time_order=time_order)
+
+    @pytest.fixture
+    def stencil(self, m, damp, u):
+        s = t.spacing
+        stencil = 1.0 / (2.0 * m + s * damp) * (
+            4.0 * m * u + (s * damp - 2.0 * m) * u.backward +
+            2.0 * s**2 * u.laplace)
+        return stencil
+
+    @pytest.fixture
+    def src(self, model, time_params):
+        time_values = np.linspace(*time_params)  # Discretized time axis
         # Define source geometry (center of domain, just below surface)
         src = RickerSource(name='src', ndim=model.dim, f0=0.01, time=time_values)
         src.coordinates.data[0, :] = np.array(model.domain_size) * .5
         src.coordinates.data[0, -1] = 30.
+        return src
 
-        # Define receiver geometry (same as source, but spread across x)
-        rec = Receiver(name='nrec', ntime=nt, npoint=nrec, ndim=model.dim)
+    @pytest.fixture
+    def rec(self, model, time_params, src):
+        nrec = 130  # Number of receivers
+        t0, tn, nt = time_params
+        rec = Receiver(name='rec', ntime=nt, npoint=nrec, ndim=model.dim)
         rec.coordinates.data[:, 0] = np.linspace(0., model.domain_size[0], num=nrec)
         rec.coordinates.data[:, 1:] = src.coordinates.data[0, 1:]
+        return rec
+
+    @pytest.fixture
+    def subs(self, model, u):
+        dt = model.critical_dt
+        return dict([(t.spacing, dt)] + [(time.spacing, dt)] +
+                    [(i.spacing, model.get_spacing()[j]) for i, j
+                     in zip(u.indices[1:], range(len(model.shape)))])
+
+    def test_acoustic_wo_src_wo_rec(self, model, stencil, subs, m, damp, u):
+        """
+        Test that the acoustic wave equation runs without crashing in absence
+        of sources and receivers.
+        """
+        u.data[:] = 0.0
+        eqn = [Eq(u.forward, stencil)]
+        op = Operator(eqn, subs=subs)
+        op.apply(u=u, m=m, damp=damp, t=10)
 
-        # Create solver object to provide relevant operators
-        solver = AcousticWaveSolver(model, source=src, receiver=rec,
-                                    time_order=time_order,
-                                    space_order=space_order)
-
-        # Create adjoint receiver symbol
-        srca = Receiver(name='srca', ntime=solver.source.nt,
-                        coordinates=solver.source.coordinates.data)
-
-        # Run forward and adjoint operators
-        rec, _, _ = solver.forward(save=False)
-        solver.adjoint(rec=rec, srca=srca)
-
-        # Adjoint test: Verify <Ax,y> matches  <x, A^Ty> closely
-        term1 = np.dot(srca.data.reshape(-1), solver.source.data)
-        term2 = np.linalg.norm(rec.data) ** 2
-        info('<Ax,y>: %f, <x, A^Ty>: %f, difference: %12.12f, ratio: %f'
-             % (term1, term2, term1 - term2, term1 / term2))
-        assert np.isclose(term1, term2, rtol=1.e-5)
+    def test_acoustic_w_src_wo_rec(self, model, stencil, subs, m, damp, u, src):
+        """
+        Test that the acoustic wave equation runs without crashing in absence
+        of receivers.
+        """
+        dt = model.critical_dt
+        u.data[:] = 0.0
+        eqns = [Eq(u.forward, stencil)]
+        eqns += src.inject(field=u.forward, expr=src * dt**2 / m, offset=model.nbpml)
+        op = Operator(eqns, subs=subs)
+        op.apply(u=u, m=m, damp=damp, src=src, t=1)
+
+    def test_acoustic_w_src_w_rec(self, model, stencil, subs, m, damp, u, src, rec):
+        """
+        Test that the acoustic wave equation forward operator produces the correct
+        results when running a 3D model also used in ``test_adjointA.py``.
+        """
+        dt = model.critical_dt
+        u.data[:] = 0.0
+        eqns = [Eq(u.forward, stencil)]
+        eqns += src.inject(field=u.forward, expr=src * dt**2 / m, offset=model.nbpml)
+        eqns += rec.interpolate(expr=u, offset=model.nbpml)
+        op = Operator(eqns, subs=subs)
+        op.apply(u=u, m=m, damp=damp, src=src, rec=rec, t=1)
+
+        # TODO: the following "hacky" way of asserting correctness will be replaced
+        # once adjoint operators could be run through YASK. At the moment, the following
+        # expected norms have been "manually" derived from an analogous test (same
+        # equation, same model, ...) in test_adjointA.py
+        exp_u = 152.76
+        exp_rec = 212.00
+        assert np.isclose(np.linalg.norm(u.data[:]), exp_u, atol=exp_u*1.e-2)
+        assert np.isclose(np.linalg.norm(rec.data), exp_rec, atol=exp_rec*1.e-2)