From aff44e4ca1f54cc0d840191775e76e4a72f76c3f Mon Sep 17 00:00:00 2001 From: Jose Baiocchi Date: Tue, 26 May 2020 14:48:04 -0700 Subject: [PATCH] Speedup python TraceMe PiperOrigin-RevId: 313271773 Change-Id: I6358253077190f43059fed416399852bab29dae6 --- tensorflow/compiler/xla/python/BUILD | 2 +- tensorflow/compiler/xla/python/xla.cc | 20 +++---- tensorflow/python/profiler/internal/BUILD | 6 +- .../profiler/internal/traceme_wrapper.cc | 18 +++--- ...me_context_manager.h => traceme_wrapper.h} | 58 +++++++++---------- tensorflow/python/profiler/trace.py | 8 +-- 6 files changed, 51 insertions(+), 61 deletions(-) rename tensorflow/python/profiler/internal/{traceme_context_manager.h => traceme_wrapper.h} (53%) diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD index 863296c681c..5b4182b75e1 100644 --- a/tensorflow/compiler/xla/python/BUILD +++ b/tensorflow/compiler/xla/python/BUILD @@ -261,7 +261,7 @@ pybind_extension( "//tensorflow/core/profiler/lib:profiler_backends", "//tensorflow/core/profiler/lib:profiler_session", "//tensorflow/core/profiler/rpc:profiler_server", - "//tensorflow/python/profiler/internal:traceme_context_manager", + "//tensorflow/python/profiler/internal:traceme_wrapper", "//tensorflow/stream_executor:device_memory_allocator", "//tensorflow/stream_executor:platform", ] + select({ diff --git a/tensorflow/compiler/xla/python/xla.cc b/tensorflow/compiler/xla/python/xla.cc index 4cf2b36db27..abf0937d057 100644 --- a/tensorflow/compiler/xla/python/xla.cc +++ b/tensorflow/compiler/xla/python/xla.cc @@ -64,7 +64,7 @@ limitations under the License. #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/errors.h" #include "tensorflow/core/profiler/rpc/profiler_server.h" -#include "tensorflow/python/profiler/internal/traceme_context_manager.h" +#include "tensorflow/python/profiler/internal/traceme_wrapper.h" #include "tensorflow/stream_executor/platform.h" namespace xla { @@ -72,7 +72,7 @@ namespace { namespace py = pybind11; -using ::tensorflow::profiler::TraceMeContextManager; +using ::tensorflow::profiler::TraceMeWrapper; struct Uniquer { absl::Mutex mu; @@ -637,23 +637,19 @@ void BuildProfilerSubmodule(py::module* m) { }, py::arg("port")); - py::class_ traceme_class(profiler, "TraceMe", - py::module_local()); + py::class_ traceme_class(profiler, "TraceMe", + py::module_local()); traceme_class.def(py::init()) - .def("__enter__", - [](py::object self) -> py::object { - py::cast(self)->Enter(); - return self; - }) + .def("__enter__", [](py::object self) -> py::object { return self; }) .def("__exit__", [](py::object self, const py::object& ex_type, const py::object& ex_value, const py::object& traceback) -> py::object { - py::cast(self)->Exit(); + py::cast(self)->Stop(); return py::none(); }) - .def("set_metadata", &TraceMeContextManager::SetMetadata) - .def_static("is_enabled", &TraceMeContextManager::IsEnabled); + .def("set_metadata", &TraceMeWrapper::SetMetadata) + .def_static("is_enabled", &TraceMeWrapper::IsEnabled); } } // namespace diff --git a/tensorflow/python/profiler/internal/BUILD b/tensorflow/python/profiler/internal/BUILD index b6648462224..6f7193b3207 100644 --- a/tensorflow/python/profiler/internal/BUILD +++ b/tensorflow/python/profiler/internal/BUILD @@ -86,14 +86,14 @@ tf_python_pybind_extension( "//tensorflow/python/profiler:__subpackages__", ], deps = [ - ":traceme_context_manager", + ":traceme_wrapper", "@pybind11", ], ) cc_library( - name = "traceme_context_manager", - hdrs = ["traceme_context_manager.h"], + name = "traceme_wrapper", + hdrs = ["traceme_wrapper.h"], features = ["-layering_check"], visibility = [ "//tensorflow/compiler/xla/python:__pkg__", diff --git a/tensorflow/python/profiler/internal/traceme_wrapper.cc b/tensorflow/python/profiler/internal/traceme_wrapper.cc index b3403fa298f..32a1f423918 100644 --- a/tensorflow/python/profiler/internal/traceme_wrapper.cc +++ b/tensorflow/python/profiler/internal/traceme_wrapper.cc @@ -13,18 +13,18 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/python/profiler/internal/traceme_wrapper.h" + #include "pybind11/attr.h" #include "pybind11/pybind11.h" -#include "tensorflow/python/profiler/internal/traceme_context_manager.h" -using ::tensorflow::profiler::TraceMeContextManager; +namespace py = ::pybind11; + +using ::tensorflow::profiler::TraceMeWrapper; PYBIND11_MODULE(_pywrap_traceme, m) { - py::class_ traceme_class(m, "TraceMe", - py::module_local()); - traceme_class.def(py::init()) - .def("Enter", &TraceMeContextManager::Enter) - .def("Exit", &TraceMeContextManager::Exit) - .def("SetMetadata", &TraceMeContextManager::SetMetadata) - .def_static("IsEnabled", &TraceMeContextManager::IsEnabled); + py::class_(m, "TraceMe", py::module_local()) + .def(py::init()) + .def("SetMetadata", &TraceMeWrapper::SetMetadata) + .def_static("IsEnabled", &TraceMeWrapper::IsEnabled); }; diff --git a/tensorflow/python/profiler/internal/traceme_context_manager.h b/tensorflow/python/profiler/internal/traceme_wrapper.h similarity index 53% rename from tensorflow/python/profiler/internal/traceme_context_manager.h rename to tensorflow/python/profiler/internal/traceme_wrapper.h index fd281684de8..c074e909640 100644 --- a/tensorflow/python/profiler/internal/traceme_context_manager.h +++ b/tensorflow/python/profiler/internal/traceme_wrapper.h @@ -12,46 +12,41 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_PYTHON_PROFILER_INTERNAL_TRACEME_CONTEXT_MANAGER_ -#define TENSORFLOW_PYTHON_PROFILER_INTERNAL_TRACEME_CONTEXT_MANAGER_ +#ifndef TENSORFLOW_PYTHON_PROFILER_INTERNAL_TRACEME_WRAPPER_ +#define TENSORFLOW_PYTHON_PROFILER_INTERNAL_TRACEME_WRAPPER_ #include #include #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" -#include "absl/types/optional.h" #include "pybind11/pytypes.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/core/profiler/lib/traceme.h" -namespace py = pybind11; - namespace tensorflow { namespace profiler { -// Helper to implement TraceMe as a context manager in Python. -class TraceMeContextManager { +// Wraps TraceMe with an interface that takes python types. +class TraceMeWrapper { public: - explicit TraceMeContextManager(py::str name, py::kwargs kwargs) - : name_(std::move(name)), kwargs_(std::move(kwargs)) {} + // pybind11::str and pybind11::kwargs are taken by const reference to avoid + // python reference-counting overhead. + TraceMeWrapper(const pybind11::str& name, const pybind11::kwargs& kwargs) + : traceme_([&]() { + std::string name_and_metadata(name); + if (!kwargs.empty()) { + AppendMetadata(&name_and_metadata, kwargs); + } + return name_and_metadata; + }) {} - void Enter() { - if (IsEnabled()) { - traceme_.emplace([this]() { - std::string name(name_); - if (!kwargs_.empty()) { - AppendMetadata(&name, kwargs_); - } - return name; - }); - } - } - - void SetMetadata(py::kwargs kwargs) { - if (TF_PREDICT_TRUE(traceme_.has_value() && !kwargs.empty())) { - traceme_->AppendMetadata([&kwargs]() { + // pybind11::kwargs is taken by const reference to avoid python + // reference-counting overhead. + void SetMetadata(const pybind11::kwargs& kwargs) { + if (TF_PREDICT_FALSE(!kwargs.empty())) { + traceme_.AppendMetadata([&]() { std::string metadata; AppendMetadata(&metadata, kwargs); return metadata; @@ -59,28 +54,27 @@ class TraceMeContextManager { } } - void Exit() { traceme_.reset(); } + void Stop() { traceme_.Stop(); } static bool IsEnabled() { return tensorflow::profiler::TraceMe::Active(); } private: // Converts kwargs to strings and appends them to name encoded as TraceMe // metadata. - static void AppendMetadata(std::string* name, const py::kwargs& kwargs) { + static void AppendMetadata(std::string* name, + const pybind11::kwargs& kwargs) { name->push_back('#'); for (const auto& kv : kwargs) { - absl::StrAppend(name, std::string(py::str(kv.first)), "=", - std::string(py::str(kv.second)), ","); + absl::StrAppend(name, std::string(pybind11::str(kv.first)), "=", + std::string(pybind11::str(kv.second)), ","); } name->back() = '#'; } - py::str name_; - py::kwargs kwargs_; - absl::optional traceme_; + tensorflow::profiler::TraceMe traceme_; }; } // namespace profiler } // namespace tensorflow -#endif // TENSORFLOW_PYTHON_PROFILER_INTERNAL_TRACEME_CONTEXT_MANAGER_ +#endif // TENSORFLOW_PYTHON_PROFILER_INTERNAL_TRACEME_WRAPPER_ diff --git a/tensorflow/python/profiler/trace.py b/tensorflow/python/profiler/trace.py index 2cdbad5118c..ea4eb060488 100644 --- a/tensorflow/python/profiler/trace.py +++ b/tensorflow/python/profiler/trace.py @@ -73,13 +73,13 @@ class Trace(object): training step being traced. """ if _pywrap_traceme.TraceMe.IsEnabled(): + # Creating _pywrap_traceme.TraceMe starts the clock. self._traceme = _pywrap_traceme.TraceMe(name, **kwargs) else: self._traceme = None def __enter__(self): - if self._traceme: - self._traceme.Enter() + # Starting the TraceMe clock here would require an extra Python->C++ call. return self def set_metadata(self, **kwargs): @@ -117,5 +117,5 @@ class Trace(object): self._traceme.SetMetadata(**kwargs) def __exit__(self, exc_type, exc_val, exc_tb): - if self._traceme: - self._traceme.Exit() + # Deallocating _pywrap_traceme.TraceMe stops the clock. + self._traceme = None