def test_new_symbol_avoids_conflicts(self): namer = naming.Namer({'temp': 1}) # temp is reserved in the global namespace self.assertEqual('temp_1', namer.new_symbol('temp', set())) # temp_2 is reserved in the local namespace self.assertEqual('temp_3', namer.new_symbol('temp', set(('temp_2', )))) self.assertItemsEqual(('temp_1', 'temp_3'), namer.generated_names)
def prepare(self, test_fn, namespace, recursive=True): namespace['ConversionOptions'] = converter.ConversionOptions future_features = ('print_function', 'division') node, source = parser.parse_entity(test_fn, future_features=future_features) namer = naming.Namer(namespace) program_ctx = converter.ProgramContext( options=converter.ConversionOptions(recursive=recursive), autograph_module=None) entity_info = transformer.EntityInfo( name=test_fn.__name__, source_code=source, source_file='<fragment>', future_features=future_features, namespace=namespace) ctx = transformer.Context(entity_info, namer, program_ctx) origin_info.resolve_entity(node, source, test_fn) graphs = cfg.build(node) node = qual_names.resolve(node) node = activity.resolve(node, ctx, None) node = reaching_definitions.resolve(node, ctx, graphs) anno.dup( node, { anno.Static.DEFINITIONS: anno.Static.ORIG_DEFINITIONS, }, ) return node, ctx
def convert_func_to_ast(f, program_ctx, do_rename=True): """Specialization of `convert_entity_to_ast` for callable functions.""" future_features = inspect_utils.getfutureimports(f) node, source = parser.parse_entity(f, future_features=future_features) logging.log(3, 'Source code of %s:\n\n%s\n', f, source) # Parsed AST should contain future imports and one function def node. # In general, the output of inspect.getsource is inexact for lambdas because # it uses regex matching to adjust the exact location around the line number # that CPython records. Then, the entire containing line is returned, which # we may have trouble disambiguating. For example: # x, y = lambda: 1, lambda: 2 if f.__name__ == '<lambda>': nodes = ast_util.find_matching_definitions(node, f) if len(nodes) != 1: raise ValueError( 'Unable to identify source code of lambda function {}. It was' ' defined on this line: {}, which must contain a single lambda with' ' matching signature. To avoid ambiguity, define each lambda' ' in a separate expression.'.format(f, source)) node, = nodes # TODO(znado): Place inside standard_analysis. origin_info.resolve_entity(node, source, f) namespace = inspect_utils.getnamespace(f) _add_self_references(namespace, program_ctx.autograph_module) namer = naming.Namer(namespace) if isinstance(node, gast.Lambda): new_name = namer.new_symbol('tf__lambda', ()) elif do_rename: new_name = namer.new_symbol('tf__' + f.__name__, ()) else: new_name = f.__name__ entity_info = transformer.EntityInfo(source_code=source, source_file='<fragment>', future_features=future_features, namespace=namespace) context = converter.EntityContext(namer, entity_info, program_ctx, new_name) node = node_to_graph(node, context) if isinstance(node, gast.Lambda): node = gast.Assign(targets=[ gast.Name(new_name, ctx=gast.Store(), annotation=None, type_comment=None) ], value=node) elif do_rename: node.name = new_name else: assert node.name == new_name return (node, ), new_name, entity_info
def _transform_function(self, fn, user_context): """Performs source code transformation on a function.""" future_features = inspect_utils.getfutureimports(fn) node, source = parser.parse_entity(fn, future_features=future_features) logging.log(3, 'Source code of %s:\n\n%s\n', fn, source) # In general, the output of inspect.getsource is inexact for lambdas # because it uses regex matching to adjust the exact location around # the line number that CPython records. Then, the entire containing line # is returned, which we may have trouble disambiguating. # For example: # x, y = lambda: 1, lambda: 2 is_lambda = fn.__name__ == '<lambda>' if is_lambda: nodes = ast_util.find_matching_definitions(node, fn) if len(nodes) != 1: raise ValueError( 'Unable to identify source code of lambda function {}.' ' It was defined in this code:\n' '{}\n' 'This code must contain a single distinguishable lambda.' ' To avoid this problem, define each lambda in a separate' ' expression.'.format(fn, source)) node, = nodes origin_info.resolve_entity(node, source, fn) namespace = inspect_utils.getnamespace(fn) namer = naming.Namer(namespace) new_name = namer.new_symbol(self.get_transformed_name(node), ()) entity_info = transformer.EntityInfo(name=new_name, source_code=source, source_file='<fragment>', future_features=future_features, namespace=namespace) context = transformer.Context(entity_info, namer, user_context) node = self._erase_arg_defaults(node) node = self.transform_ast(node, context) if is_lambda: node = gast.Assign(targets=[ gast.Name(new_name, ctx=gast.Store(), annotation=None, type_comment=None) ], value=node) else: node.name = new_name return node, context
def _parse_and_analyze(self, test_fn): # TODO(mdan): Use a custom FunctionTransformer here. node, source = parser.parse_entity(test_fn, future_features=()) entity_info = transformer.EntityInfo(name=test_fn.__name__, source_code=source, source_file=None, future_features=(), namespace={}) node = qual_names.resolve(node) namer = naming.Namer({}) ctx = transformer.Context(entity_info, namer, None) node = activity.resolve(node, ctx) return node, entity_info
def mlir_gen_internal(node, entity_info): """Returns mlir module for unprocessed node `node`.""" namer = naming.Namer({}) graphs = cfg.build(node) ctx = transformer.Context(entity_info, namer, None) node = qual_names.resolve(node) node = activity.resolve(node, ctx) node = reaching_definitions.resolve(node, ctx, graphs) node = reaching_fndefs.resolve(node, ctx, graphs) node = liveness.resolve(node, ctx, graphs) mlir_generator = MLIRGen(ctx) mlir_generator.visit(node) return mlir_generator.prog
def transform_function(self, fn, user_context): """Transforms a function. Subclasses may override this method. The return value is opaque. The method receives the original AST. The result is passed as-is to the output of `transform`. Args: fn: A function or lambda. user_context: An opaque object (may be None) that is forwarded to transform_ast, through the ctx.user_context argument. Returns: Any. By default it returns the output of transform_ast. """ future_features = inspect_utils.getfutureimports(fn) node, source = parser.parse_entity(fn, future_features=future_features) logging.log(3, 'Source code of %s:\n\n%s\n', fn, source) origin_info.resolve_entity(node, source, fn) namespace = inspect_utils.getnamespace(fn) namer = naming.Namer(namespace) new_name = namer.new_symbol(self.get_transformed_name(node), ()) entity_info = transformer.EntityInfo(name=new_name, source_code=source, source_file='<fragment>', future_features=future_features, namespace=namespace) context = transformer.Context(entity_info, namer, user_context) node = self._erase_arg_defaults(node) node = self.transform_ast(node, context) if isinstance(node, gast.Lambda): node = gast.Assign(targets=[ gast.Name(new_name, ctx=gast.Store(), annotation=None, type_comment=None) ], value=node) else: node.name = new_name return node, context
def prepare(self, test_fn, namespace, recursive=True): namespace['ConversionOptions'] = converter.ConversionOptions future_features = ('print_function', 'division') node, source = parser.parse_entity(test_fn, future_features=future_features) namer = naming.Namer(namespace) program_ctx = converter.ProgramContext( options=converter.ConversionOptions(recursive=recursive), autograph_module=None) entity_info = transformer.EntityInfo(source_code=source, source_file='<fragment>', future_features=future_features, namespace=namespace) ctx = converter.EntityContext(namer, entity_info, program_ctx, 'test_fn') origin_info.resolve_entity(node, source, test_fn) node = converter.standard_analysis(node, ctx, is_initial=True) return node, ctx
def _convert_with_cache(entity, program_ctx, free_nonglobal_var_names): """Returns a (possibly cached) factory for the converted result of entity.""" # The cache subkey encompasses any conversion options on which the generated # code may depend. # The cached factory includes the necessary definitions to distinguish # between the global and non-global free variables. For this reason, the # cache subkey includes the names of the free non-globals. subkey = (program_ctx.options, frozenset(free_nonglobal_var_names)) with _CACHE_LOCK: # The cache values are _ConvertedEntityFactoryInfo objects. if _CACHE.has(entity, subkey): # TODO(mdan): Check whether the module is still loaded. converted_entity_info = _CACHE[entity][subkey] logging.log(3, 'Cache hit for entity %s subkey %s: %s', entity, subkey, converted_entity_info) return converted_entity_info logging.log(1, 'Entity %s is not cached for subkey %s', entity, subkey) nodes, converted_name, entity_info = convert_entity_to_ast( entity, program_ctx) namer = naming.Namer(entity_info.namespace) factory_factory_name = namer.new_symbol( 'create_converted_entity_factory', ()) factory_name = namer.new_symbol('create_converted_entity', ()) nodes = _wrap_into_dynamic_factory(nodes, converted_name, factory_factory_name, factory_name, free_nonglobal_var_names, entity_info.future_features) module, _, source_map = loader.load_ast(nodes, include_source_map=True) module_name = module.__name__ converted_entity_info = _ConvertedEntityFactoryInfo( module_name=module_name, converted_name=converted_name, factory_factory_name=factory_factory_name, source_map=source_map) _CACHE[entity][subkey] = converted_entity_info return converted_entity_info
def transform_function(self, fn, user_context): """Transforms a function. Subclasses may override this method. The return value is opaque. The method receives the original AST. The result is passed as-is to the output of `transform`. Args: fn: A function or lambda. user_context: An opaque object (may be None) that is forwarded to transform_ast, through the ctx.user_context argument. Returns: Tuple[Any, Any]. By default it returns the output of transform_ast, together with a `transformer.Context` containing information about the transformation process. """ future_features = inspect_utils.getfutureimports(fn) node, source = parser.parse_entity(fn, future_features=future_features) logging.log(3, 'Source code of %s:\n\n%s\n', fn, source) origin_info.resolve_entity(node, source, fn) namespace = inspect_utils.getnamespace(fn) namer = naming.Namer(namespace) new_name = namer.new_symbol(self.get_transformed_name(node), ()) entity_info = transformer.EntityInfo( name=new_name, source_code=source, source_file='<fragment>', future_features=future_features, namespace=namespace) context = transformer.Context(entity_info, namer, user_context) node = self._erase_arg_defaults(node) result = self.transform_ast(node, context) return result, context
def _transform_function(self, fn, user_context): """Performs source code transformation on a function.""" future_features = inspect_utils.getfutureimports(fn) node, source = parser.parse_entity(fn, future_features=future_features) logging.log(3, 'Source code of %s:\n\n%s\n', fn, source) origin_info.resolve_entity(node, source, fn) namespace = inspect_utils.getnamespace(fn) namer = naming.Namer(namespace) new_name = namer.new_symbol(self.get_transformed_name(node), ()) entity_info = transformer.EntityInfo( name=new_name, source_code=source, source_file='<fragment>', future_features=future_features, namespace=namespace) context = transformer.Context(entity_info, namer, user_context) node = self._erase_arg_defaults(node) node = self.transform_ast(node, context) if isinstance(node, gast.Lambda): node = gast.Assign( targets=[ gast.Name( new_name, ctx=gast.Store(), annotation=None, type_comment=None) ], value=node) else: node.name = new_name return node, context
def test_new_symbol_avoids_duplicates(self): namer = naming.Namer({}) self.assertEqual('temp', namer.new_symbol('temp', set())) self.assertEqual('temp_1', namer.new_symbol('temp', set())) self.assertItemsEqual(('temp', 'temp_1'), namer.generated_names)
def test_new_symbol_tracks_names(self): namer = naming.Namer({}) self.assertEqual('temp', namer.new_symbol('temp', set())) self.assertItemsEqual(('temp', ), namer.generated_names)
def _parse_and_analyze(f, autobatch_functions): """Performs preliminary analyses and transformations. The goal is to massage the source program into a form on which the `_AutoBatchingTransformer` below will be successful. Args: f: Function to analyze autobatch_functions: List of Python `str` names of autobatched functions. Arguments to these functions will be canonicalized to variable references, but others will not. Returns: node: A Python AST node representing the function, suitable for passing to `_AutoBatchingTransformer.visit` entity_info: An AutoGraph `EntityInfo` object, with some information about `f`. Required for initializing `_AutoBatchingTransformer`. """ namespace = {} # Get the AST of the function future_features = inspect_utils.getfutureimports(f) node, _ = parser.parse_entity(f, future_features=future_features) # Boilerplate for AutoGraph transforms entity_info = transformer.EntityInfo(source_code='', source_file=None, future_features=future_features, namespace=namespace) program_ctx = converter.ProgramContext( options=converter.ConversionOptions(recursive=True), autograph_module=None) ctx = converter.EntityContext(namer=naming.Namer(namespace), entity_info=entity_info, program_ctx=program_ctx) # Canonicalize away break statements node = converter.standard_analysis(node, ctx, is_initial=True) node = break_statements.transform(node, ctx) # Canonicalize away continue statements node = converter.standard_analysis(node, ctx, is_initial=False) node = continue_statements.transform(node, ctx) # Force single returns node = converter.standard_analysis(node, ctx, is_initial=False) node = return_statements.transform(node, ctx, default_to_null_return=False) # Transform into ANF # Replacing if tests and autobatched function call arguments because # that's where divergence can happen. # Replacing all function calls because the downstream transformation # expects calls to lead directly to assignments. def maybe_replace_function_argument(parent, field_name, child): del field_name, child if not anno.hasanno(parent.func, anno.Basic.QN): return False func_name = anno.getanno(parent.func, anno.Basic.QN) if str(func_name) in autobatch_functions: return True return False anf_config = [ (anf.ASTEdgePattern(gast.If, 'test', anf.ANY), anf.REPLACE), (anf.ASTEdgePattern(anf.ANY, anf.ANY, gast.Call), anf.REPLACE), (anf.ASTEdgePattern(gast.Call, 'args', anf.ANY), maybe_replace_function_argument), ] node = anf.transform(node, ctx, config=anf_config) node = converter.standard_analysis(node, ctx, is_initial=False) return node, ctx