# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual # property and proprietary rights in and to this material, related # documentation and any modifications thereto. Any use, reproduction, # disclosure or distribution of this material and related documentation # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. import glob import hashlib import importlib import os import re import shutil import uuid import torch import torch.utils.cpp_extension from torch.utils.file_baton import FileBaton #---------------------------------------------------------------------------- # Global options. verbosity = 'brief' # Verbosity level: 'none', 'brief', 'full' #---------------------------------------------------------------------------- # Internal helper funcs. def _find_compiler_bindir(): patterns = [ 'C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64', 'C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64', 'C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64', 'C:/Program Files (x86)/Microsoft Visual Studio */vc/bin', ] for pattern in patterns: matches = sorted(glob.glob(pattern)) if len(matches): return matches[-1] return None #---------------------------------------------------------------------------- def _get_mangled_gpu_name(): name = torch.cuda.get_device_name().lower() out = [] for c in name: if re.match('[a-z0-9_-]+', c): out.append(c) else: out.append('-') return ''.join(out) #---------------------------------------------------------------------------- # Main entry point for compiling and loading C++/CUDA plugins. _cached_plugins = dict() def get_plugin(module_name, sources, headers=None, source_dir=None, **build_kwargs): assert verbosity in ['none', 'brief', 'full'] if headers is None: headers = [] if source_dir is not None: sources = [os.path.join(source_dir, fname) for fname in sources] headers = [os.path.join(source_dir, fname) for fname in headers] # Already cached? if module_name in _cached_plugins: return _cached_plugins[module_name] # Print status. if verbosity == 'full': print(f'Setting up PyTorch plugin "{module_name}"...') elif verbosity == 'brief': print(f'Setting up PyTorch plugin "{module_name}"... ', end='', flush=True) verbose_build = (verbosity == 'full') # Compile and load. try: # pylint: disable=too-many-nested-blocks # Make sure we can find the necessary compiler binaries. if os.name == 'nt' and os.system("where cl.exe >nul 2>nul") != 0: compiler_bindir = _find_compiler_bindir() if compiler_bindir is None: raise RuntimeError( f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in "{__file__}".' ) os.environ['PATH'] += ';' + compiler_bindir # Some containers set TORCH_CUDA_ARCH_LIST to a list that can either # break the build or unnecessarily restrict what's available to nvcc. # Unset it to let nvcc decide based on what's available on the # machine. os.environ['TORCH_CUDA_ARCH_LIST'] = '' # Incremental build md5sum trickery. Copies all the input source files # into a cached build directory under a combined md5 digest of the input # source files. Copying is done only if the combined digest has changed. # This keeps input file timestamps and filenames the same as in previous # extension builds, allowing for fast incremental rebuilds. # # This optimization is done only in case all the source files reside in # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR # environment variable is set (we take this as a signal that the user # actually cares about this.) # # EDIT: We now do it regardless of TORCH_EXTENSIOS_DIR, in order to work # around the *.cu dependency bug in ninja config. # all_source_files = sorted(sources + headers) all_source_dirs = set( os.path.dirname(fname) for fname in all_source_files) if len(all_source_dirs ) == 1: # and ('TORCH_EXTENSIONS_DIR' in os.environ): # Compute combined hash digest for all source files. hash_md5 = hashlib.md5() for src in all_source_files: with open(src, 'rb') as f: hash_md5.update(f.read()) # Select cached build directory name. source_digest = hash_md5.hexdigest() build_top_dir = torch.utils.cpp_extension._get_build_directory( module_name, verbose=verbose_build) # pylint: disable=protected-access cached_build_dir = os.path.join( build_top_dir, f'{source_digest}-{_get_mangled_gpu_name()}') if not os.path.isdir(cached_build_dir): tmpdir = f'{build_top_dir}/srctmp-{uuid.uuid4().hex}' os.makedirs(tmpdir) for src in all_source_files: shutil.copyfile( src, os.path.join(tmpdir, os.path.basename(src))) try: os.replace(tmpdir, cached_build_dir) # atomic except OSError: # source directory already exists, delete tmpdir and its contents. shutil.rmtree(tmpdir) if not os.path.isdir(cached_build_dir): raise # Compile. cached_sources = [ os.path.join(cached_build_dir, os.path.basename(fname)) for fname in sources ] torch.utils.cpp_extension.load(name=module_name, build_directory=cached_build_dir, verbose=verbose_build, sources=cached_sources, **build_kwargs) else: torch.utils.cpp_extension.load(name=module_name, verbose=verbose_build, sources=sources, **build_kwargs) # Load. module = importlib.import_module(module_name) except: if verbosity == 'brief': print('Failed!') raise # Print status and add to cache dict. if verbosity == 'full': print(f'Done setting up PyTorch plugin "{module_name}".') elif verbosity == 'brief': print('Done.') _cached_plugins[module_name] = module return module #----------------------------------------------------------------------------