[ROCm] enable aoti tests, forward fix 162353 (#162827)

Forward fix for tests added by #162353.  Enables aoti tests on rocm.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162827
Approved by: https://github.com/dolpm, https://github.com/huydhn
This commit is contained in:
Jeff Daily
2025-09-12 20:05:50 +00:00
committed by PyTorch MergeBot
parent fa4d5e76ea
commit 65d642d6db
3 changed files with 5 additions and 3 deletions

View File

@ -551,9 +551,11 @@ if(USE_CUDA OR USE_ROCM)
endif()
if(USE_CUDA)
# eventually do rocm
append_filelist("libtorch_nativert_cuda_sources" Caffe2_GPU_SRCS)
endif()
if(USE_ROCM)
append_filelist("libtorch_nativert_cuda_sources" Caffe2_HIP_SRCS)
endif()
if(USE_CUDA)
list(APPEND Caffe2_GPU_CU_SRCS ${Caffe2_GPU_HIP_JIT_FUSERS_SRCS})

View File

@ -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()

View File

@ -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));