test: add benchmarks for kernel launch for reproducibility (#528)

This PR establishes a benchmarks directory where pytest-benchmark-based
benchmarks can live. This is to serve as a baseline for #510.

Author: cpcloud

.gitignore

@@ -19,3 +19,4 @@ testing/*.ptx
 .pixi/*
 !.pixi/config.toml
 *.log
+.benchmarks

README.md

@@ -40,14 +40,14 @@ cd testing
 make -j $(nproc)
 export NUMBA_CUDA_TEST_BIN_DIR=`pwd`
 # Execute tests
-pytest -n auto -v
+pytest -n auto -v --dist loadscope
 ```
 
 Alternatively, you can use [pixi](https://pixi.sh/latest/installation/) to wrap all of that up for you:
 
 ```
 # run tests against CUDA 13
-pixi run -e cu13 test -n auto -v
+pixi run -e cu13 test -n auto -v --dist loadscope
 ```
 
 

ci/test_conda.sh

@@ -35,6 +35,7 @@ DEPENDENCIES=(
     "psutil"
     "pytest"
     "pytest-xdist"
+    "pytest-benchmark"
     "cffi"
     "ml_dtypes"
     "python=${RAPIDS_PY_VERSION}"

ci/test_simulator.sh

@@ -12,6 +12,7 @@ DEPENDENCIES=(
     "psutil"
     "pytest"
     "pytest-xdist"
+    "pytest-benchmark"
     "cffi"
     "ml_dtypes"
     "python=${RAPIDS_PY_VERSION}"

numba_cuda/numba/cuda/tests/benchmarks/__init__.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import string
+from numba import cuda
+import numpy as np
+import pytest
+
+
+@pytest.fixture
+def many_arrs():
+    return [
+        cuda.device_array(10000, dtype=np.float32)
+        for _ in range(len(string.ascii_lowercase))
+    ]
+
+
+@pytest.fixture
+def one_arr():
+    return cuda.device_array(10000, dtype=np.float32)
+
+
+def test_one_arg(benchmark, one_arr):
+    @cuda.jit("void(float32[:])")
+    def one_arg(arr1):
+        return
+
+    benchmark(one_arg[1, 1], one_arr)
+
+
+def test_many_args(benchmark, many_arrs):
+    @cuda.jit("void({})".format(", ".join(["float32[:]"] * len(many_arrs))))
+    def many_args(
+        a,
+        b,
+        c,
+        d,
+        e,
+        f,
+        g,
+        h,
+        i,
+        j,
+        k,
+        l,
+        m,
+        n,
+        o,
+        p,
+        q,
+        r,
+        s,
+        t,
+        u,
+        v,
+        w,
+        x,
+        y,
+        z,
+    ):
+        return
+
+    benchmark(many_args[1, 1], *many_arrs)

numba_cuda/numba/cuda/tests/benchmarks/test_kernel_launch.py

@@ -53,8 +53,10 @@ test = [
     "cffi",
     "pytest",
     "pytest-xdist",
+    "pytest-benchmark",
     "filecheck",
     "ml_dtypes",
+    "statistics",
 ]
 test-cu12 = [
     "nvidia-curand-cu12",
@@ -176,6 +178,26 @@ build-tests = { cmd = [
     "-C",
     "$PIXI_PROJECT_ROOT/testing",
 ] }
+bench = { cmd = [
+    "pytest",
+    "$PIXI_PROJECT_ROOT/testing",
+    "--pyargs",
+    "numba.cuda.tests.benchmarks",
+    "--benchmark-only",
+    "--benchmark-enable",
+    "--benchmark-autosave",
+    "--benchmark-group-by=name",
+] }
+benchcmp = { cmd = [
+    "pytest",
+    "$PIXI_PROJECT_ROOT/testing",
+    "--pyargs",
+    "numba.cuda.tests.benchmarks",
+    "--benchmark-only",
+    "--benchmark-enable",
+    "--benchmark-group-by=name",
+    "--benchmark-compare",
+] }
 
 [tool.pixi.target.linux.tasks.test]
 cmd = ["pytest", "$PIXI_PROJECT_ROOT/testing"]

pixi.lock

@@ -6,4 +6,4 @@
 minversion = 8.0
 consider_namespace_packages = true
 # loadscope ensures the grouping required by CUDATestCase
-addopts = --dist loadscope --pyargs numba.cuda.tests
+addopts = --pyargs numba.cuda.tests