lcompilers · ubaidsk · Apr 20, 2024 · Feb 29, 2024 · Mar 8, 2024 · Mar 8, 2024
diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt
@@ -76,8 +76,9 @@ message("LPYTHON_RTLIB_DIR: ${LPYTHON_RTLIB_DIR}")
 message("LPYTHON_RTLIB_LIBRARY: ${LPYTHON_RTLIB_LIBRARY}")
 
 
-macro(RUN_UTIL RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOMOD RUN_EXTRA_ARGS RUN_COPY_TO_BIN)
+macro(RUN_UTIL RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOMOD RUN_EXTRA_ARGS RUN_COPY_TO_BIN RUN_SKIP_ON_JIT)
     set(fail ${${RUN_FAIL}})
+    set(skip_jit ${${RUN_SKIP_ON_JIT}})
     set(name ${${RUN_NAME}})
     set(file_name ${${RUN_FILE_NAME}})
     set(labels ${${RUN_LABELS}})
@@ -101,11 +102,21 @@ macro(RUN_UTIL RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOM
             set_target_properties(${name} PROPERTIES LINKER_LANGUAGE C)
             target_link_libraries(${name} lpython_rtlib)
             add_test(${name} ${CMAKE_CURRENT_BINARY_DIR}/${name})
+            add_test(
+                NAME "${name}_jit"
+                COMMAND ${LPYTHON} --jit ${extra_args} ${CMAKE_CURRENT_SOURCE_DIR}/${file_name}.py)
             if (labels)
                 set_tests_properties(${name} PROPERTIES LABELS "${labels}")
+                set_tests_properties("${name}_jit" PROPERTIES LABELS "${labels}")
             endif()
             if (${fail})
                 set_tests_properties(${name} PROPERTIES WILL_FAIL TRUE)
+                set_tests_properties("${name}_jit" PROPERTIES WILL_FAIL TRUE)
+            endif()
+            if (${skip_jit})
+                set_tests_properties("${name}_jit" PROPERTIES SKIP_RETURN_CODE ${skip_jit})
+            elseif (DEFINED ${RUN_EXTRAFILES} AND NOT ${RUN_FAIL})
+                set_tests_properties("${name}_jit" PROPERTIES SKIP_RETURN_CODE 1)
             endif()
         elseif (KIND STREQUAL "llvm_py")
             add_custom_command(
@@ -118,11 +129,21 @@ macro(RUN_UTIL RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOM
             set_target_properties(${name} PROPERTIES LINKER_LANGUAGE C)
             target_link_libraries(${name} lpython_rtlib Python::Python)
             add_test(${name} ${CMAKE_CURRENT_BINARY_DIR}/${name})
+            add_test(
+                NAME "${name}_jit"
+                COMMAND ${LPYTHON} --jit ${extra_args} ${CMAKE_CURRENT_SOURCE_DIR}/${file_name}.py)
             if (labels)
                 set_tests_properties(${name} PROPERTIES LABELS "${labels}")
+                set_tests_properties("${name}_jit" PROPERTIES LABELS "${labels}")
             endif()
             if (${fail})
                 set_tests_properties(${name} PROPERTIES WILL_FAIL TRUE)
+                set_tests_properties("${name}_jit" PROPERTIES WILL_FAIL TRUE)
+            endif()
+            if (${skip_jit})
+                set_tests_properties("${name}_jit" PROPERTIES SKIP_RETURN_CODE ${skip_jit})
+            elseif (DEFINED ${RUN_EXTRAFILES} AND NOT ${RUN_FAIL})
+                set_tests_properties("${name}_jit" PROPERTIES SKIP_RETURN_CODE 1)
             endif()
         elseif(KIND STREQUAL "llvm_sym")
             add_custom_command(
@@ -139,11 +160,21 @@ macro(RUN_UTIL RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOM
             endif()
             target_link_libraries(${name} lpython_rtlib ${SYMENGINE_LIB})
             add_test(${name} ${CMAKE_CURRENT_BINARY_DIR}/${name})
+            add_test(
+                NAME "${name}_jit"
+                COMMAND ${LPYTHON} --enable-symengine --jit ${extra_args} ${CMAKE_CURRENT_SOURCE_DIR}/${file_name}.py)
             if (labels)
                 set_tests_properties(${name} PROPERTIES LABELS "${labels}")
+                set_tests_properties("${name}_jit" PROPERTIES LABELS "${labels}")
             endif()
             if (${fail})
                 set_tests_properties(${name} PROPERTIES WILL_FAIL TRUE)
+                set_tests_properties("${name}_jit" PROPERTIES WILL_FAIL TRUE)
+            endif()
+            if (${skip_jit})
+                set_tests_properties("${name}_jit" PROPERTIES SKIP_RETURN_CODE ${skip_jit})
+            elseif (DEFINED ${RUN_EXTRAFILES} AND NOT ${RUN_FAIL})
+                set_tests_properties("${name}_jit" PROPERTIES SKIP_RETURN_CODE 1)
             endif()
         elseif(KIND STREQUAL "c")
             add_custom_command(
@@ -312,7 +343,7 @@ endmacro(RUN_UTIL)
 
 macro(RUN)
     set(options FAIL NOFAST NOMOD ENABLE_CPYTHON LINK_NUMPY NO_WARNINGS)
-    set(oneValueArgs NAME IMPORT_PATH COPY_TO_BIN REQ_PY_VER)
+    set(oneValueArgs NAME IMPORT_PATH COPY_TO_BIN REQ_PY_VER SKIP_ON_JIT)
     set(multiValueArgs LABELS EXTRAFILES)
     cmake_parse_arguments(RUN "${options}" "${oneValueArgs}"
                         "${multiValueArgs}" ${ARGN} )
@@ -351,14 +382,14 @@ macro(RUN)
     endif()
 
     if (NOT FAST)
-        RUN_UTIL(RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOMOD RUN_EXTRA_ARGS RUN_COPY_TO_BIN)
+        RUN_UTIL(RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOMOD RUN_EXTRA_ARGS RUN_COPY_TO_BIN RUN_SKIP_ON_JIT)
     endif()
 
     if ((FAST) AND (NOT RUN_NOFAST))
         set(RUN_EXTRA_ARGS ${RUN_EXTRA_ARGS} --fast)
         set(RUN_NAME "${RUN_NAME}_FAST")
         list(REMOVE_ITEM RUN_LABELS cpython cpython_sym) # remove cpython, cpython_sym, from --fast test
-        RUN_UTIL(RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOMOD RUN_EXTRA_ARGS RUN_COPY_TO_BIN)
+        RUN_UTIL(RUN_FAIL RUN_NAME RUN_FILE_NAME RUN_LABELS RUN_EXTRAFILES RUN_NOMOD RUN_EXTRA_ARGS RUN_COPY_TO_BIN RUN_SKIP_ON_JIT)
     endif()
 endmacro(RUN)
 
@@ -806,7 +837,7 @@ RUN(NAME comp_01 LABELS cpython llvm c wasm wasm_x64)
 RUN(NAME bit_operations_i32 LABELS cpython llvm c wasm wasm_x64)
 RUN(NAME bit_operations_i64 LABELS cpython llvm c wasm)
 
-RUN(NAME test_argv_01      LABELS cpython llvm NOFAST)
+RUN(NAME test_argv_01      LABELS cpython llvm NOFAST SKIP_ON_JIT 1)
 RUN(NAME global_syms_01    LABELS cpython llvm c)
 RUN(NAME global_syms_02    LABELS cpython llvm c)
 RUN(NAME global_syms_03_b  LABELS cpython llvm c)

diff --git a/src/bin/lpython.cpp b/src/bin/lpython.cpp
@@ -792,13 +792,17 @@ int emit_llvm(const std::string &infile,
     return 0;
 }
 
-int compile_python_to_object_file(
+/*
+    Compiles python to object file, if `to_jit` is false
+    otherwise execute python code using llvm JIT
+*/
+int compile_python_using_llvm(
         const std::string &infile,
         const std::string &outfile,
         const std::string &runtime_library_dir,
         LCompilers::PassManager& pass_manager,
         CompilerOptions &compiler_options,
-        bool time_report, bool arg_c=false)
+        bool time_report, bool arg_c=false, bool to_jit=false)
 {
     Allocator al(4*1024);
     LCompilers::diag::Diagnostics diagnostics;
@@ -869,7 +873,6 @@ int compile_python_to_object_file(
     }
     LCompilers::PythonCompiler fe(compiler_options);
     LCompilers::LLVMEvaluator e(compiler_options.target);
-    std::unique_ptr<LCompilers::LLVMModule> m;
     auto asr_to_llvm_start = std::chrono::high_resolution_clock::now();
     LCompilers::Result<std::unique_ptr<LCompilers::LLVMModule>>
         res = fe.get_llvm3(*asr, pass_manager, diagnostics, infile);
@@ -882,12 +885,55 @@ int compile_python_to_object_file(
         print_time_report(times, time_report);
         return 3;
     }
-    m = std::move(res.result);
-    auto llvm_start = std::chrono::high_resolution_clock::now();
-    e.save_object_file(*(m->m_m), outfile);
-    auto llvm_end = std::chrono::high_resolution_clock::now();
-    times.push_back(std::make_pair("LLVM to binary", std::chrono::duration<double, std::milli>(llvm_end - llvm_start).count()));
-    print_time_report(times, time_report);
+    std::unique_ptr<LCompilers::LLVMModule> m = std::move(res.result);
+
+    if (to_jit) {
+        LCompilers::LPython::DynamicLibrary cpython_lib;
+        LCompilers::LPython::DynamicLibrary symengine_lib;
+
+        if (compiler_options.enable_cpython) {
+            LCompilers::LPython::open_cpython_library(cpython_lib);
+        }
+        if (compiler_options.enable_symengine) {
+            LCompilers::LPython::open_symengine_library(symengine_lib);
+        }
+
+        auto llvm_start = std::chrono::high_resolution_clock::now();
+
+        bool call_init = false;
+        bool call_stmts = false;
+        if (m->get_return_type("__module___main_____main__global_init") == "void") {
+            call_init = true;
+        }
+        if (m->get_return_type("__module___main_____main__global_stmts") == "void") {
+            call_stmts = true;
+        }
+
+        e.add_module(std::move(m));
+        if (call_init) {
+            e.voidfn("__module___main_____main__global_init");
+        }
+        if (call_stmts) {
+            e.voidfn("__module___main_____main__global_stmts");
+        }
+
+        if (compiler_options.enable_cpython) {
+            LCompilers::LPython::close_cpython_library(cpython_lib);
+        }
+        if (compiler_options.enable_symengine) {
+            LCompilers::LPython::close_symengine_library(symengine_lib);
+        }
+
+        auto llvm_end = std::chrono::high_resolution_clock::now();
+        times.push_back(std::make_pair("LLVM JIT execution", std::chrono::duration<double, std::milli>(llvm_end - llvm_start).count()));
+        print_time_report(times, time_report);
+    } else {
+        auto llvm_start = std::chrono::high_resolution_clock::now();
+        e.save_object_file(*(m->m_m), outfile);
+        auto llvm_end = std::chrono::high_resolution_clock::now();
+        times.push_back(std::make_pair("LLVM to binary", std::chrono::duration<double, std::milli>(llvm_end - llvm_start).count()));
+        print_time_report(times, time_report);
+    }
     return 0;
 }
 
@@ -1560,6 +1606,7 @@ int main(int argc, char *argv[])
         bool print_rtl_header_dir = false;
         bool print_rtl_dir = false;
         bool separate_compilation = false;
+        bool to_jit = false;
 
         std::string arg_fmt_file;
         // int arg_fmt_indent = 4;
@@ -1593,6 +1640,7 @@ int main(int argc, char *argv[])
         app.add_option("-I", compiler_options.import_paths, "Specify the paths"
             "to look for the module")->allow_extra_args(false);
         // app.add_option("-J", arg_J, "Where to save mod files");
+        app.add_flag("--jit", to_jit, "Execute the program using just-in-time (JIT) compiler");
         app.add_flag("-g", compiler_options.emit_debug_info, "Compile with debugging information");
         app.add_flag("--debug-with-line-column", compiler_options.emit_debug_line_column,
             "Convert the linear location info into line + column in the debugging information");
@@ -1894,10 +1942,10 @@ int main(int argc, char *argv[])
             }
         }
 
-        if (arg_c) {
+        if (arg_c && !to_jit) {
             if (backend == Backend::llvm) {
 #ifdef HAVE_LFORTRAN_LLVM
-                return compile_python_to_object_file(arg_file, outfile, runtime_library_dir, lpython_pass_manager, compiler_options, time_report,
+                return compile_python_using_llvm(arg_file, outfile, runtime_library_dir, lpython_pass_manager, compiler_options, time_report,
                                                      arg_c);
 #else
                 std::cerr << "The -c option requires the LLVM backend to be enabled. Recompile with `WITH_LLVM=yes`." << std::endl;
@@ -1911,6 +1959,23 @@ int main(int argc, char *argv[])
         if (endswith(arg_file, ".py"))
         {
             int err = 0;
+            if (to_jit) {
+#ifdef HAVE_LFORTRAN_LLVM
+                if (backend != Backend::llvm) {
+                    std::cerr << "JIT option is only available with LLVM backend" << std::endl;
+                    return 1;
+                }
+                compiler_options.emit_debug_info = false;
+                compiler_options.emit_debug_line_column = false;
+                compiler_options.generate_object_code = false;
+                return compile_python_using_llvm(arg_file, "", runtime_library_dir,
+                        lpython_pass_manager, compiler_options, time_report, false, true);
+#else
+                std::cerr << "Just-In-Time Compilation of Python files requires the LLVM backend to be enabled."
+                             " Recompile with `WITH_LLVM=yes`." << std::endl;
+                return 1;
+#endif
+        }
             if (backend == Backend::x86) {
                 err = compile_to_binary_x86(arg_file, outfile,
                         runtime_library_dir, compiler_options, time_report);
@@ -1931,7 +1996,7 @@ int main(int argc, char *argv[])
             } else if (backend == Backend::llvm) {
 #ifdef HAVE_LFORTRAN_LLVM
                 std::string tmp_o = outfile + ".tmp.o";
-                err = compile_python_to_object_file(arg_file, tmp_o, runtime_library_dir,
+                err = compile_python_using_llvm(arg_file, tmp_o, runtime_library_dir,
                     lpython_pass_manager, compiler_options, time_report);
                 if (err != 0) return err;
                 err = link_executable({tmp_o}, outfile, runtime_library_dir,

diff --git a/src/libasr/codegen/llvm_utils.cpp b/src/libasr/codegen/llvm_utils.cpp
@@ -2120,17 +2120,15 @@ namespace LCompilers {
             throw LCompilersException("list for " + type_code + " not declared yet.");
         }
         int32_t type_size = std::get<1>(typecode2listtype[type_code]);
-        llvm::Value* arg_size = llvm::ConstantInt::get(context,
-                                    llvm::APInt(32, type_size * initial_capacity));
-
-        llvm::Value* list_data = LLVM::lfortran_malloc(context, module, *builder,
-                                                       arg_size);
+        llvm::Value* llvm_type_size = llvm::ConstantInt::get(context, llvm::APInt(32, type_size));
+        llvm::Value* current_capacity = llvm::ConstantInt::get(context, llvm::APInt(32, initial_capacity));
+        llvm::Value* list_data = LLVM::lfortran_calloc(context, module, *builder,
+                                                       current_capacity, llvm_type_size);
         llvm::Type* el_type = std::get<2>(typecode2listtype[type_code]);
         list_data = builder->CreateBitCast(list_data, el_type->getPointerTo());
         llvm::Value* list_data_ptr = get_pointer_to_list_data(list);
         builder->CreateStore(list_data, list_data_ptr);
         llvm::Value* current_end_point = llvm::ConstantInt::get(context, llvm::APInt(32, n));
-        llvm::Value* current_capacity = llvm::ConstantInt::get(context, llvm::APInt(32, initial_capacity));
         builder->CreateStore(current_end_point, get_pointer_to_current_end_point(list));
         builder->CreateStore(current_capacity, get_pointer_to_current_capacity(list));
     }
@@ -2143,8 +2141,8 @@ namespace LCompilers {
         }
         int32_t type_size = std::get<1>(typecode2listtype[type_code]);
         llvm::Value* llvm_type_size = llvm::ConstantInt::get(context, llvm::APInt(32, type_size));
-        llvm::Value* arg_size = builder->CreateMul(llvm_type_size, initial_capacity);
-        llvm::Value* list_data = LLVM::lfortran_malloc(context, module, *builder, arg_size);
+        llvm::Value* list_data = LLVM::lfortran_calloc(context, module, *builder,
+                                                        initial_capacity, llvm_type_size);
 
         llvm::Type* el_type = std::get<2>(typecode2listtype[type_code]);
         list_data = builder->CreateBitCast(list_data, el_type->getPointerTo());
@@ -2288,10 +2286,9 @@ namespace LCompilers {
         builder->CreateStore(src_capacity, dest_capacity_ptr);
         llvm::Value* src_data = LLVM::CreateLoad(*builder, get_pointer_to_list_data(src));
         int32_t type_size = std::get<1>(typecode2listtype[src_type_code]);
-        llvm::Value* arg_size = builder->CreateMul(llvm::ConstantInt::get(context,
-                                                   llvm::APInt(32, type_size)), src_capacity);
-        llvm::Value* copy_data = LLVM::lfortran_malloc(context, *module, *builder,
-                                                       arg_size);
+        llvm::Value* llvm_type_size = llvm::ConstantInt::get(context, llvm::APInt(32, type_size));
+        llvm::Value* copy_data = LLVM::lfortran_calloc(context, *module, *builder,
+                                                       src_capacity, llvm_type_size);
         llvm::Type* el_type = std::get<2>(typecode2listtype[src_type_code]);
         copy_data = builder->CreateBitCast(copy_data, el_type->getPointerTo());
 
@@ -2346,6 +2343,8 @@ namespace LCompilers {
             // end
             llvm_utils->start_new_block(loopend);
         } else {
+            llvm::Value* arg_size = builder->CreateMul(llvm::ConstantInt::get(context,
+                                                   llvm::APInt(32, type_size)), src_capacity);
             builder->CreateMemCpy(copy_data, llvm::MaybeAlign(), src_data,
                                   llvm::MaybeAlign(), arg_size);
             builder->CreateStore(copy_data, get_pointer_to_list_data(dest));

diff --git a/src/lpython/utils.cpp b/src/lpython/utils.cpp
@@ -3,6 +3,8 @@
 #define NOMINMAX
 #endif // NOMINMAX
 #include <windows.h>
+#else
+#include <dlfcn.h>
 #endif
 
 #include <fstream>
@@ -126,6 +128,58 @@ bool path_exists(std::string path) {
     }
 }
 
+#ifdef HAVE_LFORTRAN_LLVM
+
+void open_cpython_library(DynamicLibrary &l) {
+    std::string conda_prefix = std::getenv("CONDA_PREFIX");
+#if defined (__linux__)
+    l.l = dlopen((conda_prefix + "/lib/libpython3.so").c_str(), RTLD_DEEPBIND | RTLD_GLOBAL | RTLD_NOW);
+#elif defined (__APPLE__)
+    l.l = dlopen((conda_prefix + "/lib/libpython3.dylib").c_str(), RTLD_GLOBAL | RTLD_NOW);
+#else
+    l.l = LoadLibrary((conda_prefix + "\\python3.dll").c_str());
+#endif
+
+    if (l.l == nullptr)
+        throw "Could not open CPython library";
+}
+
+void close_cpython_library(DynamicLibrary &l) {
+#if (defined (__linux__)) or (defined (__APPLE__))
+    dlclose(l.l);
+    l.l = nullptr;
+#else
+    FreeLibrary((HMODULE)l.l);
+    l.l = nullptr;
+#endif
+}
+
+void open_symengine_library(DynamicLibrary &l) {
+    std::string conda_prefix = std::getenv("CONDA_PREFIX");
+#if defined (__linux__)
+    l.l = dlopen((conda_prefix + "/lib/libsymengine.so").c_str(), RTLD_DEEPBIND | RTLD_GLOBAL | RTLD_NOW);
+#elif defined (__APPLE__)
+    l.l = dlopen((conda_prefix + "/lib/libsymengine.dylib").c_str(), RTLD_GLOBAL | RTLD_NOW);
+#else
+    l.l = LoadLibrary((conda_prefix + "\\Library\\bin\\symengine-0.11.dll").c_str());
+#endif
+
+    if (l.l == nullptr)
+        throw "Could not open SymEngine library";
+}
+
+void close_symengine_library(DynamicLibrary &l) {
+#if (defined (__linux__)) or (defined (__APPLE__))
+    dlclose(l.l);
+    l.l = nullptr;
+#else
+    FreeLibrary((HMODULE)l.l);
+    l.l = nullptr;
+#endif
+}
+
+#endif
+
 // Decodes the exit status code of the process (in Unix)
 // See `WEXITSTATUS` for more information.
 // https://stackoverflow.com/a/27117435/15913193