Faster C-implementation of Join

Reuse outputs when possible, and avoid numpy overhead on public facing function

Author: ricardoV94

pytensor/tensor/basic.py

@@ -2537,58 +2537,123 @@ def perform(self, node, inputs, output_storage):
         )
 
     def c_code_cache_version(self):
-        return (5,)
+        return (6,)
 
     def c_code(self, node, name, inputs, outputs, sub):
-        axis, tens = inputs[0], inputs[1:]
-        view = -1
-        non_empty_tensor = tens[view]
-        input_1 = tens[0]
-        l = len(tens)
-        (out,) = outputs
-        fail = sub["fail"]
-        adtype = node.inputs[0].type.dtype_specs()[1]
+        axis, *arrays = inputs
+        [out] = outputs
 
-        copy_to_list = (
-            f"""Py_INCREF({inp}); PyList_SetItem(list, {i}, (PyObject*){inp});"""
-            for i, inp in enumerate(tens)
-        )
+        n = len(arrays)
+        out_dtype = node.outputs[0].type.dtype_specs()[2]
+        ndim = node.outputs[0].type.ndim
+        fail = sub["fail"]
 
-        copy_inputs_to_list = "\n".join(copy_to_list)
-        n = len(tens)
+        # Most times axis is constant, inline it
+        # This is safe to do because the hash of the c_code includes the constant signature
+        if isinstance(node.inputs[0], Constant):
+            static_axis = int(node.inputs[0].data)
+            static_axis = normalize_axis_index(static_axis, ndim)
+            axis_def = f"{static_axis};"
+            axis_check = ""
+        else:
+            axis_dtype = node.inputs[0].type.dtype_specs()[1]
+            axis_def = f"(({axis_dtype} *)PyArray_DATA({axis}))[0];"
+            axis_check = f"""
+                if (axis < 0){{
+                    axis = {ndim} + axis;
+                }}
+                if (axis >= {ndim} || axis < 0) {{
+                    PyErr_SetString(PyExc_ValueError, "Join axis is out of bounds");
+                    {fail}
+                }}
+            """
 
         code = f"""
-        int axis = (({adtype} *)PyArray_DATA({axis}))[0];
-        PyObject* list = PyList_New({l});
-        {copy_inputs_to_list}
-        int tensors_lens_sum;
-        if({view} != -1) {{
-            tensors_lens_sum = 0;
-
-            for(int i=0; i < {n}; i++){{
-                tensors_lens_sum += PyArray_DIM((PyArrayObject *)(PyList_GetItem(list, i)), axis);
+            int axis = {axis_def}
+            PyArrayObject* arrays[{n}] = {{{','.join(arrays)}}};
+            int out_is_valid = 0;
+            npy_intp join_size = 0;
+            npy_intp offset = 0;
+
+            // Validate input shapes and compute join size
+            npy_intp *shape = PyArray_SHAPE(arrays[0]);
+
+            {axis_check}
+
+            for (int i = 0; i < {n}; i++) {{
+                if (PyArray_NDIM(arrays[i]) != {ndim}) {{
+                    PyErr_SetString(PyExc_ValueError, "Input to join has wrong ndim");
+                    {fail}
+                }}
+
+                for (int j = 0; j < {ndim}; j++) {{
+                    if (j == axis){{
+                        join_size += PyArray_DIM(arrays[i], j);
+                    }}
+                    else if(PyArray_DIM(arrays[i], j) != shape[j]) {{
+                        PyErr_SetString(PyExc_ValueError, "Arrays shape must match along non join axis");
+                        {fail}
+                    }}
+                }}
             }}
-            tensors_lens_sum -= PyArray_DIM({non_empty_tensor}, axis);
-        }}
-        if({view} != -1 && tensors_lens_sum == 0) {{
-            Py_XDECREF({out});
-            Py_INCREF({non_empty_tensor});
-            {out} = {non_empty_tensor};
-        }}else{{
-            //PyObject* PyArray_Concatenate(PyObject* obj, int axis)
-            int ndim = PyArray_NDIM({input_1});
-            if( axis < -ndim ){{
-                PyErr_Format(PyExc_IndexError,
-                             "Join axis %d out of bounds [0, %d)", axis, ndim);
-                {fail}
+
+            // Define dimensions of output array
+            npy_intp out_dims[{ndim}];
+            memcpy(out_dims, shape, {ndim} * sizeof(npy_intp));
+            out_dims[axis] = join_size;
+
+            // Reuse output or allocate new one
+            if ({out} != NULL) {{
+                out_is_valid = (PyArray_NDIM({out}) == {ndim});
+                for (int i = 0; i < {ndim}; i++) {{
+                    out_is_valid &= (PyArray_DIM({out}, i) == out_dims[i]);
+                }}
             }}
-            Py_XDECREF({out});
-            {out} = (PyArrayObject *)PyArray_Concatenate(list, axis);
-            Py_DECREF(list);
-            if(!{out}){{
-                {fail}
+
+            if (!out_is_valid) {{
+                Py_XDECREF({out});
+                {out} = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type,
+                                                             PyArray_DescrFromType({out_dtype}),
+                                                             {ndim},
+                                                             out_dims,
+                                                             NULL,  /* strides */
+                                                             NULL,  /* data */
+                                                             NPY_ARRAY_DEFAULT,
+                                                             NULL);
+
+                if ({out} == NULL) {{
+                    {fail}
+                }}
+            }}
+
+            // Copy data into output buffer
+            for (int i = 0; i < {n}; i++) {{
+                PyArrayObject *arr = arrays[i];
+
+                // Create view into output buffer
+                // PyArray_NewFromDescr steals a reference to descr, so we need to increase it
+                Py_INCREF(PyArray_DESCR({out}));
+                PyArrayObject *view = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type,
+                                                                          PyArray_DESCR({out}),
+                                                                          {ndim},
+                                                                          PyArray_SHAPE(arr),
+                                                                          PyArray_STRIDES({out}),
+                                                                          PyArray_BYTES({out}) + (offset * PyArray_STRIDES({out})[axis]),
+                                                                          NPY_ARRAY_WRITEABLE,
+                                                                          NULL);
+                if (view == NULL) {{
+                    {fail}
+                }}
+
+                // Write to it
+                int success = PyArray_CopyInto(view, arr);
+                Py_DECREF(view);
+                if (success != 0) {{
+                    {fail}
+                }}
+
+                offset += PyArray_DIM(arr, axis);
             }}
-        }}
         """
         return code