kernel_namespace.py

# 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} already exists in namespace {self._name}")

        self._asts[astname] = ast

        return SfgKernelHandle(self._ctx, astname, self, ast.get_parameters())

    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, parameters):
        self._ctx = ctx
        self._name = name
        self._namespace = namespace
        self._parameters = parameters

    @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.name}::{self.kernel_name}"
    
    @property
    def parameters(self):
        return self._parameters