Make CAReduce more SIMD friendly and do better allocation of output memory

Author: ricardoV94

pytensor/tensor/elemwise.py

@@ -1422,7 +1422,7 @@ def infer_shape(self, fgraph, node, shapes):
     def _c_all(self, node, name, input_names, output_names, sub):
         [inp] = node.inputs
         [out] = node.outputs
-        ndim = inp.type.ndim
+        inp_ndim = inp.type.ndim
 
         [inp_name] = input_names
         [out_name] = output_names
@@ -1454,10 +1454,10 @@ def _c_all(self, node, name, input_names, output_names, sub):
             assert var.dtype == node.outputs[0].dtype
             return var.owner.op._c_all(var.owner, name, input_names, output_names, sub)
 
-        inp_dims = list(range(ndim))
+        inp_dims = list(range(inp_ndim))
         non_reduced_dims = [i for i in inp_dims if i not in axis]
-        counter = iter(range(ndim))
-        acc_dims = ["x" if i in axis else next(counter) for i in range(ndim)]
+        counter = iter(range(inp_ndim))
+        acc_dims = ["x" if i in axis else next(counter) for i in range(inp_ndim)]
 
         sub = sub.copy()
         sub["lv0"] = inp_name
@@ -1484,7 +1484,9 @@ def _c_all(self, node, name, input_names, output_names, sub):
             cgen.make_declare(
                 [acc_dims], [out_dtype], out_sub, compute_stride_jump=False
             )
-            + cgen.make_alloc([non_reduced_dims], out_dtype, sub)
+            + cgen.make_careduce_alloc(
+                inp_name, out_name, inp_ndim, axis, out_dtype, sub["fail"]
+            )
             + cgen.make_checks(
                 [acc_dims], [out_dtype], out_sub, compute_stride_jump=False
             )
@@ -1500,7 +1502,10 @@ def _c_all(self, node, name, input_names, output_names, sub):
                 cgen.make_declare(
                     [acc_dims], [acc_dtype], acc_sub, compute_stride_jump=False
                 )
-                + cgen.make_alloc([non_reduced_dims], acc_dtype, sub)
+                + cgen.make_careduce_alloc(
+                    inp_name, acc_name, inp_ndim, axis, out_dtype, sub["fail"]
+                )
+                + cgen.make_careduce_alloc([non_reduced_dims], acc_dtype, sub)
                 + cgen.make_checks(
                     [acc_dims], [acc_dtype], acc_sub, compute_stride_jump=False
                 )
@@ -1524,8 +1529,6 @@ def _c_all(self, node, name, input_names, output_names, sub):
         elif identity is None:
             raise TypeError(f"The {self.scalar_op} does not define an identity.")
 
-        initial_value = f"{acc_name}_i = {identity};"
-
         inner_task = self.scalar_op.c_code(
             Apply(
                 self.scalar_op,
@@ -1544,28 +1547,16 @@ def _c_all(self, node, name, input_names, output_names, sub):
             sub,
         )
 
-        if out.type.ndim == 0:
-            # Simple case where everything is reduced, no need for loop ordering
-            loop = cgen.make_complete_loop_careduce(
-                inp_var=inp_name,
-                acc_var=acc_name,
-                inp_dtype=inp_dtype,
-                acc_dtype=acc_dtype,
-                initial_value=initial_value,
-                inner_task=inner_task,
-                fail_code=sub["fail"],
-            )
-        else:
-            loop = cgen.make_reordered_loop_careduce(
-                inp_var=inp_name,
-                acc_var=acc_name,
-                inp_dtype=inp_dtype,
-                acc_dtype=acc_dtype,
-                inp_ndim=ndim,
-                reduction_axes=axis,
-                initial_value=initial_value,
-                inner_task=inner_task,
-            )
+        loop = cgen.make_reordered_loop_careduce(
+            inp_var=inp_name,
+            acc_var=acc_name,
+            inp_dtype=inp_dtype,
+            acc_dtype=acc_dtype,
+            inp_ndim=inp_ndim,
+            reduction_axes=axis,
+            initial_value=identity,
+            inner_task=inner_task,
+        )
 
         if acc_dtype != out_dtype:
             cast = dedent(
@@ -1589,7 +1580,7 @@ def c_headers(self, **kwargs):
 
     def c_code_cache_version_apply(self, node):
         # the version corresponding to the c code in this Op
-        version = [10]
+        version = [11]
 
         # now we insert versions for the ops on which we depend...
         scalar_node = Apply(