Implement loop fusion optimization inside kernels

mypy.ini

@@ -21,12 +21,24 @@ ignore_errors = True
 [mypy-pykokkos.interface.views]
 ignore_errors = True
 
+[mypy-pykokkos.core.optimizations.loop_fuse]
+ignore_errors = True
+
+[mypy-pykokkos.core.optimizations.memory_ops_fuse]
+ignore_errors = True
+
+[mypy-pykokkos.core.optimizations.util]
+ignore_errors = True
+
 [mypy-pykokkos.core.fusion.fuse]
 ignore_errors = True
 
 [mypy-pykokkos.core.fusion.trace]
 ignore_errors = True
 
+[mypy-pykokkos.core.fusion.util]
+ignore_errors = True
+
 [mypy-pykokkos.core.fusion.access_modes]
 ignore_errors = True
 

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

@@ -20,6 +20,10 @@ def __sub__(self, other):
         self.flush_trace()
         return self.value - other
 
+    def __rsub__(self, other):
+        self.flush_trace()
+        return other - self.value
+
     def __mul__(self, other):
         self.flush_trace()
         return self.value * other
@@ -28,6 +32,10 @@ def __truediv__(self, other):
         self.flush_trace()
         return self.value / other
 
+    def __rtruediv__(self, other):
+        self.flush_trace()
+        return other / self.value
+
     def __floordiv__(self, other):
         self.flush_trace()
         return self.value // other
@@ -36,8 +44,23 @@ def __str__(self):
         self.flush_trace()
         return str(self.value)
 
+    def __eq__(self, other):
+        self.flush_trace()
+        if isinstance(other, Future):
+            return self.value == other.value
+
+        return self.value == other
+
+    def __lt__(self, other):
+        self.flush_trace()
+        return self.value < other
+
+    def __gt__(self, other):
+        self.flush_trace()
+        return self.value > other
+
     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/trace.py

@@ -1,4 +1,5 @@
 import ast
+import hashlib
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
@@ -129,6 +130,10 @@ def get_operations(self, data: Union[Future, ViewType]) -> List[TracerOperation]
         version: int = self.data_version.get(id(data), 0)
         dependency = DataDependency(None, id(data), version)
 
+        if dependency not in self.data_operation:
+            # The data does not depend on any prior operation
+            return []
+
         operation: TracerOperation = self.data_operation[dependency]
         if operation not in self.operations:
             # This means that the dependency was already updated
@@ -247,7 +252,12 @@ def fuse_operations(self, operations: List[TracerOperation]) -> TracerOperation:
             args[f"args_{index}"] = op.args
             dependencies.update(op.dependencies)
 
-        fused_name: str = "_".join(names)
+        fused_name: str
+        if len(names) < 5:
+            fused_name = "_".join(names)
+        else:
+            # Avoid long names
+            fused_name = "_".join(names[:5]) + hashlib.md5(("".join(names)).encode()).hexdigest()
 
         return TracerOperation(None, None, fused_name, policy, workunits, operation, parser, fused_name, args, dependencies)
 

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/module_setup.py

@@ -163,10 +163,18 @@ def get_entity_dir(self, main: Path, metadata: List[EntityMetadata]) -> Path:
 
         entity_dir: str = ""
 
-        for m in metadata:
+        for m in metadata[:5]:
             filename: str = m.path.split("/")[-1].split(".")[0]
             entity_dir += f"{filename}_{m.name}"
 
+        remaining: str = ""
+        for m in metadata[5:]:
+            filename: str = m.path.split("/")[-1].split(".")[0]
+            remaining += f"{filename}_{m.name}"
+
+        if remaining != "":
+            entity_dir += hashlib.md5(("".join(remaining)).encode()).hexdigest()
+
         return self.get_main_dir(main) / Path(entity_dir)
 
     @staticmethod

pykokkos/core/optimizations/__init__.py

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