From ebfb1eed6beb2952189af126aa58b8bfd835eda7 Mon Sep 17 00:00:00 2001
From: George Bisbas <g.bisbas18@imperial.ac.uk>
Date: Wed, 15 Nov 2023 12:00:27 +0000
Subject: [PATCH] ci-mlir: add Ninja

---
 .github/workflows/ci-mlir.yml |  6 ++-
 tests/test_xdsl_base.py       | 96 ++++++++++++++++++++++++++++++++++-
 2 files changed, 99 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/ci-mlir.yml b/.github/workflows/ci-mlir.yml
index 9b99564b77..c829921b26 100644
--- a/.github/workflows/ci-mlir.yml
+++ b/.github/workflows/ci-mlir.yml
@@ -3,6 +3,10 @@
 
 name: CI - MLIR-based Testing
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 on:
   # Trigger the workflow on push or pull request,
   # but only for the master branch
@@ -91,7 +95,7 @@ jobs:
       if: steps.cache-binary.outputs.cache-hit != 'true'
       run: |
         cd llvm-project/build
-        cmake --build . --target mlir-opt mlir-cpu-runner
+        ninja
 
     - name: Test with pytest and generate code coverage
       run: |
diff --git a/tests/test_xdsl_base.py b/tests/test_xdsl_base.py
index e246467810..a95d8e834e 100644
--- a/tests/test_xdsl_base.py
+++ b/tests/test_xdsl_base.py
@@ -1,4 +1,7 @@
-from devito import Grid, TimeFunction, Eq, XDSLOperator
+import numpy as np
+import pytest
+
+from devito import Grid, TimeFunction, Eq, XDSLOperator, Operator, solve
 
 
 def test_xdsl_I():
@@ -34,4 +37,93 @@ def test_xdsl_III():
     op = XDSLOperator([eq], opt=None)
     op.apply(time_M=1)
     assert (u.data[1, :] == 1.).all()
-    assert (u.data[0, :] == 2.).all()
\ No newline at end of file
+    assert (u.data[0, :] == 2.).all()
+
+
+def test_diffusion_2D():
+    # Define a simple Devito Operator
+    grid = Grid(shape=(3, 3))
+
+    # Devito setup
+    f = TimeFunction(name='f', grid=grid, space_order=2)
+    f.data[:] = 1
+    eqn = Eq(f.dt, 0.5 * f.laplace)
+    op = Operator(Eq(f.forward, solve(eqn, f.forward)))
+    op.apply(time_M=1, dt=0.1)
+
+    # xDSL-Devito setup
+    f2 = TimeFunction(name='f2', grid=grid, space_order=2)
+    f2.data[:] = 1
+    eqn = Eq(f2.dt, 0.5 * f2.laplace)
+    op = XDSLOperator(Eq(f2.forward, solve(eqn, f2.forward)))
+    op.apply(time_M=1, dt=0.1)
+   
+    assert np.isclose(f.data, f2.data, rtol=1e-06).all()
+
+
+@pytest.mark.parametrize('shape', [(11, 11), (31, 31), (51, 51), (101, 101)])
+def test_diffusion_2D_II(shape):
+    # Define a simple Devito Operator
+    grid = Grid(shape=shape)
+    rng = np.random.default_rng(123)
+    dx = 2. / (shape[0] - 1)
+    dy = 2. / (shape[1] - 1)
+    sigma = .1
+    dt = sigma * dx * dy
+    nt = 10
+    
+    arr1 = rng.random(shape)
+
+    # Devito setup
+    f = TimeFunction(name='f', grid=grid, space_order=2)
+    f.data[:] = arr1
+    eqn = Eq(f.dt, 0.5 * f.laplace)
+    op = Operator(Eq(f.forward, solve(eqn, f.forward)))
+    op.apply(time_M=nt, dt=dt)
+
+    # xDSL-Devito setup
+    f2 = TimeFunction(name='f2', grid=grid, space_order=2)
+    f2.data[:] = arr1
+    eqn = Eq(f2.dt, 0.5 * f2.laplace)
+    op = XDSLOperator(Eq(f2.forward, solve(eqn, f2.forward)))
+    op.apply(time_M=nt, dt=dt)
+ 
+    max_error = np.max(np.abs(f.data - f2.data))
+    assert np.isclose(max_error, 0.0, atol=1e-04)
+    assert np.isclose(f.data, f2.data, rtol=1e-05).all()
+
+
+@pytest.mark.parametrize('shape', [(11, 11, 11), (31, 31, 31), (51, 51, 51), (101, 101, 101)])
+def test_diffusion_3D_II(shape):
+    shape = (10, 10, 10)
+    grid = Grid(shape=shape)
+    rng = np.random.default_rng(123)
+    dx = 2. / (shape[0] - 1)
+    dy = 2. / (shape[1] - 1)
+    dz = 2. / (shape[2] - 1)
+
+    sigma = .1
+    dt = sigma * dx * dy * dz
+    nt = 50
+
+    rng = np.random.default_rng(123)
+    
+    arr1 = rng.random(shape)
+
+    # Devito setup
+    f = TimeFunction(name='f', grid=grid, space_order=2)
+    f.data[:] = arr1
+    eqn = Eq(f.dt, 0.5 * f.laplace)
+    op = Operator(Eq(f.forward, solve(eqn, f.forward)))
+    op.apply(time_M=nt, dt=dt)
+
+    # xDSL-Devito setup
+    f2 = TimeFunction(name='f2', grid=grid, space_order=2)
+    f2.data[:] = arr1
+    eqn = Eq(f2.dt, 0.5 * f2.laplace)
+    op = XDSLOperator(Eq(f2.forward, solve(eqn, f2.forward)))
+    op.apply(time_M=50, dt=dt)
+ 
+    max_error = np.max(np.abs(f.data - f2.data))
+    assert np.isclose(max_error, 0.0, atol=1e-04)
+    assert np.isclose(f.data, f2.data, rtol=1e-05).all()
\ No newline at end of file