Try helper classes with bitset (reduce dict access)

Author: ricardoV94

pytensor/tensor/rewriting/elemwise.py

@@ -675,7 +675,9 @@ def push(self, node: Apply, node_bitflag: int, is_ancestor: bool):
 
                 def pop(self) -> tuple[Apply, int, bool]:
                     node_bitflag, node = heappop(self.queue)
-                    return node, node_bitflag < 0
+                    if node_bitflag < 0:
+                        return node, -node_bitflag, True
+                    return node, node_bitflag, False
 
                 def __bool__(self):
                     return bool(self.queue)
@@ -685,7 +687,6 @@ class NonConvexError(Exception):
 
             class ConvexSubgraph:
                 __slots__ = (
-                    "nodes_bitflags",
                     "ancestors_bitset",
                     "nodes",
                     "nodes_bitset",
@@ -694,8 +695,7 @@ class ConvexSubgraph:
                     "inputs_and_outputs",
                 )
 
-                def __init__(self, nodes_bitflags, ancestors_bitset):
-                    self.nodes_bitflags = nodes_bitflags
+                def __init__(self, ancestors_bitset):
                     self.ancestors_bitset = ancestors_bitset
                     self.nodes = {}
                     self.nodes_bitset = 0
@@ -706,11 +706,10 @@ def __init__(self, nodes_bitflags, ancestors_bitset):
                 def __len__(self):
                     return len(self.nodes)
 
-                def __contains__(self, node: int):
-                    return bool(self.nodes_bitset & self.nodes_bitflags[node])
+                def __contains__(self, node_bitflag: int):
+                    return bool(self.nodes_bitset & node_bitflag)
 
-                def add_node(self, node: Apply, is_ancestor: bool):
-                    node_bitflag = self.nodes_bitflags[node]
+                def add_node(self, node: Apply, node_bitflag, is_ancestor: bool):
                     if is_ancestor:
                         if node_bitflag & self.unfuseable_ancestors_bitset:
                             raise NonConvexError
@@ -720,14 +719,14 @@ def add_node(self, node: Apply, is_ancestor: bool):
                     self.nodes[node] = None
                     self.inputs_and_outputs = None  # clear cache
 
-                def add_unfuseable_ancestor(self, ancestor: Apply):
-                    # If an ancestor is unfuseable, so are all its ancestors
-                    self.unfuseable_ancestors_bitset |= self.ancestors_bitset[ancestor]
-
-                def add_unfuseable_client(self, client: Apply):
-                    # If a client is unfuseable, so are all its clients, but we don't need to keep track of those
-                    # Any downstream client will also depend on this unfuseable client and will be rejected when visited
-                    self.unfuseable_clients_bitset |= self.nodes_bitflags[client]
+                # def add_unfuseable_ancestor(self, ancestor: Apply):
+                #     # If an ancestor is unfuseable, so are all its ancestors
+                #     self.unfuseable_ancestors_bitset |= self.ancestors_bitset[ancestor]
+                #
+                # def add_unfuseable_client(self, client: Apply):
+                #     # If a client is unfuseable, so are all its clients, but we don't need to keep track of those
+                #     # Any downstream client will also depend on this unfuseable client and will be rejected when visited
+                #     self.unfuseable_clients_bitset |= self.nodes_bitflags[client]
 
                 def get_inputs_and_outputs(self):
                     if self.inputs_and_outputs is not None:
@@ -842,20 +841,20 @@ def __iter__(self):
                 ):
                     continue
 
-                subgraph = ConvexSubgraph(nodes_bitflags, ancestors_bitset)
+                subgraph = ConvexSubgraph(ancestors_bitset)
 
                 fuseable_nodes_queue = SortedFuseableNodesQueue()
                 fuseable_nodes_queue.push(
                     starting_node, starting_bitflag, is_ancestor=True
                 )
                 while fuseable_nodes_queue:
-                    node, is_ancestor = fuseable_nodes_queue.pop()
+                    node, node_bitflag, is_ancestor = fuseable_nodes_queue.pop()
 
-                    if node in subgraph:
+                    if node_bitflag in subgraph:
                         continue
 
                     try:
-                        subgraph.add_node(node, is_ancestor=is_ancestor)
+                        subgraph.add_node(node, node_bitflag, is_ancestor=is_ancestor)
                     except NonConvexError:
                         continue
 
@@ -866,29 +865,33 @@ def __iter__(self):
                     #  - unfuseable (add to respective unfuseable bitset)
                     for ancestor in node.inputs:
                         ancestor_node = ancestor.owner
-                        if ancestor_node in subgraph:
+                        ancestor_bitset = nodes_bitflags[ancestor_node]
+                        if ancestor_bitset in subgraph:
                             continue
                         if node in fuseable_clients[ancestor_node]:
                             fuseable_nodes_queue.push(
                                 ancestor_node,
-                                nodes_bitflags[ancestor_node],
+                                ancestor_bitset,
                                 is_ancestor=True,
                             )
                         else:
-                            subgraph.add_unfuseable_ancestor(ancestor_node)
+                            subgraph.unfuseable_ancestors_bitset |= ancestors_bitset[
+                                ancestor_node
+                            ]
 
                     next_fuseable_clients = fuseable_clients[node]
                     for client_node, _ in fg_clients[node.outputs[0]]:
-                        if client_node in subgraph:
+                        client_bitflag = nodes_bitflags[client_node]
+                        if client_bitflag in subgraph:
                             continue
                         if client_node in next_fuseable_clients:
                             fuseable_nodes_queue.push(
                                 client_node,
-                                nodes_bitflags[client_node],
+                                client_bitflag,
                                 is_ancestor=False,
                             )
                         else:
-                            subgraph.add_unfuseable_client(client_node)
+                            subgraph.unfuseable_clients_bitset |= client_bitflag
 
                 # Finished exploring this subgraph
                 if len(subgraph) == 1: