diff --git a/python/cuda/nvbench/__init__.pyi b/python/cuda/nvbench/__init__.pyi index 5e7b390..0e2690e 100644 --- a/python/cuda/nvbench/__init__.pyi +++ b/python/cuda/nvbench/__init__.pyi @@ -169,7 +169,10 @@ class State: "Get Benchmark this configuration is a part of" ... def get_throttle_threshold(self) -> float: - "Get throttle threshold value" + "Get throttle threshold value, as fraction of maximal frequency" + ... + def set_throttle_threshold(self, threshold_fraction: float) -> None: + "Set throttle threshold fraction to specified value, expected to be between 0 and 1" ... def get_min_samples(self) -> int: "Get the number of benchmark timings NVBench performs before stopping criterion begins being used" diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp index a3cba65..604746e 100644 --- a/python/src/py_nvbench.cpp +++ b/python/src/py_nvbench.cpp @@ -481,6 +481,9 @@ PYBIND11_MODULE(_nvbench, m) [](const nvbench::state &state) { return std::ref(state.get_benchmark()); }, py::return_value_policy::reference); pystate_cls.def("get_throttle_threshold", &nvbench::state::get_throttle_threshold); + pystate_cls.def("set_throttle_threshold", + &nvbench::state::set_throttle_threshold, + py::arg("throttle_fraction")); pystate_cls.def("get_min_samples", &nvbench::state::get_min_samples); pystate_cls.def("set_min_samples",