Added test cases to vectorisation

pystencils_tests/test_vectorization.py

@@ -4,6 +4,7 @@ import sympy as sp
 import pystencils as ps
 from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets
 from pystencils.cpu.vectorization import vectorize
+from pystencils.fast_approximation import insert_fast_sqrts, insert_fast_divisions
 from pystencils.transformations import replace_inner_stride_with_one
 
 
@@ -109,7 +110,6 @@ def test_piecewise1():
 
 
 def test_piecewise2():
-
     arr = np.zeros((20, 20))
 
     @ps.kernel
@@ -128,7 +128,6 @@ def test_piecewise2():
 
 
 def test_piecewise3():
-
     arr = np.zeros((22, 22))
 
     @ps.kernel
@@ -146,12 +145,32 @@ def test_logical_operators():
     arr = np.zeros((22, 22))
 
     @ps.kernel
-    def test_kernel(s):
+    def kernel_and(s):
         f, g = ps.fields(f=arr, g=arr)
         s.c @= sp.And(f[0, 1] < 0.0, f[1, 0] < 0.0)
         g[0, 0] @= sp.Piecewise([1.0 / f[1, 0], s.c], [1.0, True])
 
-    ast = ps.create_kernel(test_kernel)
+    ast = ps.create_kernel(kernel_and)
+    vectorize(ast)
+    ast.compile()
+
+    @ps.kernel
+    def kernel_or(s):
+        f, g = ps.fields(f=arr, g=arr)
+        s.c @= sp.Or(f[0, 1] < 0.0, f[1, 0] < 0.0)
+        g[0, 0] @= sp.Piecewise([1.0 / f[1, 0], s.c], [1.0, True])
+
+    ast = ps.create_kernel(kernel_or)
+    vectorize(ast)
+    ast.compile()
+
+    @ps.kernel
+    def kernel_equal(s):
+        f, g = ps.fields(f=arr, g=arr)
+        s.c @= sp.Eq(f[0, 1], 2.0)
+        g[0, 0] @= sp.Piecewise([1.0 / f[1, 0], s.c], [1.0, True])
+
+    ast = ps.create_kernel(kernel_equal)
     vectorize(ast)
     ast.compile()
 
@@ -159,3 +178,61 @@ def test_logical_operators():
 def test_hardware_query():
     instruction_sets = get_supported_instruction_sets()
     assert 'sse' in instruction_sets
+
+
+def test_vectorised_pow():
+    arr = np.zeros((24, 24))
+    f, g = ps.fields(f=arr, g=arr)
+
+    as1 = ps.Assignment(g[0, 0], sp.Pow(f[0, 0], 2))
+    as2 = ps.Assignment(g[0, 0], sp.Pow(f[0, 0], 0.5))
+    as3 = ps.Assignment(g[0, 0], sp.Pow(f[0, 0], -0.5))
+    as4 = ps.Assignment(g[0, 0], sp.Pow(f[0, 0], 4))
+    as5 = ps.Assignment(g[0, 0], sp.Pow(f[0, 0], -4))
+    as6 = ps.Assignment(g[0, 0], sp.Pow(f[0, 0], -1))
+
+    ast = ps.create_kernel(as1)
+    vectorize(ast)
+    ast.compile()
+
+    ast = ps.create_kernel(as2)
+    vectorize(ast)
+    ast.compile()
+
+    ast = ps.create_kernel(as3)
+    vectorize(ast)
+    ast.compile()
+
+    ast = ps.create_kernel(as4)
+    vectorize(ast)
+    ast.compile()
+
+    ast = ps.create_kernel(as5)
+    vectorize(ast)
+    ast.compile()
+
+    ast = ps.create_kernel(as6)
+    vectorize(ast)
+    ast.compile()
+
+
+def test_vectorised_fast_approximations():
+    arr = np.zeros((24, 24))
+    f, g = ps.fields(f=arr, g=arr)
+
+    expr = sp.sqrt(f[0, 0] + f[1, 0])
+    assignment = ps.Assignment(g[0, 0], insert_fast_sqrts(expr))
+    ast = ps.create_kernel(assignment)
+    vectorize(ast)
+    ast.compile()
+
+    expr = f[0, 0] / f[1, 0]
+    assignment = ps.Assignment(g[0, 0], insert_fast_divisions(expr))
+    ast = ps.create_kernel(assignment)
+    vectorize(ast)
+    ast.compile()
+
+    assignment = ps.Assignment(sp.Symbol("tmp"), 3 / sp.sqrt(f[0, 0] + f[1, 0]))
+    ast = ps.create_kernel(insert_fast_sqrts(assignment))
+    vectorize(ast)
+    ast.compile()