Fix a bug in profiler.

PiperOrigin-RevId: 194838948

3 files changed, 26 insertions, 9 deletions

diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h
index 6f3433abcf..1074f64263 100644
--- a/tensorflow/contrib/lite/interpreter.h
+++ b/tensorflow/contrib/lite/interpreter.h
@@ -325,9 +325,7 @@ class Interpreter {
 
   void SetProfiler(profiling::Profiler* profiler) { profiler_ = profiler; }
 
-  profiling::Profiler* GetProfiler(profiling::Profiler* profiler) {
-    return profiler_;
-  }
+  profiling::Profiler* GetProfiler() { return profiler_; }
 
   // The default capacity of `tensors_` vector.
   static constexpr int kTensorsReservedCapacity = 128;
diff --git a/tensorflow/contrib/lite/profiling/profiler.h b/tensorflow/contrib/lite/profiling/profiler.h
index dfa98a6708..8c3e4dc76d 100644
--- a/tensorflow/contrib/lite/profiling/profiler.h
+++ b/tensorflow/contrib/lite/profiling/profiler.h
@@ -85,7 +85,7 @@ class Profiler {
   std::vector<const ProfileEvent*> GetProfileEvents() {
     std::vector<const ProfileEvent*> profile_events;
     profile_events.reserve(buffer_.Size());
-    for (int i = 0; i < buffer_.Size(); i++) {
+    for (size_t i = 0; i < buffer_.Size(); i++) {
       profile_events.push_back(buffer_.At(i));
     }
     return profile_events;
@@ -103,7 +103,9 @@ class ScopedProfile {
   // Adds a profile event to profile that begins with the construction
   // of object and ends when the object goes out of scope.
   // The lifetime of tag should be at least the lifetime of profiler.
-  ScopedProfile(Profiler* profiler, const char* tag) {
+
+  ScopedProfile(Profiler* profiler, const char* tag)
+      : buffer_(nullptr), event_handle_(0) {
     if (profiler) {
       buffer_ = profiler->GetProfileBuffer();
       event_handle_ =
@@ -126,7 +128,8 @@ class ScopedOperatorProfile {
   // Adds a profile event to profile that begins with the construction
   // of object and ends when the object goes out of scope.
   // The lifetime of tag should be at least the lifetime of profiler.
-  ScopedOperatorProfile(Profiler* profiler, const char* tag, int node_index) {
+  ScopedOperatorProfile(Profiler* profiler, const char* tag, int node_index)
+      : buffer_(nullptr), event_handle_(0) {
     if (profiler) {
       buffer_ = profiler->GetProfileBuffer();
       event_handle_ = buffer_->BeginEvent(
@@ -148,9 +151,11 @@ class ScopedOperatorProfile {
 }  // namespace profiling
 }  // namespace tflite
 
-#define SCOPED_OPERATOR_PROFILE(profiler, node_index)                       \
-  tflite::profiling::ScopedOperatorProfile _profile((profiler), "OpInvoke", \
-                                                    (node_index))
+#define VARNAME_UNIQ(name, ctr) name##ctr
+
+#define SCOPED_OPERATOR_PROFILE(profiler, node_index)    \
+  tflite::profiling::ScopedOperatorProfile VARNAME_UNIQ( \
+      _profile_, __COUNTER__)((profiler), "OpInvoke", (node_index))
 #else
 
 namespace tflite {
diff --git a/tensorflow/contrib/lite/profiling/profiler_test.cc b/tensorflow/contrib/lite/profiling/profiler_test.cc
index 7ea1d8f7d3..0fba0450a0 100644
--- a/tensorflow/contrib/lite/profiling/profiler_test.cc
+++ b/tensorflow/contrib/lite/profiling/profiler_test.cc
@@ -93,6 +93,20 @@ TEST(ProfilingTest, ProfilesAreCollected) {
 #endif
 }
 
+TEST(ProfilingTest, NullProfiler) {
+  Profiler* profiler = nullptr;
+  { SCOPED_OPERATOR_PROFILE(profiler, 1); }
+}
+
+TEST(ProfilingTest, ScopedProfile) {
+  Profiler profiler;
+  profiler.StartProfiling();
+  { SCOPED_OPERATOR_PROFILE(&profiler, 1); }
+  profiler.StopProfiling();
+  auto profile_events = profiler.GetProfileEvents();
+  EXPECT_EQ(1, profile_events.size());
+}
+
 }  // namespace
 }  // namespace profiling
 }  // namespace tflite