ci: Update docker used

.github/workflows/ci-lit.yml

@@ -16,7 +16,7 @@ on:
 jobs:
   build:
     runs-on: ubuntu-latest
-    container: papychacal/xdsl-llvm:98e674c9f16d677d95c67bc130e267fae331e43c
+    container: gbisbas/llvm-xdsl:0.4
     steps:
     - name: Checkout Devito
       uses: actions/checkout@v4
@@ -25,11 +25,6 @@ jobs:
       run: |
         apt-get update && apt install curl mpich -y
 
-    - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: '3.10'
-
     - name: Upgrade pip
       run: |
         pip install --upgrade pip

tests/test_xdsl_base.py

@@ -1119,3 +1119,115 @@ 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
+
+        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_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)

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