[Inductor-FX] Support unbacked symbol definitions

[blaine-rister]

Problem

Inductor sometimes generates unbacked symints to handle things like mismatched branches of torch.cond. This code is represented by pytree.KeyPath, with special codegen logic to convert it to Python and C++. This was not previously supported by the FX backend.

Feature

This PR adds support for unbacked symbol declarations to the FX backend. The implementation is fairly straightforward.

1. Instead of raw Python/C++, update the wrapper codegen method to emit a new Wrapper IR line called UnbackedSymbolDefsLine. This contains all the information needed to generate the Python and C++ code.
2. Move the existing Python/C++ codegen to a private method, which is invoked by UnbackedSymbolDefsLine.codegen().
3. Implement a method to generate FX IR from unbacked symbol definitions. The implementation is based on recursive descent, consuming some keypath entries, emitting an FX IR node, and recursing to the rest of the keypath. It is conceptually identical to the existing algorithm for Python and C++, except it generates FX nodes.
4. The FX backend currently relies on size hints to generate autotuning arguments, and consequently autotuning does not support unbacked SymInts. At some point, we would like to generalize the autotuning logic to support these. But for now, simply emit a warning and skip autotuning when we see them.
5. The new test case exposed some tricky issues reconciling Triton call args with constants stored in triton_meta. This PR rewrites the relevant helper function to do this in a more principled way.

Test plan

This PR imports an existing control flow test to the FX backend's test suite. The test uses unbacked symbol definitions to handle mismatched dynamic shapes coming from torch.cond branches.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/163729

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 3ab437c with merge base 5f90e8c ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[blaine-rister]

The test case in this PR exposed a very tricky bug related to arg/kwarg handling. Although this should no longer happen with the updated method, this assert guards against that bug if by some off chance we missed something.

[blaine-rister]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here