From fa347fe5934e9adcf0b7404f4de50cd742306f86 Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Thu, 6 Feb 2025 17:22:30 +0100
Subject: [PATCH] updating the composer WIP
---
src/pystencilssfg/composer/basic_composer.py | 210 +++++++++------
src/pystencilssfg/context.py | 261 +++++--------------
src/pystencilssfg/ir/source_components.py | 21 +-
3 files changed, 215 insertions(+), 277 deletions(-)
diff --git a/src/pystencilssfg/composer/basic_composer.py b/src/pystencilssfg/composer/basic_composer.py
index db671b9..c0b420f 100644
--- a/src/pystencilssfg/composer/basic_composer.py
+++ b/src/pystencilssfg/composer/basic_composer.py
@@ -6,7 +6,7 @@ import sympy as sp
from functools import reduce
from warnings import warn
-from pystencils import Field
+from pystencils import Field, CreateKernelConfig, create_kernel
from pystencils.codegen import Kernel
from pystencils.types import create_type, UserTypeSpec
@@ -31,13 +31,15 @@ from ..ir.postprocessing import (
)
from ..ir.source_components import (
SfgFunction,
- SfgHeaderInclude,
SfgKernelNamespace,
SfgKernelHandle,
SfgClass,
SfgConstructor,
SfgMemberVariable,
SfgClassKeyword,
+ SfgEntityDecl,
+ SfgEntityDef,
+ SfgNamespaceBlock,
)
from ..lang import (
VarLike,
@@ -61,6 +63,7 @@ from ..exceptions import SfgException
class SfgIComposer(ABC):
def __init__(self, ctx: SfgContext):
self._ctx = ctx
+ self._cursor = ctx.cursor
@property
def context(self):
@@ -80,6 +83,66 @@ SequencerArg: TypeAlias = tuple | ExprLike | SfgCallTreeNode | SfgNodeBuilder
"""Valid arguments to `make_sequence` and any sequencer that uses it."""
+class KernelsAdder:
+ def __init__(self, ctx: SfgContext, loc: SfgNamespaceBlock):
+ self._ctx = ctx
+ self._loc = SfgNamespaceBlock
+ assert isinstance(loc.namespace, SfgKernelNamespace)
+ self._kernel_namespace = loc.namespace
+
+ def add(self, kernel: Kernel, name: str | None = None):
+ """Adds an existing pystencils AST to this namespace.
+ If a name is specified, the AST's function name is changed."""
+ if name is None:
+ kernel_name = kernel.name
+ else:
+ kernel_name = name
+
+ if self._kernel_namespace.find_kernel(kernel_name) is not None:
+ raise ValueError(
+ f"Duplicate kernels: A kernel called {kernel_name} already exists "
+ f"in namespace {self._kernel_namespace.fqname}"
+ )
+
+ if name is not None:
+ kernel.name = kernel_name
+
+ khandle = SfgKernelHandle(kernel_name, self._kernel_namespace, kernel)
+ self._kernel_namespace.add_kernel(khandle)
+
+ for header in kernel.required_headers:
+ # TODO: Find current source file by traversing namespace blocks upward?
+ self._ctx.impl_file.includes.append(HeaderFile.parse(header))
+
+ return khandle
+
+ def create(
+ self,
+ assignments,
+ name: str | None = None,
+ config: CreateKernelConfig | None = None,
+ ):
+ """Creates a new pystencils kernel from a list of assignments and a configuration.
+ This is a wrapper around `pystencils.create_kernel`
+ with a subsequent call to `add`.
+ """
+ if config is None:
+ config = CreateKernelConfig()
+
+ if name is not None:
+ if self._kernel_namespace.find_kernel(name) is not None:
+ raise ValueError(
+ f"Duplicate kernels: A kernel called {name} already exists "
+ f"in namespace {self._kernel_namespace.fqname}"
+ )
+
+ config.function_name = name
+
+ # type: ignore
+ kernel = create_kernel(assignments, config=config)
+ return self.add(kernel)
+
+
class SfgBasicComposer(SfgIComposer):
"""Composer for basic source components, and base class for all composer mix-ins."""
@@ -87,7 +150,7 @@ class SfgBasicComposer(SfgIComposer):
ctx: SfgContext = sfg if isinstance(sfg, SfgContext) else sfg.context
super().__init__(ctx)
- def prelude(self, content: str):
+ def prelude(self, content: str, end: str = "\n"):
"""Append a string to the prelude comment, to be printed at the top of both generated files.
The string should not contain C/C++ comment delimiters, since these will be added automatically
@@ -105,7 +168,11 @@ class SfgBasicComposer(SfgIComposer):
*/
"""
- self._ctx.append_to_prelude(content)
+ for f in self._ctx.files:
+ if f.prelude is None:
+ f.prelude = content + end
+ else:
+ f.prelude += content + end
def code(self, *code: str):
"""Add arbitrary lines of code to the generated header file.
@@ -126,7 +193,7 @@ class SfgBasicComposer(SfgIComposer):
"""
for c in code:
- self._ctx.add_definition(c)
+ self._cursor.write_header(c)
def define(self, *definitions: str):
from warnings import warn
@@ -139,34 +206,9 @@ class SfgBasicComposer(SfgIComposer):
self.code(*definitions)
- def define_once(self, *definitions: str):
- """Add unique definitions to the header file.
-
- Each code string given to `define_once` will only be added if the exact same string
- was not already added before.
- """
- for definition in definitions:
- if all(d != definition for d in self._ctx.definitions()):
- self._ctx.add_definition(definition)
-
def namespace(self, namespace: str):
- """Set the inner code namespace. Throws an exception if a namespace was already set.
-
- :Example:
-
- After adding the following to your generator script:
-
- >>> sfg.namespace("codegen_is_awesome")
-
- All generated code will be placed within that namespace:
-
- .. code-block:: C++
-
- namespace codegen_is_awesome {
- /* all generated code */
- }
- """
- self._ctx.set_namespace(namespace)
+ # TODO: Enter into a new namespace context
+ raise NotImplementedError()
def generate(self, generator: CustomGenerator):
"""Invoke a custom code generator with the underlying context."""
@@ -183,18 +225,16 @@ class SfgBasicComposer(SfgIComposer):
sfg.kernels.add(ast, "kernel_name")
sfg.kernels.create(assignments, "kernel_name", config)
"""
- return self._ctx._default_kernel_namespace
+ return self.kernel_namespace("kernels")
def kernel_namespace(self, name: str) -> SfgKernelNamespace:
"""Return the kernel namespace of the given name, creating it if it does not exist yet."""
- kns = self._ctx.get_kernel_namespace(name)
- if kns is None:
- kns = SfgKernelNamespace(self._ctx, name)
- self._ctx.add_kernel_namespace(kns)
+ # TODO: Find the default kernel namespace as a child entity of the current
+ # namespace, or create it if it does not exist
+ # Then create a new namespace block, place it at the cursor position, and expose
+ # it to the user via an adder
- return kns
-
- def include(self, header_file: str, private: bool = False):
+ def include(self, header_file: str | HeaderFile, private: bool = False):
"""Include a header file.
Args:
@@ -214,7 +254,14 @@ class SfgBasicComposer(SfgIComposer):
#include <vector>
#include "custom.h"
"""
- self._ctx.add_include(SfgHeaderInclude(HeaderFile.parse(header_file), private))
+ header_file = HeaderFile.parse(header_file)
+
+ if private:
+ if self._ctx.impl_file is None:
+ raise ValueError("Cannot emit a private include since no implementation file is being generated")
+ self._ctx.impl_file.includes.append(header_file)
+ else:
+ self._ctx.header_file.includes.append(header_file)
def numpy_struct(
self, name: str, dtype: np.dtype, add_constructor: bool = True
@@ -224,11 +271,9 @@ class SfgBasicComposer(SfgIComposer):
Returns:
The created class object
"""
- if self._ctx.get_class(name) is not None:
- raise SfgException(f"Class with name {name} already exists.")
-
- cls = _struct_from_numpy_dtype(name, dtype, add_constructor=add_constructor)
- self._ctx.add_class(cls)
+ cls = self._struct_from_numpy_dtype(name, dtype, add_constructor=add_constructor)
+ self._ctx.add_entity(cls)
+ self._cursor.write_header(SfgEntityDecl(cls))
return cls
def kernel_function(
@@ -281,15 +326,18 @@ class SfgBasicComposer(SfgIComposer):
)
returns = return_type
- if self._ctx.get_function(name) is not None:
- raise ValueError(f"Function {name} already exists.")
-
def sequencer(*args: SequencerArg):
tree = make_sequence(*args)
func = SfgFunction(
- name, tree, return_type=create_type(returns), inline=inline
+ name, self._cursor.current_namespace, tree, return_type=create_type(returns), inline=inline
)
- self._ctx.add_function(func)
+ self._ctx.add_entity(func)
+
+ if inline:
+ self._cursor.write_header(SfgEntityDef(func))
+ else:
+ self._cursor.write_header(SfgEntityDecl(func))
+ self._cursor.write_impl(SfgEntityDef(func))
return sequencer
@@ -482,6 +530,36 @@ class SfgBasicComposer(SfgIComposer):
(asvar(c) if isinstance(c, _VarLike) else c) for c in lhs_components
]
return SfgDeferredVectorMapping(components, rhs)
+
+ def _struct_from_numpy_dtype(
+ self, struct_name: str, dtype: np.dtype, add_constructor: bool = True
+ ):
+ cls = SfgClass(struct_name, self._cursor.current_namespace, class_keyword=SfgClassKeyword.STRUCT)
+
+ fields = dtype.fields
+ if fields is None:
+ raise SfgException(f"Numpy dtype {dtype} is not a structured type.")
+
+ constr_params = []
+ constr_inits = []
+
+ for member_name, type_info in fields.items():
+ member_type = create_type(type_info[0])
+
+ member = SfgMemberVariable(member_name, member_type)
+
+ arg = SfgVar(f"{member_name}_", member_type)
+
+ cls.default.append_member(member)
+
+ constr_params.append(arg)
+ constr_inits.append(f"{member}({arg})")
+
+ if add_constructor:
+ cls.default.append_member(SfgEntityDef(SfgConstructor(constr_params, constr_inits)))
+
+ return cls
+
def make_statements(arg: ExprLike) -> SfgStatements:
@@ -628,33 +706,3 @@ class SfgSwitchBuilder(SfgNodeBuilder):
def resolve(self) -> SfgCallTreeNode:
return SfgSwitch(make_statements(self._switch_arg), self._cases, self._default)
-
-
-def _struct_from_numpy_dtype(
- struct_name: str, dtype: np.dtype, add_constructor: bool = True
-):
- cls = SfgClass(struct_name, class_keyword=SfgClassKeyword.STRUCT)
-
- fields = dtype.fields
- if fields is None:
- raise SfgException(f"Numpy dtype {dtype} is not a structured type.")
-
- constr_params = []
- constr_inits = []
-
- for member_name, type_info in fields.items():
- member_type = create_type(type_info[0])
-
- member = SfgMemberVariable(member_name, member_type)
-
- arg = SfgVar(f"{member_name}_", member_type)
-
- cls.default.append_member(member)
-
- constr_params.append(arg)
- constr_inits.append(f"{member}({arg})")
-
- if add_constructor:
- cls.default.append_member(SfgConstructor(constr_params, constr_inits))
-
- return cls
diff --git a/src/pystencilssfg/context.py b/src/pystencilssfg/context.py
index 17537a2..577dcbd 100644
--- a/src/pystencilssfg/context.py
+++ b/src/pystencilssfg/context.py
@@ -1,83 +1,45 @@
-from typing import Generator, Sequence, Any
+from __future__ import annotations
+from typing import Sequence, Any, Generator
from .config import CodeStyle
from .ir.source_components import (
- SfgHeaderInclude,
+ SfgSourceFile,
+ SfgNamespace,
SfgKernelNamespace,
- SfgFunction,
+ SfgNamespaceBlock,
+ SfgNamespaceElement,
+ SfgCodeEntity,
SfgClass,
)
from .exceptions import SfgException
class SfgContext:
- """Represents a header/implementation file pair in the code generator.
-
- **Source File Properties and Components**
-
- The SfgContext collects all properties and components of a header/implementation
- file pair (or just the header file, if header-only generation is used).
- These are:
-
- - The code namespace, which is combined from the `outer_namespace`
- and the `pystencilssfg.SfgContext.inner_namespace`. The outer namespace is meant to be set
- externally e.g. by the project configuration, while the inner namespace is meant to be set by the generator
- script.
- - The `prelude comment` is a block of text printed as a comment block
- at the top of both generated files. Typically, it contains authorship and licence information.
- - The set of included header files (`pystencilssfg.SfgContext.includes`).
- - Custom `definitions`, which are just arbitrary code strings.
- - Any number of kernel namespaces (`pystencilssfg.SfgContext.kernel_namespaces`), within which *pystencils*
- kernels are managed.
- - Any number of functions (`pystencilssfg.SfgContext.functions`), which are meant to serve as wrappers
- around kernel calls.
- - Any number of classes (`pystencilssfg.SfgContext.classes`), which can be used to build more extensive wrappers
- around kernels.
-
- **Order of Definitions**
-
- To honor C/C++ use-after-declare rules, the context preserves the order in which definitions, functions and classes
- are added to it.
- The header file printers implemented in *pystencils-sfg* will print the declarations accordingly.
- The declarations can retrieved in order of definition via `declarations_ordered`.
- """
+ """Manages context information during the execution of a generator script."""
def __init__(
self,
+ header_file: SfgSourceFile,
+ impl_file: SfgSourceFile,
outer_namespace: str | None = None,
codestyle: CodeStyle | None = None,
argv: Sequence[str] | None = None,
project_info: Any = None,
):
- """
- Args:
- outer_namespace: Qualified name of the outer code namespace
- codestyle: Code style that should be used by the code emitter
- argv: The generator script's command line arguments.
- Reserved for internal use by the [SourceFileGenerator][pystencilssfg.SourceFileGenerator].
- project_info: Project-specific information provided by a build system.
- Reserved for internal use by the [SourceFileGenerator][pystencilssfg.SourceFileGenerator].
- """
self._argv = argv
self._project_info = project_info
- self._default_kernel_namespace = SfgKernelNamespace(self, "kernels")
self._outer_namespace = outer_namespace
self._inner_namespace: str | None = None
self._codestyle = codestyle if codestyle is not None else CodeStyle()
- # Source Components
- self._prelude: str = ""
- self._includes: list[SfgHeaderInclude] = []
- self._definitions: list[str] = []
- self._kernel_namespaces = {
- self._default_kernel_namespace.name: self._default_kernel_namespace
- }
- self._functions: dict[str, SfgFunction] = dict()
- self._classes: dict[str, SfgClass] = dict()
+ self._header_file = header_file
+ self._impl_file = impl_file
- self._declarations_ordered: list[str | SfgFunction | SfgClass] = list()
+ self._entities: dict[str, SfgCodeEntity] = dict()
+
+ self._cursor: SfgCursor
@property
def argv(self) -> Sequence[str]:
@@ -100,163 +62,74 @@ class SfgContext:
"""Outer code namespace. Set by constructor argument `outer_namespace`."""
return self._outer_namespace
- @property
- def inner_namespace(self) -> str | None:
- """Inner code namespace. Set by `set_namespace`."""
- return self._inner_namespace
-
- @property
- def fully_qualified_namespace(self) -> str | None:
- """Combined outer and inner namespaces, as `outer_namespace::inner_namespace`."""
- match (self.outer_namespace, self.inner_namespace):
- case None, None:
- return None
- case outer, None:
- return outer
- case None, inner:
- return inner
- case outer, inner:
- return f"{outer}::{inner}"
- case _:
- assert False
-
@property
def codestyle(self) -> CodeStyle:
"""The code style object for this generation context."""
return self._codestyle
- # ----------------------------------------------------------------------------------------------
- # Prelude, Includes, Definitions, Namespace
- # ----------------------------------------------------------------------------------------------
-
@property
- def prelude_comment(self) -> str:
- """The prelude is a comment block printed at the top of both generated files."""
- return self._prelude
-
- def append_to_prelude(self, code_str: str):
- """Append a string to the prelude comment.
-
- The string should not contain
- C/C++ comment delimiters, since these will be added automatically during
- code generation.
- """
- if self._prelude:
- self._prelude += "\n"
-
- self._prelude += code_str
-
- if not code_str.endswith("\n"):
- self._prelude += "\n"
-
- def includes(self) -> Generator[SfgHeaderInclude, None, None]:
- """Includes of headers. Public includes are added to the header file, private includes
- are added to the implementation file."""
- yield from self._includes
-
- def add_include(self, include: SfgHeaderInclude):
- self._includes.append(include)
-
- def definitions(self) -> Generator[str, None, None]:
- """Definitions are arbitrary custom lines of code."""
- yield from self._definitions
-
- def add_definition(self, definition: str):
- """Add a custom code string to the header file."""
- self._definitions.append(definition)
- self._declarations_ordered.append(definition)
-
- def set_namespace(self, namespace: str):
- """Set the inner code namespace.
-
- Throws an exception if the namespace was already set.
- """
- if self._inner_namespace is not None:
- raise SfgException("The code namespace was already set.")
-
- self._inner_namespace = namespace
-
- # ----------------------------------------------------------------------------------------------
- # Kernel Namespaces
- # ----------------------------------------------------------------------------------------------
+ def header_file(self) -> SfgSourceFile:
+ return self._header_file
@property
- def default_kernel_namespace(self) -> SfgKernelNamespace:
- """The default kernel namespace."""
- return self._default_kernel_namespace
-
- def kernel_namespaces(self) -> Generator[SfgKernelNamespace, None, None]:
- """Iterator over all registered kernel namespaces."""
- yield from self._kernel_namespaces.values()
-
- def get_kernel_namespace(self, str) -> SfgKernelNamespace | None:
- """Retrieve a kernel namespace by name, or `None` if it does not exist."""
- return self._kernel_namespaces.get(str)
-
- def add_kernel_namespace(self, namespace: SfgKernelNamespace):
- """Adds a new kernel namespace.
-
- If a kernel namespace of the same name already exists, throws an exception.
- """
- if namespace.name in self._kernel_namespaces:
- raise ValueError(f"Duplicate kernel namespace: {namespace.name}")
+ def impl_file(self) -> SfgSourceFile | None:
+ return self._impl_file
- self._kernel_namespaces[namespace.name] = namespace
-
- # ----------------------------------------------------------------------------------------------
- # Functions
- # ----------------------------------------------------------------------------------------------
-
- def functions(self) -> Generator[SfgFunction, None, None]:
- """Iterator over all registered functions."""
- yield from self._functions.values()
-
- def get_function(self, name: str) -> SfgFunction | None:
- """Retrieve a function by name. Returns `None` if no function of the given name exists."""
- return self._functions.get(name, None)
-
- def add_function(self, func: SfgFunction):
- """Adds a new function.
-
- If a function or class with the same name exists already, throws an exception.
- """
- if func.name in self._functions or func.name in self._classes:
- raise SfgException(f"Duplicate function: {func.name}")
+ @property
+ def cursor(self) -> SfgCursor:
+ return self._cursor
- self._functions[func.name] = func
- self._declarations_ordered.append(func)
+ @property
+ def files(self) -> Generator[SfgSourceFile, None, None]:
+ yield self._header_file
+ if self._impl_file is not None:
+ yield self._impl_file
- # ----------------------------------------------------------------------------------------------
- # Classes
- # ----------------------------------------------------------------------------------------------
+ def get_entity(self, fqname: str) -> SfgCodeEntity | None:
+ # TODO: Only track top-level entities here, traverse namespaces to find qualified entities
+ return self._entities.get(fqname, None)
- def classes(self) -> Generator[SfgClass, None, None]:
- """Iterator over all registered classes."""
- yield from self._classes.values()
+ def add_entity(self, entity: SfgCodeEntity) -> None:
+ fqname = entity.fqname
+ if fqname in self._entities:
+ raise ValueError(f"Another entity with name {fqname} already exists")
+ self._entities[fqname] = entity
- def get_class(self, name: str) -> SfgClass | None:
- """Retrieve a class by name, or `None` if the class does not exist."""
- return self._classes.get(name, None)
- def add_class(self, cls: SfgClass):
- """Add a class.
+class SfgCursor:
+ """Cursor that tracks the current location in the source file(s) during execution of the generator script."""
- Throws an exception if a class or function of the same name exists already.
- """
- if cls.class_name in self._classes or cls.class_name in self._functions:
- raise SfgException(f"Duplicate class: {cls.class_name}")
+ def __init__(self, ctx: SfgContext, namespace: str | None = None) -> None:
+ self._ctx = ctx
- self._classes[cls.class_name] = cls
- self._declarations_ordered.append(cls)
+ self._cur_namespace: SfgNamespace | None
+ if namespace is not None:
+ self._cur_namespace = ctx.get_namespace(namespace)
+ else:
+ self._cur_namespace = None
- # ----------------------------------------------------------------------------------------------
- # Declarations in order of addition
- # ----------------------------------------------------------------------------------------------
+ self._loc: dict[SfgSourceFile, list[SfgNamespaceElement]]
+ for f in self._ctx.files:
+ if self._cur_namespace is not None:
+ block = SfgNamespaceBlock(self._cur_namespace)
+ f.elements.append(block)
+ self._loc[f] = block.elements
+ else:
+ self._loc[f] = f.elements
- def declarations_ordered(
- self,
- ) -> Generator[str | SfgFunction | SfgClass, None, None]:
- """All declared definitions, classes and functions in the order they were added.
+ # TODO: Enter and exit namespace blocks
- Awareness about order is necessary due to the C++ declare-before-use rules."""
- yield from self._declarations_ordered
+ @property
+ def current_namespace(self) -> SfgNamespace | None:
+ return self._cur_namespace
+
+ def write_header(self, elem: SfgNamespaceElement) -> None:
+ self._loc[self._ctx.header_file].append(elem)
+
+ def write_impl(self, elem: SfgNamespaceElement) -> None:
+ impl_file = self._ctx.impl_file
+ if impl_file is None:
+ raise SfgException(
+ f"Cannot write element {elem} to implemenation file since no implementation file is being generated."
+ )
+ self._loc[impl_file].append(elem)
diff --git a/src/pystencilssfg/ir/source_components.py b/src/pystencilssfg/ir/source_components.py
index d98e2d8..d8fb0f5 100644
--- a/src/pystencilssfg/ir/source_components.py
+++ b/src/pystencilssfg/ir/source_components.py
@@ -73,6 +73,8 @@ class SfgNamespace(SfgCodeEntity):
parent: Parent namespace enclosing this namespace
"""
+ # TODO: Namespaces must keep track of their child entities
+
class SfgKernelHandle(SfgCodeEntity):
"""Handle to a pystencils kernel."""
@@ -127,6 +129,17 @@ class SfgKernelNamespace(SfgNamespace):
def kernels(self) -> tuple[SfgKernelHandle, ...]:
return tuple(self._kernels.values())
+ def find_kernel(self, name: str) -> SfgKernelHandle | None:
+ return self._kernels.get(name, None)
+
+ def add_kernel(self, kernel: SfgKernelHandle):
+ if kernel.name in self._kernels:
+ raise ValueError(
+ f"Duplicate kernels: A kernel called {kernel.name} already exists "
+ f"in namespace {self.fqname}"
+ )
+ self._kernels[kernel.name] = kernel
+
def add(self, kernel: Kernel, name: str | None = None):
"""Adds an existing pystencils AST to this namespace.
If a name is specified, the AST's function name is changed."""
@@ -623,7 +636,7 @@ class SfgVisibilityBlock:
self._cls = cls
-class SfgNamespaceDef:
+class SfgNamespaceBlock:
"""A C++ namespace.
Each namespace has a `name` and a `parent`; its fully qualified name is given as
@@ -638,6 +651,10 @@ class SfgNamespaceDef:
self._namespace = namespace
self._elements: list[SfgNamespaceElement] = []
+ @property
+ def namespace(self) -> SfgNamespace:
+ return self._namespace
+
@property
def elements(self) -> list[SfgNamespaceElement]:
"""Sequence of source elements that make up the body of this namespace"""
@@ -648,7 +665,7 @@ class SfgNamespaceDef:
self._elements = list(elems)
-SfgNamespaceElement = str | SfgNamespaceDef | SfgEntityDecl | SfgEntityDef
+SfgNamespaceElement = str | SfgNamespaceBlock | SfgEntityDecl | SfgEntityDef
"""Elements that may be placed inside a namespace, including the global namespace."""
--
GitLab