[RFC][Codemod] Rename device/host _fn to _function (#62)

There are so places where we use device_fn, host_fn, device_function and host_function. Lets be consistent everywhere. I'm also ok if we want to swap to _fn versions, as long as we are consistent.
Pull Request resolved: #62
Approved by: https://github.com/jansel

Author: oulgen

helion/_compiler/device_function.py

@@ -245,13 +245,13 @@ def tensor_arg(
     def tensor_descriptor_arg(
         self, fake_value: torch.Tensor, block_size: list[int | torch.SymInt]
     ) -> TensorArg:
-        host_fn = HostFunction.current()
+        host_function = HostFunction.current()
         block_size_expr = ", ".join(
             map(HostFunction.current().literal_expr, block_size)
         )
         key = (fake_value, block_size_expr)
         if key not in self._tensor_descriptor_args:
-            origin = host_fn.tensor_to_origin[fake_value]
+            origin = host_function.tensor_to_origin[fake_value]
             arg = TensorDescriptorArg(
                 self.new_var(origin.suggest_var_name() + "_desc"),
                 fake_value,

helion/_compiler/generate_ast.py

@@ -33,7 +33,7 @@
 class GenerateAST(NodeVisitor):
     def __init__(self, func: HostFunction, config: Config) -> None:
         super().__init__()
-        self.host_fn = func
+        self.host_function = func
         self.host_statements: list[ast.AST] = []
         self.statements_stack: list[list[ast.AST]] = [self.host_statements]
         self.on_device = False
@@ -175,7 +175,7 @@ def visit_For(self, node: ast.For) -> ast.AST:
                     )
                 codegen_call_with_graph(
                     self,
-                    self.host_fn.device_ir.get_root(self.device_function.config),
+                    self.host_function.device_ir.get_root(self.device_function.config),
                     [],
                 )
             self.device_function.dead_code_elimination()
@@ -188,9 +188,11 @@ def visit_Name(self, node: ast.Name) -> ast.AST:
             origin = node._type_info.origin
             if (
                 isinstance(origin, ArgumentOrigin)
-                and origin.name in self.host_fn.constexpr_args
+                and origin.name in self.host_function.constexpr_args
             ):
-                return expr_from_string(repr(self.host_fn.constexpr_args[origin.name]))
+                return expr_from_string(
+                    repr(self.host_function.constexpr_args[origin.name])
+                )
             if origin.needs_rename():
                 # `x` => `_original_globals.x`
                 return expr_from_string(origin.host_str())
@@ -207,7 +209,7 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
         elif isinstance(type_info := node._type_info, TileIndexType):
             block_info = env.block_sizes[type_info.block_size_idx]
             return expr_from_string(
-                self.host_fn.literal_expr(
+                self.host_function.literal_expr(
                     block_info.from_config(self.device_function.config)
                 )
             )
@@ -216,7 +218,7 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
             if all(isinstance(x, TileIndexType) for x in values):
                 block_infos = [env.block_sizes[x.block_size_idx] for x in values]
                 return expr_from_string(
-                    self.host_fn.literal_expr(
+                    self.host_function.literal_expr(
                         [
                             x.from_config(self.device_function.config)
                             for x in block_infos
@@ -248,15 +250,15 @@ def has_mask(self) -> bool:
 
 
 def codegen_precompile_def(
-    host_def: ast.FunctionDef, device_fn_name: str
+    host_def: ast.FunctionDef, device_function_name: str
 ) -> ast.FunctionDef:
     """
     Generate a precompile function definition for the given host function.
     The precompile function is the same as the normal function, but the call to the
     kernel is replaced with a call to make_precompiler.
 
     :param host_def: The host function definition to that is used to call the kernel.
-    :param device_fn_name: The name of the device function to be called.
+    :param device_function_name: The name of the device function to be called.
     :return: A transformed function definition with the kernel call replaced.
     """
 
@@ -285,7 +287,7 @@ def transform(node: ExtendedAST) -> ExtendedAST:
                                     ast.Return,
                                     value=value.copy(
                                         func=expr_from_string(
-                                            f"make_precompiler({device_fn_name})"
+                                            f"make_precompiler({device_function_name})"
                                         )
                                     ),
                                 )

helion/_compiler/lift_closures.py

@@ -36,11 +36,13 @@ def lift_closures(func: FunctionType, origin: Origin) -> FunctionType:
     def wrapper(*args: object, **kwargs: object) -> object:
         nonlocal new_func, closure_contents
         if new_func is None:
-            host_fn = HostFunction.current()
+            host_function = HostFunction.current()
             closure = None
             if func.__closure__ is not None:
                 closure_contents = [
-                    host_fn.register_fake(obj.cell_contents, ClosureOrigin(origin, i))
+                    host_function.register_fake(
+                        obj.cell_contents, ClosureOrigin(origin, i)
+                    )
                     for i, obj in enumerate(func.__closure__)
                 ]
                 closure = (*map(make_cell, closure_contents),)

helion/_compiler/reduction_strategy.py

@@ -235,7 +235,7 @@ def __init__(
     def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
         env = CompileEnvironment.current()
         block_index = self.block_index
-        device_fn = state.device_function
+        device_function = state.device_function
         numel = env.block_sizes[block_index].numel
         offset_var = self.offset_var(block_index)
         index_var = self.index_var(block_index)
@@ -253,14 +253,14 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
         if (mask_var := self._mask_var) is not None:
             body.append(
                 statement_from_string(
-                    f"{mask_var} = {index_var} < {device_fn.sympy_expr(numel)}"
+                    f"{mask_var} = {index_var} < {device_function.sympy_expr(numel)}"
                 )
             )
         for_node = create(
             ast.For,
             target=create(ast.Name, id=offset_var, ctx=ast.Store()),
             iter=expr_from_string(
-                f"range(0, ({device_fn.sympy_expr(numel)}), {block_size_var})"
+                f"range(0, ({device_function.sympy_expr(numel)}), {block_size_var})"
             ),
             body=body,
             orelse=[],

helion/_compiler/source_location.py

@@ -67,14 +67,14 @@ def __init__(
     def from_ast(node: ast.AST) -> SourceLocation:
         from .host_function import HostFunction
 
-        host_fn = HostFunction.current()
-        code = host_fn.fn.__code__
+        host_function = HostFunction.current()
+        code = host_function.fn.__code__
         offset = code.co_firstlineno - 1
         return SourceLocation(
             node.lineno + offset,
-            node.col_offset + host_fn.column_offset,
+            node.col_offset + host_function.column_offset,
             node.end_lineno + offset,
-            node.end_col_offset + host_fn.column_offset,
+            node.end_col_offset + host_function.column_offset,
             filename=code.co_filename,
             name=code.co_name,
         )

helion/_compiler/tile_strategy.py

@@ -211,7 +211,7 @@ def _codegen_common(
         block_indices = self.block_indices
         env = CompileEnvironment.current()
         total_numel = sympy.S.One
-        device_fn = state.device_function
+        device_function = state.device_function
         offsets_var = self.new_var("offsets", dce=True)
         block_size_var = self.block_size_var(-1)
         statements = []
@@ -226,17 +226,17 @@ def _codegen_common(
             block_index_var = self.index_var(block_idx)
             expr = offsets_var
             if total_numel != sympy.S.One:
-                expr = f"({expr}) // ({device_fn.sympy_expr(total_numel)})"
+                expr = f"({expr}) // ({device_function.sympy_expr(total_numel)})"
             if i + 1 < len(block_indices):
-                expr = f"({expr}) % ({device_fn.sympy_expr(numel)})"
+                expr = f"({expr}) % ({device_function.sympy_expr(numel)})"
             statements.append(statement_from_string(f"{block_index_var} = {expr}"))
             total_numel = total_numel * numel
 
         mask_var = self.mask_var(-1)
         if mask_var is not None:
             statements.append(
                 statement_from_string(
-                    f"{mask_var} = {offsets_var} < ({device_fn.sympy_expr(total_numel)})"
+                    f"{mask_var} = {offsets_var} < ({device_function.sympy_expr(total_numel)})"
                 )
             )
         return block_size_var, offsets_var, total_numel, statements
@@ -375,7 +375,7 @@ def mask_var(self, block_idx: int) -> str | None:
     def codegen_grid(self, state: CodegenState) -> None:
         block_indices = self.block_indices
         env = CompileEnvironment.current()
-        device_fn = state.device_function
+        device_function = state.device_function
         dtype = env.triton_index_type()
         block_sizes = self.block_size
         assert len(block_sizes) == len(block_indices)
@@ -386,7 +386,7 @@ def codegen_grid(self, state: CodegenState) -> None:
             numel = env.block_sizes[block_idx].numel
             offset_var = self.offset_var(block_idx)
             index_var = self.index_var(block_idx)
-            pid_var = device_fn.new_var(f"pid_{i}", dce=True)
+            pid_var = device_function.new_var(f"pid_{i}", dce=True)
             if block_size != 1:
                 block_size_var = self.block_size_var(block_idx)
                 assert block_size_var is not None
@@ -444,7 +444,7 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
         # TODO(jansel): refactor this to share code with codegen_grid
         block_indices = self.block_indices
         env = CompileEnvironment.current()
-        device_fn = state.device_function
+        device_function = state.device_function
         dtype = env.triton_index_type()
         block_sizes = self.block_size
         body = innermost_body = []
@@ -471,7 +471,7 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
                 ast.For,
                 target=create(ast.Name, id=offset_var, ctx=ast.Store()),
                 iter=expr_from_string(
-                    f"range(0, ({device_fn.sympy_expr(numel)}), {block_size_var})"
+                    f"range(0, ({device_function.sympy_expr(numel)}), {block_size_var})"
                 ),
                 body=body,
                 orelse=[],

helion/runtime/kernel.py

@@ -251,7 +251,7 @@ def __init__(self, kernel: Kernel, args: tuple[object, ...]) -> None:
                 else:
                     self.fake_args.append(self.env.to_fake(arg, ArgumentOrigin(name)))
             with _maybe_skip_dtype_check_in_meta_registrations():
-                self.host_fn: HostFunction = HostFunction(
+                self.host_function: HostFunction = HostFunction(
                     self.kernel.fn, self.fake_args, constexpr_args
                 )
         if len(kernel.configs) == 1:
@@ -301,7 +301,7 @@ def to_triton_code(self, config: ConfigLike) -> str:
                 # pyre-ignore[6]
                 config = Config(**config)
             self.env.config_spec.normalize(config)
-            root = generate_ast(self.host_fn, config)
+            root = generate_ast(self.host_function, config)
             return get_needed_imports(root) + unparse(root)
 
     def compile_config(self, config: ConfigLike) -> CompiledConfig:
@@ -334,7 +334,7 @@ def _debug_str(self) -> str:
         :rtype: str
         """
         with self.env:
-            return self.host_fn.debug_str()
+            return self.host_function.debug_str()
 
     def autotune(
         self,