compiler: separate the compilation pipelines and improve CI

.github/workflows/ci-mlir.yml

@@ -40,8 +40,18 @@ jobs:
         pip install mpi4py
         pip install git+https://github.com/xdslproject/xdsl@2825897b87443c9369abf89871f4721e2fce2da9
 
-    - name: Test non-MPI
+    - name: Test no-MPI, no-Openmp
       run: |
+        export DEVITO_MPI=0
+        export DEVITO_LANGUAGE=C
         # Add mlir-opt to the path
         export PATH=/xdsl-sc/llvm-project/build/bin/:$PATH
-        pytest -m "not parallel" -k "not adjoint" tests/test_xdsl_*
+        pytest -m "not parallel" -k "not adjoint" tests/test_xdsl_*
+
+    - name: Test no-MPI, Openmp
+      run: |
+          export DEVITO_MPI=0
+          export DEVITO_LANGUAGE=openmp
+          # Add mlir-opt to the path
+          export PATH=/xdsl-sc/llvm-project/build/bin/:$PATH
+          pytest -m "not parallel" -k "not adjoint" tests/test_xdsl_*

devito/operator/xdsl_operator.py

@@ -124,8 +124,9 @@ def mpi_shape(self) -> tuple:
     def _jit_compile(self):
         """
         JIT-compile the C code generated by the Operator.
-        It is ensured that JIT compilation will only be performed once per
-        Operator, reagardless of how many times this method is invoked.
+        It is ensured that JIT compilation will only be performed
+        once per Operator, reagardless of how many times this method
+        is invoked.
         """
 
         with self._profiler.timer_on('jit-compile'):
@@ -187,7 +188,8 @@ def _jit_compile(self):
             source_file = open(source_name, "w")
             source_file.write(module_str)
             source_file.close()
-            # compile IR using xdsl-opt | mlir-opt | mlir-translate | clang
+
+            # Compile IR using xdsl-opt | mlir-opt | mlir-translate | clang
             try:
                 cflags = CFLAGS
                 cc = "clang"
@@ -202,35 +204,63 @@ def _jit_compile(self):
 
                 # TODO More detailed error handling manually,
                 # instead of relying on a bash-only feature.
-                cmd = 'set -eo pipefail; '\
-                    f'xdsl-opt {source_name} -p {xdsl_pipeline} |' \
-                    f'mlir-opt -p {mlir_pipeline} | ' \
-                    f'mlir-translate --mlir-to-llvmir | ' \
-                    f'{cc} {cflags} -shared -o {self._tf.name} {self._interop_tf.name} -xir -'
-                print(cmd)
-                res = subprocess.run(
-                    cmd,
-                    shell=True,
-                    text=True,
-                    capture_output=True,
-                    executable="/bin/bash"
-                )
-
-                if res.returncode != 0:
-                    print("compilation failed with output:")
-                    print(res.stderr)
-
-                assert res.returncode == 0
+
+                xdsl_cmd = f'xdsl-opt {source_name} -p {xdsl_pipeline}'
+                mlir_cmd = f'mlir-opt -p {mlir_pipeline}'
+                mlir_translate_cmd = 'mlir-translate --mlir-to-llvmir'
+                clang_cmd = f'{cc} {cflags} -shared -o {self._tf.name} {self._interop_tf.name} -xir -'
+
+
+                comp_steps = [
+                              xdsl_cmd,
+                              mlir_cmd,
+                              mlir_translate_cmd,
+                              clang_cmd 
+                             ]
+
+                # Execute each command and store the outputs
+                outputs = []
+                stdout = None
+                for cmd in comp_steps:
+                    return_code, stdout, stderr = self._cmd_compile(cmd, stdout)
+                    # Use DEVITO_LOGGING=DEBUG to print
+                    debug(cmd)
+                    outputs.append({
+                        'command': cmd,
+                        'return_code': return_code,
+                        'stdout': stdout,
+                        'stderr': stderr
+                    })
+
             except Exception as ex:
                 print("error")
                 raise ex
-            # print(res.stderr)
 
         elapsed = self._profiler.py_timers['jit-compile']
 
         perf("XDSLOperator `%s` jit-compiled `%s` in %.2f s with `mlir-opt`" %
              (self.name, source_name, elapsed))
 
+
+    def _cmd_compile(self, cmd, input=None):
+        stdin = subprocess.PIPE if input is not None else None
+
+        res = subprocess.run(
+            cmd,
+            input=input,
+            shell=True,
+            text=True,
+            capture_output=True,
+            executable="/bin/bash"
+        )
+
+        if res.returncode != 0:
+            print("compilation failed with output:")
+            print(res.stderr)
+
+        assert res.returncode == 0
+        return res.returncode, res.stdout, res.stderr
+
     @property
     def _soname(self):
         return self._tf.name
@@ -258,12 +288,12 @@ def _check_kwargs(cls, **kwargs):
 
     @classmethod
     def _build(cls, expressions, **kwargs) -> Callable:
-        info("-Building operator")
+        debug("-Building operator")
         # Python- (i.e., compile-) and C-level (i.e., run-time) performance
         profiler = create_profile('timers')
 
         # Lower the input expressions into an IET
-        info("-Lower expressions")
+        debug("-Lower expressions")
         irs, _, module = cls._lower(expressions, profiler=profiler, **kwargs)
 
         # Make it an actual Operator

tests/test_xdsl_iet.py

@@ -63,6 +63,3 @@ def test_mfe2():
         a2 := arith.Addi(a1, cst1),
         memref.Store.get(a2, ref, [cst1, cst1])
     ])
-
-    printer = Printer()
-    printer.print_op(mod)

tests/test_xdsl_mpi.py

@@ -7,7 +7,7 @@
 import pytest
 import numpy as np
 
-from devito import (Grid, TimeFunction, Eq, norm, XDSLOperator)
+from devito import (Grid, TimeFunction, Eq, norm, XDSLOperator, solve, Operator)
 
 
 pytestmark = skipif(['nompi'], whole_module=True)
@@ -27,3 +27,38 @@ def test_trivial_eq_1d(self):
         op = XDSLOperator(Eq(f.forward, f[t, x-1, y] + f[t, x+1, y] + 1))
         op.apply(time=2)
         assert np.isclose(norm(f), 515.9845, rtol=1e-4)
+
+    @pytest.mark.parallel(mode=[2])
+    @pytest.mark.parametrize('shape', [(101, 101, 101), (202, 10, 45)])
+    @pytest.mark.parametrize('so', [2, 4, 8])
+    @pytest.mark.parametrize('to', [2])
+    @pytest.mark.parametrize('nt', [10, 20, 100])
+    def test_acoustic_3D(self, shape, so, to, nt):
+
+        grid = Grid(shape=shape)
+        dt = 0.0001
+
+        # Define the wavefield with the size of the model and the time dimension
+        u = TimeFunction(name="u", grid=grid, time_order=to, space_order=so)
+
+        pde = u.dt2 - u.laplace
+        eq0 = solve(pde, u.forward)
+
+        stencil = Eq(u.forward, eq0)
+        u.data[:, :, :] = 0
+        u.data[:, 40:50, 40:50] = 1
+
+        # Devito Operator
+        op = Operator([stencil])
+        op.apply(time=nt, dt=dt)
+        devito_norm = norm(u)
+
+        u.data[:, :, :] = 0
+        u.data[:, 40:50, 40:50] = 1
+
+        # XDSL Operator
+        xdslop = XDSLOperator([stencil])
+        xdslop.apply(time=nt, dt=dt)
+        xdsl_norm = norm(u)
+
+        assert np.isclose(devito_norm, xdsl_norm, rtol=1e-04).all()

tests/test_xdsl_op_correctness.py

@@ -3,8 +3,7 @@
 import pytest
 # flake8: noqa
 
-@pytest.mark.xfail(reason="Deprecated, will be dropped")
-def test_udx_conversion():
+def test_udx():
 
     # Define a simple Devito Operator
     grid = Grid(shape=(5, 5))
@@ -18,7 +17,6 @@ def test_udx_conversion():
     u.data[:] = 0.1
 
     xdsl_op = XDSLOperator([eq])
-    xdsl_op.__class__ = XDSLOperator
     xdsl_op.apply(time_M=5)
     norm2 = norm(u)
 
@@ -43,7 +41,7 @@ def test_u_plus1_conversion():
     assert np.isclose(norm1, norm2, atol=1e-5, rtol=0)
     assert np.isclose(norm1, 23.43075, atol=1e-5, rtol=0)
 
-@pytest.mark.xfail(reason="Deprecated, will be dropped")
+@pytest.mark.xfail(reason="Needs a fix in offsets")
 def test_u_and_v_conversion():
     # Define a simple Devito Operator
     grid = Grid(shape=(3, 3))
@@ -61,7 +59,6 @@ def test_u_and_v_conversion():
     u.data[:] = 0.0001
     v.data[:] = 0.0001
     xdsl_op = XDSLOperator([eq0, eq1])
-    xdsl_op.__class__ = XDSLOperator
     xdsl_op.apply(time_M=5, dt=0.1)
     norm_u2 = norm(u)
     norm_v2 = norm(v)