mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
9465c0e0b5
We define specializations for pybind11 defined templates (in particular, PYBIND11_DECLARE_HOLDER_TYPE) and consequently it is important that these specializations *always* be #include'd when making use of pybind11 templates whose behavior depends on these specializations, otherwise we can cause an ODR violation. The easiest way to ensure that all the specializations are always loaded is to designate a header (in this case, torch/csrc/util/pybind.h) that ensures the specializations are defined, and then add a lint to ensure this header is included whenever pybind11 headers are included. The existing grep linter didn't have enough knobs to do this conveniently, so I added some features. I'm open to suggestions for how to structure the features better. The main changes: - Added an --allowlist-pattern flag, which turns off the grep lint if some other line exists. This is used to stop the grep lint from complaining about pybind11 includes if the util include already exists. - Added --match-first-only flag, which lets grep only match against the first matching line. This is because, even if there are multiple includes that are problematic, I only need to fix one of them. We don't /really/ need this, but when I was running lintrunner -a to fixup the preexisting codebase it was annoying without this, as the lintrunner overall driver fails if there are multiple edits on the same file. I excluded any files that didn't otherwise have a dependency on torch/ATen, this was mostly caffe2 and the valgrind wrapper compat bindings. Note the grep replacement is kind of crappy, but clang-tidy lint cleaned it up in most cases. See also https://github.com/pybind/pybind11/issues/4099 Signed-off-by: Edward Z. Yang <ezyang@fb.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/82552 Approved by: https://github.com/albanD
54 lines
2.2 KiB
C++
54 lines
2.2 KiB
C++
#include <ATen/core/ivalue.h>
|
|
#include <torch/csrc/utils/init.h>
|
|
#include <torch/csrc/utils/throughput_benchmark.h>
|
|
|
|
#include <pybind11/functional.h>
|
|
#include <torch/csrc/utils/pybind.h>
|
|
|
|
namespace torch {
|
|
namespace throughput_benchmark {
|
|
|
|
void initThroughputBenchmarkBindings(PyObject* module) {
|
|
auto m = py::handle(module).cast<py::module>();
|
|
using namespace torch::throughput_benchmark;
|
|
py::class_<BenchmarkConfig>(m, "BenchmarkConfig")
|
|
.def(py::init<>())
|
|
.def_readwrite(
|
|
"num_calling_threads", &BenchmarkConfig::num_calling_threads)
|
|
.def_readwrite("num_worker_threads", &BenchmarkConfig::num_worker_threads)
|
|
.def_readwrite("num_warmup_iters", &BenchmarkConfig::num_warmup_iters)
|
|
.def_readwrite("num_iters", &BenchmarkConfig::num_iters)
|
|
.def_readwrite(
|
|
"profiler_output_path", &BenchmarkConfig::profiler_output_path);
|
|
|
|
py::class_<BenchmarkExecutionStats>(m, "BenchmarkExecutionStats")
|
|
.def_readonly("latency_avg_ms", &BenchmarkExecutionStats::latency_avg_ms)
|
|
.def_readonly("num_iters", &BenchmarkExecutionStats::num_iters);
|
|
|
|
py::class_<ThroughputBenchmark>(m, "ThroughputBenchmark", py::dynamic_attr())
|
|
.def(py::init<jit::Module>())
|
|
.def(py::init<py::object>())
|
|
.def(
|
|
"add_input",
|
|
[](ThroughputBenchmark& self, py::args args, py::kwargs kwargs) {
|
|
self.addInput(std::move(args), std::move(kwargs));
|
|
})
|
|
.def(
|
|
"run_once",
|
|
[](ThroughputBenchmark& self, py::args args, py::kwargs kwargs) {
|
|
// Depending on this being ScriptModule of nn.Module we will release
|
|
// the GIL or not further down in the stack
|
|
return self.runOnce(std::move(args), std::move(kwargs));
|
|
})
|
|
.def("benchmark", [](ThroughputBenchmark& self, BenchmarkConfig config) {
|
|
// The benchmark always runs without the GIL. GIL will be used where
|
|
// needed. This will happen only in the nn.Module mode when manipulating
|
|
// inputs and running actual inference
|
|
pybind11::gil_scoped_release no_gil_guard;
|
|
return self.benchmark(config);
|
|
});
|
|
}
|
|
|
|
} // namespace throughput_benchmark
|
|
} // namespace torch
|