Improve the speed of the pointwise fusion graph pass

[ptrendx]

Description

Fixes #17105

This PR fixes 2 problems (that together lead to the huge time spent doing fusion reported in #17105).

The first problem was the get_subsets function - a part of the pointwise fusion graph pass. It involves a step that ensures there are no cycles created in a graph after the fusion step. It does so in 2 steps:

• first it looks at each incompatible node and separates its inputs and outputs into disjoint sets
• then for each node n it looks at those separations and creates a set of nodes that are incompatible to be in the same fusion with n

Both of those steps were improved:

• the first step was optimized via observation that if all inputs to an incompatible node n are also incompatible, then the separation set produced by that node is strictly smaller than the ones produced by those inputs, so there is no need to create it
• the second step was greatly optimized by the observation, that when considering which nodes are incompatible for fusion with a given node n, if the incompatible node producing the separation set was already part of the set of nodes n is incompatible with, then also every node in that separation set was already included there as well. In the repro script from the issue, this change cut the number of insertion trials in the forward fusion pass from ~1M to ~2000.

The other problem fixed in this PR is that due to variable scope of state_ptr in CachedOp::Forward method, the state of CachedOp was taken twice for any context, resulting in 2 calls into graph optimization methods instead of 1.

In my local environment, the time spent on fusion pass in the repro script came down from over 200s to 4.5s on a single GPU.

@leezu @zburning @Caenorst @samskalicky FYI

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• Changes are complete (i.e. I finished coding on this PR)
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Improved the speed of finding the subsets compatible with fusion
• Fixed a bug in CachedOp::Forward method where the shared_ptr to CachedOpState was taken twice, resulting in launching the optimization step twice.

[ptrendx]

@zburning Were you able to validate that this fixes the issue you reported?

[samskalicky]

LGTM, thanks for the quick fix!

[zhreshold]

looks good

[zburning]

@ptrendx Yes! I test it on a p3 instance, the first batch costs 11.4s for 4 gpus. Thank you for your quick fix.