code/day85-python/four_ways.pyThis is the source used by the lesson and its recorded evidence. Compile commands and expected output live in the directory README.
# SPDX-License-Identifier: MIT
#
# Day 85: one kernel, four Python libraries.
#
# The same vector add is launched by cuda.core, CuPy's RawKernel, Numba and
# PyCUDA against one set of device buffers, checked bit for bit against a
# NumPy reference, and timed twice: once as a first call including whatever
# compilation the library defers, and once as a steady-state mean after a
# warm-up. The two columns are the point of the day. A Python GPU benchmark
# that reports only the first call is reporting a compiler.
#
# Three of the four rows read their CUDA source out of vector_add.cu by its
# snippet markers, so there is one copy of the kernel in this directory.
# Numba's row is a reimplementation in Python and is labelled as one.
#
# CuPy owns every buffer. That is deliberate: with one allocator and one set
# of pointers, a difference between rows cannot be a difference in where the
# data sits. It also means CuPy is required even to run the other three rows.
#
# Requirements verified together on CUDA 12.6. CuPy 14.2 requires NumPy 2,
# while numba-cuda 0.30.4 still imports np.row_stack, which NumPy 2 removed.
# The compatible pair is therefore CuPy 13.6.0 and NumPy 1.26.4:
#
# pip install "cupy-cuda12x==13.6.0" "cuda-core[cu12]==1.1.1" \
# "numba-cuda[cu12]==0.30.4" "pycuda==2026.1" "numpy==1.26.4"
#
# PyCUDA was built from source against the base CUDA 12.6 toolkit. CuPy's
# NVRTC path needed CUDA_PATH to name that base environment's
# targets/x86_64-linux directory so toolkit headers were visible.
#
# Every row is optional at run time. A library that is not installed prints a
# skip line and the rest of the program continues, so a learner with three of
# the four still gets a table.
#
# Run: python3 four_ways.py
#
# VERIFIED: Tesla T4, driver 580.173.02, CUDA 12.6, 2026-09-02. Both retained
# processes produced four bit-identical outputs; see evidence/.
import re
import sys
import time
from importlib.metadata import PackageNotFoundError, version
from pathlib import Path
import numpy as np
# Same element count as kElems in vector_add.cu, so the two programs' GB/s
# columns mean the same thing. 4 Mi + 611: the 611 keeps the kernel's bounds
# check live, and 16 MiB per buffer keeps three buffers inside a free tier.
N = (1 << 22) + 611
THREADS = 256
BLOCKS = (N + THREADS - 1) // THREADS
WARMUP = 3
TIMED = 10
# 3n floats move per launch: two reads and one write.
MOVED_BYTES = 3.0 * N * 4.0
# snippet: shared-source
def kernel_source() -> str:
"""The kernel text from vector_add.cu, wrapped for the JIT front ends.
Reading it out of the .cu by marker is what makes "the same kernel" true
rather than a claim. extern "C" turns off name mangling, so all three
front ends can look the kernel up as plain "vectorAdd".
"""
src = (Path(__file__).resolve().parent / "vector_add.cu").read_text()
region = re.search(r"^// snippet: kernel$(.*?)^// end snippet$",
src, re.M | re.S)
if region is None:
raise SystemExit("no `kernel` snippet region in vector_add.cu")
return 'extern "C" {\n' + region.group(1).strip("\n") + "\n}\n"
# end snippet
def sample(k: np.ndarray) -> np.ndarray:
"""Knuth's multiplicative hash, one xor-shift, low 23 bits as a mantissa.
Mirrors sample() in vector_add.cu. Values land in [1, 2) with all 23
mantissa bits set from the index, so every sum lands in [2, 4) and rounds
off a real bit. The 2**-23 scale is a power of two, so the multiply is
exact on both sides and the two programs see identical bytes.
"""
h = (k * np.uint32(2654435761)).astype(np.uint32)
h ^= h >> np.uint32(15)
mantissa = (h & np.uint32(0x7FFFFF)).astype(np.float32)
return np.float32(1.0) + mantissa * np.float32(1.0 / 8388608.0)
def dist_version(name: str) -> str:
try:
return version(name)
except PackageNotFoundError:
return "unknown"
def gbps(ms: float) -> float:
return MOVED_BYTES / (ms * 1.0e-3) / 1.0e9
def run_cuda_core(src, d_a, d_b, d_out):
"""NVIDIA's own binding. Compilation is a call you write."""
# snippet: py-cuda-core
from cuda.core import (Device, EventOptions, LaunchConfig, Program,
ProgramOptions, launch)
dev = Device()
dev.set_current()
stream = dev.create_stream()
build = time.perf_counter()
options = ProgramOptions(std="c++17", arch=f"sm_{dev.arch}")
program = Program(src, code_type="c++", options=options)
module = program.compile("cubin", name_expressions=("vectorAdd",))
kernel = module.get_kernel("vectorAdd")
build_ms = (time.perf_counter() - build) * 1.0e3
config = LaunchConfig(grid=BLOCKS, block=THREADS)
args = (d_a.data.ptr, d_b.data.ptr, d_out.data.ptr, np.uint64(N))
# CuPy filled the inputs on a different stream from this one.
dev.sync()
# end snippet
first = time.perf_counter()
launch(stream, config, kernel, *args)
stream.sync()
first_ms = (time.perf_counter() - first) * 1.0e3
for _ in range(WARMUP):
launch(stream, config, kernel, *args)
stream.sync()
start = dev.create_event(EventOptions(timing_enabled=True))
stop = dev.create_event(EventOptions(timing_enabled=True))
stream.record(start)
for _ in range(TIMED):
launch(stream, config, kernel, *args)
stream.record(stop)
stop.sync()
mean_ms = (stop - start) / TIMED
stream.close()
return dist_version("cuda-core"), build_ms, first_ms, mean_ms
def run_cupy(src, d_a, d_b, d_out):
"""RawKernel: construction is free, the first call compiles."""
import cupy as cp
build = time.perf_counter()
kernel = cp.RawKernel(src, "vectorAdd")
build_ms = (time.perf_counter() - build) * 1.0e3
args = (d_a, d_b, d_out, np.uint64(N))
first = time.perf_counter()
kernel((BLOCKS,), (THREADS,), args)
cp.cuda.get_current_stream().synchronize()
first_ms = (time.perf_counter() - first) * 1.0e3
for _ in range(WARMUP):
kernel((BLOCKS,), (THREADS,), args)
cp.cuda.get_current_stream().synchronize()
start, stop = cp.cuda.Event(), cp.cuda.Event()
start.record()
for _ in range(TIMED):
kernel((BLOCKS,), (THREADS,), args)
stop.record()
stop.synchronize()
mean_ms = cp.cuda.get_elapsed_time(start, stop) / TIMED
return cp.__version__, build_ms, first_ms, mean_ms
def run_numba(_src, d_a, d_b, d_out):
"""The one row that is not this directory's CUDA source.
Numba compiles Python to PTX through NVVM, so the kernel below is a
reimplementation of vectorAdd, not the same characters. Any difference in
the output bits therefore blames this function first.
"""
from numba import cuda as nbcuda
# snippet: py-numba
@nbcuda.jit
def vector_add(a, b, out):
i = nbcuda.grid(1)
if i < out.size:
out[i] = a[i] + b[i]
# end snippet
# Zero-copy views over CuPy's buffers through __cuda_array_interface__.
n_a = nbcuda.as_cuda_array(d_a)
n_b = nbcuda.as_cuda_array(d_b)
n_out = nbcuda.as_cuda_array(d_out)
# Decorating does not compile. The first launch does, for these argument
# types, so build_ms here is the cost of building a dispatcher object.
build = time.perf_counter()
launcher = vector_add[BLOCKS, THREADS]
build_ms = (time.perf_counter() - build) * 1.0e3
first = time.perf_counter()
launcher(n_a, n_b, n_out)
nbcuda.synchronize()
first_ms = (time.perf_counter() - first) * 1.0e3
for _ in range(WARMUP):
launcher(n_a, n_b, n_out)
nbcuda.synchronize()
start, stop = nbcuda.event(timing=True), nbcuda.event(timing=True)
start.record()
for _ in range(TIMED):
launcher(n_a, n_b, n_out)
stop.record()
stop.synchronize()
mean_ms = nbcuda.event_elapsed_time(start, stop) / TIMED
return dist_version("numba-cuda"), build_ms, first_ms, mean_ms
def run_pycuda(src, d_a, d_b, d_out):
"""SourceModule shells out to nvcc, so the compile is eager and a process.
autoprimaryctx, not autoinit. autoinit calls Device.make_context(), which
creates a second context on the same device; every other library here
lives in the primary one, and the pointers would not be valid across the
two. autoprimaryctx retains the primary context instead.
"""
import pycuda
import pycuda.autoprimaryctx # noqa: F401
from pycuda.compiler import SourceModule
import pycuda.driver as drv
build = time.perf_counter()
module = SourceModule(src, options=["-std=c++17"])
kernel = module.get_function("vectorAdd")
build_ms = (time.perf_counter() - build) * 1.0e3
# PyCUDA reads __cuda_array_interface__ off each argument, so the CuPy
# arrays go in as device pointers with no copy and no wrapper.
args = (d_a, d_b, d_out, np.uint64(N))
shape = {"block": (THREADS, 1, 1), "grid": (BLOCKS, 1)}
first = time.perf_counter()
kernel(*args, **shape)
drv.Context.synchronize()
first_ms = (time.perf_counter() - first) * 1.0e3
for _ in range(WARMUP):
kernel(*args, **shape)
drv.Context.synchronize()
start, stop = drv.Event(), drv.Event()
start.record()
for _ in range(TIMED):
kernel(*args, **shape)
stop.record()
stop.synchronize()
mean_ms = start.time_till(stop) / TIMED
return pycuda.VERSION_TEXT, build_ms, first_ms, mean_ms
def main() -> int:
try:
import cupy as cp
except ImportError:
print("cupy is not installed; it owns the buffers every row uses",
file=sys.stderr)
return 1
device = cp.cuda.Device(0)
name = cp.cuda.runtime.getDeviceProperties(0)["name"].decode()
print(f"GPU: {name} (compute capability {device.compute_capability})")
print(f"n = {N}, {THREADS} threads per block, {BLOCKS} blocks, "
f"{N * 4 >> 20} MiB per buffer")
print(f"warm-up {WARMUP}, timed {TIMED}, copies not included\n")
src = kernel_source()
index = np.arange(N, dtype=np.uint32)
h_a = sample(2 * index)
h_b = sample(2 * index + np.uint32(1))
# The reference is float32 on both sides. One IEEE-754 addition per
# element, no reduction and no multiply to fuse with, so the correctly
# rounded answer is unique and the comparison below is exact rather than
# toleranced. A mismatch is an index bug, a flushed denormal or a
# fast-math flag, and none of those should be silent.
want = h_a + h_b
d_a = cp.asarray(h_a)
d_b = cp.asarray(h_b)
d_out = cp.empty_like(d_a)
rows = []
failed = False
for label, runner in (("cuda.core", run_cuda_core),
("CuPy RawKernel", run_cupy),
("Numba @cuda.jit", run_numba),
("PyCUDA SourceModule", run_pycuda)):
# Poison the output first. A launcher that quietly does nothing then
# fails the check instead of inheriting the previous row's answer.
d_out.fill(np.float32("nan"))
try:
version_text, build_ms, first_ms, mean_ms = runner(
src, d_a, d_b, d_out)
except ImportError as exc:
print(f"skipping {label}: {exc}")
continue
got = cp.asnumpy(d_out)
if not np.array_equal(got, want):
bad = int(np.argmax(got != want))
print(f"{label} wrong at {bad}: got {got[bad]:.9g}, "
f"want {want[bad]:.9g}", file=sys.stderr)
failed = True
rows.append((label, version_text, build_ms, first_ms, mean_ms))
print(f"{'library':<22}{'version':>10}{'build ms':>11}"
f"{'first ms':>11}{'mean ms':>10}{'GB/s':>9}")
print(f"{'-' * 21:<22}{'-' * 9:>10}{'-' * 10:>11}"
f"{'-' * 10:>11}{'-' * 9:>10}{'-' * 8:>9}")
for label, version_text, build_ms, first_ms, mean_ms in rows:
print(f"{label:<22}{version_text:>10}{build_ms:>11.1f}"
f"{first_ms:>11.1f}{mean_ms:>10.4f}{gbps(mean_ms):>9.1f}")
print("\nbuild ms is wall clock around whatever compilation step the "
"library makes you write.")
print("first ms is wall clock around the first launch, so it carries "
"any compilation the library deferred.")
print("mean ms is CUDA events around a loop of "
f"{TIMED}, after {WARMUP} warm-up launches.")
if failed:
return 1
if not rows:
print("no library ran", file=sys.stderr)
return 1
print(f"\nall {len(rows)} libraries produced bit-identical output")
return 0
if __name__ == "__main__":
sys.exit(main())