Fix issue causing segfault

Author: shwina

ci/build_multi_cuda_wheel.sh

@@ -123,14 +123,24 @@ cd wheelhouse_merged
 $PYTHON -m wheel unpack "$cu12_wheel"
 base_dir=$(find . -maxdepth 1 -type d -name "pynvbench-*" | head -1)
 
+# Rename libnvbench.so to libnvbench_cu12.so in the base (CUDA 12) wheel
+mv "$base_dir/cuda/bench/libnvbench.so" "$base_dir/cuda/bench/libnvbench_cu12.so"
+
+# Update _nvbench_cu12.so to link to libnvbench_cu12.so instead of libnvbench.so
+patchelf --replace-needed libnvbench.so libnvbench_cu12.so "$base_dir"/cuda/bench/_nvbench_cu12*.so
+
 # Unpack CUDA 13 wheel into a temporary subdirectory
 mkdir cu13_tmp
 cd cu13_tmp
 $PYTHON -m wheel unpack "$cu13_wheel"
 cu13_dir=$(find . -maxdepth 1 -type d -name "pynvbench-*" | head -1)
 
-# Copy the CUDA 13 extension into the base wheel
+# Copy the CUDA 13 extension AND library into the base wheel
 cp "$cu13_dir"/cuda/bench/_nvbench_cu13*.so "../$base_dir/cuda/bench/"
+cp "$cu13_dir"/cuda/bench/libnvbench.so "../$base_dir/cuda/bench/libnvbench_cu13.so"
+
+# Update _nvbench_cu13.so to link to libnvbench_cu13.so instead of libnvbench.so
+patchelf --replace-needed libnvbench.so libnvbench_cu13.so "../$base_dir"/cuda/bench/_nvbench_cu13*.so
 
 # Go back and clean up
 cd ..

ci/test_pynvbench_inner.sh

@@ -40,9 +40,10 @@ fi
 
 # Determine which CUDA extra to install (defaults to cu12 if not specified)
 CUDA_EXTRA="${cuda_extra:-cu${cuda_version}}"
+TEST_EXTRA="test-cu${cuda_version}"
 
-echo "Installing wheel: $PYNVBENCH_WHEEL_PATH with extras: ${CUDA_EXTRA},test"
-python -m pip install "${PYNVBENCH_WHEEL_PATH}[${CUDA_EXTRA},test]"
+echo "Installing wheel: $PYNVBENCH_WHEEL_PATH with extras: ${CUDA_EXTRA},${TEST_EXTRA}"
+python -m pip install "${PYNVBENCH_WHEEL_PATH}[${CUDA_EXTRA},${TEST_EXTRA}]"
 
 # Run tests
 cd "/workspace/python/test/"

python/pyproject.toml

@@ -35,7 +35,13 @@ cu12 = ["cuda-bindings>=12.0.0,<13.0.0", "nvidia-cuda-cupti-cu12"]
 # CUDA 13.x dependencies
 cu13 = ["cuda-bindings>=13.0.0,<14.0.0", "nvidia-cuda-cupti>=13.0"]
 
-# Test dependencies (using CUDA 12 by default)
+# Test dependencies for CUDA 12
+test-cu12 = ["pynvbench[cu12]", "pytest", "cupy-cuda12x", "numba"]
+
+# Test dependencies for CUDA 13
+test-cu13 = ["pynvbench[cu13]", "pytest", "cupy-cuda13x", "numba"]
+
+# Generic test dependencies (defaults to CUDA 12)
 test = ["pytest", "cupy-cuda12x", "numba"]
 
 [project.urls]