185 lines · plain
1// Benchmark for Python.2 3#include "benchmark/benchmark.h"4 5#include "nanobind/nanobind.h"6#include "nanobind/operators.h"7#include "nanobind/stl/bind_map.h"8#include "nanobind/stl/string.h"9#include "nanobind/stl/vector.h"10 11NB_MAKE_OPAQUE(benchmark::UserCounters);12 13namespace {14namespace nb = nanobind;15 16std::vector<std::string> Initialize(const std::vector<std::string>& argv) {17 // The `argv` pointers here become invalid when this function returns, but18 // benchmark holds the pointer to `argv[0]`. We create a static copy of it19 // so it persists, and replace the pointer below.20 static std::string executable_name(argv[0]);21 std::vector<char*> ptrs;22 ptrs.reserve(argv.size());23 for (auto& arg : argv) {24 ptrs.push_back(const_cast<char*>(arg.c_str()));25 }26 ptrs[0] = const_cast<char*>(executable_name.c_str());27 int argc = static_cast<int>(argv.size());28 benchmark::Initialize(&argc, ptrs.data());29 std::vector<std::string> remaining_argv;30 remaining_argv.reserve(argc);31 for (int i = 0; i < argc; ++i) {32 remaining_argv.emplace_back(ptrs[i]);33 }34 return remaining_argv;35}36 37benchmark::internal::Benchmark* RegisterBenchmark(const std::string& name,38 nb::callable f) {39 return benchmark::RegisterBenchmark(40 name, [f](benchmark::State& state) { f(&state); });41}42 43NB_MODULE(_benchmark, m) {44 45 using benchmark::TimeUnit;46 nb::enum_<TimeUnit>(m, "TimeUnit")47 .value("kNanosecond", TimeUnit::kNanosecond)48 .value("kMicrosecond", TimeUnit::kMicrosecond)49 .value("kMillisecond", TimeUnit::kMillisecond)50 .value("kSecond", TimeUnit::kSecond)51 .export_values();52 53 using benchmark::BigO;54 nb::enum_<BigO>(m, "BigO")55 .value("oNone", BigO::oNone)56 .value("o1", BigO::o1)57 .value("oN", BigO::oN)58 .value("oNSquared", BigO::oNSquared)59 .value("oNCubed", BigO::oNCubed)60 .value("oLogN", BigO::oLogN)61 .value("oNLogN", BigO::oNLogN)62 .value("oAuto", BigO::oAuto)63 .value("oLambda", BigO::oLambda)64 .export_values();65 66 using benchmark::internal::Benchmark;67 nb::class_<Benchmark>(m, "Benchmark")68 // For methods returning a pointer to the current object, reference69 // return policy is used to ask nanobind not to take ownership of the70 // returned object and avoid calling delete on it.71 // https://pybind11.readthedocs.io/en/stable/advanced/functions.html#return-value-policies72 //73 // For methods taking a const std::vector<...>&, a copy is created74 // because a it is bound to a Python list.75 // https://pybind11.readthedocs.io/en/stable/advanced/cast/stl.html76 .def("unit", &Benchmark::Unit, nb::rv_policy::reference)77 .def("arg", &Benchmark::Arg, nb::rv_policy::reference)78 .def("args", &Benchmark::Args, nb::rv_policy::reference)79 .def("range", &Benchmark::Range, nb::rv_policy::reference,80 nb::arg("start"), nb::arg("limit"))81 .def("dense_range", &Benchmark::DenseRange,82 nb::rv_policy::reference, nb::arg("start"),83 nb::arg("limit"), nb::arg("step") = 1)84 .def("ranges", &Benchmark::Ranges, nb::rv_policy::reference)85 .def("args_product", &Benchmark::ArgsProduct,86 nb::rv_policy::reference)87 .def("arg_name", &Benchmark::ArgName, nb::rv_policy::reference)88 .def("arg_names", &Benchmark::ArgNames,89 nb::rv_policy::reference)90 .def("range_pair", &Benchmark::RangePair,91 nb::rv_policy::reference, nb::arg("lo1"), nb::arg("hi1"),92 nb::arg("lo2"), nb::arg("hi2"))93 .def("range_multiplier", &Benchmark::RangeMultiplier,94 nb::rv_policy::reference)95 .def("min_time", &Benchmark::MinTime, nb::rv_policy::reference)96 .def("min_warmup_time", &Benchmark::MinWarmUpTime,97 nb::rv_policy::reference)98 .def("iterations", &Benchmark::Iterations,99 nb::rv_policy::reference)100 .def("repetitions", &Benchmark::Repetitions,101 nb::rv_policy::reference)102 .def("report_aggregates_only", &Benchmark::ReportAggregatesOnly,103 nb::rv_policy::reference, nb::arg("value") = true)104 .def("display_aggregates_only", &Benchmark::DisplayAggregatesOnly,105 nb::rv_policy::reference, nb::arg("value") = true)106 .def("measure_process_cpu_time", &Benchmark::MeasureProcessCPUTime,107 nb::rv_policy::reference)108 .def("use_real_time", &Benchmark::UseRealTime,109 nb::rv_policy::reference)110 .def("use_manual_time", &Benchmark::UseManualTime,111 nb::rv_policy::reference)112 .def(113 "complexity",114 (Benchmark * (Benchmark::*)(benchmark::BigO)) & Benchmark::Complexity,115 nb::rv_policy::reference,116 nb::arg("complexity") = benchmark::oAuto);117 118 using benchmark::Counter;119 nb::class_<Counter> py_counter(m, "Counter");120 121 nb::enum_<Counter::Flags>(py_counter, "Flags")122 .value("kDefaults", Counter::Flags::kDefaults)123 .value("kIsRate", Counter::Flags::kIsRate)124 .value("kAvgThreads", Counter::Flags::kAvgThreads)125 .value("kAvgThreadsRate", Counter::Flags::kAvgThreadsRate)126 .value("kIsIterationInvariant", Counter::Flags::kIsIterationInvariant)127 .value("kIsIterationInvariantRate",128 Counter::Flags::kIsIterationInvariantRate)129 .value("kAvgIterations", Counter::Flags::kAvgIterations)130 .value("kAvgIterationsRate", Counter::Flags::kAvgIterationsRate)131 .value("kInvert", Counter::Flags::kInvert)132 .export_values()133 .def(nb::self | nb::self);134 135 nb::enum_<Counter::OneK>(py_counter, "OneK")136 .value("kIs1000", Counter::OneK::kIs1000)137 .value("kIs1024", Counter::OneK::kIs1024)138 .export_values();139 140 py_counter141 .def(nb::init<double, Counter::Flags, Counter::OneK>(),142 nb::arg("value") = 0., nb::arg("flags") = Counter::kDefaults,143 nb::arg("k") = Counter::kIs1000)144 .def("__init__", ([](Counter *c, double value) { new (c) Counter(value); }))145 .def_rw("value", &Counter::value)146 .def_rw("flags", &Counter::flags)147 .def_rw("oneK", &Counter::oneK)148 .def(nb::init_implicit<double>());149 150 nb::implicitly_convertible<nb::int_, Counter>();151 152 nb::bind_map<benchmark::UserCounters>(m, "UserCounters");153 154 using benchmark::State;155 nb::class_<State>(m, "State")156 .def("__bool__", &State::KeepRunning)157 .def_prop_ro("keep_running", &State::KeepRunning)158 .def("pause_timing", &State::PauseTiming)159 .def("resume_timing", &State::ResumeTiming)160 .def("skip_with_error", &State::SkipWithError)161 .def_prop_ro("error_occurred", &State::error_occurred)162 .def("set_iteration_time", &State::SetIterationTime)163 .def_prop_rw("bytes_processed", &State::bytes_processed,164 &State::SetBytesProcessed)165 .def_prop_rw("complexity_n", &State::complexity_length_n,166 &State::SetComplexityN)167 .def_prop_rw("items_processed", &State::items_processed,168 &State::SetItemsProcessed)169 .def("set_label", &State::SetLabel)170 .def("range", &State::range, nb::arg("pos") = 0)171 .def_prop_ro("iterations", &State::iterations)172 .def_prop_ro("name", &State::name)173 .def_rw("counters", &State::counters)174 .def_prop_ro("thread_index", &State::thread_index)175 .def_prop_ro("threads", &State::threads);176 177 m.def("Initialize", Initialize);178 m.def("RegisterBenchmark", RegisterBenchmark,179 nb::rv_policy::reference);180 m.def("RunSpecifiedBenchmarks",181 []() { benchmark::RunSpecifiedBenchmarks(); });182 m.def("ClearRegisteredBenchmarks", benchmark::ClearRegisteredBenchmarks);183};184} // namespace185