diff --git a/src/pystencils/__init__.py b/src/pystencils/__init__.py
index 0003a8b9a990e462ab64384f749b20ed7877127e..161530f884fe155660ca352d072135a8af05c93d 100644
--- a/src/pystencils/__init__.py
+++ b/src/pystencils/__init__.py
@@ -2,18 +2,17 @@
from .enums import Backend, Target
from . import fd
from . import stencil as stencil
-from .assignment import Assignment, AddAugmentedAssignment, assignment_from_stencil
-from .typing.typed_sympy import TypedSymbol
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment, AddAugmentedAssignment, assignment_from_stencil
from .display_utils import get_code_obj, get_code_str, show_code, to_dot
from .field import Field, FieldType, fields
from .config import CreateKernelConfig
from .cache import clear_cache
from .kernel_decorator import kernel, kernel_config
-from .kernelcreation import create_kernel, create_staggered_kernel
-from .simp import AssignmentCollection
+from .kernelcreation import create_kernel
+from pystencils.sympyextensions.assignmentcollection import AssignmentCollection
from .slicing import make_slice
from .spatial_coordinates import x_, x_staggered, x_staggered_vector, x_vector, y_, y_staggered, z_, z_staggered
-from .sympyextensions import SymbolCreator
+from .sympyextensions.math import SymbolCreator
from .datahandling import create_data_handling
__all__ = ['Field', 'FieldType', 'fields',
diff --git a/src/pystencils/nbackend/__init__.py b/src/pystencils/backend/__init__.py
similarity index 100%
rename from src/pystencils/nbackend/__init__.py
rename to src/pystencils/backend/__init__.py
diff --git a/src/pystencils/nbackend/arrays.py b/src/pystencils/backend/arrays.py
similarity index 100%
rename from src/pystencils/nbackend/arrays.py
rename to src/pystencils/backend/arrays.py
diff --git a/src/pystencils/nbackend/ast/__init__.py b/src/pystencils/backend/ast/__init__.py
similarity index 100%
rename from src/pystencils/nbackend/ast/__init__.py
rename to src/pystencils/backend/ast/__init__.py
diff --git a/src/pystencils/nbackend/ast/collectors.py b/src/pystencils/backend/ast/collectors.py
similarity index 100%
rename from src/pystencils/nbackend/ast/collectors.py
rename to src/pystencils/backend/ast/collectors.py
diff --git a/src/pystencils/nbackend/ast/dispatcher.py b/src/pystencils/backend/ast/dispatcher.py
similarity index 99%
rename from src/pystencils/nbackend/ast/dispatcher.py
rename to src/pystencils/backend/ast/dispatcher.py
index e38485cb068c67dd01fb1a5958a9c689eab60b53..f5a23b4c7778edc270b202af67df4581ca1c643c 100644
--- a/src/pystencils/nbackend/ast/dispatcher.py
+++ b/src/pystencils/backend/ast/dispatcher.py
@@ -1,9 +1,10 @@
from __future__ import annotations
-from typing import Callable
-from types import MethodType
from functools import wraps
+from typing import Callable
+from types import MethodType
+
from .nodes import PsAstNode
diff --git a/src/pystencils/nbackend/ast/kernelfunction.py b/src/pystencils/backend/ast/kernelfunction.py
similarity index 100%
rename from src/pystencils/nbackend/ast/kernelfunction.py
rename to src/pystencils/backend/ast/kernelfunction.py
diff --git a/src/pystencils/nbackend/ast/nodes.py b/src/pystencils/backend/ast/nodes.py
similarity index 100%
rename from src/pystencils/nbackend/ast/nodes.py
rename to src/pystencils/backend/ast/nodes.py
diff --git a/src/pystencils/nbackend/ast/transformations.py b/src/pystencils/backend/ast/transformations.py
similarity index 100%
rename from src/pystencils/nbackend/ast/transformations.py
rename to src/pystencils/backend/ast/transformations.py
diff --git a/src/pystencils/nbackend/ast/tree_iteration.py b/src/pystencils/backend/ast/tree_iteration.py
similarity index 100%
rename from src/pystencils/nbackend/ast/tree_iteration.py
rename to src/pystencils/backend/ast/tree_iteration.py
diff --git a/src/pystencils/nbackend/ast/util.py b/src/pystencils/backend/ast/util.py
similarity index 100%
rename from src/pystencils/nbackend/ast/util.py
rename to src/pystencils/backend/ast/util.py
diff --git a/src/pystencils/nbackend/constraints.py b/src/pystencils/backend/constraints.py
similarity index 100%
rename from src/pystencils/nbackend/constraints.py
rename to src/pystencils/backend/constraints.py
diff --git a/src/pystencils/nbackend/emission.py b/src/pystencils/backend/emission.py
similarity index 100%
rename from src/pystencils/nbackend/emission.py
rename to src/pystencils/backend/emission.py
diff --git a/src/pystencils/nbackend/exceptions.py b/src/pystencils/backend/exceptions.py
similarity index 100%
rename from src/pystencils/nbackend/exceptions.py
rename to src/pystencils/backend/exceptions.py
diff --git a/src/pystencils/nbackend/functions.py b/src/pystencils/backend/functions.py
similarity index 100%
rename from src/pystencils/nbackend/functions.py
rename to src/pystencils/backend/functions.py
diff --git a/src/pystencils/nbackend/jit/__init__.py b/src/pystencils/backend/jit/__init__.py
similarity index 100%
rename from src/pystencils/nbackend/jit/__init__.py
rename to src/pystencils/backend/jit/__init__.py
diff --git a/src/pystencils/nbackend/jit/cpu_extension_module.py b/src/pystencils/backend/jit/cpu_extension_module.py
similarity index 100%
rename from src/pystencils/nbackend/jit/cpu_extension_module.py
rename to src/pystencils/backend/jit/cpu_extension_module.py
diff --git a/src/pystencils/nbackend/jit/jit.py b/src/pystencils/backend/jit/jit.py
similarity index 100%
rename from src/pystencils/nbackend/jit/jit.py
rename to src/pystencils/backend/jit/jit.py
diff --git a/src/pystencils/nbackend/kernelcreation/__init__.py b/src/pystencils/backend/kernelcreation/__init__.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/__init__.py
rename to src/pystencils/backend/kernelcreation/__init__.py
diff --git a/src/pystencils/nbackend/kernelcreation/analysis.py b/src/pystencils/backend/kernelcreation/analysis.py
similarity index 81%
rename from src/pystencils/nbackend/kernelcreation/analysis.py
rename to src/pystencils/backend/kernelcreation/analysis.py
index 98f4887f74da3b9c4e56fd34a4d6841a30c0998a..1b498427d4a30a71ecef9882e636b9adbdb467b7 100644
--- a/src/pystencils/nbackend/kernelcreation/analysis.py
+++ b/src/pystencils/backend/kernelcreation/analysis.py
@@ -9,9 +9,8 @@ import sympy as sp
from .context import KernelCreationContext
from ...field import Field
-from ...assignment import Assignment
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment
from ...simp import AssignmentCollection
-from ...transformations import NestedScopes
from ..exceptions import PsInternalCompilerError, KernelConstraintsError
@@ -156,3 +155,57 @@ class KernelAnalysis:
rec(arg)
rec(rhs)
+
+
+class NestedScopes:
+ """Symbol visibility model using nested scopes
+
+ - every accessed symbol that was not defined before, is added as a "free parameter"
+ - free parameters are global, i.e. they are not in scopes
+ - push/pop adds or removes a scope
+
+ >>> s = NestedScopes()
+ >>> s.access_symbol("a")
+ >>> s.is_defined("a")
+ False
+ >>> s.free_parameters
+ {'a'}
+ >>> s.define_symbol("b")
+ >>> s.is_defined("b")
+ True
+ >>> s.push()
+ >>> s.is_defined_locally("b")
+ False
+ >>> s.define_symbol("c")
+ >>> s.pop()
+ >>> s.is_defined("c")
+ False
+ """
+
+ def __init__(self):
+ self.free_parameters = set()
+ self._defined = [set()]
+
+ def access_symbol(self, symbol):
+ if not self.is_defined(symbol):
+ self.free_parameters.add(symbol)
+
+ def define_symbol(self, symbol):
+ self._defined[-1].add(symbol)
+
+ def is_defined(self, symbol):
+ return any(symbol in scopes for scopes in self._defined)
+
+ def is_defined_locally(self, symbol):
+ return symbol in self._defined[-1]
+
+ def push(self):
+ self._defined.append(set())
+
+ def pop(self):
+ self._defined.pop()
+ assert self.depth >= 1
+
+ @property
+ def depth(self):
+ return len(self._defined)
\ No newline at end of file
diff --git a/src/pystencils/nbackend/kernelcreation/config.py b/src/pystencils/backend/kernelcreation/config.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/config.py
rename to src/pystencils/backend/kernelcreation/config.py
diff --git a/src/pystencils/nbackend/kernelcreation/context.py b/src/pystencils/backend/kernelcreation/context.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/context.py
rename to src/pystencils/backend/kernelcreation/context.py
diff --git a/src/pystencils/nbackend/kernelcreation/defaults.py b/src/pystencils/backend/kernelcreation/defaults.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/defaults.py
rename to src/pystencils/backend/kernelcreation/defaults.py
diff --git a/src/pystencils/nbackend/kernelcreation/freeze.py b/src/pystencils/backend/kernelcreation/freeze.py
similarity index 98%
rename from src/pystencils/nbackend/kernelcreation/freeze.py
rename to src/pystencils/backend/kernelcreation/freeze.py
index 94ecc63315349ca694b1035a11eb6ea92d42431d..06d6f1adca91cf1232fd7112360009219f6f68a5 100644
--- a/src/pystencils/nbackend/kernelcreation/freeze.py
+++ b/src/pystencils/backend/kernelcreation/freeze.py
@@ -4,7 +4,7 @@ import sympy as sp
import pymbolic.primitives as pb
from pymbolic.interop.sympy import SympyToPymbolicMapper
-from ...assignment import Assignment
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment
from ...simp import AssignmentCollection
from ...field import Field, FieldType
from ...typing import BasicType
diff --git a/src/pystencils/nbackend/kernelcreation/iteration_space.py b/src/pystencils/backend/kernelcreation/iteration_space.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/iteration_space.py
rename to src/pystencils/backend/kernelcreation/iteration_space.py
diff --git a/src/pystencils/nbackend/kernelcreation/transformations.py b/src/pystencils/backend/kernelcreation/transformations.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/transformations.py
rename to src/pystencils/backend/kernelcreation/transformations.py
diff --git a/src/pystencils/nbackend/kernelcreation/typification.py b/src/pystencils/backend/kernelcreation/typification.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/typification.py
rename to src/pystencils/backend/kernelcreation/typification.py
diff --git a/src/pystencils/nbackend/kernelcreation/platform/__init__.py b/src/pystencils/backend/platforms/__init__.py
similarity index 100%
rename from src/pystencils/nbackend/kernelcreation/platform/__init__.py
rename to src/pystencils/backend/platforms/__init__.py
diff --git a/src/pystencils/nbackend/kernelcreation/platform/basic_cpu.py b/src/pystencils/backend/platforms/basic_cpu.py
similarity index 90%
rename from src/pystencils/nbackend/kernelcreation/platform/basic_cpu.py
rename to src/pystencils/backend/platforms/basic_cpu.py
index f5deaf0d65a0a19d568c7dead77bace5efbac517..27d32818a3393e77f856eb1da54c504b35f2412e 100644
--- a/src/pystencils/nbackend/kernelcreation/platform/basic_cpu.py
+++ b/src/pystencils/backend/platforms/basic_cpu.py
@@ -1,14 +1,14 @@
from .platform import Platform
-from ..iteration_space import (
+from ..kernelcreation.iteration_space import (
IterationSpace,
FullIterationSpace,
SparseIterationSpace,
)
-from ...ast import PsDeclaration, PsSymbolExpr, PsExpression, PsLoop, PsBlock
-from ...typed_expressions import PsTypedConstant
-from ...arrays import PsArrayAccess
+from ..ast import PsDeclaration, PsSymbolExpr, PsExpression, PsLoop, PsBlock
+from ..typed_expressions import PsTypedConstant
+from ..arrays import PsArrayAccess
class BasicCpu(Platform):
diff --git a/src/pystencils/nbackend/kernelcreation/platform/platform.py b/src/pystencils/backend/platforms/platform.py
similarity index 82%
rename from src/pystencils/nbackend/kernelcreation/platform/platform.py
rename to src/pystencils/backend/platforms/platform.py
index b2c7f899ea0842bf77e4d4cb759a4b5a9b3493b6..3abeebbe69599cd0181b061ffeeef96eed62dca7 100644
--- a/src/pystencils/nbackend/kernelcreation/platform/platform.py
+++ b/src/pystencils/backend/platforms/platform.py
@@ -1,9 +1,9 @@
from abc import ABC, abstractmethod
-from ...ast import PsBlock
+from ..ast import PsBlock
-from ..context import KernelCreationContext
-from ..iteration_space import IterationSpace
+from ..kernelcreation.context import KernelCreationContext
+from ..kernelcreation.iteration_space import IterationSpace
class Platform(ABC):
diff --git a/src/pystencils/nbackend/typed_expressions.py b/src/pystencils/backend/typed_expressions.py
similarity index 100%
rename from src/pystencils/nbackend/typed_expressions.py
rename to src/pystencils/backend/typed_expressions.py
diff --git a/src/pystencils/nbackend/types/__init__.py b/src/pystencils/backend/types/__init__.py
similarity index 100%
rename from src/pystencils/nbackend/types/__init__.py
rename to src/pystencils/backend/types/__init__.py
diff --git a/src/pystencils/nbackend/types/basic_types.py b/src/pystencils/backend/types/basic_types.py
similarity index 100%
rename from src/pystencils/nbackend/types/basic_types.py
rename to src/pystencils/backend/types/basic_types.py
diff --git a/src/pystencils/nbackend/types/exception.py b/src/pystencils/backend/types/exception.py
similarity index 100%
rename from src/pystencils/nbackend/types/exception.py
rename to src/pystencils/backend/types/exception.py
diff --git a/src/pystencils/nbackend/types/parsing.py b/src/pystencils/backend/types/parsing.py
similarity index 100%
rename from src/pystencils/nbackend/types/parsing.py
rename to src/pystencils/backend/types/parsing.py
diff --git a/src/pystencils/nbackend/types/quick.py b/src/pystencils/backend/types/quick.py
similarity index 100%
rename from src/pystencils/nbackend/types/quick.py
rename to src/pystencils/backend/types/quick.py
diff --git a/src/pystencils/boundaries/boundaryconditions.py b/src/pystencils/boundaries/boundaryconditions.py
index 5fd8480b65539994bfae3e2c472a81fd4fa86f51..9aa4319fb3cb596da4f97f7b4a4ea57fda205670 100644
--- a/src/pystencils/boundaries/boundaryconditions.py
+++ b/src/pystencils/boundaries/boundaryconditions.py
@@ -1,6 +1,6 @@
from typing import Any, List, Tuple
-from pystencils.astnodes import SympyAssignment
+from pystencils.sympyextensions.astnodes import SympyAssignment
from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
from pystencils.typing import create_type
diff --git a/src/pystencils/boundaries/boundaryhandling.py b/src/pystencils/boundaries/boundaryhandling.py
index 53c3980e2d2c9c947d45410d03017da135d4a843..96c82e75dacfc8d510b1d953127b971fc1ead9ac 100644
--- a/src/pystencils/boundaries/boundaryhandling.py
+++ b/src/pystencils/boundaries/boundaryhandling.py
@@ -4,7 +4,7 @@ import numpy as np
import sympy as sp
from pystencils import create_kernel, CreateKernelConfig, Target
-from pystencils.astnodes import SympyAssignment
+from pystencils.sympyextensions.astnodes import SympyAssignment
from pystencils.backends.cbackend import CustomCodeNode
from pystencils.boundaries.createindexlist import (
create_boundary_index_array, numpy_data_type_for_boundary_object)
diff --git a/src/pystencils/display_utils.py b/src/pystencils/display_utils.py
index f6c32ac88ae68d684e860b33c0e5185ccc030e4e..2fba12da606c975ed7704d111421cc200b8d79d4 100644
--- a/src/pystencils/display_utils.py
+++ b/src/pystencils/display_utils.py
@@ -2,14 +2,14 @@ from typing import Any, Dict, Optional, Union
import sympy as sp
-from pystencils.astnodes import KernelFunction
+from pystencils.sympyextensions.astnodes import KernelFunction
from pystencils.enums import Backend
from pystencils.kernel_wrapper import KernelWrapper
def to_dot(expr: sp.Expr, graph_style: Optional[Dict[str, Any]] = None, short=True):
"""Show a sympy or pystencils AST as dot graph"""
- from pystencils.astnodes import Node
+ from pystencils.sympyextensions.astnodes import Node
try:
import graphviz
except ImportError:
diff --git a/src/pystencils/fast_approximation.py b/src/pystencils/fast_approximation.py
index ab0dc59740e9ec7fcd3e59eb826979cd5350aa3f..1cbbf52f99f8661f7c8727c5176264b3066478ac 100644
--- a/src/pystencils/fast_approximation.py
+++ b/src/pystencils/fast_approximation.py
@@ -2,9 +2,9 @@ from typing import List, Union
import sympy as sp
-from pystencils.astnodes import Node
+from pystencils.sympyextensions.astnodes import Node
from pystencils.simp import AssignmentCollection
-from pystencils.assignment import Assignment
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment
# noinspection PyPep8Naming
diff --git a/src/pystencils/fd/derivative.py b/src/pystencils/fd/derivative.py
index c119d1e2ec34c32c67f18f7837c43dee05cfc65b..2396fe1a988acdfa39420d75ffd0557c97fd15fd 100644
--- a/src/pystencils/fd/derivative.py
+++ b/src/pystencils/fd/derivative.py
@@ -3,7 +3,7 @@ from collections import defaultdict, namedtuple
import sympy as sp
from pystencils.field import Field
-from pystencils.sympyextensions import normalize_product, prod
+from pystencils.sympyextensions.math import normalize_product, prod
def _default_diff_sort_key(d):
diff --git a/src/pystencils/fd/spatial.py b/src/pystencils/fd/spatial.py
index 387a03bac6c0c6f88b92851bc18b7d752d64b036..17f708ed8f6e835294f5e898e6f2070b933e13ce 100644
--- a/src/pystencils/fd/spatial.py
+++ b/src/pystencils/fd/spatial.py
@@ -3,10 +3,10 @@ from typing import Tuple
import sympy as sp
-from pystencils.astnodes import LoopOverCoordinate
+from pystencils.sympyextensions.astnodes import LoopOverCoordinate
from pystencils.fd import Diff
from pystencils.field import Field
-from pystencils.transformations import generic_visit
+from pystencils.sympyextensions import generic_visit
from .derivation import FiniteDifferenceStencilDerivation
from .derivative import diff_args
diff --git a/src/pystencils/field.py b/src/pystencils/field.py
index 7630f4cd8a50d44fc21f24164eb9afcc8c90bbc0..389413f2530a791ba923e5eac5cff5f0b4352982 100644
--- a/src/pystencils/field.py
+++ b/src/pystencils/field.py
@@ -11,13 +11,13 @@ import numpy as np
import sympy as sp
from sympy.core.cache import cacheit
-import pystencils
from pystencils.alignedarray import aligned_empty
-from pystencils.typing import StructType, TypedSymbol, BasicType, create_type
-from pystencils.typing.typed_sympy import FieldShapeSymbol, FieldStrideSymbol
-from pystencils.stencil import (
- direction_string_to_offset, inverse_direction, offset_to_direction_string)
-from pystencils.sympyextensions import is_integer_sequence
+from pystencils.spatial_coordinates import x_staggered_vector, x_vector
+from pystencils.stencil import direction_string_to_offset, inverse_direction, offset_to_direction_string
+from pystencils.sympyextensions.typed_sympy import (FieldShapeSymbol, FieldStrideSymbol, StructType,
+ TypedSymbol, BasicType, create_type)
+from pystencils.sympyextensions.math import is_integer_sequence
+
__all__ = ['Field', 'fields', 'FieldType', 'Field']
@@ -510,14 +510,14 @@ class Field:
@property
def physical_coordinates(self):
if hasattr(self.coordinate_transform, '__call__'):
- return self.coordinate_transform(self.coordinate_origin + pystencils.x_vector(self.spatial_dimensions))
+ return self.coordinate_transform(self.coordinate_origin + x_vector(self.spatial_dimensions))
else:
- return self.coordinate_transform @ (self.coordinate_origin + pystencils.x_vector(self.spatial_dimensions))
+ return self.coordinate_transform @ (self.coordinate_origin + x_vector(self.spatial_dimensions))
@property
def physical_coordinates_staggered(self):
return self.coordinate_transform @ \
- (self.coordinate_origin + pystencils.x_staggered_vector(self.spatial_dimensions))
+ (self.coordinate_origin + x_staggered_vector(self.spatial_dimensions))
def index_to_physical(self, index_coordinates: sp.Matrix, staggered=False):
if staggered:
diff --git a/src/pystencils/functions.py b/src/pystencils/functions.py
index 722c2c5d410dd81968cc593871c6714bbbba1645..550a2724f1f3366d8fe3bb583f177d62a2616d22 100644
--- a/src/pystencils/functions.py
+++ b/src/pystencils/functions.py
@@ -1,5 +1,5 @@
import sympy as sp
-from pystencils.typing import PointerType
+from .sympyextensions.typed_sympy import PointerType
class DivFunc(sp.Function):
diff --git a/src/pystencils/integer_set_analysis.py b/src/pystencils/integer_set_analysis.py
index 00fc1cb960c6fc1fd718bdc669df607285af8749..e170c076d43ffd9f03aa2223e5712fef30c46011 100644
--- a/src/pystencils/integer_set_analysis.py
+++ b/src/pystencils/integer_set_analysis.py
@@ -3,7 +3,7 @@
import islpy as isl
import sympy as sp
-import pystencils.astnodes as ast
+import pystencils.sympyextensions.astnodes as ast
from pystencils.typing import parents_of_type
from pystencils.backends.cbackend import CustomSympyPrinter
diff --git a/src/pystencils/kernel_contrains_check.py b/src/pystencils/kernel_contrains_check.py
index f1fa4b8a141400c0880672f4fdbcd356b59d4ccd..319443cd7b2a243a8b03c813905631dcc43c6133 100644
--- a/src/pystencils/kernel_contrains_check.py
+++ b/src/pystencils/kernel_contrains_check.py
@@ -5,7 +5,8 @@ import sympy as sp
from sympy.codegen import Assignment
from pystencils.simp import AssignmentCollection
-from pystencils import astnodes as ast, TypedSymbol
+from pystencils import TypedSymbol
+from pystencils.sympyextensions import astnodes as ast
from pystencils.field import Field
from pystencils.node_collection import NodeCollection
from pystencils.transformations import NestedScopes
diff --git a/src/pystencils/kernel_decorator.py b/src/pystencils/kernel_decorator.py
index ad5d625929058ef383402e2a1d97c0dfadbf5fda..a8246de07ab43b1de695f9bb2a872360f9b891b8 100644
--- a/src/pystencils/kernel_decorator.py
+++ b/src/pystencils/kernel_decorator.py
@@ -5,7 +5,7 @@ from typing import Callable, Union, List, Dict, Tuple
import sympy as sp
-from pystencils.assignment import Assignment
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment
from pystencils.sympyextensions import SymbolCreator
from pystencils.config import CreateKernelConfig
diff --git a/src/pystencils/kernelcreation.py b/src/pystencils/kernelcreation.py
index 385f42d2f01c26b633fcc25799665cab2072dceb..52c70fbfbb6b69505681a60a1c78c5d0dd4ec8d2 100644
--- a/src/pystencils/kernelcreation.py
+++ b/src/pystencils/kernelcreation.py
@@ -1,393 +1,61 @@
-import itertools
-import warnings
-from typing import Union, List
-
-import sympy as sp
-from pystencils.config import CreateKernelConfig
-
-from pystencils.assignment import Assignment, AddAugmentedAssignment
-from pystencils.astnodes import Node, Block, Conditional, LoopOverCoordinate, SympyAssignment
-from pystencils.cpu.vectorization import vectorize
-from pystencils.enums import Target, Backend
-from pystencils.field import Field, FieldType
-from pystencils.node_collection import NodeCollection
-from pystencils.simp.assignment_collection import AssignmentCollection
-from pystencils.kernel_contrains_check import KernelConstraintsCheck
-from pystencils.simplificationfactory import create_simplification_strategy
-from pystencils.stencil import direction_string_to_offset, inverse_direction_string
-from pystencils.transformations import (
- loop_blocking, move_constants_before_loop, remove_conditionals_in_staggered_kernel)
-
-
-def create_kernel(assignments: Union[Assignment, List[Assignment],
- AddAugmentedAssignment, List[AddAugmentedAssignment],
- AssignmentCollection, List[Node], NodeCollection],
- *,
- config: CreateKernelConfig = None, **kwargs):
- """
- Creates abstract syntax tree (AST) of kernel, using a list of update equations.
- This function forms the general API and delegates the kernel creation to others depending on the CreateKernelConfig.
- Args:
- assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection`
- config: CreateKernelConfig which includes the needed configuration
- kwargs: Arguments for updating the config
-
- Returns:
- abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
- can be compiled with through its 'compile()' member
-
- Example:
- >>> import pystencils as ps
- >>> import numpy as np
- >>> s, d = ps.fields('s, d: [2D]')
- >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
- >>> kernel_ast = ps.create_kernel(assignment, config=ps.CreateKernelConfig(cpu_openmp=True))
- >>> kernel = kernel_ast.compile()
- >>> d_arr = np.zeros([5, 5])
- >>> kernel(d=d_arr, s=np.ones([5, 5]))
- >>> d_arr
- array([[0., 0., 0., 0., 0.],
- [0., 4., 4., 4., 0.],
- [0., 4., 4., 4., 0.],
- [0., 4., 4., 4., 0.],
- [0., 0., 0., 0., 0.]])
- """
- # ---- Updating configuration from kwargs
- if not config:
- config = CreateKernelConfig(**kwargs)
- else:
- for k, v in kwargs.items():
- if not hasattr(config, k):
- raise KeyError(f'{v} is not a valid kwarg. Please look in CreateKernelConfig for valid settings')
- setattr(config, k, v)
-
- # ---- Normalizing parameters
- if isinstance(assignments, (Assignment, AddAugmentedAssignment)):
- assignments = [assignments]
- assert assignments, "Assignments must not be empty!"
- if isinstance(assignments, list):
- assignments = NodeCollection(assignments)
- elif isinstance(assignments, AssignmentCollection):
- # TODO Markus check and doku
- # --- applying first default simplifications
- try:
- if config.default_assignment_simplifications:
- simplification = create_simplification_strategy()
- assignments = simplification(assignments)
- 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.from_assignment_collection(assignments)
- assignments.simplification_hints = simplification_hints
-
- if config.index_fields:
- return create_indexed_kernel(assignments, config=config)
- else:
- return create_domain_kernel(assignments, config=config)
-
-
-def create_domain_kernel(assignments: NodeCollection, *, config: CreateKernelConfig):
- """
- Creates abstract syntax tree (AST) of kernel, using a NodeCollection.
-
- Note that `create_domain_kernel` is a lower level function which shoul be accessed by not providing `index_fields`
- to create_kernel
-
- Args:
- assignments: `pystencils.node_collection.NodeCollection` containing all assignements and nodes to be processed
- config: CreateKernelConfig which includes the needed configuration
-
- Returns:
- abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
- can be compiled with through its 'compile()' member
-
- Example:
- >>> import pystencils as ps
- >>> import numpy as np
- >>> from pystencils.kernelcreation import create_domain_kernel
- >>> from pystencils.node_collection import NodeCollection
- >>> s, d = ps.fields('s, d: [2D]')
- >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
- >>> kernel_config = ps.CreateKernelConfig(cpu_openmp=True)
- >>> kernel_ast = create_domain_kernel(NodeCollection([assignment]), config=kernel_config)
- >>> kernel = kernel_ast.compile()
- >>> d_arr = np.zeros([5, 5])
- >>> kernel(d=d_arr, s=np.ones([5, 5]))
- >>> d_arr
- array([[0., 0., 0., 0., 0.],
- [0., 4., 4., 4., 0.],
- [0., 4., 4., 4., 0.],
- [0., 4., 4., 4., 0.],
- [0., 0., 0., 0., 0.]])
- """
- # --- eval
- assignments.evaluate_terms()
-
- # FUTURE WORK from here we shouldn't NEED sympy
- # --- check constrains
- check = KernelConstraintsCheck(check_independence_condition=not config.skip_independence_check,
- check_double_write_condition=not config.allow_double_writes)
- check.visit(assignments)
-
- assignments.bound_fields = check.fields_written
- assignments.rhs_fields = check.fields_read
-
- # ---- Creating ast
- ast = None
- if config.target == Target.CPU:
- if config.backend == Backend.C:
- from pystencils.cpu import add_openmp, create_kernel
- ast = create_kernel(assignments, config=config)
- for optimization in config.cpu_prepend_optimizations:
- optimization(ast)
- omp_collapse = None
- if config.cpu_blocking:
- omp_collapse = loop_blocking(ast, config.cpu_blocking)
- if config.cpu_openmp:
- add_openmp(ast, num_threads=config.cpu_openmp, collapse=omp_collapse,
- assume_single_outer_loop=config.omp_single_loop)
- if config.cpu_vectorize_info:
- if config.cpu_vectorize_info is True:
- vectorize(ast)
- elif isinstance(config.cpu_vectorize_info, dict):
- vectorize(ast, **config.cpu_vectorize_info)
- if config.cpu_openmp and config.cpu_blocking and 'nontemporal' in config.cpu_vectorize_info and \
- config.cpu_vectorize_info['nontemporal'] and 'cachelineZero' in ast.instruction_set:
- # This condition is stricter than it needs to be: if blocks along the fastest axis start on a
- # cache line boundary, it's okay. But we cannot determine that here.
- # We don't need to disallow OpenMP collapsing because it is never applied to the inner loop.
- raise ValueError("Blocking cannot be combined with cacheline-zeroing")
- else:
- raise ValueError("Invalid value for cpu_vectorize_info")
- elif config.target == Target.GPU:
- if config.backend == Backend.CUDA:
- from pystencils.gpu import create_cuda_kernel
- ast = create_cuda_kernel(assignments, config=config)
-
- if not ast:
- raise NotImplementedError(
- f'{config.target} together with {config.backend} is not supported by `create_domain_kernel`')
-
- if config.use_auto_for_assignments:
- for a in ast.atoms(SympyAssignment):
- a.use_auto = True
-
- return ast
-
-
-def create_indexed_kernel(assignments: NodeCollection, *, config: CreateKernelConfig):
- """
- Similar to :func:`create_kernel`, but here not all cells of a field are updated but only cells with
- coordinates which are stored in an index field. This traversal method can e.g. be used for boundary handling.
-
- The coordinates are stored in a separated index_field, which is a one dimensional array with struct data type.
- This struct has to contain fields named 'x', 'y' and for 3D fields ('z'). These names are configurable with the
- '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: `pystencils.node_collection.NodeCollection` containing all assignements and nodes to be processed
- config: CreateKernelConfig which includes the needed configuration
-
- Returns:
- abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
- can be compiled with through its 'compile()' member
-
- Example:
- >>> import pystencils as ps
- >>> from pystencils.node_collection import NodeCollection
- >>> import numpy as np
- >>> from pystencils.kernelcreation import create_indexed_kernel
- >>>
- >>> # Index field stores the indices of the cell to visit together with optional values
- >>> index_arr_dtype = np.dtype([('x', np.int32), ('y', np.int32), ('val', np.double)])
- >>> index_arr = np.array([(1, 1, 0.1), (2, 2, 0.2), (3, 3, 0.3)], dtype=index_arr_dtype)
- >>> idx_field = ps.fields(idx=index_arr)
- >>>
- >>> # Additional values stored in index field can be accessed in the kernel as well
- >>> 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(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. ]])
-
- """
- # --- eval
- assignments.evaluate_terms()
-
- # FUTURE WORK from here we shouldn't NEED sympy
- # --- check constrains
- check = KernelConstraintsCheck(check_independence_condition=not config.skip_independence_check,
- check_double_write_condition=not config.allow_double_writes)
- check.visit(assignments)
-
- assignments.bound_fields = check.fields_written
- assignments.rhs_fields = check.fields_read
-
- ast = None
- if config.target == Target.CPU and config.backend == Backend.C:
- from pystencils.cpu import add_openmp, create_indexed_kernel
- ast = create_indexed_kernel(assignments, config=config)
- if config.cpu_openmp:
- add_openmp(ast, num_threads=config.cpu_openmp)
- elif config.target == Target.GPU:
- if config.backend == Backend.CUDA:
- from pystencils.gpu import created_indexed_cuda_kernel
- ast = created_indexed_cuda_kernel(assignments, config=config)
-
- if not ast:
- raise NotImplementedError(f'Indexed kernels are not yet supported for {config.target} with {config.backend}')
- return ast
-
-
-def create_staggered_kernel(assignments, target: Target = Target.CPU, gpu_exclusive_conditions=False, **kwargs):
- """Kernel that updates a staggered field.
-
- .. image:: /img/staggered_grid.svg
-
- For a staggered field, the first index coordinate defines the location of the staggered value.
- Further index coordinates can be used to store vectors/tensors at each point.
-
- Args:
- assignments: a sequence of assignments or an AssignmentCollection.
- Assignments to staggered field are processed specially, while subexpressions and assignments to
- regular fields are passed through to `create_kernel`. Multiple different staggered fields can be
- used, but they all need to use the same stencil (i.e. the same number of staggered points) and
- shape.
- target: 'CPU' or 'GPU'
- gpu_exclusive_conditions: disable the use of multiple conditionals inside the loop. The outer layers are then
- handled in an else branch.
- kwargs: passed directly to create_kernel, iteration_slice and ghost_layers parameters are not allowed
-
- Returns:
- AST, see `create_kernel`
- """
- # TODO: Add doku like in the other kernels
- if 'ghost_layers' in kwargs:
- assert kwargs['ghost_layers'] is None
- del kwargs['ghost_layers']
- if 'iteration_slice' in kwargs:
- assert kwargs['iteration_slice'] is None
- del kwargs['iteration_slice']
- if 'omp_single_loop' in kwargs:
- assert kwargs['omp_single_loop'] is False
- del kwargs['omp_single_loop']
-
- if isinstance(assignments, AssignmentCollection):
- subexpressions = assignments.subexpressions + [a for a in assignments.main_assignments
- if not hasattr(a, 'lhs')
- or type(a.lhs) is not Field.Access
- or not FieldType.is_staggered(a.lhs.field)]
- assignments = [a for a in assignments.main_assignments if hasattr(a, 'lhs')
- and type(a.lhs) is Field.Access
- and FieldType.is_staggered(a.lhs.field)]
+from .backend.ast import PsKernelFunction
+from .backend.kernelcreation import KernelCreationContext, KernelAnalysis, FreezeExpressions, Typifier
+from .backend.kernelcreation.iteration_space import (
+ create_sparse_iteration_space,
+ create_full_iteration_space,
+)
+from .backend.kernelcreation.transformations import EraseAnonymousStructTypes
+
+from .enums import Target
+from .config import CreateKernelConfig
+from pystencils.sympyextensions.assignmentcollection import AssignmentCollection
+
+
+def create_kernel(
+ assignments: AssignmentCollection,
+ config: CreateKernelConfig = CreateKernelConfig(),
+):
+ """Create a kernel AST from an assignment collection."""
+ ctx = KernelCreationContext(config)
+
+ analysis = KernelAnalysis(ctx)
+ analysis(assignments)
+
+ if len(ctx.fields.index_fields) > 0 or ctx.options.index_field is not None:
+ ispace = create_sparse_iteration_space(ctx, assignments)
else:
- subexpressions = [a for a in assignments if not hasattr(a, 'lhs')
- or type(a.lhs) is not Field.Access
- or not FieldType.is_staggered(a.lhs.field)]
- assignments = [a for a in assignments if hasattr(a, 'lhs')
- and type(a.lhs) is Field.Access
- and FieldType.is_staggered(a.lhs.field)]
- if len(set([tuple(a.lhs.field.staggered_stencil) for a in assignments])) != 1:
- raise ValueError("All assignments need to be made to staggered fields with the same stencil")
- if len(set([a.lhs.field.shape for a in assignments])) != 1:
- raise ValueError("All assignments need to be made to staggered fields with the same shape")
-
- staggered_field = assignments[0].lhs.field
- stencil = staggered_field.staggered_stencil
- dim = staggered_field.spatial_dimensions
- shape = staggered_field.shape
-
- counters = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)]
-
- final_assignments = []
+ ispace = create_full_iteration_space(ctx, assignments)
- # find out whether any of the ghost layers is not needed
- common_exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
- for direction in stencil:
- exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
- for elementary_direction in direction:
- exclusions.remove(inverse_direction_string(elementary_direction))
- common_exclusions.intersection_update(exclusions)
- ghost_layers = [[0, 0] for d in range(dim)]
- for direction in common_exclusions:
- direction = direction_string_to_offset(direction)
- for d, s in enumerate(direction):
- if s == 1:
- ghost_layers[d][1] = 1
- elif s == -1:
- ghost_layers[d][0] = 1
+ ctx.set_iteration_space(ispace)
- def condition(direction):
- """exclude those staggered points that correspond to fluxes between ghost cells"""
- exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
+ freeze = FreezeExpressions(ctx)
+ kernel_body = freeze(assignments)
- for elementary_direction in direction:
- exclusions.remove(inverse_direction_string(elementary_direction))
- conditions = []
- for e in exclusions:
- if e in common_exclusions:
- continue
- offset = direction_string_to_offset(e)
- for i, o in enumerate(offset):
- if o == 1:
- conditions.append(counters[i] < shape[i] - 1)
- elif o == -1:
- conditions.append(counters[i] > 0)
- return sp.And(*conditions)
+ typify = Typifier(ctx)
+ kernel_body = typify(kernel_body)
- if gpu_exclusive_conditions:
- outer_assignment = None
- conditions = {direction: condition(direction) for direction in stencil}
- for num_conditions in range(len(stencil)):
- for combination in itertools.combinations(conditions.values(), num_conditions):
- for assignment in assignments:
- direction = stencil[assignment.lhs.index[0]]
- if conditions[direction] in combination:
- assignment = SympyAssignment(assignment.lhs, assignment.rhs)
- outer_assignment = Conditional(sp.And(*combination), Block([assignment]), outer_assignment)
+ match config.target:
+ case Target.CPU:
+ from .backend.platforms import BasicCpu
- inner_assignment = []
- for assignment in assignments:
- inner_assignment.append(SympyAssignment(assignment.lhs, assignment.rhs))
- last_conditional = Conditional(sp.And(*[condition(d) for d in stencil]),
- Block(inner_assignment), outer_assignment)
- final_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
- [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
- [last_conditional]
+ # TODO: CPU platform should incorporate instruction set info, OpenMP, etc.
+ platform = BasicCpu(ctx)
+ case _:
+ # TODO: CUDA/HIP platform
+ # TODO: SYCL platform (?)
+ raise NotImplementedError("Target platform not implemented")
- config = CreateKernelConfig(target=target, ghost_layers=ghost_layers, omp_single_loop=False, **kwargs)
- ast = create_kernel(final_assignments, config=config)
- return ast
+ kernel_ast = platform.materialize_iteration_space(kernel_body, ispace)
+ kernel_ast = EraseAnonymousStructTypes(ctx)(kernel_ast)
- for assignment in assignments:
- direction = stencil[assignment.lhs.index[0]]
- sp_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
- [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
- [SympyAssignment(assignment.lhs, assignment.rhs)]
- last_conditional = Conditional(condition(direction), Block(sp_assignments))
- final_assignments.append(last_conditional)
+ # 7. Apply optimizations
+ # - Vectorization
+ # - OpenMP
+ # - Loop Splitting, Tiling, Blocking
+ kernel_ast = platform.optimize(kernel_ast)
- remove_start_conditional = any([gl[0] == 0 for gl in ghost_layers])
- prepend_optimizations = [lambda ast: remove_conditionals_in_staggered_kernel(ast, remove_start_conditional),
- move_constants_before_loop]
- if 'cpu_prepend_optimizations' in kwargs:
- prepend_optimizations += kwargs['cpu_prepend_optimizations']
- del kwargs['cpu_prepend_optimizations']
+ assert config.jit is not None
+ function = PsKernelFunction(kernel_ast, config.target, name=config.function_name, jit=config.jit)
+ function.add_constraints(*ctx.constraints)
- config = CreateKernelConfig(ghost_layers=ghost_layers, target=target, omp_single_loop=False,
- cpu_prepend_optimizations=prepend_optimizations, **kwargs)
- ast = create_kernel(final_assignments, config=config)
- return ast
+ return function
diff --git a/src/pystencils/nbackend/kernelcreation/kernelcreation.py b/src/pystencils/nbackend/kernelcreation/kernelcreation.py
deleted file mode 100644
index 8abf9c7ef0e81601ef15ee76ce82602de9557f63..0000000000000000000000000000000000000000
--- a/src/pystencils/nbackend/kernelcreation/kernelcreation.py
+++ /dev/null
@@ -1,65 +0,0 @@
-from ...simp import AssignmentCollection
-
-from ..ast import PsKernelFunction
-from ...enums import Target
-
-from .context import KernelCreationContext
-from .analysis import KernelAnalysis
-from .freeze import FreezeExpressions
-from .typification import Typifier
-from .config import CreateKernelConfig
-from .iteration_space import (
- create_sparse_iteration_space,
- create_full_iteration_space,
-)
-from .transformations import EraseAnonymousStructTypes
-
-
-def create_kernel(
- assignments: AssignmentCollection,
- config: CreateKernelConfig = CreateKernelConfig(),
-):
- """Create a kernel AST from an assignment collection."""
- ctx = KernelCreationContext(config)
-
- analysis = KernelAnalysis(ctx)
- analysis(assignments)
-
- if len(ctx.fields.index_fields) > 0 or ctx.options.index_field is not None:
- ispace = create_sparse_iteration_space(ctx, assignments)
- else:
- ispace = create_full_iteration_space(ctx, assignments)
-
- ctx.set_iteration_space(ispace)
-
- freeze = FreezeExpressions(ctx)
- kernel_body = freeze(assignments)
-
- typify = Typifier(ctx)
- kernel_body = typify(kernel_body)
-
- match config.target:
- case Target.CPU:
- from .platform import BasicCpu
-
- # TODO: CPU platform should incorporate instruction set info, OpenMP, etc.
- platform = BasicCpu(ctx)
- case _:
- # TODO: CUDA/HIP platform
- # TODO: SYCL platform (?)
- raise NotImplementedError("Target platform not implemented")
-
- kernel_ast = platform.materialize_iteration_space(kernel_body, ispace)
- kernel_ast = EraseAnonymousStructTypes(ctx)(kernel_ast)
-
- # 7. Apply optimizations
- # - Vectorization
- # - OpenMP
- # - Loop Splitting, Tiling, Blocking
- kernel_ast = platform.optimize(kernel_ast)
-
- assert config.jit is not None
- function = PsKernelFunction(kernel_ast, config.target, name=config.function_name, jit=config.jit)
- function.add_constraints(*ctx.constraints)
-
- return function
diff --git a/src/pystencils/node_collection.py b/src/pystencils/node_collection.py
index 61a2f400a440b05a137a660991da4e53f2af8294..2dc35a975ec5a743d56d3af426bf41c68a7af0ee 100644
--- a/src/pystencils/node_collection.py
+++ b/src/pystencils/node_collection.py
@@ -4,8 +4,8 @@ import sympy
import sympy as sp
from sympy.codegen.rewriting import ReplaceOptim, optimize
-from pystencils.assignment import Assignment, AddAugmentedAssignment
-import pystencils.astnodes as ast
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment, AddAugmentedAssignment
+import pystencils.sympyextensions.astnodes as ast
from pystencils.backends.cbackend import CustomCodeNode
from pystencils.functions import DivFunc
from pystencils.simp import AssignmentCollection
diff --git a/src/pystencils/backends/__init__.py b/src/pystencils/old/backends/__init__.py
similarity index 100%
rename from src/pystencils/backends/__init__.py
rename to src/pystencils/old/backends/__init__.py
diff --git a/src/pystencils/backends/arm_instruction_sets.py b/src/pystencils/old/backends/arm_instruction_sets.py
similarity index 100%
rename from src/pystencils/backends/arm_instruction_sets.py
rename to src/pystencils/old/backends/arm_instruction_sets.py
diff --git a/src/pystencils/backends/cbackend.py b/src/pystencils/old/backends/cbackend.py
similarity index 99%
rename from src/pystencils/backends/cbackend.py
rename to src/pystencils/old/backends/cbackend.py
index 7dbf84d378d768530cfe9c706186f7d2a684581f..04305a3e309e013b4282a2e1375f60d1de387c07 100644
--- a/src/pystencils/backends/cbackend.py
+++ b/src/pystencils/old/backends/cbackend.py
@@ -11,7 +11,7 @@ from sympy.logic.boolalg import BooleanFalse, BooleanTrue
from sympy.functions.elementary.trigonometric import TrigonometricFunction, InverseTrigonometricFunction
from sympy.functions.elementary.hyperbolic import HyperbolicFunction
-from pystencils.astnodes import KernelFunction, LoopOverCoordinate, Node
+from pystencils.sympyextensions.astnodes import KernelFunction, LoopOverCoordinate, Node
from pystencils.cpu.vectorization import vec_all, vec_any, CachelineSize
from pystencils.typing import (
PointerType, VectorType, CastFunc, create_type, get_type_of_expression,
diff --git a/src/pystencils/backends/cuda_backend.py b/src/pystencils/old/backends/cuda_backend.py
similarity index 98%
rename from src/pystencils/backends/cuda_backend.py
rename to src/pystencils/old/backends/cuda_backend.py
index f8fdb16dac8911811915d6d69e5af7af8f149f4f..64321007d67c2988e94c1793a704de41861eb65f 100644
--- a/src/pystencils/backends/cuda_backend.py
+++ b/src/pystencils/old/backends/cuda_backend.py
@@ -1,4 +1,4 @@
-from pystencils.astnodes import Node
+from pystencils.sympyextensions.astnodes import Node
from pystencils.backends.cbackend import CBackend, CustomSympyPrinter, generate_c
from pystencils.enums import Backend
from pystencils.fast_approximation import fast_division, fast_inv_sqrt, fast_sqrt
diff --git a/src/pystencils/backends/dot.py b/src/pystencils/old/backends/dot.py
similarity index 96%
rename from src/pystencils/backends/dot.py
rename to src/pystencils/old/backends/dot.py
index 3a5562b65fc20a5c9b10a6fcc8416cb49d0b0396..c21e7724f140c9a02fd70ee285f90305a89c2dd4 100644
--- a/src/pystencils/backends/dot.py
+++ b/src/pystencils/old/backends/dot.py
@@ -55,7 +55,7 @@ class DotPrinter(Printer):
def __shortened(node):
- from pystencils.astnodes import LoopOverCoordinate, KernelFunction, SympyAssignment, Conditional
+ from pystencils.sympyextensions.astnodes import LoopOverCoordinate, KernelFunction, SympyAssignment, Conditional
if isinstance(node, LoopOverCoordinate):
return "Loop over dim %d" % (node.coordinate_to_loop_over,)
elif isinstance(node, KernelFunction):
diff --git a/src/pystencils/backends/json.py b/src/pystencils/old/backends/json.py
similarity index 97%
rename from src/pystencils/backends/json.py
rename to src/pystencils/old/backends/json.py
index 954965c6e626ebc2fc63bba1457c53a325e052e7..4b51fb8088345a3793f9a59b1f070db0a157d357 100644
--- a/src/pystencils/backends/json.py
+++ b/src/pystencils/old/backends/json.py
@@ -9,7 +9,7 @@
"""
import json
-from pystencils.astnodes import NodeOrExpr
+from pystencils.sympyextensions.astnodes import NodeOrExpr
from pystencils.backends.cbackend import CustomSympyPrinter, generate_c
try:
diff --git a/src/pystencils/backends/ppc_instruction_sets.py b/src/pystencils/old/backends/ppc_instruction_sets.py
similarity index 100%
rename from src/pystencils/backends/ppc_instruction_sets.py
rename to src/pystencils/old/backends/ppc_instruction_sets.py
diff --git a/src/pystencils/backends/riscv_instruction_sets.py b/src/pystencils/old/backends/riscv_instruction_sets.py
similarity index 100%
rename from src/pystencils/backends/riscv_instruction_sets.py
rename to src/pystencils/old/backends/riscv_instruction_sets.py
diff --git a/src/pystencils/backends/simd_instruction_sets.py b/src/pystencils/old/backends/simd_instruction_sets.py
similarity index 98%
rename from src/pystencils/backends/simd_instruction_sets.py
rename to src/pystencils/old/backends/simd_instruction_sets.py
index e9bce873751fb6639cbb77ba2427d6c68c0b3f8f..955156ca2afae9ca539b48f00faf12207c7ff2a5 100644
--- a/src/pystencils/backends/simd_instruction_sets.py
+++ b/src/pystencils/old/backends/simd_instruction_sets.py
@@ -99,7 +99,7 @@ def get_cacheline_size(instruction_set):
return None
import pystencils as ps
- from pystencils.astnodes import SympyAssignment
+ from pystencils.sympyextensions.astnodes import SympyAssignment
import numpy as np
from pystencils.cpu.vectorization import CachelineSize
diff --git a/src/pystencils/backends/x86_instruction_sets.py b/src/pystencils/old/backends/x86_instruction_sets.py
similarity index 100%
rename from src/pystencils/backends/x86_instruction_sets.py
rename to src/pystencils/old/backends/x86_instruction_sets.py
diff --git a/src/pystencils/cpu/__init__.py b/src/pystencils/old/cpu/__init__.py
similarity index 100%
rename from src/pystencils/cpu/__init__.py
rename to src/pystencils/old/cpu/__init__.py
diff --git a/src/pystencils/cpu/cpujit.py b/src/pystencils/old/cpu/cpujit.py
similarity index 99%
rename from src/pystencils/cpu/cpujit.py
rename to src/pystencils/old/cpu/cpujit.py
index deea37c9536fea1b388273f41d7c9f8ac02c1f5e..f67bff84986e053b0ff141a70f4aceb986ed289f 100644
--- a/src/pystencils/cpu/cpujit.py
+++ b/src/pystencils/old/cpu/cpujit.py
@@ -61,7 +61,7 @@ import warnings
import numpy as np
from pystencils import FieldType
-from pystencils.astnodes import LoopOverCoordinate
+from pystencils.sympyextensions.astnodes import LoopOverCoordinate
from pystencils.backends.cbackend import generate_c, get_headers, CFunction
from pystencils.cpu.msvc_detection import get_environment
from pystencils.include import get_pystencils_include_path
diff --git a/src/pystencils/cpu/kernelcreation.py b/src/pystencils/old/cpu/kernelcreation.py
similarity index 98%
rename from src/pystencils/cpu/kernelcreation.py
rename to src/pystencils/old/cpu/kernelcreation.py
index 14e4ea3b7c1a04630f6d124566b29f81512bd944..56262a6e48227ebb1c05291738f0a2d232f3e29a 100644
--- a/src/pystencils/cpu/kernelcreation.py
+++ b/src/pystencils/old/cpu/kernelcreation.py
@@ -1,9 +1,9 @@
import sympy as sp
-import pystencils.astnodes as ast
+import pystencils.sympyextensions.astnodes as ast
from pystencils.config import CreateKernelConfig
from pystencils.enums import Target, Backend
-from pystencils.astnodes import Block, KernelFunction, LoopOverCoordinate, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, KernelFunction, LoopOverCoordinate, SympyAssignment
from pystencils.cpu.cpujit import make_python_function
from pystencils.typing import StructType, TypedSymbol, create_type
from pystencils.typing.transformations import add_types
diff --git a/src/pystencils/cpu/msvc_detection.py b/src/pystencils/old/cpu/msvc_detection.py
similarity index 100%
rename from src/pystencils/cpu/msvc_detection.py
rename to src/pystencils/old/cpu/msvc_detection.py
diff --git a/src/pystencils/cpu/vectorization.py b/src/pystencils/old/cpu/vectorization.py
similarity index 99%
rename from src/pystencils/cpu/vectorization.py
rename to src/pystencils/old/cpu/vectorization.py
index 872f0b3c45983a38b5b1be8fd7f425eb422b570d..0f214a2c4d7d1e3a5a8b0774f5d7282ba70245f9 100644
--- a/src/pystencils/cpu/vectorization.py
+++ b/src/pystencils/old/cpu/vectorization.py
@@ -5,7 +5,7 @@ import numpy as np
import sympy as sp
from sympy.logic.boolalg import BooleanFunction, BooleanAtom
-import pystencils.astnodes as ast
+import pystencils.sympyextensions.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)
diff --git a/src/pystencils/gpu/__init__.py b/src/pystencils/old/gpu/__init__.py
similarity index 100%
rename from src/pystencils/gpu/__init__.py
rename to src/pystencils/old/gpu/__init__.py
diff --git a/src/pystencils/gpu/gpu_array_handler.py b/src/pystencils/old/gpu/gpu_array_handler.py
similarity index 100%
rename from src/pystencils/gpu/gpu_array_handler.py
rename to src/pystencils/old/gpu/gpu_array_handler.py
diff --git a/src/pystencils/gpu/gpujit.py b/src/pystencils/old/gpu/gpujit.py
similarity index 100%
rename from src/pystencils/gpu/gpujit.py
rename to src/pystencils/old/gpu/gpujit.py
diff --git a/src/pystencils/gpu/indexing.py b/src/pystencils/old/gpu/indexing.py
similarity index 99%
rename from src/pystencils/gpu/indexing.py
rename to src/pystencils/old/gpu/indexing.py
index 843e77bb87f60ea4c509d4ec8827e41ee204c42f..75c81604d463a6ea0760b452db9716d6be90d18d 100644
--- a/src/pystencils/gpu/indexing.py
+++ b/src/pystencils/old/gpu/indexing.py
@@ -6,7 +6,7 @@ from typing import List, Tuple
import sympy as sp
from sympy.core.cache import cacheit
-from pystencils.astnodes import Block, Conditional, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, Conditional, SympyAssignment
from pystencils.typing import TypedSymbol, create_type
from pystencils.integer_functions import div_ceil, div_floor
from pystencils.sympyextensions import is_integer_sequence, prod
diff --git a/src/pystencils/gpu/kernelcreation.py b/src/pystencils/old/gpu/kernelcreation.py
similarity index 98%
rename from src/pystencils/gpu/kernelcreation.py
rename to src/pystencils/old/gpu/kernelcreation.py
index c2e6143bcd2e8dc98742c89a2ca0477afcade4b7..88c268809c65ccce1cb80e6c58f8ed304c40bbc4 100644
--- a/src/pystencils/gpu/kernelcreation.py
+++ b/src/pystencils/old/gpu/kernelcreation.py
@@ -1,6 +1,6 @@
import sympy as sp
-from pystencils.astnodes import Block, KernelFunction, LoopOverCoordinate, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, KernelFunction, LoopOverCoordinate, SympyAssignment
from pystencils.config import CreateKernelConfig
from pystencils.typing import StructType, TypedSymbol
from pystencils.typing.transformations import add_types
diff --git a/src/pystencils/gpu/periodicity.py b/src/pystencils/old/gpu/periodicity.py
similarity index 100%
rename from src/pystencils/gpu/periodicity.py
rename to src/pystencils/old/gpu/periodicity.py
diff --git a/src/pystencils/old/kernelcreation.py b/src/pystencils/old/kernelcreation.py
new file mode 100644
index 0000000000000000000000000000000000000000..12dddee1678a810632ad958693e60c4cc881c186
--- /dev/null
+++ b/src/pystencils/old/kernelcreation.py
@@ -0,0 +1,393 @@
+import itertools
+import warnings
+from typing import Union, List
+
+import sympy as sp
+from pystencils.config import CreateKernelConfig
+
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment, AddAugmentedAssignment
+from pystencils.sympyextensions.astnodes import Node, Block, Conditional, LoopOverCoordinate, SympyAssignment
+from pystencils.cpu.vectorization import vectorize
+from pystencils.enums import Target, Backend
+from pystencils.field import Field, FieldType
+from pystencils.node_collection import NodeCollection
+from pystencils.simp.assignment_collection import AssignmentCollection
+from pystencils.kernel_contrains_check import KernelConstraintsCheck
+from pystencils.simplificationfactory import create_simplification_strategy
+from pystencils.stencil import direction_string_to_offset, inverse_direction_string
+from pystencils.transformations import (
+ loop_blocking, move_constants_before_loop, remove_conditionals_in_staggered_kernel)
+
+
+def create_kernel(assignments: Union[Assignment, List[Assignment],
+ AddAugmentedAssignment, List[AddAugmentedAssignment],
+ AssignmentCollection, List[Node], NodeCollection],
+ *,
+ config: CreateKernelConfig = None, **kwargs):
+ """
+ Creates abstract syntax tree (AST) of kernel, using a list of update equations.
+ This function forms the general API and delegates the kernel creation to others depending on the CreateKernelConfig.
+ Args:
+ assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection`
+ config: CreateKernelConfig which includes the needed configuration
+ kwargs: Arguments for updating the config
+
+ Returns:
+ abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
+ can be compiled with through its 'compile()' member
+
+ Example:
+ >>> import pystencils as ps
+ >>> import numpy as np
+ >>> s, d = ps.fields('s, d: [2D]')
+ >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
+ >>> kernel_ast = ps.create_kernel(assignment, config=ps.CreateKernelConfig(cpu_openmp=True))
+ >>> kernel = kernel_ast.compile()
+ >>> d_arr = np.zeros([5, 5])
+ >>> kernel(d=d_arr, s=np.ones([5, 5]))
+ >>> d_arr
+ array([[0., 0., 0., 0., 0.],
+ [0., 4., 4., 4., 0.],
+ [0., 4., 4., 4., 0.],
+ [0., 4., 4., 4., 0.],
+ [0., 0., 0., 0., 0.]])
+ """
+ # ---- Updating configuration from kwargs
+ if not config:
+ config = CreateKernelConfig(**kwargs)
+ else:
+ for k, v in kwargs.items():
+ if not hasattr(config, k):
+ raise KeyError(f'{v} is not a valid kwarg. Please look in CreateKernelConfig for valid settings')
+ setattr(config, k, v)
+
+ # ---- Normalizing parameters
+ if isinstance(assignments, (Assignment, AddAugmentedAssignment)):
+ assignments = [assignments]
+ assert assignments, "Assignments must not be empty!"
+ if isinstance(assignments, list):
+ assignments = NodeCollection(assignments)
+ elif isinstance(assignments, AssignmentCollection):
+ # TODO Markus check and doku
+ # --- applying first default simplifications
+ try:
+ if config.default_assignment_simplifications:
+ simplification = create_simplification_strategy()
+ assignments = simplification(assignments)
+ 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.from_assignment_collection(assignments)
+ assignments.simplification_hints = simplification_hints
+
+ if config.index_fields:
+ return create_indexed_kernel(assignments, config=config)
+ else:
+ return create_domain_kernel(assignments, config=config)
+
+
+def create_domain_kernel(assignments: NodeCollection, *, config: CreateKernelConfig):
+ """
+ Creates abstract syntax tree (AST) of kernel, using a NodeCollection.
+
+ Note that `create_domain_kernel` is a lower level function which shoul be accessed by not providing `index_fields`
+ to create_kernel
+
+ Args:
+ assignments: `pystencils.node_collection.NodeCollection` containing all assignements and nodes to be processed
+ config: CreateKernelConfig which includes the needed configuration
+
+ Returns:
+ abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
+ can be compiled with through its 'compile()' member
+
+ Example:
+ >>> import pystencils as ps
+ >>> import numpy as np
+ >>> from pystencils.kernelcreation import create_domain_kernel
+ >>> from pystencils.node_collection import NodeCollection
+ >>> s, d = ps.fields('s, d: [2D]')
+ >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
+ >>> kernel_config = ps.CreateKernelConfig(cpu_openmp=True)
+ >>> kernel_ast = create_domain_kernel(NodeCollection([assignment]), config=kernel_config)
+ >>> kernel = kernel_ast.compile()
+ >>> d_arr = np.zeros([5, 5])
+ >>> kernel(d=d_arr, s=np.ones([5, 5]))
+ >>> d_arr
+ array([[0., 0., 0., 0., 0.],
+ [0., 4., 4., 4., 0.],
+ [0., 4., 4., 4., 0.],
+ [0., 4., 4., 4., 0.],
+ [0., 0., 0., 0., 0.]])
+ """
+ # --- eval
+ assignments.evaluate_terms()
+
+ # FUTURE WORK from here we shouldn't NEED sympy
+ # --- check constrains
+ check = KernelConstraintsCheck(check_independence_condition=not config.skip_independence_check,
+ check_double_write_condition=not config.allow_double_writes)
+ check.visit(assignments)
+
+ assignments.bound_fields = check.fields_written
+ assignments.rhs_fields = check.fields_read
+
+ # ---- Creating ast
+ ast = None
+ if config.target == Target.CPU:
+ if config.backend == Backend.C:
+ from pystencils.cpu import add_openmp, create_kernel
+ ast = create_kernel(assignments, config=config)
+ for optimization in config.cpu_prepend_optimizations:
+ optimization(ast)
+ omp_collapse = None
+ if config.cpu_blocking:
+ omp_collapse = loop_blocking(ast, config.cpu_blocking)
+ if config.cpu_openmp:
+ add_openmp(ast, num_threads=config.cpu_openmp, collapse=omp_collapse,
+ assume_single_outer_loop=config.omp_single_loop)
+ if config.cpu_vectorize_info:
+ if config.cpu_vectorize_info is True:
+ vectorize(ast)
+ elif isinstance(config.cpu_vectorize_info, dict):
+ vectorize(ast, **config.cpu_vectorize_info)
+ if config.cpu_openmp and config.cpu_blocking and 'nontemporal' in config.cpu_vectorize_info and \
+ config.cpu_vectorize_info['nontemporal'] and 'cachelineZero' in ast.instruction_set:
+ # This condition is stricter than it needs to be: if blocks along the fastest axis start on a
+ # cache line boundary, it's okay. But we cannot determine that here.
+ # We don't need to disallow OpenMP collapsing because it is never applied to the inner loop.
+ raise ValueError("Blocking cannot be combined with cacheline-zeroing")
+ else:
+ raise ValueError("Invalid value for cpu_vectorize_info")
+ elif config.target == Target.GPU:
+ if config.backend == Backend.CUDA:
+ from pystencils.gpu import create_cuda_kernel
+ ast = create_cuda_kernel(assignments, config=config)
+
+ if not ast:
+ raise NotImplementedError(
+ f'{config.target} together with {config.backend} is not supported by `create_domain_kernel`')
+
+ if config.use_auto_for_assignments:
+ for a in ast.atoms(SympyAssignment):
+ a.use_auto = True
+
+ return ast
+
+
+def create_indexed_kernel(assignments: NodeCollection, *, config: CreateKernelConfig):
+ """
+ Similar to :func:`create_kernel`, but here not all cells of a field are updated but only cells with
+ coordinates which are stored in an index field. This traversal method can e.g. be used for boundary handling.
+
+ The coordinates are stored in a separated index_field, which is a one dimensional array with struct data type.
+ This struct has to contain fields named 'x', 'y' and for 3D fields ('z'). These names are configurable with the
+ '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: `pystencils.node_collection.NodeCollection` containing all assignements and nodes to be processed
+ config: CreateKernelConfig which includes the needed configuration
+
+ Returns:
+ abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
+ can be compiled with through its 'compile()' member
+
+ Example:
+ >>> import pystencils as ps
+ >>> from pystencils.node_collection import NodeCollection
+ >>> import numpy as np
+ >>> from pystencils.kernelcreation import create_indexed_kernel
+ >>>
+ >>> # Index field stores the indices of the cell to visit together with optional values
+ >>> index_arr_dtype = np.dtype([('x', np.int32), ('y', np.int32), ('val', np.double)])
+ >>> index_arr = np.array([(1, 1, 0.1), (2, 2, 0.2), (3, 3, 0.3)], dtype=index_arr_dtype)
+ >>> idx_field = ps.fields(idx=index_arr)
+ >>>
+ >>> # Additional values stored in index field can be accessed in the kernel as well
+ >>> 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(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. ]])
+
+ """
+ # --- eval
+ assignments.evaluate_terms()
+
+ # FUTURE WORK from here we shouldn't NEED sympy
+ # --- check constrains
+ check = KernelConstraintsCheck(check_independence_condition=not config.skip_independence_check,
+ check_double_write_condition=not config.allow_double_writes)
+ check.visit(assignments)
+
+ assignments.bound_fields = check.fields_written
+ assignments.rhs_fields = check.fields_read
+
+ ast = None
+ if config.target == Target.CPU and config.backend == Backend.C:
+ from pystencils.cpu import add_openmp, create_indexed_kernel
+ ast = create_indexed_kernel(assignments, config=config)
+ if config.cpu_openmp:
+ add_openmp(ast, num_threads=config.cpu_openmp)
+ elif config.target == Target.GPU:
+ if config.backend == Backend.CUDA:
+ from pystencils.gpu import created_indexed_cuda_kernel
+ ast = created_indexed_cuda_kernel(assignments, config=config)
+
+ if not ast:
+ raise NotImplementedError(f'Indexed kernels are not yet supported for {config.target} with {config.backend}')
+ return ast
+
+
+def create_staggered_kernel(assignments, target: Target = Target.CPU, gpu_exclusive_conditions=False, **kwargs):
+ """Kernel that updates a staggered field.
+
+ .. image:: /img/staggered_grid.svg
+
+ For a staggered field, the first index coordinate defines the location of the staggered value.
+ Further index coordinates can be used to store vectors/tensors at each point.
+
+ Args:
+ assignments: a sequence of assignments or an AssignmentCollection.
+ Assignments to staggered field are processed specially, while subexpressions and assignments to
+ regular fields are passed through to `create_kernel`. Multiple different staggered fields can be
+ used, but they all need to use the same stencil (i.e. the same number of staggered points) and
+ shape.
+ target: 'CPU' or 'GPU'
+ gpu_exclusive_conditions: disable the use of multiple conditionals inside the loop. The outer layers are then
+ handled in an else branch.
+ kwargs: passed directly to create_kernel, iteration_slice and ghost_layers parameters are not allowed
+
+ Returns:
+ AST, see `create_kernel`
+ """
+ # TODO: Add doku like in the other kernels
+ if 'ghost_layers' in kwargs:
+ assert kwargs['ghost_layers'] is None
+ del kwargs['ghost_layers']
+ if 'iteration_slice' in kwargs:
+ assert kwargs['iteration_slice'] is None
+ del kwargs['iteration_slice']
+ if 'omp_single_loop' in kwargs:
+ assert kwargs['omp_single_loop'] is False
+ del kwargs['omp_single_loop']
+
+ if isinstance(assignments, AssignmentCollection):
+ subexpressions = assignments.subexpressions + [a for a in assignments.main_assignments
+ if not hasattr(a, 'lhs')
+ or type(a.lhs) is not Field.Access
+ or not FieldType.is_staggered(a.lhs.field)]
+ assignments = [a for a in assignments.main_assignments if hasattr(a, 'lhs')
+ and type(a.lhs) is Field.Access
+ and FieldType.is_staggered(a.lhs.field)]
+ else:
+ subexpressions = [a for a in assignments if not hasattr(a, 'lhs')
+ or type(a.lhs) is not Field.Access
+ or not FieldType.is_staggered(a.lhs.field)]
+ assignments = [a for a in assignments if hasattr(a, 'lhs')
+ and type(a.lhs) is Field.Access
+ and FieldType.is_staggered(a.lhs.field)]
+ if len(set([tuple(a.lhs.field.staggered_stencil) for a in assignments])) != 1:
+ raise ValueError("All assignments need to be made to staggered fields with the same stencil")
+ if len(set([a.lhs.field.shape for a in assignments])) != 1:
+ raise ValueError("All assignments need to be made to staggered fields with the same shape")
+
+ staggered_field = assignments[0].lhs.field
+ stencil = staggered_field.staggered_stencil
+ dim = staggered_field.spatial_dimensions
+ shape = staggered_field.shape
+
+ counters = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)]
+
+ final_assignments = []
+
+ # find out whether any of the ghost layers is not needed
+ common_exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
+ for direction in stencil:
+ exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
+ for elementary_direction in direction:
+ exclusions.remove(inverse_direction_string(elementary_direction))
+ common_exclusions.intersection_update(exclusions)
+ ghost_layers = [[0, 0] for d in range(dim)]
+ for direction in common_exclusions:
+ direction = direction_string_to_offset(direction)
+ for d, s in enumerate(direction):
+ if s == 1:
+ ghost_layers[d][1] = 1
+ elif s == -1:
+ ghost_layers[d][0] = 1
+
+ def condition(direction):
+ """exclude those staggered points that correspond to fluxes between ghost cells"""
+ exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
+
+ for elementary_direction in direction:
+ exclusions.remove(inverse_direction_string(elementary_direction))
+ conditions = []
+ for e in exclusions:
+ if e in common_exclusions:
+ continue
+ offset = direction_string_to_offset(e)
+ for i, o in enumerate(offset):
+ if o == 1:
+ conditions.append(counters[i] < shape[i] - 1)
+ elif o == -1:
+ conditions.append(counters[i] > 0)
+ return sp.And(*conditions)
+
+ if gpu_exclusive_conditions:
+ outer_assignment = None
+ conditions = {direction: condition(direction) for direction in stencil}
+ for num_conditions in range(len(stencil)):
+ for combination in itertools.combinations(conditions.values(), num_conditions):
+ for assignment in assignments:
+ direction = stencil[assignment.lhs.index[0]]
+ if conditions[direction] in combination:
+ assignment = SympyAssignment(assignment.lhs, assignment.rhs)
+ outer_assignment = Conditional(sp.And(*combination), Block([assignment]), outer_assignment)
+
+ inner_assignment = []
+ for assignment in assignments:
+ inner_assignment.append(SympyAssignment(assignment.lhs, assignment.rhs))
+ last_conditional = Conditional(sp.And(*[condition(d) for d in stencil]),
+ Block(inner_assignment), outer_assignment)
+ final_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
+ [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
+ [last_conditional]
+
+ config = CreateKernelConfig(target=target, ghost_layers=ghost_layers, omp_single_loop=False, **kwargs)
+ ast = create_kernel(final_assignments, config=config)
+ return ast
+
+ for assignment in assignments:
+ direction = stencil[assignment.lhs.index[0]]
+ sp_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
+ [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
+ [SympyAssignment(assignment.lhs, assignment.rhs)]
+ last_conditional = Conditional(condition(direction), Block(sp_assignments))
+ final_assignments.append(last_conditional)
+
+ remove_start_conditional = any([gl[0] == 0 for gl in ghost_layers])
+ prepend_optimizations = [lambda ast: remove_conditionals_in_staggered_kernel(ast, remove_start_conditional),
+ move_constants_before_loop]
+ if 'cpu_prepend_optimizations' in kwargs:
+ prepend_optimizations += kwargs['cpu_prepend_optimizations']
+ del kwargs['cpu_prepend_optimizations']
+
+ config = CreateKernelConfig(ghost_layers=ghost_layers, target=target, omp_single_loop=False,
+ cpu_prepend_optimizations=prepend_optimizations, **kwargs)
+ ast = create_kernel(final_assignments, config=config)
+ return ast
diff --git a/src/pystencils/transformations.py b/src/pystencils/old/transformations.py
similarity index 97%
rename from src/pystencils/transformations.py
rename to src/pystencils/old/transformations.py
index 79c24d1467e4f2cde3c4c7bd9866ca4fd56cb344..76d356f9fb85caed903c0247de337670bf106430 100644
--- a/src/pystencils/transformations.py
+++ b/src/pystencils/old/transformations.py
@@ -9,8 +9,8 @@ from typing import Set
import sympy as sp
import pystencils as ps
-import pystencils.astnodes as ast
-from pystencils.assignment import Assignment
+import pystencils.sympyextensions.astnodes as ast
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment
from pystencils.typing import (CastFunc, PointerType, StructType, TypedSymbol, get_base_type,
ReinterpretCastFunc, get_next_parent_of_type, parents_of_type)
from pystencils.field import Field, FieldType
@@ -21,58 +21,7 @@ from pystencils.slicing import normalize_slice
from pystencils.integer_functions import int_div
-class NestedScopes:
- """Symbol visibility model using nested scopes
-
- - every accessed symbol that was not defined before, is added as a "free parameter"
- - free parameters are global, i.e. they are not in scopes
- - push/pop adds or removes a scope
-
- >>> s = NestedScopes()
- >>> s.access_symbol("a")
- >>> s.is_defined("a")
- False
- >>> s.free_parameters
- {'a'}
- >>> s.define_symbol("b")
- >>> s.is_defined("b")
- True
- >>> s.push()
- >>> s.is_defined_locally("b")
- False
- >>> s.define_symbol("c")
- >>> s.pop()
- >>> s.is_defined("c")
- False
- """
-
- def __init__(self):
- self.free_parameters = set()
- self._defined = [set()]
-
- def access_symbol(self, symbol):
- if not self.is_defined(symbol):
- self.free_parameters.add(symbol)
-
- def define_symbol(self, symbol):
- self._defined[-1].add(symbol)
-
- def is_defined(self, symbol):
- return any(symbol in scopes for scopes in self._defined)
-
- def is_defined_locally(self, symbol):
- return symbol in self._defined[-1]
-
- def push(self):
- self._defined.append(set())
-
- def pop(self):
- self._defined.pop()
- assert self.depth >= 1
- @property
- def depth(self):
- return len(self._defined)
def filtered_tree_iteration(node, node_type, stop_type=None):
@@ -110,7 +59,7 @@ def iterate_loops_by_depth(node, nesting_depth):
A depth of -1 indicates the innermost loops.
Returns: Iterable listing all loop nodes of given nesting depth.
"""
- from pystencils.astnodes import LoopOverCoordinate
+ from pystencils.sympyextensions.astnodes import LoopOverCoordinate
def _internal_default(node, nesting_depth):
isloop = isinstance(node, LoopOverCoordinate)
diff --git a/src/pystencils/typing/__init__.py b/src/pystencils/old/typing/__init__.py
similarity index 100%
rename from src/pystencils/typing/__init__.py
rename to src/pystencils/old/typing/__init__.py
diff --git a/src/pystencils/old/typing/cast_functions.py b/src/pystencils/old/typing/cast_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..52e63773143eb26d2d68bdd80a47ccda8f513387
--- /dev/null
+++ b/src/pystencils/old/typing/cast_functions.py
@@ -0,0 +1,9 @@
+import numpy as np
+import sympy as sp
+from sympy.logic.boolalg import Boolean
+
+from pystencils.typing.types import AbstractType, BasicType
+from pystencils.typing.typed_sympy import TypedSymbol
+
+
+
diff --git a/src/pystencils/typing/leaf_typing.py b/src/pystencils/old/typing/leaf_typing.py
similarity index 99%
rename from src/pystencils/typing/leaf_typing.py
rename to src/pystencils/old/typing/leaf_typing.py
index 9e7065b0a5960febbcc954f70040b9097031041e..5959e8ee676b16ecb49a2c8e35d4bcab48d571e6 100644
--- a/src/pystencils/typing/leaf_typing.py
+++ b/src/pystencils/old/typing/leaf_typing.py
@@ -14,7 +14,7 @@ from sympy.functions.elementary.hyperbolic import HyperbolicFunction
from sympy.logic.boolalg import BooleanFunction
from sympy.logic.boolalg import BooleanAtom
-from pystencils import astnodes as ast
+from pystencils.sympyextensions import astnodes as ast
from pystencils.functions import DivFunc, AddressOf
from pystencils.cpu.vectorization import vec_all, vec_any
from pystencils.field import Field
@@ -107,7 +107,7 @@ class TypeAdder:
# Trivial cases
from pystencils.field import Field
import pystencils.integer_functions
- from pystencils.bit_masks import flag_cond
+ from pystencils.sympyextensions.bit_masks import flag_cond
bool_type = BasicType('bool')
# TOOO: check the access
diff --git a/src/pystencils/typing/transformations.py b/src/pystencils/old/typing/transformations.py
similarity index 94%
rename from src/pystencils/typing/transformations.py
rename to src/pystencils/old/typing/transformations.py
index 43e69eb286072160180536bf26e4f922960efc68..2cd477e125bc6ceb19993344562087ac2afa4fd3 100644
--- a/src/pystencils/typing/transformations.py
+++ b/src/pystencils/old/typing/transformations.py
@@ -1,6 +1,6 @@
from typing import List
-from pystencils.astnodes import Node
+from pystencils.sympyextensions.astnodes import Node
from pystencils.config import CreateKernelConfig
from pystencils.typing.leaf_typing import TypeAdder
diff --git a/src/pystencils/typing/utilities.py b/src/pystencils/old/typing/utilities.py
similarity index 99%
rename from src/pystencils/typing/utilities.py
rename to src/pystencils/old/typing/utilities.py
index 223da701a4d5c133715eb30f99366c44b13f16b2..086be03f84886cb4206cb0dd967025b75eb8fa43 100644
--- a/src/pystencils/typing/utilities.py
+++ b/src/pystencils/old/typing/utilities.py
@@ -111,7 +111,7 @@ def get_type_of_expression(expr,
default_float_type='double',
default_int_type='int',
symbol_type_dict=None):
- from pystencils.astnodes import ResolvedFieldAccess
+ from pystencils.sympyextensions.astnodes import ResolvedFieldAccess
from pystencils.cpu.vectorization import vec_all, vec_any
if default_float_type == 'float':
diff --git a/src/pystencils/placeholder_function.py b/src/pystencils/placeholder_function.py
index 8b8aa6ed0ed39d59f97c960d7efdd73f07cf38ef..8d675f0338de2607438069495dffb9bd5a2a725a 100644
--- a/src/pystencils/placeholder_function.py
+++ b/src/pystencils/placeholder_function.py
@@ -2,8 +2,8 @@ from typing import List
import sympy as sp
-from pystencils.assignment import Assignment
-from pystencils.astnodes import Node
+from pystencils.sympyextensions.assignmentcollection.assignment import Assignment
+from pystencils.sympyextensions.astnodes import Node
from pystencils.sympyextensions import is_constant
from pystencils.transformations import generic_visit
diff --git a/src/pystencils/rng.py b/src/pystencils/rng.py
index 84155b00c28f685b584ddb42e806e425e21df486..5f74b62a4191d92af9ec5097bef6c5e092a41ed7 100644
--- a/src/pystencils/rng.py
+++ b/src/pystencils/rng.py
@@ -3,7 +3,7 @@ import numpy as np
import sympy as sp
from pystencils.typing import TypedSymbol, CastFunc
-from pystencils.astnodes import LoopOverCoordinate
+from pystencils.sympyextensions.astnodes import LoopOverCoordinate
from pystencils.backends.cbackend import CustomCodeNode
from pystencils.sympyextensions import fast_subs
diff --git a/src/pystencils/simp/__init__.py b/src/pystencils/simp/__init__.py
deleted file mode 100644
index 0e5ff0e6435bf08f3d305ff7cc9087775a09c6e2..0000000000000000000000000000000000000000
--- a/src/pystencils/simp/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-from .assignment_collection import AssignmentCollection
-from .simplifications import (
- add_subexpressions_for_constants,
- add_subexpressions_for_divisions, add_subexpressions_for_field_reads,
- add_subexpressions_for_sums, apply_on_all_subexpressions, apply_to_all_assignments,
- subexpression_substitution_in_existing_subexpressions,
- subexpression_substitution_in_main_assignments, sympy_cse, sympy_cse_on_assignment_list)
-from .subexpression_insertion import (
- insert_aliases, insert_zeros, insert_constants,
- insert_constant_additions, insert_constant_multiples,
- insert_squares, insert_symbol_times_minus_one)
-from .simplificationstrategy import SimplificationStrategy
-
-__all__ = ['AssignmentCollection', 'SimplificationStrategy',
- 'sympy_cse', 'sympy_cse_on_assignment_list', 'apply_to_all_assignments',
- 'apply_on_all_subexpressions', 'subexpression_substitution_in_existing_subexpressions',
- 'subexpression_substitution_in_main_assignments', 'add_subexpressions_for_constants',
- 'add_subexpressions_for_divisions', 'add_subexpressions_for_sums', 'add_subexpressions_for_field_reads',
- 'insert_aliases', 'insert_zeros', 'insert_constants',
- 'insert_constant_additions', 'insert_constant_multiples',
- 'insert_squares', 'insert_symbol_times_minus_one']
diff --git a/src/pystencils/spatial_coordinates.py b/src/pystencils/spatial_coordinates.py
index 6c3ba4db7d98b50bce9442619fbbce07782004a1..a8be92c94bc16f5707ff2c93b39bb340158af129 100644
--- a/src/pystencils/spatial_coordinates.py
+++ b/src/pystencils/spatial_coordinates.py
@@ -2,17 +2,18 @@
import sympy
import pystencils
-import pystencils.astnodes
+import pystencils.sympyextensions.astnodes
-x_, y_, z_ = tuple(pystencils.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(3))
+x_, y_, z_ = tuple(pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(3))
x_staggered, y_staggered, z_staggered = x_ + 0.5, y_ + 0.5, z_ + 0.5
def x_vector(ndim):
- return sympy.Matrix(tuple(pystencils.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(ndim)))
+ return sympy.Matrix(tuple(
+ pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(ndim)))
def x_staggered_vector(ndim):
return sympy.Matrix(tuple(
- pystencils.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) + 0.5 for i in range(ndim)
+ pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) + 0.5 for i in range(ndim)
))
diff --git a/src/pystencils/sympyextensions/__init__.py b/src/pystencils/sympyextensions/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d61e705d133f263449cd2f94e21f9f0886dabe5d
--- /dev/null
+++ b/src/pystencils/sympyextensions/__init__.py
@@ -0,0 +1,23 @@
+from .assignment import Assignment, AugmentedAssignment, AddAugmentedAssignment, assignment_from_stencil
+from pystencils.sympyextensions.assignment_collection import AssignmentCollection
+from .simplificationstrategy import SimplificationStrategy
+from .simplifications import (sympy_cse, sympy_cse_on_assignment_list, apply_to_all_assignments,
+ apply_on_all_subexpressions, subexpression_substitution_in_existing_subexpressions,
+ subexpression_substitution_in_main_assignments, add_subexpressions_for_constants,
+ add_subexpressions_for_divisions, add_subexpressions_for_sums,
+ add_subexpressions_for_field_reads)
+from .subexpression_insertion import (
+ insert_aliases, insert_zeros, insert_constants,
+ insert_constant_additions, insert_constant_multiples,
+ insert_squares, insert_symbol_times_minus_one)
+
+
+__all__ = ['Assignment', 'AugmentedAssignment', 'AddAugmentedAssignment', 'assignment_from_stencil',
+ 'AssignmentCollection', 'SimplificationStrategy',
+ 'sympy_cse', 'sympy_cse_on_assignment_list', 'apply_to_all_assignments',
+ 'apply_on_all_subexpressions', 'subexpression_substitution_in_existing_subexpressions',
+ 'subexpression_substitution_in_main_assignments', 'add_subexpressions_for_constants',
+ 'add_subexpressions_for_divisions', 'add_subexpressions_for_sums', 'add_subexpressions_for_field_reads',
+ 'insert_aliases', 'insert_zeros', 'insert_constants',
+ 'insert_constant_additions', 'insert_constant_multiples',
+ 'insert_squares', 'insert_symbol_times_minus_one']
\ No newline at end of file
diff --git a/src/pystencils/assignment.py b/src/pystencils/sympyextensions/assignment.py
similarity index 97%
rename from src/pystencils/assignment.py
rename to src/pystencils/sympyextensions/assignment.py
index d0e849954cebe4726e4a03e991fbad0e0931f315..591bde9a5765ae49707ed0db960337c80374996a 100644
--- a/src/pystencils/assignment.py
+++ b/src/pystencils/sympyextensions/assignment.py
@@ -3,8 +3,6 @@ import sympy as sp
from sympy.codegen.ast import Assignment, AugmentedAssignment, AddAugmentedAssignment
from sympy.printing.latex import LatexPrinter
-__all__ = ['Assignment', 'AugmentedAssignment', 'AddAugmentedAssignment', 'assignment_from_stencil']
-
def print_assignment_latex(printer, expr):
binop = f"{expr.binop}=" if isinstance(expr, AugmentedAssignment) else ''
diff --git a/src/pystencils/simp/assignment_collection.py b/src/pystencils/sympyextensions/assignment_collection.py
similarity index 98%
rename from src/pystencils/simp/assignment_collection.py
rename to src/pystencils/sympyextensions/assignment_collection.py
index cb89e58bbd58b4b5344aa3a59d741f1618b3a574..de8f7b4e5e5f310140198cbc75e7fd0980674a8e 100644
--- a/src/pystencils/simp/assignment_collection.py
+++ b/src/pystencils/sympyextensions/assignment_collection.py
@@ -5,9 +5,9 @@ from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Set,
import sympy as sp
import pystencils
-from pystencils.assignment import Assignment
-from pystencils.simp.simplifications import (sort_assignments_topologically, transform_lhs_and_rhs, transform_rhs)
-from pystencils.sympyextensions import count_operations, fast_subs
+from .assignment import Assignment
+from .simplifications import (sort_assignments_topologically, transform_lhs_and_rhs, transform_rhs)
+from pystencils.sympyextensions.math import count_operations, fast_subs
class AssignmentCollection:
diff --git a/src/pystencils/astnodes.py b/src/pystencils/sympyextensions/astnodes.py
similarity index 77%
rename from src/pystencils/astnodes.py
rename to src/pystencils/sympyextensions/astnodes.py
index f399287ed02ec4eb0d3d0e295d72ee1cf5ecc14b..0cd82322c5cd303d5698ea00d2e084ce0e44ebcd 100644
--- a/src/pystencils/astnodes.py
+++ b/src/pystencils/sympyextensions/astnodes.py
@@ -3,14 +3,15 @@ import itertools
import uuid
from typing import Any, List, Optional, Sequence, Set, Union
+from .assignment import Assignment
+from pystencils.enums import Target, Backend
+
import sympy as sp
-import pystencils
-from pystencils.typing.utilities import create_type, get_next_parent_of_type
-from pystencils.enums import Target, Backend
-from pystencils.field import Field
-from pystencils.typing.typed_sympy import FieldPointerSymbol, FieldShapeSymbol, FieldStrideSymbol, TypedSymbol
-from pystencils.sympyextensions import fast_subs
+from .math import fast_subs
+from .typed_sympy import (create_type, CastFunc,
+ FieldPointerSymbol, FieldShapeSymbol, FieldStrideSymbol, TypedSymbol)
+
NodeOrExpr = Union['Node', sp.Expr]
@@ -136,158 +137,6 @@ class Conditional(Node):
self.parent.replace(self, [self.false_block] if self.false_block else [])
-class KernelFunction(Node):
- class Parameter:
- """Function parameter.
-
- Each undefined symbol in a `KernelFunction` node becomes a parameter to the function.
- Parameters are either symbols introduced by the user that never occur on the left hand side of an
- Assignment, or are related to fields/arrays passed to the function.
-
- A parameter consists of the typed symbol (symbol property). For field related parameters this is a symbol
- defined in pystencils.kernelparameters.
- If the parameter is related to one or multiple fields, these fields are referenced in the fields property.
- """
-
- def __init__(self, symbol, fields):
- self.symbol = symbol # type: TypedSymbol
- self.fields = fields # type: Sequence[Field]
-
- def __repr__(self):
- return repr(self.symbol)
-
- @property
- def is_field_stride(self):
- return isinstance(self.symbol, FieldStrideSymbol)
-
- @property
- def is_field_shape(self):
- return isinstance(self.symbol, FieldShapeSymbol)
-
- @property
- def is_field_pointer(self):
- return isinstance(self.symbol, FieldPointerSymbol)
-
- @property
- def is_field_parameter(self):
- return self.is_field_pointer or self.is_field_shape or self.is_field_stride
-
- @property
- def field_name(self):
- return self.fields[0].name
-
- def __init__(self, body, target: Target, backend: Backend, compile_function, ghost_layers,
- function_name: str = "kernel",
- assignments=None):
- super(KernelFunction, self).__init__()
- self._body = body
- body.parent = self
- self.function_name = function_name
- self._body.parent = self
- self.ghost_layers = ghost_layers
- self._target = target
- self._backend = backend
- # these variables are assumed to be global, so no automatic parameter is generated for them
- self.global_variables = set()
- self.instruction_set = None # used in `vectorize` function to tell the backend which i.s. (SSE,AVX) to use
- # function that compiles the node to a Python callable, is set by the backends
- self._compile_function = compile_function
- self.assignments = assignments
- # If nontemporal stores are activated together with the Neon instruction set it results in cacheline zeroing
- # For cacheline zeroing the information of the field size for each field is needed. Thus, in this case
- # all field sizes are kernel parameters and not just the common field size used for the loops
- self.use_all_written_field_sizes = False
-
- @property
- def target(self):
- """See pystencils.Target"""
- return self._target
-
- @property
- def backend(self):
- """Backend for generating the code: `Backend`"""
- return self._backend
-
- @property
- def symbols_defined(self):
- return set()
-
- @property
- def undefined_symbols(self):
- return set()
-
- @property
- def body(self):
- return self._body
-
- @body.setter
- def body(self, value):
- self._body = value
- self._body.parent = self
-
- @property
- def args(self):
- return self._body,
-
- @property
- def fields_accessed(self) -> Set[Field]:
- """Set of Field instances: fields which are accessed inside this kernel function"""
- return set(o.field for o in itertools.chain(self.atoms(ResolvedFieldAccess)))
-
- @property
- def fields_written(self) -> Set[Field]:
- assignments = self.atoms(SympyAssignment)
- return set().union(itertools.chain.from_iterable([f.field for f in a.lhs.free_symbols if hasattr(f, 'field')]
- for a in assignments))
-
- @property
- def fields_read(self) -> Set[Field]:
- assignments = self.atoms(SympyAssignment)
- return set().union(itertools.chain.from_iterable([f.field for f in a.rhs.free_symbols if hasattr(f, 'field')]
- for a in assignments))
-
- def get_parameters(self) -> Sequence['KernelFunction.Parameter']:
- """Returns list of parameters for this function.
-
- This function is expensive, cache the result where possible!
- """
- field_map = {f.name: f for f in self.fields_accessed}
- sizes = set()
-
- if self.use_all_written_field_sizes:
- sizes = set().union(*(a.shape[:a.spatial_dimensions] for a in self.fields_written))
- sizes = filter(lambda s: isinstance(s, FieldShapeSymbol), sizes)
-
- def get_fields(symbol):
- if hasattr(symbol, 'field_name'):
- return field_map[symbol.field_name],
- elif hasattr(symbol, 'field_names'):
- return tuple(field_map[fn] for fn in symbol.field_names)
- return ()
-
- argument_symbols = self._body.undefined_symbols - self.global_variables
- argument_symbols.update(sizes)
- parameters = [self.Parameter(symbol, get_fields(symbol)) for symbol in argument_symbols]
- if hasattr(self, 'indexing'):
- parameters += [self.Parameter(s, []) for s in self.indexing.symbolic_parameters()]
- parameters.sort(key=lambda p: p.symbol.name)
- return parameters
-
- def __str__(self):
- params = [p.symbol for p in self.get_parameters()]
- return '{0} {1}({2})\n{3}'.format(type(self).__name__, self.function_name, params,
- ("\t" + "\t".join(str(self.body).splitlines(True))))
-
- def __repr__(self):
- params = [p.symbol for p in self.get_parameters()]
- return f'{type(self).__name__} {self.function_name}({params})'
-
- def compile(self, *args, **kwargs):
- if self._compile_function is None:
- raise ValueError("No compile-function provided for this KernelFunction node")
- return self._compile_function(self, *args, **kwargs)
-
-
class SkipIteration(Node):
@property
def args(self):
@@ -387,7 +236,7 @@ class Block(Node):
def symbols_defined(self):
result = set()
for a in self.args:
- if isinstance(a, pystencils.Assignment):
+ if isinstance(a, Assignment):
result.update(a.free_symbols)
else:
result.update(a.symbols_defined)
@@ -398,7 +247,7 @@ class Block(Node):
result = set()
defined_symbols = set()
for a in self.args:
- if isinstance(a, pystencils.Assignment):
+ if isinstance(a, Assignment):
result.update(a.free_symbols)
defined_symbols.update({a.lhs})
else:
@@ -574,7 +423,6 @@ class SympyAssignment(Node):
self._use_auto = use_auto
def __is_declaration(self):
- from pystencils.typing import CastFunc
if isinstance(self._lhs_symbol, CastFunc):
return False
if any(isinstance(self._lhs_symbol, c) for c in (Field.Access, sp.Indexed, TemporaryMemoryAllocation)):
@@ -785,7 +633,6 @@ class TemporaryMemoryFree(Node):
def early_out(condition):
- from pystencils.cpu.vectorization import vec_all
return Conditional(vec_all(condition), Block([SkipIteration()]))
diff --git a/src/pystencils/bit_masks.py b/src/pystencils/sympyextensions/bit_masks.py
similarity index 100%
rename from src/pystencils/bit_masks.py
rename to src/pystencils/sympyextensions/bit_masks.py
diff --git a/src/pystencils/sympyextensions.py b/src/pystencils/sympyextensions/math.py
similarity index 99%
rename from src/pystencils/sympyextensions.py
rename to src/pystencils/sympyextensions/math.py
index 680b58670f61a188c0474a1cf9137fa35a552a01..4a6f01a5c69673ec47cec8eda80f23d1a82d31e8 100644
--- a/src/pystencils/sympyextensions.py
+++ b/src/pystencils/sympyextensions/math.py
@@ -10,10 +10,9 @@ from sympy import PolynomialError
from sympy.functions import Abs
from sympy.core.numbers import Zero
-from pystencils.assignment import Assignment
+from .assignment import Assignment
from pystencils.functions import DivFunc
-from pystencils.typing import CastFunc, get_type_of_expression, PointerType, VectorType
-from pystencils.typing.typed_sympy import FieldPointerSymbol
+from .typed_sympy import CastFunc, PointerType, VectorType, FieldPointerSymbol
T = TypeVar('T')
@@ -656,7 +655,7 @@ def count_operations(term: Union[sp.Expr, List[sp.Expr], List[Assignment]],
def count_operations_in_ast(ast) -> Dict[str, int]:
"""Counts number of operations in an abstract syntax tree, see also :func:`count_operations`"""
- from pystencils.astnodes import SympyAssignment
+ from pystencils.sympyextensions.astnodes import SympyAssignment
result = defaultdict(int)
def visit(node):
diff --git a/src/pystencils/simp/simplifications.py b/src/pystencils/sympyextensions/simplifications.py
similarity index 87%
rename from src/pystencils/simp/simplifications.py
rename to src/pystencils/sympyextensions/simplifications.py
index e1f42dc60467c2cd8a7fa874afd79323e77111a7..c16e42f85f2340bf061aa06d2875f9b1b69bde7f 100644
--- a/src/pystencils/simp/simplifications.py
+++ b/src/pystencils/sympyextensions/simplifications.py
@@ -4,11 +4,10 @@ from collections import defaultdict
import sympy as sp
-from pystencils.assignment import Assignment
-from pystencils.astnodes import Node
-from pystencils.field import Field
-from pystencils.sympyextensions import subs_additive, is_constant, recursive_collect
-from pystencils.typing import TypedSymbol
+from .assignment import Assignment
+from pystencils.sympyextensions.astnodes import Node
+from .math import subs_additive, is_constant, recursive_collect
+from .typed_sympy import TypedSymbol
def sort_assignments_topologically(assignments: Sequence[Union[Assignment, Node]]) -> List[Union[Assignment, Node]]:
@@ -55,7 +54,7 @@ def sympy_cse(ac, **kwargs):
def sympy_cse_on_assignment_list(assignments: List[Assignment]) -> List[Assignment]:
"""Extracts common subexpressions from a list of assignments."""
- from pystencils.simp.assignment_collection import AssignmentCollection
+ from pystencils.sympyextensions import AssignmentCollection
ec = AssignmentCollection([], assignments)
return sympy_cse(ec).all_assignments
@@ -163,6 +162,7 @@ def add_subexpressions_for_sums(ac):
for eq in ac.all_assignments:
search_addends(eq.rhs)
+ from pystencils.field import Field
addends = [a for a in addends if not isinstance(a, sp.Symbol) or isinstance(a, Field.Access)]
new_symbol_gen = ac.subexpression_symbol_generator
substitutions = {addend: new_symbol for new_symbol, addend in zip(new_symbol_gen, addends)}
@@ -185,6 +185,7 @@ def add_subexpressions_for_field_reads(ac, subexpressions=True, main_assignments
if main_assignments:
to_iterate = chain(to_iterate, ac.main_assignments)
+ from pystencils.field import Field
for assignment in to_iterate:
if hasattr(assignment, 'lhs') and hasattr(assignment, 'rhs'):
field_reads.update(assignment.rhs.atoms(Field.Access))
@@ -236,31 +237,3 @@ def apply_on_all_subexpressions(operation: Callable[[sp.Expr], sp.Expr]):
f.__name__ = operation.__name__
return f
-
-# TODO Markus
-# make this really work for Assignmentcollections
-# this function should ONLY evaluate
-# do the optims_c99 elsewhere optionally
-
-# def apply_sympy_optimisations(ac: AssignmentCollection):
-# """ Evaluates constant expressions (e.g. :math:`\\sqrt{3}` will be replaced by its floating point representation)
-# and applies the default sympy optimisations. See sympy.codegen.rewriting
-# """
-#
-# # Evaluates all constant terms
-#
-# assignments = ac.all_assignments
-#
-# evaluate_constant_terms = ReplaceOptim(lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer,
-# lambda p: p.evalf())
-#
-# sympy_optimisations = [evaluate_constant_terms] + list(optims_c99)
-#
-# assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations))
-# if hasattr(a, 'lhs')
-# else a for a in assignments]
-# assignments_nodes = [a.atoms(SympyAssignment) for a in assignments]
-# for a in chain.from_iterable(assignments_nodes):
-# a.optimize(sympy_optimisations)
-#
-# return AssignmentCollection(assignments)
diff --git a/src/pystencils/simp/simplificationstrategy.py b/src/pystencils/sympyextensions/simplificationstrategy.py
similarity index 98%
rename from src/pystencils/simp/simplificationstrategy.py
rename to src/pystencils/sympyextensions/simplificationstrategy.py
index cc601aa56981124dd75199f333a8c005fbdcb95d..cd2fa8faa790f8b0d93c1e683c60437f93362aa9 100644
--- a/src/pystencils/simp/simplificationstrategy.py
+++ b/src/pystencils/sympyextensions/simplificationstrategy.py
@@ -3,7 +3,7 @@ from typing import Any, Callable, Optional, Sequence
import sympy as sp
-from pystencils.simp.assignment_collection import AssignmentCollection
+from .assignment_collection import AssignmentCollection
class SimplificationStrategy:
diff --git a/src/pystencils/simp/subexpression_insertion.py b/src/pystencils/sympyextensions/subexpression_insertion.py
similarity index 98%
rename from src/pystencils/simp/subexpression_insertion.py
rename to src/pystencils/sympyextensions/subexpression_insertion.py
index 27ac798269bb043c5c8361ea019492a2be5dd9a3..8cedad4665033e91a6812d4c9c23634c8d208c70 100644
--- a/src/pystencils/simp/subexpression_insertion.py
+++ b/src/pystencils/sympyextensions/subexpression_insertion.py
@@ -1,5 +1,5 @@
import sympy as sp
-from pystencils.sympyextensions import is_constant
+from .math import is_constant
# Subexpression Insertion
diff --git a/src/pystencils/sympyextensions/typed_sympy.py b/src/pystencils/sympyextensions/typed_sympy.py
new file mode 100644
index 0000000000000000000000000000000000000000..458fbae20d538a5ff88a1246cb621814f5e1a68b
--- /dev/null
+++ b/src/pystencils/sympyextensions/typed_sympy.py
@@ -0,0 +1,630 @@
+from abc import abstractmethod
+from typing import Union
+
+import numpy as np
+import sympy as sp
+
+
+class AbstractType(sp.Atom):
+ # TODO: Is it necessary to ineherit from sp.Atom?
+ def __new__(cls, *args, **kwargs):
+ return sp.Basic.__new__(cls)
+
+ def _sympystr(self, *args, **kwargs):
+ return str(self)
+
+ @property
+ @abstractmethod
+ def base_type(self) -> Union[None, 'BasicType']:
+ """
+ Returns: Returns BasicType of a Vector or Pointer type, None otherwise
+ """
+ pass
+
+ @property
+ @abstractmethod
+ def item_size(self) -> int:
+ """
+ Returns: Number of items.
+ E.g. width * item_size(basic_type) in vector's case, or simple numpy itemsize in Struct's case.
+ """
+ pass
+
+
+def is_supported_type(dtype: np.dtype):
+ scalar = dtype.type
+ c = np.issctype(dtype)
+ subclass = issubclass(scalar, np.floating) or issubclass(scalar, np.integer) or issubclass(scalar, np.bool_)
+ additional_checks = dtype.fields is None and dtype.hasobject is False and dtype.subdtype is None
+ return c and subclass and additional_checks
+
+
+def numpy_name_to_c(name: str) -> str:
+ """
+ Converts a np.dtype.name into a C type
+ Args:
+ name: np.dtype.name string
+ Returns:
+ type as a C string
+ """
+ if name == 'float64':
+ return 'double'
+ elif name == 'float32':
+ return 'float'
+ elif name == 'float16' or name == 'half':
+ return 'half'
+ elif name.startswith('int'):
+ width = int(name[len("int"):])
+ return f"int{width}_t"
+ elif name.startswith('uint'):
+ width = int(name[len("uint"):])
+ return f"uint{width}_t"
+ elif name == 'bool':
+ return 'bool'
+ else:
+ raise NotImplementedError(f"Can't map numpy to C name for {name}")
+
+
+def create_type(specification: Union[type, AbstractType, str]) -> AbstractType:
+ # TODO: Deprecated Use the constructor of BasicType or StructType instead
+ """Creates a subclass of Type according to a string or an object of subclass Type.
+
+ Args:
+ specification: Type object, or a string
+
+ Returns:
+ Type object, or a new Type object parsed from the string
+ """
+ if isinstance(specification, AbstractType):
+ return specification
+ else:
+ numpy_dtype = np.dtype(specification)
+ if numpy_dtype.fields is None:
+ return BasicType(numpy_dtype, const=False)
+ else:
+ return StructType(numpy_dtype, const=False)
+
+
+def get_base_type(data_type):
+ """
+ Returns the BasicType of a Pointer or a Vector
+ """
+ while data_type.base_type is not None:
+ data_type = data_type.base_type
+ return data_type
+
+
+class BasicType(AbstractType):
+ """
+ BasicType is defined with a const qualifier and a np.dtype.
+ """
+
+ def __init__(self, dtype: Union[type, 'BasicType', str], const: bool = False):
+ if isinstance(dtype, BasicType):
+ self.numpy_dtype = dtype.numpy_dtype
+ self.const = dtype.const
+ else:
+ self.numpy_dtype = np.dtype(dtype)
+ self.const = const
+ assert is_supported_type(self.numpy_dtype), f'Type {self.numpy_dtype} is currently not supported!'
+
+ def __getnewargs__(self):
+ return self.numpy_dtype, self.const
+
+ def __getnewargs_ex__(self):
+ return (self.numpy_dtype, self.const), {}
+
+ @property
+ def base_type(self):
+ return None
+
+ @property
+ def item_size(self): # TODO: Do we want self.numpy_type.itemsize????
+ return 1
+
+ def is_float(self):
+ return issubclass(self.numpy_dtype.type, np.floating)
+
+ def is_half(self):
+ return issubclass(self.numpy_dtype.type, np.half)
+
+ def is_int(self):
+ return issubclass(self.numpy_dtype.type, np.integer)
+
+ def is_uint(self):
+ return issubclass(self.numpy_dtype.type, np.unsignedinteger)
+
+ def is_sint(self):
+ return issubclass(self.numpy_dtype.type, np.signedinteger)
+
+ def is_bool(self):
+ return issubclass(self.numpy_dtype.type, np.bool_)
+
+ def dtype_eq(self, other):
+ if not isinstance(other, BasicType):
+ return False
+ else:
+ return self.numpy_dtype == other.numpy_dtype
+
+ @property
+ def c_name(self) -> str:
+ return numpy_name_to_c(self.numpy_dtype.name)
+
+ def __str__(self):
+ return f'{self.c_name}{" const" if self.const else ""}'
+
+ def __repr__(self):
+ return f'BasicType( {str(self)} )'
+
+ def _repr_html_(self):
+ return f'BasicType( {str(self)} )'
+
+ def __eq__(self, other):
+ return self.dtype_eq(other) and self.const == other.const
+
+ def __hash__(self):
+ return hash(str(self))
+
+
+class VectorType(AbstractType):
+ """
+ VectorType consists of a BasicType and a width.
+ """
+ instruction_set = None
+
+ def __init__(self, base_type: BasicType, width: int):
+ self._base_type = base_type
+ self.width = width
+
+ @property
+ def base_type(self):
+ return self._base_type
+
+ @property
+ def item_size(self):
+ return self.width * self.base_type.item_size
+
+ def __eq__(self, other):
+ if not isinstance(other, VectorType):
+ return False
+ else:
+ return (self.base_type, self.width) == (other.base_type, other.width)
+
+ def __str__(self):
+ if self.instruction_set is None:
+ return f"{self.base_type}[{self.width}]"
+ else:
+ # TODO VectorizationRevamp: this seems super weird. the instruction_set should know how to print a type out!
+ # TODO VectorizationRevamp: this is error prone. base_type could be cons=True. Use dtype instead
+ if self.base_type == create_type("int64") or self.base_type == create_type("int32"):
+ return self.instruction_set['int']
+ elif self.base_type == create_type("float64"):
+ return self.instruction_set['double']
+ elif self.base_type == create_type("float32"):
+ return self.instruction_set['float']
+ elif self.base_type == create_type("bool"):
+ return self.instruction_set['bool']
+ else:
+ raise NotImplementedError()
+
+ def __hash__(self):
+ return hash((self.base_type, self.width))
+
+ def __getnewargs__(self):
+ return self._base_type, self.width
+
+ def __getnewargs_ex__(self):
+ return (self._base_type, self.width), {}
+
+
+class PointerType(AbstractType):
+ def __init__(self, base_type: BasicType, const: bool = False, restrict: bool = True, double_pointer: bool = False):
+ self._base_type = base_type
+ self.const = const
+ self.restrict = restrict
+ self.double_pointer = double_pointer
+
+ def __getnewargs__(self):
+ return self.base_type, self.const, self.restrict, self.double_pointer
+
+ def __getnewargs_ex__(self):
+ return (self.base_type, self.const, self.restrict, self.double_pointer), {}
+
+ @property
+ def alias(self):
+ return not self.restrict
+
+ @property
+ def base_type(self):
+ return self._base_type
+
+ @property
+ def item_size(self):
+ if self.double_pointer:
+ raise NotImplementedError("The item_size for double_pointer is not implemented")
+ else:
+ return self.base_type.item_size
+
+ def __eq__(self, other):
+ if not isinstance(other, PointerType):
+ return False
+ else:
+ own = (self.base_type, self.const, self.restrict, self.double_pointer)
+ return own == (other.base_type, other.const, other.restrict, other.double_pointer)
+
+ def __str__(self):
+ restrict_str = "RESTRICT" if self.restrict else ""
+ const_str = "const" if self.const else ""
+ if self.double_pointer:
+ return f'{str(self.base_type)} ** {restrict_str} {const_str}'
+ else:
+ return f'{str(self.base_type)} * {restrict_str} {const_str}'
+
+ def __repr__(self):
+ return str(self)
+
+ def _repr_html_(self):
+ return str(self)
+
+ def __hash__(self):
+ return hash((self._base_type, self.const, self.restrict, self.double_pointer))
+
+
+class StructType(AbstractType):
+ """
+ A list of types (with C offsets).
+ It is implemented with uint8_t and casts to the correct datatype.
+ """
+ def __init__(self, numpy_type, const=False):
+ self.const = const
+ self._dtype = np.dtype(numpy_type)
+
+ def __getnewargs__(self):
+ return self.numpy_dtype, self.const
+
+ def __getnewargs_ex__(self):
+ return (self.numpy_dtype, self.const), {}
+
+ @property
+ def base_type(self):
+ return None
+
+ @property
+ def numpy_dtype(self):
+ return self._dtype
+
+ @property
+ def item_size(self):
+ return self.numpy_dtype.itemsize
+
+ def get_element_offset(self, element_name):
+ return self.numpy_dtype.fields[element_name][1]
+
+ def get_element_type(self, element_name):
+ np_element_type = self.numpy_dtype.fields[element_name][0]
+ return BasicType(np_element_type, self.const)
+
+ def has_element(self, element_name):
+ return element_name in self.numpy_dtype.fields
+
+ def __eq__(self, other):
+ if not isinstance(other, StructType):
+ return False
+ else:
+ return (self.numpy_dtype, self.const) == (other.numpy_dtype, other.const)
+
+ def __str__(self):
+ # structs are handled byte-wise
+ result = "uint8_t"
+ if self.const:
+ result += " const"
+ return result
+
+ def __repr__(self):
+ return str(self)
+
+ def _repr_html_(self):
+ return str(self)
+
+ def __hash__(self):
+ return hash((self.numpy_dtype, self.const))
+
+
+def assumptions_from_dtype(dtype: Union[BasicType, np.dtype]):
+ """Derives SymPy assumptions from :class:`BasicType` or a Numpy dtype
+
+ Args:
+ dtype (BasicType, np.dtype): a Numpy data type
+ Returns:
+ A dict of SymPy assumptions
+ """
+ if hasattr(dtype, 'numpy_dtype'):
+ dtype = dtype.numpy_dtype
+
+ assumptions = dict()
+
+ try:
+ if np.issubdtype(dtype, np.integer):
+ assumptions.update({'integer': True})
+
+ if np.issubdtype(dtype, np.unsignedinteger):
+ assumptions.update({'negative': False})
+
+ if np.issubdtype(dtype, np.integer) or \
+ np.issubdtype(dtype, np.floating):
+ assumptions.update({'real': True})
+ except Exception: # TODO this is dirty
+ pass
+
+ return assumptions
+
+
+class TypedSymbol(sp.Symbol):
+ def __new__(cls, *args, **kwds):
+ obj = TypedSymbol.__xnew_cached_(cls, *args, **kwds)
+ return obj
+
+ def __new_stage2__(cls, name, dtype, **kwargs): # TODO does not match signature of sp.Symbol???
+ # TODO: also Symbol should be allowed ---> see sympy Variable
+ assumptions = assumptions_from_dtype(dtype)
+ assumptions.update(kwargs)
+ obj = super(TypedSymbol, cls).__xnew__(cls, name, **assumptions)
+ try:
+ obj.numpy_dtype = create_type(dtype)
+ except (TypeError, ValueError):
+ # on error keep the string
+ obj.numpy_dtype = dtype
+ return obj
+
+ __xnew__ = staticmethod(__new_stage2__)
+ __xnew_cached_ = staticmethod(sp.core.cacheit(__new_stage2__))
+
+ @property
+ def dtype(self):
+ return self.numpy_dtype
+
+ def _hashable_content(self):
+ return super()._hashable_content(), hash(self.numpy_dtype)
+
+ def __getnewargs__(self):
+ return self.name, self.dtype
+
+ def __getnewargs_ex__(self):
+ return (self.name, self.dtype), self.assumptions0
+
+ @property
+ def canonical(self):
+ return self
+
+ @property
+ def reversed(self):
+ return self
+
+ @property
+ def headers(self):
+ headers = []
+ try:
+ if np.issubdtype(self.dtype.numpy_dtype, np.complexfloating):
+ headers.append('"cuda_complex.hpp"')
+ except Exception:
+ pass
+ try:
+ if np.issubdtype(self.dtype.base_type.numpy_dtype, np.complexfloating):
+ headers.append('"cuda_complex.hpp"')
+ except Exception:
+ pass
+
+ return headers
+
+
+SHAPE_DTYPE = BasicType('int64', const=True)
+STRIDE_DTYPE = BasicType('int64', const=True)
+
+
+class FieldStrideSymbol(TypedSymbol):
+ """Sympy symbol representing the stride value of a field in a specific coordinate."""
+ def __new__(cls, *args, **kwds):
+ obj = FieldStrideSymbol.__xnew_cached_(cls, *args, **kwds)
+ return obj
+
+ def __new_stage2__(cls, field_name, coordinate):
+ name = f"_stride_{field_name}_{coordinate}"
+ obj = super(FieldStrideSymbol, cls).__xnew__(cls, name, STRIDE_DTYPE, positive=True)
+ obj.field_name = field_name
+ obj.coordinate = coordinate
+ return obj
+
+ def __getnewargs__(self):
+ return self.field_name, self.coordinate
+
+ def __getnewargs_ex__(self):
+ return (self.field_name, self.coordinate), {}
+
+ __xnew__ = staticmethod(__new_stage2__)
+ __xnew_cached_ = staticmethod(sp.core.cacheit(__new_stage2__))
+
+ def _hashable_content(self):
+ return super()._hashable_content(), self.coordinate, self.field_name
+
+
+class FieldShapeSymbol(TypedSymbol):
+ """Sympy symbol representing the shape value of a sequence of fields. In a kernel iterating over multiple fields
+ there is only one set of `FieldShapeSymbol`s since all the fields have to be of equal size."""
+ def __new__(cls, *args, **kwds):
+ obj = FieldShapeSymbol.__xnew_cached_(cls, *args, **kwds)
+ return obj
+
+ def __new_stage2__(cls, field_names, coordinate):
+ names = "_".join([field_name for field_name in field_names])
+ name = f"_size_{names}_{coordinate}"
+ obj = super(FieldShapeSymbol, cls).__xnew__(cls, name, SHAPE_DTYPE, positive=True)
+ obj.field_names = tuple(field_names)
+ obj.coordinate = coordinate
+ return obj
+
+ def __getnewargs__(self):
+ return self.field_names, self.coordinate
+
+ def __getnewargs_ex__(self):
+ return (self.field_names, self.coordinate), {}
+
+ __xnew__ = staticmethod(__new_stage2__)
+ __xnew_cached_ = staticmethod(sp.core.cacheit(__new_stage2__))
+
+ def _hashable_content(self):
+ return super()._hashable_content(), self.coordinate, self.field_names
+
+
+class FieldPointerSymbol(TypedSymbol):
+ """Sympy symbol representing the pointer to the beginning of the field data."""
+ def __new__(cls, *args, **kwds):
+ obj = FieldPointerSymbol.__xnew_cached_(cls, *args, **kwds)
+ return obj
+
+ def __new_stage2__(cls, field_name, field_dtype, const):
+ name = f"_data_{field_name}"
+ dtype = PointerType(get_base_type(field_dtype), const=const, restrict=True)
+ obj = super(FieldPointerSymbol, cls).__xnew__(cls, name, dtype)
+ obj.field_name = field_name
+ return obj
+
+ def __getnewargs__(self):
+ return self.field_name, self.dtype, self.dtype.const
+
+ def __getnewargs_ex__(self):
+ return (self.field_name, self.dtype, self.dtype.const), {}
+
+ def _hashable_content(self):
+ return super()._hashable_content(), self.field_name
+
+ __xnew__ = staticmethod(__new_stage2__)
+ __xnew_cached_ = staticmethod(sp.core.cacheit(__new_stage2__))
+
+
+class CastFunc(sp.Function):
+ """
+ CastFunc is used in order to introduce static casts. They are especially useful as a way to signal what type
+ a certain node should have, if it is impossible to add a type to a node, e.g. a sp.Number.
+ """
+ is_Atom = True
+
+ def __new__(cls, *args, **kwargs):
+ if len(args) != 2:
+ pass
+ expr, dtype, *other_args = args
+
+ # If we have two consecutive casts, throw the inner one away.
+ # This optimisation is only available for simple casts. Thus the == is intended here!
+ if expr.__class__ == CastFunc:
+ expr = expr.args[0]
+ if not isinstance(dtype, AbstractType):
+ dtype = BasicType(dtype)
+ # to work in conditions of sp.Piecewise cast_func has to be of type Boolean as well
+ # however, a cast_function should only be a boolean if its argument is a boolean, otherwise this leads
+ # to problems when for example comparing cast_func's for equality
+ #
+ # lhs = bitwise_and(a, cast_func(1, 'int'))
+ # rhs = cast_func(0, 'int')
+ # print( sp.Ne(lhs, rhs) ) # would give true if all cast_funcs are booleans
+ # -> thus a separate class boolean_cast_func is introduced
+ if (isinstance(expr, sp.logic.boolalg.Boolean) and
+ (not isinstance(expr, TypedSymbol) or expr.dtype == BasicType('bool'))):
+ cls = BooleanCastFunc
+
+ return sp.Function.__new__(cls, expr, dtype, *other_args, **kwargs)
+
+ @property
+ def canonical(self):
+ if hasattr(self.args[0], 'canonical'):
+ return self.args[0].canonical
+ else:
+ raise NotImplementedError()
+
+ @property
+ def is_commutative(self):
+ return self.args[0].is_commutative
+
+ @property
+ def dtype(self):
+ return self.args[1]
+
+ @property
+ def expr(self):
+ return self.args[0]
+
+ @property
+ def is_integer(self):
+ """
+ Uses Numpy type hierarchy to determine :func:`sympy.Expr.is_integer` predicate
+
+ For reference: Numpy type hierarchy https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.scalars.html
+ """
+ if hasattr(self.dtype, 'numpy_dtype'):
+ return np.issubdtype(self.dtype.numpy_dtype, np.integer) or super().is_integer
+ else:
+ return super().is_integer
+
+ @property
+ def is_negative(self):
+ """
+ See :func:`.TypedSymbol.is_integer`
+ """
+ if hasattr(self.dtype, 'numpy_dtype'):
+ if np.issubdtype(self.dtype.numpy_dtype, np.unsignedinteger):
+ return False
+
+ return super().is_negative
+
+ @property
+ def is_nonnegative(self):
+ """
+ See :func:`.TypedSymbol.is_integer`
+ """
+ if self.is_negative is False:
+ return True
+ else:
+ return super().is_nonnegative
+
+ @property
+ def is_real(self):
+ """
+ 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
+ else:
+ return super().is_real
+
+
+class BooleanCastFunc(CastFunc, sp.logic.boolalg.Boolean):
+ # TODO: documentation
+ pass
+
+
+class VectorMemoryAccess(CastFunc):
+ """
+ Special memory access for vectorized kernel.
+ Arguments: read/write expression, type, aligned, non-temporal, mask (or none), stride
+ """
+ nargs = (6,)
+
+
+class ReinterpretCastFunc(CastFunc):
+ """
+ Reinterpret cast is necessary for the StructType
+ """
+ pass
+
+
+class PointerArithmeticFunc(sp.Function, sp.logic.boolalg.Boolean):
+ # TODO: documentation, or deprecate!
+ @property
+ def canonical(self):
+ if hasattr(self.args[0], 'canonical'):
+ return self.args[0].canonical
+ else:
+ raise NotImplementedError()
+
+
+
diff --git a/src/pystencils/typing/cast_functions.py b/src/pystencils/typing/cast_functions.py
deleted file mode 100644
index 1b83d223cbff2ce08c1fc0516d2ce53dc2ec350a..0000000000000000000000000000000000000000
--- a/src/pystencils/typing/cast_functions.py
+++ /dev/null
@@ -1,131 +0,0 @@
-import numpy as np
-import sympy as sp
-from sympy.logic.boolalg import Boolean
-
-from pystencils.typing.types import AbstractType, BasicType
-from pystencils.typing.typed_sympy import TypedSymbol
-
-
-class CastFunc(sp.Function):
- """
- CastFunc is used in order to introduce static casts. They are especially useful as a way to signal what type
- a certain node should have, if it is impossible to add a type to a node, e.g. a sp.Number.
- """
- is_Atom = True
-
- def __new__(cls, *args, **kwargs):
- if len(args) != 2:
- pass
- expr, dtype, *other_args = args
-
- # If we have two consecutive casts, throw the inner one away.
- # This optimisation is only available for simple casts. Thus the == is intended here!
- if expr.__class__ == CastFunc:
- expr = expr.args[0]
- if not isinstance(dtype, AbstractType):
- dtype = BasicType(dtype)
- # to work in conditions of sp.Piecewise cast_func has to be of type Boolean as well
- # however, a cast_function should only be a boolean if its argument is a boolean, otherwise this leads
- # to problems when for example comparing cast_func's for equality
- #
- # lhs = bitwise_and(a, cast_func(1, 'int'))
- # rhs = cast_func(0, 'int')
- # print( sp.Ne(lhs, rhs) ) # would give true if all cast_funcs are booleans
- # -> thus a separate class boolean_cast_func is introduced
- if isinstance(expr, Boolean) and (not isinstance(expr, TypedSymbol) or expr.dtype == BasicType('bool')):
- cls = BooleanCastFunc
-
- return sp.Function.__new__(cls, expr, dtype, *other_args, **kwargs)
-
- @property
- def canonical(self):
- if hasattr(self.args[0], 'canonical'):
- return self.args[0].canonical
- else:
- raise NotImplementedError()
-
- @property
- def is_commutative(self):
- return self.args[0].is_commutative
-
- @property
- def dtype(self):
- return self.args[1]
-
- @property
- def expr(self):
- return self.args[0]
-
- @property
- def is_integer(self):
- """
- Uses Numpy type hierarchy to determine :func:`sympy.Expr.is_integer` predicate
-
- For reference: Numpy type hierarchy https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.scalars.html
- """
- if hasattr(self.dtype, 'numpy_dtype'):
- return np.issubdtype(self.dtype.numpy_dtype, np.integer) or super().is_integer
- else:
- return super().is_integer
-
- @property
- def is_negative(self):
- """
- See :func:`.TypedSymbol.is_integer`
- """
- if hasattr(self.dtype, 'numpy_dtype'):
- if np.issubdtype(self.dtype.numpy_dtype, np.unsignedinteger):
- return False
-
- return super().is_negative
-
- @property
- def is_nonnegative(self):
- """
- See :func:`.TypedSymbol.is_integer`
- """
- if self.is_negative is False:
- return True
- else:
- return super().is_nonnegative
-
- @property
- def is_real(self):
- """
- 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
- else:
- return super().is_real
-
-
-class BooleanCastFunc(CastFunc, Boolean):
- # TODO: documentation
- pass
-
-
-class VectorMemoryAccess(CastFunc):
- """
- Special memory access for vectorized kernel.
- Arguments: read/write expression, type, aligned, non-temporal, mask (or none), stride
- """
- nargs = (6,)
-
-
-class ReinterpretCastFunc(CastFunc):
- """
- Reinterpret cast is necessary for the StructType
- """
- pass
-
-
-class PointerArithmeticFunc(sp.Function, Boolean):
- # TODO: documentation, or deprecate!
- @property
- def canonical(self):
- if hasattr(self.args[0], 'canonical'):
- return self.args[0].canonical
- else:
- raise NotImplementedError()
diff --git a/src/pystencils/typing/typed_sympy.py b/src/pystencils/typing/typed_sympy.py
deleted file mode 100644
index 302c2f9987b2db1a907710678ddbb7234668cfc6..0000000000000000000000000000000000000000
--- a/src/pystencils/typing/typed_sympy.py
+++ /dev/null
@@ -1,180 +0,0 @@
-from typing import Union
-
-import numpy as np
-import sympy as sp
-from sympy.core.cache import cacheit
-
-from pystencils.typing.types import BasicType, create_type, PointerType
-
-
-def assumptions_from_dtype(dtype: Union[BasicType, np.dtype]):
- """Derives SymPy assumptions from :class:`BasicType` or a Numpy dtype
-
- Args:
- dtype (BasicType, np.dtype): a Numpy data type
- Returns:
- A dict of SymPy assumptions
- """
- if hasattr(dtype, 'numpy_dtype'):
- dtype = dtype.numpy_dtype
-
- assumptions = dict()
-
- try:
- if np.issubdtype(dtype, np.integer):
- assumptions.update({'integer': True})
-
- if np.issubdtype(dtype, np.unsignedinteger):
- assumptions.update({'negative': False})
-
- if np.issubdtype(dtype, np.integer) or \
- np.issubdtype(dtype, np.floating):
- assumptions.update({'real': True})
- except Exception: # TODO this is dirty
- pass
-
- return assumptions
-
-
-class TypedSymbol(sp.Symbol):
- def __new__(cls, *args, **kwds):
- obj = TypedSymbol.__xnew_cached_(cls, *args, **kwds)
- return obj
-
- def __new_stage2__(cls, name, dtype, **kwargs): # TODO does not match signature of sp.Symbol???
- # TODO: also Symbol should be allowed ---> see sympy Variable
- assumptions = assumptions_from_dtype(dtype)
- assumptions.update(kwargs)
- obj = super(TypedSymbol, cls).__xnew__(cls, name, **assumptions)
- try:
- obj.numpy_dtype = create_type(dtype)
- except (TypeError, ValueError):
- # on error keep the string
- obj.numpy_dtype = dtype
- return obj
-
- __xnew__ = staticmethod(__new_stage2__)
- __xnew_cached_ = staticmethod(cacheit(__new_stage2__))
-
- @property
- def dtype(self):
- return self.numpy_dtype
-
- def _hashable_content(self):
- return super()._hashable_content(), hash(self.numpy_dtype)
-
- def __getnewargs__(self):
- return self.name, self.dtype
-
- def __getnewargs_ex__(self):
- return (self.name, self.dtype), self.assumptions0
-
- @property
- def canonical(self):
- return self
-
- @property
- def reversed(self):
- return self
-
- @property
- def headers(self):
- headers = []
- try:
- if np.issubdtype(self.dtype.numpy_dtype, np.complexfloating):
- headers.append('"cuda_complex.hpp"')
- except Exception:
- pass
- try:
- if np.issubdtype(self.dtype.base_type.numpy_dtype, np.complexfloating):
- headers.append('"cuda_complex.hpp"')
- except Exception:
- pass
-
- return headers
-
-
-SHAPE_DTYPE = BasicType('int64', const=True)
-STRIDE_DTYPE = BasicType('int64', const=True)
-
-
-class FieldStrideSymbol(TypedSymbol):
- """Sympy symbol representing the stride value of a field in a specific coordinate."""
- def __new__(cls, *args, **kwds):
- obj = FieldStrideSymbol.__xnew_cached_(cls, *args, **kwds)
- return obj
-
- def __new_stage2__(cls, field_name, coordinate):
- name = f"_stride_{field_name}_{coordinate}"
- obj = super(FieldStrideSymbol, cls).__xnew__(cls, name, STRIDE_DTYPE, positive=True)
- obj.field_name = field_name
- obj.coordinate = coordinate
- return obj
-
- def __getnewargs__(self):
- return self.field_name, self.coordinate
-
- def __getnewargs_ex__(self):
- return (self.field_name, self.coordinate), {}
-
- __xnew__ = staticmethod(__new_stage2__)
- __xnew_cached_ = staticmethod(cacheit(__new_stage2__))
-
- def _hashable_content(self):
- return super()._hashable_content(), self.coordinate, self.field_name
-
-
-class FieldShapeSymbol(TypedSymbol):
- """Sympy symbol representing the shape value of a sequence of fields. In a kernel iterating over multiple fields
- there is only one set of `FieldShapeSymbol`s since all the fields have to be of equal size."""
- def __new__(cls, *args, **kwds):
- obj = FieldShapeSymbol.__xnew_cached_(cls, *args, **kwds)
- return obj
-
- def __new_stage2__(cls, field_names, coordinate):
- names = "_".join([field_name for field_name in field_names])
- name = f"_size_{names}_{coordinate}"
- obj = super(FieldShapeSymbol, cls).__xnew__(cls, name, SHAPE_DTYPE, positive=True)
- obj.field_names = tuple(field_names)
- obj.coordinate = coordinate
- return obj
-
- def __getnewargs__(self):
- return self.field_names, self.coordinate
-
- def __getnewargs_ex__(self):
- return (self.field_names, self.coordinate), {}
-
- __xnew__ = staticmethod(__new_stage2__)
- __xnew_cached_ = staticmethod(cacheit(__new_stage2__))
-
- def _hashable_content(self):
- return super()._hashable_content(), self.coordinate, self.field_names
-
-
-class FieldPointerSymbol(TypedSymbol):
- """Sympy symbol representing the pointer to the beginning of the field data."""
- def __new__(cls, *args, **kwds):
- obj = FieldPointerSymbol.__xnew_cached_(cls, *args, **kwds)
- return obj
-
- def __new_stage2__(cls, field_name, field_dtype, const):
- from pystencils.typing.utilities import get_base_type
-
- name = f"_data_{field_name}"
- dtype = PointerType(get_base_type(field_dtype), const=const, restrict=True)
- obj = super(FieldPointerSymbol, cls).__xnew__(cls, name, dtype)
- obj.field_name = field_name
- return obj
-
- def __getnewargs__(self):
- return self.field_name, self.dtype, self.dtype.const
-
- def __getnewargs_ex__(self):
- return (self.field_name, self.dtype, self.dtype.const), {}
-
- def _hashable_content(self):
- return super()._hashable_content(), self.field_name
-
- __xnew__ = staticmethod(__new_stage2__)
- __xnew_cached_ = staticmethod(cacheit(__new_stage2__))
diff --git a/src/pystencils/typing/types.py b/src/pystencils/typing/types.py
deleted file mode 100644
index 4d80daffae7bd0ca2cf4787c096ef74c51707400..0000000000000000000000000000000000000000
--- a/src/pystencils/typing/types.py
+++ /dev/null
@@ -1,321 +0,0 @@
-from abc import abstractmethod
-from typing import Union
-
-import numpy as np
-import sympy as sp
-
-
-def is_supported_type(dtype: np.dtype):
- scalar = dtype.type
- c = np.issctype(dtype)
- subclass = issubclass(scalar, np.floating) or issubclass(scalar, np.integer) or issubclass(scalar, np.bool_)
- additional_checks = dtype.fields is None and dtype.hasobject is False and dtype.subdtype is None
- return c and subclass and additional_checks
-
-
-def numpy_name_to_c(name: str) -> str:
- """
- Converts a np.dtype.name into a C type
- Args:
- name: np.dtype.name string
- Returns:
- type as a C string
- """
- if name == 'float64':
- return 'double'
- elif name == 'float32':
- return 'float'
- elif name == 'float16' or name == 'half':
- return 'half'
- elif name.startswith('int'):
- width = int(name[len("int"):])
- return f"int{width}_t"
- elif name.startswith('uint'):
- width = int(name[len("uint"):])
- return f"uint{width}_t"
- elif name == 'bool':
- return 'bool'
- else:
- raise NotImplementedError(f"Can't map numpy to C name for {name}")
-
-
-class AbstractType(sp.Atom):
- # TODO: Is it necessary to ineherit from sp.Atom?
- def __new__(cls, *args, **kwargs):
- return sp.Basic.__new__(cls)
-
- def _sympystr(self, *args, **kwargs):
- return str(self)
-
- @property
- @abstractmethod
- def base_type(self) -> Union[None, 'BasicType']:
- """
- Returns: Returns BasicType of a Vector or Pointer type, None otherwise
- """
- pass
-
- @property
- @abstractmethod
- def item_size(self) -> int:
- """
- Returns: Number of items.
- E.g. width * item_size(basic_type) in vector's case, or simple numpy itemsize in Struct's case.
- """
- pass
-
-
-class BasicType(AbstractType):
- """
- BasicType is defined with a const qualifier and a np.dtype.
- """
-
- def __init__(self, dtype: Union[type, 'BasicType', str], const: bool = False):
- if isinstance(dtype, BasicType):
- self.numpy_dtype = dtype.numpy_dtype
- self.const = dtype.const
- else:
- self.numpy_dtype = np.dtype(dtype)
- self.const = const
- assert is_supported_type(self.numpy_dtype), f'Type {self.numpy_dtype} is currently not supported!'
-
- def __getnewargs__(self):
- return self.numpy_dtype, self.const
-
- def __getnewargs_ex__(self):
- return (self.numpy_dtype, self.const), {}
-
- @property
- def base_type(self):
- return None
-
- @property
- def item_size(self): # TODO: Do we want self.numpy_type.itemsize????
- return 1
-
- def is_float(self):
- return issubclass(self.numpy_dtype.type, np.floating)
-
- def is_half(self):
- return issubclass(self.numpy_dtype.type, np.half)
-
- def is_int(self):
- return issubclass(self.numpy_dtype.type, np.integer)
-
- def is_uint(self):
- return issubclass(self.numpy_dtype.type, np.unsignedinteger)
-
- def is_sint(self):
- return issubclass(self.numpy_dtype.type, np.signedinteger)
-
- def is_bool(self):
- return issubclass(self.numpy_dtype.type, np.bool_)
-
- def dtype_eq(self, other):
- if not isinstance(other, BasicType):
- return False
- else:
- return self.numpy_dtype == other.numpy_dtype
-
- @property
- def c_name(self) -> str:
- return numpy_name_to_c(self.numpy_dtype.name)
-
- def __str__(self):
- return f'{self.c_name}{" const" if self.const else ""}'
-
- def __repr__(self):
- return f'BasicType( {str(self)} )'
-
- def _repr_html_(self):
- return f'BasicType( {str(self)} )'
-
- def __eq__(self, other):
- return self.dtype_eq(other) and self.const == other.const
-
- def __hash__(self):
- return hash(str(self))
-
-
-class VectorType(AbstractType):
- """
- VectorType consists of a BasicType and a width.
- """
- instruction_set = None
-
- def __init__(self, base_type: BasicType, width: int):
- self._base_type = base_type
- self.width = width
-
- @property
- def base_type(self):
- return self._base_type
-
- @property
- def item_size(self):
- return self.width * self.base_type.item_size
-
- def __eq__(self, other):
- if not isinstance(other, VectorType):
- return False
- else:
- return (self.base_type, self.width) == (other.base_type, other.width)
-
- def __str__(self):
- if self.instruction_set is None:
- return f"{self.base_type}[{self.width}]"
- else:
- # TODO VectorizationRevamp: this seems super weird. the instruction_set should know how to print a type out!
- # TODO VectorizationRevamp: this is error prone. base_type could be cons=True. Use dtype instead
- if self.base_type == create_type("int64") or self.base_type == create_type("int32"):
- return self.instruction_set['int']
- elif self.base_type == create_type("float64"):
- return self.instruction_set['double']
- elif self.base_type == create_type("float32"):
- return self.instruction_set['float']
- elif self.base_type == create_type("bool"):
- return self.instruction_set['bool']
- else:
- raise NotImplementedError()
-
- def __hash__(self):
- return hash((self.base_type, self.width))
-
- def __getnewargs__(self):
- return self._base_type, self.width
-
- def __getnewargs_ex__(self):
- return (self._base_type, self.width), {}
-
-
-class PointerType(AbstractType):
- def __init__(self, base_type: BasicType, const: bool = False, restrict: bool = True, double_pointer: bool = False):
- self._base_type = base_type
- self.const = const
- self.restrict = restrict
- self.double_pointer = double_pointer
-
- def __getnewargs__(self):
- return self.base_type, self.const, self.restrict, self.double_pointer
-
- def __getnewargs_ex__(self):
- return (self.base_type, self.const, self.restrict, self.double_pointer), {}
-
- @property
- def alias(self):
- return not self.restrict
-
- @property
- def base_type(self):
- return self._base_type
-
- @property
- def item_size(self):
- if self.double_pointer:
- raise NotImplementedError("The item_size for double_pointer is not implemented")
- else:
- return self.base_type.item_size
-
- def __eq__(self, other):
- if not isinstance(other, PointerType):
- return False
- else:
- own = (self.base_type, self.const, self.restrict, self.double_pointer)
- return own == (other.base_type, other.const, other.restrict, other.double_pointer)
-
- def __str__(self):
- restrict_str = "RESTRICT" if self.restrict else ""
- const_str = "const" if self.const else ""
- if self.double_pointer:
- return f'{str(self.base_type)} ** {restrict_str} {const_str}'
- else:
- return f'{str(self.base_type)} * {restrict_str} {const_str}'
-
- def __repr__(self):
- return str(self)
-
- def _repr_html_(self):
- return str(self)
-
- def __hash__(self):
- return hash((self._base_type, self.const, self.restrict, self.double_pointer))
-
-
-class StructType(AbstractType):
- """
- A list of types (with C offsets).
- It is implemented with uint8_t and casts to the correct datatype.
- """
- def __init__(self, numpy_type, const=False):
- self.const = const
- self._dtype = np.dtype(numpy_type)
-
- def __getnewargs__(self):
- return self.numpy_dtype, self.const
-
- def __getnewargs_ex__(self):
- return (self.numpy_dtype, self.const), {}
-
- @property
- def base_type(self):
- return None
-
- @property
- def numpy_dtype(self):
- return self._dtype
-
- @property
- def item_size(self):
- return self.numpy_dtype.itemsize
-
- def get_element_offset(self, element_name):
- return self.numpy_dtype.fields[element_name][1]
-
- def get_element_type(self, element_name):
- np_element_type = self.numpy_dtype.fields[element_name][0]
- return BasicType(np_element_type, self.const)
-
- def has_element(self, element_name):
- return element_name in self.numpy_dtype.fields
-
- def __eq__(self, other):
- if not isinstance(other, StructType):
- return False
- else:
- return (self.numpy_dtype, self.const) == (other.numpy_dtype, other.const)
-
- def __str__(self):
- # structs are handled byte-wise
- result = "uint8_t"
- if self.const:
- result += " const"
- return result
-
- def __repr__(self):
- return str(self)
-
- def _repr_html_(self):
- return str(self)
-
- def __hash__(self):
- return hash((self.numpy_dtype, self.const))
-
-
-def create_type(specification: Union[type, AbstractType, str]) -> AbstractType:
- # TODO: Deprecated Use the constructor of BasicType or StructType instead
- """Creates a subclass of Type according to a string or an object of subclass Type.
-
- Args:
- specification: Type object, or a string
-
- Returns:
- Type object, or a new Type object parsed from the string
- """
- if isinstance(specification, AbstractType):
- return specification
- else:
- numpy_dtype = np.dtype(specification)
- if numpy_dtype.fields is None:
- return BasicType(numpy_dtype, const=False)
- else:
- return StructType(numpy_dtype, const=False)
diff --git a/tests/nbackend/kernelcreation/test_domain_kernels.py b/tests/nbackend/kernelcreation/test_domain_kernels.py
index 5a3649d68ad1b2f225501990dcb6d21af0bd9956..4221bbad18849b1148139aa350d9db999f511f9d 100644
--- a/tests/nbackend/kernelcreation/test_domain_kernels.py
+++ b/tests/nbackend/kernelcreation/test_domain_kernels.py
@@ -1,13 +1,12 @@
-import pytest
-
import sympy as sp
import numpy as np
from pystencils import fields, Field, AssignmentCollection
-from pystencils.assignment import assignment_from_stencil
+from pystencils.sympyextensions.assignmentcollection.assignment import assignment_from_stencil
from pystencils.nbackend.kernelcreation import create_kernel
+
def test_filter_kernel():
weight = sp.Symbol("weight")
stencil = [
diff --git a/tests/nbackend/kernelcreation/test_index_kernels.py b/tests/nbackend/kernelcreation/test_index_kernels.py
index 4c4c6b8b816114964ae0ba7f3c77e39add292f7e..817b4a244898f3061d37a0b69483b5a79253aad2 100644
--- a/tests/nbackend/kernelcreation/test_index_kernels.py
+++ b/tests/nbackend/kernelcreation/test_index_kernels.py
@@ -1,10 +1,8 @@
-import pytest
-
-import sympy as sp
import numpy as np
from pystencils import Assignment, Field, FieldType, AssignmentCollection
-from pystencils.nbackend.kernelcreation import create_kernel, CreateKernelConfig
+from pystencils.kernelcreation import create_kernel, CreateKernelConfig
+
def test_indexed_kernel():
arr = np.zeros((3, 4))
diff --git a/tests/nbackend/test_basic_printing.py b/tests/nbackend/test_basic_printing.py
index 7d1966882db3090de68b7766412ead3557537ef2..36a51b44b5ba4d259dcfbf9f755d68d400d744d2 100644
--- a/tests/nbackend/test_basic_printing.py
+++ b/tests/nbackend/test_basic_printing.py
@@ -1,12 +1,11 @@
-import pytest
-
from pystencils import Target
-from pystencils.nbackend.ast import *
-from pystencils.nbackend.typed_expressions import *
-from pystencils.nbackend.arrays import PsLinearizedArray, PsArrayBasePointer, PsArrayAccess
-from pystencils.nbackend.types.quick import *
-from pystencils.nbackend.emission import CAstPrinter
+from pystencils.backend.ast import *
+from pystencils.backend.typed_expressions import *
+from pystencils.backend.arrays import PsLinearizedArray, PsArrayBasePointer, PsArrayAccess
+from pystencils.backend.types.quick import *
+from pystencils.backend.emission import CAstPrinter
+
def test_basic_kernel():
diff --git a/tests/test_assignment_collection.py b/tests/test_assignment_collection.py
index f0c1f2a91e93739c1f40d8e2223641bd2f8b9bbb..69ca33e746d10e4e77c46292dc1611946f7bef21 100644
--- a/tests/test_assignment_collection.py
+++ b/tests/test_assignment_collection.py
@@ -3,7 +3,7 @@ import sympy as sp
import pystencils as ps
from pystencils import Assignment, AssignmentCollection
-from pystencils.astnodes import Conditional
+from pystencils.sympyextensions.astnodes import Conditional
from pystencils.simp.assignment_collection import SymbolGen
a, b, c = sp.symbols("a b c")
diff --git a/tests/test_assignment_from_stencil.py b/tests/test_assignment_from_stencil.py
index a4bb705d693403af3ea8cc3d44142f52d92b8668..6658628e1ac26ce6815cee05cc38aa9eadba2e97 100644
--- a/tests/test_assignment_from_stencil.py
+++ b/tests/test_assignment_from_stencil.py
@@ -13,9 +13,9 @@ def test_assignment_from_stencil():
x, y = pystencils.fields('x, y: [2D]')
- assignment = pystencils.assignment.assignment_from_stencil(stencil, x, y)
+ assignment = pystencils.sympyextensions.assignmentcollection.assignment.assignment_from_stencil(stencil, x, y)
assert isinstance(assignment, pystencils.Assignment)
assert assignment.rhs == x[0, 1] + 4 * x[-1, 1] + 2 * x[0, 0] + 3 * x[0, -1]
- assignment = pystencils.assignment.assignment_from_stencil(stencil, x, y, normalization_factor=1 / np.sum(stencil))
+ assignment = pystencils.sympyextensions.assignmentcollection.assignment.assignment_from_stencil(stencil, x, y, normalization_factor=1 / np.sum(stencil))
assert isinstance(assignment, pystencils.Assignment)
diff --git a/tests/test_astnodes.py b/tests/test_astnodes.py
index 91c11d8ecb1c76bfe19a0594663dfaa2e5b7ca4e..014e3574aedcfc5e49abc709d7b18f5d56aee140 100644
--- a/tests/test_astnodes.py
+++ b/tests/test_astnodes.py
@@ -1,12 +1,10 @@
-import pytest
import sys
-import pystencils.config
import sympy as sp
import pystencils as ps
from pystencils import Assignment
-from pystencils.astnodes import Block, LoopOverCoordinate, SkipIteration, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, LoopOverCoordinate, SkipIteration
dst = ps.fields('dst(8): double[2D]')
s = sp.symbols('s_:8')
diff --git a/tests/test_bit_masks.py b/tests/test_bit_masks.py
index 423fc13cc63569d3b6277983ca1e9210a3bbe9c9..3844a7a9f881401de7d787ae8e76bf070e92c358 100644
--- a/tests/test_bit_masks.py
+++ b/tests/test_bit_masks.py
@@ -3,7 +3,7 @@ import numpy as np
import pystencils as ps
from pystencils import Field, Assignment, create_kernel
-from pystencils.bit_masks import flag_cond
+from pystencils.sympyextensions.bit_masks import flag_cond
@pytest.mark.parametrize('mask_type', [np.uint8, np.uint16, np.uint32, np.uint64])
diff --git a/tests/test_buffer_gpu.py b/tests/test_buffer_gpu.py
index 944488d97f182dfd6b961e38e648b4d16b3d19ef..63017332eadb74702a18a61d097289bc402dd462 100644
--- a/tests/test_buffer_gpu.py
+++ b/tests/test_buffer_gpu.py
@@ -6,7 +6,7 @@ import pytest
import pystencils
from pystencils import Assignment, Field, FieldType, Target, CreateKernelConfig, create_kernel, fields
-from pystencils.bit_masks import flag_cond
+from pystencils.sympyextensions.bit_masks import flag_cond
from pystencils.field import create_numpy_array_with_layout, layout_string_to_tuple
from pystencils.slicing import (
add_ghost_layers, get_ghost_region_slice, get_slice_before_ghost_layer)
diff --git a/tests/test_conditional_field_access.py b/tests/test_conditional_field_access.py
index f8026c7dc22acf4e3664f637855aab3c029d0e26..e872164b4339d97307c0cefdff5fd0b41a30616e 100644
--- a/tests/test_conditional_field_access.py
+++ b/tests/test_conditional_field_access.py
@@ -15,7 +15,7 @@ import sympy as sp
import pystencils as ps
from pystencils import Field, x_vector
-from pystencils.astnodes import ConditionalFieldAccess
+from pystencils.sympyextensions.astnodes import ConditionalFieldAccess
from pystencils.simp import sympy_cse
diff --git a/tests/test_conditional_vec.py b/tests/test_conditional_vec.py
index ebb25d50d58b2a1ba0c8a5b7a5e089f1e6a7ea5d..0d1b4d6edc6482dc74376e2cbd899b111be43d1b 100644
--- a/tests/test_conditional_vec.py
+++ b/tests/test_conditional_vec.py
@@ -3,7 +3,7 @@ import sympy as sp
import pytest
import pystencils as ps
-from pystencils.astnodes import Block, Conditional, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, Conditional, SympyAssignment
from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets, get_vector_instruction_set
from pystencils.enums import Target
from pystencils.cpu.vectorization import vec_all, vec_any
diff --git a/tests/test_dot_printer.py b/tests/test_dot_printer.py
index a9d362c4fc2be1b46e969fe943c8179c532fdb36..3301589b5a7d293d6ff9185580d878a38860ae12 100644
--- a/tests/test_dot_printer.py
+++ b/tests/test_dot_printer.py
@@ -1,6 +1,6 @@
import pystencils as ps
-from pystencils.astnodes import Block, Conditional, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, Conditional, SympyAssignment
def test_dot_print():
diff --git a/tests/test_dtype_check.py b/tests/test_dtype_check.py
index 194a6338fb32bb95b1cd6d5df3cd12640abdb7d8..0d73224c60cc2406c630183c55407b468739c6df 100644
--- a/tests/test_dtype_check.py
+++ b/tests/test_dtype_check.py
@@ -11,7 +11,7 @@ def test_dtype_check_wrong_type():
stencil = [[1, 1, 1],
[1, 1, 1],
[1, 1, 1]]
- assignment = pystencils.assignment.assignment_from_stencil(stencil, x, y, normalization_factor=1 / np.sum(stencil))
+ assignment = pystencils.sympyextensions.assignmentcollection.assignment.assignment_from_stencil(stencil, x, y, normalization_factor=1 / np.sum(stencil))
kernel = pystencils.create_kernel([assignment]).compile()
with pytest.raises(ValueError) as e:
@@ -26,7 +26,7 @@ def test_dtype_check_correct_type():
stencil = [[1, 1, 1],
[1, 1, 1],
[1, 1, 1]]
- assignment = pystencils.assignment.assignment_from_stencil(stencil, x, y, normalization_factor=1 / np.sum(stencil))
+ assignment = pystencils.sympyextensions.assignmentcollection.assignment.assignment_from_stencil(stencil, x, y, normalization_factor=1 / np.sum(stencil))
kernel = pystencils.create_kernel([assignment]).compile()
kernel(x=array, y=output)
assert np.allclose(output[1:-1, 1:-1], np.ones_like(output[1:-1, 1:-1]))
diff --git a/tests/test_finite_differences.py b/tests/test_finite_differences.py
index 3ea65267fc99f1895d38d578cc06109ec1655c80..854c65ba17e6e49a491853517b921ecee6bce06a 100644
--- a/tests/test_finite_differences.py
+++ b/tests/test_finite_differences.py
@@ -2,7 +2,7 @@ import sympy as sp
import pytest
import pystencils as ps
-from pystencils.astnodes import LoopOverCoordinate
+from pystencils.sympyextensions.astnodes import LoopOverCoordinate
from pystencils.fd import diff, diffusion, Discretization2ndOrder
from pystencils.fd.spatial import discretize_spatial, fd_stencils_isotropic, fd_stencils_standard, \
fd_stencils_forth_order_isotropic
diff --git a/tests/test_fvm.py b/tests/test_fvm.py
index e4e3cacd53d34a6ea2079e1acbc82f93b8e23585..0b103e52554f767a0948509e8cae084c1af9e124 100644
--- a/tests/test_fvm.py
+++ b/tests/test_fvm.py
@@ -4,7 +4,7 @@ import numpy as np
import pytest
from itertools import product
from pystencils.rng import random_symbol
-from pystencils.astnodes import SympyAssignment
+from pystencils.sympyextensions.astnodes import SympyAssignment
from pystencils.node_collection import NodeCollection
@@ -182,8 +182,8 @@ def advection_diffusion_fluctuations(dim: int):
flux.main_assignments[i] = ps.Assignment(flux.main_assignments[i].lhs, flux.main_assignments[i].rhs + fluct)
# Add the folding to the flux, so that the random numbers persist through the ghostlayers.
- fold = {ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
- ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
+ fold = {pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
+ pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
flux.subs(fold)
flux_kernel = ps.create_staggered_kernel(flux).compile()
@@ -320,8 +320,8 @@ def diffusion_reaction(fluctuations: bool):
flux.main_assignments[i] = ps.Assignment(flux.main_assignments[i].lhs, flux.main_assignments[i].rhs + fluct)
# Add the folding to the flux, so that the random numbers persist through the ghostlayers.
- fold = {ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
- ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
+ fold = {pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
+ pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
flux.subs(fold)
r_flux = NodeCollection([SympyAssignment(j_fields[i].center, 0) for i in range(species)])
diff --git a/tests/test_global_definitions.py b/tests/test_global_definitions.py
index 8b6ee1b5bfb030cddfed2d7e0e70f91d8ccdfc04..365ce491cbe745e193475baee665956b78e2c83d 100644
--- a/tests/test_global_definitions.py
+++ b/tests/test_global_definitions.py
@@ -1,11 +1,9 @@
-import sympy
-
-import pystencils.astnodes
+import pystencils.sympyextensions.astnodes
from pystencils.backends.cbackend import CBackend
from pystencils.typing import TypedSymbol
-class BogusDeclaration(pystencils.astnodes.Node):
+class BogusDeclaration(pystencils.sympyextensions.astnodes.Node):
"""Base class for all AST nodes."""
def __init__(self, parent=None):
@@ -45,7 +43,7 @@ class BogusDeclaration(pystencils.astnodes.Node):
return result
-class BogusUsage(pystencils.astnodes.Node):
+class BogusUsage(pystencils.sympyextensions.astnodes.Node):
"""Base class for all AST nodes."""
def __init__(self, requires_global: bool, parent=None):
diff --git a/tests/test_helpful_errors.py b/tests/test_helpful_errors.py
index a13178afebd424c0f3e5e3355f8b0ecf97eab870..23215a52874abf24e6f50ba1323de1209870976b 100644
--- a/tests/test_helpful_errors.py
+++ b/tests/test_helpful_errors.py
@@ -4,7 +4,7 @@
import pytest
-from pystencils.astnodes import Block
+from pystencils.sympyextensions.astnodes import Block
from pystencils.backends.cbackend import CustomCodeNode, get_headers
diff --git a/tests/test_jacobi_cbackend.py b/tests/test_jacobi_cbackend.py
index 2ad6c007aa730a9c8eecd73b33f9d34ce9930395..a64919743202c1296d1db0a63b4dc37e5e3b9afe 100644
--- a/tests/test_jacobi_cbackend.py
+++ b/tests/test_jacobi_cbackend.py
@@ -1,7 +1,7 @@
import numpy as np
from pystencils import get_code_obj
-from pystencils.astnodes import Block, KernelFunction, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, KernelFunction, SympyAssignment
from pystencils.cpu import make_python_function
from pystencils.field import Field
from pystencils.enums import Target, Backend
diff --git a/tests/test_loop_cutting.py b/tests/test_loop_cutting.py
index 0372c57398f9dcf4493baf8153d5f3fd7faf27e2..3e44e59da3d7b7cd054a46058cc26d7fc07c3ab2 100644
--- a/tests/test_loop_cutting.py
+++ b/tests/test_loop_cutting.py
@@ -4,10 +4,10 @@ import sympy as sp
import pytest
import pystencils as ps
-import pystencils.astnodes as ast
+import pystencils.sympyextensions.astnodes as ast
from pystencils.field import Field, FieldType
-from pystencils.astnodes import Conditional, LoopOverCoordinate, SympyAssignment
-from pystencils.cpu import create_kernel, make_python_function
+from pystencils.sympyextensions.astnodes import Conditional, LoopOverCoordinate, SympyAssignment
+from pystencils.cpu import make_python_function
from pystencils.kernelcreation import create_staggered_kernel
from pystencils.transformations import (
cleanup_blocks, cut_loop, move_constants_before_loop, simplify_conditionals)
diff --git a/tests/test_modulo.py b/tests/test_modulo.py
index 66a3772c477808b08101cae930f4847fc55b4732..1a1d4c3548af0b396fd6fdf70ec80f75151ee2a9 100644
--- a/tests/test_modulo.py
+++ b/tests/test_modulo.py
@@ -3,7 +3,7 @@ import pytest
import numpy as np
import sympy as sp
import pystencils as ps
-from pystencils.astnodes import LoopOverCoordinate, Conditional, Block, SympyAssignment
+from pystencils.sympyextensions.astnodes import LoopOverCoordinate, Conditional, Block, SympyAssignment
SLICE_LIST = [False,
ps.make_slice[1:-1:2, 1:-1:2],
diff --git a/tests/test_move_constant_before_loop.py b/tests/test_move_constant_before_loop.py
index 0f5300f1b8c14747456efb6f1834a3fc3f3bc46f..e5aebc298399d2927a7bc529f9ea2ac6a855d4c0 100644
--- a/tests/test_move_constant_before_loop.py
+++ b/tests/test_move_constant_before_loop.py
@@ -1,7 +1,7 @@
import numpy as np
import pystencils as ps
-from pystencils.astnodes import Block, LoopOverCoordinate, SympyAssignment, TypedSymbol
+from pystencils.sympyextensions.astnodes import Block, LoopOverCoordinate, SympyAssignment, TypedSymbol
from pystencils.transformations import move_constants_before_loop
@@ -47,7 +47,7 @@ def test_keep_order_of_accesses():
new_loops = ps.transformations.cut_loop(loop, [n - 1])
ps.transformations.move_constants_before_loop(new_loops.args[1])
- kernel_func = ps.astnodes.KernelFunction(
+ kernel_func = pystencils.sympyextensions.astnodes.KernelFunction(
block, ps.Target.CPU, ps.Backend.C, ps.cpu.cpujit.make_python_function, None
)
kernel = kernel_func.compile()
diff --git a/tests/test_nodecollection.py b/tests/test_nodecollection.py
index ab24e58e7dc4a70985e6a9c631b9085b39a3f00f..beaf592b02e7ecad779b22c0663babf7402813ad 100644
--- a/tests/test_nodecollection.py
+++ b/tests/test_nodecollection.py
@@ -2,7 +2,7 @@ import sympy as sp
from pystencils import AssignmentCollection, Assignment
from pystencils.node_collection import NodeCollection
-from pystencils.astnodes import SympyAssignment
+from pystencils.sympyextensions.astnodes import SympyAssignment
def test_node_collection_from_assignment_collection():
diff --git a/tests/test_random.py b/tests/test_random.py
index e82bff309e636c3c07f4b1f82b552d67e6ccb3d7..63cd61494b675acbc1675277ff9eb9b1f57ddb3e 100644
--- a/tests/test_random.py
+++ b/tests/test_random.py
@@ -3,7 +3,7 @@ import numpy as np
import pytest
import pystencils as ps
-from pystencils.astnodes import SympyAssignment
+from pystencils.sympyextensions.astnodes import SympyAssignment
from pystencils.node_collection import NodeCollection
from pystencils.rng import PhiloxFourFloats, PhiloxTwoDoubles, AESNIFourFloats, AESNITwoDoubles, random_symbol
from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets
diff --git a/tests/test_source_code_comment.py b/tests/test_source_code_comment.py
index b1006a941f5c5922b1168944f17d2c9aceba66ba..90323754b7ee619337c38e2181242ba5e24c2619 100644
--- a/tests/test_source_code_comment.py
+++ b/tests/test_source_code_comment.py
@@ -8,7 +8,7 @@
"""
import pystencils
-import pystencils.astnodes
+import pystencils.sympyextensions.astnodes
import pystencils.config
@@ -23,9 +23,9 @@ def test_source_code_comment():
config = pystencils.config.CreateKernelConfig(target=pystencils.Target.CPU)
ast = pystencils.create_kernel(assignments, config=config)
- ast.body.append(pystencils.astnodes.SourceCodeComment("Hallo"))
- ast.body.append(pystencils.astnodes.EmptyLine())
- ast.body.append(pystencils.astnodes.SourceCodeComment("World!"))
+ ast.body.append(pystencils.sympyextensions.astnodes.SourceCodeComment("Hallo"))
+ ast.body.append(pystencils.sympyextensions.astnodes.EmptyLine())
+ ast.body.append(pystencils.sympyextensions.astnodes.SourceCodeComment("World!"))
print(ast)
compiled = ast.compile()
assert compiled is not None
diff --git a/tests/test_transformations.py b/tests/test_transformations.py
index ba660a115d6438da31b7bf653730c751da9920a7..3585fe42a61012cad3595d268c41ed38ab49c4ca 100644
--- a/tests/test_transformations.py
+++ b/tests/test_transformations.py
@@ -3,11 +3,11 @@ import numpy as np
import pystencils as ps
from pystencils import fields, TypedSymbol
-from pystencils.astnodes import LoopOverCoordinate, SympyAssignment
+from pystencils.sympyextensions.astnodes import LoopOverCoordinate, SympyAssignment
from pystencils.typing import create_type
from pystencils.transformations import (
filtered_tree_iteration, get_loop_hierarchy, get_loop_counter_symbol_hierarchy,
- iterate_loops_by_depth, split_inner_loop, loop_blocking
+ iterate_loops_by_depth, split_inner_loop
)
from pystencils.cpu import add_pragmas
diff --git a/tests/test_vectorization.py b/tests/test_vectorization.py
index 302d8cbd9de5c53e60bb0b43d24e2ab1d8155228..fd416ab4cd05c7b8891aae4da91cc9aeae425698 100644
--- a/tests/test_vectorization.py
+++ b/tests/test_vectorization.py
@@ -6,7 +6,7 @@ import pystencils.config
import sympy as sp
import pystencils as ps
-import pystencils.astnodes as ast
+import pystencils.sympyextensions.astnodes as ast
from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets, get_vector_instruction_set
from pystencils.cpu.vectorization import vectorize
from pystencils.enums import Target