diff --git a/pystencils_tests/test_vectorization.py b/pystencils_tests/test_vectorization.py index e2e4a396e14a71398a68bbe6959d006e10242d0a..04e505cd4044b6a3a54f0372e743fd5febde7101 100644 --- a/pystencils_tests/test_vectorization.py +++ b/pystencils_tests/test_vectorization.py @@ -7,6 +7,8 @@ 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 +supported_instruction_sets = get_supported_instruction_sets() if get_supported_instruction_sets() else [] +instruction_set = "avx" if "avx" in supported_instruction_sets else "sse" def test_vector_type_propagation(): a, b, c, d, e = sp.symbols("a b c d e") @@ -19,7 +21,7 @@ def test_vector_type_propagation(): ps.Assignment(g[0, 0], b + 3 + f[0, 1])] ast = ps.create_kernel(update_rule) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) func = ast.compile() dst = np.zeros_like(arr) @@ -43,7 +45,6 @@ def test_inplace_update(): f2 @= 2 * s.tmp0 ast = ps.create_kernel(update_rule, cpu_vectorize_info={'instruction_set': 'sse'}) - print(ps.show_code(ast)) kernel = ast.compile() kernel(f=arr) np.testing.assert_equal(arr, 2) @@ -68,7 +69,7 @@ def test_vectorization_fixed_size(): update_rule = [ps.Assignment(g[0, 0], f[0, 0] + f[-1, 0] + f[1, 0] + f[0, 1] + f[0, -1] + 42.0)] ast = ps.create_kernel(update_rule) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) func = ast.compile() dst = np.zeros_like(arr) @@ -82,7 +83,7 @@ def test_vectorization_variable_size(): ast = ps.create_kernel(update_rule) replace_inner_stride_with_one(ast) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) func = ast.compile() arr = np.ones((23 + 2, 17 + 2)) * 5.0 @@ -103,7 +104,7 @@ def test_piecewise1(): ps.Assignment(g[0, 0], sp.Piecewise((b + 3 + f[0, 1], c), (0.0, True)))] ast = ps.create_kernel(update_rule) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) func = ast.compile() dst = np.zeros_like(arr) func(g=dst, f=arr) @@ -122,7 +123,7 @@ def test_piecewise2(): g[0, 0] @= s.result ast = ps.create_kernel(test_kernel) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) func = ast.compile() func(f=arr, g=arr) np.testing.assert_equal(arr, np.ones_like(arr)) @@ -138,7 +139,7 @@ def test_piecewise3(): g[0, 0] @= 1.0 / (s.b + s.k) if f[0, 0] > 0.0 else 1.0 ast = ps.create_kernel(test_kernel) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() @@ -152,7 +153,7 @@ def test_logical_operators(): g[0, 0] @= sp.Piecewise([1.0 / f[1, 0], s.c], [1.0, True]) ast = ps.create_kernel(kernel_and) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() @ps.kernel @@ -162,7 +163,7 @@ def test_logical_operators(): g[0, 0] @= sp.Piecewise([1.0 / f[1, 0], s.c], [1.0, True]) ast = ps.create_kernel(kernel_or) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() @ps.kernel @@ -172,7 +173,7 @@ def test_logical_operators(): g[0, 0] @= sp.Piecewise([1.0 / f[1, 0], s.c], [1.0, True]) ast = ps.create_kernel(kernel_equal) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() @@ -193,27 +194,27 @@ def test_vectorised_pow(): as6 = ps.Assignment(g[0, 0], sp.Pow(f[0, 0], -1)) ast = ps.create_kernel(as1) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() ast = ps.create_kernel(as2) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() ast = ps.create_kernel(as3) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() ast = ps.create_kernel(as4) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() ast = ps.create_kernel(as5) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() ast = ps.create_kernel(as6) - vectorize(ast) + vectorize(ast, instruction_set=instruction_set) ast.compile() @@ -224,16 +225,16 @@ def test_vectorised_fast_approximations(): 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) + vectorize(ast, instruction_set=instruction_set) 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) + vectorize(ast, instruction_set=instruction_set) 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) + vectorize(ast, instruction_set=instruction_set) ast.compile()