[nativert] aoti

[dolpm]

Summary: att

Test Plan:
ci

Rollback Plan:

Differential Revision: D81731425

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 5d1f70b with merge base 43d9b5e ():
💚 Looks good so far! There are no failures yet. 💚

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

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D81731425

[facebook-github-bot]

@pytorchbot merge

(Initiating merge automatically since Phabricator Diff has merged)

[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

[jeffdaily]

causing rocm failures

https://hud.pytorch.org/failure?name=rocm-mi300%20%2F%20linux-noble-rocm-py3.12-mi300%20%2F%20test%20(default%2C%203%2C%206%2C%20linux.rocm.gpu.gfx942.1)&jobName=undefined&failureCaptures=%27test%2Fexport%2Ftest_nativert.py%3A%3ATestNativeRT%3A%3Atest_aoti_3_cuda%27%2C%20%27test%2Fexport%2Ftest_nativert.py%3A%3ATestNativeRT%3A%3Atest_aoti_4_cuda%27%2C%20%27test%2Fexport%2Ftest_nativert.py%3A%3ATestNativeRT%3A%3Atest_aoti_5_cuda%27

[dolpm]

causing rocm failures

https://hud.pytorch.org/failure?name=rocm-mi300%20%2F%20linux-noble-rocm-py3.12-mi300%20%2F%20test%20(default%2C%203%2C%206%2C%20linux.rocm.gpu.gfx942.1)&jobName=undefined&failureCaptures=%27test%2Fexport%2Ftest_nativert.py%3A%3ATestNativeRT%3A%3Atest_aoti_3_cuda%27%2C%20%27test%2Fexport%2Ftest_nativert.py%3A%3ATestNativeRT%3A%3Atest_aoti_4_cuda%27%2C%20%27test%2Fexport%2Ftest_nativert.py%3A%3ATestNativeRT%3A%3Atest_aoti_5_cuda%27

let me fix. can revert if not simple

[jeffdaily]

I've taken a quick look. There are some changes that seem necessary for ROCm but they don't avoid the seg fault.

diff --git a/test/cpp/nativert/CMakeLists.txt b/test/cpp/nativert/CMakeLists.txt
index 30c68283282..77636a2d109 100644
--- a/test/cpp/nativert/CMakeLists.txt
+++ b/test/cpp/nativert/CMakeLists.txt
@@ -48,7 +48,7 @@ set(NATIVERT_TEST_SRCS
   ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp
 )

-if(USE_CUDA)
+if(USE_CUDA OR USE_ROCM)
   list(APPEND NATIVERT_TEST_SRCS ${TORCH_ROOT}/torch/nativert/executor/triton/CudaTritonKernelManager.cpp)
   list(APPEND NATIVERT_TEST_SRCS ${TORCH_ROOT}/torch/nativert/executor/AOTInductorModelContainerCudaShim.cpp)
 endif()
diff --git a/test/cpp/nativert/test_aoti_model_container_registration.cpp b/test/cpp/nativert/test_aoti_model_container_registration.cpp
index 94b19bf8e75..9a5d67a9c3c 100644
--- a/test/cpp/nativert/test_aoti_model_container_registration.cpp
+++ b/test/cpp/nativert/test_aoti_model_container_registration.cpp
@@ -8,7 +8,7 @@ using namespace torch::nativert;
 TEST(AOTIModelContainerRegistrationTests, TestRegister) {
   EXPECT_TRUE(AOTIModelContainerRunnerRegistry()->Has(at::kCPU));

-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_ROCM)
   EXPECT_TRUE(AOTIModelContainerRunnerRegistry()->Has(at::kCUDA));
 #else
   EXPECT_FALSE(AOTIModelContainerRunnerRegistry()->Has(at::kCUDA));

I'm doing a build with BUILD_LIBTORCH_CPU_WITH_DEBUG:BOOL=ON to get a backtrace for the seg fault.

[dolpm]

I've taken a quick look. There are some changes that seem necessary for ROCm but they don't avoid the seg fault.

diff --git a/test/cpp/nativert/CMakeLists.txt b/test/cpp/nativert/CMakeLists.txt
index 30c68283282..77636a2d109 100644
--- a/test/cpp/nativert/CMakeLists.txt
+++ b/test/cpp/nativert/CMakeLists.txt
@@ -48,7 +48,7 @@ set(NATIVERT_TEST_SRCS
   ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp
 )

-if(USE_CUDA)
+if(USE_CUDA OR USE_ROCM)
   list(APPEND NATIVERT_TEST_SRCS ${TORCH_ROOT}/torch/nativert/executor/triton/CudaTritonKernelManager.cpp)
   list(APPEND NATIVERT_TEST_SRCS ${TORCH_ROOT}/torch/nativert/executor/AOTInductorModelContainerCudaShim.cpp)
 endif()
diff --git a/test/cpp/nativert/test_aoti_model_container_registration.cpp b/test/cpp/nativert/test_aoti_model_container_registration.cpp
index 94b19bf8e75..9a5d67a9c3c 100644
--- a/test/cpp/nativert/test_aoti_model_container_registration.cpp
+++ b/test/cpp/nativert/test_aoti_model_container_registration.cpp
@@ -8,7 +8,7 @@ using namespace torch::nativert;
 TEST(AOTIModelContainerRegistrationTests, TestRegister) {
   EXPECT_TRUE(AOTIModelContainerRunnerRegistry()->Has(at::kCPU));

-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_ROCM)
   EXPECT_TRUE(AOTIModelContainerRunnerRegistry()->Has(at::kCUDA));
 #else
   EXPECT_FALSE(AOTIModelContainerRunnerRegistry()->Has(at::kCUDA));

I'm doing a build with BUILD_LIBTORCH_CPU_WITH_DEBUG:BOOL=ON to get a backtrace for the seg fault.

it will take me a few mins to repro this locally on amd machine, but i didn't expect these tests to work on rocm -- we can just disable them for now if this is the case? the reason this isn't changing anything is USE_ROCM and USE_CUDA are set when it's an amd gpu build so this isn't doing anything.

[jeffdaily]

Found one more diff I'm trying out. The earlier diffs were for the test binary only but the sources weren't being included for rocm builds into libtorch.

diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
index b5d47bb4b5d..bcf8fffa4ac 100644
--- a/caffe2/CMakeLists.txt
+++ b/caffe2/CMakeLists.txt
@@ -550,8 +550,7 @@ if(USE_CUDA OR USE_ROCM)
   append_filelist("libtorch_cuda_core_sources" Caffe2_GPU_HIP_JIT_FUSERS_SRCS)
 endif()

-if(USE_CUDA)
-  # eventually do rocm
+if(USE_CUDA OR USE_ROCM)
   append_filelist("libtorch_nativert_cuda_sources" Caffe2_GPU_SRCS)
 endif()

[dolpm]

Found one more diff I'm trying out. The earlier diffs were for the test binary only but the sources weren't being included for rocm builds into libtorch.

diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
index b5d47bb4b5d..bcf8fffa4ac 100644
--- a/caffe2/CMakeLists.txt
+++ b/caffe2/CMakeLists.txt
@@ -550,8 +550,7 @@ if(USE_CUDA OR USE_ROCM)
   append_filelist("libtorch_cuda_core_sources" Caffe2_GPU_HIP_JIT_FUSERS_SRCS)
 endif()

-if(USE_CUDA)
-  # eventually do rocm
+if(USE_CUDA OR USE_ROCM)
   append_filelist("libtorch_nativert_cuda_sources" Caffe2_GPU_SRCS)
 endif()

can try 186290f as well.

[jeffdaily]

Besides skipping the tests, is there some technical reason why you don't think these tests should work on ROCm?

[dolpm]

Besides skipping the tests, is there some technical reason why you don't think these tests should work on ROCm?

it will get the registry by the device type of the input, so we need to change the registry code as well based on USE_ROCM -- i think i know how to fix the issue you are seeing. give me 5 mins.

for some other reason i can't get my rocm build to work -- mind trying 52cf1c7 ? assuming i understand the problem, this should fix it.

[jeffdaily]

I figured it out finally. See my PR here. Tests were passing for me.
#162827

[dolpm]

I figured it out finally. See my PR here. Tests were passing for me. #162827

left a comment