mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-28 18:56:51 +08:00
Jeff Daily
f2e5b9121f
test/cpp_extensions/setup.py only built the cublas and cusolver extensions when CUDA_HOME was set, and CUDA_HOME is None on ROCm builds (ROCM_HOME is set instead), so neither module existed on ROCm and both tests had to be skipped there. Build them when either variable is set, and link hipblas and hipsolver explicitly, mirroring -lcublas and -lcusolver on CUDA: the previous empty HIP libraries list would only have resolved the symbols through libtorch_hip's own dependency, which is exactly the re-export assumption these tests exist to rule out.
The cusolver extension also needed a hipify include mapping. Every cusolverDn identifier the source uses is already mapped, but cusolverDn.h itself is not, so the include line survived hipification and the build failed on the missing header. Map it to hipsolver/hipsolver.h. The only file in the tree that includes cusolverDn.h is ATen/cuda/CUDAContextLight.h, which guards it with ifdef CUDART_VERSION and includes hipsolver/hipsolver.h separately under USE_ROCM, so the mapping only rewrites a line that ROCm builds never compile.
The dlink extension stays CUDA-only: its build path resolves CUDA arch flags with no ROCm branch. Its skipIfRocm now carries that reason.
Test Plan:
On a ROCm 10.0 gfx950 host, the CI entry point for this file builds every extension and runs the tests:
python test/run_test.py --include test_cpp_extensions_aot_ninja --verbose
# 20 passed, 4 skipped
test_cublas_extension and test_cusolver_extension both pass on ROCm (they were skipped before), and the dlink tests still skip with their reason. readelf on the built cublas module shows libhipblas as a direct NEEDED entry.
This PR was authored with the assistance of an AI agent (Claude). The analysis and test results were verified by hand.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/197816
Approved by: https://github.com/jeffdaily
Co-authored-by: Jeff Daily <jeff.daily@amd.com>
140 lines
3.9 KiB
Python
140 lines
3.9 KiB
Python
import os
|
|
import sys
|
|
|
|
from setuptools import setup
|
|
|
|
import torch.cuda
|
|
from torch.testing._internal.common_utils import IS_WINDOWS
|
|
from torch.utils.cpp_extension import (
|
|
BuildExtension,
|
|
CppExtension,
|
|
CUDA_HOME,
|
|
CUDAExtension,
|
|
ROCM_HOME,
|
|
SyclExtension,
|
|
)
|
|
|
|
|
|
if sys.platform == "win32":
|
|
vc_version = os.getenv("VCToolsVersion", "")
|
|
if vc_version.startswith("14.16."):
|
|
CXX_FLAGS = ["/sdl"]
|
|
else:
|
|
CXX_FLAGS = ["/sdl", "/permissive-"]
|
|
else:
|
|
CXX_FLAGS = ["-g"]
|
|
|
|
USE_NINJA = os.getenv("USE_NINJA") == "1"
|
|
|
|
ext_modules = [
|
|
CppExtension(
|
|
"torch_test_cpp_extension.cpp", ["extension.cpp"], extra_compile_args=CXX_FLAGS
|
|
),
|
|
CppExtension(
|
|
"torch_test_cpp_extension.maia",
|
|
["maia_extension.cpp"],
|
|
extra_compile_args=CXX_FLAGS,
|
|
),
|
|
CppExtension(
|
|
"torch_test_cpp_extension.rng",
|
|
["rng_extension.cpp"],
|
|
extra_compile_args=CXX_FLAGS,
|
|
),
|
|
]
|
|
|
|
NVCC_FLAGS = ["-O2"] + (["-DUSE_CUDA"] if IS_WINDOWS else [])
|
|
|
|
if torch.cuda.is_available() and (CUDA_HOME is not None or ROCM_HOME is not None):
|
|
extension = CUDAExtension(
|
|
"torch_test_cpp_extension.cuda",
|
|
[
|
|
"cuda_extension.cpp",
|
|
"cuda_extension_kernel.cu",
|
|
"cuda_extension_kernel2.cu",
|
|
],
|
|
extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
|
|
)
|
|
ext_modules.append(extension)
|
|
|
|
if torch.cuda.is_available() and (CUDA_HOME is not None or ROCM_HOME is not None):
|
|
extension = CUDAExtension(
|
|
"torch_test_cpp_extension.torch_library",
|
|
["torch_library.cu"],
|
|
extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
|
|
)
|
|
ext_modules.append(extension)
|
|
|
|
if torch.backends.mps.is_available():
|
|
extension = CppExtension(
|
|
"torch_test_cpp_extension.mps",
|
|
["mps_extension.mm"],
|
|
extra_compile_args=CXX_FLAGS,
|
|
)
|
|
ext_modules.append(extension)
|
|
|
|
if torch.xpu.is_available() and USE_NINJA:
|
|
extension = SyclExtension(
|
|
"torch_test_cpp_extension.sycl",
|
|
["xpu_extension.sycl"],
|
|
extra_compile_args={"cxx": CXX_FLAGS, "sycl": ["-O2"]},
|
|
)
|
|
ext_modules.append(extension)
|
|
|
|
|
|
# todo(mkozuki): Figure out the root cause
|
|
if (
|
|
(not IS_WINDOWS)
|
|
and torch.cuda.is_available()
|
|
and (CUDA_HOME is not None or ROCM_HOME is not None)
|
|
):
|
|
# malfet: One should not assume that PyTorch re-exports CUDA dependencies
|
|
cublas_extension = CUDAExtension(
|
|
name="torch_test_cpp_extension.cublas_extension",
|
|
sources=["cublas_extension.cpp"],
|
|
libraries=["cublas"] if torch.version.hip is None else ["hipblas"],
|
|
)
|
|
ext_modules.append(cublas_extension)
|
|
|
|
if (
|
|
(not IS_WINDOWS)
|
|
and torch.cuda.is_available()
|
|
and (CUDA_HOME is not None or ROCM_HOME is not None)
|
|
):
|
|
cusolver_extension = CUDAExtension(
|
|
name="torch_test_cpp_extension.cusolver_extension",
|
|
sources=["cusolver_extension.cpp"],
|
|
libraries=["cusolver"] if torch.version.hip is None else ["hipsolver"],
|
|
)
|
|
ext_modules.append(cusolver_extension)
|
|
|
|
if (
|
|
USE_NINJA
|
|
and (not IS_WINDOWS)
|
|
and torch.cuda.is_available()
|
|
and CUDA_HOME is not None
|
|
):
|
|
extension = CUDAExtension(
|
|
name="torch_test_cpp_extension.cuda_dlink",
|
|
sources=[
|
|
"cuda_dlink_extension.cpp",
|
|
"cuda_dlink_extension_kernel.cu",
|
|
"cuda_dlink_extension_add.cu",
|
|
],
|
|
dlink=True,
|
|
extra_compile_args={"cxx": CXX_FLAGS, "nvcc": ["-O2", "-dc"]},
|
|
)
|
|
ext_modules.append(extension)
|
|
|
|
setup(
|
|
name="torch_test_cpp_extension",
|
|
packages=["torch_test_cpp_extension"],
|
|
ext_modules=ext_modules,
|
|
include_dirs="self_compiler_include_dirs_test",
|
|
cmdclass={"build_ext": BuildExtension.with_options(use_ninja=USE_NINJA)},
|
|
entry_points={
|
|
"torch.backends": [
|
|
"device_backend = torch_test_cpp_extension:_autoload",
|
|
],
|
|
},
|
|
)
|