frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Inductor] do loop reordering in a separate final round (#162355)

Browse Source 
Previous LOAF after fusion algorithm is not guaranteed to create more fusion opportunities even if loop reordering happens. I can not find an example that LOAF reduce the amount of fusion, but here is an example that reordering loops does not add more fusions:

a1f7639922/test/inductor/test_loop_ordering.py (L612-L641)

Move LOAF to a separate final round of fusion so that we are guaranteed to not reducing the amount of fusions. Hopefully this also helps compilation time since LOAF kicks in when there are less nodes.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162355
Approved by: https://github.com/eellison, https://github.com/jansel
ghstack dependencies: #162101, #162126
This commit is contained in:
Shunting Zhang
2025-09-10 18:10:42 -07:00
committed by PyTorch MergeBot
parent e88460f453
commit 248156ed06
2 changed files with 55 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
torch/_inductor/scheduler.py
View File

@ -2978,7 +2978,7 @@ class Scheduler:
i + 1,
old_len,
)
nodes = self.fuse_nodes_once(nodes)
nodes = self.fuse_nodes_once(nodes, is_reorder_round=False)
new_len = len(nodes)
fusion_log.debug(
"completed fusion round (%d/10): fused %d nodes into %d nodes\n",
@ -2991,6 +2991,9 @@ class Scheduler:
"===== fusion complete (%d iterations) =====", i + 1
)
break
if config.loop_ordering_after_fusion:
nodes = self.fuse_nodes_once(nodes, is_reorder_round=True)
return nodes
def process_grouped_nodes(self) -> None:
@ -3497,7 +3500,9 @@ class Scheduler:
return self.name_to_fused_node[node.get_first_name()]
def fuse_nodes_once(
self, nodes: list[BaseSchedulerNode]
self,
nodes: list[BaseSchedulerNode],
is_reorder_round: bool,
) -> list[BaseSchedulerNode]:
"""
Combine eligible nodes into FusedSchedulerNodes.
@ -3560,7 +3565,7 @@ class Scheduler:
fuse_two_nodes(node_key1, node_key2)
for node1, node2 in self.get_possible_fusions(nodes):
for node1, node2 in self.get_possible_fusions(nodes, is_reorder_round):
# if either node is in a pending fusion, resolve it.
# since we iterate on potential fusions based on profitability
# the first potential fusion should take precedence.
@ -3568,9 +3573,9 @@ class Scheduler:
node1 = self.get_fused_node(node1)
node2 = self.get_fused_node(node2)
if self.can_fuse(node1, node2) and not self.will_fusion_create_cycle(
node1, node2
):
if self.can_fuse(
node1, node2, is_reorder_round
) and not self.will_fusion_create_cycle(node1, node2):
speedup = self.speedup_by_fusion(node1, node2)
if callable(speedup):
pending_fusions[node1] = (speedup, node1, node2)
@ -3655,7 +3660,9 @@ class Scheduler:
node.prune_redundant_deps(self.name_to_fused_node)
def get_possible_fusions(
self, nodes: list[BaseSchedulerNode]
self,
nodes: list[BaseSchedulerNode],
is_reorder_round: bool,
) -> list[tuple[BaseSchedulerNode, BaseSchedulerNode]]:
"""
Helper to find all legal fusion opportunities, sorted by self.score_fusion()
@ -3675,10 +3682,10 @@ class Scheduler:
continue
seen.add(key)
if self.can_fuse(node1, node2):
if self.can_fuse(node1, node2, is_reorder_round):
possible_fusions.append(key)
elif (node2.is_template() or node2.is_foreach()) and self.can_fuse(
node2, node1
node2, node1, is_reorder_round
):
# foreach fusions and epilogue fusions are order dependent
possible_fusions.append((node2, node1))
@ -4148,7 +4155,12 @@ class Scheduler:
else:
return None
def can_fuse(self, node1: BaseSchedulerNode, node2: BaseSchedulerNode) -> bool:
def can_fuse(
self,
node1: BaseSchedulerNode,
node2: BaseSchedulerNode,
can_reorder: bool = False,
) -> bool:
"""
Determine if it is possible to combine node1 and node2 into a
single fused node.
@ -4265,7 +4277,8 @@ class Scheduler:
shared_data_score = self.score_fusion_memory(node1, node2)
if (
shared_data_score < config.score_fusion_memory_threshold
can_reorder
and shared_data_score < config.score_fusion_memory_threshold
and config.loop_ordering_after_fusion
):
new_shared_data_score = self.shared_data_after_reordering_loop(node1, node2)