Add a fusion rewrite for CAReduces with Elemwise inputs

Author: brandonwillard

aesara/tensor/rewriting/elemwise.py

@@ -11,11 +11,16 @@
 from aesara.graph.basic import Apply, Constant, io_toposort
 from aesara.graph.features import ReplaceValidate
 from aesara.graph.op import compute_test_value, get_test_value
-from aesara.graph.rewriting.basic import GraphRewriter, copy_stack_trace, node_rewriter
+from aesara.graph.rewriting.basic import (
+    GraphRewriter,
+    copy_stack_trace,
+    in2out,
+    node_rewriter,
+)
 from aesara.graph.rewriting.db import SequenceDB
 from aesara.graph.utils import InconsistencyError, MethodNotDefined, TestValueError
 from aesara.tensor.basic import MakeVector, alloc, cast, get_scalar_constant_value
-from aesara.tensor.elemwise import DimShuffle, Elemwise
+from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from aesara.tensor.exceptions import NotScalarConstantError
 from aesara.tensor.rewriting.basic import register_canonicalize, register_specialize
 from aesara.tensor.shape import shape_padleft
@@ -944,3 +949,54 @@ def local_useless_composite(fgraph, node):
         c = aes.Composite(inputs=comp.inputs, outputs=new_outputs)
         e = Elemwise(scalar_op=c)(*node.inputs, return_list=True)
         return dict(zip([node.outputs[i] for i in idx], e))
+
+
+@node_rewriter([CAReduce])
+def local_careduce_fusion(fgraph, node):
+    """Fuse a `CAReduce` applied to an `Elemwise`."""
+
+    (car_input,) = node.inputs
+    elm_node = car_input.owner
+
+    if elm_node is None or not isinstance(elm_node.op, Elemwise):
+        return False
+
+    elm_inputs = elm_node.inputs
+
+    if len(elm_inputs) > 1:
+        # TODO: Implement the multiple inputs case
+        raise False
+
+    # TODO: What about the `dtype`s and other properties in `CAReduceDtype` types?
+    car_axis = node.op.axis
+    car_scalar_op = node.op.scalar_op
+    elm_scalar_op = elm_node.op.scalar_op
+
+    scalar_elm_inputs = [
+        aes.get_scalar_type(inp.type.dtype).make_variable() for inp in elm_inputs
+    ]
+    elm_output = elm_scalar_op(*scalar_elm_inputs)
+    # This input represents the previous value in the `CAReduce` binary reduction
+    carried_car_input = elm_output.type()
+    scalar_fused_outputs = [car_scalar_op(elm_output, carried_car_input)]
+
+    fused_scalar_op = aes.Composite(
+        inputs=[carried_car_input] + scalar_elm_inputs, outputs=scalar_fused_outputs
+    )
+    # The fused `Op` needs to look and behave like a `BinaryScalarOp`
+    # TODO: Generate a new `type` and make this relationship official?
+    fused_scalar_op.identity = car_scalar_op.identity
+    fused_scalar_op.nin = 2
+    fused_scalar_op.nout = 1
+
+    new_car_op = CAReduce(fused_scalar_op, car_axis)
+
+    return [new_car_op(*elm_inputs)]
+
+
+compile.optdb.register(  # type: ignore
+    "local_careduce_fusion",
+    in2out(local_careduce_fusion),
+    # "fusion",
+    position=49,
+)

tests/tensor/rewriting/test_elemwise.py

@@ -1105,6 +1105,57 @@ def test_test_values(self, test_value):
                 f.maker.fgraph.outputs[0].tag.test_value, np.c_[[2.0]]
             )
 
+    def test_CAReduce_single_input(self):
+        """Make sure that `CAReduce` and `Elemwise` fusions work with a single input."""
+
+        mode = Mode(linker="cvm")
+        mode._optimizer = mode._optimizer.including(
+            "local_careduce_fusion",
+            "canonicalize",
+            "inplace",
+        )
+
+        x = matrix("x")
+        out = exp(x).sum(axis=1)
+
+        out_fn = function([x], out, mode=mode)
+        (out_node,) = out_fn.maker.fgraph.toposort()
+
+        assert isinstance(getattr(out_node.op, "scalar_op"), aes.basic.Composite)
+
+        rng = np.random.default_rng(29320)
+        x_val = rng.random((4, 3), dtype=config.floatX)
+        exp_res = np.exp(x_val).sum(axis=1)
+        out_val = out_fn(x_val)
+        assert np.array_equal(out_val, exp_res)
+
+    @pytest.mark.xfail(reason="Not implemented")
+    def test_CAReduce_multiple_inputs(self):
+        """Make sure that `CAReduce` and `Elemwise` fusions work with multiple inputs."""
+
+        mode = Mode(linker="cvm")
+        mode._optimizer = mode._optimizer.including(
+            "local_careduce_fusion",
+            "canonicalize",
+            "inplace",
+        )
+
+        x = vector("x")
+        y = vector("y")
+        out = (x + y).sum()
+
+        out_fn = function([x, y], out, mode=mode)
+        (out_node,) = out_fn.maker.fgraph.toposort()
+
+        assert isinstance(getattr(out_node.op, "scalar_op"), aes.basic.Composite)
+
+        rng = np.random.default_rng(29320)
+        x_val = rng.random((4, 3), dtype=config.floatX)
+        y_val = rng.random((4, 3), dtype=config.floatX)
+        exp_res = (x_val + y_val).sum(axis=1)
+        out_val = out_fn(x_val, y_val)
+        assert np.array_equal(out_val, exp_res)
+
 
 class TimesN(aes.basic.UnaryScalarOp):
     """