From 5e7adc5c3f3906a61ff7795dc64a590fc817a4b7 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 28 Jan 2026 19:13:24 -0600 Subject: [PATCH] Build multi architecture cuda wheels (#302) * Add cuda architectures to build wheel for * Package scripts in wheel * Separate cuda major version extraction to fix architecutre selection logic * Add back statement printing cuda version * [pre-commit.ci] auto code formatting --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- ci/build_pynvbench_wheel_for_cuda.sh | 25 +++++++++++++++++-- nvbench/json_printer.cu | 2 +- python/pyproject.toml | 15 +++++++++++ python/scripts/__init__.py | 1 + .../scripts}/nvbench_compare.py | 6 ++++- .../scripts}/nvbench_histogram.py | 6 ++++- .../scripts}/nvbench_json/__init__.py | 0 .../scripts}/nvbench_json/reader.py | 0 .../scripts}/nvbench_json/version.py | 0 .../scripts}/nvbench_walltime.py | 6 ++++- {scripts => python/scripts}/test_cmp.json | 0 {scripts => python/scripts}/test_ref.json | 0 scripts/nvbench_json/.gitignore | 1 - 13 files changed, 55 insertions(+), 7 deletions(-) create mode 100644 python/scripts/__init__.py rename {scripts => python/scripts}/nvbench_compare.py (99%) mode change 100755 => 100644 rename {scripts => python/scripts}/nvbench_histogram.py (96%) mode change 100755 => 100644 rename {scripts => python/scripts}/nvbench_json/__init__.py (100%) rename {scripts => python/scripts}/nvbench_json/reader.py (100%) rename {scripts => python/scripts}/nvbench_json/version.py (100%) rename {scripts => python/scripts}/nvbench_walltime.py (98%) rename {scripts => python/scripts}/test_cmp.json (100%) rename {scripts => python/scripts}/test_ref.json (100%) delete mode 100644 scripts/nvbench_json/.gitignore diff --git a/ci/build_pynvbench_wheel_for_cuda.sh b/ci/build_pynvbench_wheel_for_cuda.sh index c0289ca..b95d43b 100755 --- a/ci/build_pynvbench_wheel_for_cuda.sh +++ b/ci/build_pynvbench_wheel_for_cuda.sh @@ -23,9 +23,30 @@ set -euxo pipefail # The /workspace pathnames are hard-wired here. # Determine CUDA version from nvcc early (needed for dev package installation) -cuda_version=$(nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+' | cut -d. -f1) +cuda_version=$(nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+') +cuda_version_major=$(echo "${cuda_version}" | cut -d. -f1) echo "Detected CUDA version: ${cuda_version}" +# Select CUDA architectures for multi-arch cubins + PTX fallback (if not set) +if [[ -z "${CUDAARCHS:-}" ]]; then + version_ge() { + [[ "$(printf '%s\n' "$1" "$2" | sort -V | head -n1)" == "$2" ]] + } + + if version_ge "${cuda_version}" "13.0"; then + CUDAARCHS="75-real;80-real;86-real;90a-real;100f-real;120a-real;120-virtual" + elif version_ge "${cuda_version}" "12.9"; then + CUDAARCHS="70-real;75-real;80-real;86-real;90a-real;100f-real;120a-real;120-virtual" + else + CUDAARCHS="70-real;75-real;80-real;86-real;90a-real;90-virtual" + if version_ge "${cuda_version}" "12.8"; then + CUDAARCHS="70-real;75-real;80-real;86-real;90a-real;100-real;120a-real;120-virtual" + fi + fi +fi +export CUDAARCHS +echo "Using CUDAARCHS: ${CUDAARCHS}" + # Install GCC 13 toolset (needed for the build) /workspace/ci/util/retry.sh 5 30 dnf -y install gcc-toolset-13-gcc gcc-toolset-13-gcc-c++ echo -e "#!/bin/bash\nsource /opt/rh/gcc-toolset-13/enable" >/etc/profile.d/enable_devtools.sh @@ -60,7 +81,7 @@ python -m pip wheel --no-deps --verbose --wheel-dir dist . for wheel in dist/pynvbench-*.whl; do if [[ -f "$wheel" ]]; then base_name=$(basename "$wheel" .whl) - new_name="${base_name}.cu${cuda_version}.whl" + new_name="${base_name}.cu${cuda_version_major}.whl" mv "$wheel" "dist/${new_name}" echo "Renamed wheel to: ${new_name}" fi diff --git a/nvbench/json_printer.cu b/nvbench/json_printer.cu index 4d3bdf0..085faae 100644 --- a/nvbench/json_printer.cu +++ b/nvbench/json_printer.cu @@ -153,7 +153,7 @@ namespace nvbench json_printer::version_t json_printer::get_json_file_version() { // This version number should stay in sync with `file_version` in - // scripts/nvbench_json/version.py. + // python/scripts/nvbench_json/version.py. // // Use semantic versioning: // Major version: backwards incompatible changes diff --git a/python/pyproject.toml b/python/pyproject.toml index e752eca..284af82 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -43,6 +43,20 @@ test-cu13 = ["pynvbench[cu13]", "pytest", "cupy-cuda13x", "numba"] # Generic test dependencies (defaults to CUDA 12) test = ["pytest", "cupy-cuda12x", "numba"] +tools = [ + "colorama", + "jsondiff", + "matplotlib", + "numpy", + "pandas", + "seaborn", + "tabulate", +] + +[project.scripts] +nvbench-compare = "scripts.nvbench_compare:main" +nvbench-histogram = "scripts.nvbench_histogram:main" +nvbench-walltime = "scripts.nvbench_walltime:main" [project.urls] Homepage = "https://github.com/NVIDIA/nvbench" @@ -75,3 +89,4 @@ fallback_version = "0.0.0" [tool.scikit-build.wheel.packages] "cuda" = "cuda" "cuda/bench" = "cuda/bench" +"scripts" = "scripts" diff --git a/python/scripts/__init__.py b/python/scripts/__init__.py new file mode 100644 index 0000000..82d7b8e --- /dev/null +++ b/python/scripts/__init__.py @@ -0,0 +1 @@ +# Package placeholder for nvbench CLI tools. diff --git a/scripts/nvbench_compare.py b/python/scripts/nvbench_compare.py old mode 100755 new mode 100644 similarity index 99% rename from scripts/nvbench_compare.py rename to python/scripts/nvbench_compare.py index bf07fe6..bfc7a32 --- a/scripts/nvbench_compare.py +++ b/python/scripts/nvbench_compare.py @@ -8,7 +8,11 @@ import sys import jsondiff import tabulate from colorama import Fore -from nvbench_json import reader + +try: + from nvbench_json import reader +except ImportError: + from scripts.nvbench_json import reader # Parse version string into tuple, "x.y.z" -> (x, y, z) diff --git a/scripts/nvbench_histogram.py b/python/scripts/nvbench_histogram.py old mode 100755 new mode 100644 similarity index 96% rename from scripts/nvbench_histogram.py rename to python/scripts/nvbench_histogram.py index 4eb6155..240e7cd --- a/scripts/nvbench_histogram.py +++ b/python/scripts/nvbench_histogram.py @@ -8,7 +8,11 @@ import matplotlib.pyplot as plt import numpy as np import pandas as pd import seaborn as sns -from nvbench_json import reader + +try: + from nvbench_json import reader +except ImportError: + from scripts.nvbench_json import reader def parse_files(): diff --git a/scripts/nvbench_json/__init__.py b/python/scripts/nvbench_json/__init__.py similarity index 100% rename from scripts/nvbench_json/__init__.py rename to python/scripts/nvbench_json/__init__.py diff --git a/scripts/nvbench_json/reader.py b/python/scripts/nvbench_json/reader.py similarity index 100% rename from scripts/nvbench_json/reader.py rename to python/scripts/nvbench_json/reader.py diff --git a/scripts/nvbench_json/version.py b/python/scripts/nvbench_json/version.py similarity index 100% rename from scripts/nvbench_json/version.py rename to python/scripts/nvbench_json/version.py diff --git a/scripts/nvbench_walltime.py b/python/scripts/nvbench_walltime.py similarity index 98% rename from scripts/nvbench_walltime.py rename to python/scripts/nvbench_walltime.py index ad78e8c..ee27fbd 100644 --- a/scripts/nvbench_walltime.py +++ b/python/scripts/nvbench_walltime.py @@ -6,7 +6,11 @@ import os import sys import tabulate -from nvbench_json import reader + +try: + from nvbench_json import reader +except ImportError: + from scripts.nvbench_json import reader # Parse version string into tuple, "x.y.z" -> (x, y, z) diff --git a/scripts/test_cmp.json b/python/scripts/test_cmp.json similarity index 100% rename from scripts/test_cmp.json rename to python/scripts/test_cmp.json diff --git a/scripts/test_ref.json b/python/scripts/test_ref.json similarity index 100% rename from scripts/test_ref.json rename to python/scripts/test_ref.json diff --git a/scripts/nvbench_json/.gitignore b/scripts/nvbench_json/.gitignore deleted file mode 100644 index 96403d3..0000000 --- a/scripts/nvbench_json/.gitignore +++ /dev/null @@ -1 +0,0 @@ -__pycache__/*