From 5676c66e9108dd6bb609b6e032860ba9c2e80078 Mon Sep 17 00:00:00 2001
From: Kibeom Kim <kkb@google.com>
Date: Sun, 9 Aug 2020 09:48:49 -0700
Subject: [PATCH] Add convert_to_tensor KPI and introduce
 @trace.trace_wrapper(...) API

Note that this is a superset of tf.convert_to_tensor, which includes internal convert_to_tensor calls.

Also, introduced @trace.trace_wrapper(...) API that's faster than `with trace.Trace(...):` API.

Benchmark:

`with trace.Trace(...):`
time               : 0.67 us

`@trace.trace_wrapper(...)`
time               : 0.21 us

Direct `if trace.enabled:` inlining
time               : 0.17 us
PiperOrigin-RevId: 325690563
Change-Id: I3251e38b7543e3121be6fad0266706a1b3b5c389
---
 tensorflow/python/BUILD             |  1 +
 tensorflow/python/framework/ops.py  |  2 ++
 tensorflow/python/profiler/trace.py | 42 +++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+)

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 4efe769c59d..a8a70566ab7 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -1702,6 +1702,7 @@ py_library(
         "//tensorflow/python/eager:core",
         "//tensorflow/python/eager:monitoring",
         "//tensorflow/python/eager:tape",
+        "//tensorflow/python/profiler:traceme",
         "@six_archive//:six",
     ],
 )
diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index 75a36f83fc5..f07bca17061 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -62,6 +62,7 @@ from tensorflow.python.framework import versions
 from tensorflow.python.ops import control_flow_util
 from tensorflow.python.platform import app
 from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.profiler import trace
 from tensorflow.python.types import core as core_tf_types
 from tensorflow.python.types import internal
 from tensorflow.python.util import compat
@@ -1472,6 +1473,7 @@ def pack_eager_tensors(tensors, ctx=None):
   return packed_tensor
 
 
+@trace.trace_wrapper("convert_to_tensor")
 def convert_to_tensor(value,
                       dtype=None,
                       name=None,
diff --git a/tensorflow/python/profiler/trace.py b/tensorflow/python/profiler/trace.py
index 0591d90fa43..e4cf581bd25 100644
--- a/tensorflow/python/profiler/trace.py
+++ b/tensorflow/python/profiler/trace.py
@@ -18,6 +18,8 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 from tensorflow.python.profiler.internal import _pywrap_traceme
 from tensorflow.python.util.tf_export import tf_export
 
@@ -123,3 +125,43 @@ class Trace(object):
   def __exit__(self, exc_type, exc_val, exc_tb):
     if self._traceme:
       self._traceme.Stop()
+
+
+def trace_wrapper(trace_name, **trace_kwargs):
+  """Decorator alternative to `with Trace(): ...`.  It's faster.
+
+  Args:
+    trace_name: The name of the trace event.
+    **trace_kwargs: Keyword arguments added to the trace event. Both the key and
+      value are of types that can be converted to strings, which will be
+      interpreted by the profiler according to the traceme name.
+
+  Returns:
+    A decorator that can wrap a function and apply `Trace` scope if needed.
+
+  Example usage:
+    ```python
+
+    @trace_wrapper('trace_name')
+    def func(x, y, z):
+      pass  # code to execute and apply `Trace` if needed.
+
+    # Equivalent to
+    # with Trace('trace_name'):
+    #   func(1, 2, 3)
+    func(1, 2, 3)
+    ```
+  """
+
+  def inner_wrapper(func):
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+      if enabled:
+        with Trace(trace_name, **trace_kwargs):
+          return func(*args, **kwargs)
+      return func(*args, **kwargs)
+
+    return wrapped
+
+  return inner_wrapper