diff --git a/src/pystencils/nbackend/ast/__init__.py b/src/pystencils/nbackend/ast/__init__.py
index c7731f37a77927acca8b7ecd1b1d7ff153d6d855..80e710c482370138840e224ab00b7bdfb8774539 100644
--- a/src/pystencils/nbackend/ast/__init__.py
+++ b/src/pystencils/nbackend/ast/__init__.py
@@ -8,11 +8,12 @@ from .nodes import (
PsDeclaration,
PsLoop,
PsConditional,
+ PsComment,
)
from .kernelfunction import PsKernelFunction
+from .tree_iteration import dfs_preorder, dfs_postorder
from .dispatcher import ast_visitor
-from .transformations import ast_subs
__all__ = [
"ast_visitor",
@@ -26,5 +27,7 @@ __all__ = [
"PsDeclaration",
"PsLoop",
"PsConditional",
- "ast_subs"
+ "PsComment",
+ "dfs_preorder",
+ "dfs_postorder",
]
diff --git a/src/pystencils/nbackend/ast/nodes.py b/src/pystencils/nbackend/ast/nodes.py
index 865fa6dec049bd413cc0fd268a9b5939d0765b84..d11be448dfc8f5f5f208062494f123ee17502f78 100644
--- a/src/pystencils/nbackend/ast/nodes.py
+++ b/src/pystencils/nbackend/ast/nodes.py
@@ -97,7 +97,7 @@ class PsExpression(PsLeafNode):
self._expr = expr
def __repr__(self) -> str:
- return repr(self._expr)
+ return f"Expr({repr(self._expr)})"
def __eq__(self, other: object) -> bool:
if not isinstance(other, PsExpression):
@@ -351,3 +351,17 @@ class PsConditional(PsAstNode):
self._branch_false = failing_cast((PsBlock, NoneType), c)
case _:
assert False, "unreachable code"
+
+
+class PsComment(PsLeafNode):
+ def __init__(self, text: str) -> None:
+ self._text = text
+ self._lines = tuple(text.splitlines())
+
+ @property
+ def text(self) -> str:
+ return self._text
+
+ @property
+ def lines(self) -> tuple[str, ...]:
+ return self._lines
diff --git a/src/pystencils/nbackend/ast/tree_iteration.py b/src/pystencils/nbackend/ast/tree_iteration.py
new file mode 100644
index 0000000000000000000000000000000000000000..2019e7d02b59950f24f272d33c154f8174c68b5f
--- /dev/null
+++ b/src/pystencils/nbackend/ast/tree_iteration.py
@@ -0,0 +1,37 @@
+from typing import Callable, Generator
+
+from .nodes import PsAstNode
+
+
+def dfs_preorder(
+ node: PsAstNode,
+ yield_pred: Callable[[PsAstNode], bool] = lambda _: True
+) -> Generator[PsAstNode, None, None]:
+ """Pre-Order depth-first traversal of an abstract syntax tree.
+
+ Args:
+ node: The tree's root node
+ yield_pred: Filter predicate; a node is only yielded to the caller if `yield_pred(node)` returns True
+ """
+ if yield_pred(node):
+ yield node
+
+ for c in node.children:
+ yield from dfs_preorder(c, yield_pred)
+
+
+def dfs_postorder(
+ node: PsAstNode,
+ yield_pred: Callable[[PsAstNode], bool] = lambda _: True
+) -> Generator[PsAstNode, None, None]:
+ """Post-Order depth-first traversal of an abstract syntax tree.
+
+ Args:
+ node: The tree's root node
+ yield_pred: Filter predicate; a node is only yielded to the caller if `yield_pred(node)` returns True
+ """
+ for c in node.children:
+ yield from dfs_postorder(c, yield_pred)
+
+ if yield_pred(node):
+ yield node
diff --git a/src/pystencils/nbackend/emission.py b/src/pystencils/nbackend/emission.py
index b89b58224594da80176a3d8b8cfa8ed41c73dd56..3fe16e9c1f82ee15db57bc126332fef5c8155d3c 100644
--- a/src/pystencils/nbackend/emission.py
+++ b/src/pystencils/nbackend/emission.py
@@ -11,6 +11,7 @@ from .ast import (
PsAssignment,
PsLoop,
PsConditional,
+ PsComment
)
from .ast.kernelfunction import PsKernelFunction
from .typed_expressions import PsTypedVariable
@@ -122,3 +123,12 @@ class CAstPrinter:
code += f"\nelse\n{else_code}"
return self.indent(code)
+
+ @visit.case(PsComment)
+ def comment(self, node: PsComment):
+ lines = list(node.lines)
+ lines[0] = "/* " + lines[0]
+ for i in range(1, len(lines)):
+ lines[i] = " " + lines[i]
+ lines[-1] = lines[-1] + " */"
+ return self.indent("\n".join(lines))
diff --git a/src/pystencils/nbackend/kernelcreation/__init__.py b/src/pystencils/nbackend/kernelcreation/__init__.py
index 110acb8187e55f4d5c151ae4a06cf16f9fd051e6..18b1acc48f6b455d648a9e9f9c22630433cf200f 100644
--- a/src/pystencils/nbackend/kernelcreation/__init__.py
+++ b/src/pystencils/nbackend/kernelcreation/__init__.py
@@ -6,7 +6,12 @@ from .analysis import KernelAnalysis
from .freeze import FreezeExpressions
from .typification import Typifier
-from .iteration_space import FullIterationSpace, SparseIterationSpace
+from .iteration_space import (
+ FullIterationSpace,
+ SparseIterationSpace,
+ create_full_iteration_space,
+ create_sparse_iteration_space,
+)
__all__ = [
"KernelCreationOptions",
@@ -17,4 +22,6 @@ __all__ = [
"Typifier",
"FullIterationSpace",
"SparseIterationSpace",
+ "create_full_iteration_space",
+ "create_sparse_iteration_space",
]
diff --git a/src/pystencils/nbackend/kernelcreation/iteration_space.py b/src/pystencils/nbackend/kernelcreation/iteration_space.py
index 22375f7cb82cd5c4eeb8f2bd85ed0400702e4b6e..43ac685e5c59acc805e54fa22612d056bd804101 100644
--- a/src/pystencils/nbackend/kernelcreation/iteration_space.py
+++ b/src/pystencils/nbackend/kernelcreation/iteration_space.py
@@ -75,10 +75,11 @@ class FullIterationSpace(IterationSpace):
archetype_field: Field,
ghost_layers: int | Sequence[int | tuple[int, int]],
) -> FullIterationSpace:
- """Create an iteration space for a collection of fields with ghost layers."""
+ """Create an iteration space over an archetype field with ghost layers."""
archetype_array = ctx.get_array(archetype_field)
dim = archetype_field.spatial_dimensions
+
counters = [
PsTypedVariable(name, ctx.index_dtype)
for name in Defaults.spatial_counter_names[:dim]
@@ -112,7 +113,9 @@ class FullIterationSpace(IterationSpace):
)
]
- # TODO: Reorder dimensions according to optimal loop layout (?)
+ # Determine loop order by permuting dimensions
+ loop_order = archetype_field.layout
+ dimensions = [dimensions[coordinate] for coordinate in loop_order]
return FullIterationSpace(ctx, dimensions)
diff --git a/tests/nbackend/kernelcreation/platform/__init__.py b/tests/nbackend/kernelcreation/platform/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tests/nbackend/kernelcreation/platform/test_basic_cpu.py b/tests/nbackend/kernelcreation/platform/test_basic_cpu.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0fb706255f59f21cca25c411af6a27b9fb2a23b
--- /dev/null
+++ b/tests/nbackend/kernelcreation/platform/test_basic_cpu.py
@@ -0,0 +1,34 @@
+import pytest
+
+from pystencils.field import Field
+
+from pystencils.nbackend.kernelcreation import (
+ KernelCreationContext,
+ KernelCreationOptions,
+ FullIterationSpace
+)
+
+from pystencils.nbackend.ast import PsBlock, PsLoop, PsComment, dfs_preorder
+
+from pystencils.nbackend.kernelcreation.platform import BasicCpu
+
+@pytest.mark.parametrize("layout", ["fzyx", "zyxf", "c", "f"])
+def test_loop_nest(layout):
+ ctx = KernelCreationContext(KernelCreationOptions())
+
+ body = PsBlock([PsComment("Loop body goes here")])
+ platform = BasicCpu(ctx)
+
+ # FZYX Order
+ archetype_field = Field.create_generic("fzyx_field", spatial_dimensions=3, layout=layout)
+ ispace = FullIterationSpace.create_with_ghost_layers(ctx, archetype_field, 0)
+
+ loop_nest = platform.materialize_iteration_space(body, ispace)
+
+ loops = dfs_preorder(loop_nest, lambda n: isinstance(n, PsLoop))
+ for loop, dim in zip(loops, ispace.dimensions, strict=True):
+ assert isinstance(loop, PsLoop)
+ assert loop.start.expression == dim.start
+ assert loop.stop.expression == dim.stop
+ assert loop.step.expression == dim.step
+ assert loop.counter.expression == dim.counter
diff --git a/tests/nbackend/kernelcreation/test_domain_kernels.py b/tests/nbackend/kernelcreation/test_domain_kernels.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f75df8568d2260bbd37f056c9211a2480dfb163
--- /dev/null
+++ b/tests/nbackend/kernelcreation/test_domain_kernels.py
@@ -0,0 +1,64 @@
+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.nbackend.kernelcreation import create_kernel
+from pystencils.cpu.cpujit import compile_and_load
+
+def test_filter_kernel():
+ weight = sp.Symbol("weight")
+ stencil = [
+ [1, 1, 1],
+ [1, 1, 1],
+ [1, 1, 1]
+ ]
+
+ src, dst = fields("src, dst: [2D]")
+ asm = assignment_from_stencil(stencil, src, dst, normalization_factor=weight)
+ asms = AssignmentCollection([asm])
+
+ ast = create_kernel(asms)
+
+ kernel = compile_and_load(ast)
+
+ src_arr = np.ones((42, 42))
+ dst_arr = np.zeros_like(src_arr)
+
+ kernel(src=src_arr, dst=dst_arr, weight=2.0)
+
+ expected = np.zeros_like(src_arr)
+ expected[1:-1, 1:-1].fill(18.0)
+
+ np.testing.assert_allclose(dst_arr, expected)
+
+
+def test_filter_kernel_fixedsize():
+ weight = sp.Symbol("weight")
+ stencil = [
+ [1, 1, 1],
+ [1, 1, 1],
+ [1, 1, 1]
+ ]
+
+ src_arr = np.ones((42, 42))
+ dst_arr = np.zeros_like(src_arr)
+
+ src = Field.create_from_numpy_array("src", src_arr)
+ dst = Field.create_from_numpy_array("dst", dst_arr)
+
+ asm = assignment_from_stencil(stencil, src, dst, normalization_factor=weight)
+ asms = AssignmentCollection([asm])
+
+ ast = create_kernel(asms)
+ kernel = compile_and_load(ast)
+
+ kernel(src=src_arr, dst=dst_arr, weight=2.0)
+
+ expected = np.zeros_like(src_arr)
+ expected[1:-1, 1:-1].fill(18.0)
+
+ np.testing.assert_allclose(dst_arr, expected)
diff --git a/tests/nbackend/kernelcreation/test_iteration_space.py b/tests/nbackend/kernelcreation/test_iteration_space.py
new file mode 100644
index 0000000000000000000000000000000000000000..935bf4e703ce2dd0e41256ae7d08048042a0e6ea
--- /dev/null
+++ b/tests/nbackend/kernelcreation/test_iteration_space.py
@@ -0,0 +1,49 @@
+from pystencils.field import Field
+
+from pystencils.nbackend.kernelcreation import (
+ KernelCreationContext,
+ KernelCreationOptions,
+ FullIterationSpace
+)
+
+from pystencils.nbackend.kernelcreation.defaults import Pymbolic as PbDefaults
+
+
+def test_loop_order():
+ ctx = KernelCreationContext(KernelCreationOptions())
+ ctr_symbols = PbDefaults.spatial_counters
+
+ # FZYX Order
+ archetype_field = Field.create_generic("fzyx_field", spatial_dimensions=3, layout='fzyx')
+ ispace = FullIterationSpace.create_with_ghost_layers(ctx, archetype_field, 0)
+
+ for dim, ctr in zip(ispace.dimensions, ctr_symbols[::-1]):
+ assert dim.counter == ctr
+
+ # ZYXF Order
+ archetype_field = Field.create_generic("zyxf_field", spatial_dimensions=3, layout='zyxf')
+ ispace = FullIterationSpace.create_with_ghost_layers(ctx, archetype_field, 0)
+
+ for dim, ctr in zip(ispace.dimensions, ctr_symbols[::-1]):
+ assert dim.counter == ctr
+
+ # C Order
+ archetype_field = Field.create_generic("c_field", spatial_dimensions=3, layout='c')
+ ispace = FullIterationSpace.create_with_ghost_layers(ctx, archetype_field, 0)
+
+ for dim, ctr in zip(ispace.dimensions, ctr_symbols):
+ assert dim.counter == ctr
+
+ # Fortran Order
+ archetype_field = Field.create_generic("fortran_field", spatial_dimensions=3, layout='f')
+ ispace = FullIterationSpace.create_with_ghost_layers(ctx, archetype_field, 0)
+
+ for dim, ctr in zip(ispace.dimensions, ctr_symbols[::-1]):
+ assert dim.counter == ctr
+
+ # Scrambled Layout
+ archetype_field = Field.create_generic("scrambled_field", spatial_dimensions=3, layout=(2, 0, 1))
+ ispace = FullIterationSpace.create_with_ghost_layers(ctx, archetype_field, 0)
+
+ for dim, ctr in zip(ispace.dimensions, [ctr_symbols[2], ctr_symbols[0], ctr_symbols[1]]):
+ assert dim.counter == ctr