Add hl.associative_scan

stack-info: PR: #239, branch: jansel/stack/78

Author: jansel

helion/_compiler/compile_environment.py

@@ -222,10 +222,15 @@ def to_fake(self, obj: object, origin: Origin) -> object:
             ),
         ):
             return obj
-        if isinstance(obj, types.FunctionType):
+        # Handle functions and Kernel objects
+        from ..runtime.kernel import Kernel
+
+        if isinstance(obj, (types.FunctionType, Kernel)):
+            from .helper_function import extract_helper_function
             from .lift_closures import lift_closures
 
-            return lift_closures(obj, origin)
+            fn = extract_helper_function(obj)
+            return lift_closures(fn, origin)
         if isinstance(obj, ConstExpr):
             return obj.value
         if isinstance(obj, list):

helion/_compiler/device_function.py

@@ -37,6 +37,7 @@
 
 if TYPE_CHECKING:
     from ..runtime.config import Config
+    from .device_ir import HelperFunctionGraphInfo
     from .generate_ast import GenerateAST
     from .program_id import ProgramIDs
 
@@ -185,6 +186,8 @@ def __init__(self, name: str, config: Config, codegen: GenerateAST) -> None:
         self.block_size_var_cache: dict[tuple[int, ...], str] = {}
         self.expr_to_var_info: dict[sympy.Expr, VarInfo] = {}
 
+        self.helper_functions: dict[str, HelperFunctionGraphInfo] = {}
+
         from .indexing_strategy import IndexingStrategy
         from .tile_dispatch import TileStrategyDispatch
 
@@ -488,6 +491,90 @@ def dead_code_elimination(self) -> None:
                     if v.name in args_to_remove:
                         del cache[k]
 
+    def register_helper_function(
+        self, helper_graph_info: HelperFunctionGraphInfo
+    ) -> None:
+        """Register a helper function to be generated at global scope."""
+        self.helper_functions[helper_graph_info.name] = helper_graph_info
+
+    def codegen_helper_functions(self) -> list[ast.stmt]:
+        """Generate helper function definitions at global scope."""
+        helper_defs = []
+        for helper_graph_info in self.helper_functions.values():
+            # Determine the number of parameters from the graph
+            input_nodes = helper_graph_info.find_input_nodes()
+
+            # Generate argument list with consistent names
+            args = []
+            param_names = []
+            for i in range(len(input_nodes)):
+                arg_name = f"param_{i}"
+                args.append(create_arg(arg_name))
+                param_names.append(arg_name)
+
+            # Store parameter names for use in body generation
+            helper_graph_info._param_names = param_names
+
+            # Process the FX graph to generate the correct helper function body
+            func_body = self._codegen_helper_function_body(helper_graph_info)
+
+            # Generate the function structure with @triton.jit decorator
+            func_def = create(
+                ast.FunctionDef,
+                name=helper_graph_info.name,
+                args=create_arguments(args),
+                body=func_body,
+                decorator_list=[expr_from_string("triton.jit")],
+                type_params=[],
+            )
+
+            helper_defs.append(func_def)
+
+        return helper_defs
+
+    def _codegen_helper_function_body(
+        self, helper_graph_info: HelperFunctionGraphInfo
+    ) -> list[ast.stmt]:
+        """Generate the body of a helper function by processing its FX graph."""
+        from .helper_function import HelperCodegen
+        from .inductor_lowering import GraphInterpreter
+
+        # Create a temporary DeviceFunction for generating the helper function
+        temp_device_function = DeviceFunction(
+            name=f"temp_{helper_graph_info.name}",
+            config=self.config,
+            codegen=self.codegen,
+        )
+
+        # Use the parameter names from the function definition
+        param_names = helper_graph_info._param_names
+
+        # Create parameter AST nodes for the helper function
+        param_args = []
+        for param_name in param_names:
+            param_args.append(expr_from_string(param_name))
+
+        # Process the graph using the existing interpreter infrastructure
+        with temp_device_function:
+            helper_codegen = HelperCodegen(temp_device_function)
+            interpreter = GraphInterpreter(helper_graph_info.graph, helper_codegen)
+            results = interpreter.run(*param_args)
+
+            # Get the generated statements from the temporary device function
+            statements = temp_device_function.body.copy()
+
+            # Ensure there's a return statement
+            if not statements or not isinstance(statements[-1], ast.Return):
+                if isinstance(results, ast.AST):
+                    statements.append(create(ast.Return, value=results))
+                else:
+                    # This should not happen in normal operation
+                    raise RuntimeError(
+                        f"Helper function {helper_graph_info.name} did not produce a valid result"
+                    )
+
+        return cast("list[ast.stmt]", statements)
+
     def __enter__(self) -> None:
         try:
             tls.functions.append(self)

helion/_compiler/device_ir.py

@@ -2,7 +2,6 @@
 
 import ast
 import builtins
-from collections.abc import Callable
 import contextlib
 import dataclasses
 import functools
@@ -807,9 +806,94 @@ def visit_Call(self, node: ast.Call) -> object:
         else:
             func = self.visit(node.func)
 
+        # Special handling for associative_scan
+        import helion.language as hl
+
+        if isinstance(
+            (func_type_info := node.func._type_info),
+            CallableType,
+        ) and (
+            func_type_info.value is hl.associative_scan or func is hl.associative_scan
+        ):
+            return self._handle_associative_scan(node, args, kwargs)
+
         # pyre-ignore[6]
         return _CheckForIndexCalls.retry_call(func, args, kwargs)
 
+    def _handle_associative_scan(
+        self, node: ast.Call, args: list[object], kwargs: dict[str, object]
+    ) -> object:
+        """Handle associative_scan calls by tracing the combine function as a subgraph."""
+        from ..language import _tracing_ops
+
+        combine_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = cast(
+            "Callable[[torch.Tensor, torch.Tensor], torch.Tensor]", args[0]
+        )  # The combine function
+        input_tensor = args[1]  # The input tensor
+
+        # Extract other arguments from kwargs
+        dim = kwargs.get("dim", 0)
+        reverse = kwargs.get("reverse", False)
+
+        # Create a subgraph for the combine function
+        def run_combine_subgraph(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+            # This will trace the combine function
+            # Check if combine_fn is a Kernel object and extract the underlying function
+            from .helper_function import extract_helper_function
+
+            actual_fn = extract_helper_function(combine_fn)
+            return actual_fn(a, b)
+
+        # Create fake inputs for the combine function based on input tensor element type
+        if isinstance(input_tensor, torch.Tensor):
+            fake_a = torch.empty(
+                [1], dtype=input_tensor.dtype, device=input_tensor.device
+            )
+            fake_b = torch.empty(
+                [1], dtype=input_tensor.dtype, device=input_tensor.device
+            )
+        else:
+            # Fallback for when input_tensor is a proxy
+            fake_a = torch.empty([1], dtype=torch.float32)
+            fake_b = torch.empty([1], dtype=torch.float32)
+
+        with self.disable_tracing() as tracer:
+            combine_graph = proxy_tensor.make_fx(
+                run_combine_subgraph, decomposition_table=select_decomp_table()
+            )(fake_a, fake_b).graph
+
+            combine_graph_id = self.device_ir.add_graph(
+                combine_graph,
+                HelperFunctionGraphInfo,
+                node_args=[],  # The combine function doesn't use external args
+            )
+
+            # Create the associative_scan tracing operation
+            scan_args = (
+                combine_graph_id,
+                input_tensor,
+                dim,
+                reverse,
+            )
+
+            proxy_args, proxy_kwargs = args_to_proxies(tracer, scan_args)
+            proxy_out = tracer.create_proxy(
+                "call_function",
+                _tracing_ops._associative_scan,
+                proxy_args,
+                proxy_kwargs,
+            )
+
+            # The output has the same shape as the input
+            proxy_tensor.track_tensor_tree(
+                input_tensor,
+                proxy_out,
+                constant=None,
+                tracer=tracer,
+            )
+
+        return proxy_out
+
     def visit_Attribute(self, node: ast.Attribute) -> object:
         return getattr(self.visit(node.value), node.attr)
 
@@ -898,6 +982,37 @@ def lower_to_device_ir(func: HostFunction) -> DeviceIR:
         return device_ir
 
 
+@dataclasses.dataclass
+class HelperFunctionGraphInfo(NodeArgsGraphInfo):
+    """Graph info for helper functions in higher-order operations like associative_scan."""
+
+    _param_names: list[str] = dataclasses.field(default_factory=list)
+
+    @property
+    def name(self) -> str:
+        return f"helper_function_{self.graph_id}"
+
+    def find_input_nodes(self) -> list[torch.fx.Node]:
+        """Find all placeholder nodes (inputs) in the graph."""
+        return self.graph.find_nodes(op="placeholder")
+
+    def codegen(self, state: CodegenState) -> list[object]:
+        # For helper functions, we need to inline the function body
+        # The helper function takes variable arguments and returns their combination
+
+        # Generate temporary variable names for the helper function arguments
+        # Use the graph's input nodes to determine the number of parameters
+        input_nodes = self.find_input_nodes()
+        args: list[ast.AST] = []
+
+        for i in range(len(input_nodes)):
+            var_name = state.codegen.tmpvar(prefix=f"helper_arg_{i}")
+            args.append(create(ast.Name, id=var_name, ctx=ast.Load()))
+
+        # Generate the helper function call
+        return codegen_call_with_graph(state.codegen, self.graph, args)
+
+
 def remove_unnecessary_tile_index(graph: torch.fx.Graph) -> None:
     """
     Remove unnecessary tile_index nodes from the graph.

helion/_compiler/generate_ast.py

@@ -17,6 +17,7 @@
 from .ast_extension import statement_from_string
 from .compile_environment import CompileEnvironment
 from .device_function import DeviceFunction
+from .helper_function import CodegenInterface
 from .inductor_lowering import CodegenState
 from .inductor_lowering import codegen_call_with_graph
 from .program_id import ForEachProgramID
@@ -32,19 +33,25 @@
     from .type_propagation import TensorType
 
 
-class GenerateAST(NodeVisitor):
+class GenerateAST(NodeVisitor, CodegenInterface):
     def __init__(self, func: HostFunction, config: Config) -> None:
-        super().__init__()
+        # Initialize NodeVisitor first
+        NodeVisitor.__init__(self)
+
+        # Initialize our attributes
         self.host_function = func
         self.host_statements: list[ast.AST] = []
         self.statements_stack: list[list[ast.AST]] = [self.host_statements]
         self.on_device = False
-        self.device_function = DeviceFunction(f"_{func.name}_kernel", config, self)
         self.active_device_loops: dict[int, list[DeviceLoopOrGridState]] = (
             collections.defaultdict(list)
         )
         self.next_else_block: list[ast.AST] | None = None
 
+        # Now create device function and initialize CodegenInterface
+        self.device_function = DeviceFunction(f"_{func.name}_kernel", config, self)
+        CodegenInterface.__init__(self, self.device_function)
+
     def offset_var(self, block_idx: int) -> str:
         return self.active_device_loops[block_idx][-1].strategy.offset_var(block_idx)
 
@@ -63,9 +70,6 @@ def add_statement(self, stmt: ast.AST | str | None) -> None:
             stmt = statement_from_string(stmt)
         self.statements_stack[-1].append(stmt)
 
-    def tmpvar(self, *, dce: bool = False, prefix: str = "v") -> str:
-        return self.device_function.unique_name(prefix, dce=dce)
-
     def lift(self, expr: ast.AST, *, dce: bool = False, prefix: str = "v") -> ast.Name:
         if isinstance(expr, ast.Name):
             return expr
@@ -413,6 +417,7 @@ def generate_ast(func: HostFunction, config: Config) -> ast.AST:
             result = ast.Module(
                 [
                     *func.codegen_imports(),
+                    *codegen.device_function.codegen_helper_functions(),
                     *kernel_def,
                     host_def,
                     precompile_def,

helion/_compiler/helper_function.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+import ast
+from typing import TYPE_CHECKING
+
+from .ast_extension import create
+from .ast_extension import statement_from_string
+
+if TYPE_CHECKING:
+    import types
+
+    from .device_function import DeviceFunction
+
+
+class CodegenInterface(ABC):
+    """Abstract base class for codegen interfaces used by GraphInterpreter."""
+
+    def __init__(self, device_function: DeviceFunction) -> None:
+        self.device_function = device_function
+
+    @abstractmethod
+    def add_statement(self, stmt: ast.AST | str | None) -> None:
+        """Add a statement to the generated code."""
+
+    def tmpvar(self, *, dce: bool = False, prefix: str = "v") -> str:
+        """Generate a temporary variable name."""
+        return self.device_function.unique_name(prefix, dce=dce)
+
+    def lift(self, expr: ast.AST, *, dce: bool = False, prefix: str = "v") -> ast.Name:
+        """Lift an expression to a temporary variable if needed."""
+        if isinstance(expr, ast.Name):
+            return expr
+        varname = self.tmpvar(dce=dce, prefix=prefix)
+        self.add_statement(statement_from_string(f"{varname} = expr", expr=expr))
+        return create(ast.Name, id=varname, ctx=ast.Load())
+
+
+def extract_helper_function(helper_fn: object) -> types.FunctionType:
+    """Extract the actual function from a Kernel object or return as-is.
+
+    This utility function centralizes the logic for handling both regular functions
+    and Kernel objects that wrap functions.
+    """
+    from ..runtime.kernel import Kernel
+
+    # pyre-ignore[16]: We check isinstance before accessing .fn
+    return helper_fn.fn if isinstance(helper_fn, Kernel) else helper_fn
+
+
+class HelperCodegen(CodegenInterface):
+    """Codegen wrapper for helper function generation."""
+
+    def __init__(self, device_function: DeviceFunction) -> None:
+        super().__init__(device_function)
+
+    def add_statement(self, stmt: ast.AST | str | None) -> None:
+        if stmt is not None:
+            if isinstance(stmt, str):
+                stmt = statement_from_string(stmt)
+            self.device_function.body.append(stmt)