NVIDIA/nvbench
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvbench.git synced 2026-05-25 15:24:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add decorators for registering benchmarks and adding axis

Browse Source 
cuda.bench.register(fn) continues returning Benchmark, and supports
legacy use.

New signature added:
   cuda.bench.register():
      Returns a decorator

```
@bench.register()
@bench.axis.float64("Duration (s)", [7e-5, 1e-4, 5e-4])
@bench.option.min_samples(120)
def single_float64_axis(state: bench.State):
   ...
```
This commit is contained in:
Oleksandr Pavlyk
2026-05-04 08:21:41 -05:00
parent 338936b6fe
commit 63019a7e42
15 changed files with 592 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/examples/cuda_compute_segmented_reduce.py
View File

@@ -1,4 +1,4 @@
# Copyright 2025 NVIDIA Corporation
# Copyright 2025-2026 NVIDIA Corporation
#
# Licensed under the Apache License, Version 2.0 with the LLVM exception
# (the "License"); you may not use this file except in compliance with
@@ -33,6 +33,9 @@ def as_cp_ExternalStream(cs: bench.CudaStream) -> cp.cuda.ExternalStream:
return cp.cuda.Stream.from_external(cs)
@bench.register()
@bench.axis.int64("numElems", [2**20, 2**22, 2**24])
@bench.axis.int64("numCols", [1024, 2048, 4096, 8192])
def segmented_reduce(state: bench.State):
"Benchmark segmented_reduce example"
n_elems = state.get_int64("numElems")
@@ -111,8 +114,4 @@ def segmented_reduce(state: bench.State):
if __name__ == "__main__":
b = bench.register(segmented_reduce)
b.add_int64_axis("numElems", [2**20, 2**22, 2**24])
b.add_int64_axis("numCols", [1024, 2048, 4096, 8192])
bench.run_all_benchmarks(sys.argv)