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

@@ -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

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

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"

python/scripts/__init__.py

@@ -0,0 +1 @@
+# Package placeholder for nvbench CLI tools.

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)

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():

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)

scripts/nvbench_json/.gitignore

@@ -1 +0,0 @@
-__pycache__/*