| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425 |
- # mypy: allow-untyped-defs
- import copy
- import glob
- import importlib
- import importlib.abc
- import os
- import re
- import shlex
- import shutil
- import setuptools
- import subprocess
- import sys
- import sysconfig
- import warnings
- import collections
- from pathlib import Path
- import errno
- import torch
- import torch._appdirs
- from .file_baton import FileBaton
- from ._cpp_extension_versioner import ExtensionVersioner
- from .hipify import hipify_python
- from .hipify.hipify_python import GeneratedFileCleaner
- from typing import Dict, List, Optional, Union, Tuple
- from torch.torch_version import TorchVersion, Version
- from setuptools.command.build_ext import build_ext
- IS_WINDOWS = sys.platform == 'win32'
- IS_MACOS = sys.platform.startswith('darwin')
- IS_LINUX = sys.platform.startswith('linux')
- LIB_EXT = '.pyd' if IS_WINDOWS else '.so'
- EXEC_EXT = '.exe' if IS_WINDOWS else ''
- CLIB_PREFIX = '' if IS_WINDOWS else 'lib'
- CLIB_EXT = '.dll' if IS_WINDOWS else '.so'
- SHARED_FLAG = '/DLL' if IS_WINDOWS else '-shared'
- _HERE = os.path.abspath(__file__)
- _TORCH_PATH = os.path.dirname(os.path.dirname(_HERE))
- TORCH_LIB_PATH = os.path.join(_TORCH_PATH, 'lib')
- SUBPROCESS_DECODE_ARGS = ('oem',) if IS_WINDOWS else ()
- MINIMUM_GCC_VERSION = (5, 0, 0)
- MINIMUM_MSVC_VERSION = (19, 0, 24215)
- VersionRange = Tuple[Tuple[int, ...], Tuple[int, ...]]
- VersionMap = Dict[str, VersionRange]
- # The following values were taken from the following GitHub gist that
- # summarizes the minimum valid major versions of g++/clang++ for each supported
- # CUDA version: https://gist.github.com/ax3l/9489132
- # Or from include/crt/host_config.h in the CUDA SDK
- # The second value is the exclusive(!) upper bound, i.e. min <= version < max
- CUDA_GCC_VERSIONS: VersionMap = {
- '11.0': (MINIMUM_GCC_VERSION, (10, 0)),
- '11.1': (MINIMUM_GCC_VERSION, (11, 0)),
- '11.2': (MINIMUM_GCC_VERSION, (11, 0)),
- '11.3': (MINIMUM_GCC_VERSION, (11, 0)),
- '11.4': ((6, 0, 0), (12, 0)),
- '11.5': ((6, 0, 0), (12, 0)),
- '11.6': ((6, 0, 0), (12, 0)),
- '11.7': ((6, 0, 0), (12, 0)),
- }
- MINIMUM_CLANG_VERSION = (3, 3, 0)
- CUDA_CLANG_VERSIONS: VersionMap = {
- '11.1': (MINIMUM_CLANG_VERSION, (11, 0)),
- '11.2': (MINIMUM_CLANG_VERSION, (12, 0)),
- '11.3': (MINIMUM_CLANG_VERSION, (12, 0)),
- '11.4': (MINIMUM_CLANG_VERSION, (13, 0)),
- '11.5': (MINIMUM_CLANG_VERSION, (13, 0)),
- '11.6': (MINIMUM_CLANG_VERSION, (14, 0)),
- '11.7': (MINIMUM_CLANG_VERSION, (14, 0)),
- }
- __all__ = ["get_default_build_root", "check_compiler_ok_for_platform", "get_compiler_abi_compatibility_and_version", "BuildExtension",
- "CppExtension", "CUDAExtension", "include_paths", "library_paths", "load", "load_inline", "is_ninja_available",
- "verify_ninja_availability", "remove_extension_h_precompiler_headers", "get_cxx_compiler", "check_compiler_is_gcc"]
- # Taken directly from python stdlib < 3.9
- # See https://github.com/pytorch/pytorch/issues/48617
- def _nt_quote_args(args: Optional[List[str]]) -> List[str]:
- """Quote command-line arguments for DOS/Windows conventions.
- Just wraps every argument which contains blanks in double quotes, and
- returns a new argument list.
- """
- # Cover None-type
- if not args:
- return []
- return [f'"{arg}"' if ' ' in arg else arg for arg in args]
- def _find_cuda_home() -> Optional[str]:
- """Find the CUDA install path."""
- # Guess #1
- cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
- if cuda_home is None:
- # Guess #2
- nvcc_path = shutil.which("nvcc")
- if nvcc_path is not None:
- cuda_home = os.path.dirname(os.path.dirname(nvcc_path))
- else:
- # Guess #3
- if IS_WINDOWS:
- cuda_homes = glob.glob(
- 'C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*')
- if len(cuda_homes) == 0:
- cuda_home = ''
- else:
- cuda_home = cuda_homes[0]
- else:
- cuda_home = '/usr/local/cuda'
- if not os.path.exists(cuda_home):
- cuda_home = None
- if cuda_home and not torch.cuda.is_available():
- print(f"No CUDA runtime is found, using CUDA_HOME='{cuda_home}'",
- file=sys.stderr)
- return cuda_home
- def _find_rocm_home() -> Optional[str]:
- """Find the ROCm install path."""
- # Guess #1
- rocm_home = os.environ.get('ROCM_HOME') or os.environ.get('ROCM_PATH')
- if rocm_home is None:
- # Guess #2
- hipcc_path = shutil.which('hipcc')
- if hipcc_path is not None:
- rocm_home = os.path.dirname(os.path.dirname(
- os.path.realpath(hipcc_path)))
- # can be either <ROCM_HOME>/hip/bin/hipcc or <ROCM_HOME>/bin/hipcc
- if os.path.basename(rocm_home) == 'hip':
- rocm_home = os.path.dirname(rocm_home)
- else:
- # Guess #3
- fallback_path = '/opt/rocm'
- if os.path.exists(fallback_path):
- rocm_home = fallback_path
- if rocm_home and torch.version.hip is None:
- print(f"No ROCm runtime is found, using ROCM_HOME='{rocm_home}'",
- file=sys.stderr)
- return rocm_home
- def _join_rocm_home(*paths) -> str:
- """
- Join paths with ROCM_HOME, or raises an error if it ROCM_HOME is not set.
- This is basically a lazy way of raising an error for missing $ROCM_HOME
- only once we need to get any ROCm-specific path.
- """
- if ROCM_HOME is None:
- raise OSError('ROCM_HOME environment variable is not set. '
- 'Please set it to your ROCm install root.')
- elif IS_WINDOWS:
- raise OSError('Building PyTorch extensions using '
- 'ROCm and Windows is not supported.')
- return os.path.join(ROCM_HOME, *paths)
- ABI_INCOMPATIBILITY_WARNING = '''
- !! WARNING !!
- !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
- Your compiler ({}) may be ABI-incompatible with PyTorch!
- Please use a compiler that is ABI-compatible with GCC 5.0 and above.
- See https://gcc.gnu.org/onlinedocs/libstdc++/manual/abi.html.
- See https://gist.github.com/goldsborough/d466f43e8ffc948ff92de7486c5216d6
- for instructions on how to install GCC 5 or higher.
- !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
- !! WARNING !!
- '''
- WRONG_COMPILER_WARNING = '''
- !! WARNING !!
- !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
- Your compiler ({user_compiler}) is not compatible with the compiler Pytorch was
- built with for this platform, which is {pytorch_compiler} on {platform}. Please
- use {pytorch_compiler} to to compile your extension. Alternatively, you may
- compile PyTorch from source using {user_compiler}, and then you can also use
- {user_compiler} to compile your extension.
- See https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md for help
- with compiling PyTorch from source.
- !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
- !! WARNING !!
- '''
- CUDA_MISMATCH_MESSAGE = '''
- The detected CUDA version ({0}) mismatches the version that was used to compile
- PyTorch ({1}). Please make sure to use the same CUDA versions.
- '''
- CUDA_MISMATCH_WARN = "The detected CUDA version ({0}) has a minor version mismatch with the version that was used to compile PyTorch ({1}). Most likely this shouldn't be a problem."
- CUDA_NOT_FOUND_MESSAGE = '''
- CUDA was not found on the system, please set the CUDA_HOME or the CUDA_PATH
- environment variable or add NVCC to your system PATH. The extension compilation will fail.
- '''
- ROCM_HOME = _find_rocm_home()
- HIP_HOME = _join_rocm_home('hip') if ROCM_HOME else None
- IS_HIP_EXTENSION = True if ((ROCM_HOME is not None) and (torch.version.hip is not None)) else False
- ROCM_VERSION = None
- if torch.version.hip is not None:
- ROCM_VERSION = tuple(int(v) for v in torch.version.hip.split('.')[:2])
- CUDA_HOME = _find_cuda_home() if torch.cuda._is_compiled() else None
- CUDNN_HOME = os.environ.get('CUDNN_HOME') or os.environ.get('CUDNN_PATH')
- # PyTorch releases have the version pattern major.minor.patch, whereas when
- # PyTorch is built from source, we append the git commit hash, which gives
- # it the below pattern.
- BUILT_FROM_SOURCE_VERSION_PATTERN = re.compile(r'\d+\.\d+\.\d+\w+\+\w+')
- COMMON_MSVC_FLAGS = ['/MD', '/wd4819', '/wd4251', '/wd4244', '/wd4267', '/wd4275', '/wd4018', '/wd4190', '/wd4624', '/wd4067', '/wd4068', '/EHsc']
- MSVC_IGNORE_CUDAFE_WARNINGS = [
- 'base_class_has_different_dll_interface',
- 'field_without_dll_interface',
- 'dll_interface_conflict_none_assumed',
- 'dll_interface_conflict_dllexport_assumed'
- ]
- COMMON_NVCC_FLAGS = [
- '-D__CUDA_NO_HALF_OPERATORS__',
- '-D__CUDA_NO_HALF_CONVERSIONS__',
- '-D__CUDA_NO_BFLOAT16_CONVERSIONS__',
- '-D__CUDA_NO_HALF2_OPERATORS__',
- '--expt-relaxed-constexpr'
- ]
- COMMON_HIP_FLAGS = [
- '-fPIC',
- '-D__HIP_PLATFORM_AMD__=1',
- '-DUSE_ROCM=1',
- '-DHIPBLAS_V2',
- ]
- COMMON_HIPCC_FLAGS = [
- '-DCUDA_HAS_FP16=1',
- '-D__HIP_NO_HALF_OPERATORS__=1',
- '-D__HIP_NO_HALF_CONVERSIONS__=1',
- ]
- JIT_EXTENSION_VERSIONER = ExtensionVersioner()
- PLAT_TO_VCVARS = {
- 'win32' : 'x86',
- 'win-amd64' : 'x86_amd64',
- }
- def get_cxx_compiler():
- if IS_WINDOWS:
- compiler = os.environ.get('CXX', 'cl')
- else:
- compiler = os.environ.get('CXX', 'c++')
- return compiler
- def _is_binary_build() -> bool:
- return not BUILT_FROM_SOURCE_VERSION_PATTERN.match(torch.version.__version__)
- def _accepted_compilers_for_platform() -> List[str]:
- # gnu-c++ and gnu-cc are the conda gcc compilers
- return ['clang++', 'clang'] if IS_MACOS else ['g++', 'gcc', 'gnu-c++', 'gnu-cc', 'clang++', 'clang']
- def _maybe_write(filename, new_content):
- r'''
- Equivalent to writing the content into the file but will not touch the file
- if it already had the right content (to avoid triggering recompile).
- '''
- if os.path.exists(filename):
- with open(filename) as f:
- content = f.read()
- if content == new_content:
- # The file already contains the right thing!
- return
- with open(filename, 'w') as source_file:
- source_file.write(new_content)
- def get_default_build_root() -> str:
- """
- Return the path to the root folder under which extensions will built.
- For each extension module built, there will be one folder underneath the
- folder returned by this function. For example, if ``p`` is the path
- returned by this function and ``ext`` the name of an extension, the build
- folder for the extension will be ``p/ext``.
- This directory is **user-specific** so that multiple users on the same
- machine won't meet permission issues.
- """
- return os.path.realpath(torch._appdirs.user_cache_dir(appname='torch_extensions'))
- def check_compiler_ok_for_platform(compiler: str) -> bool:
- """
- Verify that the compiler is the expected one for the current platform.
- Args:
- compiler (str): The compiler executable to check.
- Returns:
- True if the compiler is gcc/g++ on Linux or clang/clang++ on macOS,
- and always True for Windows.
- """
- if IS_WINDOWS:
- return True
- which = subprocess.check_output(['which', compiler], stderr=subprocess.STDOUT)
- # Use os.path.realpath to resolve any symlinks, in particular from 'c++' to e.g. 'g++'.
- compiler_path = os.path.realpath(which.decode(*SUBPROCESS_DECODE_ARGS).strip())
- # Check the compiler name
- if any(name in compiler_path for name in _accepted_compilers_for_platform()):
- return True
- # If compiler wrapper is used try to infer the actual compiler by invoking it with -v flag
- env = os.environ.copy()
- env['LC_ALL'] = 'C' # Don't localize output
- version_string = subprocess.check_output([compiler, '-v'], stderr=subprocess.STDOUT, env=env).decode(*SUBPROCESS_DECODE_ARGS)
- if IS_LINUX:
- # Check for 'gcc' or 'g++' for sccache wrapper
- pattern = re.compile("^COLLECT_GCC=(.*)$", re.MULTILINE)
- results = re.findall(pattern, version_string)
- if len(results) != 1:
- # Clang is also a supported compiler on Linux
- # Though on Ubuntu it's sometimes called "Ubuntu clang version"
- return 'clang version' in version_string
- compiler_path = os.path.realpath(results[0].strip())
- # On RHEL/CentOS c++ is a gcc compiler wrapper
- if os.path.basename(compiler_path) == 'c++' and 'gcc version' in version_string:
- return True
- return any(name in compiler_path for name in _accepted_compilers_for_platform())
- if IS_MACOS:
- # Check for 'clang' or 'clang++'
- return version_string.startswith("Apple clang")
- return False
- def get_compiler_abi_compatibility_and_version(compiler) -> Tuple[bool, TorchVersion]:
- """
- Determine if the given compiler is ABI-compatible with PyTorch alongside its version.
- Args:
- compiler (str): The compiler executable name to check (e.g. ``g++``).
- Must be executable in a shell process.
- Returns:
- A tuple that contains a boolean that defines if the compiler is (likely) ABI-incompatible with PyTorch,
- followed by a `TorchVersion` string that contains the compiler version separated by dots.
- """
- if not _is_binary_build():
- return (True, TorchVersion('0.0.0'))
- if os.environ.get('TORCH_DONT_CHECK_COMPILER_ABI') in ['ON', '1', 'YES', 'TRUE', 'Y']:
- return (True, TorchVersion('0.0.0'))
- # First check if the compiler is one of the expected ones for the particular platform.
- if not check_compiler_ok_for_platform(compiler):
- warnings.warn(WRONG_COMPILER_WARNING.format(
- user_compiler=compiler,
- pytorch_compiler=_accepted_compilers_for_platform()[0],
- platform=sys.platform))
- return (False, TorchVersion('0.0.0'))
- if IS_MACOS:
- # There is no particular minimum version we need for clang, so we're good here.
- return (True, TorchVersion('0.0.0'))
- try:
- if IS_LINUX:
- minimum_required_version = MINIMUM_GCC_VERSION
- versionstr = subprocess.check_output([compiler, '-dumpfullversion', '-dumpversion'])
- version = versionstr.decode(*SUBPROCESS_DECODE_ARGS).strip().split('.')
- else:
- minimum_required_version = MINIMUM_MSVC_VERSION
- compiler_info = subprocess.check_output(compiler, stderr=subprocess.STDOUT)
- match = re.search(r'(\d+)\.(\d+)\.(\d+)', compiler_info.decode(*SUBPROCESS_DECODE_ARGS).strip())
- version = ['0', '0', '0'] if match is None else list(match.groups())
- except Exception:
- _, error, _ = sys.exc_info()
- warnings.warn(f'Error checking compiler version for {compiler}: {error}')
- return (False, TorchVersion('0.0.0'))
- if tuple(map(int, version)) >= minimum_required_version:
- return (True, TorchVersion('.'.join(version)))
- compiler = f'{compiler} {".".join(version)}'
- warnings.warn(ABI_INCOMPATIBILITY_WARNING.format(compiler))
- return (False, TorchVersion('.'.join(version)))
- def _check_cuda_version(compiler_name: str, compiler_version: TorchVersion) -> None:
- if not CUDA_HOME:
- raise RuntimeError(CUDA_NOT_FOUND_MESSAGE)
- nvcc = os.path.join(CUDA_HOME, 'bin', 'nvcc')
- cuda_version_str = subprocess.check_output([nvcc, '--version']).strip().decode(*SUBPROCESS_DECODE_ARGS)
- cuda_version = re.search(r'release (\d+[.]\d+)', cuda_version_str)
- if cuda_version is None:
- return
- cuda_str_version = cuda_version.group(1)
- cuda_ver = Version(cuda_str_version)
- if torch.version.cuda is None:
- return
- torch_cuda_version = Version(torch.version.cuda)
- if cuda_ver != torch_cuda_version:
- # major/minor attributes are only available in setuptools>=49.4.0
- if getattr(cuda_ver, "major", None) is None:
- raise ValueError("setuptools>=49.4.0 is required")
- if cuda_ver.major != torch_cuda_version.major:
- raise RuntimeError(CUDA_MISMATCH_MESSAGE.format(cuda_str_version, torch.version.cuda))
- warnings.warn(CUDA_MISMATCH_WARN.format(cuda_str_version, torch.version.cuda))
- if not (sys.platform.startswith('linux') and
- os.environ.get('TORCH_DONT_CHECK_COMPILER_ABI') not in ['ON', '1', 'YES', 'TRUE', 'Y'] and
- _is_binary_build()):
- return
- cuda_compiler_bounds: VersionMap = CUDA_CLANG_VERSIONS if compiler_name.startswith('clang') else CUDA_GCC_VERSIONS
- if cuda_str_version not in cuda_compiler_bounds:
- warnings.warn(f'There are no {compiler_name} version bounds defined for CUDA version {cuda_str_version}')
- else:
- min_compiler_version, max_excl_compiler_version = cuda_compiler_bounds[cuda_str_version]
- # Special case for 11.4.0, which has lower compiler bounds than 11.4.1
- if "V11.4.48" in cuda_version_str and cuda_compiler_bounds == CUDA_GCC_VERSIONS:
- max_excl_compiler_version = (11, 0)
- min_compiler_version_str = '.'.join(map(str, min_compiler_version))
- max_excl_compiler_version_str = '.'.join(map(str, max_excl_compiler_version))
- version_bound_str = f'>={min_compiler_version_str}, <{max_excl_compiler_version_str}'
- if compiler_version < TorchVersion(min_compiler_version_str):
- raise RuntimeError(
- f'The current installed version of {compiler_name} ({compiler_version}) is less '
- f'than the minimum required version by CUDA {cuda_str_version} ({min_compiler_version_str}). '
- f'Please make sure to use an adequate version of {compiler_name} ({version_bound_str}).'
- )
- if compiler_version >= TorchVersion(max_excl_compiler_version_str):
- raise RuntimeError(
- f'The current installed version of {compiler_name} ({compiler_version}) is greater '
- f'than the maximum required version by CUDA {cuda_str_version}. '
- f'Please make sure to use an adequate version of {compiler_name} ({version_bound_str}).'
- )
- class BuildExtension(build_ext):
- """
- A custom :mod:`setuptools` build extension .
- This :class:`setuptools.build_ext` subclass takes care of passing the
- minimum required compiler flags (e.g. ``-std=c++17``) as well as mixed
- C++/CUDA compilation (and support for CUDA files in general).
- When using :class:`BuildExtension`, it is allowed to supply a dictionary
- for ``extra_compile_args`` (rather than the usual list) that maps from
- languages (``cxx`` or ``nvcc``) to a list of additional compiler flags to
- supply to the compiler. This makes it possible to supply different flags to
- the C++ and CUDA compiler during mixed compilation.
- ``use_ninja`` (bool): If ``use_ninja`` is ``True`` (default), then we
- attempt to build using the Ninja backend. Ninja greatly speeds up
- compilation compared to the standard ``setuptools.build_ext``.
- Fallbacks to the standard distutils backend if Ninja is not available.
- .. note::
- By default, the Ninja backend uses #CPUS + 2 workers to build the
- extension. This may use up too many resources on some systems. One
- can control the number of workers by setting the `MAX_JOBS` environment
- variable to a non-negative number.
- """
- @classmethod
- def with_options(cls, **options):
- """Return a subclass with alternative constructor that extends any original keyword arguments to the original constructor with the given options."""
- class cls_with_options(cls): # type: ignore[misc, valid-type]
- def __init__(self, *args, **kwargs):
- kwargs.update(options)
- super().__init__(*args, **kwargs)
- return cls_with_options
- def __init__(self, *args, **kwargs) -> None:
- super().__init__(*args, **kwargs)
- self.no_python_abi_suffix = kwargs.get("no_python_abi_suffix", False)
- self.use_ninja = kwargs.get('use_ninja', True)
- if self.use_ninja:
- # Test if we can use ninja. Fallback otherwise.
- msg = ('Attempted to use ninja as the BuildExtension backend but '
- '{}. Falling back to using the slow distutils backend.')
- if not is_ninja_available():
- warnings.warn(msg.format('we could not find ninja.'))
- self.use_ninja = False
- def finalize_options(self) -> None:
- super().finalize_options()
- if self.use_ninja:
- self.force = True
- def build_extensions(self) -> None:
- compiler_name, compiler_version = self._check_abi()
- cuda_ext = False
- extension_iter = iter(self.extensions)
- extension = next(extension_iter, None)
- while not cuda_ext and extension:
- for source in extension.sources:
- _, ext = os.path.splitext(source)
- if ext == '.cu':
- cuda_ext = True
- break
- extension = next(extension_iter, None)
- if cuda_ext and not IS_HIP_EXTENSION:
- _check_cuda_version(compiler_name, compiler_version)
- for extension in self.extensions:
- # Ensure at least an empty list of flags for 'cxx' and 'nvcc' when
- # extra_compile_args is a dict. Otherwise, default torch flags do
- # not get passed. Necessary when only one of 'cxx' and 'nvcc' is
- # passed to extra_compile_args in CUDAExtension, i.e.
- # CUDAExtension(..., extra_compile_args={'cxx': [...]})
- # or
- # CUDAExtension(..., extra_compile_args={'nvcc': [...]})
- if isinstance(extension.extra_compile_args, dict):
- for ext in ['cxx', 'nvcc']:
- if ext not in extension.extra_compile_args:
- extension.extra_compile_args[ext] = []
- self._add_compile_flag(extension, '-DTORCH_API_INCLUDE_EXTENSION_H')
- # See note [Pybind11 ABI constants]
- for name in ["COMPILER_TYPE", "STDLIB", "BUILD_ABI"]:
- val = getattr(torch._C, f"_PYBIND11_{name}")
- if val is not None and not IS_WINDOWS:
- self._add_compile_flag(extension, f'-DPYBIND11_{name}="{val}"')
- self._define_torch_extension_name(extension)
- self._add_gnu_cpp_abi_flag(extension)
- if 'nvcc_dlink' in extension.extra_compile_args:
- assert self.use_ninja, f"With dlink=True, ninja is required to build cuda extension {extension.name}."
- # Register .cu, .cuh, .hip, and .mm as valid source extensions.
- self.compiler.src_extensions += ['.cu', '.cuh', '.hip']
- if torch.backends.mps.is_built():
- self.compiler.src_extensions += ['.mm']
- # Save the original _compile method for later.
- if self.compiler.compiler_type == 'msvc':
- self.compiler._cpp_extensions += ['.cu', '.cuh']
- original_compile = self.compiler.compile
- original_spawn = self.compiler.spawn
- else:
- original_compile = self.compiler._compile
- def append_std17_if_no_std_present(cflags) -> None:
- # NVCC does not allow multiple -std to be passed, so we avoid
- # overriding the option if the user explicitly passed it.
- cpp_format_prefix = '/{}:' if self.compiler.compiler_type == 'msvc' else '-{}='
- cpp_flag_prefix = cpp_format_prefix.format('std')
- cpp_flag = cpp_flag_prefix + 'c++17'
- if not any(flag.startswith(cpp_flag_prefix) for flag in cflags):
- cflags.append(cpp_flag)
- def unix_cuda_flags(cflags):
- cflags = (COMMON_NVCC_FLAGS +
- ['--compiler-options', "'-fPIC'"] +
- cflags + _get_cuda_arch_flags(cflags))
- # NVCC does not allow multiple -ccbin/--compiler-bindir to be passed, so we avoid
- # overriding the option if the user explicitly passed it.
- _ccbin = os.getenv("CC")
- if (
- _ccbin is not None
- and not any(flag.startswith(('-ccbin', '--compiler-bindir')) for flag in cflags)
- ):
- cflags.extend(['-ccbin', _ccbin])
- return cflags
- def convert_to_absolute_paths_inplace(paths):
- # Helper function. See Note [Absolute include_dirs]
- if paths is not None:
- for i in range(len(paths)):
- if not os.path.isabs(paths[i]):
- paths[i] = os.path.abspath(paths[i])
- def unix_wrap_single_compile(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None:
- # Copy before we make any modifications.
- cflags = copy.deepcopy(extra_postargs)
- try:
- original_compiler = self.compiler.compiler_so
- if _is_cuda_file(src):
- nvcc = [_join_rocm_home('bin', 'hipcc') if IS_HIP_EXTENSION else _join_cuda_home('bin', 'nvcc')]
- self.compiler.set_executable('compiler_so', nvcc)
- if isinstance(cflags, dict):
- cflags = cflags['nvcc']
- if IS_HIP_EXTENSION:
- cflags = COMMON_HIPCC_FLAGS + cflags + _get_rocm_arch_flags(cflags)
- else:
- cflags = unix_cuda_flags(cflags)
- elif isinstance(cflags, dict):
- cflags = cflags['cxx']
- if IS_HIP_EXTENSION:
- cflags = COMMON_HIP_FLAGS + cflags
- append_std17_if_no_std_present(cflags)
- original_compile(obj, src, ext, cc_args, cflags, pp_opts)
- finally:
- # Put the original compiler back in place.
- self.compiler.set_executable('compiler_so', original_compiler)
- def unix_wrap_ninja_compile(sources,
- output_dir=None,
- macros=None,
- include_dirs=None,
- debug=0,
- extra_preargs=None,
- extra_postargs=None,
- depends=None):
- r"""Compiles sources by outputting a ninja file and running it."""
- # NB: I copied some lines from self.compiler (which is an instance
- # of distutils.UnixCCompiler). See the following link.
- # https://github.com/python/cpython/blob/f03a8f8d5001963ad5b5b28dbd95497e9cc15596/Lib/distutils/ccompiler.py#L564-L567
- # This can be fragile, but a lot of other repos also do this
- # (see https://github.com/search?q=_setup_compile&type=Code)
- # so it is probably OK; we'll also get CI signal if/when
- # we update our python version (which is when distutils can be
- # upgraded)
- # Use absolute path for output_dir so that the object file paths
- # (`objects`) get generated with absolute paths.
- output_dir = os.path.abspath(output_dir)
- # See Note [Absolute include_dirs]
- convert_to_absolute_paths_inplace(self.compiler.include_dirs)
- _, objects, extra_postargs, pp_opts, _ = \
- self.compiler._setup_compile(output_dir, macros,
- include_dirs, sources,
- depends, extra_postargs)
- common_cflags = self.compiler._get_cc_args(pp_opts, debug, extra_preargs)
- extra_cc_cflags = self.compiler.compiler_so[1:]
- with_cuda = any(map(_is_cuda_file, sources))
- # extra_postargs can be either:
- # - a dict mapping cxx/nvcc to extra flags
- # - a list of extra flags.
- if isinstance(extra_postargs, dict):
- post_cflags = extra_postargs['cxx']
- else:
- post_cflags = list(extra_postargs)
- if IS_HIP_EXTENSION:
- post_cflags = COMMON_HIP_FLAGS + post_cflags
- append_std17_if_no_std_present(post_cflags)
- cuda_post_cflags = None
- cuda_cflags = None
- if with_cuda:
- cuda_cflags = common_cflags
- if isinstance(extra_postargs, dict):
- cuda_post_cflags = extra_postargs['nvcc']
- else:
- cuda_post_cflags = list(extra_postargs)
- if IS_HIP_EXTENSION:
- cuda_post_cflags = cuda_post_cflags + _get_rocm_arch_flags(cuda_post_cflags)
- cuda_post_cflags = COMMON_HIP_FLAGS + COMMON_HIPCC_FLAGS + cuda_post_cflags
- else:
- cuda_post_cflags = unix_cuda_flags(cuda_post_cflags)
- append_std17_if_no_std_present(cuda_post_cflags)
- cuda_cflags = [shlex.quote(f) for f in cuda_cflags]
- cuda_post_cflags = [shlex.quote(f) for f in cuda_post_cflags]
- if isinstance(extra_postargs, dict) and 'nvcc_dlink' in extra_postargs:
- cuda_dlink_post_cflags = unix_cuda_flags(extra_postargs['nvcc_dlink'])
- else:
- cuda_dlink_post_cflags = None
- _write_ninja_file_and_compile_objects(
- sources=sources,
- objects=objects,
- cflags=[shlex.quote(f) for f in extra_cc_cflags + common_cflags],
- post_cflags=[shlex.quote(f) for f in post_cflags],
- cuda_cflags=cuda_cflags,
- cuda_post_cflags=cuda_post_cflags,
- cuda_dlink_post_cflags=cuda_dlink_post_cflags,
- build_directory=output_dir,
- verbose=True,
- with_cuda=with_cuda)
- # Return *all* object filenames, not just the ones we just built.
- return objects
- def win_cuda_flags(cflags):
- return (COMMON_NVCC_FLAGS +
- cflags + _get_cuda_arch_flags(cflags))
- def win_wrap_single_compile(sources,
- output_dir=None,
- macros=None,
- include_dirs=None,
- debug=0,
- extra_preargs=None,
- extra_postargs=None,
- depends=None):
- self.cflags = copy.deepcopy(extra_postargs)
- extra_postargs = None
- def spawn(cmd):
- # Using regex to match src, obj and include files
- src_regex = re.compile('/T(p|c)(.*)')
- src_list = [
- m.group(2) for m in (src_regex.match(elem) for elem in cmd)
- if m
- ]
- obj_regex = re.compile('/Fo(.*)')
- obj_list = [
- m.group(1) for m in (obj_regex.match(elem) for elem in cmd)
- if m
- ]
- include_regex = re.compile(r'((\-|\/)I.*)')
- include_list = [
- m.group(1)
- for m in (include_regex.match(elem) for elem in cmd) if m
- ]
- if len(src_list) >= 1 and len(obj_list) >= 1:
- src = src_list[0]
- obj = obj_list[0]
- if _is_cuda_file(src):
- nvcc = _join_cuda_home('bin', 'nvcc')
- if isinstance(self.cflags, dict):
- cflags = self.cflags['nvcc']
- elif isinstance(self.cflags, list):
- cflags = self.cflags
- else:
- cflags = []
- cflags = win_cuda_flags(cflags) + ['-std=c++17', '--use-local-env']
- for flag in COMMON_MSVC_FLAGS:
- cflags = ['-Xcompiler', flag] + cflags
- for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
- cflags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cflags
- cmd = [nvcc, '-c', src, '-o', obj] + include_list + cflags
- elif isinstance(self.cflags, dict):
- cflags = COMMON_MSVC_FLAGS + self.cflags['cxx']
- append_std17_if_no_std_present(cflags)
- cmd += cflags
- elif isinstance(self.cflags, list):
- cflags = COMMON_MSVC_FLAGS + self.cflags
- append_std17_if_no_std_present(cflags)
- cmd += cflags
- return original_spawn(cmd)
- try:
- self.compiler.spawn = spawn
- return original_compile(sources, output_dir, macros,
- include_dirs, debug, extra_preargs,
- extra_postargs, depends)
- finally:
- self.compiler.spawn = original_spawn
- def win_wrap_ninja_compile(sources,
- output_dir=None,
- macros=None,
- include_dirs=None,
- debug=0,
- extra_preargs=None,
- extra_postargs=None,
- depends=None):
- if not self.compiler.initialized:
- self.compiler.initialize()
- output_dir = os.path.abspath(output_dir)
- # Note [Absolute include_dirs]
- # Convert relative path in self.compiler.include_dirs to absolute path if any,
- # For ninja build, the build location is not local, the build happens
- # in a in script created build folder, relative path lost their correctness.
- # To be consistent with jit extension, we allow user to enter relative include_dirs
- # in setuptools.setup, and we convert the relative path to absolute path here
- convert_to_absolute_paths_inplace(self.compiler.include_dirs)
- _, objects, extra_postargs, pp_opts, _ = \
- self.compiler._setup_compile(output_dir, macros,
- include_dirs, sources,
- depends, extra_postargs)
- common_cflags = extra_preargs or []
- cflags = []
- if debug:
- cflags.extend(self.compiler.compile_options_debug)
- else:
- cflags.extend(self.compiler.compile_options)
- common_cflags.extend(COMMON_MSVC_FLAGS)
- cflags = cflags + common_cflags + pp_opts
- with_cuda = any(map(_is_cuda_file, sources))
- # extra_postargs can be either:
- # - a dict mapping cxx/nvcc to extra flags
- # - a list of extra flags.
- if isinstance(extra_postargs, dict):
- post_cflags = extra_postargs['cxx']
- else:
- post_cflags = list(extra_postargs)
- append_std17_if_no_std_present(post_cflags)
- cuda_post_cflags = None
- cuda_cflags = None
- if with_cuda:
- cuda_cflags = ['-std=c++17', '--use-local-env']
- for common_cflag in common_cflags:
- cuda_cflags.append('-Xcompiler')
- cuda_cflags.append(common_cflag)
- for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
- cuda_cflags.append('-Xcudafe')
- cuda_cflags.append('--diag_suppress=' + ignore_warning)
- cuda_cflags.extend(pp_opts)
- if isinstance(extra_postargs, dict):
- cuda_post_cflags = extra_postargs['nvcc']
- else:
- cuda_post_cflags = list(extra_postargs)
- cuda_post_cflags = win_cuda_flags(cuda_post_cflags)
- cflags = _nt_quote_args(cflags)
- post_cflags = _nt_quote_args(post_cflags)
- if with_cuda:
- cuda_cflags = _nt_quote_args(cuda_cflags)
- cuda_post_cflags = _nt_quote_args(cuda_post_cflags)
- if isinstance(extra_postargs, dict) and 'nvcc_dlink' in extra_postargs:
- cuda_dlink_post_cflags = win_cuda_flags(extra_postargs['nvcc_dlink'])
- else:
- cuda_dlink_post_cflags = None
- _write_ninja_file_and_compile_objects(
- sources=sources,
- objects=objects,
- cflags=cflags,
- post_cflags=post_cflags,
- cuda_cflags=cuda_cflags,
- cuda_post_cflags=cuda_post_cflags,
- cuda_dlink_post_cflags=cuda_dlink_post_cflags,
- build_directory=output_dir,
- verbose=True,
- with_cuda=with_cuda)
- # Return *all* object filenames, not just the ones we just built.
- return objects
- # Monkey-patch the _compile or compile method.
- # https://github.com/python/cpython/blob/dc0284ee8f7a270b6005467f26d8e5773d76e959/Lib/distutils/ccompiler.py#L511
- if self.compiler.compiler_type == 'msvc':
- if self.use_ninja:
- self.compiler.compile = win_wrap_ninja_compile
- else:
- self.compiler.compile = win_wrap_single_compile
- else:
- if self.use_ninja:
- self.compiler.compile = unix_wrap_ninja_compile
- else:
- self.compiler._compile = unix_wrap_single_compile
- build_ext.build_extensions(self)
- def get_ext_filename(self, ext_name):
- # Get the original shared library name. For Python 3, this name will be
- # suffixed with "<SOABI>.so", where <SOABI> will be something like
- # cpython-37m-x86_64-linux-gnu.
- ext_filename = super().get_ext_filename(ext_name)
- # If `no_python_abi_suffix` is `True`, we omit the Python 3 ABI
- # component. This makes building shared libraries with setuptools that
- # aren't Python modules nicer.
- if self.no_python_abi_suffix:
- # The parts will be e.g. ["my_extension", "cpython-37m-x86_64-linux-gnu", "so"].
- ext_filename_parts = ext_filename.split('.')
- # Omit the second to last element.
- without_abi = ext_filename_parts[:-2] + ext_filename_parts[-1:]
- ext_filename = '.'.join(without_abi)
- return ext_filename
- def _check_abi(self) -> Tuple[str, TorchVersion]:
- # On some platforms, like Windows, compiler_cxx is not available.
- if hasattr(self.compiler, 'compiler_cxx'):
- compiler = self.compiler.compiler_cxx[0]
- else:
- compiler = get_cxx_compiler()
- _, version = get_compiler_abi_compatibility_and_version(compiler)
- # Warn user if VC env is activated but `DISTUILS_USE_SDK` is not set.
- if IS_WINDOWS and 'VSCMD_ARG_TGT_ARCH' in os.environ and 'DISTUTILS_USE_SDK' not in os.environ:
- msg = ('It seems that the VC environment is activated but DISTUTILS_USE_SDK is not set.'
- 'This may lead to multiple activations of the VC env.'
- 'Please set `DISTUTILS_USE_SDK=1` and try again.')
- raise UserWarning(msg)
- return compiler, version
- def _add_compile_flag(self, extension, flag):
- extension.extra_compile_args = copy.deepcopy(extension.extra_compile_args)
- if isinstance(extension.extra_compile_args, dict):
- for args in extension.extra_compile_args.values():
- args.append(flag)
- else:
- extension.extra_compile_args.append(flag)
- def _define_torch_extension_name(self, extension):
- # pybind11 doesn't support dots in the names
- # so in order to support extensions in the packages
- # like torch._C, we take the last part of the string
- # as the library name
- names = extension.name.split('.')
- name = names[-1]
- define = f'-DTORCH_EXTENSION_NAME={name}'
- self._add_compile_flag(extension, define)
- def _add_gnu_cpp_abi_flag(self, extension):
- # use the same CXX ABI as what PyTorch was compiled with
- self._add_compile_flag(extension, '-D_GLIBCXX_USE_CXX11_ABI=' + str(int(torch._C._GLIBCXX_USE_CXX11_ABI)))
- def CppExtension(name, sources, *args, **kwargs):
- """
- Create a :class:`setuptools.Extension` for C++.
- Convenience method that creates a :class:`setuptools.Extension` with the
- bare minimum (but often sufficient) arguments to build a C++ extension.
- All arguments are forwarded to the :class:`setuptools.Extension`
- constructor. Full list arguments can be found at
- https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference
- Example:
- >>> # xdoctest: +SKIP
- >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT)
- >>> from setuptools import setup
- >>> from torch.utils.cpp_extension import BuildExtension, CppExtension
- >>> setup(
- ... name='extension',
- ... ext_modules=[
- ... CppExtension(
- ... name='extension',
- ... sources=['extension.cpp'],
- ... extra_compile_args=['-g'],
- ... extra_link_flags=['-Wl,--no-as-needed', '-lm'])
- ... ],
- ... cmdclass={
- ... 'build_ext': BuildExtension
- ... })
- """
- include_dirs = kwargs.get('include_dirs', [])
- include_dirs += include_paths()
- kwargs['include_dirs'] = include_dirs
- library_dirs = kwargs.get('library_dirs', [])
- library_dirs += library_paths()
- kwargs['library_dirs'] = library_dirs
- libraries = kwargs.get('libraries', [])
- libraries.append('c10')
- libraries.append('torch')
- libraries.append('torch_cpu')
- libraries.append('torch_python')
- if IS_WINDOWS:
- libraries.append("sleef")
- kwargs['libraries'] = libraries
- kwargs['language'] = 'c++'
- return setuptools.Extension(name, sources, *args, **kwargs)
- def CUDAExtension(name, sources, *args, **kwargs):
- """
- Create a :class:`setuptools.Extension` for CUDA/C++.
- Convenience method that creates a :class:`setuptools.Extension` with the
- bare minimum (but often sufficient) arguments to build a CUDA/C++
- extension. This includes the CUDA include path, library path and runtime
- library.
- All arguments are forwarded to the :class:`setuptools.Extension`
- constructor. Full list arguments can be found at
- https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference
- Example:
- >>> # xdoctest: +SKIP
- >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT)
- >>> from setuptools import setup
- >>> from torch.utils.cpp_extension import BuildExtension, CUDAExtension
- >>> setup(
- ... name='cuda_extension',
- ... ext_modules=[
- ... CUDAExtension(
- ... name='cuda_extension',
- ... sources=['extension.cpp', 'extension_kernel.cu'],
- ... extra_compile_args={'cxx': ['-g'],
- ... 'nvcc': ['-O2']},
- ... extra_link_flags=['-Wl,--no-as-needed', '-lcuda'])
- ... ],
- ... cmdclass={
- ... 'build_ext': BuildExtension
- ... })
- Compute capabilities:
- By default the extension will be compiled to run on all archs of the cards visible during the
- building process of the extension, plus PTX. If down the road a new card is installed the
- extension may need to be recompiled. If a visible card has a compute capability (CC) that's
- newer than the newest version for which your nvcc can build fully-compiled binaries, Pytorch
- will make nvcc fall back to building kernels with the newest version of PTX your nvcc does
- support (see below for details on PTX).
- You can override the default behavior using `TORCH_CUDA_ARCH_LIST` to explicitly specify which
- CCs you want the extension to support:
- ``TORCH_CUDA_ARCH_LIST="6.1 8.6" python build_my_extension.py``
- ``TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX" python build_my_extension.py``
- The +PTX option causes extension kernel binaries to include PTX instructions for the specified
- CC. PTX is an intermediate representation that allows kernels to runtime-compile for any CC >=
- the specified CC (for example, 8.6+PTX generates PTX that can runtime-compile for any GPU with
- CC >= 8.6). This improves your binary's forward compatibility. However, relying on older PTX to
- provide forward compat by runtime-compiling for newer CCs can modestly reduce performance on
- those newer CCs. If you know exact CC(s) of the GPUs you want to target, you're always better
- off specifying them individually. For example, if you want your extension to run on 8.0 and 8.6,
- "8.0+PTX" would work functionally because it includes PTX that can runtime-compile for 8.6, but
- "8.0 8.6" would be better.
- Note that while it's possible to include all supported archs, the more archs get included the
- slower the building process will be, as it will build a separate kernel image for each arch.
- Note that CUDA-11.5 nvcc will hit internal compiler error while parsing torch/extension.h on Windows.
- To workaround the issue, move python binding logic to pure C++ file.
- Example use:
- #include <ATen/ATen.h>
- at::Tensor SigmoidAlphaBlendForwardCuda(....)
- Instead of:
- #include <torch/extension.h>
- torch::Tensor SigmoidAlphaBlendForwardCuda(...)
- Currently open issue for nvcc bug: https://github.com/pytorch/pytorch/issues/69460
- Complete workaround code example: https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48
- Relocatable device code linking:
- If you want to reference device symbols across compilation units (across object files),
- the object files need to be built with `relocatable device code` (-rdc=true or -dc).
- An exception to this rule is "dynamic parallelism" (nested kernel launches) which is not used a lot anymore.
- `Relocatable device code` is less optimized so it needs to be used only on object files that need it.
- Using `-dlto` (Device Link Time Optimization) at the device code compilation step and `dlink` step
- help reduce the protentional perf degradation of `-rdc`.
- Note that it needs to be used at both steps to be useful.
- If you have `rdc` objects you need to have an extra `-dlink` (device linking) step before the CPU symbol linking step.
- There is also a case where `-dlink` is used without `-rdc`:
- when an extension is linked against a static lib containing rdc-compiled objects
- like the [NVSHMEM library](https://developer.nvidia.com/nvshmem).
- Note: Ninja is required to build a CUDA Extension with RDC linking.
- Example:
- >>> # xdoctest: +SKIP
- >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT)
- >>> CUDAExtension(
- ... name='cuda_extension',
- ... sources=['extension.cpp', 'extension_kernel.cu'],
- ... dlink=True,
- ... dlink_libraries=["dlink_lib"],
- ... extra_compile_args={'cxx': ['-g'],
- ... 'nvcc': ['-O2', '-rdc=true']})
- """
- library_dirs = kwargs.get('library_dirs', [])
- library_dirs += library_paths(cuda=True)
- kwargs['library_dirs'] = library_dirs
- libraries = kwargs.get('libraries', [])
- libraries.append('c10')
- libraries.append('torch')
- libraries.append('torch_cpu')
- libraries.append('torch_python')
- if IS_HIP_EXTENSION:
- libraries.append('amdhip64')
- libraries.append('c10_hip')
- libraries.append('torch_hip')
- else:
- libraries.append('cudart')
- libraries.append('c10_cuda')
- libraries.append('torch_cuda')
- kwargs['libraries'] = libraries
- include_dirs = kwargs.get('include_dirs', [])
- if IS_HIP_EXTENSION:
- build_dir = os.getcwd()
- hipify_result = hipify_python.hipify(
- project_directory=build_dir,
- output_directory=build_dir,
- header_include_dirs=include_dirs,
- includes=[os.path.join(build_dir, '*')], # limit scope to build_dir only
- extra_files=[os.path.abspath(s) for s in sources],
- show_detailed=True,
- is_pytorch_extension=True,
- hipify_extra_files_only=True, # don't hipify everything in includes path
- )
- hipified_sources = set()
- for source in sources:
- s_abs = os.path.abspath(source)
- hipified_s_abs = (hipify_result[s_abs].hipified_path if (s_abs in hipify_result and
- hipify_result[s_abs].hipified_path is not None) else s_abs)
- # setup() arguments must *always* be /-separated paths relative to the setup.py directory,
- # *never* absolute paths
- hipified_sources.add(os.path.relpath(hipified_s_abs, build_dir))
- sources = list(hipified_sources)
- include_dirs += include_paths(cuda=True)
- kwargs['include_dirs'] = include_dirs
- kwargs['language'] = 'c++'
- dlink_libraries = kwargs.get('dlink_libraries', [])
- dlink = kwargs.get('dlink', False) or dlink_libraries
- if dlink:
- extra_compile_args = kwargs.get('extra_compile_args', {})
- extra_compile_args_dlink = extra_compile_args.get('nvcc_dlink', [])
- extra_compile_args_dlink += ['-dlink']
- extra_compile_args_dlink += [f'-L{x}' for x in library_dirs]
- extra_compile_args_dlink += [f'-l{x}' for x in dlink_libraries]
- if (torch.version.cuda is not None) and TorchVersion(torch.version.cuda) >= '11.2':
- extra_compile_args_dlink += ['-dlto'] # Device Link Time Optimization started from cuda 11.2
- extra_compile_args['nvcc_dlink'] = extra_compile_args_dlink
- kwargs['extra_compile_args'] = extra_compile_args
- return setuptools.Extension(name, sources, *args, **kwargs)
- def include_paths(cuda: bool = False) -> List[str]:
- """
- Get the include paths required to build a C++ or CUDA extension.
- Args:
- cuda: If `True`, includes CUDA-specific include paths.
- Returns:
- A list of include path strings.
- """
- lib_include = os.path.join(_TORCH_PATH, 'include')
- paths = [
- lib_include,
- # Remove this once torch/torch.h is officially no longer supported for C++ extensions.
- os.path.join(lib_include, 'torch', 'csrc', 'api', 'include'),
- # Some internal (old) Torch headers don't properly prefix their includes,
- # so we need to pass -Itorch/lib/include/TH as well.
- os.path.join(lib_include, 'TH'),
- os.path.join(lib_include, 'THC')
- ]
- if cuda and IS_HIP_EXTENSION:
- paths.append(os.path.join(lib_include, 'THH'))
- paths.append(_join_rocm_home('include'))
- elif cuda:
- cuda_home_include = _join_cuda_home('include')
- # if we have the Debian/Ubuntu packages for cuda, we get /usr as cuda home.
- # but gcc doesn't like having /usr/include passed explicitly
- if cuda_home_include != '/usr/include':
- paths.append(cuda_home_include)
- # Support CUDA_INC_PATH env variable supported by CMake files
- if (cuda_inc_path := os.environ.get("CUDA_INC_PATH", None)) and \
- cuda_inc_path != '/usr/include':
- paths.append(cuda_inc_path)
- if CUDNN_HOME is not None:
- paths.append(os.path.join(CUDNN_HOME, 'include'))
- return paths
- def library_paths(cuda: bool = False) -> List[str]:
- """
- Get the library paths required to build a C++ or CUDA extension.
- Args:
- cuda: If `True`, includes CUDA-specific library paths.
- Returns:
- A list of library path strings.
- """
- # We need to link against libtorch.so
- paths = [TORCH_LIB_PATH]
- if cuda and IS_HIP_EXTENSION:
- lib_dir = 'lib'
- paths.append(_join_rocm_home(lib_dir))
- if HIP_HOME is not None:
- paths.append(os.path.join(HIP_HOME, 'lib'))
- elif cuda:
- if IS_WINDOWS:
- lib_dir = os.path.join('lib', 'x64')
- else:
- lib_dir = 'lib64'
- if (not os.path.exists(_join_cuda_home(lib_dir)) and
- os.path.exists(_join_cuda_home('lib'))):
- # 64-bit CUDA may be installed in 'lib' (see e.g. gh-16955)
- # Note that it's also possible both don't exist (see
- # _find_cuda_home) - in that case we stay with 'lib64'.
- lib_dir = 'lib'
- paths.append(_join_cuda_home(lib_dir))
- if CUDNN_HOME is not None:
- paths.append(os.path.join(CUDNN_HOME, lib_dir))
- return paths
- def load(name,
- sources: Union[str, List[str]],
- extra_cflags=None,
- extra_cuda_cflags=None,
- extra_ldflags=None,
- extra_include_paths=None,
- build_directory=None,
- verbose=False,
- with_cuda: Optional[bool] = None,
- is_python_module=True,
- is_standalone=False,
- keep_intermediates=True):
- """
- Load a PyTorch C++ extension just-in-time (JIT).
- To load an extension, a Ninja build file is emitted, which is used to
- compile the given sources into a dynamic library. This library is
- subsequently loaded into the current Python process as a module and
- returned from this function, ready for use.
- By default, the directory to which the build file is emitted and the
- resulting library compiled to is ``<tmp>/torch_extensions/<name>``, where
- ``<tmp>`` is the temporary folder on the current platform and ``<name>``
- the name of the extension. This location can be overridden in two ways.
- First, if the ``TORCH_EXTENSIONS_DIR`` environment variable is set, it
- replaces ``<tmp>/torch_extensions`` and all extensions will be compiled
- into subfolders of this directory. Second, if the ``build_directory``
- argument to this function is supplied, it overrides the entire path, i.e.
- the library will be compiled into that folder directly.
- To compile the sources, the default system compiler (``c++``) is used,
- which can be overridden by setting the ``CXX`` environment variable. To pass
- additional arguments to the compilation process, ``extra_cflags`` or
- ``extra_ldflags`` can be provided. For example, to compile your extension
- with optimizations, pass ``extra_cflags=['-O3']``. You can also use
- ``extra_cflags`` to pass further include directories.
- CUDA support with mixed compilation is provided. Simply pass CUDA source
- files (``.cu`` or ``.cuh``) along with other sources. Such files will be
- detected and compiled with nvcc rather than the C++ compiler. This includes
- passing the CUDA lib64 directory as a library directory, and linking
- ``cudart``. You can pass additional flags to nvcc via
- ``extra_cuda_cflags``, just like with ``extra_cflags`` for C++. Various
- heuristics for finding the CUDA install directory are used, which usually
- work fine. If not, setting the ``CUDA_HOME`` environment variable is the
- safest option.
- Args:
- name: The name of the extension to build. This MUST be the same as the
- name of the pybind11 module!
- sources: A list of relative or absolute paths to C++ source files.
- extra_cflags: optional list of compiler flags to forward to the build.
- extra_cuda_cflags: optional list of compiler flags to forward to nvcc
- when building CUDA sources.
- extra_ldflags: optional list of linker flags to forward to the build.
- extra_include_paths: optional list of include directories to forward
- to the build.
- build_directory: optional path to use as build workspace.
- verbose: If ``True``, turns on verbose logging of load steps.
- with_cuda: Determines whether CUDA headers and libraries are added to
- the build. If set to ``None`` (default), this value is
- automatically determined based on the existence of ``.cu`` or
- ``.cuh`` in ``sources``. Set it to `True`` to force CUDA headers
- and libraries to be included.
- is_python_module: If ``True`` (default), imports the produced shared
- library as a Python module. If ``False``, behavior depends on
- ``is_standalone``.
- is_standalone: If ``False`` (default) loads the constructed extension
- into the process as a plain dynamic library. If ``True``, build a
- standalone executable.
- Returns:
- If ``is_python_module`` is ``True``:
- Returns the loaded PyTorch extension as a Python module.
- If ``is_python_module`` is ``False`` and ``is_standalone`` is ``False``:
- Returns nothing. (The shared library is loaded into the process as
- a side effect.)
- If ``is_standalone`` is ``True``.
- Return the path to the executable. (On Windows, TORCH_LIB_PATH is
- added to the PATH environment variable as a side effect.)
- Example:
- >>> # xdoctest: +SKIP
- >>> from torch.utils.cpp_extension import load
- >>> module = load(
- ... name='extension',
- ... sources=['extension.cpp', 'extension_kernel.cu'],
- ... extra_cflags=['-O2'],
- ... verbose=True)
- """
- return _jit_compile(
- name,
- [sources] if isinstance(sources, str) else sources,
- extra_cflags,
- extra_cuda_cflags,
- extra_ldflags,
- extra_include_paths,
- build_directory or _get_build_directory(name, verbose),
- verbose,
- with_cuda,
- is_python_module,
- is_standalone,
- keep_intermediates=keep_intermediates)
- def _get_pybind11_abi_build_flags():
- # Note [Pybind11 ABI constants]
- #
- # Pybind11 before 2.4 used to build an ABI strings using the following pattern:
- # f"__pybind11_internals_v{PYBIND11_INTERNALS_VERSION}{PYBIND11_INTERNALS_KIND}{PYBIND11_BUILD_TYPE}__"
- # Since 2.4 compier type, stdlib and build abi parameters are also encoded like this:
- # f"__pybind11_internals_v{PYBIND11_INTERNALS_VERSION}{PYBIND11_INTERNALS_KIND}{PYBIND11_COMPILER_TYPE}{PYBIND11_STDLIB}{PYBIND11_BUILD_ABI}{PYBIND11_BUILD_TYPE}__"
- #
- # This was done in order to further narrow down the chances of compiler ABI incompatibility
- # that can cause a hard to debug segfaults.
- # For PyTorch extensions we want to relax those restrictions and pass compiler, stdlib and abi properties
- # captured during PyTorch native library compilation in torch/csrc/Module.cpp
- abi_cflags = []
- for pname in ["COMPILER_TYPE", "STDLIB", "BUILD_ABI"]:
- pval = getattr(torch._C, f"_PYBIND11_{pname}")
- if pval is not None and not IS_WINDOWS:
- abi_cflags.append(f'-DPYBIND11_{pname}=\\"{pval}\\"')
- return abi_cflags
- def _get_glibcxx_abi_build_flags():
- glibcxx_abi_cflags = ['-D_GLIBCXX_USE_CXX11_ABI=' + str(int(torch._C._GLIBCXX_USE_CXX11_ABI))]
- return glibcxx_abi_cflags
- def check_compiler_is_gcc(compiler):
- if not IS_LINUX:
- return False
- env = os.environ.copy()
- env['LC_ALL'] = 'C' # Don't localize output
- try:
- version_string = subprocess.check_output([compiler, '-v'], stderr=subprocess.STDOUT, env=env).decode(*SUBPROCESS_DECODE_ARGS)
- except Exception as e:
- try:
- version_string = subprocess.check_output([compiler, '--version'], stderr=subprocess.STDOUT, env=env).decode(*SUBPROCESS_DECODE_ARGS)
- except Exception as e:
- return False
- # Check for 'gcc' or 'g++' for sccache wrapper
- pattern = re.compile("^COLLECT_GCC=(.*)$", re.MULTILINE)
- results = re.findall(pattern, version_string)
- if len(results) != 1:
- return False
- compiler_path = os.path.realpath(results[0].strip())
- # On RHEL/CentOS c++ is a gcc compiler wrapper
- if os.path.basename(compiler_path) == 'c++' and 'gcc version' in version_string:
- return True
- return False
- def _check_and_build_extension_h_precompiler_headers(
- extra_cflags,
- extra_include_paths,
- is_standalone=False):
- r'''
- Precompiled Headers(PCH) can pre-build the same headers and reduce build time for pytorch load_inline modules.
- GCC offical manual: https://gcc.gnu.org/onlinedocs/gcc-4.0.4/gcc/Precompiled-Headers.html
- PCH only works when built pch file(header.h.gch) and build target have the same build parameters. So, We need
- add a signature file to record PCH file parameters. If the build parameters(signature) changed, it should rebuild
- PCH file.
- Note:
- 1. Windows and MacOS have different PCH mechanism. We only support Linux currently.
- 2. It only works on GCC/G++.
- '''
- if not IS_LINUX:
- return
- compiler = get_cxx_compiler()
- b_is_gcc = check_compiler_is_gcc(compiler)
- if b_is_gcc is False:
- return
- head_file = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h')
- head_file_pch = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.gch')
- head_file_signature = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.sign')
- def listToString(s):
- # initialize an empty string
- string = ""
- if s is None:
- return string
- # traverse in the string
- for element in s:
- string += (element + ' ')
- # return string
- return string
- def format_precompiler_header_cmd(compiler, head_file, head_file_pch, common_cflags, torch_include_dirs, extra_cflags, extra_include_paths):
- return re.sub(
- r"[ \n]+",
- " ",
- f"""
- {compiler} -x c++-header {head_file} -o {head_file_pch} {torch_include_dirs} {extra_include_paths} {extra_cflags} {common_cflags}
- """,
- ).strip()
- def command_to_signature(cmd):
- signature = cmd.replace(' ', '_')
- return signature
- def check_pch_signature_in_file(file_path, signature):
- b_exist = os.path.isfile(file_path)
- if b_exist is False:
- return False
- with open(file_path) as file:
- # read all content of a file
- content = file.read()
- # check if string present in a file
- if signature == content:
- return True
- else:
- return False
- def _create_if_not_exist(path_dir):
- if not os.path.exists(path_dir):
- try:
- Path(path_dir).mkdir(parents=True, exist_ok=True)
- except OSError as exc: # Guard against race condition
- if exc.errno != errno.EEXIST:
- raise RuntimeError(f"Fail to create path {path_dir}") from exc
- def write_pch_signature_to_file(file_path, pch_sign):
- _create_if_not_exist(os.path.dirname(file_path))
- with open(file_path, "w") as f:
- f.write(pch_sign)
- f.close()
- def build_precompile_header(pch_cmd):
- try:
- subprocess.check_output(pch_cmd, shell=True, stderr=subprocess.STDOUT)
- except subprocess.CalledProcessError as e:
- raise RuntimeError(f"Compile PreCompile Header fail, command: {pch_cmd}") from e
- extra_cflags_str = listToString(extra_cflags)
- extra_include_paths_str = " ".join(
- [f"-I{include}" for include in extra_include_paths] if extra_include_paths else []
- )
- lib_include = os.path.join(_TORCH_PATH, 'include')
- torch_include_dirs = [
- f"-I {lib_include}",
- # Python.h
- "-I {}".format(sysconfig.get_path("include")),
- # torch/all.h
- "-I {}".format(os.path.join(lib_include, 'torch', 'csrc', 'api', 'include')),
- ]
- torch_include_dirs_str = listToString(torch_include_dirs)
- common_cflags = []
- if not is_standalone:
- common_cflags += ['-DTORCH_API_INCLUDE_EXTENSION_H']
- common_cflags += ['-std=c++17', '-fPIC']
- common_cflags += [f"{x}" for x in _get_pybind11_abi_build_flags()]
- common_cflags += [f"{x}" for x in _get_glibcxx_abi_build_flags()]
- common_cflags_str = listToString(common_cflags)
- pch_cmd = format_precompiler_header_cmd(compiler, head_file, head_file_pch, common_cflags_str, torch_include_dirs_str, extra_cflags_str, extra_include_paths_str)
- pch_sign = command_to_signature(pch_cmd)
- if os.path.isfile(head_file_pch) is not True:
- build_precompile_header(pch_cmd)
- write_pch_signature_to_file(head_file_signature, pch_sign)
- else:
- b_same_sign = check_pch_signature_in_file(head_file_signature, pch_sign)
- if b_same_sign is False:
- build_precompile_header(pch_cmd)
- write_pch_signature_to_file(head_file_signature, pch_sign)
- def remove_extension_h_precompiler_headers():
- def _remove_if_file_exists(path_file):
- if os.path.exists(path_file):
- os.remove(path_file)
- head_file_pch = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.gch')
- head_file_signature = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.sign')
- _remove_if_file_exists(head_file_pch)
- _remove_if_file_exists(head_file_signature)
- def load_inline(name,
- cpp_sources,
- cuda_sources=None,
- functions=None,
- extra_cflags=None,
- extra_cuda_cflags=None,
- extra_ldflags=None,
- extra_include_paths=None,
- build_directory=None,
- verbose=False,
- with_cuda=None,
- is_python_module=True,
- with_pytorch_error_handling=True,
- keep_intermediates=True,
- use_pch=False):
- r'''
- Load a PyTorch C++ extension just-in-time (JIT) from string sources.
- This function behaves exactly like :func:`load`, but takes its sources as
- strings rather than filenames. These strings are stored to files in the
- build directory, after which the behavior of :func:`load_inline` is
- identical to :func:`load`.
- See `the
- tests <https://github.com/pytorch/pytorch/blob/master/test/test_cpp_extensions_jit.py>`_
- for good examples of using this function.
- Sources may omit two required parts of a typical non-inline C++ extension:
- the necessary header includes, as well as the (pybind11) binding code. More
- precisely, strings passed to ``cpp_sources`` are first concatenated into a
- single ``.cpp`` file. This file is then prepended with ``#include
- <torch/extension.h>``.
- Furthermore, if the ``functions`` argument is supplied, bindings will be
- automatically generated for each function specified. ``functions`` can
- either be a list of function names, or a dictionary mapping from function
- names to docstrings. If a list is given, the name of each function is used
- as its docstring.
- The sources in ``cuda_sources`` are concatenated into a separate ``.cu``
- file and prepended with ``torch/types.h``, ``cuda.h`` and
- ``cuda_runtime.h`` includes. The ``.cpp`` and ``.cu`` files are compiled
- separately, but ultimately linked into a single library. Note that no
- bindings are generated for functions in ``cuda_sources`` per se. To bind
- to a CUDA kernel, you must create a C++ function that calls it, and either
- declare or define this C++ function in one of the ``cpp_sources`` (and
- include its name in ``functions``).
- See :func:`load` for a description of arguments omitted below.
- Args:
- cpp_sources: A string, or list of strings, containing C++ source code.
- cuda_sources: A string, or list of strings, containing CUDA source code.
- functions: A list of function names for which to generate function
- bindings. If a dictionary is given, it should map function names to
- docstrings (which are otherwise just the function names).
- with_cuda: Determines whether CUDA headers and libraries are added to
- the build. If set to ``None`` (default), this value is
- automatically determined based on whether ``cuda_sources`` is
- provided. Set it to ``True`` to force CUDA headers
- and libraries to be included.
- with_pytorch_error_handling: Determines whether pytorch error and
- warning macros are handled by pytorch instead of pybind. To do
- this, each function ``foo`` is called via an intermediary ``_safe_foo``
- function. This redirection might cause issues in obscure cases
- of cpp. This flag should be set to ``False`` when this redirect
- causes issues.
- Example:
- >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT)
- >>> from torch.utils.cpp_extension import load_inline
- >>> source = """
- at::Tensor sin_add(at::Tensor x, at::Tensor y) {
- return x.sin() + y.sin();
- }
- """
- >>> module = load_inline(name='inline_extension',
- ... cpp_sources=[source],
- ... functions=['sin_add'])
- .. note::
- By default, the Ninja backend uses #CPUS + 2 workers to build the
- extension. This may use up too many resources on some systems. One
- can control the number of workers by setting the `MAX_JOBS` environment
- variable to a non-negative number.
- '''
- build_directory = build_directory or _get_build_directory(name, verbose)
- if isinstance(cpp_sources, str):
- cpp_sources = [cpp_sources]
- cuda_sources = cuda_sources or []
- if isinstance(cuda_sources, str):
- cuda_sources = [cuda_sources]
- cpp_sources.insert(0, '#include <torch/extension.h>')
- if use_pch is True:
- # Using PreCompile Header('torch/extension.h') to reduce compile time.
- _check_and_build_extension_h_precompiler_headers(extra_cflags, extra_include_paths)
- else:
- remove_extension_h_precompiler_headers()
- # If `functions` is supplied, we create the pybind11 bindings for the user.
- # Here, `functions` is (or becomes, after some processing) a map from
- # function names to function docstrings.
- if functions is not None:
- module_def = []
- module_def.append('PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {')
- if isinstance(functions, str):
- functions = [functions]
- if isinstance(functions, list):
- # Make the function docstring the same as the function name.
- functions = {f: f for f in functions}
- elif not isinstance(functions, dict):
- raise ValueError(f"Expected 'functions' to be a list or dict, but was {type(functions)}")
- for function_name, docstring in functions.items():
- if with_pytorch_error_handling:
- module_def.append(f'm.def("{function_name}", torch::wrap_pybind_function({function_name}), "{docstring}");')
- else:
- module_def.append(f'm.def("{function_name}", {function_name}, "{docstring}");')
- module_def.append('}')
- cpp_sources += module_def
- cpp_source_path = os.path.join(build_directory, 'main.cpp')
- _maybe_write(cpp_source_path, "\n".join(cpp_sources))
- sources = [cpp_source_path]
- if cuda_sources:
- cuda_sources.insert(0, '#include <torch/types.h>')
- cuda_sources.insert(1, '#include <cuda.h>')
- cuda_sources.insert(2, '#include <cuda_runtime.h>')
- cuda_source_path = os.path.join(build_directory, 'cuda.cu')
- _maybe_write(cuda_source_path, "\n".join(cuda_sources))
- sources.append(cuda_source_path)
- return _jit_compile(
- name,
- sources,
- extra_cflags,
- extra_cuda_cflags,
- extra_ldflags,
- extra_include_paths,
- build_directory,
- verbose,
- with_cuda,
- is_python_module,
- is_standalone=False,
- keep_intermediates=keep_intermediates)
- def _jit_compile(name,
- sources,
- extra_cflags,
- extra_cuda_cflags,
- extra_ldflags,
- extra_include_paths,
- build_directory: str,
- verbose: bool,
- with_cuda: Optional[bool],
- is_python_module,
- is_standalone,
- keep_intermediates=True) -> None:
- if is_python_module and is_standalone:
- raise ValueError("`is_python_module` and `is_standalone` are mutually exclusive.")
- if with_cuda is None:
- with_cuda = any(map(_is_cuda_file, sources))
- with_cudnn = any('cudnn' in f for f in extra_ldflags or [])
- old_version = JIT_EXTENSION_VERSIONER.get_version(name)
- version = JIT_EXTENSION_VERSIONER.bump_version_if_changed(
- name,
- sources,
- build_arguments=[extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths],
- build_directory=build_directory,
- with_cuda=with_cuda,
- is_python_module=is_python_module,
- is_standalone=is_standalone,
- )
- if version > 0:
- if version != old_version and verbose:
- print(f'The input conditions for extension module {name} have changed. ' +
- f'Bumping to version {version} and re-building as {name}_v{version}...',
- file=sys.stderr)
- name = f'{name}_v{version}'
- baton = FileBaton(os.path.join(build_directory, 'lock'))
- if baton.try_acquire():
- try:
- if version != old_version:
- with GeneratedFileCleaner(keep_intermediates=keep_intermediates) as clean_ctx:
- if IS_HIP_EXTENSION and (with_cuda or with_cudnn):
- hipify_result = hipify_python.hipify(
- project_directory=build_directory,
- output_directory=build_directory,
- header_include_dirs=(extra_include_paths if extra_include_paths is not None else []),
- extra_files=[os.path.abspath(s) for s in sources],
- ignores=[_join_rocm_home('*'), os.path.join(_TORCH_PATH, '*')], # no need to hipify ROCm or PyTorch headers
- show_detailed=verbose,
- show_progress=verbose,
- is_pytorch_extension=True,
- clean_ctx=clean_ctx
- )
- hipified_sources = set()
- for source in sources:
- s_abs = os.path.abspath(source)
- hipified_sources.add(hipify_result[s_abs].hipified_path if s_abs in hipify_result else s_abs)
- sources = list(hipified_sources)
- _write_ninja_file_and_build_library(
- name=name,
- sources=sources,
- extra_cflags=extra_cflags or [],
- extra_cuda_cflags=extra_cuda_cflags or [],
- extra_ldflags=extra_ldflags or [],
- extra_include_paths=extra_include_paths or [],
- build_directory=build_directory,
- verbose=verbose,
- with_cuda=with_cuda,
- is_standalone=is_standalone)
- elif verbose:
- print('No modifications detected for re-loaded extension '
- f'module {name}, skipping build step...', file=sys.stderr)
- finally:
- baton.release()
- else:
- baton.wait()
- if verbose:
- print(f'Loading extension module {name}...', file=sys.stderr)
- if is_standalone:
- return _get_exec_path(name, build_directory)
- return _import_module_from_library(name, build_directory, is_python_module)
- def _write_ninja_file_and_compile_objects(
- sources: List[str],
- objects,
- cflags,
- post_cflags,
- cuda_cflags,
- cuda_post_cflags,
- cuda_dlink_post_cflags,
- build_directory: str,
- verbose: bool,
- with_cuda: Optional[bool]) -> None:
- verify_ninja_availability()
- compiler = get_cxx_compiler()
- get_compiler_abi_compatibility_and_version(compiler)
- if with_cuda is None:
- with_cuda = any(map(_is_cuda_file, sources))
- build_file_path = os.path.join(build_directory, 'build.ninja')
- if verbose:
- print(f'Emitting ninja build file {build_file_path}...', file=sys.stderr)
- _write_ninja_file(
- path=build_file_path,
- cflags=cflags,
- post_cflags=post_cflags,
- cuda_cflags=cuda_cflags,
- cuda_post_cflags=cuda_post_cflags,
- cuda_dlink_post_cflags=cuda_dlink_post_cflags,
- sources=sources,
- objects=objects,
- ldflags=None,
- library_target=None,
- with_cuda=with_cuda)
- if verbose:
- print('Compiling objects...', file=sys.stderr)
- _run_ninja_build(
- build_directory,
- verbose,
- # It would be better if we could tell users the name of the extension
- # that failed to build but there isn't a good way to get it here.
- error_prefix='Error compiling objects for extension')
- def _write_ninja_file_and_build_library(
- name,
- sources: List[str],
- extra_cflags,
- extra_cuda_cflags,
- extra_ldflags,
- extra_include_paths,
- build_directory: str,
- verbose: bool,
- with_cuda: Optional[bool],
- is_standalone: bool = False) -> None:
- verify_ninja_availability()
- compiler = get_cxx_compiler()
- get_compiler_abi_compatibility_and_version(compiler)
- if with_cuda is None:
- with_cuda = any(map(_is_cuda_file, sources))
- extra_ldflags = _prepare_ldflags(
- extra_ldflags or [],
- with_cuda,
- verbose,
- is_standalone)
- build_file_path = os.path.join(build_directory, 'build.ninja')
- if verbose:
- print(f'Emitting ninja build file {build_file_path}...', file=sys.stderr)
- # NOTE: Emitting a new ninja build file does not cause re-compilation if
- # the sources did not change, so it's ok to re-emit (and it's fast).
- _write_ninja_file_to_build_library(
- path=build_file_path,
- name=name,
- sources=sources,
- extra_cflags=extra_cflags or [],
- extra_cuda_cflags=extra_cuda_cflags or [],
- extra_ldflags=extra_ldflags or [],
- extra_include_paths=extra_include_paths or [],
- with_cuda=with_cuda,
- is_standalone=is_standalone)
- if verbose:
- print(f'Building extension module {name}...', file=sys.stderr)
- _run_ninja_build(
- build_directory,
- verbose,
- error_prefix=f"Error building extension '{name}'")
- def is_ninja_available():
- """Return ``True`` if the `ninja <https://ninja-build.org/>`_ build system is available on the system, ``False`` otherwise."""
- try:
- subprocess.check_output('ninja --version'.split())
- except Exception:
- return False
- else:
- return True
- def verify_ninja_availability():
- """Raise ``RuntimeError`` if `ninja <https://ninja-build.org/>`_ build system is not available on the system, does nothing otherwise."""
- if not is_ninja_available():
- raise RuntimeError("Ninja is required to load C++ extensions")
- def _prepare_ldflags(extra_ldflags, with_cuda, verbose, is_standalone):
- if IS_WINDOWS:
- python_lib_path = os.path.join(sys.base_exec_prefix, 'libs')
- extra_ldflags.append('c10.lib')
- if with_cuda:
- extra_ldflags.append('c10_cuda.lib')
- extra_ldflags.append('torch_cpu.lib')
- if with_cuda:
- extra_ldflags.append('torch_cuda.lib')
- # /INCLUDE is used to ensure torch_cuda is linked against in a project that relies on it.
- # Related issue: https://github.com/pytorch/pytorch/issues/31611
- extra_ldflags.append('-INCLUDE:?warp_size@cuda@at@@YAHXZ')
- extra_ldflags.append('torch.lib')
- extra_ldflags.append(f'/LIBPATH:{TORCH_LIB_PATH}')
- if not is_standalone:
- extra_ldflags.append('torch_python.lib')
- extra_ldflags.append(f'/LIBPATH:{python_lib_path}')
- else:
- extra_ldflags.append(f'-L{TORCH_LIB_PATH}')
- extra_ldflags.append('-lc10')
- if with_cuda:
- extra_ldflags.append('-lc10_hip' if IS_HIP_EXTENSION else '-lc10_cuda')
- extra_ldflags.append('-ltorch_cpu')
- if with_cuda:
- extra_ldflags.append('-ltorch_hip' if IS_HIP_EXTENSION else '-ltorch_cuda')
- extra_ldflags.append('-ltorch')
- if not is_standalone:
- extra_ldflags.append('-ltorch_python')
- if is_standalone:
- extra_ldflags.append(f"-Wl,-rpath,{TORCH_LIB_PATH}")
- if with_cuda:
- if verbose:
- print('Detected CUDA files, patching ldflags', file=sys.stderr)
- if IS_WINDOWS:
- extra_ldflags.append(f'/LIBPATH:{_join_cuda_home("lib", "x64")}')
- extra_ldflags.append('cudart.lib')
- if CUDNN_HOME is not None:
- extra_ldflags.append(f'/LIBPATH:{os.path.join(CUDNN_HOME, "lib", "x64")}')
- elif not IS_HIP_EXTENSION:
- extra_lib_dir = "lib64"
- if (not os.path.exists(_join_cuda_home(extra_lib_dir)) and
- os.path.exists(_join_cuda_home("lib"))):
- # 64-bit CUDA may be installed in "lib"
- # Note that it's also possible both don't exist (see _find_cuda_home) - in that case we stay with "lib64"
- extra_lib_dir = "lib"
- extra_ldflags.append(f'-L{_join_cuda_home(extra_lib_dir)}')
- extra_ldflags.append('-lcudart')
- if CUDNN_HOME is not None:
- extra_ldflags.append(f'-L{os.path.join(CUDNN_HOME, "lib64")}')
- elif IS_HIP_EXTENSION:
- extra_ldflags.append(f'-L{_join_rocm_home("lib")}')
- extra_ldflags.append('-lamdhip64')
- return extra_ldflags
- def _get_cuda_arch_flags(cflags: Optional[List[str]] = None) -> List[str]:
- """
- Determine CUDA arch flags to use.
- For an arch, say "6.1", the added compile flag will be
- ``-gencode=arch=compute_61,code=sm_61``.
- For an added "+PTX", an additional
- ``-gencode=arch=compute_xx,code=compute_xx`` is added.
- See select_compute_arch.cmake for corresponding named and supported arches
- when building with CMake.
- """
- # If cflags is given, there may already be user-provided arch flags in it
- # (from `extra_compile_args`)
- if cflags is not None:
- for flag in cflags:
- if 'TORCH_EXTENSION_NAME' in flag:
- continue
- if 'arch' in flag:
- return []
- # Note: keep combined names ("arch1+arch2") above single names, otherwise
- # string replacement may not do the right thing
- named_arches = collections.OrderedDict([
- ('Kepler+Tesla', '3.7'),
- ('Kepler', '3.5+PTX'),
- ('Maxwell+Tegra', '5.3'),
- ('Maxwell', '5.0;5.2+PTX'),
- ('Pascal', '6.0;6.1+PTX'),
- ('Volta+Tegra', '7.2'),
- ('Volta', '7.0+PTX'),
- ('Turing', '7.5+PTX'),
- ('Ampere+Tegra', '8.7'),
- ('Ampere', '8.0;8.6+PTX'),
- ('Ada', '8.9+PTX'),
- ('Hopper', '9.0+PTX'),
- ])
- supported_arches = ['3.5', '3.7', '5.0', '5.2', '5.3', '6.0', '6.1', '6.2',
- '7.0', '7.2', '7.5', '8.0', '8.6', '8.7', '8.9', '9.0', '9.0a']
- valid_arch_strings = supported_arches + [s + "+PTX" for s in supported_arches]
- # The default is sm_30 for CUDA 9.x and 10.x
- # First check for an env var (same as used by the main setup.py)
- # Can be one or more architectures, e.g. "6.1" or "3.5;5.2;6.0;6.1;7.0+PTX"
- # See cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake
- _arch_list = os.environ.get('TORCH_CUDA_ARCH_LIST', None)
- # If not given, determine what's best for the GPU / CUDA version that can be found
- if not _arch_list:
- warnings.warn(
- "TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. \n"
- "If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].")
- arch_list = []
- # the assumption is that the extension should run on any of the currently visible cards,
- # which could be of different types - therefore all archs for visible cards should be included
- for i in range(torch.cuda.device_count()):
- capability = torch.cuda.get_device_capability(i)
- supported_sm = [int(arch.split('_')[1])
- for arch in torch.cuda.get_arch_list() if 'sm_' in arch]
- max_supported_sm = max((sm // 10, sm % 10) for sm in supported_sm)
- # Capability of the device may be higher than what's supported by the user's
- # NVCC, causing compilation error. User's NVCC is expected to match the one
- # used to build pytorch, so we use the maximum supported capability of pytorch
- # to clamp the capability.
- capability = min(max_supported_sm, capability)
- arch = f'{capability[0]}.{capability[1]}'
- if arch not in arch_list:
- arch_list.append(arch)
- arch_list = sorted(arch_list)
- arch_list[-1] += '+PTX'
- else:
- # Deal with lists that are ' ' separated (only deal with ';' after)
- _arch_list = _arch_list.replace(' ', ';')
- # Expand named arches
- for named_arch, archval in named_arches.items():
- _arch_list = _arch_list.replace(named_arch, archval)
- arch_list = _arch_list.split(';')
- flags = []
- for arch in arch_list:
- if arch not in valid_arch_strings:
- raise ValueError(f"Unknown CUDA arch ({arch}) or GPU not supported")
- else:
- num = arch[0] + arch[2:].split("+")[0]
- flags.append(f'-gencode=arch=compute_{num},code=sm_{num}')
- if arch.endswith('+PTX'):
- flags.append(f'-gencode=arch=compute_{num},code=compute_{num}')
- return sorted(set(flags))
- def _get_rocm_arch_flags(cflags: Optional[List[str]] = None) -> List[str]:
- # If cflags is given, there may already be user-provided arch flags in it
- # (from `extra_compile_args`)
- if cflags is not None:
- for flag in cflags:
- if 'amdgpu-target' in flag or 'offload-arch' in flag:
- return ['-fno-gpu-rdc']
- # Use same defaults as used for building PyTorch
- # Allow env var to override, just like during initial cmake build.
- _archs = os.environ.get('PYTORCH_ROCM_ARCH', None)
- if not _archs:
- archFlags = torch._C._cuda_getArchFlags()
- if archFlags:
- archs = archFlags.split()
- else:
- archs = []
- else:
- archs = _archs.replace(' ', ';').split(';')
- flags = [f'--offload-arch={arch}' for arch in archs]
- flags += ['-fno-gpu-rdc']
- return flags
- def _get_build_directory(name: str, verbose: bool) -> str:
- root_extensions_directory = os.environ.get('TORCH_EXTENSIONS_DIR')
- if root_extensions_directory is None:
- root_extensions_directory = get_default_build_root()
- cu_str = ('cpu' if torch.version.cuda is None else
- f'cu{torch.version.cuda.replace(".", "")}') # type: ignore[attr-defined]
- python_version = f'py{sys.version_info.major}{sys.version_info.minor}'
- build_folder = f'{python_version}_{cu_str}'
- root_extensions_directory = os.path.join(
- root_extensions_directory, build_folder)
- if verbose:
- print(f'Using {root_extensions_directory} as PyTorch extensions root...', file=sys.stderr)
- build_directory = os.path.join(root_extensions_directory, name)
- if not os.path.exists(build_directory):
- if verbose:
- print(f'Creating extension directory {build_directory}...', file=sys.stderr)
- # This is like mkdir -p, i.e. will also create parent directories.
- os.makedirs(build_directory, exist_ok=True)
- return build_directory
- def _get_num_workers(verbose: bool) -> Optional[int]:
- max_jobs = os.environ.get('MAX_JOBS')
- if max_jobs is not None and max_jobs.isdigit():
- if verbose:
- print(f'Using envvar MAX_JOBS ({max_jobs}) as the number of workers...',
- file=sys.stderr)
- return int(max_jobs)
- if verbose:
- print('Allowing ninja to set a default number of workers... '
- '(overridable by setting the environment variable MAX_JOBS=N)',
- file=sys.stderr)
- return None
- def _run_ninja_build(build_directory: str, verbose: bool, error_prefix: str) -> None:
- command = ['ninja', '-v']
- num_workers = _get_num_workers(verbose)
- if num_workers is not None:
- command.extend(['-j', str(num_workers)])
- env = os.environ.copy()
- # Try to activate the vc env for the users
- if IS_WINDOWS and 'VSCMD_ARG_TGT_ARCH' not in env:
- from setuptools import distutils
- plat_name = distutils.util.get_platform()
- plat_spec = PLAT_TO_VCVARS[plat_name]
- vc_env = distutils._msvccompiler._get_vc_env(plat_spec)
- vc_env = {k.upper(): v for k, v in vc_env.items()}
- for k, v in env.items():
- uk = k.upper()
- if uk not in vc_env:
- vc_env[uk] = v
- env = vc_env
- try:
- sys.stdout.flush()
- sys.stderr.flush()
- # Warning: don't pass stdout=None to subprocess.run to get output.
- # subprocess.run assumes that sys.__stdout__ has not been modified and
- # attempts to write to it by default. However, when we call _run_ninja_build
- # from ahead-of-time cpp extensions, the following happens:
- # 1) If the stdout encoding is not utf-8, setuptools detachs __stdout__.
- # https://github.com/pypa/setuptools/blob/7e97def47723303fafabe48b22168bbc11bb4821/setuptools/dist.py#L1110
- # (it probably shouldn't do this)
- # 2) subprocess.run (on POSIX, with no stdout override) relies on
- # __stdout__ not being detached:
- # https://github.com/python/cpython/blob/c352e6c7446c894b13643f538db312092b351789/Lib/subprocess.py#L1214
- # To work around this, we pass in the fileno directly and hope that
- # it is valid.
- stdout_fileno = 1
- subprocess.run(
- command,
- stdout=stdout_fileno if verbose else subprocess.PIPE,
- stderr=subprocess.STDOUT,
- cwd=build_directory,
- check=True,
- env=env)
- except subprocess.CalledProcessError as e:
- # Python 2 and 3 compatible way of getting the error object.
- _, error, _ = sys.exc_info()
- # error.output contains the stdout and stderr of the build attempt.
- message = error_prefix
- # `error` is a CalledProcessError (which has an `output`) attribute, but
- # mypy thinks it's Optional[BaseException] and doesn't narrow
- if hasattr(error, 'output') and error.output: # type: ignore[union-attr]
- message += f": {error.output.decode(*SUBPROCESS_DECODE_ARGS)}" # type: ignore[union-attr]
- raise RuntimeError(message) from e
- def _get_exec_path(module_name, path):
- if IS_WINDOWS and TORCH_LIB_PATH not in os.getenv('PATH', '').split(';'):
- torch_lib_in_path = any(
- os.path.exists(p) and os.path.samefile(p, TORCH_LIB_PATH)
- for p in os.getenv('PATH', '').split(';')
- )
- if not torch_lib_in_path:
- os.environ['PATH'] = f"{TORCH_LIB_PATH};{os.getenv('PATH', '')}"
- return os.path.join(path, f'{module_name}{EXEC_EXT}')
- def _import_module_from_library(module_name, path, is_python_module):
- filepath = os.path.join(path, f"{module_name}{LIB_EXT}")
- if is_python_module:
- # https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path
- spec = importlib.util.spec_from_file_location(module_name, filepath)
- assert spec is not None
- module = importlib.util.module_from_spec(spec)
- assert isinstance(spec.loader, importlib.abc.Loader)
- spec.loader.exec_module(module)
- return module
- else:
- torch.ops.load_library(filepath)
- def _write_ninja_file_to_build_library(path,
- name,
- sources,
- extra_cflags,
- extra_cuda_cflags,
- extra_ldflags,
- extra_include_paths,
- with_cuda,
- is_standalone) -> None:
- extra_cflags = [flag.strip() for flag in extra_cflags]
- extra_cuda_cflags = [flag.strip() for flag in extra_cuda_cflags]
- extra_ldflags = [flag.strip() for flag in extra_ldflags]
- extra_include_paths = [flag.strip() for flag in extra_include_paths]
- # Turn into absolute paths so we can emit them into the ninja build
- # file wherever it is.
- user_includes = [os.path.abspath(file) for file in extra_include_paths]
- # include_paths() gives us the location of torch/extension.h
- system_includes = include_paths(with_cuda)
- # sysconfig.get_path('include') gives us the location of Python.h
- # Explicitly specify 'posix_prefix' scheme on non-Windows platforms to workaround error on some MacOS
- # installations where default `get_path` points to non-existing `/Library/Python/M.m/include` folder
- python_include_path = sysconfig.get_path('include', scheme='nt' if IS_WINDOWS else 'posix_prefix')
- if python_include_path is not None:
- system_includes.append(python_include_path)
- common_cflags = []
- if not is_standalone:
- common_cflags.append(f'-DTORCH_EXTENSION_NAME={name}')
- common_cflags.append('-DTORCH_API_INCLUDE_EXTENSION_H')
- common_cflags += [f"{x}" for x in _get_pybind11_abi_build_flags()]
- # Windows does not understand `-isystem` and quotes flags later.
- if IS_WINDOWS:
- common_cflags += [f'-I{include}' for include in user_includes + system_includes]
- else:
- common_cflags += [f'-I{shlex.quote(include)}' for include in user_includes]
- common_cflags += [f'-isystem {shlex.quote(include)}' for include in system_includes]
- common_cflags += [f"{x}" for x in _get_glibcxx_abi_build_flags()]
- if IS_WINDOWS:
- cflags = common_cflags + COMMON_MSVC_FLAGS + ['/std:c++17'] + extra_cflags
- cflags = _nt_quote_args(cflags)
- else:
- cflags = common_cflags + ['-fPIC', '-std=c++17'] + extra_cflags
- if with_cuda and IS_HIP_EXTENSION:
- cuda_flags = ['-DWITH_HIP'] + cflags + COMMON_HIP_FLAGS + COMMON_HIPCC_FLAGS
- cuda_flags += extra_cuda_cflags
- cuda_flags += _get_rocm_arch_flags(cuda_flags)
- elif with_cuda:
- cuda_flags = common_cflags + COMMON_NVCC_FLAGS + _get_cuda_arch_flags()
- if IS_WINDOWS:
- for flag in COMMON_MSVC_FLAGS:
- cuda_flags = ['-Xcompiler', flag] + cuda_flags
- for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
- cuda_flags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cuda_flags
- cuda_flags = cuda_flags + ['-std=c++17']
- cuda_flags = _nt_quote_args(cuda_flags)
- cuda_flags += _nt_quote_args(extra_cuda_cflags)
- else:
- cuda_flags += ['--compiler-options', "'-fPIC'"]
- cuda_flags += extra_cuda_cflags
- if not any(flag.startswith('-std=') for flag in cuda_flags):
- cuda_flags.append('-std=c++17')
- cc_env = os.getenv("CC")
- if cc_env is not None:
- cuda_flags = ['-ccbin', cc_env] + cuda_flags
- else:
- cuda_flags = None
- def object_file_path(source_file: str) -> str:
- # '/path/to/file.cpp' -> 'file'
- file_name = os.path.splitext(os.path.basename(source_file))[0]
- if _is_cuda_file(source_file) and with_cuda:
- # Use a different object filename in case a C++ and CUDA file have
- # the same filename but different extension (.cpp vs. .cu).
- target = f'{file_name}.cuda.o'
- else:
- target = f'{file_name}.o'
- return target
- objects = [object_file_path(src) for src in sources]
- ldflags = ([] if is_standalone else [SHARED_FLAG]) + extra_ldflags
- # The darwin linker needs explicit consent to ignore unresolved symbols.
- if IS_MACOS:
- ldflags.append('-undefined dynamic_lookup')
- elif IS_WINDOWS:
- ldflags = _nt_quote_args(ldflags)
- ext = EXEC_EXT if is_standalone else LIB_EXT
- library_target = f'{name}{ext}'
- _write_ninja_file(
- path=path,
- cflags=cflags,
- post_cflags=None,
- cuda_cflags=cuda_flags,
- cuda_post_cflags=None,
- cuda_dlink_post_cflags=None,
- sources=sources,
- objects=objects,
- ldflags=ldflags,
- library_target=library_target,
- with_cuda=with_cuda)
- def _write_ninja_file(path,
- cflags,
- post_cflags,
- cuda_cflags,
- cuda_post_cflags,
- cuda_dlink_post_cflags,
- sources,
- objects,
- ldflags,
- library_target,
- with_cuda) -> None:
- r"""Write a ninja file that does the desired compiling and linking.
- `path`: Where to write this file
- `cflags`: list of flags to pass to $cxx. Can be None.
- `post_cflags`: list of flags to append to the $cxx invocation. Can be None.
- `cuda_cflags`: list of flags to pass to $nvcc. Can be None.
- `cuda_postflags`: list of flags to append to the $nvcc invocation. Can be None.
- `sources`: list of paths to source files
- `objects`: list of desired paths to objects, one per source.
- `ldflags`: list of flags to pass to linker. Can be None.
- `library_target`: Name of the output library. Can be None; in that case,
- we do no linking.
- `with_cuda`: If we should be compiling with CUDA.
- """
- def sanitize_flags(flags):
- if flags is None:
- return []
- else:
- return [flag.strip() for flag in flags]
- cflags = sanitize_flags(cflags)
- post_cflags = sanitize_flags(post_cflags)
- cuda_cflags = sanitize_flags(cuda_cflags)
- cuda_post_cflags = sanitize_flags(cuda_post_cflags)
- cuda_dlink_post_cflags = sanitize_flags(cuda_dlink_post_cflags)
- ldflags = sanitize_flags(ldflags)
- # Sanity checks...
- assert len(sources) == len(objects)
- assert len(sources) > 0
- compiler = get_cxx_compiler()
- # Version 1.3 is required for the `deps` directive.
- config = ['ninja_required_version = 1.3']
- config.append(f'cxx = {compiler}')
- if with_cuda or cuda_dlink_post_cflags:
- if "PYTORCH_NVCC" in os.environ:
- nvcc = os.getenv("PYTORCH_NVCC") # user can set nvcc compiler with ccache using the environment variable here
- else:
- if IS_HIP_EXTENSION:
- nvcc = _join_rocm_home('bin', 'hipcc')
- else:
- nvcc = _join_cuda_home('bin', 'nvcc')
- config.append(f'nvcc = {nvcc}')
- if IS_HIP_EXTENSION:
- post_cflags = COMMON_HIP_FLAGS + post_cflags
- flags = [f'cflags = {" ".join(cflags)}']
- flags.append(f'post_cflags = {" ".join(post_cflags)}')
- if with_cuda:
- flags.append(f'cuda_cflags = {" ".join(cuda_cflags)}')
- flags.append(f'cuda_post_cflags = {" ".join(cuda_post_cflags)}')
- flags.append(f'cuda_dlink_post_cflags = {" ".join(cuda_dlink_post_cflags)}')
- flags.append(f'ldflags = {" ".join(ldflags)}')
- # Turn into absolute paths so we can emit them into the ninja build
- # file wherever it is.
- sources = [os.path.abspath(file) for file in sources]
- # See https://ninja-build.org/build.ninja.html for reference.
- compile_rule = ['rule compile']
- if IS_WINDOWS:
- compile_rule.append(
- ' command = cl /showIncludes $cflags -c $in /Fo$out $post_cflags')
- compile_rule.append(' deps = msvc')
- else:
- compile_rule.append(
- ' command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags')
- compile_rule.append(' depfile = $out.d')
- compile_rule.append(' deps = gcc')
- if with_cuda:
- cuda_compile_rule = ['rule cuda_compile']
- nvcc_gendeps = ''
- # --generate-dependencies-with-compile is not supported by ROCm
- # Nvcc flag `--generate-dependencies-with-compile` is not supported by sccache, which may increase build time.
- if torch.version.cuda is not None and os.getenv('TORCH_EXTENSION_SKIP_NVCC_GEN_DEPENDENCIES', '0') != '1':
- cuda_compile_rule.append(' depfile = $out.d')
- cuda_compile_rule.append(' deps = gcc')
- # Note: non-system deps with nvcc are only supported
- # on Linux so use --generate-dependencies-with-compile
- # to make this work on Windows too.
- nvcc_gendeps = '--generate-dependencies-with-compile --dependency-output $out.d'
- cuda_compile_rule.append(
- f' command = $nvcc {nvcc_gendeps} $cuda_cflags -c $in -o $out $cuda_post_cflags')
- # Emit one build rule per source to enable incremental build.
- build = []
- for source_file, object_file in zip(sources, objects):
- is_cuda_source = _is_cuda_file(source_file) and with_cuda
- rule = 'cuda_compile' if is_cuda_source else 'compile'
- if IS_WINDOWS:
- source_file = source_file.replace(':', '$:')
- object_file = object_file.replace(':', '$:')
- source_file = source_file.replace(" ", "$ ")
- object_file = object_file.replace(" ", "$ ")
- build.append(f'build {object_file}: {rule} {source_file}')
- if cuda_dlink_post_cflags:
- devlink_out = os.path.join(os.path.dirname(objects[0]), 'dlink.o')
- devlink_rule = ['rule cuda_devlink']
- devlink_rule.append(' command = $nvcc $in -o $out $cuda_dlink_post_cflags')
- devlink = [f'build {devlink_out}: cuda_devlink {" ".join(objects)}']
- objects += [devlink_out]
- else:
- devlink_rule, devlink = [], []
- if library_target is not None:
- link_rule = ['rule link']
- if IS_WINDOWS:
- cl_paths = subprocess.check_output(['where',
- 'cl']).decode(*SUBPROCESS_DECODE_ARGS).split('\r\n')
- if len(cl_paths) >= 1:
- cl_path = os.path.dirname(cl_paths[0]).replace(':', '$:')
- else:
- raise RuntimeError("MSVC is required to load C++ extensions")
- link_rule.append(f' command = "{cl_path}/link.exe" $in /nologo $ldflags /out:$out')
- else:
- link_rule.append(' command = $cxx $in $ldflags -o $out')
- link = [f'build {library_target}: link {" ".join(objects)}']
- default = [f'default {library_target}']
- else:
- link_rule, link, default = [], [], []
- # 'Blocks' should be separated by newlines, for visual benefit.
- blocks = [config, flags, compile_rule]
- if with_cuda:
- blocks.append(cuda_compile_rule) # type: ignore[possibly-undefined]
- blocks += [devlink_rule, link_rule, build, devlink, link, default]
- content = "\n\n".join("\n".join(b) for b in blocks)
- # Ninja requires a new lines at the end of the .ninja file
- content += "\n"
- _maybe_write(path, content)
- def _join_cuda_home(*paths) -> str:
- """
- Join paths with CUDA_HOME, or raises an error if it CUDA_HOME is not set.
- This is basically a lazy way of raising an error for missing $CUDA_HOME
- only once we need to get any CUDA-specific path.
- """
- if CUDA_HOME is None:
- raise OSError('CUDA_HOME environment variable is not set. '
- 'Please set it to your CUDA install root.')
- return os.path.join(CUDA_HOME, *paths)
- def _is_cuda_file(path: str) -> bool:
- valid_ext = ['.cu', '.cuh']
- if IS_HIP_EXTENSION:
- valid_ext.append('.hip')
- return os.path.splitext(path)[1] in valid_ext
|