Improve naming for generated helper functions

Author: jansel

helion/_compiler/device_function.py

@@ -504,7 +504,8 @@ def register_helper_function(
         self, helper_graph_info: HelperFunctionGraphInfo
     ) -> None:
         """Register a helper function to be generated at global scope."""
-        self.helper_manager.register_helper_function(helper_graph_info)
+        name = self.namespace.create_name(self.name, None)
+        self.helper_manager.register_helper_function(helper_graph_info, name)
 
     def codegen_helper_functions(self) -> list[ast.stmt]:
         """Generate helper function definitions at global scope."""

helion/_compiler/device_ir.py

@@ -937,9 +937,14 @@ class HelperFunctionGraphInfo(NodeArgsGraphInfo):
     """Graph info for helper functions in higher-order operations like associative_scan."""
 
     _param_names: list[str] = dataclasses.field(default_factory=list)
+    original_function_name: str | None = dataclasses.field(default=None)
 
     @property
     def name(self) -> str:
+        # This property should only be used during registration, not for final names
+        # Final names are generated in codegen using the namespace below
+        if self.original_function_name:
+            return f"{self.original_function_name}_{self.graph_id}"
         return f"helper_function_{self.graph_id}"
 
     def find_input_nodes(self) -> list[torch.fx.Node]:

helion/_compiler/helper_function.py

@@ -58,6 +58,11 @@ def extract_helper_function(helper_fn: object) -> types.FunctionType:
     return helper_fn.fn if isinstance(helper_fn, Kernel) else helper_fn  # pyright: ignore[reportReturnType]
 
 
+def extract_helper_function_name(helper_fn: object) -> str:
+    """Extract the function name from a Kernel object or regular function."""
+    return extract_helper_function(helper_fn).__name__
+
+
 CombineFunctionBasic = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
 CombineFunctionTuple = Callable[..., tuple[torch.Tensor, ...]]
 CombineFunction = CombineFunctionBasic | CombineFunctionTuple
@@ -153,17 +158,22 @@ class HelperFunctionManager:
 
     def __init__(self) -> None:
         self.helper_functions: dict[str, HelperFunctionGraphInfo] = {}
+        self._final_names: dict[str, str] = {}
 
     def register_helper_function(
-        self, helper_graph_info: HelperFunctionGraphInfo
+        self, helper_graph_info: HelperFunctionGraphInfo, final_name: str
     ) -> None:
         """Register a helper function to be generated at global scope."""
         self.helper_functions[helper_graph_info.name] = helper_graph_info
+        self._final_names[helper_graph_info.name] = final_name
 
     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():
+            # Get the final name that was already determined during registration
+            final_name = self._final_names[helper_graph_info.name]
+
             # Determine the number of parameters from the graph
             input_nodes = helper_graph_info.find_input_nodes()
 
@@ -184,7 +194,7 @@ def codegen_helper_functions(self) -> list[ast.stmt]:
             # Generate the function structure with @triton.jit decorator
             func_def = create(
                 ast.FunctionDef,
-                name=helper_graph_info.name,
+                name=final_name,
                 args=create_arguments(args),
                 body=func_body,
                 decorator_list=[expr_from_string("triton.jit")],
@@ -195,6 +205,10 @@ def codegen_helper_functions(self) -> list[ast.stmt]:
 
         return helper_defs
 
+    def get_final_name(self, helper_graph_info: HelperFunctionGraphInfo) -> str:
+        """Get the final generated name for a helper function."""
+        return self._final_names.get(helper_graph_info.name, helper_graph_info.name)
+
     def _codegen_helper_function_body(
         self, helper_graph_info: HelperFunctionGraphInfo
     ) -> list[ast.stmt]:

helion/language/reduce_ops.py

@@ -137,6 +137,7 @@ def _(
     from .._compiler.device_ir import args_to_proxies
     from .._compiler.device_ir import select_decomp_table
     from .._compiler.helper_function import create_combine_function_wrapper
+    from .._compiler.helper_function import extract_helper_function_name
 
     is_tuple_input = isinstance(input_tensor, (tuple, list))
     if is_tuple_input:
@@ -147,6 +148,8 @@ def _(
         assert isinstance(input_tensor, torch.Tensor), "reduce input must be a tensor"
 
     assert callable(combine_fn), "combine_fn must be callable"
+    # Extract the function name before wrapping
+    original_function_name = extract_helper_function_name(combine_fn)
     combine_fn = create_combine_function_wrapper(
         combine_fn, is_tuple_input=is_tuple_input, target_format="tuple"
     )
@@ -182,6 +185,7 @@ def _(
         combine_graph,
         HelperFunctionGraphInfo,
         node_args=[],
+        original_function_name=original_function_name,
     )
 
     # Validate other parameter for mask_node_inputs
@@ -334,12 +338,17 @@ def _(state: CodegenState) -> ast.AST | list[ast.AST]:
 
 
 def _register_helper_function(state: CodegenState, combine_graph_id: int) -> str:
-    """Register the helper function and return its name."""
+    """Register the helper function and return its final name."""
+    from .._compiler.device_ir import HelperFunctionGraphInfo
     from .._compiler.host_function import HostFunction
 
     helper_graph_info = HostFunction.current().device_ir.graphs[combine_graph_id]
-    state.codegen.device_function.register_helper_function(helper_graph_info)  # pyright: ignore[reportArgumentType]
-    return helper_graph_info.name
+    assert isinstance(helper_graph_info, HelperFunctionGraphInfo)
+    state.codegen.device_function.register_helper_function(helper_graph_info)
+    # Get the final name from the helper manager (which uses the namespace)
+    return state.codegen.device_function.helper_manager.get_final_name(
+        helper_graph_info
+    )
 
 
 def _create_reduce_expression(

helion/language/scan_ops.py

@@ -117,6 +117,7 @@ def _(
     from .._compiler.device_ir import args_to_proxies
     from .._compiler.device_ir import select_decomp_table
     from .._compiler.helper_function import create_combine_function_wrapper
+    from .._compiler.helper_function import extract_helper_function_name
 
     is_tuple_input = isinstance(input_tensor, (tuple, list))
     if is_tuple_input:
@@ -130,6 +131,8 @@ def _(
     assert isinstance(dim, int), "associative_scan dim must be an integer"
 
     assert callable(combine_fn), "combine_fn must be callable"
+    # Extract the function name before wrapping
+    original_function_name = extract_helper_function_name(combine_fn)
     combine_fn = create_combine_function_wrapper(
         combine_fn, is_tuple_input=is_tuple_input, target_format="unpacked"
     )
@@ -151,6 +154,7 @@ def _(
         combine_graph,
         HelperFunctionGraphInfo,
         node_args=[],
+        original_function_name=original_function_name,
     )
 
     # Create the associative_scan tracing operation
@@ -349,12 +353,17 @@ def _get_input_tensor_ast(state: CodegenState, is_tuple_input: bool) -> ast.AST:
 
 
 def _register_helper_function(state: CodegenState, combine_graph_id: int) -> str:
-    """Register the helper function and return its name."""
+    """Register the helper function and return its final name."""
+    from .._compiler.device_ir import HelperFunctionGraphInfo
     from .._compiler.host_function import HostFunction
 
     helper_graph_info = HostFunction.current().device_ir.graphs[combine_graph_id]
-    state.codegen.device_function.register_helper_function(helper_graph_info)  # pyright: ignore[reportArgumentType]
-    return helper_graph_info.name
+    assert isinstance(helper_graph_info, HelperFunctionGraphInfo)
+    state.codegen.device_function.register_helper_function(helper_graph_info)
+    # Get the final name from the helper manager (which uses the namespace)
+    return state.codegen.device_function.helper_manager.get_final_name(
+        helper_graph_info
+    )
 
 
 def _create_scan_expression(