NVIDIA Toolkit¶

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
NVTX + CUDA 测试程序（Python）
- 自动选择后端：PyTorch -> CuPy -> Numba
- 在关键阶段打 NVTX 区段（range）与点标记（marker）
- 进行一次 warmup，再进行多次重复以便观察
"""

import time
import math

# ---- NVTX（CPU 侧区段/标记）----
try:
    import nvtx
except Exception as e:
    raise SystemExit(
        "未找到 nvtx 库，请先安装：\n  pip install nvtx\n原始错误：%s" % e
    )

REPEATS = 5
N = 1_000_000  # 元素数量，可调大以拉开时间
DEVICE_ID = 0

def now_ms():
    return time.perf_counter() * 1000.0

def with_nvtx_range(name):
    # 轻量包装，统一风格
    return nvtx.annotate(message=name, domain="demo")

def torch_backend():
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("PyTorch 可用，但未检测到 CUDA 设备。")

    torch.cuda.set_device(DEVICE_ID)
    # 也用 torch 自带的 NVTX（作用等价于 nvtx.annotate）
    from torch.cuda import nvtx as tnvtx

    with with_nvtx_range("torch:init"):
        a = torch.rand(N, device="cuda")
        b = torch.rand(N, device="cuda")
        c = torch.empty_like(a)

    # warmup
    with tnvtx.range("torch:warmup"):
        c = a + b
        torch.cuda.synchronize()

    # 重复多次，便于观察
    for i in range(REPEATS):
        with tnvtx.range(f"torch:iter_{i}:compute(+sync)"):
            c = a + b
            torch.cuda.synchronize()
        nvtx.mark(message=f"torch:iter_{i}:done", domain="demo")

    # 额外做一次小的归约，拉起 cuBLAS/cuDNN 之外的内核
    with tnvtx.range("torch:reduce(sum)+sync"):
        s = torch.sum(c)
        torch.cuda.synchronize()

    # 防止优化器消除
    print("torch sum:", float(s))

def cupy_backend():
    import cupy as cp
    from cupy.cuda import nvtx as cnvtx

    cp.cuda.Device(DEVICE_ID).use()

    with with_nvtx_range("cupy:init"):
        a = cp.random.random(N, dtype=cp.float32)
        b = cp.random.random(N, dtype=cp.float32)
        c = cp.empty_like(a)

    # warmup
    with cnvtx.RangePush("cupy:warmup"):
        c = a + b
        cp.cuda.runtime.deviceSynchronize()
    cnvtx.RangePop()

    for i in range(REPEATS):
        cnvtx.RangePush(f"cupy:iter_{i}:compute(+sync)")
        c = a + b
        cp.cuda.runtime.deviceSynchronize()
        cnvtx.RangePop()
        nvtx.mark(message=f"cupy:iter_{i}:done", domain="demo")

    with cnvtx.RangePush("cupy:reduce(sum)+sync"):
        s = float(cp.sum(c).get())
        cp.cuda.runtime.deviceSynchronize()
    cnvtx.RangePop()

    print("cupy sum:", s)

def numba_backend():
    import numpy as np
    from numba import cuda

    cuda.select_device(DEVICE_ID)

    @cuda.jit
    def vadd(x, y, z):
        i = cuda.grid(1)
        if i < x.size:
            z[i] = x[i] + y[i]

    threads = 256
    blocks = math.ceil(N / threads)

    with with_nvtx_range("numba:host_init"):
        ha = np.random.rand(N).astype(np.float32)
        hb = np.random.rand(N).astype(np.float32)

    with with_nvtx_range("numba:to_device"):
        da = cuda.to_device(ha)
        db = cuda.to_device(hb)
        dc = cuda.device_array_like(da)

    # warmup
    with with_nvtx_range("numba:warmup(+sync)"):
        vadd[blocks, threads](da, db, dc)
        cuda.synchronize()

    for i in range(REPEATS):
        with with_nvtx_range(f"numba:iter_{i}:kernel(+sync)"):
            vadd[blocks, threads](da, db, dc)
            cuda.synchronize()
        nvtx.mark(message=f"numba:iter_{i}:done", domain="demo")

    with with_nvtx_range("numba:copy_back"):
        hc = dc.copy_to_host()

    print("numba sum sample:", float(hc[:5].sum()))

def main():
    started = now_ms()
    print("NVTX CUDA demo start")

    # 优先顺序：torch -> cupy -> numba
    for name, runner in [
        ("PyTorch", torch_backend),
        ("CuPy", cupy_backend),
        ("Numba", numba_backend),
    ]:
        try:
            with with_nvtx_range(f"backend:{name}"):
                runner()
            print(f"使用后端：{name}")
            break
        except Exception as e:
            print(f"[跳过 {name}] 原因：{e}")
    else:
        raise SystemExit(
            "未找到可用的 CUDA Python 后端（PyTorch / CuPy / Numba 任一）。\n"
            "请先安装其中之一，并确保 CUDA 驱动可用。"
        )

    print("Total elapsed: %.2f ms" % (now_ms() - started))

if __name__ == "__main__":
    main()