Allow passing backend optimisation flags in lpython decorator (#2201)

Author: czgdp1807

.github/workflows/CI.yml

@@ -349,7 +349,7 @@ jobs:
       - name: Test Linux
         shell: bash -e -l {0}
         run: |
-            ctest
+            ctest --rerun-failed --output-on-failure
             ./run_tests.py -s
             cd integration_tests
             ./run_tests.py -b llvm c

integration_tests/lpython_decorator_01.py

@@ -1,7 +1,7 @@
 from numpy import array
 from lpython import i32, f64, lpython
 
-@lpython
+@lpython(backend="c", backend_optimisation_flags=["-ffast-math", "-funroll-loops", "-O3"])
 def fast_sum(n: i32, x: f64[:]) -> f64:
     s: f64 = 0.0
     i: i32

integration_tests/lpython_decorator_02.py

@@ -3,7 +3,7 @@
 
 n = TypeVar("n")
 
-@lpython
+@lpython(backend="c", backend_optimisation_flags=["-ffast-math", "-funroll-loops"])
 def multiply_01(n: i32, x: f64[:]) -> f64[n]:
     i: i32
     for i in range(n):

src/runtime/lpython/lpython.py

@@ -3,6 +3,7 @@
 import ctypes
 import platform
 from dataclasses import dataclass as py_dataclass, is_dataclass as py_is_dataclass
+import functools
 
 
 # TODO: this does not seem to restrict other imports
@@ -647,37 +648,49 @@ def ccallable(f):
 def ccallback(f):
     return f
 
-class lpython:
-    """
-    The @lpython decorator compiles a given function using LPython.
+class LpythonJITCache:
 
-    The decorator should be used from CPython mode, i.e., when the module is
-    being run using CPython. When possible, it is recommended to use LPython
-    for the main program, and use the @cpython decorator from the LPython mode
-    to access CPython features that are not supported by LPython.
-    """
+    def __init__(self):
+        self.pyfunc2compiledfunc = {}
+
+    def compile(self, function, backend, optimisation_flags):
+        if function in self.pyfunc2compiledfunc:
+            return self.pyfunc2compiledfunc[function]
+
+        if optimisation_flags is not None and backend is None:
+            raise ValueError("backend must be specified if backend_optimisation_flags are provided.")
+
+        if backend is None:
+            backend = "c"
 
-    def __init__(self, function):
         def get_rtlib_dir():
             current_dir = os.path.dirname(os.path.abspath(__file__))
             return os.path.join(current_dir, "..")
 
-        self.fn_name = function.__name__
+        fn_name = function.__name__
         # Get the source code of the function
         source_code = getsource(function)
         source_code = source_code[source_code.find('\n'):]
 
-        dir_name = "./lpython_decorator_" + self.fn_name
+        dir_name = "./lpython_decorator_" + fn_name
         if not os.path.exists(dir_name):
             os.mkdir(dir_name)
-        filename = dir_name + "/" + self.fn_name
+        filename = dir_name + "/" + fn_name
 
         # Open the file for writing
         with open(filename + ".py", "w") as file:
             # Write the Python source code to the file
             file.write("@pythoncallable")
             file.write(source_code)
 
+        if backend != "c":
+            raise NotImplementedError("Backend %s is not supported with @lpython yet."%(backend))
+
+        opt_flags = " "
+        if optimisation_flags is not None:
+            for opt_flag in optimisation_flags:
+                opt_flags += opt_flag + " "
+
         # ----------------------------------------------------------------------
         # Generate the shared library
         # TODO: Use LLVM instead of C backend
@@ -687,12 +700,14 @@ def get_rtlib_dir():
 
         gcc_flags = ""
         if platform.system() == "Linux":
-            gcc_flags = " -shared -fPIC "
+            gcc_flags = " -shared -fPIC"
         elif platform.system() == "Darwin":
-            gcc_flags = " -bundle -flat_namespace -undefined suppress "
+            gcc_flags = " -bundle -flat_namespace -undefined suppress"
         else:
             raise NotImplementedError("Platform not implemented")
 
+        gcc_flags += opt_flags
+
         from numpy import get_include
         from distutils.sysconfig import get_python_inc, get_python_lib, \
             get_python_version
@@ -706,17 +721,38 @@ def get_rtlib_dir():
 
         # ----------------------------------------------------------------------
         # Compile the C file and create a shared library
+        shared_library_name = "lpython_module_" + fn_name
         r = os.system("gcc -g" +  gcc_flags + python_path + numpy_path +
-            filename + ".c -o lpython_module_" + self.fn_name + ".so " +
+            filename + ".c -o " + shared_library_name + ".so " +
             rt_path_01 + rt_path_02 + python_lib)
         assert r == 0, "Failed to create the shared library"
+        self.pyfunc2compiledfunc[function] = (shared_library_name, fn_name)
+        return self.pyfunc2compiledfunc[function]
 
-    def __call__(self, *args, **kwargs):
-        import sys; sys.path.append('.')
-        # import the symbol from the shared library
-        function = getattr(__import__("lpython_module_" + self.fn_name),
-            self.fn_name)
-        return function(*args, **kwargs)
+lpython_jit_cache = LpythonJITCache()
+
+# Taken from https://stackoverflow.com/a/24617244
+def lpython(original_function=None, backend=None, backend_optimisation_flags=None):
+    """
+    The @lpython decorator compiles a given function using LPython.
+
+    The decorator should be used from CPython mode, i.e., when the module is
+    being run using CPython. When possible, it is recommended to use LPython
+    for the main program, and use the @cpython decorator from the LPython mode
+    to access CPython features that are not supported by LPython.
+    """
+    def _lpython(function):
+        @functools.wraps(function)
+        def __lpython(*args, **kwargs):
+            import sys; sys.path.append('.')
+            lib_name, fn_name = lpython_jit_cache.compile(
+                function, backend, backend_optimisation_flags)
+            return getattr(__import__(lib_name), fn_name)(*args, **kwargs)
+        return __lpython
+
+    if original_function:
+        return _lpython(original_function)
+    return _lpython
 
 def bitnot(x, bitsize):
     return (~x) % (2 ** bitsize)