Implement loop fusion optimization inside kernels

pykokkos/core/compiler.py

@@ -10,6 +10,7 @@
 from typing import Dict, List, Optional, Tuple
 
 from pykokkos.core.fusion import fuse_workunits
+from pykokkos.core.optimizations import loop_fuse, memory_ops_fuse
 from pykokkos.core.parsers import Parser, PyKokkosEntity, PyKokkosStyles
 from pykokkos.core.translators import PyKokkosMembers, StaticTranslator
 from pykokkos.core.type_inference import UpdatedTypes, UpdatedDecorator
@@ -223,6 +224,11 @@ def compile_entity(
         bindings: List[str]
         cast: List[str]
 
+        if entity.style in {PyKokkosStyles.workunit, PyKokkosStyles.fused}:
+            if "PK_LOOP_FUSE" in os.environ:
+                loop_fuse(entity.AST)
+            if "PK_MEM_FUSE" in os.environ:
+                memory_ops_fuse(entity.AST, entity.pk_import)
         functor, bindings, cast = translator.translate(entity, classtypes)
 
         t_end: float = time.perf_counter() - t_start

pykokkos/core/fusion/fuse.py

@@ -2,52 +2,7 @@
 import os
 from typing import Any, Dict, List, Set, Tuple, Union
 
-def get_node_name(node: Union[ast.Attribute, ast.Name]) -> str:
-    """
-    Copied from visitors_util.py due to circular import
-    """
-
-    name: str
-    if isinstance(node, ast.Attribute):
-        name = node.attr
-    else:
-        name = node.id
-
-    return name
-
-
-class DeclarationsVisitor(ast.NodeVisitor):
-    """
-    Get all variable declarations
-    """
-
-    def __init__(self) -> None:
-        self.declarations: Set[str] = set()
-
-    def visit_AnnAssign(self, node: ast.AnnAssign) -> Any:
-        self.declarations.add(get_node_name(node.target))
-
-
-class VariableRenamer(ast.NodeTransformer):
-    """
-    Renames variables in a fused ast according to a map
-    """
-
-    def __init__(self, name_map: Dict[Tuple[str, int], str], workunit_idx: int):
-        self.name_map = name_map
-        self.workunit_idx = workunit_idx
-
-    def visit_Name(self, node: ast.Name) -> Any:
-        key = (node.id, self.workunit_idx)
-        # If the name is not mapped, keep the original name
-        node.id = self.name_map.get(key, node.id)
-        return node
-
-    def visit_keyword(self, node: ast.keyword) -> Any:
-        key = (node.id, self.workunit_idx)
-        # If the name is not mapped, keep the original name
-        node.arg = self.name_map.get(key, node.arg)
-        return node
+from .util import DeclarationsVisitor, VariableRenamer
 
 
 def fuse_workunit_kwargs_and_params(

pykokkos/core/fusion/future.py

@@ -37,7 +37,7 @@ def __str__(self):
         return str(self.value)
 
     def __repr__(self) -> str:
-        return str(f"Future(value={self.value})")
+        return f"Future(value={self.value})"
 
     def flush_trace(self) -> None:
         runtime_singleton.runtime.flush_data(self)

pykokkos/core/fusion/util.py

@@ -0,0 +1,54 @@
+import ast
+from typing import Dict, Set, Tuple, Union
+
+
+def get_node_name(node: Union[ast.Attribute, ast.Name]) -> str:
+    """
+    Copied from visitors_util.py due to circular import
+    """
+
+    name: str
+    if isinstance(node, ast.Attribute):
+        name = node.attr
+    else:
+        name = node.id
+
+    return name
+
+
+class DeclarationsVisitor(ast.NodeVisitor):
+    """
+    Get all variable declarations
+    """
+
+    def __init__(self) -> None:
+        self.declarations: Set[str] = set()
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        self.declarations.add(get_node_name(node.target))
+
+    def visit_For(self, node: ast.For) -> None:
+        self.declarations.add(get_node_name(node.target))
+        for n in node.body:
+            self.visit(n)
+
+class VariableRenamer(ast.NodeTransformer):
+    """
+    Renames variables in a fused ast according to a map
+    """
+
+    def __init__(self, name_map: Dict[Tuple[str, int], str], workunit_idx: int):
+        self.name_map = name_map
+        self.workunit_idx = workunit_idx
+
+    def visit_Name(self, node: ast.Name) -> None:
+        key = (node.id, self.workunit_idx)
+        # If the name is not mapped, keep the original name
+        node.id = self.name_map.get(key, node.id)
+        return node
+
+    def visit_keyword(self, node: ast.keyword) -> None:
+        key = (node.id, self.workunit_idx)
+        # If the name is not mapped, keep the original name
+        node.arg = self.name_map.get(key, node.arg)
+        return node

pykokkos/core/optimizations/__init__.py

@@ -0,0 +1,2 @@
+from .loop_fuse import loop_fuse
+from .memory_ops_fuse import memory_ops_fuse