expose profiler locker for other kind of profiler.

PiperOrigin-RevId: 264937433

tensorflow/contrib/makefile/Makefile

@@ -655,8 +655,7 @@ $(wildcard tensorflow/core/util/*/*.cc) \
 $(wildcard tensorflow/contrib/makefile/downloads/double_conversion/double-conversion/*.cc) \
 tensorflow/core/profiler/internal/profiler_interface.cc \
 tensorflow/core/profiler/internal/traceme_recorder.cc \
-tensorflow/core/profiler/lib/profiler_session.cc \
-tensorflow/core/profiler/lib/traceme.cc \
+$(wildcard tensorflow/core/profiler/lib/*.cc) \
 tensorflow/core/util/version_info.cc
 # Remove duplicates (for version_info.cc)
 CORE_CC_ALL_SRCS := $(sort $(CORE_CC_ALL_SRCS))

tensorflow/core/platform/BUILD

@@ -309,6 +309,7 @@ filegroup(
             "**/stream_executor.h",
             "**/env_time.cc",
             "**/device_tracer.cc",
+            "**/tpu_tracer.cc",
             "**/logger.cc",
             "**/logging.cc",
             "**/human_readable_json.cc",
@@ -405,6 +406,7 @@ filegroup(
             "**/monitoring.cc",
             "**/cuda_libdevice_path.cc",
             "**/device_tracer.cc",
+            "**/tpu_tracer.cc",
             "**/logger.cc",
             "**/logging.cc",
             "**/human_readable_json.cc",

tensorflow/core/platform/default/build_config.bzl

@@ -608,13 +608,6 @@ def tf_additional_device_tracer_test_flags():
 def tf_additional_cupti_test_flags():
     return []
 
-def tf_additional_profiler_lib_deps():
-    return [
-        "//tensorflow/core/profiler/internal/cpu:host_tracer",
-    ] + if_cuda([
-        "//tensorflow/core/profiler/internal/gpu:device_tracer",
-    ])
-
 def tf_additional_libdevice_data():
     return []
 

tensorflow/core/platform/default/device_tracer.cc

@@ -653,7 +653,7 @@ Status DeviceTracer::CollectData(RunMetadata* run_metadata) {
 }  // namespace
 
 // Not in anonymous namespace for testing purposes.
-std::unique_ptr<profiler::ProfilerInterface> CreateDeviceTracer() {
+std::unique_ptr<profiler::ProfilerInterface> CreateGpuTracer() {
   auto status = cuInit(0);
   if (status != CUDA_SUCCESS) {
     LogIfError(ToStatus(status));
@@ -666,7 +666,7 @@ auto register_device_tracer_factory = [] {
   bool enable;
   TF_CHECK_OK(ReadBoolFromEnvVar("TF_ENABLE_OSS_GPU_PROFILER", true, &enable));
   if (enable) {
-    RegisterProfilerFactory(&CreateDeviceTracer);
+    RegisterProfilerFactory(&CreateGpuTracer);
   }
   return 0;
 }();

tensorflow/core/platform/device_tracer_test.cc

@@ -41,10 +41,10 @@ limitations under the License.
 namespace tensorflow {
 
 #if GOOGLE_CUDA
-std::unique_ptr<profiler::ProfilerInterface> CreateDeviceTracer();
+std::unique_ptr<profiler::ProfilerInterface> CreateGpuTracer();
 #else
 // We don't have device tracer for non-cuda case.
-std::unique_ptr<profiler::ProfilerInterface> CreateDeviceTracer() {
+std::unique_ptr<profiler::ProfilerInterface> CreateGpuTracer() {
   return nullptr;
 }
 #endif
@@ -108,21 +108,21 @@ class DeviceTracerTest : public ::testing::Test {
 };
 
 TEST_F(DeviceTracerTest, StartStop) {
-  auto tracer = CreateDeviceTracer();
+  auto tracer = CreateGpuTracer();
   if (!tracer) return;
   TF_EXPECT_OK(tracer->Start());
   TF_EXPECT_OK(tracer->Stop());
 }
 
 TEST_F(DeviceTracerTest, StopBeforeStart) {
-  auto tracer = CreateDeviceTracer();
+  auto tracer = CreateGpuTracer();
   if (!tracer) return;
   TF_EXPECT_OK(tracer->Stop());
   TF_EXPECT_OK(tracer->Stop());
 }
 
 TEST_F(DeviceTracerTest, CollectBeforeStart) {
-  auto tracer = CreateDeviceTracer();
+  auto tracer = CreateGpuTracer();
   if (!tracer) return;
   RunMetadata run_metadata;
   TF_EXPECT_OK(tracer->CollectData(&run_metadata));
@@ -130,7 +130,7 @@ TEST_F(DeviceTracerTest, CollectBeforeStart) {
 }
 
 TEST_F(DeviceTracerTest, CollectBeforeStop) {
-  auto tracer = CreateDeviceTracer();
+  auto tracer = CreateGpuTracer();
   if (!tracer) return;
   TF_EXPECT_OK(tracer->Start());
   RunMetadata run_metadata;
@@ -140,8 +140,8 @@ TEST_F(DeviceTracerTest, CollectBeforeStop) {
 }
 
 TEST_F(DeviceTracerTest, StartTwoTracers) {
-  auto tracer1 = CreateDeviceTracer();
-  auto tracer2 = CreateDeviceTracer();
+  auto tracer1 = CreateGpuTracer();
+  auto tracer2 = CreateGpuTracer();
   if (!tracer1 || !tracer2) return;
 
   TF_EXPECT_OK(tracer1->Start());
@@ -154,7 +154,7 @@ TEST_F(DeviceTracerTest, StartTwoTracers) {
 
 TEST_F(DeviceTracerTest, RunWithTracer) {
   // On non-GPU platforms, we may not support DeviceTracer.
-  auto tracer = CreateDeviceTracer();
+  auto tracer = CreateGpuTracer();
   if (!tracer) return;
 
   Initialize({3, 2, -1, 0});
@@ -181,7 +181,7 @@ TEST_F(DeviceTracerTest, RunWithTracer) {
 }
 
 TEST_F(DeviceTracerTest, TraceToStepStatsCollector) {
-  auto tracer = CreateDeviceTracer();
+  auto tracer = CreateGpuTracer();
   if (!tracer) return;
 
   Initialize({3, 2, -1, 0});

tensorflow/core/profiler/lib/BUILD

@@ -2,10 +2,7 @@ load(
     "//tensorflow:tensorflow.bzl",
     "tf_cuda_library",
 )
-load(
-    "//tensorflow/core/platform:default/build_config.bzl",
-    "tf_additional_profiler_lib_deps",
-)
+load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 
 package(
     default_visibility = [
@@ -26,6 +23,7 @@ tf_cuda_library(
     visibility = ["//tensorflow:internal"],
     deps = [
         "//tensorflow/core/profiler/internal:profiler_interface",
+        "//tensorflow/core/profiler/lib:profiler_utils",
         "//tensorflow/core/profiler:protos_all_cc",
         "@com_google_absl//absl/strings",
     ] + select({
@@ -45,7 +43,11 @@ tf_cuda_library(
 tf_cuda_library(
     name = "profiler_lib",
     visibility = ["//tensorflow:internal"],
-    deps = tf_additional_profiler_lib_deps(),
+    deps = [
+        "//tensorflow/core/profiler/internal/cpu:host_tracer",
+    ] + if_cuda([
+        "//tensorflow/core/profiler/internal/gpu:device_tracer",
+    ]),
     alwayslink = 1,
 )
 
@@ -61,6 +63,14 @@ tf_cuda_library(
     ],
 )
 
+cc_library(
+    name = "profiler_utils",
+    srcs = ["profiler_utils.cc"],
+    hdrs = ["profiler_utils.h"],
+    visibility = ["//visibility:public"],
+    alwayslink = 1,
+)
+
 filegroup(
     name = "mobile_srcs",
     srcs = glob(["*"]),

tensorflow/core/profiler/lib/profiler_session.cc

@@ -26,18 +26,13 @@ limitations under the License.
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/profiler/lib/profiler_utils.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/protobuf/trace_events.pb.h"
 
 namespace tensorflow {
 namespace {
 
-// Track whether there's an active ProfilerSession.
-// Prevents another ProfilerSession from creating ProfilerInterface(s), as they
-// use singletons that do not allow concurrent profiling request (e.g.,
-// DeviceTracer).
-std::atomic<bool> session_active = ATOMIC_VAR_INIT(false);
-
 // Given a node_name in the format "op_name:op_type", returns the "op_type".
 // If the "op_type" is missing, returns the node_name.
 // This is done so all ops with the same type appear in the same color in trace
@@ -173,7 +168,7 @@ Status ProfilerSession::CollectData(RunMetadata* run_metadata) {
 
   if (active_) {
     // Allow another session to start.
-    session_active.store(false);
+    profiler::ReleaseProfilerLock();
     active_ = false;
   }
 
@@ -194,7 +189,7 @@ Status ProfilerSession::SerializeToString(string* content) {
 }
 
 ProfilerSession::ProfilerSession()
-    : active_(!session_active.exchange(true)),
+    : active_(profiler::AcquireProfilerLock()),
       start_time_micros_(Env::Default()->NowNanos() / EnvTime::kMicrosToNanos) {
   if (!active_) {
     status_ = tensorflow::Status(error::UNAVAILABLE,
@@ -223,7 +218,7 @@ ProfilerSession::~ProfilerSession() {
 
   if (active_) {
     // Allow another session to start.
-    session_active.store(false);
+    profiler::ReleaseProfilerLock();
   }
 }
 }  // namespace tensorflow

tensorflow/core/profiler/lib/profiler_utils.cc

@@ -0,0 +1,31 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#include "tensorflow/core/profiler/lib/profiler_utils.h"
+
+#include <atomic>
+
+namespace tensorflow {
+namespace profiler {
+
+// Track whether there's an active profiler session.
+// Prevents another profiler session from creating ProfilerInterface(s).
+std::atomic<bool> session_active = ATOMIC_VAR_INIT(false);
+
+bool AcquireProfilerLock() { return !session_active.exchange(true); }
+
+void ReleaseProfilerLock() { session_active.store(false); }
+
+}  // namespace profiler
+}  // namespace tensorflow

tensorflow/core/profiler/lib/profiler_utils.h

@@ -0,0 +1,31 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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_CORE_PROFILER_LIB_PROFILER_UTILS_H_
+#define TENSORFLOW_CORE_PROFILER_LIB_PROFILER_UTILS_H_
+
+namespace tensorflow {
+namespace profiler {
+
+// If return false, other profiler session is active right now.
+// Otherwise the profiler lock is acquired.
+bool AcquireProfilerLock();
+
+// Release the acquired profiler lock.
+void ReleaseProfilerLock();
+
+}  // namespace profiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_PROFILER_LIB_PROFILER_UTILS_H_