Skip to content
Snippets Groups Projects
kernel_namespace.py 1.39 KiB
Newer Older
# from .context import SfgContext
from pystencils import CreateKernelConfig, create_kernel
from pystencils.astnodes import KernelFunction

class SfgKernelNamespace:
    def __init__(self, ctx, name: str):
        self._ctx = ctx
        self._name = name
        self._asts = dict()

    @property
    def name(self):
        return self._name

    @property
    def asts(self):
        yield from self._asts.values()

    def add(self, ast: KernelFunction):
        """Adds an existing pystencils AST to this namespace."""
        astname = ast.function_name
        if astname in self._asts:
            raise ValueError(f"Duplicate ASTs: An AST with name {astname} was already registered in namespace {self._name}")

        self._asts[astname] = ast

        return SfgKernelHandle(self._ctx, astname, self)

    def create(self, assignments, config: CreateKernelConfig):
        ast = create_kernel(assignments, config)
        return self.add(ast)


class SfgKernelHandle:
    def __init__(self, ctx, name: str, namespace: SfgKernelNamespace):
        self._ctx = ctx
        self._name = name
        self._namespace = namespace

    @property
    def kernel_name(self):
        return self._name

    @property
    def kernel_namespace(self):
        return self._namespace

    @property
    def fully_qualified_name(self):
        return f"{self.ctx.root_namespace}::{self.kernel_namespace}::{self.kernel_name}"