Skip to content

Commit

Permalink
Inside a Map there can not be a library node for fusion.
Browse files Browse the repository at this point in the history
  • Loading branch information
philip-paul-mueller committed Jul 31, 2024
1 parent bcfbd68 commit 03f4b1a
Showing 1 changed file with 7 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -477,7 +477,7 @@ def partition_first_outputs(
# Certain nodes need more than one element as input. As explained
# above, in this situation we assume that we can naturally decompose
# them iff the node does not consume that whole intermediate.
# Furthermore, it can not be a dynamic map range.
# Furthermore, it can not be a dynamic map range or a library node.
intermediate_size = functools.reduce(lambda a, b: a * b, intermediate_desc.shape)
consumers = util.find_downstream_consumers(state=state, begin=intermediate_node)
for consumer_node, feed_edge in consumers:
Expand All @@ -488,6 +488,9 @@ def partition_first_outputs(
return None
if consumer_node is map_entry_2: # Dynamic map range.
return None
if isinstance(consumer_node, nodes.LibraryNode):
# TODO(phimuell): Allow some library nodes.
return None

# Note that "remove" has a special meaning here, regardless of the
# output of the check function, from within the second map we remove
Expand Down Expand Up @@ -520,6 +523,9 @@ def partition_first_outputs(
return None
if consumer_node is map_entry_2: # Dynamic map range
return None
if isinstance(consumer_node, nodes.LibraryNode):
# TODO(phimuell): Allow some library nodes.
return None
else:
# Ensure that there is no path that leads to the second map.
after_intermdiate_node = util.all_nodes_between(
Expand Down

0 comments on commit 03f4b1a

Please sign in to comment.