Merge pull request #1891 from Thirumalai-Shaktivel/lpython_02

Handle array as return variable using the lpython decorator

Author: certik

integration_tests/CMakeLists.txt

@@ -619,5 +619,6 @@ RUN(NAME intrinsics_01     LABELS cpython llvm NOFAST) # any
 
 # lpython decorator
 RUN(NAME lpython_decorator_01    LABELS cpython)
+RUN(NAME lpython_decorator_02    LABELS cpython)
 
 COMPILE(NAME import_order_01     LABELS cpython llvm c) # any

integration_tests/lpython_decorator_02.py

@@ -0,0 +1,19 @@
+from numpy import array
+from lpython import i32, f64, lpython, TypeVar
+
+n = TypeVar("n")
+
+@lpython
+def multiply(n: i32, x: f64[:]) -> f64[n]:
+    i: i32
+    for i in range(n):
+        x[i] *= 5.0
+    return x
+
+def test():
+    size: i32 = 5
+    x: f64[5] = array([11.0, 12.0, 13.0, 14.0, 15.0])
+    y: f64[5] = (multiply(size, x))
+    assert y[2] == 65.
+
+test()

src/runtime/lpython/lpython.py

@@ -637,7 +637,8 @@ def get_rtlib_dir():
 
         def get_type_info(arg):
             # return_type -> (`type_format`, `variable type`, `array struct name`)
-            # See: https://docs.python.org/3/c-api/arg.html for more info on type_format
+            # See: https://docs.python.org/3/c-api/arg.html for more info on `type_format`
+            # `array struct name`: used by the C backend
             if arg == f64:
                 return ('d', "double", 'r64')
             elif arg == f32:
@@ -652,7 +653,10 @@ def get_type_info(arg):
                 t = get_type_info(arg._type)
                 if t[2] == '':
                     raise NotImplementedError("Type %r not implemented" % arg)
-                return ('O', ["PyArrayObject *", "struct "+t[2]+" *", t[1]+" *"], '')
+                n = ''
+                if not isinstance(arg._dims, slice):
+                    n = arg._dims._name
+                return ('O', ["PyArrayObject *", "struct "+t[2]+" *", t[1]+" *", n], '')
             else:
                 raise NotImplementedError("Type %r not implemented" % arg)
 
@@ -662,6 +666,18 @@ def get_data_type(t):
             else:
                 return t + " "
 
+        def get_typenum(t):
+            if t == "int":
+                return "NPY_INT"
+            elif t == "long int":
+                return "NPY_LONG"
+            elif t == "float":
+                return "NPY_FLOAT"
+            elif t == "double":
+                return "NPY_DOUBLE"
+            else:
+                raise NotImplementedError("Type %s not implemented" % t)
+
         self.fn_name = function.__name__
         # Get the source code of the function
         source_code = getsource(function)
@@ -682,46 +698,55 @@ def get_data_type(t):
         self.arg_type_formats = ""
         self.return_type = ""
         self.return_type_format = ""
+        self.array_as_return_type = ()
         self.arg_types = {}
-        counter = 1
         for t in types.keys():
             if t == "return":
                 type = get_type_info(types[t])
-                self.return_type_format = type[0]
-                self.return_type = type[1]
+                if type[0] == 'O':
+                    self.array_as_return_type = type
+                    continue
+                else:
+                    self.return_type_format = type[0]
+                    self.return_type = type[1]
             else:
                 type = get_type_info(types[t])
                 self.arg_type_formats += type[0]
-                self.arg_types[counter] = type[1]
-                counter += 1
+                self.arg_types[t] = type[1]
         # ----------------------------------------------------------------------
-        # `arg_0`: used as the return variables
-        # arguments are declared as `arg_1`, `arg_2`, ...
-        variables_decl = ""
+        # `_<fn_name>_return_value`: used as the return variables
+        variables_decl = "// Declare return variables and arguments\n"
         if self.return_type != "":
-            variables_decl = "// Declare return variables and arguments\n"
-            variables_decl += "    " + get_data_type(self.return_type) + "arg_" \
-                + str(0) + ";\n"
+            variables_decl += "    " + get_data_type(self.return_type)  \
+                + "_" + self.fn_name + "_return_value;\n"
+        elif self.array_as_return_type:
+            variables_decl += "    " + self.array_as_return_type[1][1] + "_" \
+                + self.fn_name + "_return_value = malloc(sizeof(" \
+                + self.array_as_return_type[1][1][:-2] + "));\n"
+        else:
+            variables_decl = ""
         # ----------------------------------------------------------------------
         # `PyArray_AsCArray` is used to convert NumPy Arrays to C Arrays
-        # `fill_array_details` contains arrays operations to be
+        # `fill_array_details` contains array operations to be
         # performed on the arguments
         # `parse_args` are used to capture the args from CPython
         # `pass_args` are the args that are passed to the shared library function
         fill_array_details = ""
         parse_args = ""
         pass_args = ""
         numpy_init = ""
+        prefix_comma = False
         for i, t in self.arg_types.items():
-            if i > 1:
+            if prefix_comma:
                 parse_args += ", "
                 pass_args += ", "
+            prefix_comma = True
             if isinstance(t, list):
                 if numpy_init == "":
                     numpy_init = "// Initialize NumPy\n    import_array();\n\n    "
                 fill_array_details += f"""\n
-    // fill array details for args[{i-1}]
-    if (PyArray_NDIM(arg_{i}) != 1) {{
+    // fill array details for {i}
+    if (PyArray_NDIM({i}) != 1) {{
         PyErr_SetString(PyExc_TypeError,
             "Only 1 dimension is implemented for now.");
         return NULL;
@@ -731,9 +756,9 @@ def get_data_type(t):
     {{
         {t[2]}array;
         // Create C arrays from numpy objects:
-        PyArray_Descr *descr = PyArray_DescrFromType(PyArray_TYPE(arg_{i}));
+        PyArray_Descr *descr = PyArray_DescrFromType(PyArray_TYPE({i}));
         npy_intp dims[1];
-        if (PyArray_AsCArray((PyObject **)&arg_{i}, (void *)&array, dims, 1, descr) < 0) {{
+        if (PyArray_AsCArray((PyObject **)&{i}, (void *)&array, dims, 1, descr) < 0) {{
             PyErr_SetString(PyExc_TypeError, "error converting to c array");
             return NULL;
         }}
@@ -744,11 +769,11 @@ def get_data_type(t):
         s_array_{i}->dims[0].length = dims[0];
         s_array_{i}->is_allocated = false;
     }}"""
-                pass_args += "s_array_" + str(i)
+                pass_args += "s_array_" + i
             else:
-                pass_args += "arg_" + str(i)
-            variables_decl += "    " + get_data_type(t) + "arg_" + str(i) + ";\n"
-            parse_args += "&arg_" + str(i)
+                pass_args += i
+            variables_decl += "    " + get_data_type(t) + i + ";\n"
+            parse_args += "&" + i
 
         if parse_args != "":
             parse_args = f"""\n    // Parse the arguments from Python
@@ -761,12 +786,38 @@ def get_data_type(t):
         fill_return_details = ""
         if self.return_type != "":
             fill_return_details = f"""\n\n    // Call the C function
-    arg_0 = {self.fn_name}({pass_args});
+    _{self.fn_name}_return_value = {self.fn_name}({pass_args});
 
     // Build and return the result as a Python object
-    return Py_BuildValue("{self.return_type_format}", arg_0);"""
+    return Py_BuildValue("{self.return_type_format}", _{self.fn_name}_return_value);"""
         else:
-            fill_return_details = f"""{self.fn_name}({pass_args});
+            if self.array_as_return_type:
+                fill_return_details = f"""\n
+    _{self.fn_name}_return_value->data = malloc({self.array_as_return_type[1][3]
+        } * sizeof({self.array_as_return_type[1][2][:-2]}));
+    _{self.fn_name}_return_value->n_dims = 1;
+    _{self.fn_name}_return_value->dims[0].lower_bound = 0;
+    _{self.fn_name}_return_value->dims[0].length = {
+        self.array_as_return_type[1][3]};
+    _{self.fn_name}_return_value->is_allocated = false;
+
+    // Call the C function
+    {self.fn_name}({pass_args}, &_{self.fn_name}_return_value[0]);
+
+    // Build and return the result as a Python object
+    {{
+        npy_intp dims[] = {{{self.array_as_return_type[1][3]}}};
+        PyObject* numpy_array = PyArray_SimpleNewFromData(1, dims, {
+            get_typenum(self.array_as_return_type[1][2][:-2])},
+            _{self.fn_name}_return_value->data);
+        if (numpy_array == NULL) {{
+            PyErr_SetString(PyExc_TypeError, "error creating an array");
+            return NULL;
+        }}
+        return numpy_array;
+    }}"""
+            else:
+                fill_return_details = f"""{self.fn_name}({pass_args});
     Py_RETURN_NONE;"""
 
         # ----------------------------------------------------------------------