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
Initialize(const std::vector<std::string> & argv)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(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
RegisterBenchmark(const std::string & name,nb::callable f)37 benchmark::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
NB_MODULE(_benchmark,m)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", nb::is_arithmetic())
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("__or__", [](Counter::Flags a, Counter::Flags b) {
134 return static_cast<int>(a) | static_cast<int>(b);
135 });
136
137 nb::enum_<Counter::OneK>(py_counter, "OneK")
138 .value("kIs1000", Counter::OneK::kIs1000)
139 .value("kIs1024", Counter::OneK::kIs1024)
140 .export_values();
141
142 py_counter
143 .def(
144 "__init__",
145 [](Counter* c, double value, int flags, Counter::OneK oneK) {
146 new (c) Counter(value, static_cast<Counter::Flags>(flags), oneK);
147 },
148 nb::arg("value") = 0., nb::arg("flags") = Counter::kDefaults,
149 nb::arg("k") = Counter::kIs1000)
150 .def("__init__",
151 ([](Counter* c, double value) { new (c) Counter(value); }))
152 .def_rw("value", &Counter::value)
153 .def_rw("flags", &Counter::flags)
154 .def_rw("oneK", &Counter::oneK)
155 .def(nb::init_implicit<double>());
156
157 nb::implicitly_convertible<nb::int_, Counter>();
158
159 nb::bind_map<benchmark::UserCounters>(m, "UserCounters");
160
161 using benchmark::State;
162 nb::class_<State>(m, "State")
163 .def("__bool__", &State::KeepRunning)
164 .def_prop_ro("keep_running", &State::KeepRunning)
165 .def("pause_timing", &State::PauseTiming)
166 .def("resume_timing", &State::ResumeTiming)
167 .def("skip_with_error", &State::SkipWithError)
168 .def_prop_ro("error_occurred", &State::error_occurred)
169 .def("set_iteration_time", &State::SetIterationTime)
170 .def_prop_rw("bytes_processed", &State::bytes_processed,
171 &State::SetBytesProcessed)
172 .def_prop_rw("complexity_n", &State::complexity_length_n,
173 &State::SetComplexityN)
174 .def_prop_rw("items_processed", &State::items_processed,
175 &State::SetItemsProcessed)
176 .def("set_label", &State::SetLabel)
177 .def("range", &State::range, nb::arg("pos") = 0)
178 .def_prop_ro("iterations", &State::iterations)
179 .def_prop_ro("name", &State::name)
180 .def_rw("counters", &State::counters)
181 .def_prop_ro("thread_index", &State::thread_index)
182 .def_prop_ro("threads", &State::threads);
183
184 m.def("Initialize", Initialize);
185 m.def("RegisterBenchmark", RegisterBenchmark,
186 nb::rv_policy::reference);
187 m.def("RunSpecifiedBenchmarks",
188 []() { benchmark::RunSpecifiedBenchmarks(); });
189 m.def("ClearRegisteredBenchmarks", benchmark::ClearRegisteredBenchmarks);
190 };
191 } // namespace
192