From 0da250c0e70ed5d88b0e8044eafb4eda59773f22 Mon Sep 17 00:00:00 2001 From: Stephan Seitz <stephan.seitz@fau.de> Date: Sun, 18 Aug 2019 12:46:54 +0200 Subject: [PATCH] Regard case that pystencils.astnodes.Conditional is in assignments or sympy.codegen.rewriting does not exist --- pystencils/astnodes.py | 7 ++++++ pystencils/cpu/kernelcreation.py | 6 ++--- pystencils/gpucuda/kernelcreation.py | 9 +++----- pystencils/optimizations.py | 23 ++++++++++++++------ pystencils_tests/test_sympy_optimizations.py | 5 +++++ 5 files changed, 34 insertions(+), 16 deletions(-) diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py index 393396aa..b2413828 100644 --- a/pystencils/astnodes.py +++ b/pystencils/astnodes.py @@ -509,6 +509,13 @@ class SympyAssignment(Node): self.lhs = fast_subs(self.lhs, subs_dict) self.rhs = fast_subs(self.rhs, subs_dict) + def optimize(self, optimizations): + try: + from sympy.codegen.rewriting import optimize + self.rhs = optimize(self.rhs, optimizations) + except Exception: + pass + @property def args(self): return [self._lhs_symbol, self.rhs] diff --git a/pystencils/cpu/kernelcreation.py b/pystencils/cpu/kernelcreation.py index 27d4f196..9a5374d1 100644 --- a/pystencils/cpu/kernelcreation.py +++ b/pystencils/cpu/kernelcreation.py @@ -8,7 +8,7 @@ from pystencils.astnodes import Block, KernelFunction, LoopOverCoordinate, Sympy from pystencils.cpu.cpujit import make_python_function from pystencils.data_types import BasicType, StructType, TypedSymbol, create_type from pystencils.field import Field, FieldType -from pystencils.optimizations import optims_pystencils_cpu +from pystencils.optimizations import optims_pystencils_cpu, optimize_assignments from pystencils.transformations import ( add_types, filtered_tree_iteration, get_base_buffer_index, get_optimal_loop_ordering, make_loop_over_domain, move_constants_before_loop, parse_base_pointer_info, @@ -57,7 +57,7 @@ def create_kernel(assignments: AssignmentOrAstNodeList, function_name: str = "ke if optimizations is None: optimizations = optims_pystencils_cpu - assignments = [Assignment(a.lhs, sp.codegen.rewriting.optimize(a.rhs, optimizations)) for a in assignments] + assignments = optimize_assignments(assignments, optimizations) fields_read, fields_written, assignments = add_types(assignments, type_info, not skip_independence_check) all_fields = fields_read.union(fields_written) @@ -119,7 +119,7 @@ def create_indexed_kernel(assignments: AssignmentOrAstNodeList, if optimizations is None: optimizations = optims_pystencils_cpu - assignments = [Assignment(a.lhs, sp.codegen.rewriting.optimize(a.rhs, optimizations)) for a in assignments] + assignments = optimize_assignments(assignments, optimizations) fields_read, fields_written, assignments = add_types(assignments, type_info, check_independence_condition=False) all_fields = fields_read.union(fields_written) diff --git a/pystencils/gpucuda/kernelcreation.py b/pystencils/gpucuda/kernelcreation.py index 65f39956..84908336 100644 --- a/pystencils/gpucuda/kernelcreation.py +++ b/pystencils/gpucuda/kernelcreation.py @@ -1,12 +1,9 @@ -import sympy as sp - -from pystencils import Assignment from pystencils.astnodes import Block, KernelFunction, LoopOverCoordinate, SympyAssignment from pystencils.data_types import BasicType, StructType, TypedSymbol from pystencils.field import Field, FieldType from pystencils.gpucuda.cudajit import make_python_function from pystencils.gpucuda.indexing import BlockIndexing -from pystencils.optimizations import optims_pystencils_gpu +from pystencils.optimizations import optims_pystencils_gpu, optimize_assignments from pystencils.transformations import ( add_types, get_base_buffer_index, get_common_shape, parse_base_pointer_info, resolve_buffer_accesses, resolve_field_accesses, unify_shape_symbols) @@ -17,7 +14,7 @@ def create_cuda_kernel(assignments, function_name="kernel", type_info=None, inde if optimizations is None: optimizations = optims_pystencils_gpu - assignments = [Assignment(a.lhs, sp.codegen.rewriting.optimize(a.rhs, optimizations)) for a in assignments] + assignments = optimize_assignments(assignments, optimizations) fields_read, fields_written, assignments = add_types(assignments, type_info, not skip_independence_check) all_fields = fields_read.union(fields_written) @@ -99,7 +96,7 @@ def created_indexed_cuda_kernel(assignments, index_fields, function_name="kernel coordinate_names=('x', 'y', 'z'), indexing_creator=BlockIndexing, optimizations=None): if optimizations is None: optimizations = optims_pystencils_gpu - assignments = [Assignment(a.lhs, sp.codegen.rewriting.optimize(a.rhs, optimizations)) for a in assignments] + assignments = optimize_assignments(assignments, optimizations) fields_read, fields_written, assignments = add_types(assignments, type_info, check_independence_condition=False) all_fields = fields_read.union(fields_written) diff --git a/pystencils/optimizations.py b/pystencils/optimizations.py index 01aba5d8..44fda53e 100644 --- a/pystencils/optimizations.py +++ b/pystencils/optimizations.py @@ -5,6 +5,11 @@ See :func:`sympy.codegen.rewriting.optimize`. """ +import itertools + +from pystencils import Assignment +from pystencils.astnodes import SympyAssignment + try: from sympy.codegen.rewriting import optims_c99, optimize from sympy.codegen.rewriting import ReplaceOptim @@ -20,14 +25,18 @@ try: optims_pystencils_gpu = [evaluate_constant_terms] + list(optims_c99) except ImportError: optims_c99 = [] + optims_pystencils_cpu = [] + optims_pystencils_gpu = [] + HAS_REWRITING = False -# Evaluates all constant terms -evaluate_constant_terms = ReplaceOptim( - lambda e: e.is_constant, - lambda p: p.evalf() -) +def optimize_assignments(assignments, optimizations): + if HAS_REWRITING: + assignments = [Assignment(a.lhs, optimize(a.rhs, optimizations)) + if hasattr(a, 'lhs') else a for a in assignments] + assignments_nodes = [a.atoms(SympyAssignment) for a in assignments] + for a in itertools.chain.from_iterable(assignments_nodes): + a.optimize(optimizations) -optims_pystencils_cpu = [evaluate_constant_terms] + list(optims_c99) -optims_pystencils_gpu = [evaluate_constant_terms] + list(optims_c99) + return assignments diff --git a/pystencils_tests/test_sympy_optimizations.py b/pystencils_tests/test_sympy_optimizations.py index 7673accb..3e5fd718 100644 --- a/pystencils_tests/test_sympy_optimizations.py +++ b/pystencils_tests/test_sympy_optimizations.py @@ -1,8 +1,11 @@ +import pytest import sympy as sp import pystencils +from pystencils.optimizations import HAS_REWRITING +@pytest.mark.skipif(not HAS_REWRITING, reason="need sympy.codegen.rewriting") def test_sympy_optimizations(): for target in ('cpu', 'gpu'): x, y, z = pystencils.fields('x, y, z: float32[2d]') @@ -17,6 +20,7 @@ def test_sympy_optimizations(): assert 'expm1(' in code +@pytest.mark.skipif(not HAS_REWRITING, reason="need sympy.codegen.rewriting") def test_evaluate_constant_terms(): for target in ('cpu', 'gpu'): x, y, z = pystencils.fields('x, y, z: float32[2d]') @@ -32,6 +36,7 @@ def test_evaluate_constant_terms(): print(code) +@pytest.mark.skipif(not HAS_REWRITING, reason="need sympy.codegen.rewriting") def test_do_not_evaluate_constant_terms(): optimizations = pystencils.optimizations.optims_pystencils_cpu optimizations.remove(pystencils.optimizations.evaluate_constant_terms) -- GitLab