From 827c100592cf1d79554ede81c22b8c12a7a0faaa Mon Sep 17 00:00:00 2001 From: Dan Moldovan Date: Thu, 18 Jul 2019 14:19:28 -0700 Subject: [PATCH] Cache information about unconverted functions for faster subsequent execution. This resulted in about 3x speedup in a basic benchmark operating on NumPy code, although it remained 3-4 slower compared to pure Python code. It will also avoid log spam when conversion fails in most cases. PiperOrigin-RevId: 258843781 --- tensorflow/python/autograph/impl/api.py | 30 +++++++++++-------- .../python/autograph/impl/conversion.py | 17 +++++++++++ 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/tensorflow/python/autograph/impl/api.py b/tensorflow/python/autograph/impl/api.py index 6ad11f10f11..d850937423c 100644 --- a/tensorflow/python/autograph/impl/api.py +++ b/tensorflow/python/autograph/impl/api.py @@ -340,8 +340,11 @@ def _attach_metadata(e, f, converted): e.ag_error_metadata = _ErrorMetadata(cause_tb, metadata, message, source_map) -def _call_unconverted(f, args, kwargs): +def _call_unconverted(f, args, kwargs, options, update_cache=True): """Calls the original function without converting with AutoGraph.""" + if update_cache: + conversion.cache_unconverted(f, options) + if inspect_utils.istfmethodtarget(f): return f.__self__.call(args, kwargs) @@ -387,6 +390,9 @@ def converted_call(f, options, args, kwargs): logging.log(1, 'Converted call: %s\n args: %s\n kwargs: %s\n', f, args, kwargs) + if conversion.check_cached_unconverted(f, options): + return _call_unconverted(f, args, kwargs, options, False) + if inspect_utils.isbuiltin(f): if f is eval: return py_builtins.eval_in_original_context(f, args, 1) @@ -398,7 +404,7 @@ def converted_call(f, options, args, kwargs): # TODO(mdan): Clean up the naming inconsistency. if hasattr(f, 'autograph_info__') or hasattr(f, '__ag_compiled'): logging.log(2, 'Permanently whitelisted: %s: already converted', f) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) # TODO(b/122265385): Remove this bypass. if (_is_known_loaded_type(f, 'wrapt', 'FunctionWrapper') or @@ -408,42 +414,42 @@ def converted_call(f, options, args, kwargs): ' by AutoGraph. The function will be called without transformation.' ' You may however apply AutoGraph before the decorator.'.format(f)) logging.log(2, 'Permanently whitelisted: %s: wrapt decorated', f) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) if _is_known_loaded_type(f, 'functools', '_lru_cache_wrapper'): logging.log(2, 'Permanently whitelisted: %s: lru_cache', f) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) # Constructors are permanently whitelisted. # TODO(mdan): Toggle as experimental feature instead. # TODO(b/124016764): Remove this limitation. if tf_inspect.isclass(f): logging.log(2, 'Permanently whitelisted: %s: constructor', f) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) # Other built-in modules are permanently whitelisted. # TODO(mdan): Figure out how to do this consistently for all stdlib modules. if any( f in m.__dict__.values() for m in (collections, pdb, copy, inspect, re)): logging.log(2, 'Permanently whitelisted: %s: part of builtin module', f) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) # Custom ops and kernels are also permanently whitelisted. # See tensorflow.framework.load_library. if (hasattr(f, '__module__') and hasattr(f.__module__, '_IS_TENSORFLOW_PLUGIN')): logging.log(2, 'Permanently whitelisted: %s: TensorFlow plugin', f) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) if not options.force_conversion and conversion.is_whitelisted_for_graph(f): - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) # internal_convert_user_code is for example turned off when issuing a dynamic # call conversion from generated code while in nonrecursive mode. In that # case we evidently don't want to recurse, but we still have to convert # things like builtins. if not options.internal_convert_user_code: - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) # TODO(mdan): Move this entire block inside to_graph. try: # Begin of transformation error guards @@ -493,13 +499,13 @@ def converted_call(f, options, args, kwargs): if not hasattr(target_entity, '__code__'): logging.log(2, 'Permanently whitelisted: %s: native binding', target_entity) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) elif (hasattr(target_entity.__code__, 'co_filename') and target_entity.__code__.co_filename == ''): # TODO(mdan): __globals__['txt'] might work in Py3. logging.log(2, 'Permanently whitelisted: %s: dynamic code (exec?)', target_entity) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) converted_f = to_graph( target_entity, @@ -532,7 +538,7 @@ def converted_call(f, options, args, kwargs): ' Please report this to the AutoGraph team. When filing the bug, set' ' the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and' ' attach the full output. Cause: %s', target_entity, e) - return _call_unconverted(f, args, kwargs) + return _call_unconverted(f, args, kwargs, options) with StackTraceMapper(converted_f), tf_stack.CurrentModuleFilter(): try: diff --git a/tensorflow/python/autograph/impl/conversion.py b/tensorflow/python/autograph/impl/conversion.py index 5c8cf8c362a..b97c7e5d066 100644 --- a/tensorflow/python/autograph/impl/conversion.py +++ b/tensorflow/python/autograph/impl/conversion.py @@ -129,6 +129,7 @@ _CACHE_LOCK = threading.RLock() _CACHE = _ConversionCache() +_UNCONVERTED_CACHE = _ConversionCache() # Note: strictly speaking, a simple factory might have been sufficient for @@ -414,6 +415,22 @@ def is_whitelisted_for_graph(o, check_call_override=True): return False +def check_cached_unconverted(entity, options): + try: + # Catch-all for entities that are unhashable or don't allow weakrefs. + return _UNCONVERTED_CACHE.has(entity, options) + except TypeError: + return False + + +def cache_unconverted(entity, options): + try: + # Catch-all for entities that are unhashable or don't allow weakrefs. + _UNCONVERTED_CACHE[entity][options] = True + except TypeError: + pass + + # TODO(mdan): Rename to convert_*_node to avoid confusion with convert. def convert_entity_to_ast(o, program_ctx): """Compile a Python entity into equivalent TensorFlow. -- GitLab