Adds default compiler pipeline for TFF runtime.

PiperOrigin-RevId: 321423023

tensorflow_federated/python/core/backends/native/BUILD

@@ -10,7 +10,24 @@ py_library(
     srcs = ["__init__.py"],
     srcs_version = "PY3",
     visibility = ["//tensorflow_federated/python/core/backends:__pkg__"],
-    deps = [":execution_contexts"],
+    deps = [
+        ":compiler",
+        ":execution_contexts",
+    ],
+)
+
+py_library(
+    name = "compiler",
+    srcs = ["compiler.py"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorflow_federated/python/core/api:computation_base",
+        "//tensorflow_federated/python/core/impl:computation_impl",
+        "//tensorflow_federated/python/core/impl/compiler:building_blocks",
+        "//tensorflow_federated/python/core/impl/compiler:transformations",
+        "//tensorflow_federated/python/core/impl/compiler:tree_transformations",
+        "//tensorflow_federated/python/core/impl/wrappers:computation_wrapper_instances",
+    ],
 )
 
 py_library(
@@ -18,6 +35,7 @@ py_library(
     srcs = ["execution_contexts.py"],
     srcs_version = "PY3",
     deps = [
+        ":compiler",
         "//tensorflow_federated/python/core/impl/context_stack:context_stack_impl",
         "//tensorflow_federated/python/core/impl/executors:execution_context",
         "//tensorflow_federated/python/core/impl/executors:executor_stacks",

tensorflow_federated/python/core/backends/native/__init__.py

@@ -13,4 +13,5 @@
 # limitations under the License.
 """Utilities for interacting with a native backend."""
 
+from tensorflow_federated.python.core.backends.native.compiler import transform_to_native_form
 from tensorflow_federated.python.core.backends.native.execution_contexts import set_local_execution_context

tensorflow_federated/python/core/backends/native/compiler.py

+# Copyright 2020, The TensorFlow Federated Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Library of compiler functions for usage in the native execution context."""
+from absl import logging
+
+from tensorflow_federated.python.core.api import computation_base
+from tensorflow_federated.python.core.impl import computation_impl
+from tensorflow_federated.python.core.impl.compiler import building_blocks
+from tensorflow_federated.python.core.impl.compiler import transformations
+from tensorflow_federated.python.core.impl.compiler import tree_transformations
+from tensorflow_federated.python.core.impl.wrappers import computation_wrapper_instances
+
+
+def transform_to_native_form(
+    comp: computation_base.Computation) -> computation_base.Computation:
+  """Compiles a computation for execution in the TFF native runtime.
+
+  This function transforms the proto underlying `comp` by first transforming it
+  to call-dominant form (see `tff.framework.transform_to_call_dominant` for
+  definition), then computing information on the dependency structure of the
+  bindings and remapping them into tuples, such that every computation is
+  evaluated as early as possible, and parallelized with any other computation
+  with which it shares dependency structure.
+
+  Args:
+    comp: Instance of `computation_base.Computation` to compile.
+
+  Returns:
+    A new `computation_base.Computation` representing the compiled version of
+    `comp`.
+  """
+  proto = computation_impl.ComputationImpl.get_proto(comp)
+  computation_building_block = building_blocks.ComputationBuildingBlock.from_proto(
+      proto)
+  try:
+    logging.debug('Compiling TFF computation.')
+    call_dominant_form, _ = transformations.transform_to_call_dominant(
+        computation_building_block)
+    locals_rearranged, _ = tree_transformations.group_block_locals_by_dependency(
+        call_dominant_form)
+    logging.debug('Computation compiled to:')
+    logging.debug(locals_rearranged.formatted_representation())
+    return computation_wrapper_instances.building_block_to_computation(
+        locals_rearranged)
+  except ValueError as e:
+    logging.debug('Compilation for native runtime failed with error %s', e)
+    logging.debug('computation: %s',
+                  computation_building_block.compact_representation())
+    return comp

tensorflow_federated/python/core/backends/native/execution_contexts.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Execution contexts for the native backend."""
 
+from tensorflow_federated.python.core.backends.native import compiler
 from tensorflow_federated.python.core.impl.context_stack import context_stack_impl
 from tensorflow_federated.python.core.impl.executors import execution_context
 from tensorflow_federated.python.core.impl.executors import executor_stacks
@@ -30,5 +31,6 @@ def set_local_execution_context(num_clients=None,
       num_client_executors=num_client_executors,
       server_tf_device=server_tf_device,
       client_tf_devices=client_tf_devices)
-  context = execution_context.ExecutionContext(factory)
+  context = execution_context.ExecutionContext(
+      executor_fn=factory, compiler_fn=compiler.transform_to_native_form)
   context_stack_impl.context_stack.set_default_context(context)