| 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 | |
| 11 | NB_MAKE_OPAQUE(benchmark::UserCounters); |
| 12 | |
| 13 | namespace { |
| 14 | namespace nb = nanobind; |
| 15 | |
| 16 | std::vector<std::string> Initialize(const std::vector<std::string>& argv) { |
| 17 | // The `argv` pointers here become invalid when this function returns, but |
| 18 | // benchmark holds the pointer to `argv[0]`. We create a static copy of it |
| 19 | // so it persists, and replace the pointer below. |
| 20 | static std::string executable_name(argv[0]); |
| 21 | std::vector<char*> ptrs; |
| 22 | ptrs.reserve(n: argv.size()); |
| 23 | for (auto& arg : argv) { |
| 24 | ptrs.push_back(x: 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: &argc, argv: ptrs.data()); |
| 29 | std::vector<std::string> remaining_argv; |
| 30 | remaining_argv.reserve(n: argc); |
| 31 | for (int i = 0; i < argc; ++i) { |
| 32 | remaining_argv.emplace_back(args&: ptrs[i]); |
| 33 | } |
| 34 | return remaining_argv; |
| 35 | } |
| 36 | |
| 37 | benchmark::internal::Benchmark* RegisterBenchmark(const std::string& name, |
| 38 | nb::callable f) { |
| 39 | return benchmark::RegisterBenchmark( |
| 40 | name, fn: [f](benchmark::State& state) { f(&state); }); |
| 41 | } |
| 42 | |
| 43 | NB_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, reference |
| 69 | // return policy is used to ask nanobind not to take ownership of the |
| 70 | // returned object and avoid calling delete on it. |
| 71 | // https://pybind11.readthedocs.io/en/stable/advanced/functions.html#return-value-policies |
| 72 | // |
| 73 | // For methods taking a const std::vector<...>&, a copy is created |
| 74 | // because a it is bound to a Python list. |
| 75 | // https://pybind11.readthedocs.io/en/stable/advanced/cast/stl.html |
| 76 | .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_counter |
| 141 | .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 | } // namespace |
| 185 | |