Make TCC JIT functional on Linux and mac

Author: metab0t

CMakeLists.txt

@@ -220,10 +220,6 @@ nanobind_add_module(
 target_link_libraries(ipopt_model_ext PUBLIC ipopt_model)
 install(TARGETS ipopt_model_ext LIBRARY DESTINATION ${POI_INSTALL_DIR})
 
-# Set RPATH
-# set_rpath(nlcore_ext)
-set_rpath(ipopt_model_ext)
-
 # stub
 nanobind_add_stub(
   core_ext_stub

src/pyoptinterface/_src/codegen_c.py

@@ -38,8 +38,8 @@
 def generate_csrc_prelude(io: IO[str]):
     io.write(
         """// includes
-# include <stddef.h>
-# include <math.h>
+#include <stddef.h>
+#include <math.h>
 
 // typedefs
 typedef double float_point_t;
@@ -63,24 +63,6 @@ def generate_csrc_prelude(io: IO[str]):
     )
 
 
-def generate_csrc_prelude_declaration(io: IO[str]):
-    io.write(
-        """// includes
-# include <stddef.h>
-# include <math.h>
-
-// typedefs
-typedef double float_point_t;
-
-// externals
-// azmul
-extern float_point_t azmul(float_point_t x, float_point_t y);
-// sign
-extern float_point_t sign(float_point_t x);
-"""
-    )
-
-
 def generate_csrc_from_graph(
     io: IO[str],
     graph_obj,
@@ -299,7 +281,3 @@ def get_node_name(node: int) -> str:
 }
 """
     )
-
-    extern_function_declaration = "extern " + function_prototype
-
-    return extern_function_declaration

src/pyoptinterface/_src/ipopt.py

@@ -194,6 +194,7 @@ def compile_functions_llvm(model: "Model", jit_compiler: LLJITCompiler):
             indirect_x=True,
             indirect_p=True,
         )
+        export_functions.append(f_name)
         if autodiff_structure.has_jacobian:
             jacobian_name = name + "_jacobian"
             generate_llvmir_from_graph(
@@ -215,6 +216,7 @@ def compile_functions_llvm(model: "Model", jit_compiler: LLJITCompiler):
                 indirect_y=True,
                 add_y=True,
             )
+            export_functions.extend([jacobian_name, gradient_name])
         if autodiff_structure.has_hessian:
             hessian_name = name + "_hessian"
             generate_llvmir_from_graph(
@@ -229,8 +231,7 @@ def compile_functions_llvm(model: "Model", jit_compiler: LLJITCompiler):
                 indirect_y=True,
                 add_y=True,
             )
-
-        export_functions.extend([f_name, jacobian_name, gradient_name, hessian_name])
+            export_functions.append(hessian_name)
 
     rt = jit_compiler.compile_module(module, export_functions)
 

src/pyoptinterface/_src/jit_c.py

@@ -4,6 +4,8 @@
 
 import tccbox
 
+system = platform.system()
+
 libtcc_dir = tccbox.tcc_lib_dir()
 # windows: libtcc.dll
 # linux: libtcc.so
@@ -12,9 +14,18 @@
     "Windows": "dll",
     "Linux": "so",
     "Darwin": "dylib",
-}[platform.system()]
+}[system]
 libtcc_path = os.path.join(libtcc_dir, f"libtcc.{sharedlib_suffix}")
 
+# On Linux/Mac, tcc has lib/tcc/include/ and lib/tcc/libtcc1.a which must be included in compilation
+libtcc_extra_include_path = None
+libtcc_extra_lib_path = None
+libtcc_extra_lib_name = None
+if system in ["Linux", "Darwin"]:
+    libtcc_extra_include_path = os.path.join(libtcc_dir, "tcc", "include")
+    libtcc_extra_lib_path = os.path.join(libtcc_dir, "tcc")
+    libtcc_extra_lib_name = "libtcc1.a"
+
 # Define types
 TCCState = ctypes.c_void_p
 TCCReallocFunc = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_ulong)
@@ -49,6 +60,32 @@ def create_state(self):
         if self.libtcc.tcc_set_output_type(state, TCC_OUTPUT_MEMORY) == -1:
             raise Exception("Failed to set output type")
 
+        # Add extra include path and library path
+        if libtcc_extra_include_path:
+            if (
+                self.libtcc.tcc_add_include_path(
+                    state, libtcc_extra_include_path.encode()
+                )
+                == -1
+            ):
+                raise Exception("Failed to add extra include path")
+            if (
+                self.libtcc.tcc_add_sysinclude_path(
+                    state, libtcc_extra_include_path.encode()
+                )
+                == -1
+            ):
+                raise Exception("Failed to add extra sysinclude path")
+        if libtcc_extra_lib_path:
+            if (
+                self.libtcc.tcc_add_library_path(state, libtcc_extra_lib_path.encode())
+                == -1
+            ):
+                raise Exception("Failed to add extra library path")
+        if libtcc_extra_lib_name:
+            if self.libtcc.tcc_add_library(state, libtcc_extra_lib_name.encode()) == -1:
+                raise Exception("Failed to add extra library")
+
         self.states.append(state)
 
         return state
@@ -65,6 +102,14 @@ def _initialize_function_prototypes(self):
         libtcc.tcc_add_symbol.restype = ctypes.c_int
         libtcc.tcc_get_symbol.argtypes = [TCCState, ctypes.c_char_p]
         libtcc.tcc_get_symbol.restype = ctypes.c_void_p
+        libtcc.tcc_add_include_path.argtypes = [TCCState, ctypes.c_char_p]
+        libtcc.tcc_add_include_path.restype = ctypes.c_int
+        libtcc.tcc_add_sysinclude_path.argtypes = [TCCState, ctypes.c_char_p]
+        libtcc.tcc_add_sysinclude_path.restype = ctypes.c_int
+        libtcc.tcc_add_library_path.argtypes = [TCCState, ctypes.c_char_p]
+        libtcc.tcc_add_library_path.restype = ctypes.c_int
+        libtcc.tcc_add_library.argtypes = [TCCState, ctypes.c_char_p]
+        libtcc.tcc_add_library.restype = ctypes.c_int
 
     def compile_string(self, state, c_code: str):
         if self.libtcc.tcc_compile_string(state, c_code.encode()) == -1:

tests/test_nlp.py

@@ -157,5 +157,6 @@ def con(vars):
 
 
 if __name__ == "__main__":
-    test_ipopt()
-    test_nlp_param()
+    test_ipopt(ipopt.Model)
+    test_nlp_param(ipopt.Model)
+    test_nlfunc_ifelse(ipopt.Model)

tests/test_nlp_rocket.py

@@ -97,4 +97,4 @@ def test_rocket(ipopt_model_ctor):
 
 
 if __name__ == "__main__":
-    test_rocket()
+    test_rocket(ipopt.Model)