mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
10b490a3e0
Summary: This PR revamps the hipify module in PyTorch to overcome a long list of shortcomings in the original implementation. However, these improvements are applied only when using hipify to build PyTorch extensions, **not for PyTorch or Caffe2 itself**. Correspondingly, changes are made to `cpp_extension.py` to match these improvements. The list of improvements to hipify is as follows: 1. Hipify files in the same directory as the original file, unless there's a "cuda" subdirectory in the original file path, in which case the hipified file will be in the corresponding file path with "hip" subdirectory instead of "cuda". 2. Never hipify the file in-place if changes are introduced due to hipification i.e. always ensure the hipified file either resides in a different folder or has a different filename compared to the original file. 3. Prevent re-hipification of already hipified files. This avoids creation of unnecessary "hip/hip" etc. subdirectories and additional files which have no actual use. 4. Do not write out hipified versions of files if they are identical to the original file. This results in a cleaner output directory, with minimal number of hipified files created. 5. Update header rewrite logic so that it accounts for the previous improvement. 6. Update header rewrite logic so it respects the rules for finding header files depending on whether `""` or `<>` is used. 7. Return a dictionary of mappings of original file paths to hipified file paths from `hipify` function. 8. Introduce a version for hipify module to allow extensions to contain back-compatible code that targets a specific point in PyTorch where the hipify functionality changed. 9. Update `cuda_to_hip_mappings.py` to account for the ROCm component subdirectories inside `/opt/rocm/include`. This also results in cleanup of the `Caffe2_HIP_INCLUDE` path to remove unnecessary additions to the include path. The list of changes to `cpp_extension.py` is as follows: 1. Call `hipify` when building a CUDAExtension for ROCm. 2. Prune the list of source files to CUDAExtension to include only the hipified versions of any source files in the list (if both original and hipified versions of the source file are in the list) 3. Add subdirectories of /opt/rocm/include to the include path for extensions, so that ROCm headers for subcomponent libraries are found automatically cc jeffdaily sunway513 hgaspar lcskrishna ashishfarmer Pull Request resolved: https://github.com/pytorch/pytorch/pull/45451 Reviewed By: ezyang Differential Revision: D24924736 Pulled By: malfet fbshipit-source-id: 4af42b8ff4f21c3782dedb8719b8f9f86b34bd2d
74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
import sys
|
|
import torch.cuda
|
|
import os
|
|
from setuptools import setup
|
|
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
|
|
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
|
|
from torch.testing._internal.common_utils import IS_WINDOWS
|
|
|
|
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.msnpu', ['msnpu_extension.cpp'],
|
|
extra_compile_args=CXX_FLAGS),
|
|
CppExtension(
|
|
'torch_test_cpp_extension.rng', ['rng_extension.cpp'],
|
|
extra_compile_args=CXX_FLAGS),
|
|
]
|
|
|
|
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': ['-O2']})
|
|
ext_modules.append(extension)
|
|
|
|
if not IS_WINDOWS: # MSVC has bug compiling this example
|
|
if torch.cuda.is_available() and CUDA_HOME is not None:
|
|
extension = CUDAExtension(
|
|
'torch_test_cpp_extension.torch_library', [
|
|
'torch_library.cu'
|
|
],
|
|
extra_compile_args={'cxx': CXX_FLAGS,
|
|
'nvcc': ['-O2']})
|
|
ext_modules.append(extension)
|
|
elif torch.cuda.is_available() and ROCM_HOME is not None:
|
|
from torch.utils.hipify import hipify_python
|
|
hipify_python.hipify(
|
|
project_directory=this_dir,
|
|
output_directory=this_dir,
|
|
includes="./*",
|
|
show_detailed=True,
|
|
is_pytorch_extension=True,)
|
|
extension = CUDAExtension(
|
|
'torch_test_cpp_extension.torch_library', [
|
|
'hip/torch_library.hip'
|
|
],
|
|
extra_compile_args={'cxx': CXX_FLAGS,
|
|
'nvcc': ['-O2']})
|
|
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)})
|