xref: /aosp_15_r20/external/pytorch/torch/csrc/itt.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/itt_wrapper.h>
2*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/utils/pybind.h>
3*da0073e9SAndroid Build Coastguard Worker 
4*da0073e9SAndroid Build Coastguard Worker namespace torch::profiler {
initIttBindings(PyObject * module)5*da0073e9SAndroid Build Coastguard Worker void initIttBindings(PyObject* module) {
6*da0073e9SAndroid Build Coastguard Worker   auto m = py::handle(module).cast<py::module>();
7*da0073e9SAndroid Build Coastguard Worker 
8*da0073e9SAndroid Build Coastguard Worker   auto itt = m.def_submodule("_itt", "VTune ITT bindings");
9*da0073e9SAndroid Build Coastguard Worker   itt.def("is_available", itt_is_available);
10*da0073e9SAndroid Build Coastguard Worker   itt.def("rangePush", itt_range_push);
11*da0073e9SAndroid Build Coastguard Worker   itt.def("rangePop", itt_range_pop);
12*da0073e9SAndroid Build Coastguard Worker   itt.def("mark", itt_mark);
13*da0073e9SAndroid Build Coastguard Worker }
14*da0073e9SAndroid Build Coastguard Worker } // namespace torch::profiler
15