Reuse Elemwise inplace machinery for Blockwise

Author: ricardoV94

pytensor/tensor/rewriting/blockwise.py

@@ -2,7 +2,7 @@
 from pytensor.graph import Constant, node_rewriter
 from pytensor.graph.destroyhandler import inplace_candidates
 from pytensor.graph.replace import vectorize_node
-from pytensor.graph.rewriting.basic import copy_stack_trace, in2out, out2in
+from pytensor.graph.rewriting.basic import copy_stack_trace, out2in
 from pytensor.tensor.basic import Alloc, ARange, alloc, shape_padleft
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.math import Dot
@@ -11,6 +11,7 @@
     register_specialize,
     register_stabilize,
 )
+from pytensor.tensor.rewriting.elemwise import InplaceGraphOptimizer
 from pytensor.tensor.shape import Reshape
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
@@ -260,68 +261,76 @@ def local_blockwise_of_subtensor(fgraph, node):
     return [x[(*none_slices, *core_idxs)]]
 
 
-@node_rewriter(tracks=[Blockwise], inplace=True)
-def blockwise_inplace(fgraph, node):
-    blockwise_op = node.op
+class InplaceBlockwiseOptimizer(InplaceGraphOptimizer):
+    op = Blockwise
 
-    if blockwise_op.destroy_map:
-        # Op already has inplace
-        return
-
-    # Find out valid inputs for inplacing
-    batch_ndim = blockwise_op.batch_ndim(node)
-    out_batch_bcast = node.outputs[0].type.broadcastable[:batch_ndim]
+    def filter_candidate_pairs(self, fgraph, node, protected_inputs):
+        blockwise_op = node.op
+        batch_ndim = blockwise_op.batch_ndim(node)
+        out_batch_bcast = node.outputs[0].type.broadcastable[:batch_ndim]
+        inputs = node.inputs
 
-    inputs = node.inputs
-    candidate_inputs = set(
-        inplace_candidates(
-            fgraph,
-            [
-                inp
-                for inp in inputs
-                if inp.type.broadcastable[:batch_ndim] == out_batch_bcast
-            ],
+        candidate_inputs = set(
+            inplace_candidates(
+                fgraph,
+                [
+                    inp
+                    for inp in inputs
+                    if inp.type.broadcastable[:batch_ndim] == out_batch_bcast
+                ],
+            )
         )
-    )
-    allowed_inplace_inputs = [
-        i for i, inp in enumerate(inputs) if inp in candidate_inputs
-    ]
 
-    if not allowed_inplace_inputs:
-        return None
+        allowed_inplace_inputs = [
+            i for i, inp in enumerate(inputs) if inp in candidate_inputs
+        ]
+        destroy_map = blockwise_op.core_op.inplace_on_inputs(
+            allowed_inplace_inputs=allowed_inplace_inputs
+        ).destroy_map
+
+        if not destroy_map:
+            return []
+
+        outputs = node.outputs
+        return [
+            ((out_idx, outputs[out_idx]), (inp_idx, inputs[inp_idx]))
+            for out_idx, inp_idxs in destroy_map.items()
+            for inp_idx in inp_idxs
+        ]
 
-    inplace_core_op = blockwise_op.core_op.inplace_on_inputs(
-        allowed_inplace_inputs=allowed_inplace_inputs
-    )
+    def create_inplace_node(self, node, inplace_pattern):
+        blockwise_op = node.op
+        allowed_inplace_inputs = tuple(v[0] for v in inplace_pattern.values())
+        inplace_core_op = blockwise_op.core_op.inplace_on_inputs(
+            allowed_inplace_inputs=allowed_inplace_inputs
+        )
 
-    if not inplace_core_op.destroy_map:
-        return None
+        if not inplace_core_op.destroy_map:
+            return node
 
-    # Check Op is not trying to inplace on non-candidate inputs
-    for destroyed_inputs in inplace_core_op.destroy_map.values():
-        for destroyed_input in destroyed_inputs:
-            if destroyed_input not in allowed_inplace_inputs:
-                raise ValueError(
-                    f"Op {blockwise_op.core_op} destroy_map does not respect allowed_inplace_inputs {allowed_inplace_inputs}"
-                )
+        # Check Op is not trying to inplace on non-candidate inputs
+        for destroyed_inputs in inplace_core_op.destroy_map.values():
+            for destroyed_input in destroyed_inputs:
+                if destroyed_input not in allowed_inplace_inputs:
+                    raise ValueError(
+                        f"Op {blockwise_op.core_op} destroy_map does not respect allowed_inplace_inputs {allowed_inplace_inputs}"
+                    )
 
-    # Recreate core_op with inplace
-    inplace_blockwise_op = Blockwise(
-        core_op=inplace_core_op,
-        signature=blockwise_op.signature,
-        name=blockwise_op.name,
-        gufunc_spec=blockwise_op.gufunc_spec,
-        destroy_map=inplace_core_op.destroy_map,
-    )
+        # Recreate core_op with inplace
+        inplace_blockwise_op = type(blockwise_op)(
+            core_op=inplace_core_op,
+            signature=blockwise_op.signature,
+            name=blockwise_op.name,
+            gufunc_spec=blockwise_op.gufunc_spec,
+            destroy_map=inplace_core_op.destroy_map,
+        )
 
-    out = inplace_blockwise_op.make_node(*node.inputs).outputs
-    copy_stack_trace(node.outputs, out)
-    return out
+        return inplace_blockwise_op.make_node(*node.inputs)
 
 
 optdb.register(
     "blockwise_inplace",
-    in2out(blockwise_inplace),
+    InplaceBlockwiseOptimizer(),
     "fast_run",
     "inplace",
     position=50.1,

pytensor/tensor/rewriting/elemwise.py

@@ -1,8 +1,9 @@
+import abc
 import itertools
 import operator
 import sys
 from collections import defaultdict, deque
-from collections.abc import Generator
+from collections.abc import Generator, Sequence
 from functools import cache, reduce
 from typing import TypeVar
 from warnings import warn
@@ -12,7 +13,7 @@
 from pytensor.compile.function.types import Supervisor
 from pytensor.compile.mode import get_target_language
 from pytensor.configdefaults import config
-from pytensor.graph import FunctionGraph
+from pytensor.graph import FunctionGraph, Op
 from pytensor.graph.basic import Apply, Variable, ancestors
 from pytensor.graph.destroyhandler import DestroyHandler, inplace_candidates
 from pytensor.graph.features import ReplaceValidate
@@ -47,14 +48,28 @@
 from pytensor.tensor.variable import TensorConstant, TensorVariable
 
 
-class InplaceElemwiseOptimizer(GraphRewriter):
+class InplaceGraphOptimizer(GraphRewriter):
     r"""
     This is parameterized so that it works for `Elemwise` `Op`\s.
     """
 
+    op: type[Op]
+
     def add_requirements(self, fgraph):
         fgraph.attach_feature(DestroyHandler())
 
+    @abc.abstractmethod
+    def filter_candidate_pairs(
+        self, fgraph: FunctionGraph, node: Apply, protected_inputs: Sequence[Variable]
+    ) -> Sequence[tuple[tuple[int, Variable], tuple[int, Variable]]]:
+        pass
+
+    @abc.abstractmethod
+    def create_inplace_node(
+        self, node: Apply, inplace_pattern: dict[int, Sequence[int]]
+    ) -> Apply:
+        pass
+
     def apply(self, fgraph):
         r"""
 
@@ -93,30 +108,6 @@ def apply(self, fgraph):
         # tackle them in a more general way. The whole try/except approach is probably suboptimal.
         # We can consider restricting inputs with static shapes that are large enough.
 
-        def create_inplace_node(node, inplace_pattern):
-            op = node.op
-            scalar_op = op.scalar_op
-            inplace_pattern = {i: o for i, [o] in inplace_pattern.items()}
-            if hasattr(scalar_op, "make_new_inplace"):
-                new_scalar_op = scalar_op.make_new_inplace(
-                    ps.transfer_type(
-                        *[
-                            inplace_pattern.get(i, o.dtype)
-                            for i, o in enumerate(node.outputs)
-                        ]
-                    )
-                )
-            else:
-                new_scalar_op = type(scalar_op)(
-                    ps.transfer_type(
-                        *[
-                            inplace_pattern.get(i, None)
-                            for i in range(len(node.outputs))
-                        ]
-                    )
-                )
-            return type(op)(new_scalar_op, inplace_pattern).make_node(*node.inputs)
-
         if config.tensor__insert_inplace_optimizer_validate_nb != -1:
             warn(
                 "tensor__insert_inplace_optimizer_validate_nb config is deprecated. Setting it will fail in a future release.",
@@ -140,43 +131,30 @@ def create_inplace_node(node, inplace_pattern):
         protected_inputs.update(fgraph.outputs)
         root_destroyer = fgraph.destroy_handler.root_destroyer
 
+        self_op = self.op
         update_mapping = fgraph.update_mapping or {}
         op_updates: dict[TensorVariable, TensorVariable] = {
             out: fgraph.inputs[update_mapping[out_idx]]
             for out_idx, out in enumerate(fgraph.outputs)
             if (
                 out_idx in update_mapping
                 and out.owner
-                and isinstance(out.owner.op, Elemwise)
+                and isinstance(out.owner.op, self_op)
             )
         }
         set_op_updates = set(op_updates.keys())
 
         for node in fgraph.toposort():
-            if not isinstance(node.op, Elemwise) or node.op.destroy_map:
+            if not isinstance(node.op, self_op) or node.op.destroy_map:
                 continue
 
             # If big graph and the outputs are scalar, do not make it inplace.
             if large_graph and all(node.outputs[0].type.broadcastable):
                 continue
 
-            candidate_inputs = [
-                (node.inputs.index(inp), inp)
-                for inp in inplace_candidates(
-                    fgraph,
-                    node.inputs,
-                    protected_inputs=protected_inputs,
-                )
-            ]
-            if not candidate_inputs:
-                return []
-
-            candidate_pairs = [
-                ((o, out), (i, inp))
-                for o, out in enumerate(node.outputs)
-                for i, inp in candidate_inputs
-                if inp.type == out.type
-            ]
+            candidate_pairs = self.filter_candidate_pairs(
+                fgraph, node, protected_inputs
+            )
 
             if not candidate_pairs:
                 continue
@@ -216,7 +194,7 @@ def create_inplace_node(node, inplace_pattern):
                     inplace_pattern[o] = [i]
                     tried_inputs.add(i)
 
-            inplace_node = create_inplace_node(node, inplace_pattern)
+            inplace_node = self.create_inplace_node(node, inplace_pattern)
             if inplace_node.op.destroy_map == inplace_pattern:
                 replacements = tuple(zip(node.outputs, inplace_node.outputs))
                 try:
@@ -238,7 +216,7 @@ def create_inplace_node(node, inplace_pattern):
                     inplace_pattern[o] = [i]
                     tried_inputs.add(i)
 
-                    inplace_node = create_inplace_node(node, inplace_pattern)
+                    inplace_node = self.create_inplace_node(node, inplace_pattern)
                     if inplace_node.op.destroy_map != inplace_pattern:
                         # This Op can't respect this partial inplace pattern,
                         # We assume it can't support any other cases
@@ -249,6 +227,7 @@ def create_inplace_node(node, inplace_pattern):
                             fgraph.replace_all_validate(
                                 replacements, reason="inplace_elemwise_optimizer"
                             )
+                            node = inplace_node
                             replaced = True
                             node = inplace_node
                         except InconsistencyError:
@@ -278,6 +257,50 @@ def print_summary(self, stream=sys.stdout, level=0, depth=-1):
         )
 
 
+class InplaceElemwiseOptimizer(InplaceGraphOptimizer):
+    op = Elemwise
+
+    def filter_candidate_pairs(self, fgraph, node, protected_inputs):
+        candidate_inputs = [
+            (node.inputs.index(inp), inp)
+            for inp in inplace_candidates(
+                fgraph,
+                node.inputs,
+                protected_inputs=protected_inputs,
+            )
+        ]
+        if not candidate_inputs:
+            return []
+
+        return [
+            ((o, out), (i, inp))
+            for o, out in enumerate(node.outputs)
+            for i, inp in candidate_inputs
+            if inp.type == out.type
+        ]
+
+    def create_inplace_node(self, node, inplace_pattern):
+        op = node.op
+        scalar_op = op.scalar_op
+        inplace_pattern = {i: o for i, [o] in inplace_pattern.items()}
+        if hasattr(scalar_op, "make_new_inplace"):
+            new_scalar_op = scalar_op.make_new_inplace(
+                ps.transfer_type(
+                    *[
+                        inplace_pattern.get(i, o.dtype)
+                        for i, o in enumerate(node.outputs)
+                    ]
+                )
+            )
+        else:
+            new_scalar_op = type(scalar_op)(
+                ps.transfer_type(
+                    *[inplace_pattern.get(i, None) for i in range(len(node.outputs))]
+                )
+            )
+        return type(op)(new_scalar_op, inplace_pattern).make_node(*node.inputs)
+
+
 compile.optdb.register(
     "inplace_elemwise",
     InplaceElemwiseOptimizer(),