Files
paddlepaddle--paddle/test/cpp_extension/test_cpp_extension_ninja.py
T
2026-07-13 12:40:42 +08:00

1318 lines
49 KiB
Python

# Copyright (c) 2026 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 shlex
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace, SimpleNamespace as ModuleSimpleNamespace
from unittest import mock
from setuptools import Distribution
from paddle.utils.cpp_extension.cpp_extension import (
_as_command_list,
_get_num_workers,
_is_ninja_available,
_join_ninja_shell_list,
_ninja_escape_path,
_run_ninja_build,
_write_ninja_file,
)
from paddle.utils.cpp_extension.extension_utils import (
_write_setup_file,
)
class TestNinjaHelperFunctions(unittest.TestCase):
def test_is_ninja_available(self):
result = _is_ninja_available()
self.assertIsInstance(result, bool)
if sys.platform.startswith('linux'):
self.assertTrue(result, "ninja should be available on Linux CI")
def test_is_ninja_available_handles_probe_failure(self):
with mock.patch(
'paddle.utils.cpp_extension.cpp_extension.subprocess.check_output',
side_effect=OSError("ninja missing"),
):
self.assertFalse(_is_ninja_available())
def test_ninja_escape_path(self):
self.assertEqual(_ninja_escape_path("/path/to/file"), "/path/to/file")
self.assertEqual(
_ninja_escape_path("/path with space/file"),
"/path$ with$ space/file",
)
self.assertEqual(
_ninja_escape_path("C:/path/to/file"), "C$:/path/to/file"
)
self.assertEqual(
_ninja_escape_path("/path/$var/file"), "/path/$$var/file"
)
self.assertEqual(
_ninja_escape_path("C:/path with $var/file"),
"C$:/path$ with$ $$var/file",
)
def test_join_ninja_shell_list(self):
self.assertEqual(
_join_ninja_shell_list("simple string"), "simple string"
)
self.assertEqual(_join_ninja_shell_list([]), "")
result = _join_ninja_shell_list(['-c', '-O2', '-I/usr/include'])
self.assertIn('-c', result)
self.assertIn('-O2', result)
result = _join_ninja_shell_list(['/path with space', '-c'])
self.assertIn('/path with space', result)
def test_as_command_list(self):
self.assertEqual(_as_command_list("gcc"), ["gcc"])
self.assertEqual(_as_command_list(["gcc", "-c"]), ["gcc", "-c"])
self.assertEqual(_as_command_list(("gcc", "-c")), ["gcc", "-c"])
def test_write_ninja_file(self):
with tempfile.TemporaryDirectory() as tmpdir:
ninja_path = os.path.join(tmpdir, "build.ninja")
_write_ninja_file(ninja_path, "ninja_required_version = 1.5")
with open(ninja_path, 'r', encoding='utf-8') as f:
written_content = f.read()
self.assertTrue(written_content.endswith('\n'))
self.assertIn("ninja_required_version = 1.5", written_content)
content_with_newline = "rule compile\n command = $cxx -c $in\n"
_write_ninja_file(ninja_path, content_with_newline)
with open(ninja_path, 'r', encoding='utf-8') as f:
written_content = f.read()
self.assertTrue(written_content.endswith('\n'))
nested_path = os.path.join(tmpdir, "nested", "dir", "build.ninja")
_write_ninja_file(nested_path, "test content")
self.assertTrue(os.path.exists(nested_path))
def test_get_num_workers_with_max_jobs_env(self):
with mock.patch.dict(os.environ, {'MAX_JOBS': '4'}, clear=False):
result = _get_num_workers(verbose=False)
self.assertEqual(result, 4)
def test_run_ninja_build_windows_vc_env(self):
if sys.platform != 'win32':
self.skipTest("Windows-only test")
with tempfile.TemporaryDirectory() as tmpdir:
ninja_path = os.path.join(tmpdir, "build.ninja")
_write_ninja_file(ninja_path, "rule cc\n command = echo hello\n")
vc_env_mock = {'PATH': '/vc/bin', 'INCLUDE': '/vc/include'}
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.distutils.util.get_platform',
return_value='win-amd64',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_vc_env',
return_value=vc_env_mock,
),
mock.patch('subprocess.run') as mock_run,
):
_run_ninja_build(tmpdir, verbose=True, error_prefix="Test")
mock_run.assert_called_once()
def test_run_ninja_build_with_work_directory(self):
with tempfile.TemporaryDirectory() as tmpdir:
build_dir = os.path.join(tmpdir, 'build')
work_dir = os.path.join(tmpdir, 'src')
os.makedirs(build_dir)
os.makedirs(work_dir)
with (
mock.patch.dict(
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
),
mock.patch('subprocess.run') as mock_run,
):
_run_ninja_build(
build_dir,
verbose=True,
error_prefix='Test',
work_directory=work_dir,
)
command = mock_run.call_args.args[0]
self.assertIn('-f', command)
self.assertIn(os.path.join(build_dir, 'build.ninja'), command)
self.assertEqual(mock_run.call_args.kwargs['cwd'], work_dir)
def test_run_ninja_build_windows_with_vscmd_env(self):
if sys.platform != 'win32':
self.skipTest("Windows-only test")
with tempfile.TemporaryDirectory() as tmpdir:
ninja_path = os.path.join(tmpdir, "build.ninja")
_write_ninja_file(ninja_path, "rule cc\n command = echo hello\n")
with (
mock.patch.dict(
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS', True
),
mock.patch('subprocess.run') as mock_run,
):
_run_ninja_build(tmpdir, verbose=True, error_prefix="Test")
mock_run.assert_called_once()
class _FakeUnixCompiler:
compiler_type = 'unix'
def __init__(self, objects, build_map):
self.src_extensions = ['.c', '.cc', '.cpp']
self.compiler_so = ['g++']
self.compiler_cxx = ['g++']
self._objects = objects
self._build_map = build_map
self._compile_calls = []
self.verbose = False
def compile(self, *args, **kwargs):
raise AssertionError("original unix compile should be replaced")
def link_shared_object(self, *args, **kwargs):
return None
def object_filenames(self, *args, **kwargs):
return self._objects
def _setup_compile(
self,
output_dir,
macros,
include_dirs,
sources,
depends,
extra_postargs,
):
pp_opts = [f'-I{inc}' for inc in (include_dirs or [])]
return macros, self._objects, extra_postargs, pp_opts, self._build_map
def _get_cc_args(self, pp_opts, debug, extra_preargs):
return list(extra_preargs or []) + list(pp_opts)
def _compile(self, obj, src, ext, cc_args, cflags, pp_opts):
self._compile_calls.append(
{
'obj': obj,
'src': src,
'ext': ext,
'cc_args': list(cc_args),
'cflags': list(cflags),
'pp_opts': list(pp_opts),
'compiler_so': list(self.compiler_so)
if isinstance(self.compiler_so, list)
else self.compiler_so,
}
)
def set_executable(self, key, value):
setattr(self, key, value)
def set_executables(self, **kwargs):
return None
class _FakeMsvcCompiler:
compiler_type = 'msvc'
def __init__(self, objects, build_map):
self.src_extensions = ['.c', '.cc', '.cpp']
self._cpp_extensions = ['.c', '.cc', '.cpp']
self.cc = ['cl.exe']
self.compile_options = ['/nologo', '/O2', '/W3', '/MD']
self.compile_options_debug = [
'/nologo',
'/Od',
'/MDd',
'/Zi',
'/W3',
'/D_DEBUG',
]
self.initialized = True
self._objects = objects
self._build_map = build_map
self._spawned_cmds = []
self.verbose = False
self.spawn = self._spawn
def compile(self, *args, **kwargs):
macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
kwargs.get('output_dir'),
kwargs.get('macros'),
kwargs.get('include_dirs'),
args[0],
kwargs.get('depends'),
kwargs.get('extra_postargs'),
)
del macros, extra_postargs
for obj in objects:
src, ext = build[obj]
cmd = [
'cl.exe',
'/Tp' + src,
'/Fo' + obj,
*pp_opts,
]
if ext == '.c':
cmd[1] = '/Tc' + src
self.spawn(cmd)
return objects
def initialize(self):
self.initialized = True
def object_filenames(self, *args, **kwargs):
return self._objects
def _setup_compile(
self,
output_dir,
macros,
include_dirs,
sources,
depends,
extra_postargs,
):
pp_opts = [f'/I{inc}' for inc in (include_dirs or [])]
return macros, self._objects, extra_postargs, pp_opts, self._build_map
def _spawn(self, cmd):
self._spawned_cmds.append(list(cmd))
class TestBuildExtension(unittest.TestCase):
def _build_extension(self, **kwargs):
from paddle.utils.cpp_extension.cpp_extension import BuildExtension
return BuildExtension(dist=Distribution(), **kwargs)
def _run_build_with_fake_compiler(
self,
compiler,
sources,
extra_compile_args,
include_dirs=None,
):
ext = SimpleNamespace(
name='fake_extension',
_full_name='fake_extension',
sources=sources,
extra_compile_args=extra_compile_args,
)
cmd = self._build_extension(use_ninja=True)
with tempfile.TemporaryDirectory() as tmpdir:
cmd.extensions = [ext]
cmd.compiler = compiler
cmd.build_temp = tmpdir
cmd.build_lib = tmpdir
cmd.verbose = False
cmd.contain_cuda_file = any(
source.endswith('.cu') for source in sources
)
captured = {}
def fake_write_ninja_file(path, content):
captured['ninja_path'] = path
captured['ninja_content'] = content
def fake_run_ninja_build(
build_directory, verbose, error_prefix, work_directory=None
):
captured['ninja_build_directory'] = build_directory
captured['ninja_work_directory'] = work_directory
def fake_build_extensions(_):
_.compiler.compile(
sources,
output_dir=_.build_temp,
macros=[],
include_dirs=include_dirs or [],
debug=False,
extra_preargs=[],
extra_postargs=extra_compile_args,
depends=None,
)
patches = [
mock.patch.object(cmd, '_check_abi', return_value=None),
mock.patch.object(cmd, '_record_op_info', return_value=None),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.clean_object_if_change_cflags',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.define_paddle_extension_name',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._reset_so_rpath',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._write_ninja_file',
side_effect=fake_write_ninja_file,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._run_ninja_build',
side_effect=fake_run_ninja_build,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.build_ext.build_extensions',
side_effect=fake_build_extensions,
),
]
with (
patches[0],
patches[1],
patches[2],
patches[3],
patches[4],
patches[5],
patches[6],
patches[7],
):
cmd.build_extensions()
return captured
def _patch_build_extension_common(self, cmd):
return [
mock.patch.object(cmd, '_check_abi', return_value=None),
mock.patch.object(cmd, '_record_op_info', return_value=None),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.clean_object_if_change_cflags',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.define_paddle_extension_name',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._reset_so_rpath',
return_value=None,
),
]
@staticmethod
def _expected_nvcc_path(cuda_home='/opt/cuda'):
return os.path.join(cuda_home, 'bin', 'nvcc')
def test_unix_ninja_build_file_contains_multiple_sources(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_sub.cc"),
]
objects = [
str(test_dir / "build" / "custom_extension.o"),
str(test_dir / "build" / "custom_sub.o"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cc'),
}
compiler = _FakeUnixCompiler(objects, build_map)
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={'cxx': ['-w', '-g'], 'nvcc': []},
include_dirs=[str(test_dir / "include")],
)
content = captured['ninja_content']
self.assertIn('rule compile', content)
self.assertIn('deps = gcc', content)
self.assertIn(_ninja_escape_path(os.path.abspath(objects[0])), content)
self.assertIn(_ninja_escape_path(os.path.abspath(objects[1])), content)
self.assertIn(_ninja_escape_path(os.path.abspath(sources[0])), content)
self.assertIn(_ninja_escape_path(os.path.abspath(sources[1])), content)
self.assertIn('post_cflags = -w -g', content)
self.assertIn(
f'default {_ninja_escape_path(os.path.abspath(objects[0]))} {_ninja_escape_path(os.path.abspath(objects[1]))}',
content,
)
self.assertTrue(captured['ninja_path'].endswith('build.ninja'))
def test_windows_ninja_build_file_contains_multiple_sources(self):
if sys.platform != 'win32':
self.skipTest("Windows-only test")
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_sub.cc"),
]
objects = [
str(test_dir / "build" / "custom_extension.obj"),
str(test_dir / "build" / "custom_sub.obj"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cc'),
}
compiler = _FakeMsvcCompiler(objects, build_map)
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={'cxx': ['/wd4244'], 'nvcc': []},
include_dirs=[str(test_dir / "include dir")],
)
content = captured['ninja_content']
self.assertIn('rule compile', content)
self.assertIn('deps = msvc', content)
self.assertIn('command = $cxx /showIncludes', content)
self.assertIn('cl.exe', content)
self.assertIn('/wd4244', content)
self.assertIn(_ninja_escape_path(os.path.abspath(objects[0])), content)
self.assertIn(_ninja_escape_path(os.path.abspath(objects[1])), content)
self.assertIn(_ninja_escape_path(os.path.abspath(sources[0])), content)
self.assertIn(_ninja_escape_path(os.path.abspath(sources[1])), content)
self.assertIn('/I', content)
self.assertTrue(captured['ninja_path'].endswith('build.ninja'))
def test_windows_ninja_quotes_flags_once(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
include_dir = test_dir / "include dir"
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.obj"),
str(test_dir / "build" / "custom_kernel.obj"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeMsvcCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS',
True,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={
'cxx': ['/DNAME=value with space'],
'nvcc': ['--compiler-options=/Wall /WX'],
},
include_dirs=[str(include_dir)],
)
content = captured['ninja_content']
self.assertIn(subprocess.list2cmdline([f'/I{include_dir}']), content)
self.assertIn(
subprocess.list2cmdline(['/DNAME=value with space']), content
)
self.assertIn(
subprocess.list2cmdline(['--compiler-options=/Wall /WX']),
content,
)
self.assertNotIn('\\"/I', content)
self.assertNotIn('\\"/DNAME', content)
def test_unix_ninja_build_file_contains_cuda_sources(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.o"),
str(test_dir / "build" / "custom_kernel.cu.o"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeUnixCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
return_value='/usr/bin/ccache',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
return_value=False,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
return_value=False,
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={
'cxx': ['-w'],
'nvcc': ['--gpu-flag'],
},
include_dirs=[str(test_dir / "include")],
)
content = captured['ninja_content']
self.assertIn('cxx = /usr/bin/ccache g++', content)
self.assertIn(
f'nvcc = /usr/bin/ccache {self._expected_nvcc_path()}',
content,
)
self.assertIn('rule cuda_compile', content)
self.assertIn('cuda_post_cflags = --prepared --gpu-flag', content)
self.assertIn('post_cflags = -w -D_GLIBCXX_USE_CXX11_ABI=1', content)
self.assertIn('-DPADDLE_WITH_CUDA', content)
def test_unix_ninja_quotes_flags_once(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
include_dir = test_dir / "include dir"
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.o"),
str(test_dir / "build" / "custom_kernel.cu.o"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeUnixCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS',
False,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
return_value=False,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
return_value=False,
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={
'cxx': ['-DNAME=value with space'],
'nvcc': ['--compiler-options=-Wall -Wextra'],
},
include_dirs=[str(include_dir)],
)
content = captured['ninja_content']
self.assertIn(shlex.join([f'-I{include_dir}']), content)
self.assertIn(shlex.join(['-DNAME=value with space']), content)
self.assertIn(shlex.join(['--compiler-options=-Wall -Wextra']), content)
self.assertNotIn('"\'"\'"', content)
def test_unix_ninja_preserves_relative_include_flags(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.o"),
str(test_dir / "build" / "custom_kernel.cu.o"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeUnixCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.os.getcwd',
return_value=str(test_dir),
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
return_value=False,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
return_value=False,
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={
'cxx': ['-Ithird_party/cxx', '-isystem', 'system/cxx'],
'nvcc': ['-Ithird_party/cuda', '-I', 'cuda_split'],
},
include_dirs=['public/include'],
)
content = captured['ninja_content']
self.assertIn('-Ipublic/include', content)
self.assertIn('-Ithird_party/cxx', content)
self.assertIn('-isystem system/cxx', content)
self.assertIn('-Ithird_party/cuda', content)
self.assertIn('-I cuda_split', content)
self.assertEqual(captured['ninja_work_directory'], str(test_dir))
def test_unix_ninja_uses_hipcc_for_rocm_sources(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.o"),
str(test_dir / "build" / "custom_kernel.cu.o"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeUnixCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.ROCM_HOME',
'/opt/rocm',
create=True,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
return_value=True,
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={
'cxx': ['-w'],
'hipcc': ['--hip-flag'],
},
include_dirs=[str(test_dir / "include")],
)
content = captured['ninja_content']
self.assertIn('nvcc = ', content)
self.assertIn('hipcc', content)
self.assertIn('cuda_post_cflags = --prepared --hip-flag', content)
self.assertIn('-D__HIP_PLATFORM_HCC__', content)
self.assertIn(
'-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_HIP', content
)
self.assertIn('-DPADDLE_WITH_HIP', content)
def test_unix_ninja_reports_missing_corex_compiler(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [str(test_dir / "custom_kernel.cu")]
objects = [str(test_dir / "build" / "custom_kernel.cu.o")]
build_map = {objects[0]: (sources[0], '.cu')}
compiler = _FakeUnixCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
return_value=False,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
return_value=True,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.os.path.isfile',
return_value=False,
),
self.assertRaisesRegex(ValueError, 'Corex compiler'),
):
self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={'nvcc': ['--gpu-flag']},
)
def test_use_ninja_attribute_default(self):
build_ext = self._build_extension()
self.assertEqual(build_ext.use_ninja, _is_ninja_available())
def test_use_ninja_attribute_explicit_false(self):
build_ext = self._build_extension(use_ninja=False)
self.assertFalse(build_ext.use_ninja)
def test_use_ninja_attribute_explicit_true(self):
build_ext = self._build_extension(use_ninja=True)
self.assertEqual(build_ext.use_ninja, _is_ninja_available())
def test_use_ninja_falls_back_when_ninja_missing(self):
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._is_ninja_available',
return_value=False,
),
mock.patch('builtins.print') as mock_print,
):
build_ext = self._build_extension(use_ninja=True)
self.assertFalse(build_ext.use_ninja)
mock_print.assert_any_call(
"Ninja is not available, falling back to the distutils backend."
)
def test_unix_compiler_with_use_ninja_false_executes_compile(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.o"),
str(test_dir / "build" / "custom_kernel.o"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeUnixCompiler(objects, build_map)
ext = SimpleNamespace(
name='fake_extension',
_full_name='fake_extension',
sources=sources,
extra_compile_args={'cxx': ['-w'], 'nvcc': ['--gpu-flag']},
)
cmd = self._build_extension(use_ninja=False)
cmd.extensions = [ext]
cmd.compiler = compiler
cmd.build_temp = tmpdir
cmd.build_lib = tmpdir
cmd.verbose = False
cmd.contain_cuda_file = True
patches = self._patch_build_extension_common(cmd)
with (
patches[0],
patches[1],
patches[2],
patches[3],
patches[4],
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_unix_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_ccache_home',
return_value=None,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_num_workers',
return_value=4,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.os.cpu_count',
return_value=8,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_rocm',
return_value=False,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.core.is_compiled_with_custom_device',
return_value=False,
),
mock.patch('builtins.print') as mock_print,
):
def fake_build_extensions(_):
compiled = _.compiler.compile(
sources,
output_dir=_.build_temp,
macros=[],
include_dirs=[str(test_dir / "include")],
debug=False,
extra_preargs=[],
extra_postargs=ext.extra_compile_args,
depends=None,
)
self.assertEqual(compiled, objects)
with mock.patch(
'paddle.utils.cpp_extension.cpp_extension.build_ext.build_extensions',
side_effect=fake_build_extensions,
):
cmd.build_extensions()
self.assertEqual(len(compiler._compile_calls), 2)
cc_call = next(
call
for call in compiler._compile_calls
if call['src'].endswith('.cc')
)
cu_call = next(
call
for call in compiler._compile_calls
if call['src'].endswith('.cu')
)
self.assertIn('-DPADDLE_WITH_CUDA', cc_call['cflags'])
self.assertIn('-D_GLIBCXX_USE_CXX11_ABI=1', cc_call['cflags'])
self.assertIn('--prepared', cu_call['cflags'])
self.assertEqual(cu_call['compiler_so'], self._expected_nvcc_path())
mock_print.assert_any_call(
"Using 2 workers for compilation. HINT: export MAX_JOBS=n to set the number of workers"
)
def test_msvc_compiler_with_use_ninja_false_executes_compile(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.obj"),
str(test_dir / "build" / "custom_kernel.obj"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeMsvcCompiler(objects, build_map)
ext = SimpleNamespace(
name='fake_extension',
_full_name='fake_extension',
sources=sources,
extra_compile_args={
'cxx': ['/wd4244'],
'nvcc': ['--gpu-flag'],
},
)
cmd = self._build_extension(use_ninja=False)
cmd.extensions = [ext]
cmd.compiler = compiler
cmd.build_temp = tmpdir
cmd.build_lib = tmpdir
cmd.verbose = False
cmd.contain_cuda_file = True
patches = self._patch_build_extension_common(cmd)
with (
patches[0],
patches[1],
patches[2],
patches[3],
patches[4],
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
):
def fake_build_extensions(_):
_.compiler.compile(
sources,
output_dir=_.build_temp,
macros=[],
include_dirs=[str(test_dir / "include dir")],
debug=False,
extra_preargs=[],
extra_postargs=ext.extra_compile_args,
depends=None,
)
with mock.patch(
'paddle.utils.cpp_extension.cpp_extension.build_ext.build_extensions',
side_effect=fake_build_extensions,
):
cmd.build_extensions()
self.assertEqual(len(compiler._spawned_cmds), 2)
cc_cmd = next(
cmd
for cmd in compiler._spawned_cmds
if any(arg.endswith('.cc') for arg in cmd)
)
cu_cmd = next(
cmd
for cmd in compiler._spawned_cmds
if any(arg.endswith('.cu') for arg in cmd)
)
self.assertIn('-DPADDLE_WITH_CUDA', cc_cmd)
self.assertIn('/wd4244', cc_cmd)
self.assertEqual(cu_cmd[0], self._expected_nvcc_path())
self.assertIn('--prepared', cu_cmd)
self.assertIn('--use-local-env', cu_cmd)
def test_windows_ninja_build_file_contains_cuda_sources(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.obj"),
str(test_dir / "build" / "custom_kernel.obj"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeMsvcCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={
'cxx': ['/wd4244'],
'nvcc': ['--gpu-flag'],
},
include_dirs=[str(test_dir / "include dir")],
)
content = captured['ninja_content']
self.assertIn(f'nvcc = {self._expected_nvcc_path()}', content)
self.assertIn('rule cuda_compile', content)
self.assertIn('-Xcompiler /EHsc', content)
self.assertIn('--prepared', content)
self.assertIn('/DPADDLE_WITH_CUDA', content)
self.assertIn('deps = msvc', content)
def test_windows_ninja_preserves_relative_include_flags(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.obj"),
str(test_dir / "build" / "custom_kernel.obj"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeMsvcCompiler(objects, build_map)
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.os.getcwd',
return_value=str(test_dir),
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args={
'cxx': ['/Ithird_party/cxx'],
'nvcc': ['/Ithird_party/cuda'],
},
include_dirs=['public/include'],
)
content = captured['ninja_content']
self.assertIn('/Ipublic/include', content)
self.assertIn('/Ithird_party/cxx', content)
self.assertIn('/Ithird_party/cuda', content)
self.assertEqual(captured['ninja_work_directory'], str(test_dir))
def test_windows_ninja_ignores_invalid_extra_compile_args(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_dir = Path(tmpdir)
sources = [
str(test_dir / "custom_extension.cc"),
str(test_dir / "custom_kernel.cu"),
]
objects = [
str(test_dir / "build" / "custom_extension.obj"),
str(test_dir / "build" / "custom_kernel.obj"),
]
build_map = {
objects[0]: (sources[0], '.cc'),
objects[1]: (sources[1], '.cu'),
}
compiler = _FakeMsvcCompiler(objects, build_map)
invalid_extra_args = 'not-a-list'
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.CUDA_HOME',
'/opt/cuda',
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.prepare_win_cudaflags',
side_effect=lambda flags: ['--prepared', *flags],
),
):
captured = self._run_build_with_fake_compiler(
compiler,
sources,
extra_compile_args=invalid_extra_args,
include_dirs=[str(test_dir / "include dir")],
)
content = captured['ninja_content']
self.assertIn('rule cuda_compile', content)
self.assertIn('cuda_post_cflags = ', content)
self.assertIn('-Xcompiler /EHsc', content)
self.assertIn('--prepared --use-local-env', content)
self.assertIn('/DPADDLE_WITH_CUDA', content)
self.assertNotIn(invalid_extra_args, content)
class TestNinjaGeneratedSetupFile(unittest.TestCase):
def test_load_setup_file_uses_default_build_extension_options(self):
with tempfile.TemporaryDirectory() as tmpdir:
setup_path = os.path.join(tmpdir, "demo_setup.py")
build_dir = os.path.join(tmpdir, "build")
sources = ['custom_extension.cc', 'custom_sub.cc']
_write_setup_file(
'demo_extension',
sources,
setup_path,
build_dir,
['include_dir'],
['library_dir'],
['-w', '-g'],
[],
[],
)
content = Path(setup_path).read_text(encoding='utf-8')
self.assertIn('BuildExtension.with_options(', content)
self.assertIn("output_dir=r'", content)
self.assertIn('no_python_abi_suffix=True', content)
self.assertIn(
"sources=['custom_extension.cc', 'custom_sub.cc']", content
)
self.assertNotIn('use_ninja=', content)
class TestRunNinjaBuild(unittest.TestCase):
def test_run_ninja_build_uses_verbose_subprocess_streams(self):
with (
mock.patch.dict(
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.subprocess.run'
) as mock_run,
):
_run_ninja_build('/tmp/build', verbose=True, error_prefix='Test')
_, kwargs = mock_run.call_args
self.assertEqual(mock_run.call_args.args[0], ['ninja', '-v'])
self.assertEqual(kwargs['cwd'], '/tmp/build')
self.assertIsNone(kwargs['stdout'])
self.assertIsNone(kwargs['stderr'])
def test_run_ninja_build_uses_num_workers_and_quiet_streams(self):
with (
mock.patch.dict(
os.environ,
{'MAX_JOBS': '3', 'VSCMD_ARG_TGT_ARCH': 'x64'},
clear=False,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.subprocess.run'
) as mock_run,
):
_run_ninja_build('/tmp/build', verbose=False, error_prefix='Test')
_, kwargs = mock_run.call_args
self.assertEqual(mock_run.call_args.args[0], ['ninja', '-v', '-j', '3'])
self.assertEqual(kwargs['stdout'], subprocess.PIPE)
self.assertEqual(kwargs['stderr'], subprocess.STDOUT)
self.assertTrue(kwargs['text'])
def test_run_ninja_build_windows_merges_vc_env(self):
vc_env = {'path': 'vc/bin', 'include': 'vc/include'}
with (
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.IS_WINDOWS', True
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.PLAT_TO_VCVARS',
{'win-amd64': 'x86_amd64'},
create=True,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.distutils',
ModuleSimpleNamespace(
util=ModuleSimpleNamespace(get_platform=lambda: 'win-amd64')
),
create=True,
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension._get_vc_env',
return_value=vc_env,
create=True,
),
mock.patch.dict(
os.environ, {'Path': 'user/bin', 'TMP': 'tmp'}, clear=True
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.subprocess.run'
) as mock_run,
):
_run_ninja_build('/tmp/build', verbose=False, error_prefix='Test')
env = mock_run.call_args.kwargs['env']
self.assertEqual(env['PATH'], 'vc/bin')
self.assertEqual(env['INCLUDE'], 'vc/include')
self.assertEqual(env['TMP'], 'tmp')
def test_run_ninja_build_raises_runtime_error(self):
with (
mock.patch.dict(
os.environ, {'VSCMD_ARG_TGT_ARCH': 'x64'}, clear=False
),
mock.patch(
'paddle.utils.cpp_extension.cpp_extension.subprocess.run',
side_effect=subprocess.CalledProcessError(1, ['ninja', '-v']),
),
self.assertRaisesRegex(RuntimeError, 'Prefix'),
):
_run_ninja_build('/tmp/build', verbose=False, error_prefix='Prefix')
def test_get_num_workers_verbose_without_env(self):
with (
mock.patch.dict(os.environ, {}, clear=True),
mock.patch('sys.stderr') as mock_stderr,
mock.patch('builtins.print') as mock_print,
):
result = _get_num_workers(verbose=True)
self.assertIsNone(result)
mock_print.assert_called_with(
'Allowing ninja to set a default number of workers... '
'(overridable by setting the environment variable MAX_JOBS=N)',
file=mock_stderr,
)
if __name__ == '__main__':
unittest.main()