Enable vectorization in llvm elemwise

Author: aseyboldt

pytensor/link/numba/dispatch/elemwise.py

@@ -10,7 +10,6 @@
 from numba import TypingError, types
 from numba.core import cgutils
 from numba.cpython.unsafe.tuple import tuple_setitem
-from numba.extending import register_jitable
 from numba.np import arrayobj
 from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple
 
@@ -500,6 +499,14 @@ def extract_array(aryty, ary):
                 layout = aryty.layout
                 return (data, shape, strides, layout)
 
+            mod = builder.module
+            domain = mod.add_metadata([], self_ref=True)
+            input_scope = mod.add_metadata([domain], self_ref=True)
+            output_scope = mod.add_metadata([domain], self_ref=True)
+            input_scope_set = mod.add_metadata([input_scope, output_scope])
+
+            output_scope_set = mod.add_metadata([input_scope, output_scope])
+
             inputs = [
                 extract_array(aryty, ary)
                 for aryty, ary in zip(in_types, inputs, strict=True)
@@ -544,6 +551,8 @@ def extract_array(aryty, ary):
                     context, builder, *array_info, idxs_bc, *safe
                 )
                 val = builder.load(ptr)
+                val.set_metadata("alias.scope", input_scope_set)
+                val.set_metadata("noalias", output_scope_set)
                 input_vals.append(val)
 
             # Call scalar function
@@ -573,7 +582,9 @@ def extract_array(aryty, ary):
                     ptr = cgutils.get_item_pointer2(
                         context, builder, *outputs[i], idxs_bc
                     )
-                    builder.store(value, ptr)
+                    store = builder.store(value, ptr)
+                    store.set_metadata("alias.scope", output_scope_set)
+                    store.set_metadata("noalias", input_scope_set)
 
             # Close the loops and write accumulator values to the output arrays
             for depth, loop in enumerate(loop_stack[::-1]):
@@ -588,7 +599,9 @@ def extract_array(aryty, ary):
                         ptr = cgutils.get_item_pointer2(
                             context, builder, *outputs[output], idxs_bc
                         )
-                        builder.store(builder.load(accu), ptr)
+                        store = builder.store(builder.load(accu), ptr)
+                        store.set_metadata("alias.scope", output_scope_set)
+                        store.set_metadata("noalias", input_scope_set)
                 loop.__exit__(None, None, None)
             return