def remove_called_lambdas_and_blocks(comp):
"""Removes any called lambdas and blocks from `comp`.
This function first resolves any higher-order functions, so that replacing
called lambdas with blocks and then inlining the block locals cannot result
in more called lambdas. It then performs this sequence of transformations,
taking care to inline selections from tuples before inlining the rest of
the block locals to prevent possible combinatorial growth of the generated
AST.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` from which we
want to remove called lambdas and blocks.
Returns:
A transformed version of `comp` which has no called lambdas or blocks, and
no extraneous selections from tuples.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
comp, names_uniquified = tree_transformations.uniquify_reference_names(comp)
comp, fns_resolved = tree_transformations.resolve_higher_order_functions(comp)
comp, lambdas_replaced = tree_transformations.replace_called_lambda_with_block(
comp)
if fns_resolved or lambdas_replaced:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, sels_removed = tree_transformations.inline_selections_from_tuple(comp)
if sels_removed:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, locals_inlined = tree_transformations.inline_block_locals(comp)
modified = names_uniquified or fns_resolved or lambdas_replaced or sels_removed or locals_inlined
return comp, modified
def parse_tff_to_tf(comp):
comp, _ = tree_transformations.insert_called_tf_identity_at_leaves(comp)
parser_callable = tree_to_cc_transformations.TFParser()
comp, _ = tree_transformations.replace_called_lambda_with_block(comp)
comp, _ = tree_transformations.inline_block_locals(comp)
comp, _ = tree_transformations.replace_selection_from_tuple_with_element(comp)
new_comp, transformed = transformation_utils.transform_postorder(
comp, parser_callable)
return new_comp, transformed
def remove_called_lambdas_and_blocks(comp):
"""Removes any called lambdas and blocks from `comp`.
This function first resolves any higher-order functions, so that replacing
called lambdas with blocks and then inlining the block locals cannot result
in more called lambdas. It then performs this sequence of transformations,
taking care to inline selections from tuples at appropriate stages to prevent
possible combinatorial growth of the generated AST.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` from which we
want to remove called lambdas and blocks.
Returns:
A transformed version of `comp` which has no called lambdas or blocks, and
no extraneous selections from tuples.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
comp, names_uniquified = tree_transformations.uniquify_reference_names(
comp)
# TODO(b/162888191): Remove this gating when `resolve_higher_order_functions`
# is more efficient, or avoided.
if _contains_higher_order_fns(comp):
# `resolve_higher_order_functions` can be expensive, so we only call into it
# when necessary.
comp, fns_resolved = tree_transformations.resolve_higher_order_functions(
comp)
else:
# We must still inline any functional references. We first inline selections
# from tuples to prevent the AST from becoming unnecessarly large.
comp, sels_removed = tree_transformations.inline_selections_from_tuple(
comp)
if sels_removed:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, locals_inlined = tree_transformations.inline_block_locals(comp)
fns_resolved = sels_removed or locals_inlined
comp, lambdas_replaced = tree_transformations.replace_called_lambda_with_block(
comp)
if fns_resolved or lambdas_replaced:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, sels_removed = tree_transformations.inline_selections_from_tuple(
comp)
if sels_removed:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, locals_inlined = tree_transformations.inline_block_locals(comp)
if locals_inlined:
# Inlining local symbols may reintroduce selection-from-tuple pattern,
# combinatorially increasing build times in the worst case. We ensure
# here that remove_called_lambdas_and_blocks respects the postcondition that
# selections from tuples are always collapsed.
comp, _ = tree_transformations.inline_selections_from_tuple(comp)
comp, _ = tree_transformations.uniquify_reference_names(comp)
modified = names_uniquified or fns_resolved or lambdas_replaced or sels_removed or locals_inlined
return comp, modified
def prepare_for_rebinding(comp):
"""Prepares `comp` for extracting rebound variables.
Currently, this means replacing all called lambdas and inlining all blocks.
This does not necessarly guarantee that the resulting computation has no
called lambdas, it merely reduces a level of indirection here. This reduction
has proved sufficient for identifying variables which are about to be rebound
in the top-level lambda, necessarily when compiler components factor work out
from a single function into multiple functions. Since this function makes no
guarantees about sufficiency, it is the responsibility of the caller to
ensure that no unbound variables are introduced during the rebinding.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` from which all
occurrences of a given variable need to be extracted and rebound.
Returns:
Another instance of `building_blocks.ComputationBuildingBlock` which has
had all called lambdas replaced by blocks, all blocks inlined and all
selections from tuples collapsed.
"""
# TODO(b/146430051): Follow up here and consider removing or enforcing more
# strict output invariants when `remove_called_lambdas_and_blocks` is moved
# in here.
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, _ = tree_transformations.replace_called_lambda_with_block(comp)
block_inliner = tree_transformations.InlineBlock(comp)
selection_replacer = tree_transformations.ReplaceSelectionFromTuple()
transforms = [block_inliner, selection_replacer]
symbol_tree = transformation_utils.SymbolTree(
transformation_utils.ReferenceCounter)
def _transform_fn(comp, symbol_tree):
"""Transform function chaining inlining and collapsing selections."""
modified = False
for transform in transforms:
if transform.global_transform:
comp, transform_modified = transform.transform(
comp, symbol_tree)
else:
comp, transform_modified = transform.transform(comp)
modified = modified or transform_modified
return comp, modified
return transformation_utils.transform_postorder_with_symbol_bindings(
comp, _transform_fn, symbol_tree)
def _remove_functional_symbol_bindings(comp):
"""Removes symbol bindings which contain functional types."""
comp, refs_renamed = tree_transformations.uniquify_reference_names(
comp)
comp, lambdas_replaced = tree_transformations.replace_called_lambda_with_block(
comp)
comp, selections_inlined = tree_transformations.inline_selections_from_tuple(
comp)
if selections_inlined:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, functions_inlined = _inline_functions(comp)
comp, locals_removed = tree_transformations.remove_unused_block_locals(
comp)
modified = (refs_renamed or lambdas_replaced or selections_inlined
or functions_inlined or locals_removed)
return comp, modified
def _force_align_intrinsics_to_top_level_lambda(comp, uri):
"""Forcefully aligns `comp` by the intrinsics for the given `uri`.
This function transforms `comp` by extracting, grouping, and potentially
merging all the intrinsics for the given `uri`. The result of this
transformation should contain exactly one instance of the intrinsic for the
given `uri` that is bound only by the `parameter_name` of `comp`.
Args:
comp: The `building_blocks.Lambda` to align.
uri: A Python `list` of URI of intrinsics.
Returns:
A new computation with the transformation applied or the original `comp`.
"""
py_typecheck.check_type(comp, building_blocks.Lambda)
py_typecheck.check_type(uri, list)
for x in uri:
py_typecheck.check_type(x, str)
comp, _ = tree_transformations.uniquify_reference_names(comp)
if not _can_extract_intrinsics_to_top_level_lambda(comp, uri):
comp, _ = tree_transformations.replace_called_lambda_with_block(comp)
comp = _inline_block_variables_required_to_align_intrinsics(comp, uri)
comp, modified = _extract_intrinsics_to_top_level_lambda(comp, uri)
if modified:
if len(uri) > 1:
comp, _ = _group_by_intrinsics_in_top_level_lambda(comp)
modified = False
for intrinsic_uri in uri:
comp, transform_modified = transformations.dedupe_and_merge_tuple_intrinsics(
comp, intrinsic_uri)
if transform_modified:
# Required because merging called intrinsics invokes building block
# factories that do not name references uniquely.
comp, _ = tree_transformations.uniquify_reference_names(comp)
modified = modified or transform_modified
if modified:
# Required because merging called intrinsics will nest the called
# intrinsics such that they can no longer be split.
comp, _ = _extract_intrinsics_to_top_level_lambda(comp, uri)
return comp
