Rename function_name parameter for consistency

lleaves/compiler/codegen/codegen.py

@@ -39,7 +39,7 @@ class LTree:
     class_id: int
 
 
-def gen_forest(forest, module, fblocksize, function_name="forest_root"):
+def gen_forest(forest, module, fblocksize, froot_func_name):
     """
     Populate the passed IR module with code for the forest.
 
@@ -81,7 +81,7 @@ def gen_forest(forest, module, fblocksize, function_name="forest_root"):
     root_func = ir.Function(
         module,
         ir.FunctionType(ir.VoidType(), (DOUBLE_PTR, DOUBLE_PTR, INT, INT)),
-        name=function_name,
+        name=froot_func_name,
     )
 
     def make_tree(tree):

lleaves/compiler/tree_compiler.py

@@ -7,12 +7,18 @@
 from lleaves.compiler.codegen import gen_forest
 
 
-def compile_to_module(file_path, fblocksize=34, finline=True, raw_score=False, function_name="forest_root"):
+def compile_to_module(
+    file_path,
+    fblocksize=34,
+    finline=True,
+    raw_score=False,
+    froot_func_name="forest_root",
+):
     forest = parse_to_ast(file_path)
     forest.raw_score = raw_score
 
     ir = llvmlite.ir.Module(name="forest")
-    gen_forest(forest, ir, fblocksize, function_name)
+    gen_forest(forest, ir, fblocksize, froot_func_name)
 
     ir.triple = llvm.get_process_triple()
     module = llvm.parse_assembly(str(ir))

lleaves/lleaves.py

@@ -85,7 +85,7 @@ def compile(
         fblocksize=34,
         fcodemodel="large",
         finline=True,
-        function_name="forest_root",
+        froot_func_name="forest_root",
     ):
         """
         Generate the LLVM IR for this model and compile it to ASM.
@@ -108,7 +108,8 @@ def compile(
             very large forests.
         :param finline: Whether or not to inline function. Setting this to False will speed-up compilation time
             significantly but will slow down prediction.
-        :param function_name: Name of the prediction function, default to be `forest_root`.
+        :param froot_func_name: Name of entry point function in the compiled binary. This is the function to link when
+            writing a C function wrapper. Defaults to "forest_root".
         """
         assert 0 < fblocksize
         assert fcodemodel in ("small", "large")
@@ -119,7 +120,7 @@ def compile(
                 raw_score=raw_score,
                 fblocksize=fblocksize,
                 finline=finline,
-                function_name=function_name,
+                froot_func_name=froot_func_name,
             )
         else:
             # when loading binary from cache we use a dummy empty module
@@ -131,7 +132,7 @@ def compile(
         )
 
         # Drops GIL during call, re-acquires it after
-        addr = self._execution_engine.get_function_address(function_name)
+        addr = self._execution_engine.get_function_address(froot_func_name)
         self._c_entry_func = ENTRY_FUNC_TYPE(addr)
 
         self.is_compiled = True

tests/test_compile_flags.py

@@ -79,7 +79,7 @@ def test_no_inline(NYC_data):
 def test_function_name():
     llvm_model = Model(model_file="tests/models/tiniest_single_tree/model.txt")
     lgbm_model = Booster(model_file="tests/models/tiniest_single_tree/model.txt")
-    llvm_model.compile(function_name="tiniest_single_tree_123132_XXX-")
+    llvm_model.compile(froot_func_name="tiniest_single_tree_123132_XXX-")
 
     data = [
         [1.0] * 3,