Merge pull request #104 from xdslproject/elastic_deterministic

compiler: Fix non-deterministic codegen for stencil.Loads and add ela…

.github/workflows/ci-lit.yml

@@ -21,22 +21,15 @@ jobs:
     - name: Checkout Devito
       uses: actions/checkout@v4
 
-    - name: Install native dependencies
-      run: |
-        apt-get update && apt install curl mpich -y
-
     - name: Upgrade pip
       run: |
         pip install --upgrade pip
 
     - name: Install requirements and xDSL
       run: |
-        pip install -e .[tests]
-        pip install mpi4py
         pip install git+https://github.com/xdslproject/xdsl@f8bb935880276cf077e0a80f1905105d0a98eb33
+        pip install -e .[tests]
 
     - name: Execute lit tests
       run: |
-        export PYTHONPATH=$(pwd)
-        export PATH=/xdsl-sc/llvm-project/build/bin/:$PATH
         lit -v tests/filecheck/

.github/workflows/ci-mlir.yml

@@ -35,7 +35,6 @@ jobs:
     - name: Install requirements and xDSL
       run: |
         pip install -e .[tests]
-        pip install mpi4py
         pip install git+https://github.com/xdslproject/xdsl@f8bb935880276cf077e0a80f1905105d0a98eb33
         
     - name: Test no-MPI, no-Openmp

devito/ir/xdsl_iet/cluster_to_ssa.py

@@ -321,7 +321,7 @@ def build_stencil_step(self, dim: SteppingDimension, eq: LoweredEq) -> None:
         Returns:
             None
         """
-        read_functions = set()
+        read_functions = OrderedSet()
         for f in retrieve_function_carriers(eq.rhs):
             if isinstance(f.function, PointSource):
                 time_offset = 0

requirements-testing.txt

@@ -7,4 +7,4 @@ nbval
 scipy
 lit<19.0.0
 filecheck<0.0.25
-pooch; python_version >= "3.8"
+pooch; python_version >= "3.8"

tests/filecheck/version.mlir

@@ -0,0 +1,3 @@
+// RUN: xdsl-opt --version | filecheck %s
+
+//CHECK: xdsl-opt built from xdsl version {{.*}}

tests/test_xdsl_base.py

@@ -1119,3 +1119,124 @@ def test_elastic_2D(self, shape, so, nt):
         assert not np.isclose(xdsl_norm_tau1, 0.0, rtol=1e-04)
         assert not np.isclose(xdsl_norm_tau2, 0.0, rtol=1e-04)
         assert not np.isclose(xdsl_norm_tau3, 0.0, rtol=1e-04)
+
+    @pytest.mark.parametrize('shape', [(101, 101, 101), (201, 201, 201)])
+    @pytest.mark.parametrize('so', [2, 4, 8])
+    @pytest.mark.parametrize('nt', [10, 20, ])
+    def test_elastic_3D(self, shape, so, nt):
+
+        # Initial grid: km x km, with spacing
+        shape = shape  # Number of grid point (nx, nz)
+        spacing = as_tuple(10.0 for _ in range(len(shape)))
+        extent = tuple([s*(n-1) for s, n in zip(spacing, shape)])
+
+        x = SpaceDimension(name='x', spacing=Constant(name='h_x', value=extent[0]/(shape[0]-1)))  # noqa
+        y = SpaceDimension(name='y', spacing=Constant(name='h_y', value=extent[1]/(shape[1]-1)))  # noqa
+        z = SpaceDimension(name='z', spacing=Constant(name='h_z', value=extent[2]/(shape[2]-1)))  # noqa
+        grid = Grid(extent=extent, shape=shape, dimensions=(x, y, z))
+
+        # Timestep size from Eq. 7 with V_p=6000. and dx=100
+        t0, tn = 0., nt
+        dt = (10. / np.sqrt(2.)) / 6.
+        time_range = TimeAxis(start=t0, stop=tn, step=dt)
+
+        src = RickerSource(name='src', grid=grid, f0=0.01, time_range=time_range)
+        src.coordinates.data[:] = [250., 250., 250.]
+
+        # Now we create the velocity and pressure fields
+        v = VectorTimeFunction(name='v', grid=grid, space_order=so, time_order=1)
+        tau = TensorTimeFunction(name='t', grid=grid, space_order=so, time_order=1)
+
+        # We need some initial conditions
+        V_p = 2.0
+        V_s = 1.0
+        density = 1.8
+
+        # The source injection term
+        src_xx = src.inject(field=tau.forward[0, 0], expr=src)
+        src_yy = src.inject(field=tau.forward[1, 1], expr=src)
+        src_zz = src.inject(field=tau.forward[2, 2], expr=src)
+
+        # Thorbecke's parameter notation
+        cp2 = V_p*V_p
+        cs2 = V_s*V_s
+        ro = 1/density
+
+        mu = cs2*density
+        l = (cp2*density - 2*mu)
+
+        # First order elastic wave equation
+        pde_v = v.dt - ro * div(tau)
+        pde_tau = (tau.dt - l * diag(div(v.forward)) - mu * (grad(v.forward) +
+                   grad(v.forward).transpose(inner=False)))
+
+        # Time update
+        u_v = Eq(v.forward, solve(pde_v, v.forward))
+        u_t = Eq(tau.forward, solve(pde_tau, tau.forward))
+
+        # Inject sources. We use it to preinject data
+        # Up to here, let's only use Devito
+        op = Operator([u_v] + [u_t] + src_xx + src_yy + src_zz)
+        op(dt=dt)
+
+        opx = Operator([u_v] + [u_t], opt='xdsl')
+        opx(dt=dt, time_M=nt)
+
+        store_ops = [op for op in opx._module.walk() if isinstance(op, StoreOp)]
+        assert len(store_ops) == 9
+
+        xdsl_norm_v0 = norm(v[0])
+        xdsl_norm_v1 = norm(v[1])
+        xdsl_norm_tau0 = norm(tau[0])
+        xdsl_norm_tau1 = norm(tau[1])
+        xdsl_norm_tau2 = norm(tau[2])
+        xdsl_norm_tau3 = norm(tau[3])
+
+        # Reinitialize the fields to zero
+
+        # TODO: Reinitialize the fields to zero
+        v[0].data[:] = 0
+        v[1].data[:] = 0
+        v[2].data[:] = 0
+
+        # t_xx, t_xy, t_xz
+        # t_yx, t_yy, t_yz
+        # t_zx, t_zy, t_zz
+
+        tau[0].data[:] = 0
+        tau[1].data[:] = 0
+        tau[2].data[:] = 0
+        tau[3].data[:] = 0
+        tau[4].data[:] = 0
+        tau[5].data[:] = 0
+        tau[6].data[:] = 0
+        tau[7].data[:] = 0
+        tau[8].data[:] = 0
+
+        # Inject sources. We use it to preinject data
+        op = Operator([u_v] + [u_t] + src_xx + src_yy + src_zz)
+        op(dt=dt)
+
+        op = Operator([u_v] + [u_t], opt='advanced')
+        op(dt=dt, time_M=nt)
+
+        dv_norm_v0 = norm(v[0])
+        dv_norm_v1 = norm(v[1])
+        dv_norm_tau0 = norm(tau[0])
+        dv_norm_tau1 = norm(tau[1])
+        dv_norm_tau2 = norm(tau[2])
+        dv_norm_tau3 = norm(tau[3])
+
+        assert np.isclose(xdsl_norm_v0, dv_norm_v0, rtol=1e-04)
+        assert np.isclose(xdsl_norm_v1, dv_norm_v1, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau0, dv_norm_tau0, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau1, dv_norm_tau1, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau2, dv_norm_tau2, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau3, dv_norm_tau3, rtol=1e-04)
+
+        assert not np.isclose(xdsl_norm_v0, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_v1, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau0, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau1, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau2, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau3, 0.0, rtol=1e-04)

tests/test_xdsl_mpi.py

@@ -7,7 +7,12 @@
 import pytest
 import numpy as np
 
-from devito import (Grid, TimeFunction, Eq, norm, solve, Operator)
+from devito import (Grid, TimeFunction, Eq, norm, solve, Operator,
+                    TensorTimeFunction, VectorTimeFunction, div, grad, diag)
+from devito.types import SpaceDimension, Constant
+from devito.tools import as_tuple
+from examples.seismic.source import RickerSource, TimeAxis
+from xdsl.dialects.stencil import StoreOp
 
 
 pytestmark = skipif(['nompi'], whole_module=True)
@@ -135,3 +140,123 @@ def test_topology_8(self, shape, to, nt):
         assert topology == (2, 2, 2)
 
         xdslop.apply(time=nt, dt=dt)
+
+
+class TestElastic():
+
+    @pytest.mark.xfail(reason=" With MPI, data can only be set via scalars, numpy arrays or other data")  # noqa
+    @pytest.mark.parallel(mode=[1])
+    @pytest.mark.parametrize('shape', [(101, 101), (201, 201), (301, 301)])
+    @pytest.mark.parametrize('so', [2, 4, 8])
+    @pytest.mark.parametrize('nt', [10, 20, 50, 100])
+    def test_elastic_2D(self, shape, so, nt):
+
+        # Initial grid: km x km, with spacing
+        shape = shape  # Number of grid point (nx, nz)
+        spacing = as_tuple(10.0 for _ in range(len(shape)))
+        extent = tuple([s*(n-1) for s, n in zip(spacing, shape)])
+
+        x = SpaceDimension(name='x', spacing=Constant(name='h_x', value=extent[0]/(shape[0]-1)))  # noqa
+        z = SpaceDimension(name='z', spacing=Constant(name='h_z', value=extent[1]/(shape[1]-1)))  # noqa
+        grid = Grid(extent=extent, shape=shape, dimensions=(x, z))
+
+        # To be checked again in the future
+        # class DGaussSource(WaveletSource):
+
+        #    def wavelet(self, f0, t):
+        #        a = 0.004
+        #        return -2.*a*(t - 1/f0) * np.exp(-a * (t - 1/f0)**2)
+
+        # Timestep size from Eq. 7 with V_p=6000. and dx=100
+        t0, tn = 0., nt
+        dt = (10. / np.sqrt(2.)) / 6.
+        time_range = TimeAxis(start=t0, stop=tn, step=dt)
+
+        src = RickerSource(name='src', grid=grid, f0=0.01, time_range=time_range)
+        src.coordinates.data[:] = [250., 250.]
+
+        # Now we create the velocity and pressure fields
+        v = VectorTimeFunction(name='v', grid=grid, space_order=so, time_order=1)
+        tau = TensorTimeFunction(name='t', grid=grid, space_order=so, time_order=1)
+
+        # We need some initial conditions
+        V_p = 2.0
+        V_s = 1.0
+        density = 1.8
+
+        # The source injection term
+        src_xx = src.inject(field=tau.forward[0, 0], expr=src)
+        src_zz = src.inject(field=tau.forward[1, 1], expr=src)
+
+        # Thorbecke's parameter notation
+        cp2 = V_p*V_p
+        cs2 = V_s*V_s
+        ro = 1/density
+
+        mu = cs2*density
+        l = (cp2*density - 2*mu)
+
+        # First order elastic wave equation
+        pde_v = v.dt - ro * div(tau)
+        pde_tau = (tau.dt - l * diag(div(v.forward)) - mu * (grad(v.forward) +
+                   grad(v.forward).transpose(inner=False)))
+
+        # Time update
+        u_v = Eq(v.forward, solve(pde_v, v.forward))
+        u_t = Eq(tau.forward, solve(pde_tau, tau.forward))
+
+        # Inject sources. We use it to preinject data
+        # Up to here, let's only use Devito
+        op = Operator([u_v] + [u_t] + src_xx + src_zz)
+        op(dt=dt)
+
+        opx = Operator([u_v] + [u_t], opt='xdsl')
+        opx(dt=dt, time_M=nt)
+
+        store_ops = [op for op in opx._module.walk() if isinstance(op, StoreOp)]
+        assert len(store_ops) == 5
+
+        xdsl_norm_v0 = norm(v[0])
+        xdsl_norm_v1 = norm(v[1])
+        xdsl_norm_tau0 = norm(tau[0])
+        xdsl_norm_tau1 = norm(tau[1])
+        xdsl_norm_tau2 = norm(tau[2])
+        xdsl_norm_tau3 = norm(tau[3])
+
+        # Reinitialize the fields to zero
+
+        v[0].data[:] = 0
+        v[1].data[:] = 0
+
+        tau[0].data[:] = 0
+        tau[1].data[:] = 0
+        tau[2].data[:] = 0
+        tau[3].data[:] = 0
+
+        # Inject sources. We use it to preinject data
+        op = Operator([u_v] + [u_t] + src_xx + src_zz)
+        op(dt=dt)
+
+        op = Operator([u_v] + [u_t], opt='advanced')
+        op(dt=dt, time_M=nt)
+
+        dv_norm_v0 = norm(v[0])
+        dv_norm_v1 = norm(v[1])
+        dv_norm_tau0 = norm(tau[0])
+        dv_norm_tau1 = norm(tau[1])
+        dv_norm_tau2 = norm(tau[2])
+        dv_norm_tau3 = norm(tau[3])
+
+        assert np.isclose(xdsl_norm_v0, dv_norm_v0, rtol=1e-04)
+        assert np.isclose(xdsl_norm_v1, dv_norm_v1, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau0, dv_norm_tau0, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau1, dv_norm_tau1, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau2, dv_norm_tau2, rtol=1e-04)
+        assert np.isclose(xdsl_norm_tau3, dv_norm_tau3, rtol=1e-04)
+
+        assert not np.isclose(xdsl_norm_v0, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_v1, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau0, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau1, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau2, 0.0, rtol=1e-04)
+        assert not np.isclose(xdsl_norm_tau3, 0.0, rtol=1e-04)

xdsl_examples/diffusion_2D_wBCs.py

@@ -25,8 +25,8 @@
                     help="Block levels")
 parser.add_argument("-devito", "--devito", default=False, type=bool, help="Devito run")
 parser.add_argument("-xdsl", "--xdsl", default=False, type=bool, help="xDSL run")
-args = parser.parse_args()
 parser.add_argument("-plot", "--plot", default=False, type=bool, help="Plot2D")
+args = parser.parse_args()
 
 
 mpiconf = configuration['mpi']

xdsl_llvm.docker

@@ -4,7 +4,8 @@ ARG mlirhash=98e674c9f16d677d95c67bc130e267fae331e43c
 
 # base requirements
 RUN apt-get update \
- && apt-get install --yes --no-install-recommends git cmake build-essential ca-certificates ninja-build clang lld python3.10-dev pip \
+ && apt-get install --yes --no-install-recommends git cmake build-essential \
+ && apt-get install --yes ca-certificates ninja-build clang lld python3 python3-dev python3-pip python3-venv pip \
  && apt-get install --yes libstdc++-12-dev binutils-dev
 
 WORKDIR /xdsl-sc
@@ -31,3 +32,10 @@ RUN git init llvm-project \
  && ninja install \
  && cd ../../ \
  && rm -rf llvm-project
+
+
+
+# Build dockerfile
+# Example name: <gbisbas/llvm-xdsl:0.4>
+# docker build --network=host --file xdsl_llvm.docker --tag <name> .
+# docker push <name>