test cases for iteration spaces and domain kernels

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",
 ]

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

src/pystencils/nbackend/ast/tree_iteration.py

+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

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))

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",
 ]

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)
 

tests/nbackend/kernelcreation/platform/test_basic_cpu.py

+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

tests/nbackend/kernelcreation/test_domain_kernels.py

+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)

tests/nbackend/kernelcreation/test_iteration_space.py

+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