diff --git a/src/pystencilssfg/composer/basic_composer.py b/src/pystencilssfg/composer/basic_composer.py
index 97966247d4be769dfcb800e0b25bcd631bdda69e..9f0673de3c109bbdbbbfa1eaf1a3ebec51af87cf 100644
--- a/src/pystencilssfg/composer/basic_composer.py
+++ b/src/pystencilssfg/composer/basic_composer.py
@@ -1,13 +1,19 @@
from __future__ import annotations
-from typing import Sequence
+from typing import Sequence, TypeAlias
from abc import ABC, abstractmethod
import numpy as np
import sympy as sp
from functools import reduce
-from pystencils import Field
+from pystencils import Field, TypedSymbol
from pystencils.backend import KernelParameter, KernelFunction
-from pystencils.types import create_type, UserTypeSpec, PsCustomType, PsPointerType
+from pystencils.types import (
+ create_type,
+ UserTypeSpec,
+ PsCustomType,
+ PsPointerType,
+ PsType,
+)
from ..context import SfgContext
from .custom import CustomGenerator
@@ -58,8 +64,23 @@ class SfgNodeBuilder(ABC):
pass
-ExprLike = str | SfgVar | AugExpr
-SequencerArg = tuple | str | AugExpr | SfgCallTreeNode | SfgNodeBuilder
+_ExprLike = (str, AugExpr, TypedSymbol)
+ExprLike: TypeAlias = str | AugExpr | TypedSymbol
+"""Things that may act as a C++ expression.
+
+Expressions need not necesserily have a known data type.
+"""
+
+_VarLike = (TypedSymbol, AugExpr)
+VarLike: TypeAlias = TypedSymbol | AugExpr
+"""Things that may act as a variable.
+
+Variables must always define their name *and* data type.
+"""
+
+_SequencerArg = (tuple, ExprLike, SfgCallTreeNode, SfgNodeBuilder)
+SequencerArg: TypeAlias = tuple | ExprLike | SfgCallTreeNode | SfgNodeBuilder
+"""Valid arguments to `make_sequence` and any sequencer that uses it."""
class SfgBasicComposer(SfgIComposer):
@@ -208,7 +229,11 @@ class SfgBasicComposer(SfgIComposer):
num_blocks_str = str(num_blocks)
tpb_str = str(threads_per_block)
stream_str = str(stream) if stream is not None else None
- depends = _depends(num_blocks) | _depends(threads_per_block) | _depends(stream)
+
+ depends = _depends(num_blocks) | _depends(threads_per_block)
+ if stream is not None:
+ depends |= _depends(stream)
+
return SfgCudaKernelInvocation(
kernel_handle, num_blocks_str, tpb_str, stream_str, depends
)
@@ -217,9 +242,9 @@ class SfgBasicComposer(SfgIComposer):
"""Syntax sequencing. For details, see `make_sequence`"""
return make_sequence(*args)
- def params(self, *args: SfgVar) -> SfgFunctionParams:
+ def params(self, *args: AugExpr) -> SfgFunctionParams:
"""Use inside a function body to add parameters to the function."""
- return SfgFunctionParams(args)
+ return SfgFunctionParams([x.as_variable() for x in args])
def require(self, *includes: str | SfgHeaderInclude) -> SfgRequireIncludes:
return SfgRequireIncludes(
@@ -232,7 +257,7 @@ class SfgBasicComposer(SfgIComposer):
ptr: bool = False,
ref: bool = False,
const: bool = False,
- ):
+ ) -> PsType:
if ptr and ref:
raise SfgException("Create either a pointer, or a ref type, not both!")
@@ -250,11 +275,11 @@ class SfgBasicComposer(SfgIComposer):
else:
return base_type
- def var(self, name: str, dtype: UserTypeSpec) -> SfgVar:
+ def var(self, name: str, dtype: UserTypeSpec) -> AugExpr:
"""Create a variable with given name and data type."""
- return SfgVar(name, create_type(dtype))
+ return AugExpr(create_type(dtype)).var(name)
- def init(self, lhs: SfgVar) -> SfgInplaceInitBuilder:
+ def init(self, lhs: VarLike) -> SfgInplaceInitBuilder:
"""Create a C++ in-place initialization.
Usage:
@@ -270,7 +295,7 @@ class SfgBasicComposer(SfgIComposer):
SomeClass obj { arg1, arg2, arg3 };
"""
- return SfgInplaceInitBuilder(lhs)
+ return SfgInplaceInitBuilder(_asvar(lhs))
def expr(self, fmt: str, *deps, **kwdeps):
return AugExpr.format(fmt, *deps, **kwdeps)
@@ -329,15 +354,8 @@ class SfgBasicComposer(SfgIComposer):
def make_statements(arg: ExprLike) -> SfgStatements:
- match arg:
- case str():
- return SfgStatements(arg, (), ())
- case SfgVar(name, _):
- return SfgStatements(name, (), (arg,))
- case AugExpr():
- return SfgStatements(str(arg), (), arg.depends)
- case _:
- assert False
+ depends = _depends(arg)
+ return SfgStatements(str(arg), (), depends)
def make_sequence(*args: SequencerArg) -> SfgSequence:
@@ -392,10 +410,8 @@ def make_sequence(*args: SequencerArg) -> SfgSequence:
children.append(arg.resolve())
elif isinstance(arg, SfgCallTreeNode):
children.append(arg)
- elif isinstance(arg, AugExpr):
- children.append(SfgStatements(str(arg), (), arg.depends))
- elif isinstance(arg, str):
- children.append(SfgStatements(arg, (), ()))
+ elif isinstance(arg, _ExprLike):
+ children.append(make_statements(arg))
elif isinstance(arg, tuple):
# Tuples are treated as blocks
subseq = make_sequence(*arg)
@@ -407,25 +423,20 @@ def make_sequence(*args: SequencerArg) -> SfgSequence:
class SfgInplaceInitBuilder(SfgNodeBuilder):
- def __init__(self, lhs: SfgVar | AugExpr) -> None:
- if isinstance(lhs, AugExpr):
- lhs = lhs.as_variable()
-
+ def __init__(self, lhs: SfgVar) -> None:
self._lhs: SfgVar = lhs
self._depends: set[SfgVar] = set()
self._rhs: str | None = None
def __call__(
self,
- *rhs: str | AugExpr,
+ *rhs: ExprLike,
) -> SfgInplaceInitBuilder:
if self._rhs is not None:
raise SfgException("Assignment builder used multiple times.")
self._rhs = ", ".join(str(expr) for expr in rhs)
- self._depends = reduce(
- set.union, (obj.depends for obj in rhs if isinstance(obj, AugExpr)), set()
- )
+ self._depends = reduce(set.union, (_depends(obj) for obj in rhs), set())
return self
def resolve(self) -> SfgCallTreeNode:
@@ -538,12 +549,30 @@ def struct_from_numpy_dtype(
return cls
-def _depends(expr: ExprLike | Sequence[ExprLike] | None) -> set[SfgVar]:
+def _asvar(var: VarLike) -> SfgVar:
+ match var:
+ case AugExpr():
+ return var.as_variable()
+ case TypedSymbol():
+ from pystencils import DynamicType
+
+ if isinstance(var.dtype, DynamicType):
+ raise SfgException(
+ f"Unable to cast dynamically typed symbol {var} to a variable.\n"
+ f"{var} has dynamic type {var.dtype}, which cannot be resolved to a type outside of a kernel."
+ )
+
+ return SfgVar(var.name, var.dtype)
+ case _:
+ raise ValueError(f"Invalid variable: {var}")
+
+
+def _depends(expr: ExprLike) -> set[SfgVar]:
match expr:
case None | str():
return set()
- case SfgVar():
- return {expr}
+ case TypedSymbol():
+ return {_asvar(expr)}
case AugExpr():
return expr.depends
case _:
diff --git a/src/pystencilssfg/composer/class_composer.py b/src/pystencilssfg/composer/class_composer.py
index 63588b7829b2a94122e5c9d7d38770cabce9d6f5..9b8fcfb13e5cfa84349275735cab8ee25f53771f 100644
--- a/src/pystencilssfg/composer/class_composer.py
+++ b/src/pystencilssfg/composer/class_composer.py
@@ -1,8 +1,10 @@
from __future__ import annotations
from typing import Sequence
+from pystencils import TypedSymbol
from pystencils.types import PsCustomType, UserTypeSpec
+from ..lang import AugExpr
from ..ir import SfgCallTreeNode
from ..ir.source_components import (
SfgClass,
@@ -14,12 +16,18 @@ from ..ir.source_components import (
SfgClassKeyword,
SfgVisibility,
SfgVisibilityBlock,
- SfgVar,
)
from ..exceptions import SfgException
from .mixin import SfgComposerMixIn
-from .basic_composer import SfgNodeBuilder, make_sequence
+from .basic_composer import (
+ SfgNodeBuilder,
+ make_sequence,
+ _VarLike,
+ VarLike,
+ ExprLike,
+ _asvar,
+)
class SfgClassComposer(SfgComposerMixIn):
@@ -46,7 +54,7 @@ class SfgClassComposer(SfgComposerMixIn):
def __call__(
self,
*args: (
- SfgClassMember | SfgClassComposer.ConstructorBuilder | SfgVar | str
+ SfgClassMember | SfgClassComposer.ConstructorBuilder | VarLike | str
),
):
for arg in args:
@@ -63,15 +71,21 @@ class SfgClassComposer(SfgComposerMixIn):
Returned by `constructor`.
"""
- def __init__(self, *params: SfgVar):
- self._params = params
+ def __init__(self, *params: VarLike):
+ self._params = tuple(_asvar(p) for p in params)
self._initializers: list[str] = []
self._body: str | None = None
- def init(self, initializer: str) -> SfgClassComposer.ConstructorBuilder:
+ def init(self, var: VarLike):
"""Add an initialization expression to the constructor's initializer list."""
- self._initializers.append(initializer)
- return self
+
+ def init_sequencer(expr: ExprLike):
+ expr = str(expr)
+ initializer = f"{_asvar(var)}{{ {expr} }}"
+ self._initializers.append(initializer)
+ return self
+
+ return init_sequencer
def body(self, body: str):
"""Define the constructor body"""
@@ -120,7 +134,7 @@ class SfgClassComposer(SfgComposerMixIn):
"""Create a `private` visibility block in a class or struct body"""
return SfgClassComposer.VisibilityContext(SfgVisibility.PRIVATE)
- def constructor(self, *params: SfgVar):
+ def constructor(self, *params: VarLike):
"""In a class or struct body or visibility block, add a constructor.
Args:
@@ -171,7 +185,7 @@ class SfgClassComposer(SfgComposerMixIn):
SfgClassComposer.VisibilityContext
| SfgClassMember
| SfgClassComposer.ConstructorBuilder
- | SfgVar
+ | VarLike
| str
),
):
@@ -186,9 +200,9 @@ class SfgClassComposer(SfgComposerMixIn):
(
SfgClassMember,
SfgClassComposer.ConstructorBuilder,
- SfgVar,
str,
- ),
+ )
+ + _VarLike,
):
if default_ended:
raise SfgException(
@@ -204,13 +218,17 @@ class SfgClassComposer(SfgComposerMixIn):
@staticmethod
def _resolve_member(
- arg: SfgClassMember | SfgClassComposer.ConstructorBuilder | SfgVar | str,
- ):
- if isinstance(arg, SfgVar):
- return SfgMemberVariable(arg.name, arg.dtype)
- elif isinstance(arg, str):
- return SfgInClassDefinition(arg)
- elif isinstance(arg, SfgClassComposer.ConstructorBuilder):
- return arg.resolve()
- else:
- return arg
+ arg: SfgClassMember | SfgClassComposer.ConstructorBuilder | VarLike | str,
+ ) -> SfgClassMember:
+ match arg:
+ case AugExpr() | TypedSymbol():
+ var = _asvar(arg)
+ return SfgMemberVariable(var.name, var.dtype)
+ case str():
+ return SfgInClassDefinition(arg)
+ case SfgClassComposer.ConstructorBuilder():
+ return arg.resolve()
+ case SfgClassMember():
+ return arg
+ case _:
+ raise ValueError(f"Invalid class member: {arg}")
diff --git a/src/pystencilssfg/lang/expressions.py b/src/pystencilssfg/lang/expressions.py
index 579646f3ee5b96dd3a57c1ece45c6b6407c48102..ca5a78c4d0794d33a2b22827b7a5e9d5a09e47d4 100644
--- a/src/pystencilssfg/lang/expressions.py
+++ b/src/pystencilssfg/lang/expressions.py
@@ -48,7 +48,7 @@ class DependentExpression:
def __add__(self, other: DependentExpression):
return DependentExpression(self.expr + other.expr, self.depends | other.depends)
-
+
class VarExpr(DependentExpression):
def __init__(self, var: SfgVar):
@@ -61,6 +61,8 @@ class VarExpr(DependentExpression):
class AugExpr:
+ __match_args__ = ("expr", "dtype")
+
def __init__(self, dtype: PsType | None = None):
self._dtype = dtype
self._bound: DependentExpression | None = None
@@ -77,6 +79,7 @@ class AugExpr:
@staticmethod
def format(fmt: str, *deps, **kwdeps) -> AugExpr:
+ """Create a new `AugExpr` by combining existing expressions."""
return AugExpr().bind(fmt, *deps, **kwdeps)
def bind(self, fmt: str, *deps, **kwdeps):
@@ -109,11 +112,11 @@ class AugExpr:
raise SfgException("This AugExpr has no known data type.")
return self._dtype
-
+
@property
def is_variable(self) -> bool:
return isinstance(self._bound, VarExpr)
-
+
def as_variable(self) -> SfgVar:
if not isinstance(self._bound, VarExpr):
raise SfgException("This expression is not a variable")
diff --git a/tests/integration/expected/Variables.h b/tests/integration/expected/Variables.h
new file mode 100644
index 0000000000000000000000000000000000000000..96c16d7b306fe7594db0caa02e9168b2f5c86fc6
--- /dev/null
+++ b/tests/integration/expected/Variables.h
@@ -0,0 +1,13 @@
+#pragma once
+
+#include <cstdint>
+
+#define RESTRICT __restrict__
+
+class Scale {
+private:
+ float alpha;
+public:
+ Scale(float alpha) : alpha{ alpha } {}
+ void operator() (float *const _data_f, float *const _data_g);
+};
diff --git a/tests/integration/scripts/SimpleJacobi.py b/tests/integration/scripts/SimpleJacobi.py
index 199419c541fb4aef68d44a1baf7cfc2f28d67e86..e84c872b05c354f4e1473b2eab17781f0880f035 100644
--- a/tests/integration/scripts/SimpleJacobi.py
+++ b/tests/integration/scripts/SimpleJacobi.py
@@ -20,4 +20,4 @@ with SourceFileGenerator() as sfg:
sfg.map_field(u_dst, mdspan_ref(u_dst)),
sfg.map_field(f, mdspan_ref(f)),
sfg.call(poisson_kernel)
- )
\ No newline at end of file
+ )
diff --git a/tests/integration/scripts/Variables.py b/tests/integration/scripts/Variables.py
new file mode 100644
index 0000000000000000000000000000000000000000..9fd4e0027104451738abea72e15084047cf4465e
--- /dev/null
+++ b/tests/integration/scripts/Variables.py
@@ -0,0 +1,22 @@
+import sympy as sp
+from pystencils import TypedSymbol, fields, kernel
+
+from pystencilssfg import SourceFileGenerator, SfgConfiguration
+
+with SourceFileGenerator() as sfg:
+ α = TypedSymbol("alpha", "float32")
+ f, g = fields("f, g: float32[10]")
+
+ @kernel
+ def scale():
+ f[0] @= α * g.center()
+
+ khandle = sfg.kernels.create(scale)
+
+ sfg.klass("Scale")(
+ sfg.private(α),
+ sfg.public(
+ sfg.constructor(α).init(α)(α.name),
+ sfg.method("operator()")(sfg.init(α)(f"this->{α}"), sfg.call(khandle)),
+ ),
+ )
diff --git a/tests/integration/test_generator_scripts.py b/tests/integration/test_generator_scripts.py
index 33f747820b966bd7a2630e9497cbb00a7937340c..466dbf34ca6f622442090d7e7001e437a8260156 100644
--- a/tests/integration/test_generator_scripts.py
+++ b/tests/integration/test_generator_scripts.py
@@ -15,20 +15,45 @@ EXPECTED_DIR = path.join(THIS_DIR, "expected")
@dataclass
class ScriptInfo:
script_name: str
+ """Name of the generator script, without .py-extension.
+
+ Generator scripts must be located in the ``scripts`` folder.
+ """
+
expected_outputs: tuple[str, ...]
+ """List of file extensions expected to be emitted by the generator script.
+
+ Output files will all be placed in the ``out`` folder.
+ """
compilable_output: str | None = None
+ """File extension of the output file that can be compiled.
+
+ If this is set, and the expected file exists, the ``compile_cmd`` will be
+ executed to check for error-free compilation of the output.
+ """
+
compile_cmd: str = f"g++ --std=c++17 -I {THIS_DIR}/deps/mdspan/include"
+ """Command to be invoked to compile the generated source file."""
SCRIPTS = [
ScriptInfo("SimpleJacobi", ("h", "cpp"), compilable_output="cpp"),
ScriptInfo("SimpleClasses", ("h", "cpp")),
+ ScriptInfo("Variables", ("h", "cpp"), compilable_output="cpp"),
]
@pytest.mark.parametrize("script_info", SCRIPTS)
def test_generator_script(script_info: ScriptInfo):
+ """Test a generator script defined by ``script_info``.
+
+ The generator script will be run, with its output placed in the ``out`` folder.
+ If it is successful, its output files will be compared against
+ any files of the same name from the ``expected`` folder.
+ Finally, if any compilable files are specified, the test will attempt to compile them.
+ """
+
script_name = script_info.script_name
script_file = path.join(SCRIPTS_DIR, script_name + ".py")