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 | |