Files
2026-07-13 12:40:42 +08:00

362 lines
13 KiB
Python

# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import tempfile
import unittest
from unittest import mock
import paddle.base as core
from paddle.utils.cpp_extension import (
CUDA_HOME,
_get_cuda_arch_flags,
_get_num_workers,
_get_pybind11_abi_build_flags,
extension_utils,
)
class TestGetCudaArchFlags(unittest.TestCase):
def setUp(self):
if not core.is_compiled_with_cuda() or core.is_compiled_with_rocm():
self.skipTest('should compile with cuda (not rocm).')
self._old_env = dict(os.environ)
def tearDown(self):
os.environ.clear()
os.environ.update(self._old_env)
def test_with_user_cflags(self):
flags = _get_cuda_arch_flags(cflags=["-arch=sm_90"])
self.assertEqual(flags, [])
def test_with_env_hopper(self):
os.environ["PADDLE_CUDA_ARCH_LIST"] = "Hopper"
flags = _get_cuda_arch_flags()
# Hopper -> 9.0+PTX -> sm_90 + compute_90
self.assertIn("-gencode=arch=compute_90,code=sm_90", flags)
self.assertIn("-gencode=arch=compute_90,code=compute_90", flags)
def test_with_env_hopper_and_flags(self):
os.environ["PADDLE_CUDA_ARCH_LIST"] = "Hopper"
flags = _get_cuda_arch_flags("Hopper")
# Hopper -> 9.0+PTX -> sm_90 + compute_90
self.assertIn("-gencode=arch=compute_90,code=sm_90", flags)
self.assertIn("-gencode=arch=compute_90,code=compute_90", flags)
def test_with_env_multiple(self):
os.environ["PADDLE_CUDA_ARCH_LIST"] = "8.6;9.0+PTX"
flags = _get_cuda_arch_flags()
self.assertIn("-gencode=arch=compute_86,code=sm_86", flags)
self.assertIn("-gencode=arch=compute_90,code=sm_90", flags)
self.assertIn("-gencode=arch=compute_90,code=compute_90", flags)
os.environ["PADDLE_CUDA_ARCH_LIST"] = "8.6,9.0+PTX"
flags = _get_cuda_arch_flags()
self.assertIn("-gencode=arch=compute_86,code=sm_86", flags)
self.assertIn("-gencode=arch=compute_90,code=sm_90", flags)
self.assertIn("-gencode=arch=compute_90,code=compute_90", flags)
os.environ["PADDLE_CUDA_ARCH_LIST"] = "8.6 9.0+PTX"
flags = _get_cuda_arch_flags()
self.assertIn("-gencode=arch=compute_86,code=sm_86", flags)
self.assertIn("-gencode=arch=compute_90,code=sm_90", flags)
self.assertIn("-gencode=arch=compute_90,code=compute_90", flags)
def test_auto_detect(self):
if "PADDLE_CUDA_ARCH_LIST" in os.environ:
del os.environ["PADDLE_CUDA_ARCH_LIST"]
flags = _get_cuda_arch_flags()
self.assertTrue(len(flags) > 0)
def test_get_cuda_arch_flags_with_invalid_arch(self):
os.environ["PADDLE_CUDA_ARCH_LIST"] = "invalid_arch"
with self.assertRaises(ValueError) as context:
_get_cuda_arch_flags()
self.assertIn(
"Unknown CUDA arch (invalid_arch) or GPU not supported",
str(context.exception),
)
def test_skip_paddle_extension_name_flag(self):
flags = _get_cuda_arch_flags(cflags=["-DPADDLE_EXTENSION_NAME=my_ext"])
self.assertNotEqual(flags, [])
def test_rocm_returns_empty_flags(self):
with mock.patch.object(
extension_utils.core, "is_compiled_with_rocm", return_value=True
):
self.assertEqual(_get_cuda_arch_flags(), [])
class TestGetRocmArchFlags(unittest.TestCase):
def setUp(self):
self._old_env = dict(os.environ)
def tearDown(self):
os.environ.clear()
os.environ.update(self._old_env)
def test_default_arch_list(self):
if "PADDLE_ROCM_ARCH_LIST" in os.environ:
del os.environ["PADDLE_ROCM_ARCH_LIST"]
os.environ["ROCM_PATH"] = "/tmp/paddle-missing-rocm-for-test"
os.environ["ROCM_HOME"] = "/tmp/paddle-missing-rocm-for-test"
flags = extension_utils.get_rocm_arch_flags([])
self.assertIn("-fno-gpu-rdc", flags)
self.assertIn("--offload-arch=gfx906", flags)
self.assertIn("--offload-arch=gfx936", flags)
self.assertNotIn("--offload-arch=gfx950", flags)
def test_rocm70_default_arch_list(self):
if "PADDLE_ROCM_ARCH_LIST" in os.environ:
del os.environ["PADDLE_ROCM_ARCH_LIST"]
with tempfile.TemporaryDirectory() as rocm_home:
hip_include = os.path.join(rocm_home, "include", "hip")
os.makedirs(hip_include)
with open(
os.path.join(hip_include, "hip_version.h"),
"w",
encoding="utf-8",
) as f:
f.write(
"#define HIP_VERSION_MAJOR 7\n"
"#define HIP_VERSION_MINOR 0\n"
"#define HIP_VERSION_PATCH 0\n"
)
os.environ["ROCM_PATH"] = rocm_home
os.environ["ROCM_HOME"] = rocm_home
flags = extension_utils.get_rocm_arch_flags([])
self.assertIn("--offload-arch=gfx90a", flags)
self.assertIn("--offload-arch=gfx950", flags)
def test_env_arch_list_override(self):
os.environ["PADDLE_ROCM_ARCH_LIST"] = "gfx950;gfx942 gfx942,gfx908"
flags = extension_utils.get_rocm_arch_flags([])
arch_flags = [
flag for flag in flags if flag.startswith("--offload-arch=")
]
self.assertEqual(
sorted(arch_flags),
[
"--offload-arch=gfx908",
"--offload-arch=gfx942",
"--offload-arch=gfx950",
],
)
def test_user_arch_flags_keep_no_gpu_rdc(self):
flags = extension_utils.get_rocm_arch_flags(["--offload-arch=gfx950"])
self.assertEqual(flags, ["-fno-gpu-rdc"])
def test_user_split_arch_flags_keep_no_gpu_rdc(self):
flags = extension_utils.get_rocm_arch_flags(
["--offload-arch", "gfx950"]
)
self.assertEqual(flags, ["-fno-gpu-rdc"])
def test_user_arch_flags_without_duplicate_no_gpu_rdc(self):
flags = extension_utils.get_rocm_arch_flags(
["-fno-gpu-rdc", "--offload-arch=gfx950"]
)
self.assertEqual(flags, [])
def test_rocm_version_header_empty_home(self):
self.assertIsNone(extension_utils._get_rocm_version_from_header(""))
def test_rocm_version_header_oserror(self):
with tempfile.TemporaryDirectory() as rocm_home:
hip_include = os.path.join(rocm_home, "include", "hip")
os.makedirs(hip_include)
hip_version_file = os.path.join(hip_include, "hip_version.h")
with open(hip_version_file, "w", encoding="utf-8") as f:
f.write("")
with mock.patch(
"paddle.utils.cpp_extension.extension_utils.open",
side_effect=OSError("forced"),
create=True,
):
self.assertIsNone(
extension_utils._get_rocm_version_from_header(rocm_home)
)
def test_default_arch_list_falls_back_to_opt_rocm(self):
for var in ("ROCM_HOME", "ROCM_PATH"):
os.environ.pop(var, None)
with mock.patch.object(
extension_utils,
"_get_rocm_version_from_header",
return_value=None,
) as mocked:
result = extension_utils._get_default_rocm_arch_list()
mocked.assert_called_once_with("/opt/rocm")
self.assertEqual(result, extension_utils._ROCM_LEGACY_AMDGPU_TARGETS)
def test_get_rocm_arch_flags_accepts_none_cflags(self):
if "PADDLE_ROCM_ARCH_LIST" in os.environ:
del os.environ["PADDLE_ROCM_ARCH_LIST"]
os.environ["ROCM_PATH"] = "/tmp/paddle-missing-rocm-for-test"
os.environ["ROCM_HOME"] = "/tmp/paddle-missing-rocm-for-test"
flags = extension_utils.get_rocm_arch_flags(None)
self.assertIn("-fno-gpu-rdc", flags)
self.assertIn("--offload-arch=gfx906", flags)
class TestCppExtensionUtils(unittest.TestCase):
def test_cuda_home(self):
if core.is_compiled_with_cuda():
value = CUDA_HOME
self.assertTrue(value is None or isinstance(value, str))
def test_get_pybind11_abi_build_flags(self):
flags = _get_pybind11_abi_build_flags()
self.assertIsInstance(flags, list)
for f in flags:
self.assertIsInstance(f, str)
def test_get_num_workers_with_env_verbose_false(self):
os.environ["MAX_JOBS"] = "8"
num = _get_num_workers(verbose=False)
self.assertEqual(num, 8)
def test_get_num_workers_with_env_verbose_true(self):
os.environ["MAX_JOBS"] = "8"
num = _get_num_workers(verbose=True)
self.assertEqual(num, 8)
def test_get_num_workers_without_env_verbose_true(self):
if "MAX_JOBS" in os.environ:
del os.environ["MAX_JOBS"]
num = _get_num_workers(verbose=True)
self.assertEqual(num, None)
def test_normalize_extension_kwargs_add_phi_lib_on_windows(self):
with (
mock.patch.object(extension_utils, 'IS_WINDOWS', True),
mock.patch.object(
extension_utils,
'create_sym_link_if_not_exist',
return_value='libpaddle.lib',
),
mock.patch.object(
extension_utils, 'find_paddle_libraries', return_value=[]
),
mock.patch.object(
extension_utils,
'find_paddle_custom_device_includes',
return_value=[],
),
mock.patch.object(
extension_utils, 'find_paddle_includes', return_value=[]
),
mock.patch.object(
extension_utils, 'find_python_includes', return_value=[]
),
):
kwargs = extension_utils.normalize_extension_kwargs(
{'extra_link_args': ['/DEBUG']}, use_cuda=False
)
self.assertEqual(
kwargs['extra_link_args'],
[
'/DEBUG',
*extension_utils.MSVC_LINK_FLAGS,
'libpaddle.lib',
'phi.lib',
],
)
def test_normalize_extension_kwargs_keep_user_phi_lib_on_windows(self):
with (
mock.patch.object(extension_utils, 'IS_WINDOWS', True),
mock.patch.object(
extension_utils,
'create_sym_link_if_not_exist',
return_value='libpaddle.lib',
),
mock.patch.object(
extension_utils, 'find_paddle_libraries', return_value=[]
),
mock.patch.object(
extension_utils,
'find_paddle_custom_device_includes',
return_value=[],
),
mock.patch.object(
extension_utils, 'find_paddle_includes', return_value=[]
),
mock.patch.object(
extension_utils, 'find_python_includes', return_value=[]
),
):
kwargs = extension_utils.normalize_extension_kwargs(
{
'extra_link_args': [
'/DEBUG',
'phi.lib',
'libpaddle.lib',
]
},
use_cuda=False,
)
self.assertEqual(kwargs['extra_link_args'].count('phi.lib'), 1)
self.assertEqual(kwargs['extra_link_args'].count('libpaddle.lib'), 1)
def test_normalize_extension_kwargs_add_cuda_libs_on_windows(self):
with (
mock.patch.object(extension_utils, 'IS_WINDOWS', True),
mock.patch.object(
extension_utils,
'create_sym_link_if_not_exist',
return_value='libpaddle.lib',
),
mock.patch.object(
extension_utils, 'find_paddle_libraries', return_value=[]
),
mock.patch.object(
extension_utils,
'find_paddle_custom_device_includes',
return_value=[],
),
mock.patch.object(
extension_utils, 'find_paddle_includes', return_value=[]
),
mock.patch.object(
extension_utils, 'find_python_includes', return_value=[]
),
):
kwargs = extension_utils.normalize_extension_kwargs(
{
'extra_link_args': [
'/DEBUG',
'phi.lib',
'cudadevrt.lib',
]
},
use_cuda=True,
)
self.assertEqual(kwargs['extra_link_args'].count('phi.lib'), 1)
self.assertEqual(kwargs['extra_link_args'].count('libpaddle.lib'), 1)
self.assertEqual(kwargs['extra_link_args'].count('cudadevrt.lib'), 1)
self.assertEqual(
kwargs['extra_link_args'].count('cudart_static.lib'), 1
)
if __name__ == "__main__":
unittest.main()