[Relax][PyTorch] Add run_ep_decomposition flag to control PyTorch decomposition

[tlopex]

This PR adds a new run_ep_decomposition flag to the from_exported_program function to control whether PyTorch's decomposition should be run before translation. This enables gradual migration from non-decomposed to decomposed operator support. This flag will be a temporary one and will be removed until everything is fixed.

Motivation

Currently, when run_decompositions() is called on an ExportedProgram, high-level operators are decomposed into their constituent parts (e.g., torch.square → torch.pow(x, 2)). However, the current translator expects the original high-level operators, causing ~40% of tests to fail when decomposition is enabled. Previously, we couldn't enable decomposition because run_decompositions() is not an in-place operation so that it cannot be used.

This flag allows us to:

1. Keep existing tests passing (default behavior)
2. Gradually add support for decomposed operators
3. Eventually migrate all tests to use decomposition
4. Remove the flag once migration is complete

Future Work

This PR sets up the foundation for:

• Adding support for decomposed operators (e.g., pow.Tensor_Scalar, mul.Tensor)
• Gradually migrating tests to use decomposition
• Eventually removing the flag once migration is complete

Example Usage

# Current behavior (no decomposition)
mod = from_exported_program(exported_program)

# New behavior (with decomposition)
mod = from_exported_program(exported_program, run_ep_decomposition=True)

[gemini-code-assist]

Summary of Changes

Hello @tlopex, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a crucial, temporary flag to the from_exported_program function, enabling developers to control PyTorch's operator decomposition process. This change is designed to allow for a phased transition to supporting decomposed operators, which is essential for future optimizations and broader compatibility. By making decomposition optional, the PR ensures existing functionality remains stable while laying the groundwork for incremental improvements and eventual full adoption of decomposed operations.

Highlights

• New Flag for PyTorch Decomposition: Introduced a new run_ep_decomposition boolean flag to the from_exported_program function, allowing explicit control over whether PyTorch's operator decomposition is applied.
• Enabling Gradual Migration: This flag facilitates a gradual migration path from non-decomposed to decomposed operator support, addressing current translator limitations where high-level operators are expected.
• Temporary Measure: The run_ep_decomposition flag is temporary and will be removed once full migration to decomposed operators is complete and stable.
• Conditional Decomposition: The exported_program.run_decompositions() call is now conditional, executing only when run_ep_decomposition is set to True, ensuring backward compatibility by defaulting to False.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a run_ep_decomposition flag to control PyTorch operator decomposition, which is a good approach for enabling a gradual migration. The change also correctly handles the return value of run_decompositions(), fixing a subtle bug in the previous implementation. My review includes suggestions to add a TODO for the temporary flag and to include tests for the new code path to ensure its correctness.

[gemini-code-assist]

Since this flag is temporary, it would be good practice to add a TODO comment here to track its future removal. This helps ensure the codebase is cleaned up after the migration is complete.

# TODO(user): Remove this flag and always run decomposition once all operators are supported.
if run_ep_decomposition:
    exported_program = exported_program.run_decompositions()

It would also be helpful to note the temporary nature of this flag in its docstring.

[tlopex]

@tvm-bot rerun

[tlopex]

cc @mshr-h