Unroll some loops in llvm-elemwise

Author: aseyboldt

pytensor/link/numba/dispatch/elemwise.py

@@ -7,7 +7,7 @@
 import numba
 import numpy as np
 from llvmlite import ir
-from numba import TypingError, types
+from numba import TypingError, literal_unroll, types
 from numba.core import cgutils
 from numba.cpython.unsafe.tuple import tuple_setitem
 from numba.np import arrayobj
@@ -653,8 +653,8 @@ def impl(*inputs):
             iter_shape = iter_shape_template
             for i in range(ndim):
                 maxval = 1
-                for j in range(n_inputs):
-                    maxval = max(maxval, inputs[j].shape[i])
+                for inp in literal_unroll(inputs):
+                    maxval = max(maxval, inp.shape[i])
 
                 iter_shape = tuple_setitem(iter_shape, i, maxval)
 
@@ -667,12 +667,20 @@ def impl(*inputs):
                 )
 
             outputs = make_outputs(iter_shape_rep, output_bc_patterns, output_dtypes)
+            #outputs = (np.empty(inputs[0].shape),)
+            #iter_shape = inputs[0].shape
 
-            for input_, bcs in zip(inputs, input_bc_patterns):
+            i = 0
+            for input_ in literal_unroll(inputs):
+                bcs = input_bc_patterns[i]
                 check_broadcasting(input_, bcs, iter_shape)
+                i = i + 1
 
-            for out, bcs in zip(outputs, output_bc_patterns):
+            i = 0
+            for out in literal_unroll(outputs):
+                bcs = output_bc_patterns[i]
                 check_broadcasting(out, bcs, iter_shape)
+                i = i + 1
 
             loop_call(*outputs, *inputs, iter_shape)
             return outputs