mirror of
https://github.com/NVIDIA/nvbench.git
synced 2026-07-16 08:44:05 +00:00
44ec7de6bd
* Add decorators for registering benchmarks and adding axis
cuda.bench.register(fn) continues returning Benchmark, and supports
legacy use.
New signature added:
cuda.bench.register():
Returns a decorator
```
@bench.register()
@bench.axis.float64("Duration (s)", [7e-5, 1e-4, 5e-4])
@bench.option.min_samples(120)
def single_float64_axis(state: bench.State):
...
```
* Remove example/auto_throughput.py
The C++ counterpart's purpose is to demonstrate use of CUPTI
metrics, but these are not supported in Python bindings, so
this example is a duplicate of example/throughput.py
* Add wrong decorator order test for bench.axis.*
* Strengthen type annotation for register function
Acting on code rabbit nit-pick require that function being
registered take cuda.bench.State object as an argument.
Verified the fix as
```
(py313) :~/repos/nvbench/python$ python -m mypy --ignore-missing-import /tmp/t.py
/tmp/t.py:8: error: Argument 1 has incompatible type "Callable[[], None]"; expected "Callable[[State], None]" [arg-type]
Found 1 error in 1 file (checked 1 source file)
(py313) :~/repos/nvbench/python$ nl -ba /tmp/t.py
1 # /tmp/check_nvbench_register.py
2 import cuda.bench as bench
3
4 @bench.register()
5 def good(state: bench.State) -> None:
6 pass
7
8 @bench.register()
9 def bad() -> None:
10 pass
```
* Replace use of global variable with thread-safe lru_cache
This improves thread-safety of module initialization.
* Abide by RUF005 linting rule
* Expand docstrings regarding cuda.bench.register() decorator
It explains to the user what the decorator does and provides
a concise usage example.
* Sharpen wording on exception maybe-thrown by decorator
150 lines
3.9 KiB
Python
150 lines
3.9 KiB
Python
# Copyright 2025-2026 NVIDIA Corporation
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 with the LLVM exception
|
|
# (the "License"); you may not use this file except in compliance with
|
|
# the License.
|
|
#
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://llvm.org/foundation/relicensing/LICENSE.txt
|
|
#
|
|
# 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.
|
|
|
|
"""CUDA Kernel Benchmarking Library Python API."""
|
|
|
|
import functools
|
|
import importlib
|
|
import importlib.metadata
|
|
import warnings
|
|
|
|
from ._decorators import axis as axis
|
|
from ._decorators import make_register as _make_register
|
|
from ._decorators import option as option
|
|
|
|
try:
|
|
__version__ = importlib.metadata.version("cuda-bench")
|
|
except Exception as e:
|
|
__version__ = "0.0.0dev"
|
|
warnings.warn(
|
|
"Could not retrieve version of cuda-bench package dynamically from its metadata. "
|
|
f"Exception {e} was raised. "
|
|
f"Version is set to fall-back value '{__version__}' instead."
|
|
)
|
|
|
|
|
|
_NVBENCH_EXPORTS = (
|
|
"Benchmark",
|
|
"CudaStream",
|
|
"Launch",
|
|
"NVBenchRuntimeError",
|
|
"State",
|
|
"run_all_benchmarks",
|
|
)
|
|
|
|
_PUBLIC_EXPORTS = (
|
|
*_NVBENCH_EXPORTS,
|
|
"axis",
|
|
"option",
|
|
"register",
|
|
)
|
|
|
|
_NVBENCH_TEST_EXPORTS = (
|
|
"_test_cpp_exception",
|
|
"_test_py_exception",
|
|
)
|
|
|
|
__all__ = list(_PUBLIC_EXPORTS)
|
|
|
|
# Optional test override used by decorator tests.
|
|
_register = None
|
|
|
|
|
|
# Detect CUDA runtime version and load appropriate extension
|
|
def _get_cuda_major_version():
|
|
"""Detect the CUDA runtime major version."""
|
|
try:
|
|
import cuda.bindings
|
|
|
|
# Get CUDA version from cuda-bindings package version
|
|
# cuda-bindings version is in format like "12.9.1" or "13.0.0"
|
|
version_str = cuda.bindings.__version__
|
|
major = int(version_str.split(".")[0])
|
|
return major
|
|
except ImportError:
|
|
raise ImportError(
|
|
"cuda-bindings is required for runtime CUDA version detection. "
|
|
"Install with: pip install cuda-bench[cu12] or pip install cuda-bench[cu13]"
|
|
)
|
|
|
|
|
|
def _bind_nvbench_module(module):
|
|
for name in _NVBENCH_EXPORTS:
|
|
globals()[name] = getattr(module, name)
|
|
# Set module of exposed objects
|
|
globals()[name].__module__ = __name__
|
|
|
|
for name in _NVBENCH_TEST_EXPORTS:
|
|
globals()[name] = getattr(module, name)
|
|
|
|
# Expose the module as _nvbench for backward compatibility (e.g., for tests)
|
|
globals()["_nvbench"] = module
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _load_nvbench_module():
|
|
cuda_major = _get_cuda_major_version()
|
|
extra_name = f"cu{cuda_major}"
|
|
module_fullname = f"cuda.bench.{extra_name}._nvbench"
|
|
|
|
try:
|
|
module = importlib.import_module(module_fullname)
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
f"No cuda-bench extension found for CUDA {cuda_major}.x. "
|
|
f"This wheel may not include support for your CUDA version. "
|
|
f"Supported CUDA versions: 12, 13. "
|
|
f"Original error: {e}"
|
|
) from e
|
|
|
|
_bind_nvbench_module(module)
|
|
return module
|
|
|
|
|
|
def __getattr__(name):
|
|
if name == "_nvbench":
|
|
return _load_nvbench_module()
|
|
|
|
if name in _NVBENCH_EXPORTS + _NVBENCH_TEST_EXPORTS:
|
|
_load_nvbench_module()
|
|
return globals()[name]
|
|
|
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
|
|
|
|
def __dir__():
|
|
return sorted(
|
|
set(globals())
|
|
| set(_PUBLIC_EXPORTS)
|
|
| set(_NVBENCH_TEST_EXPORTS)
|
|
| {"_nvbench"}
|
|
)
|
|
|
|
|
|
__doc__ = """
|
|
CUDA Kernel Benchmarking Library Python API
|
|
"""
|
|
|
|
|
|
def _get_register():
|
|
if _register is not None:
|
|
return _register
|
|
return _load_nvbench_module().register
|
|
|
|
|
|
register = _make_register(_get_register)
|
|
register.__module__ = __name__
|