Skip to content

Commit 00b24c1

Browse files
Add license to stub fuile, add comment about syncing impl and stubs
Add comments stating the need to keep implementation and Python stub file in sync to both files. In the stub file to comment documents use of mypy's stubgen to generate stubs and calls to compare that against current stubs. It also calls out the need to keep docstrings and doctring examples in sync with implementation.
1 parent 7268fa9 commit 00b24c1

File tree

2 files changed

+34
-6
lines changed

2 files changed

+34
-6
lines changed

python/cuda/nvbench/__init__.pyi

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,29 @@
1-
# from __future__ import annotations
1+
# Copyright 2025 NVIDIA Corporation
2+
#
3+
# Licensed under the Apache License, Version 2.0 with the LLVM exception
4+
# (the "License"); you may not use this file except in compliance with
5+
# the License.
6+
#
7+
# You may obtain a copy of the License at
8+
#
9+
# http://llvm.org/foundation/relicensing/LICENSE.txt
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
# ============================================
18+
# PLEASE KEEP IN SYNC WITH py_nvbench.cpp FILE
19+
# ============================================
20+
# Please be sure to keep these type hints and docstring in sync
21+
# with the pybind11 bindings in ``../../src/py_nvbench.cpp``
22+
23+
# Use mypy's stubgen to auto-generate stubs using
24+
# ``stubgen -m cuda.nvbench._nvbench`` and compare
25+
# stubs in generated out/cuda/nvbench/_nvbench.pyi
26+
# with definitions given here.
227

328
from collections.abc import Callable, Sequence
429
from typing import Optional, Self, Union
@@ -22,11 +47,11 @@ class CudaStream:
2247
import cuda.nvbench as nvbench
2348
2449
def bench(state: nvbench.State):
25-
dev = core.Device(state.getDevice())
50+
dev = core.Device(state.get_device())
2651
dev.set_current()
2752
# converts CudaString to core.Stream
2853
# using __cuda_stream__ protocol
29-
dev.create_stream(state.getStream())
54+
dev.create_stream(state.get_stream())
3055
"""
3156
...
3257

python/src/py_nvbench.cpp

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,12 @@ constinit std::unique_ptr<GlobalBenchmarkRegistry, py::nodelete> global_registry
209209

210210
} // end of anonymous namespace
211211

212+
// ==========================================
213+
// PLEASE KEEP IN SYNC WITH __init__.pyi FILE
214+
// ==========================================
215+
// If you modify these bindings, please be sure to update the
216+
// corresponding type hints in ``../cuda/nvbench/__init__.pyi``
217+
212218
PYBIND11_MODULE(_nvbench, m)
213219
{
214220
// == STEP 1
@@ -223,9 +229,6 @@ PYBIND11_MODULE(_nvbench, m)
223229
// It is reinitialized before running all benchmarks to set devices to use
224230
nvbench::benchmark_manager::get().initialize();
225231

226-
// TODO: Use cuModuleGetLoadingMode(&mode) to confirm that (mode == CU_MODULE_EAGER_LOADING)
227-
// and issue warning otherwise
228-
229232
// == STEP 2
230233
// Define CudaStream class
231234
// ATTN: nvbench::cuda_stream is move-only class

0 commit comments

Comments
 (0)