From 3c7f36d1fdad82c8c50236691aac127234b1988a Mon Sep 17 00:00:00 2001 From: Markus Holzer <markus.holzer@fau.de> Date: Wed, 2 Feb 2022 19:48:37 +0100 Subject: [PATCH] Minor fixes and separation of test cases and quicktests --- pystencils/cpu/vectorization.py | 13 ++-- pystencils/gpucuda/kernelcreation.py | 2 +- pystencils/kernelcreation.py | 23 ++++-- pystencils/node_collection.py | 2 +- pystencils/simp/assignment_collection.py | 12 ++- pystencils/transformations.py | 2 +- pystencils/typing/__init__.py | 18 ++++- pystencils/typing/cast_functions.py | 7 +- pystencils/typing/leaf_typing.py | 2 +- pystencils/typing/transformations.py | 1 - pystencils/typing/types.py | 1 - pystencils/typing/utilities.py | 11 +-- pystencils_tests/test_datahandling.py | 2 +- pystencils_tests/test_loop_cutting.py | 7 +- pystencils_tests/test_quicktests.py | 74 +++++++++++++++++++ .../test_simplification_strategy.py | 1 + setup.py | 7 +- 17 files changed, 135 insertions(+), 50 deletions(-) create mode 100644 pystencils_tests/test_quicktests.py diff --git a/pystencils/cpu/vectorization.py b/pystencils/cpu/vectorization.py index 4d609a114..ac25639b1 100644 --- a/pystencils/cpu/vectorization.py +++ b/pystencils/cpu/vectorization.py @@ -7,8 +7,8 @@ from sympy.logic.boolalg import BooleanFunction, BooleanAtom import pystencils.astnodes as ast from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets, get_vector_instruction_set -from pystencils.typing import ( BasicType, PointerType, TypedSymbol, VectorType, CastFunc, collate_types, - get_type_of_expression, VectorMemoryAccess) +from pystencils.typing import (BasicType, PointerType, TypedSymbol, VectorType, CastFunc, collate_types, + get_type_of_expression, VectorMemoryAccess) from pystencils.fast_approximation import fast_division, fast_inv_sqrt, fast_sqrt from pystencils.functions import DivFunc from pystencils.field import Field @@ -203,9 +203,10 @@ def vectorize_inner_loops_and_adapt_load_stores(ast_node, assume_aligned, nontem loop_node.step = vector_width loop_node.subs(substitutions) vector_int_width = ast_node.instruction_set['intwidth'] - vector_loop_counter = CastFunc(loop_counter_symbol, VectorType(loop_counter_symbol.dtype, vector_int_width)) \ - + CastFunc(tuple(range(vector_int_width if type(vector_int_width) is int else 2)), - VectorType(loop_counter_symbol.dtype, vector_int_width)) + arg_1 = CastFunc(loop_counter_symbol, VectorType(loop_counter_symbol.dtype, vector_int_width)) + arg_2 = CastFunc(tuple(range(vector_int_width if type(vector_int_width) is int else 2)), + VectorType(loop_counter_symbol.dtype, vector_int_width)) + vector_loop_counter = arg_1 + arg_2 fast_subs(loop_node, {loop_counter_symbol: vector_loop_counter}, skip=lambda e: isinstance(e, ast.ResolvedFieldAccess) or isinstance(e, VectorMemoryAccess)) @@ -333,7 +334,7 @@ def insert_vector_casts(ast_node, instruction_set, default_float_type='double'): assignment = arg # If there is a remainder loop we do not vectorise it, thus lhs will indicate this # if isinstance(assignment.lhs, ast.ResolvedFieldAccess): - # continue + # continue subs_expr = fast_subs(assignment.rhs, substitution_dict, skip=lambda e: isinstance(e, ast.ResolvedFieldAccess)) assignment.rhs = visit_expr(subs_expr, default_type) diff --git a/pystencils/gpucuda/kernelcreation.py b/pystencils/gpucuda/kernelcreation.py index 21721bb7f..a50953b64 100644 --- a/pystencils/gpucuda/kernelcreation.py +++ b/pystencils/gpucuda/kernelcreation.py @@ -10,7 +10,7 @@ from pystencils.field import Field, FieldType from pystencils.enums import Target, Backend from pystencils.gpucuda.cudajit import make_python_function from pystencils.node_collection import NodeCollection -from pystencils.gpucuda.indexing import BlockIndexing, indexing_creator_from_params +from pystencils.gpucuda.indexing import indexing_creator_from_params from pystencils.simp.assignment_collection import AssignmentCollection from pystencils.transformations import ( get_base_buffer_index, get_common_shape, parse_base_pointer_info, diff --git a/pystencils/kernelcreation.py b/pystencils/kernelcreation.py index 50eaedf4f..bfb2c09a3 100644 --- a/pystencils/kernelcreation.py +++ b/pystencils/kernelcreation.py @@ -74,7 +74,9 @@ def create_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCol except Exception as e: warnings.warn(f"It was not possible to apply the default pystencils optimisations to the " f"AssignmentCollection due to the following problem :{e}") + simplification_hints = assignments.simplification_hints assignments = NodeCollection(assignments.all_assignments) + assignments.simplification_hints = simplification_hints if config.index_fields: return create_indexed_kernel(assignments, config=config) @@ -86,6 +88,9 @@ def create_domain_kernel(assignments: NodeCollection, *, config: CreateKernelCon """ Creates abstract syntax tree (AST) of kernel, using a list of update equations. + Note that `create_domain_kernel` is a lower level function which shoul be accessed by not providing `index_fields` + to create_kernel + Args: assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection` config: CreateKernelConfig which includes the needed configuration @@ -179,6 +184,9 @@ def create_indexed_kernel(assignments: NodeCollection, *, config: CreateKernelCo 'coordinate_names' parameter. The struct can have also other fields that can be read and written in the kernel, for example boundary parameters. + Note that `create_indexed_kernel` is a lower level function which shoul be accessed by providing `index_fields` + to create_kernel + Args: assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection` config: CreateKernelConfig which includes the needed configuration @@ -188,8 +196,8 @@ def create_indexed_kernel(assignments: NodeCollection, *, config: CreateKernelCo can be compiled with through its 'compile()' member Example: - >>> import pystencils.kernel_creation_config >>> import pystencils as ps + >>> from pystencils.node_collection import NodeCollection >>> import numpy as np >>> from pystencils.kernelcreation import create_indexed_kernel >>> @@ -202,16 +210,17 @@ def create_indexed_kernel(assignments: NodeCollection, *, config: CreateKernelCo >>> s, d = ps.fields('s, d: [2D]') >>> assignment = ps.Assignment(d[0, 0], 2 * s[0, 1] + 2 * s[1, 0] + idx_field('val')) >>> kernel_config = ps.CreateKernelConfig(index_fields=[idx_field], coordinate_names=('x', 'y')) - >>> kernel_ast = create_indexed_kernel(ps.AssignmentCollection([assignment]), config=kernel_config) + >>> kernel_ast = create_indexed_kernel(NodeCollection([assignment]), config=kernel_config) >>> kernel = kernel_ast.compile() >>> d_arr = np.zeros([5, 5]) >>> kernel(s=np.ones([5, 5]), d=d_arr, idx=index_arr) >>> d_arr - array([[0., 0., 0., 0., 0.], - [0., 4.1, 0., 0., 0.], - [0., 0., 4.2, 0., 0.], - [0., 0., 0., 4.3, 0.], - [0., 0., 0., 0., 0.]]) + array([[0. , 0. , 0. , 0. , 0. ], + [0. , 4.1, 0. , 0. , 0. ], + [0. , 0. , 4.2, 0. , 0. ], + [0. , 0. , 0. , 4.3, 0. ], + [0. , 0. , 0. , 0. , 0. ]]) + """ # --- eval assignments.evaluate_terms() diff --git a/pystencils/node_collection.py b/pystencils/node_collection.py index 0287b6fc8..8e396cc38 100644 --- a/pystencils/node_collection.py +++ b/pystencils/node_collection.py @@ -26,7 +26,7 @@ class NodeCollection: raise ValueError(f'The list "{assignments}" is mixed. Pass either a list of "pystencils.Assignments" ' f'or a list of "pystencils.astnodes.Node') - self.simplification_hints = () + self.simplification_hints = {} def evaluate_terms(self): evaluate_constant_terms = ReplaceOptim( diff --git a/pystencils/simp/assignment_collection.py b/pystencils/simp/assignment_collection.py index 69dcf9567..b3324e42f 100644 --- a/pystencils/simp/assignment_collection.py +++ b/pystencils/simp/assignment_collection.py @@ -136,8 +136,7 @@ class AssignmentCollection: bound_symbols_set = bound_symbols_set.union(*[ assignment.symbols_defined for assignment in self.all_assignments if isinstance(assignment, pystencils.astnodes.Node) - ] - ) + ]) return bound_symbols_set @@ -159,11 +158,9 @@ class AssignmentCollection: @property def defined_symbols(self) -> Set[sp.Symbol]: """All symbols which occur as left-hand-sides of one of the main equations""" - return (set( - [assignment.lhs for assignment in self.main_assignments if isinstance(assignment, Assignment)] - ).union(*[assignment.symbols_defined for assignment in self.main_assignments if isinstance( - assignment, pystencils.astnodes.Node)] - )) + lhs_set = set([assignment.lhs for assignment in self.main_assignments if isinstance(assignment, Assignment)]) + return (lhs_set.union(*[assignment.symbols_defined for assignment in self.main_assignments + if isinstance(assignment, pystencils.astnodes.Node)])) @property def operation_count(self): @@ -365,6 +362,7 @@ class AssignmentCollection: new_assignment = [fast_subs(eq, substitution_dict) for eq in self.main_assignments] return self.copy(new_assignment, kept_subexpressions) + # ----------------------------------------- Display and Printing ------------------------------------------------- def _repr_html_(self): diff --git a/pystencils/transformations.py b/pystencils/transformations.py index 7f864f9af..2e885904a 100644 --- a/pystencils/transformations.py +++ b/pystencils/transformations.py @@ -11,7 +11,7 @@ import pystencils.astnodes as ast from pystencils.assignment import Assignment from pystencils.typing import ( PointerType, StructType, TypedSymbol, get_base_type, ReinterpretCastFunc, get_next_parent_of_type, parents_of_type) -from pystencils.field import Field, Field, FieldType +from pystencils.field import Field, FieldType from pystencils.typing import FieldPointerSymbol from pystencils.simp.assignment_collection import AssignmentCollection from pystencils.slicing import normalize_slice diff --git a/pystencils/typing/__init__.py b/pystencils/typing/__init__.py index 2221b812b..5bb560d10 100644 --- a/pystencils/typing/__init__.py +++ b/pystencils/typing/__init__.py @@ -1,6 +1,16 @@ +from pystencils.typing.cast_functions import (CastFunc, BooleanCastFunc, VectorMemoryAccess, ReinterpretCastFunc, + PointerArithmeticFunc) +from pystencils.typing.types import (is_supported_type, numpy_name_to_c, AbstractType, BasicType, VectorType, + PointerType, StructType, create_type) +from pystencils.typing.typed_sympy import (assumptions_from_dtype, TypedSymbol, FieldStrideSymbol, FieldShapeSymbol, + FieldPointerSymbol) +from pystencils.typing.utilities import (typed_symbols, get_base_type, result_type, collate_types, + get_type_of_expression, insert_casts, get_next_parent_of_type, parents_of_type) -from pystencils.typing.types import * -from pystencils.typing.typed_sympy import * -from pystencils.typing.cast_functions import * -from pystencils.typing.utilities import * +__all__ = ['CastFunc', 'BooleanCastFunc', 'VectorMemoryAccess', 'ReinterpretCastFunc', 'PointerArithmeticFunc', + 'is_supported_type', 'numpy_name_to_c', 'AbstractType', 'BasicType', + 'VectorType', 'PointerType', 'StructType', 'create_type', + 'assumptions_from_dtype', 'TypedSymbol', 'FieldStrideSymbol', 'FieldShapeSymbol', 'FieldPointerSymbol', + 'typed_symbols', 'get_base_type', 'result_type', 'collate_types', + 'get_type_of_expression', 'insert_casts', 'get_next_parent_of_type', 'parents_of_type'] diff --git a/pystencils/typing/cast_functions.py b/pystencils/typing/cast_functions.py index 8200e9697..76686c211 100644 --- a/pystencils/typing/cast_functions.py +++ b/pystencils/typing/cast_functions.py @@ -2,7 +2,7 @@ import numpy as np import sympy as sp from sympy.logic.boolalg import Boolean -from pystencils.typing.types import AbstractType, BasicType, create_type +from pystencils.typing.types import AbstractType, BasicType from pystencils.typing.typed_sympy import TypedSymbol @@ -93,9 +93,8 @@ class CastFunc(sp.Function): See :func:`.TypedSymbol.is_integer` """ if hasattr(self.dtype, 'numpy_dtype'): - return np.issubdtype(self.dtype.numpy_dtype, np.integer) or \ - np.issubdtype(self.dtype.numpy_dtype, np.floating) or \ - super().is_real + return np.issubdtype(self.dtype.numpy_dtype, np.integer) or np.issubdtype(self.dtype.numpy_dtype, + np.floating) or super().is_real else: return super().is_real diff --git a/pystencils/typing/leaf_typing.py b/pystencils/typing/leaf_typing.py index 6ccd864e3..c62824892 100644 --- a/pystencils/typing/leaf_typing.py +++ b/pystencils/typing/leaf_typing.py @@ -185,7 +185,7 @@ class TypeAdder: collated_type = collate_types([t for _, t in args_types]) new_expressions = [a if t.dtype_eq(collated_type) else CastFunc(a, collated_type) for a, t in args_types] return expr.func(expr.args[0], expr.args[1], *new_expressions), collated_type - #elif isinstance(expr, sp.Mul): + # elif isinstance(expr, sp.Mul): # raise NotImplementedError('sp.Mul') # # TODO can we ignore this and move it to general expr handling, i.e. removing Mul? # # args_types = [self.figure_out_type(arg) for arg in expr.args if arg not in (-1, 1)] diff --git a/pystencils/typing/transformations.py b/pystencils/typing/transformations.py index f5ddcfa42..74ecf19f1 100644 --- a/pystencils/typing/transformations.py +++ b/pystencils/typing/transformations.py @@ -2,7 +2,6 @@ from typing import List from pystencils.config import CreateKernelConfig from pystencils.typing.leaf_typing import TypeAdder -from pystencils.typing import BasicType from sympy.codegen import Assignment diff --git a/pystencils/typing/types.py b/pystencils/typing/types.py index 2f45ff4af..dbe284496 100644 --- a/pystencils/typing/types.py +++ b/pystencils/typing/types.py @@ -293,4 +293,3 @@ def create_type(specification: Union[np.dtype, AbstractType, str]) -> AbstractTy return BasicType(numpy_dtype, const=False) else: return StructType(numpy_dtype, const=False) - diff --git a/pystencils/typing/utilities.py b/pystencils/typing/utilities.py index 15d0beed1..6a43c7984 100644 --- a/pystencils/typing/utilities.py +++ b/pystencils/typing/utilities.py @@ -34,8 +34,6 @@ def get_base_type(data_type): return data_type -############################# This is basically our type system ######################################################## - def result_type(*args: np.dtype): s = sorted(args, key=lambda x: x.itemsize) @@ -104,7 +102,8 @@ def get_type_of_expression(expr, # TODO: we shouldn't need to have default. AST leaves should have a type default_int_type='int', # TODO: we shouldn't need to have default. AST leaves should have a type - symbol_type_dict=None): # TODO: we shouldn't need to have default. AST leaves should have a type + # TODO: we shouldn't need to have default. AST leaves should have a type + symbol_type_dict=None): from pystencils.astnodes import ResolvedFieldAccess from pystencils.cpu.vectorization import vec_all, vec_any @@ -181,9 +180,6 @@ def get_type_of_expression(expr, raise NotImplementedError("Could not determine type for", expr, type(expr)) -# ############################# End This is basically our type system ################################################## - - # TODO this seems quite wrong... sympy_version = sp.__version__.split('.') if int(sympy_version[0]) * 100 + int(sympy_version[1]) >= 109: @@ -191,7 +187,6 @@ if int(sympy_version[0]) * 100 + int(sympy_version[1]) >= 109: sp.Number.__getstate__ = sp.Basic.__getstate__ del sp.Basic.__getstate__ - class FunctorWithStoredKwargs: def __init__(self, func, **kwargs): self.func = func @@ -200,7 +195,6 @@ if int(sympy_version[0]) * 100 + int(sympy_version[1]) >= 109: def __call__(self, *args): return self.func(*args, **self.kwargs) - # __reduce_ex__ would strip kwargs, so we override it def basic_reduce_ex(self, protocol): if hasattr(self, '__getnewargs_ex__'): @@ -213,7 +207,6 @@ if int(sympy_version[0]) * 100 + int(sympy_version[1]) >= 109: state = None return FunctorWithStoredKwargs(type(self), **kwargs), args, state - sp.Number.__reduce_ex__ = sp.Basic.__reduce_ex__ sp.Basic.__reduce_ex__ = basic_reduce_ex diff --git a/pystencils_tests/test_datahandling.py b/pystencils_tests/test_datahandling.py index be695d078..afd5f70da 100644 --- a/pystencils_tests/test_datahandling.py +++ b/pystencils_tests/test_datahandling.py @@ -132,7 +132,7 @@ def kernel_execution_jacobi(dh, target): def jacobi(): dh.fields.tmp.center @= sum(dh.fields.f.neighbors(stencil)) / len(stencil) - kernel = create_kernel(jacobi, target=target).compile() + kernel = create_kernel(jacobi, config=ps.CreateKernelConfig(target=target)).compile() for b in dh.iterate(ghost_layers=1): b['f'].fill(42) dh.run_kernel(kernel) diff --git a/pystencils_tests/test_loop_cutting.py b/pystencils_tests/test_loop_cutting.py index 9c833aca6..0372c5739 100644 --- a/pystencils_tests/test_loop_cutting.py +++ b/pystencils_tests/test_loop_cutting.py @@ -50,7 +50,9 @@ def test_staggered_iteration(): 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() + # TODO: correct type hint + config = ps.CreateKernelConfig(target=ps.Target.CPU, ghost_layers=[(1, 0), (1, 0), (1, 0)]) + func = ps.create_kernel(eqs, config=config).compile() # --- Built-in optimized expressions = [] @@ -93,7 +95,8 @@ def test_staggered_iteration_manual(): cond = sp.And(*[conditions2]) eqs.append(Conditional(cond, eq)) - kernel_ast = create_kernel(eqs, ghost_layers=[(1, 0), (1, 0), (1, 0)]) + config = ps.CreateKernelConfig(target=ps.Target.CPU, ghost_layers=[(1, 0), (1, 0), (1, 0)]) + kernel_ast = ps.create_kernel(eqs, config=config) func = make_python_function(kernel_ast) func(f=f_arr, s=s_arr_ref) diff --git a/pystencils_tests/test_quicktests.py b/pystencils_tests/test_quicktests.py new file mode 100644 index 000000000..d694b30b4 --- /dev/null +++ b/pystencils_tests/test_quicktests.py @@ -0,0 +1,74 @@ +import numpy as np + +import pystencils as ps +from pystencils.cpu.vectorization import get_supported_instruction_sets +from pystencils.cpu.vectorization import replace_inner_stride_with_one, vectorize + + +def test_basic_kernel(): + for domain_shape in [(4, 5), (3, 4, 5)]: + dh = ps.create_data_handling(domain_size=domain_shape, periodicity=True) + assert all(dh.periodicity) + + f = dh.add_array('f', values_per_cell=1) + tmp = dh.add_array('tmp', values_per_cell=1) + + stencil_2d = [(1, 0), (-1, 0), (0, 1), (0, -1)] + stencil_3d = [(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1)] + stencil = stencil_2d if dh.dim == 2 else stencil_3d + + jacobi = ps.Assignment(tmp.center, sum(f.neighbors(stencil)) / len(stencil)) + kernel = ps.create_kernel(jacobi).compile() + + for b in dh.iterate(ghost_layers=1): + b['f'].fill(42) + dh.run_kernel(kernel) + for b in dh.iterate(ghost_layers=0): + np.testing.assert_equal(b['f'], 42) + + float_seq = [1.0, 2.0, 3.0, 4.0] + int_seq = [1, 2, 3] + for op in ('min', 'max', 'sum'): + assert (dh.reduce_float_sequence(float_seq, op) == float_seq).all() + assert (dh.reduce_int_sequence(int_seq, op) == int_seq).all() + + +def test_basic_blocking_staggered(): + f = ps.fields("f: double[2D]") + stag = ps.fields("stag(2): double[2D]", field_type=ps.FieldType.STAGGERED) + terms = [ + f[0, 0] - f[-1, 0], + f[0, 0] - f[0, -1], + ] + assignments = [ps.Assignment(stag.staggered_access(d), terms[i]) for i, d in enumerate(stag.staggered_stencil)] + kernel = ps.create_staggered_kernel(assignments, cpu_blocking=(3, 16)).compile() + reference_kernel = ps.create_staggered_kernel(assignments).compile() + + f_arr = np.random.rand(80, 33) + stag_arr = np.zeros((80, 33, 3)) + stag_ref = np.zeros((80, 33, 3)) + kernel(f=f_arr, stag=stag_arr) + reference_kernel(f=f_arr, stag=stag_ref) + np.testing.assert_almost_equal(stag_arr, stag_ref) + + +def test_basic_vectorization(): + supported_instruction_sets = get_supported_instruction_sets() + if supported_instruction_sets: + instruction_set = supported_instruction_sets[-1] + else: + instruction_set = None + + f, g = ps.fields("f, g : double[2D]") + 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) + + replace_inner_stride_with_one(ast) + vectorize(ast, instruction_set=instruction_set) + func = ast.compile() + + arr = np.ones((23 + 2, 17 + 2)) * 5.0 + dst = np.zeros_like(arr) + + func(g=dst, f=arr) + np.testing.assert_equal(dst[1:-1, 1:-1], 5 * 5.0 + 42.0) \ No newline at end of file diff --git a/pystencils_tests/test_simplification_strategy.py b/pystencils_tests/test_simplification_strategy.py index 31fa43544..40b350af3 100644 --- a/pystencils_tests/test_simplification_strategy.py +++ b/pystencils_tests/test_simplification_strategy.py @@ -71,6 +71,7 @@ def test_split_inner_loop(): ast = ps.create_kernel(ac) code = ps.get_code_str(ast) + ps.show_code(ast) # we have four inner loops as indicated in split groups (4 elements) plus one outer loop assert code.count('for') == 5 ast = ps.create_kernel(ac, target=ps.Target.GPU) diff --git a/setup.py b/setup.py index a2053b422..5e2cfc866 100644 --- a/setup.py +++ b/setup.py @@ -16,10 +16,9 @@ except ImportError: USE_CYTHON = False quick_tests = [ - 'test_datahandling.test_kernel', - 'test_blocking_staggered.test_blocking_staggered', - 'test_blocking_staggered.test_blocking_staggered', - 'test_vectorization.test_vectorization_variable_size', + 'test_quicktests.test_basic_kernel', + 'test_quicktests.test_basic_blocking_staggered', + 'test_quicktests.test_basic_vectorization', ] -- GitLab