Decorator generation + Polished type Inference

pykokkos/core/compiler.py

@@ -8,7 +8,7 @@
 from typing import Dict, List, Optional
 from pykokkos.core.parsers import Parser, PyKokkosEntity, PyKokkosStyles
 from pykokkos.core.translators import PyKokkosMembers, StaticTranslator
-from pykokkos.interface import ExecutionSpace, UpdatedTypes
+from pykokkos.interface import ExecutionSpace, UpdatedTypes, UpdatedDecorator
 import pykokkos.kokkos_manager as km
 from .cpp_setup import CppSetup
 from .module_setup import EntityMetadata, ModuleSetup
@@ -51,25 +51,29 @@ def compile_object(
         module_setup: ModuleSetup,
         space: ExecutionSpace,
         force_uvm: bool,
+        updated_decorator: UpdatedDecorator,
         updated_types: Optional[UpdatedTypes] = None,
+        types_signature: Optional[str] = None
     ) -> Optional[PyKokkosMembers]:
         """
         Compile an entity object for a single execution space
 
         :param entity_object: the module_setup object containing module info
         :param space: the execution space to compile for
         :param force_uvm: whether CudaUVMSpace is enabled
+        :param updated_decorator: Object for decorator specifiers
         :param updated_types: Object with with inferred types
         :returns: the PyKokkos members obtained during translation
         """
 
         metadata = module_setup.metadata
         parser = self.get_parser(metadata.path)
-        types_signature = None if updated_types is None else updated_types.types_signature
+
         hash: str = self.members_hash(metadata.path, metadata.name, types_signature)
         entity: PyKokkosEntity = parser.get_entity(metadata.name)
 
         types_inferred: bool = updated_types is not None
+        decorator_inferred: bool = updated_decorator is not None
 
         if types_inferred and entity.style is not PyKokkosStyles.workunit:
             raise Exception(f"Types are required for style: {entity.style}")
@@ -78,6 +82,8 @@ def compile_object(
             if hash not in self.members: # True if pre-compiled
                 if types_inferred:
                     entity.AST = parser.fix_types(entity, updated_types)
+                if decorator_inferred:
+                    entity.AST = parser.fix_decorator(entity, updated_decorator)
                 self.members[hash] = self.extract_members(metadata)
 
             return self.members[hash]
@@ -88,7 +94,9 @@ def compile_object(
 
         if types_inferred:
             entity.AST = parser.fix_types(entity, updated_types)
-
+        if decorator_inferred:
+            entity.AST = parser.fix_decorator(entity, updated_decorator)
+
         if hash in self.members: # True if compiled with another execution space
             members = self.members[hash]
         else:

pykokkos/core/module_setup.py

@@ -82,6 +82,7 @@ def __init__(
         ModuleSetup constructor
 
         :param entity: the functor/workunit/workload
+        :param types_signature: hash/string to identify workunit signature against types
         """
 
         self.metadata: EntityMetadata
@@ -129,7 +130,7 @@ def get_output_dir(
         :param main: the path to the main file in the current PyKokkos application
         :param metadata: the metadata of the entity being compiled
         :param space: the execution space to compile for
-        :param types_signature: optional identifier string for inferred types of parameters
+        :param types_signature: optional identifier/hash string for types of parameters
         :returns: the path to the output directory for a specific execution space
         """
 

pykokkos/core/parsers/parser.py

@@ -1,12 +1,11 @@
 import ast
 import inspect
-import copy
 from dataclasses import dataclass
 from enum import Enum, auto
 from typing import Callable, Dict, List, Tuple, Union
 
 from pykokkos.core import cppast
-from pykokkos.interface import Decorator, UpdatedTypes, get_type_str
+from pykokkos.interface import Decorator, UpdatedTypes, UpdatedDecorator, get_type_str
 
 class PyKokkosStyles(Enum):
     """
@@ -157,11 +156,6 @@ def fix_types(self, entity: PyKokkosEntity, updated_types: UpdatedTypes) -> ast.
         if needs_reset:
             entity_tree = self.reset_entity_tree(entity_tree, updated_types)
 
-        # if modifications to layout decorator is needed, do not change original inferences
-        if len(updated_types.layout_change):
-            inferred_layouts = copy.deepcopy(updated_types.layout_change)
-            entity_tree.decorator_list = self.fix_view_layout(entity_tree, inferred_layouts)
-
         for arg_obj in entity_tree.args.args:
             # Type already provided by the user
             if arg_obj.arg not in updated_types.inferred_types:
@@ -182,23 +176,22 @@ def check_self(self, entity_tree: ast.AST) -> bool:
         :returns: True if a 'self' argument exists, False otherwise
         '''
 
-        for arg in entity_tree.args.args:
-            if arg.arg == "self":
-                return True
+        if entity_tree.args.args[0].arg == "self":
+            return True
         return False
 
-    def reset_entity_tree(self, entity_tree: ast.AST, updated_types: UpdatedTypes) -> ast.AST:
+    def reset_entity_tree(self, entity_tree: ast.AST, updated_obj: Union[UpdatedTypes, UpdatedDecorator]) -> ast.AST:
         '''
-        Remove type annotations and self argument from the entity tree. This allows
+        Remove the inferred type annotations and self argument from the entity tree. This allows
         the types to be inserted again if they change dynamically
 
-        :param entity_tree: Ast of pykokkos entity being resiet
-        :param updated_types: inferred types information
+        :param entity_tree: Ast of pykokkos entity being reset
+        :param updated_obj: inferred types/decorator object that must have original inspect param list
         :returns: updated entity ast as it would be in the first run
         '''
 
         args_list: List[ast.arg] = []
-        param_list = updated_types.param_list
+        param_list = updated_obj.param_list
         for param in param_list:
             arg_obj = ast.arg(arg=param.name)
             if param.annotation is not inspect._empty:
@@ -286,44 +279,38 @@ def get_annotation_node(self, type: str) -> ast.AST:
 
         return annotation_node
 
-    def fix_view_layout(self, node : ast.AST, layout_change: Dict[str, str]) -> List[ast.Call]:
+    def fix_decorator(self, entity : PyKokkosEntity, updated_decorator: UpdatedDecorator) -> ast.AST:
         '''
-        Construct the decorator list (as in AST) with the missing layout decorators for pykokkos views
+        Add the decorator list with the specifiers for pykokkos views to the workunit AST
 
         :param node: ast object for the entity
-        :param layout_change: Dict that maps [view -> layout]
+        :param updated_decorator: Object with dict that maps view to its layout, space and trait
         :returns: decorator list 
         '''
 
-        assert len(node.decorator_list), f"Decorator cannot be missing for pykokkos workunit {node.name}"
-        # Decorator list will have ast.Call object as first element if user has provided layout decorators
-        is_layout_given: bool = isinstance(node.decorator_list[0], ast.Call)
-
-        if is_layout_given:
-            # Filter out layouts already given by user
-            layout_change = self.filter_layout_change(node, layout_change)
-
-        if len(layout_change):
-            call_obj = None
-
-            if is_layout_given: # preserve the existing call object
-                call_obj = node.decorator_list[0]
-            else: 
-                call_obj= ast.Call()
-                call_obj.func = ast.Attribute(value=ast.Name(id=self.pk_import, ctx=ast.Load()), attr='workunit', ctx=ast.Load())
-                call_obj.args = []
-                call_obj.keywords = []
-
-            for view, layout in layout_change.items():
-                call_obj.keywords.append(self.get_keyword_node(view, layout))
-
-            return [call_obj]
-
-        # no change needed
-        return node.decorator_list
+        entity_tree = entity.AST
+        needs_reset: bool = self.check_self(entity_tree)
+        if needs_reset:
+            entity_tree = self.reset_entity_tree(entity_tree, updated_decorator)
+        assert len(entity_tree.decorator_list), f"Decorator cannot be missing for pykokkos workunit {entity_tree.name}"
+
+        if not len(updated_decorator.inferred_decorator):
+            # no change needed
+            return entity_tree.decorator_list
+
+        call_obj= ast.Call()
+        call_obj.func = ast.Attribute(value=ast.Name(id=self.pk_import, ctx=ast.Load()), attr='workunit', ctx=ast.Load())
+        call_obj.args = []
+        call_obj.keywords = []
+
+        for view, specifier_dict in updated_decorator.inferred_decorator.items():
+            call_obj.keywords.append(self.get_keyword_node(view, specifier_dict))
+
+        entity_tree.decorator_list = [call_obj]
+        return entity_tree
 
 
-    def get_keyword_node(self, view_name: str, layout: str) -> ast.keyword:
+    def get_keyword_node(self, view_name: str, specifiers: Dict[str, str]) -> ast.keyword:
         '''
         Make the ast.keyword node to be added to the decorator list
 
@@ -332,6 +319,23 @@ def get_keyword_node(self, view_name: str, layout: str) -> ast.keyword:
         :returns: corresponding ast.keyword node that can be added to decorator list 
         '''
 
+        skip_space: bool = False if specifiers['trait'] == "Unmanaged" else True
+        keywords_list: List[ast.keyword] = []
+        attr_names = {'layout' : 'Layout', 'space' : 'MemorySpace', 'trait' : 'Trait'}
+        for specifier, value in specifiers.items():
+            if specifier == "space" and skip_space:
+                continue
+            keywords_list.append(
+                ast.keyword(
+                    arg=specifier, 
+                    value=ast.Attribute(
+                        value=ast.Attribute(
+                            value=ast.Name(id=self.pk_import, ctx=ast.Load()), 
+                            attr=attr_names[specifier], ctx=ast.Load()), 
+                        attr=value, ctx=ast.Load()
+                        )
+                )
+            )
         return ast.keyword(
             arg=view_name, 
             value=ast.Call(
@@ -340,38 +344,10 @@ def get_keyword_node(self, view_name: str, layout: str) -> ast.keyword:
                     attr='ViewTypeInfo', ctx=ast.Load()
                 ), 
                 args=[], 
-                keywords=[
-                    ast.keyword(
-                        arg='layout', 
-                        value=ast.Attribute(
-                            value=ast.Attribute(
-                                value=ast.Name(id=self.pk_import, ctx=ast.Load()), 
-                                attr='Layout', ctx=ast.Load()), 
-                            attr= layout, ctx=ast.Load()
-                            )
-                    )
-                ]
+                keywords= keywords_list
             )
         )
 
-    @staticmethod
-    def filter_layout_change(node: ast.AST, working_dict: Dict[str, str]) -> Dict[str, str]:
-        '''
-        Filter out (from the working dict) the views whose layouts decorators user has already provided
-
-        :param node: AST of pykokkos entity to work with
-        :param working_dict: map of view -> layout
-        :returns updated working dict
-        '''
-
-        call_obj = node.decorator_list[0]
-        # iterating over view layout decorators in signature (of workunit)
-        for keyword_obj in call_obj.keywords:
-            if keyword_obj.arg in working_dict:
-                # user provided layout decorator for this view, remove from working dict
-                del working_dict[keyword_obj.arg]
-
-        return working_dict
 
     @staticmethod
     def is_classtype(node: ast.stmt, pk_import: str) -> bool:

pykokkos/core/runtime.py

@@ -11,8 +11,8 @@
 from pykokkos.core.visitors import visitors_util
 from pykokkos.interface import (
     DataType, ExecutionPolicy, ExecutionSpace, MemorySpace,
-    RandomPool, RangePolicy, TeamPolicy, View, ViewType, UpdatedTypes,
-    is_host_execution_space
+    RandomPool, RangePolicy, TeamPolicy, View, ViewType, UpdatedTypes, UpdatedDecorator,
+    is_host_execution_space, get_types_signature
 )
 import pykokkos.kokkos_manager as km
 
@@ -57,18 +57,25 @@ def precompile_workunit(
         self,
         workunit: Callable[..., None],
         space: ExecutionSpace,
-        updated_types: Optional[UpdatedTypes] = None
-    ) -> Optional[PyKokkosMembers]:
+        updated_decorator: UpdatedDecorator,
+        updated_types: Optional[UpdatedTypes] = None,
+        types_signature: Optional[str] = None,
+        ) -> Optional[PyKokkosMembers]:
         """
         precompile the workunit
 
         :param workunit: the workunit function object
+        :param updated_decorator: Object for decorator specifier
+        :param updated_types: Object with type inference information
         :param space: the ExecutionSpace for which the bindings are generated
         :returns: the members the functor is containing
         """
 
-        module_setup: ModuleSetup = self.get_module_setup(workunit, space, updated_types)
-        members: Optional[PyKokkosMembers] = self.compiler.compile_object(module_setup, space, km.is_uvm_enabled(), updated_types)
+        module_setup: ModuleSetup = self.get_module_setup(workunit, space, types_signature)
+        members: Optional[PyKokkosMembers] = self.compiler.compile_object(module_setup, 
+                                                                          space, km.is_uvm_enabled(), 
+                                                                          updated_decorator, 
+                                                                          updated_types, types_signature)
 
         return members
 
@@ -92,6 +99,7 @@ def run_workunit(
         name: Optional[str],
         policy: ExecutionPolicy,
         workunit: Callable[..., None],
+        updated_decorator: UpdatedDecorator,
         updated_types: Optional[UpdatedTypes] = None,
         operation: Optional[str] = None,
         initial_value: Union[float, int] = 0,
@@ -104,6 +112,7 @@ def run_workunit(
         :param policy: the execution policy of the operation
         :param workunit: the workunit function object
         :param kwargs: the keyword arguments passed to the workunit
+        :param updated_decorator: Object with decorator specifier information
         :param updated_types: UpdatedTypes object with type inferrence information
         :param operation: the name of the operation "for", "reduce", or "scan"
         :param initial_value: the initial value of the accumulator
@@ -117,11 +126,12 @@ def run_workunit(
                 raise RuntimeError("ERROR: operation cannot be None for Debug")
             return run_workunit_debug(policy, workunit, operation, initial_value, **kwargs)
 
-        members: Optional[PyKokkosMembers] = self.precompile_workunit(workunit, execution_space, updated_types)
+        types_signature: str = get_types_signature(updated_types, updated_decorator, execution_space)
+        members: Optional[PyKokkosMembers] = self.precompile_workunit(workunit, execution_space, updated_decorator, updated_types, types_signature)
         if members is None:
             raise RuntimeError("ERROR: members cannot be none")
 
-        module_setup: ModuleSetup = self.get_module_setup(workunit, execution_space, updated_types)
+        module_setup: ModuleSetup = self.get_module_setup(workunit, execution_space, types_signature)
         return self.execute(workunit, module_setup, members, execution_space, policy=policy, name=name, **kwargs)
 
     def is_debug(self, space: ExecutionSpace) -> bool:
@@ -429,19 +439,18 @@ def get_module_setup(
         self,
         entity: Union[object, Callable[..., None]],
         space: ExecutionSpace,
-        updated_types: Optional[UpdatedTypes] = None
+        types_signature: Optional[str] = None
         ) -> ModuleSetup:
         """
         Get the compiled module setup information unique to an entity + space
 
         :param entity: the workload or workunit object
         :param space: the execution space
-        :updated_types: Object with information about inferred types (if any)
+        :types_signature: Hash/identifer string for workunit module against data types
         :returns: the ModuleSetup object
         """
 
         space: ExecutionSpace = km.get_default_space() if space is ExecutionSpace.Debug else space
-        types_signature: str = None if updated_types is None else updated_types.types_signature
 
         module_setup_id = self.get_module_setup_id(entity, space, types_signature)
 
@@ -467,7 +476,7 @@ def get_module_setup_id(
 
         :param entity: the workload or workunit object
         :param space: the execution space
-        :param types_signature: optional identifier string for inferred types of parameters
+        :param types_signature: optional identifier/hash string for types of parameters against workunit module
         :returns: a unique tuple per entity and space
         """
 

pykokkos/interface/__init__.py

@@ -53,7 +53,7 @@
 )
 
 from .ext_module import compile_into_module
-from .args_type_inference import UpdatedTypes, get_annotations, get_type_str
+from .args_type_inference import UpdatedTypes, UpdatedDecorator, get_annotations, get_type_str, get_types_signature
 
 def fence():
     pass