[aoti] Add support for more custom op input/output types

[angelayi]

Stack from ghstack (oldest at bottom):

• -> [aoti] Add support for more custom op input/output types #126215
• [aoti] Add initial custom op support #125242

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

[pytorch-bot]

🔗 Helpful Links

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

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

❌ 4 New Failures, 1 Unrelated Failure

As of commit 03456a8 with merge base ff82e2e ():

NEW FAILURES - The following jobs have failed:

• pull / linux-focal-cpu-py3.10-gcc9-bazel-test / build-and-test (default, 1, 1, linux.4xlarge) (gh)
  /var/lib/jenkins/workspace/BUILD.bazel:826:11: Compiling torch/csrc/inductor/aoti_runner/model_container_runner.cpp failed: (Exit 1): gcc failed: error executing command (from target //:torch) /opt/cache/bin/gcc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wunused-but-set-parameter -Wno-free-nonheap-object -fno-omit-frame-pointer '-std=c++11' -MD -MF ... (remaining 323 arguments skipped)
• pull / linux-focal-cuda12.1-py3.10-gcc9-bazel-test / build-and-test (default, 1, 1, linux.4xlarge.nvidia.gpu) (gh)
  /var/lib/jenkins/workspace/BUILD.bazel:826:11: Compiling torch/csrc/inductor/aoti_torch/shim_common.cpp failed: (Exit 1): gcc failed: error executing command (from target //:torch) /opt/cache/bin/gcc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wunused-but-set-parameter -Wno-free-nonheap-object -fno-omit-frame-pointer '-std=c++11' -MD -MF ... (remaining 372 arguments skipped)
• pull / linux-focal-cuda12.4-py3.10-gcc9-bazel-test / build-and-test (default, 1, 1, linux.4xlarge.nvidia.gpu) (gh)
  /var/lib/jenkins/workspace/BUILD.bazel:826:11: Compiling torch/csrc/inductor/aoti_runner/model_container_runner.cpp failed: (Exit 1): gcc failed: error executing command (from target //:torch) /opt/cache/bin/gcc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wunused-but-set-parameter -Wno-free-nonheap-object -fno-omit-frame-pointer '-std=c++11' -MD -MF ... (remaining 372 arguments skipped)
• pull / win-vs2019-cpu-py3 / build (gh)
  C:\actions-runner\_work\pytorch\pytorch\torch\csrc\inductor\aoti_torch\oss_proxy_executor.cpp(228): error C2143: syntax error: missing ';' before '{'

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

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

• pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 2, 3, linux.8xlarge.nvidia.gpu) (gh) (trunk failure)
  distributed/pipelining/test_schedule.py::ScheduleTest::test_grad_with_manual_interleaved_ScheduleClass0

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

[chenyang78]

Are these couts for debugging or logging? If they are for logging, could we use c10's logging utilities?

[desertfire]

We need to make sure this kind of enum value mapping is always valid, to avoid any surprising BC breakage. I think this is achieved by defining a Thrift mapping internally (or maybe even that doesn't provide a strong guard on BC, cc @SherlockNoMad ). Because this kind of enum value should rarely change in practice, we could add some enum value checking at the compile time to make sure the mapping hasn't changed. Or maybe simply adding more tests to test/cpp/aoti_abi_check to prevent surprising enum value changes.