Handles the loading of c++ extension modules. Args: module_name: Name of the module to load. Must match the name of the relevant source directory in the `_extensions` directory. defines: Dictionary containing names and values of compilation defines. verb
(
module_name: str, defines: dict | None = None, verbose_build: bool = False, build_timeout: int = 300
)
| 47 | |
| 48 | |
| 49 | def load_module( |
| 50 | module_name: str, defines: dict | None = None, verbose_build: bool = False, build_timeout: int = 300 |
| 51 | ) -> ModuleType: |
| 52 | """ |
| 53 | Handles the loading of c++ extension modules. |
| 54 | |
| 55 | Args: |
| 56 | module_name: Name of the module to load. |
| 57 | Must match the name of the relevant source directory in the `_extensions` directory. |
| 58 | defines: Dictionary containing names and values of compilation defines. |
| 59 | verbose_build: Set to true to enable build logging. |
| 60 | build_timeout: Time in seconds before the build will throw an exception to prevent hanging. |
| 61 | """ |
| 62 | |
| 63 | # Ensuring named module exists in _extensions directory. |
| 64 | module_dir = path.join(dir_path, module_name) |
| 65 | if not path.exists(module_dir): |
| 66 | raise ValueError(f"No extension module named {module_name}") |
| 67 | |
| 68 | platform_str = f"_{platform.system()}_{platform.python_version()}_" |
| 69 | platform_str += "".join(f"{v}" for v in get_torch_version_tuple()[:2]) |
| 70 | # Adding configuration to module name. |
| 71 | if defines is not None: |
| 72 | module_name = "_".join([module_name] + [f"{v}" for v in defines.values()]) |
| 73 | |
| 74 | # Gathering source files. |
| 75 | source = glob(path.join(module_dir, "**", "*.cpp"), recursive=True) |
| 76 | if torch.cuda.is_available(): |
| 77 | source += glob(path.join(module_dir, "**", "*.cu"), recursive=True) |
| 78 | platform_str += f"_{torch.version.cuda}" |
| 79 | |
| 80 | # Constructing compilation argument list. |
| 81 | define_args = [] if not defines else [f"-D {key}={defines[key]}" for key in defines] |
| 82 | |
| 83 | # Ninja may be blocked by something out of our control. |
| 84 | # This will error if the build takes longer than expected. |
| 85 | with timeout(build_timeout, "Build appears to be blocked. Is there a stopped process building the same extension?"): |
| 86 | load, _ = optional_import("torch.utils.cpp_extension", name="load") # main trigger some JIT config in pytorch |
| 87 | # This will either run the build or return the existing .so object. |
| 88 | name = module_name + platform_str.replace(".", "_") |
| 89 | module = load( |
| 90 | name=name, sources=source, extra_cflags=define_args, extra_cuda_cflags=define_args, verbose=verbose_build |
| 91 | ) |
| 92 | |
| 93 | return module # type: ignore[no-any-return] |
searching dependent graphs…