Add support for multiple top level for loops

ghstack-source-id: 1f33fe9
Pull Request resolved: #52

Author: oulgen

helion/_compiler/ast_extension.py

@@ -49,13 +49,15 @@ def __init__(
         _type_info: TypeInfo | None = None,
         _loop_type: LoopType = LoopType.UNSET,
         _is_kernel_call: bool = False,
+        _root_id: int | None = None,
         **kwargs: object,
     ) -> None:
         super().__init__(**kwargs)
         self._type_info: TypeInfo | None = _type_info
         self._location: SourceLocation = _location
         self._loop_type: LoopType = _loop_type
         self._is_kernel_call: bool = _is_kernel_call
+        self._root_id: int | None = _root_id
 
     def new(self, fields: dict[str, object]) -> ExtendedAST:
         result = self.__class__(
@@ -64,6 +66,7 @@ def new(self, fields: dict[str, object]) -> ExtendedAST:
             _type_info=self._type_info,
             _loop_type=self._loop_type,
             _is_kernel_call=self._is_kernel_call,
+            _root_id=self._root_id,
         )
         return self._location.to_ast(result)
 

helion/_compiler/ast_read_writes.py

@@ -10,17 +10,34 @@
     _A = TypeVar("_A", bound=ast.AST)
 
 
+# TODO(oulgen): This visitor is extremely primitive, does not consider alpha renaming or scopes
 class _ReadWriteVisitor(ast.NodeVisitor):
     def __init__(self) -> None:
         super().__init__()
         self.rw = ReadWrites(collections.Counter(), collections.Counter())
 
+    def _update(self, name: str, ctx: ast.expr_context) -> None:
+        if isinstance(ctx, ast.Load):
+            self.rw.reads[name] += 1
+        elif isinstance(ctx, ast.Store):
+            self.rw.writes[name] += 1
+
     def visit_Name(self, node: ast.Name) -> None:
-        if isinstance(node.ctx, ast.Load):
-            self.rw.reads[node.id] += 1
-        elif isinstance(node.ctx, ast.Store):
-            self.rw.writes[node.id] += 1
-        self.generic_visit(node)
+        self._update(node.id, node.ctx)
+
+    def visit_Subscript(self, node: ast.Subscript) -> None:
+        if isinstance(node.value, ast.Name):
+            self._update(node.value.id, node.ctx)
+        else:
+            self.visit(node.value)
+
+    def visit_For(self, node: ast.For) -> None:
+        # Skip target
+        self.visit(node.iter)
+        for stmt in node.body:
+            self.visit(stmt)
+        for stmt in node.orelse:
+            self.visit(stmt)
 
 
 class ReadWrites(typing.NamedTuple):

helion/_compiler/compile_environment.py

@@ -20,6 +20,7 @@
 from .. import exc
 from ..language.constexpr import ConstExpr
 from .error_reporting import ErrorReporting
+from .loop_dependency_checker import LoopDependencyChecker
 from .variable_origin import BlockSizeOrigin
 from .variable_origin import Origin
 
@@ -70,6 +71,7 @@ def __init__(self, device: torch.device, settings: Settings) -> None:
             collections.Counter()
         )
         self.specialized_vars: set[sympy.Symbol] = set()
+        self.loop_dependency_checker = LoopDependencyChecker()
 
     def add_kernel_tensor_size(self, sizes: Sequence[int | torch.SymInt]) -> None:
         from .tile_strategy import TileStrategy
@@ -280,6 +282,7 @@ def __enter__(self) -> Self:
         self.fake_mode.__enter__()
         tls.env = self
         self.errors = ErrorReporting(self.settings)  # clear prior errors
+        self.loop_dependency_checker = LoopDependencyChecker()
         return self
 
     def __exit__(

helion/_compiler/device_function.py

@@ -36,6 +36,8 @@
 
 if TYPE_CHECKING:
     from ..runtime.config import Config
+    from .program_id import ProgramIDs
+    from .program_id import SharedProgramID
 
     _P = TypeVar("_P", bound="TensorPropertyArg")
 
@@ -145,7 +147,7 @@ def __init__(self, name: str, config: Config) -> None:
         self._unique_counter: dict[str, itertools.count[int]] = defaultdict(
             itertools.count
         )
-        self.grid_expr: ast.AST | None = None
+        self.pid: SharedProgramID | ProgramIDs | None = None
         self.namespace: _Namespace = _Namespace()
         self.namespace._used_names.update(reserved_names())
         self._variable_renames: dict[str, list[str]] = {}
@@ -169,9 +171,9 @@ def merge_variable_names(self, a: str, b: str) -> None:
         for n in name_group:
             self._variable_renames[n] = name_group
 
-    def set_grid_expr(self, grid_expr: ast.AST) -> None:
-        assert self.grid_expr is None, "grid_expr already set"
-        self.grid_expr = grid_expr
+    def set_pid(self, pid: SharedProgramID | ProgramIDs) -> None:
+        assert self.pid is None, "pid already set"
+        self.pid = pid
 
     def sympy_expr(self, expr: sympy.Expr) -> str:
         expr_to_origin = HostFunction.current().expr_to_origin
@@ -341,12 +343,12 @@ def codegen_function_call(self) -> ast.AST:
                 f"num_stages={self.config.num_stages}",
             ]
         )
-        grid_expr = self.grid_expr
-        assert grid_expr is not None
+        pid = self.pid
+        assert pid is not None
         # TODO(jansel): we should run CSE this statement
         call_statement = statement_from_string(
             f"{self.name}[__call_grid_expr]({', '.join(args)})",
-            __call_grid_expr=grid_expr,
+            __call_grid_expr=pid.codegen_grid(),
         )
         assert isinstance(call_statement, ExtendedAST)
         # Mark the kernel call we can find it in codegen_precompile_def

helion/_compiler/device_ir.py

@@ -249,22 +249,22 @@ class DeviceIR:
     def __init__(self) -> None:
         super().__init__()
         self.graphs: list[GraphInfo] = []
-        self.root_id: int | None = None
+        self.root_ids: list[int] = []
         self.rolled_reductions: list[RolledReductionInfo] = []
         self.grid_block_ids: list[list[int]] = []
 
-    def get_root(self, config: Config) -> torch.fx.Graph:
+    def get_root(self, config: Config, graph_id: int) -> torch.fx.Graph:
         """ " If we are using a rolled reduction, return the rolled reduction graph otherwise
         return the root graph."""
-        if (root_id := self.root_id) is None:
-            raise AssertionError("No root graph")
+        if graph_id >= len(self.graphs):
+            raise AssertionError("Invalid graph id")
         reduction_loops = config.reduction_loops
         if len(reduction_loops) > 1:
             raise NotImplementedError("Multiple reduction loops not implemented")
         if len(reduction_loops) == 0 or reduction_loops[0] is None:
-            return self.graphs[root_id].graph
+            return self.graphs[graph_id].graph
         for info in reversed(self.rolled_reductions):
-            if info.original_graph_id == root_id:
+            if info.original_graph_id == graph_id:
                 assert info.new_graph_id is not None
                 return self.graphs[info.new_graph_id].graph
         raise AssertionError("No rolled reduction graph found")
@@ -301,8 +301,7 @@ def add_reduction_loop_graph(
         )
 
     def add_root_graph(self, graph: torch.fx.Graph) -> None:
-        assert self.root_id is None
-        self.root_id = self.add_graph(graph, graph_info_cls=RootGraphInfo)
+        self.root_ids.append(self.add_graph(graph, graph_info_cls=RootGraphInfo))
 
     def build_rolled_reductions(self) -> None:
         env = CompileEnvironment.current()

helion/_compiler/generate_ast.py

@@ -19,6 +19,7 @@
 from .device_function import DeviceFunction
 from .inductor_lowering import CodegenState
 from .inductor_lowering import codegen_call_with_graph
+from .program_id import SharedProgramID
 from .variable_origin import ArgumentOrigin
 
 if TYPE_CHECKING:
@@ -42,6 +43,7 @@ def __init__(self, func: HostFunction, config: Config) -> None:
         self.active_device_loops: dict[int, list[DeviceLoopOrGridState]] = (
             collections.defaultdict(list)
         )
+        self.next_else_block: list[ast.AST] | None = None
 
     def offset_var(self, block_idx: int) -> str:
         return self.active_device_loops[block_idx][-1].strategy.offset_var(block_idx)
@@ -54,7 +56,9 @@ def mask_var(self, block_idx: int) -> str | None:
             return loops[-1].strategy.mask_var(block_idx)
         return None
 
-    def add_statement(self, stmt: ast.AST | str) -> None:
+    def add_statement(self, stmt: ast.AST | str | None) -> None:
+        if stmt is None:
+            return
         if isinstance(stmt, str):
             stmt = statement_from_string(stmt)
         self.statements_stack[-1].append(stmt)
@@ -134,13 +138,36 @@ def generic_visit(self, node: ast.AST) -> ast.AST:
                 fields[field] = old_value
         return node.new(fields)
 
-    def visit_For(self, node: ast.For) -> ast.AST:
+    def visit_For(self, node: ast.For) -> ast.AST | None:
         assert isinstance(node, ExtendedAST)
         if node._loop_type == LoopType.GRID:
             assert not node.orelse
+
+            if len(self.host_function.device_ir.root_ids) == 1:
+                body = self.device_function.body
+            else:
+                assert len(self.host_function.device_ir.root_ids) > 1
+                assert node._root_id is not None
+                # Multiple top level for loops
+
+                if node._root_id == 0:
+                    self.device_function.set_pid(
+                        SharedProgramID(
+                            self.device_function.new_var("pid_shared", dce=False)
+                        )
+                    )
+                    self.device_function.body.append(
+                        self.device_function.pid.codegen_pid_init()
+                    )
+                if node._root_id < len(self.host_function.device_ir.root_ids) - 1:
+                    body = []
+                else:
+                    # This is the last top level for, dont emit more if statements
+                    assert self.next_else_block is not None
+                    body = self.next_else_block
             with (
                 self.set_on_device(),
-                self.set_statements(self.device_function.body),
+                self.set_statements(body),
             ):
                 iter_node = node.iter
                 assert isinstance(iter_node, ExtendedAST)
@@ -166,21 +193,44 @@ def visit_For(self, node: ast.For) -> ast.AST:
 
                     from .inductor_lowering import CodegenState
 
-                    fn._codegen(
-                        CodegenState(
-                            self,
-                            fx_node=None,
-                            proxy_args=[*bound.arguments.values()],
-                            ast_args=None,
-                        ),
+                    state = CodegenState(
+                        self,
+                        fx_node=None,
+                        proxy_args=[*bound.arguments.values()],
+                        ast_args=None,
                     )
+
+                    fn._codegen(state)
+                assert node._root_id is not None
                 codegen_call_with_graph(
                     self,
-                    self.host_function.device_ir.get_root(self.device_function.config),
+                    self.host_function.device_ir.get_root(
+                        self.device_function.config,
+                        self.host_function.device_ir.root_ids[node._root_id],
+                    ),
                     [],
                 )
+                # If we are in a multi top level loop, for all loops except for the last one
+                # emit ifthenelse blocks
+                if node._root_id < len(self.host_function.device_ir.root_ids) - 1:
+                    block = (
+                        self.device_function.body
+                        if self.next_else_block is None
+                        else self.next_else_block
+                    )
+                    self.next_else_block = []
+                    block.append(
+                        create(
+                            ast.If,
+                            test=self.device_function.pid.codegen_test(state),
+                            body=body,
+                            orelse=self.next_else_block,
+                        )
+                    )
             self.device_function.dead_code_elimination()
-            return self.device_function.codegen_function_call()
+            if node._root_id == len(self.host_function.device_ir.root_ids) - 1:
+                return self.device_function.codegen_function_call()
+            return None
         return self.generic_visit(node)
 
     def visit_Name(self, node: ast.Name) -> ast.AST:

helion/_compiler/loop_dependency_checker.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import itertools
+from typing import TYPE_CHECKING
+
+from .. import exc
+from .ast_read_writes import ReadWrites
+
+if TYPE_CHECKING:
+    import ast
+
+
+class LoopDependencyChecker:
+    """
+    A class to check dependencies between top-level for loops in a Helion kernel.
+
+    This class tracks memory accesses (reads and writes) for each top-level for loop
+    and raises an error if a later loop reads or writes to anything written in a
+    previous loop.
+    """
+
+    def __init__(self) -> None:
+        self.reads: set[str] = set()
+        self.writes: set[str] = set()
+
+    def register_loop(self, loop_node: ast.For) -> None:
+        rw = ReadWrites.from_list(loop_node.body)
+
+        self._check_dependencies(rw)
+
+        self.reads |= set(rw.reads)
+        self.writes |= set(rw.writes)
+
+    def _check_dependencies(self, rw: ReadWrites) -> None:
+        """
+        Check for dependencies between the current loop and previous loops.
+
+        Raises:
+            exc.LoopDependencyError: If a dependency is detected
+        """
+        for name in sorted(itertools.chain(rw.reads, rw.writes)):
+            if name in self.writes:
+                raise exc.LoopDependencyError(name)