[hierarchical-compilation][dynamo] Support invoke_subgraph HOP

[anijain2305]

Stack from ghstack (oldest at bottom):

• [hierarchical-compilation][inductor] Support invoke_subgraph HOP #138031
• [hierarchical-compilation][invoke-subgraph] Add fake tensor caching #137808
• [hierarchical-compilation][invoke_subgraph] Use tracing context to cache artifacts of dispatch keys #137965
• [hierarchical-compilation][invoke_subgraph] Graph break on input mutation or aliasing #138036
• -> [hierarchical-compilation][dynamo] Support invoke_subgraph HOP #137557
• [hierarchical-compilation][hop] Introduce invoke_subgraph #137538
• [inductor][subgraph] Add size asserts #138424
• [dynamo][refactor] Refactor Wrap HOP to reuse it for invoke_subgraph #137555

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @rec

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❌ 3 New Failures, 20 Cancelled Jobs, 3 Unrelated Failures

As of commit 5707e4b with merge base 239ad73 ():

NEW FAILURES - The following jobs have failed:

• pull / cuda12.1-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu) (gh)
  WIN: benchmark ('basic_modules_ListOfLinears_inductor_gpu_force_shape_pad', 'compile_time_instruction_count') failed, actual result 15552666084 is -1.60% lower than expected 15806042948 ±1.50% please update the expected results.
• pull / linux-focal-py3.11-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  higher_order_ops/test_invoke_subgraph.py::TestInvokeSubgraphCompile::test_multiple
• pull / linux-focal-py3.9-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  higher_order_ops/test_invoke_subgraph.py::TestInvokeSubgraphCompile::test_multiple

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_timm, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor / cuda12.1-py3.12-gcc9-sm86 / test (inductor, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor / cuda12.1-py3.12-gcc9-sm86 / test (inductor, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx2, 1, 2, linux.10xlarge.avx2) (gh)
  ##[error]The operation was canceled.
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx2, 2, 2, linux.10xlarge.avx2) (gh)
  ##[error]The operation was canceled.
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx512, 2, 2, linux.12xlarge) (gh)
  ##[error]The operation was canceled.
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  ##[error]The operation was canceled.
• linux-binary-manywheel / manywheel-py3_9-cuda11_8-build / build (gh)
  ##[error]The operation was canceled.
• linux-binary-manywheel / manywheel-py3_9-cuda11_8-test (gh)
• linux-binary-manywheel / manywheel-py3_9-cuda12_1-build / build (gh)
  ##[error]The operation was canceled.
• linux-binary-manywheel / manywheel-py3_9-cuda12_1-test (gh)
• linux-binary-manywheel / manywheel-py3_9-cuda12_4-build / build (gh)
  ##[error]The operation was canceled.
• linux-binary-manywheel / manywheel-py3_9-cuda12_4-test (gh)
• pull / linux-focal-py3.12-clang10-experimental-split-build / test (default, 1, 3, linux.4xlarge) (gh)
  ##[error]The operation was canceled.
• trunk / macos-py3-arm64 / test (default, 1, 3, macos-m1-stable) (gh)
  inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_nn_functional_max_unpool2d_cpu_float16

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• inductor-periodic / cuda12.1-py3.10-gcc9-sm80 / test (inductor_torchbench_smoketest_perf, 1, 1, linux.gcp.a100) (gh) (detected as infra flaky with no log or failing log classifier)

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

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

• linux-binary-libtorch-cxx11-abi / libtorch-cpu-shared-with-deps-cxx11-abi-build / build (gh) (trunk failure)
  ../torch/csrc/distributed/c10d/GlooDeviceFactory.cpp:93:19: error: ‘getenv’ is not a member of ‘c10::utils’; did you mean ‘get_env’?
• linux-binary-libtorch-pre-cxx11 / libtorch-cpu-shared-with-deps-pre-cxx11-build / build (gh) (trunk failure)
  ../torch/csrc/distributed/c10d/GlooDeviceFactory.cpp:93:19: error: ‘getenv’ is not a member of ‘c10::utils’; did you mean ‘get_env’?

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

[zou3519]

We allow side effects inside the function traced by invoke_subgraph?

[anijain2305]

Yes, these side effects are same as what supported by speculate_subgraph. This test example generates different graphs for the two calls. So, it checks that Dynamo should not incorrectly de-dupe them.

[anijain2305]

Will add expect test.

[zou3519]

Does it make sense to put this in the same place as all of the other invoke_subgraph caches? (just to centralize them?)

[anijain2305]

Yes, that indeed looks better. If you don't mind, we can merge this PR with the cache in tx.output. But, I will move it to a central place in #137965

[zou3519]

sgtm

[anijain2305]

@zou3519 this is ready for another round of review

[zou3519]

what happens if you mutate counter here? Does that error out?

[zou3519]

Code looks fine to me. I do want to check on what happens with side effects here (added a question). Also your tests might be failing

[anijain2305]

Merged in #137538 because of hop_db test requirements.