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

@@ -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",
     ],
 )

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,

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