Merge branch 'emilien/coming-soon' into emilien/inplace

xdslproject · Aug 30, 2024 · 17f0291 · 17f0291
2 parents 74746f5 + 67a9772
commit 17f0291
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 27 deletions.
diff --git a/devito/ir/xdsl_iet/cluster_to_ssa.py b/devito/ir/xdsl_iet/cluster_to_ssa.py
@@ -380,7 +380,20 @@ def build_stencil_step(self, dim: SteppingDimension, eq: LoweredEq) -> None:
         self.apply_temps = {k: v for k, v in zip(read_functions, apply.region.block.args)}
 
         with ImplicitBuilder(apply.region.block):
-            stencil.ReturnOp.get([self._visit_math_nodes(dim, eq.rhs, eq.lhs)])
+            result = self._visit_math_nodes(dim, eq.rhs, eq.lhs)
+            expected_type = apply.res[0].type.get_element_type()
+            match expected_type:
+                case result.type:
+                    pass
+                case builtin.f32:
+                    if result.type == IndexType():
+                        result = arith.IndexCastOp(result, builtin.i64).result
+                    result = arith.SIToFPOp(result, builtin.f32).result
+                case builtin.IndexType:
+                    result = arith.IndexCastOp(result, IndexType()).result
+                case _:
+                    raise Exception(f"Unexpected result type {type(result)}")
+            stencil.ReturnOp.get([result])
 
         lb = stencil.IndexAttr.get(*([0] * len(shape)))
         ub = stencil.IndexAttr.get(*shape)

diff --git a/devito/xdsl_core/xdsl_gpu.py b/devito/xdsl_core/xdsl_gpu.py
@@ -140,7 +140,10 @@ def _jit_compile(self):
 
 def generate_XDSL_GPU_PIPELINE():
     passes = [
+        "canonicalize",
+        "cse",
         "shape-inference",
+        "stencil-bufferize",
         "convert-stencil-to-ll-mlir",
         "reconcile-unrealized-casts",
         "printf-to-llvm",

diff --git a/tests/test_xdsl_base.py b/tests/test_xdsl_base.py
@@ -972,6 +972,20 @@ def test_function_IV():
     assert np.isclose(norm(u), devito_norm_u)
 
 
+def test_function_V():
+    grid = Grid(shape=(5, 5))
+    x, y = grid.dimensions
+
+    f = Function(name="f", grid=grid)
+
+    eqns = [Eq(f, 2)]
+
+    op = Operator(eqns, opt="xdsl")
+    op.apply()
+
+    assert np.all(f.data == 2)
+
+
 class TestTrigonometric(object):
 
     @pytest.mark.parametrize('deg, exp', ([90.0, 3.5759869], [30.0, 3.9521265],
@@ -1028,37 +1042,20 @@ def test_tan(self, deg, exp):
         assert np.isclose(norm(u), exp, rtol=1e-4)
 
 
-class TestOperatorUnsupported(object):
+def test_forward_assignment():
+    # simple forward assignment
 
-    @pytest.mark.xfail(reason="stencil.return operation does not verify for i64")
-    def test_forward_assignment(self):
-        # simple forward assignment
-
-        grid = Grid(shape=(4, 4))
-        u = TimeFunction(name="u", grid=grid, space_order=2)
-        u.data[:, :, :] = 0
-
-        eq0 = Eq(u.forward, 1)
-
-        op = Operator([eq0], opt='xdsl')
-
-        op.apply(time_M=1)
-
-        assert np.isclose(norm(u), 5.6584, rtol=0.001)
-
-    @pytest.mark.xfail(reason="stencil.return operation does not verify for i64")
-    def test_function(self):
-        grid = Grid(shape=(5, 5))
-        x, y = grid.dimensions
+    grid = Grid(shape=(4, 4))
+    u = TimeFunction(name="u", grid=grid, space_order=2)
+    u.data[:, :, :] = 0
 
-        f = Function(name="f", grid=grid)
+    eq0 = Eq(u.forward, 1)
 
-        eqns = [Eq(f, 2)]
+    op = Operator([eq0], opt='xdsl')
 
-        op = Operator(eqns, opt='xdsl')
-        op.apply()
+    op.apply(time_M=1)
 
-        assert np.all(f.data == 4)
+    assert np.isclose(norm(u), 5.6584, rtol=0.001)
 
 
 class TestElastic():