test_loop_cutting.py 5.08 KB
Newer Older
 Martin Bauer committed Mar 21, 2019 1 ``````import numpy as np `````` Martin Bauer committed Jul 11, 2019 2 3 ``````import sympy as sp `````` Martin Bauer committed Mar 21, 2019 4 ``````import pystencils as ps `````` Martin Bauer committed Jul 11, 2019 5 ``````import pystencils.astnodes as ast `````` Michael Kuron committed Nov 15, 2019 6 ``````from pystencils.field import Field, FieldType `````` Martin Bauer committed Mar 21, 2019 7 ``````from pystencils.astnodes import Conditional, LoopOverCoordinate, SympyAssignment `````` Martin Bauer committed Jul 11, 2019 8 ``````from pystencils.cpu import create_kernel, make_python_function `````` Martin Bauer committed Mar 21, 2019 9 ``````from pystencils.kernelcreation import create_staggered_kernel `````` Martin Bauer committed Jul 11, 2019 10 11 ``````from pystencils.transformations import ( cleanup_blocks, cut_loop, move_constants_before_loop, simplify_conditionals) `````` Martin Bauer committed Mar 21, 2019 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 `````` def offsets_in_plane(normal_plane, offset_int, dimension): offset = [0] * dimension offset[normal_plane] = offset_int result = [tuple(offset)] for i in range(dimension): if i == normal_plane: continue lower = offset.copy() upper = offset.copy() lower[i] -= 1 upper[i] += 1 result.append(tuple(lower)) result.append(tuple(upper)) return result def test_staggered_iteration(): dim = 2 f_arr = np.arange(5**dim).reshape([5]*dim).astype(np.float64) s_arr = np.ones([5]*dim + [dim]) * 1234 s_arr_ref = s_arr.copy() fields_fixed = (Field.create_from_numpy_array('f', f_arr), `````` Michael Kuron committed Nov 15, 2019 37 `````` Field.create_from_numpy_array('s', s_arr, index_dimensions=1, field_type=FieldType.STAGGERED)) `````` Martin Bauer committed Mar 21, 2019 38 `````` fields_var = (Field.create_generic('f', 2), `````` Michael Kuron committed Nov 15, 2019 39 `````` Field.create_generic('s', 2, index_dimensions=1, index_shape=(dim,), field_type=FieldType.STAGGERED)) `````` Martin Bauer committed Mar 21, 2019 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 `````` for f, s in [fields_var, fields_fixed]: # --- Manual eqs = [] counters = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)] conditions = [counters[i] < f.shape[i] - 1 for i in range(dim)] for d in range(dim): eq = SympyAssignment(s(d), sum(f[o] for o in offsets_in_plane(d, 0, dim)) - sum(f[o] for o in offsets_in_plane(d, -1, dim))) cond = sp.And(*[conditions[i] for i in range(dim) if d != i]) eqs.append(Conditional(cond, eq)) func = create_kernel(eqs, ghost_layers=[(1, 0), (1, 0), (1, 0)]).compile() # --- Built-in optimized expressions = [] for d in range(dim): expressions.append(sum(f[o] for o in offsets_in_plane(d, 0, dim)) - sum(f[o] for o in offsets_in_plane(d, -1, dim))) `````` Michael Kuron committed Nov 29, 2019 58 59 `````` assignments = [ps.Assignment(s.staggered_access(d), expressions[i]) for i, d in enumerate(s.staggered_stencil)] func_optimized = create_staggered_kernel(assignments).compile() `````` Martin Bauer committed Mar 21, 2019 60 61 62 63 64 65 66 67 68 69 70 71 72 73 `````` assert not func_optimized.ast.atoms(Conditional), "Loop cutting optimization did not work" func(f=f_arr, s=s_arr_ref) func_optimized(f=f_arr, s=s_arr) np.testing.assert_almost_equal(s_arr_ref, s_arr) def test_staggered_iteration_manual(): dim = 2 f_arr = np.arange(5**dim).reshape([5]*dim) s_arr = np.ones([5]*dim + [dim]) * 1234 s_arr_ref = s_arr.copy() f = Field.create_from_numpy_array('f', f_arr) `````` Michael Kuron committed Nov 15, 2019 74 `````` s = Field.create_from_numpy_array('s', s_arr, index_dimensions=1, field_type=FieldType.STAGGERED) `````` Martin Bauer committed Mar 21, 2019 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 `````` eqs = [] counters = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)] conditions = [counters[i] < f.shape[i] - 1 for i in range(dim)] for d in range(dim): eq = SympyAssignment(s(d), sum(f[o] for o in offsets_in_plane(d, 0, dim)) - sum(f[o] for o in offsets_in_plane(d, -1, dim))) cond = sp.And(*[conditions[i] for i in range(dim) if d != i]) eqs.append(Conditional(cond, eq)) kernel_ast = create_kernel(eqs, ghost_layers=[(1, 0), (1, 0), (1, 0)]) func = make_python_function(kernel_ast) func(f=f_arr, s=s_arr_ref) inner_loop = [n for n in kernel_ast.atoms(ast.LoopOverCoordinate) if n.is_innermost_loop][0] cut_loop(inner_loop, [4]) outer_loop = [n for n in kernel_ast.atoms(ast.LoopOverCoordinate) if n.is_outermost_loop][0] cut_loop(outer_loop, [4]) simplify_conditionals(kernel_ast.body, loop_counter_simplification=True) cleanup_blocks(kernel_ast.body) move_constants_before_loop(kernel_ast.body) cleanup_blocks(kernel_ast.body) assert not kernel_ast.atoms(Conditional), "Loop cutting optimization did not work" func_optimized = make_python_function(kernel_ast) func_optimized(f=f_arr, s=s_arr) np.testing.assert_almost_equal(s_arr_ref, s_arr) def test_staggered_gpu(): dim = 2 `````` Michael Kuron committed Nov 14, 2019 111 `````` f = ps.fields("f: double[{dim}D]".format(dim=dim)) `````` Michael Kuron committed Nov 15, 2019 112 `````` s = ps.fields("s({dim}): double[{dim}D]".format(dim=dim), field_type=FieldType.STAGGERED) `````` Martin Bauer committed Mar 21, 2019 113 114 `````` expressions = [(f[0, 0] + f[-1, 0]) / 2, (f[0, 0] + f[0, -1]) / 2] `````` Michael Kuron committed Nov 29, 2019 115 116 `````` assignments = [ps.Assignment(s.staggered_access(d), expressions[i]) for i, d in enumerate(s.staggered_stencil)] kernel_ast = ps.create_staggered_kernel(assignments, target='gpu', gpu_exclusive_conditions=True) `````` Martin Bauer committed Mar 21, 2019 117 118 `````` assert len(kernel_ast.atoms(Conditional)) == 4 `````` Michael Kuron committed Nov 29, 2019 119 120 `````` assignments = [ps.Assignment(s.staggered_access(d), expressions[i]) for i, d in enumerate(s.staggered_stencil)] kernel_ast = ps.create_staggered_kernel(assignments, target='gpu', gpu_exclusive_conditions=False) `````` Martin Bauer committed Mar 21, 2019 121 `` assert len(kernel_ast.atoms(Conditional)) == 3``