Aliasing in elemwise and some fixes

Author: aseyboldt

pytensor/link/numba/dispatch/elemwise.py

@@ -447,11 +447,13 @@ def _vectorize_bc(
     noalias_outputs=False,
 ):
 
-    flags = {
+    flags = True 
+    {
         "arcp",  # Allow Reciprocal
         "contract",  # Allow floating-point contraction
         "afn",  # Approximate functions
         "reassoc",
+        "nsz",  # TODO Do we want this one?
     }
 
     n_inputs = len(input_bc_patterns)
@@ -473,6 +475,9 @@ def loop_call(typingctx, *args):
         sig = types.void(types.StarArgTuple([*out_types, *in_types, iter_shape_type]))
 
         def codegen(context, builder, signature, args):
+            for i in [0]:
+                arg = builder.function.args[i]
+                arg.add_attribute("noalias")
             safe = (boundscheck, False)
             [args] = args
             args = cgutils.unpack_tuple(builder, args)
@@ -485,9 +490,10 @@ def codegen(context, builder, signature, args):
 
             # Lower the code of the scalar function so that we can use it in the inner loop
             # Caching is set to false to avoid a numba bug TODO ref?
-            inner = context.compile_subroutine(
+            inner_func = context.compile_subroutine(
                 builder, scalar_func, scalar_signature, caching=False,
-            ).fndesc
+            )
+            inner = inner_func.fndesc
 
             # Extract shape and stride information from the array.
             # For later use in the loop body to do the indexing
@@ -499,13 +505,15 @@ 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])
+            # TODO I think this is better than the noalias attribute
+            # for the input, but self_ref isn't supported in a released
+            # llvmlite version yet
+            #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)
@@ -551,8 +559,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)
+                #val.set_metadata("alias.scope", input_scope_set)
+                #val.set_metadata("noalias", output_scope_set)
                 input_vals.append(val)
 
             # Call scalar function
@@ -572,8 +580,14 @@ def extract_array(aryty, ary):
                 zip(output_accumulator, output_values, strict=True)
             ):
                 if accu is not None:
-                    new_value = builder.fadd(builder.load(accu), value)
-                    builder.store(new_value, accu)
+                    load = builder.load(accu)
+                    #load.set_metadata("alias.scope", output_scope_set)
+                    #load.set_metadata("noalias", input_scope_set)
+                    new_value = builder.fadd(load, value)
+                    store = builder.store(new_value, accu)
+                    # TODO ?
+                    #store.set_metadata("alias.scope", output_scope_set)
+                    #store.set_metadata("noalias", input_scope_set)
                 else:
                     idxs_bc = [
                         zero if bc else idx
@@ -582,9 +596,10 @@ def extract_array(aryty, ary):
                     ptr = cgutils.get_item_pointer2(
                         context, builder, *outputs[i], idxs_bc
                     )
-                    store = builder.store(value, ptr)
-                    store.set_metadata("alias.scope", output_scope_set)
-                    store.set_metadata("noalias", input_scope_set)
+                    #store = builder.store(value, ptr)
+                    store = arrayobj.store_item(context, builder, out_types[i], 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]):
@@ -599,16 +614,20 @@ def extract_array(aryty, ary):
                         ptr = cgutils.get_item_pointer2(
                             context, builder, *outputs[output], idxs_bc
                         )
-                        store = builder.store(builder.load(accu), ptr)
-                        store.set_metadata("alias.scope", output_scope_set)
-                        store.set_metadata("noalias", input_scope_set)
+                        load = builder.load(accu)
+                        #load.set_metadata("alias.scope", output_scope_set)
+                        #load.set_metadata("noalias", input_scope_set)
+                        #store = builder.store(load, ptr)
+                        store = arrayobj.store_item(context, builder, out_types[output], load, ptr)
+                        #store.set_metadata("alias.scope", output_scope_set)
+                        #store.set_metadata("noalias", input_scope_set)
                 loop.__exit__(None, None, None)
             return
 
         return sig, codegen
 
     def vectorized(*inputs):
-        pass
+        raise NotImplementedError()
 
     @numba.extending.overload(vectorized, jit_options={"fastmath": flags})
     def impl_vectorized(*inputs):
@@ -635,17 +654,32 @@ def impl_vectorized(*inputs):
 
         iter_shape_repeated = tuple([iter_shape_template[:] for _ in range(n_outputs)])
 
-        @numba.extending.register_jitable
-        def make_output(iter_shape, bc, dtype):
-            shape = iter_shape
-            for i in range(ndim):
-                if bc[i]:
-                    shape = tuple_setitem(
-                        shape,
-                        i,
-                        1,
-                    )
-            return np.empty(shape, dtype)
+        ndim_range = tuple(range(ndim))
+
+        if ndim > 0:
+            # TODO workaround for https://github.com/numba/numba/issues/8654
+            @numba.extending.register_jitable
+            def make_output(iter_shape, bc, dtype):
+                shape = iter_shape
+                for i in literal_unroll(ndim_range):
+                    if bc[i]:
+                        shape = tuple_setitem(
+                            shape,
+                            i,
+                            1,
+                        )
+                return np.empty(shape, dtype)
+
+            check_arrays = check_broadcasting
+        else:
+            @numba.extending.register_jitable
+            def make_output(iter_shape, bc, dtype):
+                return np.empty((), dtype)
+
+            @numba.extending.register_jitable
+            def check_arrays(a, b, c):
+                pass
+
 
         make_outputs = tuple_mapper(make_output)
 
@@ -667,8 +701,6 @@ 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
 
             i = 0
             for input_ in literal_unroll(inputs):
@@ -704,21 +736,24 @@ def numba_funcify_Elemwise(op, node, **kwargs):
         scalar_inputs = [scalar(dtype=input.dtype) for input in node.inputs]
         scalar_node = op.scalar_op.make_node(*scalar_inputs)
 
+    flags = True
+    {
+        "arcp",  # Allow Reciprocal
+        "contract",  # Allow floating-point contraction
+        "afn",  # Approximate functions
+        "reassoc",
+    }
+
     scalar_op_fn = numba_funcify(
-        op.scalar_op, node=scalar_node, parent_node=node, **kwargs
+        op.scalar_op, node=scalar_node, parent_node=node, fastmath=flags, **kwargs
     )
 
-    assert not op.inplace_pattern
-
-    #scalar_wrapper = register_jitable(scalar_op_fn)
-    scalar_wrapper = scalar_op_fn
-
     ndim = node.outputs[0].ndim
     output_bc_patterns = tuple([(False,) * ndim for _ in node.outputs])
     input_bc_patterns = tuple([input_var.broadcastable for input_var in node.inputs])
 
     vectorized = _vectorize_bc(
-        scalar_wrapper,
+        scalar_op_fn,
         input_bc_patterns,
         output_bc_patterns,
         output_dtypes=tuple([
@@ -727,10 +762,22 @@ def numba_funcify_Elemwise(op, node, **kwargs):
         ]),
     )
 
+    # TODO We should do this in vectorize instead
+    if op.inplace_pattern:
+        pattern = list(op.inplace_pattern.items())
+
+        @numba_njit
+        def elemwise_inplace(*inputs):
+            outputs = vectorized(*inputs)
+            for out_idx, in_idx in literal_unroll(pattern):
+                inputs[in_idx][...] = outputs[out_idx]
+    else:
+        elemwise_inplace = vectorized
+
     if len(node.outputs) == 1:
         @numba_njit
         def elemwise_wrapper(*inputs):
-            return vectorized(*inputs)[0]
+            return elemwise_inplace(*inputs)[0]
     else:
         elemwise_wrapper = vectorized