mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
b12d645c2f
Summary: In the [official documentation](https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html), it is recommended to use `TORCH_LIBRARY` to register ops for TorchScript. However, that code is never tested with CUDA extension and is actually broken (https://github.com/pytorch/pytorch/issues/47493). This PR adds a test for it. It will not pass CI now, but it will pass when the issue https://github.com/pytorch/pytorch/issues/47493 is fixed. Pull Request resolved: https://github.com/pytorch/pytorch/pull/47524 Reviewed By: zou3519 Differential Revision: D24991839 Pulled By: ezyang fbshipit-source-id: 037196621c7ff9a6e7905efc1097ff97906a0b1c
90 lines
3.1 KiB
Python
90 lines
3.1 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:
|
|
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)
|
|
elif torch.cuda.is_available() and ROCM_HOME is not None:
|
|
from torch.utils.hipify import hipify_python
|
|
this_dir = os.path.dirname(os.path.abspath(__file__))
|
|
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.cuda', [
|
|
'cuda_extension.cpp',
|
|
'hip/hip_extension_kernel.hip',
|
|
'hip/hip_extension_kernel2.hip',
|
|
])
|
|
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)})
|