(module_name, sources, headers=None, source_dir=None, **build_kwargs)
| 59 | _cached_plugins = dict() |
| 60 | |
| 61 | def get_plugin(module_name, sources, headers=None, source_dir=None, **build_kwargs): |
| 62 | assert verbosity in ['none', 'brief', 'full'] |
| 63 | if headers is None: |
| 64 | headers = [] |
| 65 | if source_dir is not None: |
| 66 | sources = [os.path.join(source_dir, fname) for fname in sources] |
| 67 | headers = [os.path.join(source_dir, fname) for fname in headers] |
| 68 | |
| 69 | # Already cached? |
| 70 | if module_name in _cached_plugins: |
| 71 | return _cached_plugins[module_name] |
| 72 | |
| 73 | # Print status. |
| 74 | if verbosity == 'full': |
| 75 | print(f'Setting up PyTorch plugin "{module_name}"...') |
| 76 | elif verbosity == 'brief': |
| 77 | print(f'Setting up PyTorch plugin "{module_name}"... ', end='', flush=True) |
| 78 | verbose_build = (verbosity == 'full') |
| 79 | |
| 80 | # Compile and load. |
| 81 | try: # pylint: disable=too-many-nested-blocks |
| 82 | # Make sure we can find the necessary compiler binaries. |
| 83 | if os.name == 'nt' and os.system("where cl.exe >nul 2>nul") != 0: |
| 84 | compiler_bindir = _find_compiler_bindir() |
| 85 | if compiler_bindir is None: |
| 86 | raise RuntimeError(f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in "{__file__}".') |
| 87 | os.environ['PATH'] += ';' + compiler_bindir |
| 88 | |
| 89 | # Some containers set TORCH_CUDA_ARCH_LIST to a list that can either |
| 90 | # break the build or unnecessarily restrict what's available to nvcc. |
| 91 | # Unset it to let nvcc decide based on what's available on the |
| 92 | # machine. |
| 93 | os.environ['TORCH_CUDA_ARCH_LIST'] = '' |
| 94 | |
| 95 | # Incremental build md5sum trickery. Copies all the input source files |
| 96 | # into a cached build directory under a combined md5 digest of the input |
| 97 | # source files. Copying is done only if the combined digest has changed. |
| 98 | # This keeps input file timestamps and filenames the same as in previous |
| 99 | # extension builds, allowing for fast incremental rebuilds. |
| 100 | # |
| 101 | # This optimization is done only in case all the source files reside in |
| 102 | # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR |
| 103 | # environment variable is set (we take this as a signal that the user |
| 104 | # actually cares about this.) |
| 105 | # |
| 106 | # EDIT: We now do it regardless of TORCH_EXTENSIOS_DIR, in order to work |
| 107 | # around the *.cu dependency bug in ninja config. |
| 108 | # |
| 109 | all_source_files = sorted(sources + headers) |
| 110 | all_source_dirs = set(os.path.dirname(fname) for fname in all_source_files) |
| 111 | if len(all_source_dirs) == 1: # and ('TORCH_EXTENSIONS_DIR' in os.environ): |
| 112 | |
| 113 | # Compute combined hash digest for all source files. |
| 114 | hash_md5 = hashlib.md5() |
| 115 | for src in all_source_files: |
| 116 | with open(src, 'rb') as f: |
| 117 | hash_md5.update(f.read()) |
| 118 |
nothing calls this directly
no test coverage detected