# 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}"