Improve cache key graph printing performance

[aorenste]

Teach the graph printer how to allow overriding printing SymTypes (SymInt, SymFloat, SymBool) and then use that to reuse the fast SymNode printing from torch._inductor.utils.sympy_str() to make computing the cache key faster.

On my computer the repro from #151823 goes from 480s -> 80s (still terrible... but better).

Fixes #151823

Stack from ghstack (oldest at bottom):

• -> Improve cache key graph printing performance #151928

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 or our office hours

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• CI workflows being skipped on PR

❌ 1 New Failure, 5 Unrelated Failures

As of commit c9e0ad7 with merge base d57bf53 ():

NEW FAILURE - The following job has failed:

• pull / linux-jammy-py3.9-gcc11 / test (backwards_compat, 1, 1, ephemeral.linux.2xlarge) (gh)
  test_modules_can_be_imported

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• inductor / cuda12.6-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, ephemeral.linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  Process completed with exit code 1.
• inductor / cuda12.6-py3.10-gcc9-sm86 / test (inductor_torchbench, 2, 2, ephemeral.linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  Process completed with exit code 1.
• pull / linux-jammy-py3.9-gcc11-mobile-lightweight-dispatch-build / build (gh) (trunk failure)
  /var/lib/jenkins/workspace/aten/src/ATen/native/BlasKernel.cpp:741:23: error: ‘bf16_dot’ is not a member of ‘at::native::blas_impl’; did you mean ‘fp16_dot’?
• trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh) (trunk failure)
  'test/dynamo/test_unittest.py::CPythonTest_Assertions::testAssertNotRegex'

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

• inductor / unit-test / cuda12.6-py3.10-gcc9-sm86 / test (inductor_cpp_wrapper, 1, 2, ephemeral.linux.g5.4xlarge.nvidia.gpu) (gh) (#152916)
  ../aten/src/ATen/native/BlasKernel.cpp:741:23: error: ‘bf16_dot’ is not a member of ‘at::native::blas_impl’; did you mean ‘fp16_dot’?

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

[laithsakka]

looks good just one comment about pushing the override logic inside print_readable to make it easier to replicate this on other places and reduce code repetition

[aorenste]

test_modules_can_be_imported looks unrelated (appears in several different PRs)
@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 5 checks: pull / linux-jammy-py3.9-gcc11-mobile-lightweight-dispatch-build / build, pull / linux-jammy-py3.9-gcc11 / test (backwards_compat, 1, 1, ephemeral.linux.2xlarge), inductor / cuda12.6-py3.10-gcc9-sm86 / test (inductor_torchbench, 2, 2, ephemeral.linux.g5.4xlarge.nvidia.gpu), inductor / cuda12.6-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, ephemeral.linux.g5.4xlarge.nvidia.gpu), inductor / unit-test / cuda12.6-py3.10-gcc9-sm86 / test (inductor_cpp_wrapper, 1, 2, ephemeral.linux.g5.4xlarge.nvidia.gpu)

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