Add fusion_type compile option for controlling the fusion logic

[IvanYashchuk]

This PR adds an option for choosing a fusion type (requested in #1763 (comment)). Currently, there are two fusion types:

1. "dataflow" - fusions are formed based on dataflow and merged horizontally and vertically,
2. "consecutive" - fusions are formed if two symbols are fusible preserving the input order (added in fuse only consecutive symbols #1608).

thunder.dynamo.thunderfx and thunder.jit have different defaults, thunderfx uses the "dataflow" fusion type (as it was before f2d7152) and jit uses the "consecutive" fusion type introduced in f2d7152.

[IvanYashchuk]

This pull request is ready for review. I'm keeping the draft mode to prevent accidental merges into the wrong base branch.

[t-vi]

Please make it a boolean option. Thank you.

[IvanYashchuk]

Please make it a boolean option. Thank you.

Sure, I can make it. How would you prefer this option to be named and what default value would you like to see?

In the future, we may need to extend the option to disable horizontal merges like it's explored in #1361. Do you want a boolean option for each fusion type with a separate check that there are no 2+ options enabled at the same time?

[t-vi]

I think we are having similar discussions for just about any option.
There are two things that make me prefer a boolean option:

• To my mind, having keyword options like this is a discoverability problem.
• I don't think it is good to have many codepaths through the compiler, so I would prefer to keep the option space small.

My concern about the algorithm you prefer to be the default with the dynamo frontend is that it appears to reoder operations without regard to memory requirements and just happens to produce a happier outcome on some models and seems to be producing very high memory requirements in some checkpointing cases for us. As such it would be preferable if the non-dynamo-frontend codepath to use the consecutive merging exclusively until it is amended by a memory-aware reordering of operations. Whether to capture this in a boolean option set by the dynamo frontend call to the jit or just hardcode based on the frontend option is OK.

I understand that your vision for the dynamo frontend is different from that, so if you deal with that however you want, it will be fine.

[IvanYashchuk]

I understand your concerns about the default choice for ThunderFX being suboptimal for certain types of computations. We are aware of the limitations.

As such it would be preferable if the non-dynamo-frontend codepath to use the consecutive merging exclusively until it is amended by a memory-aware reordering of operations.

It's done so in the current PR. thunder.jit by default uses the fusion_type="consecutive" option.

I'm sorry I don't get your request. I thought you were asking instead of "fusion_type=string" something like "enable_dataflow_fusion=True" (which is still a keyword argument to the thunder.jit call). But then you state one reason to prefer a boolean option is that keyword arguments have a discoverability problem. I thought "boolean" refers to the type of the value passed as a keyword argument to thunder.jit. Could you please clarify what change you would like to see?

[t-vi]

Let's leave it as is then.
Can you please resolve the conflicts so we can merge?

[t-vi]

Thank you @IvanYashchuk