[Minimizer] for sequential mode, respect find_all setting (#134339)

Summary: Currently, for sequential mode, minimizer search terminates after a node is excluded via the user defined exclusion_fn. However, on some occasions we would like the search to continue past that for the remaining nodes. In this diff I am changing the termination criteria to respect the find_all setting, where we continue sequential search if it is set.

Test Plan: CI

Differential Revision: D61720262

Pull Request resolved: https://github.com/pytorch/pytorch/pull/134339
Approved by: https://github.com/jfix71
This commit is contained in:
Qiaochu Yuan 2024-08-23 19:59:43 +00:00 committed by PyTorch MergeBot
parent 58e2cf364b
commit 2ca7f0fc5c

View file

@ -490,7 +490,10 @@ class _MinimizerBase:
if len(node_list) == 0:
report.append(f"User exclusion : {node.name}")
self.print_report(report)
return culprits
if not self.settings.find_all:
return culprits
else:
continue
cur_nodes: NodeSet = {node}