(self, verbose=True)
| 445 | return self.jit_load(verbose) |
| 446 | |
| 447 | def jit_load(self, verbose=True): |
| 448 | if not self.is_compatible(verbose): |
| 449 | raise RuntimeError( |
| 450 | f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}" |
| 451 | ) |
| 452 | try: |
| 453 | import ninja # noqa: F401 |
| 454 | except ImportError: |
| 455 | raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") |
| 456 | |
| 457 | if isinstance(self, CUDAOpBuilder) and not self.is_rocm_pytorch(): |
| 458 | try: |
| 459 | assert_no_cuda_mismatch(self.name) |
| 460 | self.build_for_cpu = False |
| 461 | except BaseException: |
| 462 | self.build_for_cpu = True |
| 463 | |
| 464 | self.jit_mode = True |
| 465 | from torch.utils.cpp_extension import load |
| 466 | |
| 467 | start_build = time.time() |
| 468 | sources = [self.sat_src_path(path) for path in self.sources()] |
| 469 | extra_include_paths = [self.sat_src_path(path) for path in self.include_paths()] |
| 470 | |
| 471 | # Torch will try and apply whatever CCs are in the arch list at compile time, |
| 472 | # we have already set the intended targets ourselves we know that will be |
| 473 | # needed at runtime. This prevents CC collisions such as multiple __half |
| 474 | # implementations. Stash arch list to reset after build. |
| 475 | torch_arch_list = None |
| 476 | if "TORCH_CUDA_ARCH_LIST" in os.environ: |
| 477 | torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") |
| 478 | os.environ["TORCH_CUDA_ARCH_LIST"] = "" |
| 479 | |
| 480 | op_module = load(name=self.name, |
| 481 | sources=self.strip_empty_entries(sources), |
| 482 | extra_include_paths=self.strip_empty_entries(extra_include_paths), |
| 483 | extra_cflags=self.strip_empty_entries(self.cxx_args()), |
| 484 | extra_cuda_cflags=self.strip_empty_entries(self.nvcc_args()), |
| 485 | extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), |
| 486 | verbose=verbose) |
| 487 | |
| 488 | build_duration = time.time() - start_build |
| 489 | if verbose: |
| 490 | print(f"Time to load {self.name} op: {build_duration} seconds") |
| 491 | |
| 492 | # Reset arch list so we are not silently removing it for other possible use cases |
| 493 | if torch_arch_list: |
| 494 | os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list |
| 495 | |
| 496 | return op_module |
| 497 | |
| 498 | |
| 499 | class CUDAOpBuilder(OpBuilder): |
no test coverage detected