1 #pragma once 2 3 #include <Python.h> 4 5 #include <torch/csrc/profiler/collection.h> 6 #include <torch/csrc/profiler/python/pybind.h> 7 8 namespace pybind11::detail { 9 using torch::profiler::impl::TensorID; 10 11 #define STRONG_POINTER_TYPE_CASTER(T) \ 12 template <> \ 13 struct type_caster<T> : public strong_pointer_type_caster<T> {}; 14 15 STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::StorageImplData); 16 STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::AllocationID); 17 STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::TensorImplAddress); 18 STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::PyModuleSelf); 19 STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::PyModuleCls); 20 STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::PyOptimizerSelf); 21 #undef STRONG_POINTER_TYPE_CASTER 22 23 template <> 24 struct type_caster<TensorID> : public strong_uint_type_caster<TensorID> {}; 25 } // namespace pybind11::detail 26 27 namespace torch::profiler { 28 29 void initPythonBindings(PyObject* module); 30 31 } // namespace torch::profiler 32