diff --git a/pykokkos/core/compiler.py b/pykokkos/core/compiler.py index 64d4af75..16a9d0bf 100644 --- a/pykokkos/core/compiler.py +++ b/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( diff --git a/pykokkos/core/module_setup.py b/pykokkos/core/module_setup.py index 95fc4471..7ae809ce 100644 --- a/pykokkos/core/module_setup.py +++ b/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 "" not in entity.__qualname__: + return None + + return hashlib.md5("".join(trees).encode()).hexdigest() + + @dataclass class EntityMetadata: """ diff --git a/pykokkos/core/parsers/parser.py b/pykokkos/core/parsers/parser.py index 4724e5ad..0ad4e63b 100644 --- a/pykokkos/core/parsers/parser.py +++ b/pykokkos/core/parsers/parser.py @@ -109,7 +109,9 @@ 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 @@ -117,15 +119,71 @@ def get_entity(self, name: str) -> PyKokkosEntity: :returns: the PyKokkosEntity representation of the entity """ + is_nested = runtime_entity is not None and "" 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 diff --git a/pykokkos/core/runtime.py b/pykokkos/core/runtime.py index 046f3852..b5c7789e 100644 --- a/pykokkos/core/runtime.py +++ b/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}" diff --git a/pykokkos/core/translators/members.py b/pykokkos/core/translators/members.py index 3d8e8e64..5b4ec876 100644 --- a/pykokkos/core/translators/members.py +++ b/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) diff --git a/pykokkos/core/translators/static.py b/pykokkos/core/translators/static.py index d8287423..54e17e09 100644 --- a/pykokkos/core/translators/static.py +++ b/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: """ diff --git a/pykokkos/core/type_inference/args_type_inference.py b/pykokkos/core/type_inference/args_type_inference.py index 7abd06fb..6d16139f 100644 --- a/pykokkos/core/type_inference/args_type_inference.py +++ b/pykokkos/core/type_inference/args_type_inference.py @@ -6,7 +6,7 @@ from typing import Callable, Dict, Optional, Tuple, Union, List from pykokkos.core.fusion import fuse_workunit_kwargs_and_params -from pykokkos.core.module_setup import get_metadata +from pykokkos.core.module_setup import get_callable_ast_signature, get_metadata from pykokkos.interface import ( MDRangePolicy, TeamPolicy, @@ -489,7 +489,17 @@ def get_type_info( for this_workunit, this_parser in zip(workunit, parser): this_metadata = get_metadata(this_workunit) - this_tree = this_parser.get_entity(this_metadata.name).AST + runtime_entity = ( + this_metadata.entity + if get_callable_ast_signature(this_metadata.entity) is not None + else None + ) + this_entity = ( + this_parser.get_entity(this_metadata.name) + if runtime_entity is None + else this_parser.get_entity(this_metadata.name, runtime_entity) + ) + this_tree = this_entity.AST workunit_str = str(this_workunit) if not isinstance(this_tree, ast.FunctionDef): diff --git a/pykokkos/interface/decorators.py b/pykokkos/interface/decorators.py index e02f3cd0..bf096cc9 100644 --- a/pykokkos/interface/decorators.py +++ b/pykokkos/interface/decorators.py @@ -85,6 +85,7 @@ def classtype(func): def function(func): + func._pk_function = True return func diff --git a/tests/implicit_functions_helper.py b/tests/implicit_functions_helper.py new file mode 100644 index 00000000..cade12e3 --- /dev/null +++ b/tests/implicit_functions_helper.py @@ -0,0 +1,6 @@ +import pykokkos as pk + + +@pk.function +def external_helper(i: int) -> int: + return i + 3 diff --git a/tests/test_implicit_functions.py b/tests/test_implicit_functions.py new file mode 100644 index 00000000..0896d0ec --- /dev/null +++ b/tests/test_implicit_functions.py @@ -0,0 +1,69 @@ +import unittest + +import pykokkos as pk + +from tests.implicit_functions_helper import external_helper + + +@pk.function +def implicit_leaf(i: int) -> int: + return i + 1 + + +@pk.function +def implicit_helper(i: int) -> int: + return implicit_leaf(i) * 2 + + +def unrelated_host_function() -> None: + raise RuntimeError("This function must not be translated") + + +def log(value): + raise RuntimeError("This host wrapper must not shadow the math intrinsic") + + +@pk.workunit +def implicit_workunit(i: int, acc: pk.Acc[pk.int64]) -> None: + acc += implicit_helper(i) + + +@pk.workunit +def intrinsic_workunit(i: int, view: pk.View1D[pk.double]) -> None: + view[i] = log(view[i]) + + +@pk.workunit +def cross_module_workunit(i: int, acc: pk.Acc[pk.int64]) -> None: + acc += external_helper(i) + + +def make_implicit_workunit(functor): + @pk.workunit + def workunit(i: int, acc: pk.Acc[pk.int64]) -> None: + acc += functor(i) + + return workunit + + +class TestImplicitFunctions(unittest.TestCase): + def test_executes_implicit_functions(self): + result = pk.parallel_reduce(10, implicit_workunit) + + self.assertEqual(110, result) + + def test_executes_nested_workunit_with_captured_function(self): + workunit = make_implicit_workunit(implicit_helper) + + result = pk.parallel_reduce(10, workunit) + + self.assertEqual(110, result) + + def test_executes_function_from_another_module(self): + result = pk.parallel_reduce(10, cross_module_workunit) + + self.assertEqual(75, result) + + +if __name__ == "__main__": + unittest.main()