compiler: Init work for TTI (add sin, cos)

devito/ir/xdsl_iet/cluster_to_ssa.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass, field
 from sympy import (Add, And, Expr, Float, GreaterThan, Indexed, Integer, LessThan,
                    Number, Pow, StrictGreaterThan, StrictLessThan, Symbol, floor,
-                   Mul)
+                   Mul, sin, cos)
 from sympy.core.relational import Relational
 from sympy.logic.boolalg import BooleanFunction
 from devito.operations.interpolators import Injection
@@ -52,7 +52,6 @@
 
 
 def field_from_function(f: DiscreteFunction) -> stencil.FieldType:
-    # import pdb;pdb.set_trace()
     halo = [f.halo[d] for d in f.dimensions]
     shape = f.shape
     bounds = [(-h[0], s+h[1]) for h, s in zip(halo, shape)]
@@ -288,6 +287,19 @@ def _visit_math_nodes(self, dim: SteppingDimension, node: Expr,
             SSAargs = (self._visit_math_nodes(dim, arg, output_indexed)
                        for arg in node.args)
             return reduce(lambda x, y : arith.AndI(x, y).result, SSAargs)
+
+        elif isinstance(node, sin):
+            assert len(node.args) == 1, "Expected single argument for sin."
+            return math.SinOp(self._visit_math_nodes(dim, node.args[0],
+                              output_indexed)).result
+
+        elif isinstance(node, cos):
+            assert len(node.args) == 1, "Expected single argument for cos."
+
+            return math.CosOp(self._visit_math_nodes(dim, node.args[0],
+                              output_indexed)).result
+
+
         elif isinstance(node, Relational):
             if isinstance(node, GreaterThan):
                 mnemonic = "sge"
@@ -499,6 +511,7 @@ def build_time_loop(
     def lower_devito_Eqs(self, eqs: list[Any], **kwargs):
         # Lower devito Equations to xDSL
 
+
         for eq in eqs:
             lowered = self.operator._lower_exprs(as_tuple(eq), **kwargs)
             if isinstance(eq, Eq):
@@ -631,7 +644,7 @@ def convert(self, eqs: Iterable[Eq], **kwargs) -> ModuleOp:
                         functions.add(f.function)
 
                 elif isinstance(eq, Injection):
-                    # import pdb; pdb.set_trace()
+
                     functions.add(eq.field.function)
                     for f in retrieve_functions(eq.expr):
                         if isinstance(f, PointSource):

tests/test_xdsl_base.py

@@ -2,7 +2,7 @@
 import pytest
 
 from devito import (Grid, TensorTimeFunction, VectorTimeFunction, div, grad, diag, solve,
-                    Operator, Eq, Constant, norm, SpaceDimension, switchconfig)
+                    Operator, Eq, Constant, norm, SpaceDimension, switchconfig, sin, cos)
 from devito.types import Array, Function, TimeFunction
 from devito.tools import as_tuple
 
@@ -970,6 +970,37 @@ def test_function_IV():
     assert np.isclose(norm(u), devito_norm_u)
 
 
+class TestTrigonometric(object):
+
+    @pytest.mark.parametrize('deg, exp', ([90.0, 3.5759869], [30.0, 3.9521265],
+                             [45.0, 3.403614]))
+    def test_sine(self, deg, exp):
+        grid = Grid(shape=(4, 4))
+
+        u = Function(name="u", grid=grid)
+        u.data[:, :] = 0
+
+        eq0 = Eq(u, sin(deg))
+
+        op = Operator([eq0], opt='xdsl')
+        op.apply()
+        assert np.isclose(norm(u), exp, rtol=1e-4)
+
+    @pytest.mark.parametrize('deg, exp', ([90.0, 1.7922944], [30.0, 0.6170056],
+                             [45.0, 2.101288]))
+    def test_cosine(self, deg, exp):
+        grid = Grid(shape=(4, 4))
+
+        u = Function(name="u", grid=grid)
+        u.data[:, :] = 0
+
+        eq0 = Eq(u, cos(deg))
+
+        op = Operator([eq0], opt='xdsl')
+        op.apply()
+        assert np.isclose(norm(u), exp, rtol=1e-4)
+
+
 class TestOperatorUnsupported(object):
 
     @pytest.mark.xfail(reason="stencil.return operation does not verify for i64")

tests/test_xdsl_op_correctness.py

@@ -1,17 +1,18 @@
 import numpy as np
-from devito import Grid, TimeFunction, Eq, Operator, norm, switchconfig, Function, NODE, div, grad
+from devito import (Grid, TimeFunction, Eq, Operator, norm, switchconfig,
+                    Function, NODE, div, grad, solve)
 import pytest
-# flake8: noqa
 
 from xdsl.dialects.scf import For, Yield
 from xdsl.dialects.arith import Addi, Constant
 from xdsl.dialects.func import Call, Return
 from xdsl.dialects.stencil import FieldType, ApplyOp, LoadOp, StoreOp
-from xdsl.dialects.llvm import LLVMPointerType
-from xdsl.printer import Printer
 
 from devito.types.basic import Symbol
 
+from examples.seismic import demo_model, TimeAxis, RickerSource
+
+
 def test_udx():
 
     # Define a simple Devito Operator
@@ -28,10 +29,11 @@ def test_udx():
     xdsl_op = Operator([eq], opt='xdsl')
     xdsl_op.apply(time_M=5)
     norm2 = norm(u)
-    
-    assert np.isclose(norm1, norm2,   atol=1e-5, rtol=0)
+
+    assert np.isclose(norm1, norm2, atol=1e-5, rtol=0)
     assert np.isclose(norm1, 14636.3955, atol=1e-5, rtol=0)
 
+
 def test_u_plus1_conversion():
     # Define a simple Devito Operator
     grid = Grid(shape=(3, 3))
@@ -41,7 +43,7 @@ def test_u_plus1_conversion():
     op = Operator([eq])
     op.apply(time_M=5)
     norm1 = norm(u)
-    
+
     u.data[:] = 0
     xdsl_op = Operator([eq], opt='xdsl')
     xdsl_op.apply(time_M=5)
@@ -88,7 +90,7 @@ def test_u_and_v_conversion():
     assert type(ops[6] == For)
 
     scffor_ops = list(ops[6].regions[0].blocks[0].ops)
-    
+
     assert len(scffor_ops) == 7
 
     # First
@@ -108,6 +110,7 @@ def test_u_and_v_conversion():
     assert type(ops[8] == StoreOp)
     assert type(ops[9] == Return)
 
+
 def test_symbol_I():
     # Define a simple Devito a = 1 operator
 
@@ -169,7 +172,7 @@ def test_extraction_from_lifted_ispace(rotate):
     # Operator
     op0 = Operator(eqns, opt='xdsl')
     op1 = Operator(eqns, opt=('advanced', {'openmp': True, 'cire-mingain': 1,
-                                            'cire-rotate': rotate}))
+                              'cire-rotate': rotate}))
 
     # Check numerical output
     op0(time_M=1)
@@ -179,4 +182,52 @@ def test_extraction_from_lifted_ispace(rotate):
 
     # Also check against expected operation count to make sure
     # all redundancies have been detected correctly
-    assert summary[('section0', None)].ops == 93
+    assert summary[('section0', None)].ops == 93
+
+
+@pytest.mark.parametrize('shape', [(101, 101), (201, 201)])
+@pytest.mark.parametrize('nt', [10, 20, ])
+def test_dt2_sources_script(shape, nt):
+
+    spacing = (10., 10.)  # spacing of 10 meters
+    nbl = 1  # number of pad layers
+
+    model = demo_model('layers-tti', spacing=spacing, space_order=4,
+                       shape=shape, nbl=nbl, nlayers=nbl)
+
+    # Compute the dt and set time range
+    t0 = 0.  # Simulation time start
+    dt = model.critical_dt
+
+    time_range = TimeAxis(start=t0, stop=nt, step=dt)
+
+    p = TimeFunction(name="p", grid=model.grid, time_order=2, space_order=8)
+
+    # Main equations
+    stencil_p = solve(p.dt2, p.forward)
+    update_p = Eq(p.forward, stencil_p)
+
+    # Create stencil and boundary condition expressions
+    x, z = model.grid.dimensions
+
+    # set source and receivers
+    src = RickerSource(name='src', grid=model.grid, f0=0.02, npoint=1,
+                       time_range=time_range)
+
+    src.coordinates.data[:, 0] = model.domain_size[0] * .5
+    src.coordinates.data[:, 1] = model.domain_size[0] * .5
+    # Define the source injection
+
+    src_term = src.inject(field=p.forward, expr=src)
+
+    optime = Operator([update_p] + src_term)
+    optime.apply(time=time_range.num-1, dt=dt)
+    norm0 = norm(p)
+
+    p.data[:] = 0.
+
+    optime = Operator([update_p] + src_term, opt='xdsl')
+    optime.apply(time=time_range.num-1, dt=dt)
+    norm1 = norm(p)
+
+    assert np.isclose(norm0, norm1, atol=1e-5, rtol=0)