mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-03 08:14:29 +08:00
92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Utility functions for loading CPP / CUDA extensions."""
|
|
|
|
import os
|
|
import warnings
|
|
from pathlib import Path
|
|
from time import time
|
|
from types import ModuleType
|
|
from typing import Any
|
|
|
|
import torch
|
|
from packaging.specifiers import SpecifierSet
|
|
from packaging.version import Version
|
|
from torch.utils.cpp_extension import load
|
|
|
|
__all__ = ["load_cpp_extension"]
|
|
|
|
|
|
def load_cpp_extension(
|
|
name: str,
|
|
sources: list[str | Path],
|
|
cuda_version_specifiers: str | None,
|
|
fail_msg: str = "",
|
|
raise_if_failed: bool = False,
|
|
**load_kwargs: Any,
|
|
) -> ModuleType | None:
|
|
"""Load a C++ / CUDA extension using torch.utils.cpp_extension.load() if the current CUDA version satisfies it.
|
|
|
|
Loading first time may take a few mins because of the compilation, but subsequent loads are instantaneous.
|
|
|
|
Args:
|
|
name: Name of the extension.
|
|
sources: Source files to compile.
|
|
cuda_version_specifiers: Specifier (e.g. ">=11.8,<12") for CUDA versions required to enable the extension.
|
|
fail_msg: Additional message to display if the extension fails to load.
|
|
raise_if_failed: Raise an exception if the extension fails to load.
|
|
**load_kwargs: Keyword arguments to torch.utils.cpp_extension.load().
|
|
"""
|
|
ext = None
|
|
print(f"Loading extension {name}...")
|
|
start = time()
|
|
|
|
if not os.environ.get("TORCH_CUDA_ARCH_LIST"):
|
|
try:
|
|
device_capability = torch.cuda.get_device_capability()
|
|
os.environ["TORCH_CUDA_ARCH_LIST"] = f"{device_capability[0]}.{device_capability[1]}"
|
|
except Exception:
|
|
warnings.warn("GPU not detected. Please unset `TORCH_CUDA_ARCH_LIST` env variable.")
|
|
|
|
if torch.version.cuda is None:
|
|
fail_msg = f"Skipping extension {name} because CUDA is not available."
|
|
elif cuda_version_specifiers and Version(torch.version.cuda) not in SpecifierSet(
|
|
cuda_version_specifiers
|
|
):
|
|
fail_msg = (
|
|
f"Skipping extension {name} because the current CUDA version {torch.version.cuda}"
|
|
f" does not satisfy the specifiers {cuda_version_specifiers}."
|
|
)
|
|
else:
|
|
try:
|
|
ext = load(name, sources, **load_kwargs)
|
|
except Exception as e:
|
|
if not fail_msg:
|
|
fail_msg = f"Unable to load extension {name} and falling back to CPU version."
|
|
fail_msg = f"{e}\n{fail_msg}"
|
|
# RuntimeError can be raised if there are any errors while compiling the extension.
|
|
# OSError can be raised if CUDA_HOME path is not set correctly.
|
|
# subprocess.CalledProcessError can be raised on `-runtime` images where c++ is not installed.
|
|
|
|
if ext is None:
|
|
if raise_if_failed:
|
|
raise RuntimeError(fail_msg)
|
|
else:
|
|
warnings.warn(fail_msg)
|
|
else:
|
|
print(f"Loaded extension {name} in {time() - start:.1f} seconds")
|
|
return ext
|