Add vmodule support to TF OSS logging.

Setting the environment variable TF_CPP_VMODULE=foo=7,bar=7 will cause the
"foo" and "bar" translation units (foo.{h.cc},bar.{h.cc}) to have a VLOG level
of 7, meaning VLOG(0) to VLOG(7) in those translation units will be logged to
stderr (as LOG(ERROR) does).

Previous:

Benchmark         Time(ns) Iterations
-------------------------------------
BM_DisabledVlog          3  215821571

Lambda:

Benchmark         Time(ns) Iterations
-------------------------------------
BM_DisabledVlog          3  247013779

So change for disabled VLOGs seems in the noise
in the microbenchmark.

PiperOrigin-RevId: 163995038

5 files changed, 247 insertions, 10 deletions

diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index c2e58b456b..bbae186f57 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -2121,6 +2121,8 @@ tf_cc_tests(
         "platform/port_test.cc",
         "platform/profile_utils/cpu_utils_test.cc",
         "platform/subprocess_test.cc",
+        "platform/vmodule_benchmark_test.cc",
+        "platform/vmodule_test.cc",
     ],
     deps = [
         ":lib",
diff --git a/tensorflow/core/platform/default/logging.cc b/tensorflow/core/platform/default/logging.cc
index ebdd4b624a..ac0988e704 100644
--- a/tensorflow/core/platform/default/logging.cc
+++ b/tensorflow/core/platform/default/logging.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/platform/default/logging.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/platform/env_time.h"
 #include "tensorflow/core/platform/macros.h"
 
@@ -24,8 +25,12 @@ limitations under the License.
 #endif
 
 #include <stdlib.h>
+#include <string.h>
 #include <time.h>
 
+#include <string>
+#include <unordered_map>
+
 namespace tensorflow {
 namespace internal {
 
@@ -83,11 +88,11 @@ void LogMessage::GenerateLogMessage() {
   const size_t time_buffer_size = 30;
   char time_buffer[time_buffer_size];
   strftime(time_buffer, time_buffer_size, "%Y-%m-%d %H:%M:%S",
-	   localtime(&now_seconds));
+           localtime(&now_seconds));
 
   // TODO(jeff,sanjay): Replace this with something that logs through the env.
   fprintf(stderr, "%s.%06d: %c %s:%d] %s\n", time_buffer, micros_remainder,
-	  "IWEF"[severity_], fname_, line_, str().c_str());
+          "IWEF"[severity_], fname_, line_, str().c_str());
 }
 #endif
 
@@ -124,6 +129,48 @@ int64 MinVLogLevelFromEnv() {
   return LogLevelStrToInt(tf_env_var_val);
 }
 
+using VmoduleMap = std::unordered_map<StringPiece, int, StringPiece::Hasher>;
+
+// Returns a mapping from module name to VLOG level, derived from the
+// TF_CPP_VMOUDLE environment variable; ownership is transferred to the caller.
+VmoduleMap* VmoduleRecordsFromEnv() {
+  // The value of the env var is supposed to be of the form:
+  //    "foo=1,bar=2,baz=3"
+  const char* tf_env_var_val = getenv("TF_CPP_VMODULE");
+  auto* result = new VmoduleMap();
+  if (tf_env_var_val == nullptr) return result;
+  while (true) {
+    const char* eq = strchr(tf_env_var_val, '=');
+    if (eq == nullptr) break;
+    const char* after_eq = eq + 1;
+
+    // Comma either points at the next comma delimiter, or at a null terminator.
+    // We check that the integer we parse ends at this delimiter.
+    const char* comma = strchr(after_eq, ',');
+    const char* new_tf_env_var_val;
+    if (comma == nullptr) {
+      comma = strchr(after_eq, '\0');
+      new_tf_env_var_val = comma;
+    } else {
+      new_tf_env_var_val = comma + 1;
+    }
+
+    char* endptr = nullptr;
+    int level = strtol(after_eq, &endptr, 10);
+    if (endptr != comma) {
+      fprintf(stderr,
+              "warning: could not parse integer in vmodule specification in "
+              "\"%s\".\n",
+              after_eq);
+      break;
+    }
+    StringPiece module(tf_env_var_val, eq - tf_env_var_val);
+    tf_env_var_val = new_tf_env_var_val;
+    (*result)[module] = level;
+  }
+  return result;
+}
+
 }  // namespace
 
 LogMessage::~LogMessage() {
@@ -137,6 +184,19 @@ int64 LogMessage::MinVLogLevel() {
   return min_vlog_level;
 }
 
+bool LogMessage::VmoduleActivated(const char* fname, int lvl) {
+  static VmoduleMap* vmodule_records = VmoduleRecordsFromEnv();
+  const char* last_slash = strrchr(fname, '/');
+  const char* module_start = last_slash == nullptr ? fname : last_slash + 1;
+  const char* dot_after = strchr(module_start, '.');
+  const char* module_limit =
+      dot_after == nullptr ? strchr(fname, '\0') : dot_after;
+  StringPiece module(module_start, module_limit - module_start);
+  auto it = vmodule_records->find(module);
+  if (it == vmodule_records->end()) return false;
+  return it->second >= lvl;
+}
+
 LogMessageFatal::LogMessageFatal(const char* file, int line)
     : LogMessage(file, line, FATAL) {}
 LogMessageFatal::~LogMessageFatal() {
diff --git a/tensorflow/core/platform/default/logging.h b/tensorflow/core/platform/default/logging.h
index 04ff9e12b6..c8c9b2da11 100644
--- a/tensorflow/core/platform/default/logging.h
+++ b/tensorflow/core/platform/default/logging.h
@@ -46,6 +46,16 @@ class LogMessage : public std::basic_ostringstream<char> {
   // but VLOG(3) will not. Defaults to 0.
   static int64 MinVLogLevel();
 
+  // Returns whether VLOG level lvl is activated for the file fname.
+  //
+  // E.g. if the environment variable TF_CPP_VMODULE contains foo=3 and fname is
+  // foo.cc and lvl is <= 3, this will return true.
+  //
+  // It is expected that the result of this query will be cached in the VLOG-ing
+  // call site to avoid repeated lookups. This routine performs a hash-map
+  // access against the VLOG-ing specification provided by the env var.
+  static bool VmoduleActivated(const char* fname, int lvl);
+
  protected:
   void GenerateLogMessage();
 
@@ -76,18 +86,38 @@ class LogMessageFatal : public LogMessage {
 
 #define LOG(severity) _TF_LOG_##severity
 
-#ifdef IS_MOBILE_PLATFORM
+#if defined(IS_MOBILE_PLATFORM)
+
 // Turn VLOG off when under mobile devices for considerations of binary size.
-#define VLOG_IS_ON(lvl) ((lvl) <= 0)
-#else
-// Otherwise, Set TF_CPP_MIN_VLOG_LEVEL environment to update minimum log level
-// of VLOG
-#define VLOG_IS_ON(lvl) \
+#define _VLOG_IS_ON(lvl, file) ((lvl) <= 0)
+
+#elif defined(PLATFORM_WINDOWS)
+
+// TODO(b/64279502) The _VLOG_IS_ON definition below appears to cause MSVC to
+// fatal error, so we fall back to the vmodule-less implementation for now.
+#define _VLOG_IS_ON(lvl, file) \
   ((lvl) <= ::tensorflow::internal::LogMessage::MinVLogLevel())
+
+#else
+
+// Otherwise, set TF_CPP_MIN_VLOG_LEVEL environment to update minimum log level
+// of VLOG, or TF_CPP_VMODULE to set the minimum log level for individual
+// translation units.
+#define _VLOG_IS_ON(lvl, file)                                              \
+  (([](int level, const char* fname) {                                      \
+    if (level <= ::tensorflow::internal::LogMessage::MinVLogLevel())        \
+      return true;                                                          \
+    static bool vmodule_activated =                                         \
+        ::tensorflow::internal::LogMessage::VmoduleActivated(fname, level); \
+    return vmodule_activated;                                               \
+  })(lvl, file))
+
 #endif
 
-#define VLOG(lvl)      \
-  if (TF_PREDICT_FALSE(VLOG_IS_ON(lvl))) \
+#define VLOG_IS_ON(lvl) _VLOG_IS_ON(lvl, __FILE__)
+
+#define VLOG(lvl)                                   \
+  if (TF_PREDICT_FALSE(_VLOG_IS_ON(lvl, __FILE__))) \
   ::tensorflow::internal::LogMessage(__FILE__, __LINE__, tensorflow::INFO)
 
 // CHECK dies with a fatal error if condition is not true.  It is *not*
diff --git a/tensorflow/core/platform/vmodule_benchmark_test.cc b/tensorflow/core/platform/vmodule_benchmark_test.cc
new file mode 100644
index 0000000000..0f9e75bf9c
--- /dev/null
+++ b/tensorflow/core/platform/vmodule_benchmark_test.cc
@@ -0,0 +1,28 @@
+/* Copyright 2017 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/platform/logging.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace tensorflow {
+
+static void BM_DisabledVlog(int iters) {
+  for (int i = 0; i < iters; ++i) {
+    VLOG(1) << "Testing VLOG(1)!";
+  }
+}
+BENCHMARK(BM_DisabledVlog);
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/platform/vmodule_test.cc b/tensorflow/core/platform/vmodule_test.cc
new file mode 100644
index 0000000000..47b4b2e0e7
--- /dev/null
+++ b/tensorflow/core/platform/vmodule_test.cc
@@ -0,0 +1,117 @@
+/* Copyright 2017 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.
+==============================================================================*/
+
+// Test that popens a child process with the VLOG-ing environment variable set
+// for the logging framework, and observes VLOG_IS_ON and VLOG macro output.
+
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/platform.h"
+#include "tensorflow/core/platform/test.h"
+
+#include <string.h>
+
+namespace tensorflow {
+namespace {
+
+int RealMain(const char* argv0, bool do_vlog) {
+  if (do_vlog) {
+#if !defined(PLATFORM_GOOGLE)
+    // Note, we only test this when !defined(PLATFORM_GOOGLE) because
+    // VmoduleActivated doesn't exist in that implementation.
+    //
+    // Also, we call this internal API to simulate what would happen if
+    // differently-named translation units attempted to VLOG, so we don't need
+    // to create dummy translation unit files.
+    bool ok = internal::LogMessage::VmoduleActivated("vmodule_test.cc", 7) &&
+              internal::LogMessage::VmoduleActivated("shoobadooba.h", 3);
+    if (!ok) {
+      fprintf(stderr, "vmodule activated levels not as expected.\n");
+      return EXIT_FAILURE;
+    }
+#endif
+
+    // Print info on which VLOG levels are activated.
+    fprintf(stderr, "VLOG_IS_ON(8)? %d\n", VLOG_IS_ON(8));
+    fprintf(stderr, "VLOG_IS_ON(7)? %d\n", VLOG_IS_ON(7));
+    fprintf(stderr, "VLOG_IS_ON(6)? %d\n", VLOG_IS_ON(6));
+    // Do some VLOG-ing.
+    VLOG(8) << "VLOG(8)";
+    VLOG(7) << "VLOG(7)";
+    VLOG(6) << "VLOG(6)";
+    LOG(INFO) << "INFO";
+    return EXIT_SUCCESS;
+  }
+
+  // Popen the child process.
+  std::string command = std::string(argv0);
+#if defined(PLATFORM_GOOGLE)
+  command = command + " do_vlog --vmodule=vmodule_test=7 --alsologtostderr";
+#else
+  command =
+      "TF_CPP_VMODULE=vmodule_test=7,shoobadooba=3 " + command + " do_vlog";
+#endif
+  command += " 2>&1";
+  fprintf(stderr, "Running: \"%s\"\n", command.c_str());
+  FILE* f = popen(command.c_str(), "r");
+  if (f == nullptr) {
+    fprintf(stderr, "Failed to popen child: %s\n", strerror(errno));
+    return EXIT_FAILURE;
+  }
+
+  // Read data from the child's stdout.
+  constexpr int kBufferSizeBytes = 4096;
+  char buffer[kBufferSizeBytes];
+  size_t result = fread(buffer, sizeof(buffer[0]), kBufferSizeBytes - 1, f);
+  if (result == 0) {
+    fprintf(stderr, "Failed to read from child stdout: %zu %s\n", result,
+            strerror(errno));
+    return EXIT_FAILURE;
+  }
+  buffer[result] = '\0';
+  int status = pclose(f);
+  if (status == -1) {
+    fprintf(stderr, "Failed to close popen child: %s\n", strerror(errno));
+    return EXIT_FAILURE;
+  }
+
+  // Check output is as expected.
+  const char kExpected[] =
+      "VLOG_IS_ON(8)? 0\nVLOG_IS_ON(7)? 1\nVLOG_IS_ON(6)? 1\n";
+  if (strstr(buffer, kExpected) == nullptr) {
+    fprintf(stderr, "error: unexpected output from child: \"%.*s\"\n",
+            kBufferSizeBytes, buffer);
+    return EXIT_FAILURE;
+  }
+  bool ok = strstr(buffer, "VLOG(7)\n") != nullptr &&
+            strstr(buffer, "VLOG(6)\n") != nullptr &&
+            strstr(buffer, "VLOG(8)\n") == nullptr;
+  if (!ok) {
+    fprintf(stderr, "error: VLOG output not as expected: \"%.*s\"\n",
+            kBufferSizeBytes, buffer);
+    return EXIT_FAILURE;
+  }
+
+  // Success!
+  return EXIT_SUCCESS;
+}
+
+}  // namespace
+}  // namespace tensorflow
+
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  bool do_vlog = argc >= 2 && strcmp(argv[1], "do_vlog") == 0;
+  return tensorflow::RealMain(argv[0], do_vlog);
+}