lambda/generative-music-visualizer

1# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9import glob10import hashlib11import importlib12import os13import re14import shutil15import uuid16 17import torch18import torch.utils.cpp_extension19from torch.utils.file_baton import FileBaton20 21#----------------------------------------------------------------------------22# Global options.23 24verbosity = 'brief' # Verbosity level: 'none', 'brief', 'full'25 26#----------------------------------------------------------------------------27# Internal helper funcs.28 29def _find_compiler_bindir():30    patterns = [31        'C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64',32        'C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64',33        'C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64',34        'C:/Program Files (x86)/Microsoft Visual Studio */vc/bin',35    ]36    for pattern in patterns:37        matches = sorted(glob.glob(pattern))38        if len(matches):39            return matches[-1]40    return None41 42#----------------------------------------------------------------------------43 44def _get_mangled_gpu_name():45    name = torch.cuda.get_device_name().lower()46    out = []47    for c in name:48        if re.match('[a-z0-9_-]+', c):49            out.append(c)50        else:51            out.append('-')52    return ''.join(out)53 54#----------------------------------------------------------------------------55# Main entry point for compiling and loading C++/CUDA plugins.56 57_cached_plugins = dict()58 59def get_plugin(module_name, sources, headers=None, source_dir=None, **build_kwargs):60    assert verbosity in ['none', 'brief', 'full']61    if headers is None:62        headers = []63    if source_dir is not None:64        sources = [os.path.join(source_dir, fname) for fname in sources]65        headers = [os.path.join(source_dir, fname) for fname in headers]66 67    # Already cached?68    if module_name in _cached_plugins:69        return _cached_plugins[module_name]70 71    # Print status.72    if verbosity == 'full':73        print(f'Setting up PyTorch plugin "{module_name}"...')74    elif verbosity == 'brief':75        print(f'Setting up PyTorch plugin "{module_name}"... ', end='', flush=True)76    verbose_build = (verbosity == 'full')77 78    # Compile and load.79    try: # pylint: disable=too-many-nested-blocks80        # Make sure we can find the necessary compiler binaries.81        if os.name == 'nt' and os.system("where cl.exe >nul 2>nul") != 0:82            compiler_bindir = _find_compiler_bindir()83            if compiler_bindir is None:84                raise RuntimeError(f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in "{__file__}".')85            os.environ['PATH'] += ';' + compiler_bindir86 87        # Some containers set TORCH_CUDA_ARCH_LIST to a list that can either88        # break the build or unnecessarily restrict what's available to nvcc.89        # Unset it to let nvcc decide based on what's available on the90        # machine.91        os.environ['TORCH_CUDA_ARCH_LIST'] = ''92 93        # Incremental build md5sum trickery.  Copies all the input source files94        # into a cached build directory under a combined md5 digest of the input95        # source files.  Copying is done only if the combined digest has changed.96        # This keeps input file timestamps and filenames the same as in previous97        # extension builds, allowing for fast incremental rebuilds.98        #99        # This optimization is done only in case all the source files reside in100        # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR101        # environment variable is set (we take this as a signal that the user102        # actually cares about this.)103        #104        # EDIT: We now do it regardless of TORCH_EXTENSIOS_DIR, in order to work105        # around the *.cu dependency bug in ninja config.106        #107        all_source_files = sorted(sources + headers)108        all_source_dirs = set(os.path.dirname(fname) for fname in all_source_files)109        if len(all_source_dirs) == 1: # and ('TORCH_EXTENSIONS_DIR' in os.environ):110 111            # Compute combined hash digest for all source files.112            hash_md5 = hashlib.md5()113            for src in all_source_files:114                with open(src, 'rb') as f:115                    hash_md5.update(f.read())116 117            # Select cached build directory name.118            source_digest = hash_md5.hexdigest()119            build_top_dir = torch.utils.cpp_extension._get_build_directory(module_name, verbose=verbose_build) # pylint: disable=protected-access120            cached_build_dir = os.path.join(build_top_dir, f'{source_digest}-{_get_mangled_gpu_name()}')121 122            if not os.path.isdir(cached_build_dir):123                tmpdir = f'{build_top_dir}/srctmp-{uuid.uuid4().hex}'124                os.makedirs(tmpdir)125                for src in all_source_files:126                    shutil.copyfile(src, os.path.join(tmpdir, os.path.basename(src)))127                try:128                    os.replace(tmpdir, cached_build_dir) # atomic129                except OSError:130                    # source directory already exists, delete tmpdir and its contents.131                    shutil.rmtree(tmpdir)132                    if not os.path.isdir(cached_build_dir): raise133 134            # Compile.135            cached_sources = [os.path.join(cached_build_dir, os.path.basename(fname)) for fname in sources]136            torch.utils.cpp_extension.load(name=module_name, build_directory=cached_build_dir,137                verbose=verbose_build, sources=cached_sources, **build_kwargs)138        else:139            torch.utils.cpp_extension.load(name=module_name, verbose=verbose_build, sources=sources, **build_kwargs)140 141        # Load.142        module = importlib.import_module(module_name)143 144    except:145        if verbosity == 'brief':146            print('Failed!')147        raise148 149    # Print status and add to cache dict.150    if verbosity == 'full':151        print(f'Done setting up PyTorch plugin "{module_name}".')152    elif verbosity == 'brief':153        print('Done.')154    _cached_plugins[module_name] = module155    return module156 157#----------------------------------------------------------------------------158