362 lines
13 KiB
Python
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()
|