mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
[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:
committed by
PyTorch MergeBot
parent
fa4d5e76ea
commit
65d642d6db
@ -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})
|
||||
|
@ -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()
|
||||
|
@ -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));
|
||||
|
Reference in New Issue
Block a user