def jit_load()

in modules/SwissArmyTransformer/sat/ops/ops_builder/builder.py [0:0]
37 lines of code
13 McCabe index (conditional complexity)

    def jit_load(self, verbose=True):
        if not self.is_compatible(verbose):
            raise RuntimeError(
                f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}"
            )
        try:
            import ninja  # noqa: F401
        except ImportError:
            raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.")

        if isinstance(self, CUDAOpBuilder) and not self.is_rocm_pytorch():
            try:
                assert_no_cuda_mismatch(self.name)
                self.build_for_cpu = False
            except BaseException:
                self.build_for_cpu = True

        self.jit_mode = True
        from torch.utils.cpp_extension import load

        start_build = time.time()
        sources = [self.sat_src_path(path) for path in self.sources()]
        extra_include_paths = [self.sat_src_path(path) for path in self.include_paths()]

        # Torch will try and apply whatever CCs are in the arch list at compile time,
        # we have already set the intended targets ourselves we know that will be
        # needed at runtime. This prevents CC collisions such as multiple __half
        # implementations. Stash arch list to reset after build.
        torch_arch_list = None
        if "TORCH_CUDA_ARCH_LIST" in os.environ:
            torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST")
            os.environ["TORCH_CUDA_ARCH_LIST"] = ""

        op_module = load(name=self.name,
                         sources=self.strip_empty_entries(sources),
                         extra_include_paths=self.strip_empty_entries(extra_include_paths),
                         extra_cflags=self.strip_empty_entries(self.cxx_args()),
                         extra_cuda_cflags=self.strip_empty_entries(self.nvcc_args()),
                         extra_ldflags=self.strip_empty_entries(self.extra_ldflags()),
                         verbose=verbose)

        build_duration = time.time() - start_build
        if verbose:
            print(f"Time to load {self.name} op: {build_duration} seconds")

        # Reset arch list so we are not silently removing it for other possible use cases
        if torch_arch_list:
            os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list

        return op_module