Fix nested workunits

pykokkos/core/compiler.py

@@ -18,7 +18,7 @@
 import pykokkos.kokkos_manager as km
 
 from .cpp_setup import CppSetup
-from .module_setup import EntityMetadata, ModuleSetup
+from .module_setup import EntityMetadata, ModuleSetup, get_callable_ast_signature
 
 
 @dataclass
@@ -80,7 +80,17 @@ def fuse_objects(
         pk_imports: List[str] = []
         for m in metadata:
             parser = self.get_parser(m.path)
-            entity: PyKokkosEntity = parser.get_entity(m.name)
+            runtime_entity = (
+                m.entity
+                if callable(m.entity)
+                and get_callable_ast_signature(m.entity) is not None
+                else None
+            )
+            entity: PyKokkosEntity = (
+                parser.get_entity(m.name)
+                if runtime_entity is None
+                else parser.get_entity(m.name, runtime_entity)
+            )
 
             for c in parser.get_classtypes():
                 if c.name in pyk_classtype_ids:
@@ -155,13 +165,25 @@ def compile_object(
         parser = self.get_parser(metadata[0].path)
 
         if len(metadata) == 1:
-            entity = parser.get_entity(metadata[0].name)
+            runtime_entity = (
+                metadata[0].entity
+                if callable(metadata[0].entity)
+                and get_callable_ast_signature(metadata[0].entity) is not None
+                else None
+            )
+            entity = (
+                parser.get_entity(metadata[0].name)
+                if runtime_entity is None
+                else parser.get_entity(metadata[0].name, runtime_entity)
+            )
             classtypes = parser.get_classtypes()
         else:
             # Avoid fusing the ASTs before checking if it was already compiled
             entity, classtypes = self.fuse_objects(metadata, fuse_ASTs=False, **kwargs)
 
-        hash: str = self.members_hash(entity.path, entity.name, types_signature)
+        hash: str = self.members_hash(
+            entity.path, entity.name, module_setup.ast_signature, types_signature
+        )
 
         types_inferred: bool = updated_types is not None
         decorator_inferred: bool = updated_decorator is not None
@@ -393,21 +415,26 @@ def read_defaults(self) -> Optional[CompilationDefaults]:
         return defaults
 
     def members_hash(
-        self, path: List[str], name: str, types_signature: Optional[str]
+        self,
+        path: List[str],
+        name: str,
+        ast_signature: str,
+        types_signature: Optional[str],
     ) -> str:
         """
         Map from entity path and name to a string to index members
 
         :param path: the path to the file containing the entity
         :param name: the name of the entity
+        :param ast_signature: signature of the translated AST and dependencies
         :param types_signature: string signature of inferred parameter types
         :returns: the hash of the entity
         """
 
         return (
-            f"{path}_{name}"
+            f"{path}_{name}_{ast_signature}"
             if types_signature is None
-            else f"{path}_{name}_{types_signature}"
+            else f"{path}_{name}_{ast_signature}_{types_signature}"
         )
 
     def extract_members(

pykokkos/core/module_setup.py

@@ -1,10 +1,12 @@
 from dataclasses import dataclass
+import ast
 import hashlib
 import inspect
 import os
 from pathlib import Path
 import sys
 import sysconfig
+import textwrap
 from typing import Callable, List, Optional, Set, Union
 
 from pykokkos.interface import ExecutionSpace
@@ -15,6 +17,38 @@
 BASE_DIR: str = ".pykokkos"
 
 
+def get_callable_ast_signature(entity: Callable) -> Optional[str]:
+    """Hash a nested workunit or one with resolved @pk.function dependencies."""
+
+    pending = [entity]
+    visited = set()
+    trees = []
+    has_dependencies = False
+
+    while pending:
+        function = pending.pop()
+        if id(function) in visited:
+            continue
+        visited.add(id(function))
+
+        tree = ast.parse(textwrap.dedent(inspect.getsource(function)))
+        trees.append(ast.dump(tree))
+        context = inspect.getclosurevars(function)
+        bindings = {**context.globals, **context.nonlocals}
+        for call in ast.walk(tree):
+            if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Name):
+                continue
+            target = bindings.get(call.func.id)
+            if callable(target) and getattr(target, "_pk_function", False):
+                has_dependencies = True
+                pending.append(target)
+
+    if not has_dependencies and "<locals>" not in entity.__qualname__:
+        return None
+
+    return hashlib.md5("".join(trees).encode()).hexdigest()
+
+
 @dataclass
 class EntityMetadata:
     """

pykokkos/core/parsers/parser.py

@@ -109,23 +109,81 @@ def get_classtypes(self) -> List[PyKokkosEntity]:
 
         return list(self.classtypes.values())
 
-    def get_entity(self, name: str) -> PyKokkosEntity:
+    def get_entity(
+        self, name: str, runtime_entity: Optional[Callable] = None
+    ) -> PyKokkosEntity:
         """
         Get the parsed entity
 
         :param name: the name of the functor
         :returns: the PyKokkosEntity representation of the entity
         """
 
+        is_nested = runtime_entity is not None and "<locals>" in getattr(
+            runtime_entity, "__qualname__", ""
+        )
+        if is_nested:
+            return self.get_runtime_workunit(runtime_entity)
+
         if name in self.workloads:
-            return self.workloads[name]
+            entity = self.workloads[name]
+            if runtime_entity is not None:
+                setattr(entity, "runtime_entity", runtime_entity)
+            return entity
         if name in self.functors:
-            return self.functors[name]
+            entity = self.functors[name]
+            if runtime_entity is not None:
+                setattr(entity, "runtime_entity", runtime_entity)
+            return entity
         if name in self.workunits:
-            return self.workunits[name]
+            entity = self.workunits[name]
+            if runtime_entity is not None:
+                setattr(entity, "runtime_entity", runtime_entity)
+            return entity
+
+        if runtime_entity is not None:
+            return self.get_runtime_workunit(runtime_entity)
 
         raise RuntimeError(f"Entity '{name}' not found by parser")
 
+    def get_runtime_workunit(self, function: Callable) -> PyKokkosEntity:
+        """Create an entity for a workunit represented by a runtime callable."""
+
+        function_def = self.get_function_def(function)
+        if not self.is_workunit(function_def, self.pk_import):
+            raise RuntimeError(f"Function '{function.__qualname__}' is not a workunit")
+
+        start = function_def.lineno - 1
+        stop = function_def.end_lineno or function_def.lineno
+        entity = PyKokkosEntity(
+            PyKokkosStyles.workunit,
+            function_def.name,
+            function_def,
+            self.tree,
+            (self.lines[start:stop], start),
+            self.path,
+            self.pk_import,
+        )
+        setattr(entity, "runtime_entity", function)
+        return entity
+
+    def get_function_def(self, function: Callable) -> ast.FunctionDef:
+        """Find the AST node corresponding to a runtime function."""
+
+        first_line = function.__code__.co_firstlineno
+        candidates = [
+            node
+            for node in ast.walk(self.tree)
+            if isinstance(node, ast.FunctionDef) and node.name == function.__name__
+        ]
+        for node in candidates:
+            decorator_lines = [decorator.lineno for decorator in node.decorator_list]
+            start = min(decorator_lines, default=node.lineno)
+            if start <= first_line <= node.lineno:
+                return node
+
+        raise RuntimeError(f"Function '{function.__qualname__}' not found by parser")
+
     def get_entities(self, style: PyKokkosStyles) -> Dict[str, PyKokkosEntity]:
         """
         Get the entities from path that are of a particular style

pykokkos/core/runtime.py

@@ -44,7 +44,12 @@
 import pykokkos.kokkos_manager as km
 
 from .compiler import Compiler
-from .module_setup import EntityMetadata, get_metadata, ModuleSetup
+from .module_setup import (
+    EntityMetadata,
+    ModuleSetup,
+    get_callable_ast_signature,
+    get_metadata,
+)
 from .run_debug import run_workload_debug, run_workunit_debug
 
 
@@ -415,10 +420,20 @@ def execute_workunit(
                 parsers = [
                     self.compiler.get_parser(get_metadata(e).path) for e in workunit
                 ]
-                entity_trees = [
-                    this_parser.get_entity(get_metadata(this_entity).name).AST
-                    for this_entity, this_parser in zip(workunit, parsers)
-                ]
+                entity_trees = []
+                for this_entity, this_parser in zip(workunit, parsers):
+                    metadata = get_metadata(this_entity)
+                    runtime_entity = (
+                        this_entity
+                        if get_callable_ast_signature(this_entity) is not None
+                        else None
+                    )
+                    parsed_entity = (
+                        this_parser.get_entity(metadata.name)
+                        if runtime_entity is None
+                        else this_parser.get_entity(metadata.name, runtime_entity)
+                    )
+                    entity_trees.append(parsed_entity.AST)
                 restrict_kwargs, _ = fuse_workunit_kwargs_and_params(
                     entity_trees, workunit_kwargs, f"parallel_{operation}"
                 )
@@ -434,10 +449,20 @@ def execute_workunit(
 
         # Set ast signature
         if isinstance(parser, list):
-            ast_signature = "".join([p.signature for p in parser])
+            signatures = [
+                (
+                    get_callable_ast_signature(entity) or this_parser.signature
+                    if not hasattr(entity, "__self__")
+                    else this_parser.signature
+                )
+                for entity, this_parser in zip(workunit, parser)
+            ]
+            ast_signature = "".join(signatures)
             ast_signature = hashlib.md5(ast_signature.encode()).hexdigest()
         else:
             ast_signature = parser.signature
+            if not hasattr(workunit, "__self__"):
+                ast_signature = get_callable_ast_signature(workunit) or ast_signature
 
         execution_space: ExecutionSpace = policy.space.space
         members: PyKokkosMembers = self.precompile_workunit(
@@ -646,10 +671,20 @@ def get_arguments(
                     parsers = [
                         self.compiler.get_parser(get_metadata(e).path) for e in entity
                     ]
-                    entity_trees = [
-                        this_parser.get_entity(get_metadata(this_entity).name).AST
-                        for this_entity, this_parser in zip(entity, parsers)
-                    ]
+                    entity_trees = []
+                    for this_entity, this_parser in zip(entity, parsers):
+                        metadata = get_metadata(this_entity)
+                        runtime_entity = (
+                            this_entity
+                            if get_callable_ast_signature(this_entity) is not None
+                            else None
+                        )
+                        parsed_entity = (
+                            this_parser.get_entity(metadata.name)
+                            if runtime_entity is None
+                            else this_parser.get_entity(metadata.name, runtime_entity)
+                        )
+                        entity_trees.append(parsed_entity.AST)
 
                     kwargs, _ = fuse_workunit_kwargs_and_params(
                         entity_trees, kwargs, f"parallel_{operation}"

pykokkos/core/translators/members.py

@@ -82,7 +82,7 @@ def extract(self, entity: PyKokkosEntity, classtypes: List[PyKokkosEntity]) -> N
                     param_begin = i + 1
                     # handle last_pass param for parallel_scan
                     if (
-                        i + 1 <= len(args)
+                        i + 1 < len(args)
                         and isinstance(args[i + 1].annotation, ast.Name)
                         and args[i + 1].annotation.id == "bool"
                     ):
@@ -114,8 +114,12 @@ def extract(self, entity: PyKokkosEntity, classtypes: List[PyKokkosEntity]) -> N
             )
         else:
             self.pk_workunits[cppast.DeclRefExpr(AST.name)] = AST
-            self.pk_functions = self.get_decorated_functions(
-                entity.full_AST, Decorator.KokkosFunction
+            self.pk_functions = (
+                {}
+                if getattr(entity, "runtime_entity", None) is not None
+                else self.get_decorated_functions(
+                    entity.full_AST, Decorator.KokkosFunction
+                )
             )
 
         self.classtype_methods = self.get_classtype_methods(classtypes)

pykokkos/core/translators/static.py

@@ -1,5 +1,6 @@
 import ast
 import copy
+import inspect
 import os
 import sys
 from typing import Dict, List, Optional, Set, Tuple, Union
@@ -86,6 +87,8 @@ def translate(
         else:
             self.parser = Parser(None, pk_import=entity.pk_import)
 
+        if getattr(entity, "runtime_entity", None) is not None:
+            self.resolve_functions(entity)
         entity.AST = self.add_parent_refs(entity.AST)
         for c in classtypes:
             c.AST = self.add_parent_refs(c.AST)
@@ -142,6 +145,44 @@ def translate(
 
         return functor, bindings, cast
 
+    def resolve_functions(self, entity: PyKokkosEntity) -> None:
+        """Resolve called @pk.function objects from the callable's context."""
+
+        runtime_entity = getattr(entity, "runtime_entity", None)
+        if runtime_entity is None or not isinstance(entity.AST, ast.FunctionDef):
+            return
+
+        pending = [(entity.AST, runtime_entity)]
+        visited: Set[Tuple[str, int]] = set()
+
+        while pending:
+            function_ast, function = pending.pop()
+            context = inspect.getclosurevars(function)
+            bindings = {**context.globals, **context.nonlocals}
+
+            for call in ast.walk(function_ast):
+                if not isinstance(call, ast.Call) or not isinstance(
+                    call.func, ast.Name
+                ):
+                    continue
+
+                name = call.func.id
+                target = bindings.get(name)
+                if not callable(target) or not getattr(target, "_pk_function", False):
+                    continue
+
+                key = (name, id(target))
+                if key in visited:
+                    continue
+                visited.add(key)
+
+                parser = Parser(inspect.getfile(target))
+                target_ast = copy.deepcopy(parser.get_function_def(target))
+                target_ast.name = name
+                reference = cppast.DeclRefExpr(name)
+                self.pk_members.pk_functions[reference] = target_ast
+                pending.append((target_ast, target))
+
     @staticmethod
     def add_parent_refs(classdef: ast.ClassDef) -> ast.ClassDef:
         """