mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
5f62308739
Summary: [Refiled version of earlier PR https://github.com/pytorch/pytorch/issues/45451] 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 ezyang Pull Request resolved: https://github.com/pytorch/pytorch/pull/48715 Reviewed By: bdhirsh Differential Revision: D25272824 Pulled By: ezyang fbshipit-source-id: 8bba68b27e41ca742781e1c4d7b07c6f985f040e
58 lines
2 KiB
Python
58 lines
2 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 or ROCM_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)
|
|
|
|
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)})
|