tests: Add to specialization tests

Author: EdCaunt

tests/test_pickle.py

@@ -4,6 +4,7 @@
 import cloudpickle as pickle1
 import numpy as np
 import pytest
+import sympy
 from sympy import Symbol
 
 from devito import (
@@ -14,6 +15,7 @@
 from devito.data import LEFT, OWNED
 from devito.finite_differences.tools import centered, direct, left, right, transpose
 from devito.ir import Backward, Forward, GuardBound, GuardBoundNext
+from devito.ir.iet.visitors import Specializer
 from devito.ir.support.guards import GuardFactorEq
 from devito.mpi.halo_scheme import Halo
 from devito.mpi.routines import (
@@ -756,6 +758,28 @@ def test_unjitted_operator(self, pickle):
 
         assert str(op) == str(new_op)
 
+    def test_specialized_operator(self, pickle):
+        grid = Grid(shape=(3, 3))
+        x_m, y_m = [d.symbolic_min for d in grid.dimensions]
+
+        f = Function(name='f', grid=grid)
+
+        op = Operator(Eq(f, f + 1))
+
+        pkl_op = pickle.dumps(op)
+        new_op = pickle.loads(pkl_op)
+
+        mapper = {x_m: sympy.S.Zero, y_m: sympy.S.Zero}
+
+        # Check pickled operators can be specialized
+        op0 = Specializer(mapper).visit(new_op)
+
+        # Check that specialized operators can be pickled
+        pkl_op0 = pickle.dumps(op0)
+        new_op0 = pickle.loads(pkl_op0)
+
+        assert str(op0) == str(new_op0)
+
     def test_operator_function(self, pickle):
         grid = Grid(shape=(3, 3, 3))
         f = Function(name='f', grid=grid)

tests/test_specialization.py

@@ -161,7 +161,20 @@ def test_sizes(self):
         assert "const int x_fsz0 = 11;" in str(op1.ccode)
         assert "const int y_fsz0 = 11;" in str(op1.ccode)
 
-    # TODO: Should strides get specialized? If so, how?
+    def test_blocking(self):
+        grid = Grid(shape=(11, 11, 11))
+        f = TimeFunction(name='f', grid=grid, space_order=4)
+
+        op0 = Operator(Eq(f.forward, f.laplace))
+
+        block_params = [p for p in op0.parameters if 'blk' in str(p)]
+        mapper = {p: sympy.Integer(32) for p in block_params}
+
+        op1 = Specializer(mapper).visit(op0)
+
+        assert "blk0_size" not in str(op1.ccode)
+        assert "x0_blk0 += 32" in str(op1.ccode)
+        assert "y0_blk0 += 32" in str(op1.ccode)
 
     def test_apply_basic(self):
         """
@@ -190,11 +203,6 @@ def test_apply_basic(self):
 
         assert np.all(check == f.data)
 
-    # TODO: Need a test to check that block sizes can be specialized
-    # TODO: Need to test that tile sizes can be specialized
-    # TODO: Test pickling followed by specialization
-    # TODO: Test specialized operators can be pickled
-
 
 class TestApply:
     """Tests for specialization of operators at apply time"""
@@ -311,9 +319,9 @@ def gaussian_bump(x, y, x0=0, y0=0, sigma=1):
         assert np.all(np.isclose(check, u.data))
 
         if source:
-            assert np.isclose(np.linalg.norm(u.data), 21.698477)
+            assert np.isclose(np.linalg.norm(u.data), 21.698477, atol=1e-4, rtol=0)
         else:
-            assert np.isclose(np.linalg.norm(u.data), 10.793053)
+            assert np.isclose(np.linalg.norm(u.data), 10.793053, atol=1e-4, rtol=0)
 
     @pytest.mark.parametrize('specialize',
                              [('x_m',),
@@ -374,4 +382,4 @@ def test_elastic(self, specialize):
         norms = (1.333946, 0.4931774, 1.333946, 1.1619346, 1.1619346)
 
         for field, norm in zip(fields, norms, strict=True):
-            assert np.isclose(np.linalg.norm(field.data), norm)
+            assert np.isclose(np.linalg.norm(field.data), norm, atol=1e-4, rtol=0)