Reducing op overhead for eager.

Benchmarks:
(new) tfe_py_fastpath_execute_matmul: 7.86213080088 walltime
(old) gen_math_ops_matmul: 11.2566947937 walltime

(The slowdown is due to adding the record_gradient callback)

This will be a 3-step process:

1. (This CL) Add a function that allows execution of an op on the fastpath.
2. Update python generation code to create 2 new python functions
(_op/_op_gradient_callback and _op_slowpath_fallback):
   -- first function (_op) checks if it is in graph mode, if so, do the normal
      thing, else call out to the function added in step 1.
   -- second function does the else part similar to today (calling out to
   args_to_matching_eager etc.)
3. Rename the first function generated above to be the canonical _op function.

PiperOrigin-RevId: 182791741

1 files changed, 7 insertions, 2 deletions

diff --git a/tensorflow/python/pywrap_tfe.i b/tensorflow/python/pywrap_tfe.i
index 083931aa83..de0e5856ff 100644
--- a/tensorflow/python/pywrap_tfe.i
+++ b/tensorflow/python/pywrap_tfe.i
@@ -28,6 +28,7 @@ limitations under the License.
 %rename("%s") TFE_Py_InitEagerTensor;
 %rename("%s") TFE_Py_RegisterExceptionClass;
 %rename("%s") TFE_Py_Execute;
+%rename("%s") TFE_Py_FastPathExecute;
 %rename("%s") TFE_Py_UID;
 %rename("%s") TFE_Py_TapeSetNew;
 %rename("%s") TFE_Py_TapeSetRemove;
@@ -155,7 +156,7 @@ limitations under the License.
   }
   $1 = &temp;
   $1->resize(PyInt_AsLong($input), nullptr);
-} 
+}
 
 // Create new Status object.
 %typemap(in, numinputs=0) TF_Status *out_status {
@@ -180,10 +181,14 @@ limitations under the License.
   }
 }
 
+// SWIG usually unwraps the tuple that the native Python/C interface generates.
+// Since we wanted to have a function with a variable length of arguments, we
+// used the native Python/C interface directly (which by default supports
+// passing all arguments as a tuple).
+%native(TFE_Py_FastPathExecute) TFE_Py_FastPathExecute_C;
 
 %include "tensorflow/python/eager/pywrap_tfe.h"
 
-
 // Clear all typemaps.
 %typemap(out) TF_DataType;
 %typemap(out) int64_t;