wehub-resource-sync
772 lines
31 KiB
Python
772 lines
31 KiB
Python
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file8 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.
|
|
# ======================================
|
|
|
|
"""OutsideCompilation, TPUReplicateContext, and supporting functions."""
|
|
|
|
from typing import Any, Callable, List, Optional, Text, Tuple, Union
|
|
from absl import logging
|
|
from tensorflow.core.framework import attr_value_pb2
|
|
from tensorflow.python.distribute import device_util
|
|
from tensorflow.python.distribute import distribute_lib
|
|
from tensorflow.python.framework import device as pydev
|
|
from tensorflow.python.framework import errors
|
|
from tensorflow.python.framework import func_graph
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import control_flow_ops
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.tpu import device_assignment as device_assignment_lib
|
|
from tensorflow.python.tpu.ops import tpu_ops
|
|
from tensorflow.python.types import core as core_types
|
|
from tensorflow.python.util import compat
|
|
from tensorflow.python.util.tf_export import tf_export
|
|
|
|
_MAX_WARNING_LINES = 5
|
|
_TPU_REPLICATE_ATTR = "_tpu_replicate"
|
|
_OUTSIDE_COMPILATION_ATTR = "_xla_outside_compilation"
|
|
_MAP_OUTSIDE_COMPILATION_ATTR = "_xla_map_outside_compilation"
|
|
|
|
# Operations that indicate some error in the users graph, e.g. a placeholder
|
|
# that's introduced outside of the infeed.
|
|
_DENYLISTED_OPS = frozenset([
|
|
"Placeholder",
|
|
])
|
|
|
|
|
|
# XLA doesn't currently support reading of intermediate tensors, thus some ops
|
|
# are not supported.
|
|
_UNSUPPORTED_OPS = frozenset([
|
|
"AudioSummary",
|
|
"AudioSummaryV2",
|
|
"HistogramSummary",
|
|
"ImageSummary",
|
|
"MergeSummary",
|
|
"Print",
|
|
"ScalarSummary",
|
|
"TensorSummary",
|
|
"TensorSummaryV2",
|
|
])
|
|
|
|
|
|
def is_tpu_strategy(strategy: Any) -> bool:
|
|
is_tpu_strat = lambda k: k.__name__.startswith("TPUStrategy")
|
|
clz = strategy.__class__
|
|
return is_tpu_strat(clz) or any(map(is_tpu_strat, clz.__bases__))
|
|
|
|
|
|
def _enclosing_tpu_device_assignment(
|
|
) -> Optional[device_assignment_lib.DeviceAssignment]:
|
|
if not distribute_lib.has_strategy():
|
|
return None
|
|
strategy = distribute_lib.get_strategy()
|
|
if not is_tpu_strategy(strategy):
|
|
return None
|
|
return strategy.extended._device_assignment # pylint: disable=protected-access
|
|
|
|
|
|
class TPUReplicateContext(control_flow_ops.XLAControlFlowContext):
|
|
"""A `ControlFlowContext` for nodes inside a TPU computation.
|
|
|
|
The primary role of `TPUReplicateContext` is to mark operators inside a
|
|
tpu.replicate() computation with the attribute "_tpu_replicate=XYZ", where XYZ
|
|
is a unique name.
|
|
|
|
We use a `ControlFlowContext` to perform the annotation since it integrates
|
|
with Tensorflow constructs like ResourceVariables. For example, if a
|
|
`ResourceVariable` is constructed inside a tpu.replicate() block, the
|
|
`ResourceVariable` implementation can use
|
|
`with ops.control_dependencies(None)` to build the variable's definition
|
|
outside the replicated computation.
|
|
"""
|
|
|
|
def __init__(self, name: Text, num_replicas: int, pivot: ops.Operation):
|
|
"""Builds a new TPUReplicateContext.
|
|
|
|
Args:
|
|
name: a unique name for the context, used to populate the `_tpu_replicate`
|
|
attribute.
|
|
num_replicas: an integer that gives the number of replicas for the
|
|
computation.
|
|
pivot: a pivot node. Nodes in the TPUReplicateContext that do not have any
|
|
inputs will have a control dependency on the pivot node. This ensures
|
|
that nodes are correctly included in any enclosing control flow
|
|
contexts.
|
|
"""
|
|
super(TPUReplicateContext, self).__init__()
|
|
self._num_replicas = num_replicas
|
|
self._outer_device_function_stack = None
|
|
self._oc_dev_fn_stack = None
|
|
self._outside_compilation_cluster = None
|
|
self._is_map_outside_compilation = False
|
|
self._outside_compilation_v2_context = None
|
|
self._outside_compilation_counter = 0
|
|
self._in_gradient_colocation = None
|
|
self._gradient_colocation_stack = []
|
|
self._host_compute_core = []
|
|
self._name = name
|
|
self._tpu_replicate_attr = attr_value_pb2.AttrValue(
|
|
s=compat.as_bytes(self._name)
|
|
)
|
|
self._unsupported_ops = []
|
|
self._pivot = pivot
|
|
self._replicated_vars = {}
|
|
|
|
def get_replicated_var_handle(self,
|
|
name: Text,
|
|
handle_id: Text,
|
|
vars_: Union[List[core_types.Tensor],
|
|
List[variables.Variable]],
|
|
is_mirrored: bool = False,
|
|
is_packed: bool = False) -> core_types.Tensor:
|
|
"""Returns a variable handle for replicated TPU variable 'var'.
|
|
|
|
This is a method used by an experimental replicated variable implementation
|
|
and is not intended as a public API.
|
|
|
|
Args:
|
|
name: The common name of the variable.
|
|
handle_id: Unique ID of the variable handle, used as the cache key.
|
|
vars_: The replicated TPU variables or handles.
|
|
is_mirrored: Whether the variables are mirrored, which guarantees the
|
|
values in each replica are always the same.
|
|
is_packed: Whether the replicated variables are packed into one variable.
|
|
|
|
Returns:
|
|
The handle of the TPU replicated input node.
|
|
"""
|
|
device_assignment = _enclosing_tpu_device_assignment()
|
|
# We don't need to put device assignment as part of the replicated_vars key
|
|
# because each TPUReplicateContext will only have one device assignment.
|
|
handle = self._replicated_vars.get(handle_id)
|
|
if handle is not None:
|
|
return handle
|
|
|
|
if device_assignment is not None and not is_packed:
|
|
# Find a variable copy for each replica in the device assignment.
|
|
# Note that the order of devices for replicas for the variable and the
|
|
# device assignment might not match.
|
|
job_name = pydev.DeviceSpec.from_string(vars_[0].device).job
|
|
devices_to_vars = {device_util.canonicalize(v.device): v for v in vars_}
|
|
replicated_vars = []
|
|
for replica_id in range(device_assignment.num_replicas):
|
|
for logical_core in range(device_assignment.num_cores_per_replica):
|
|
device = device_util.canonicalize(
|
|
device_assignment.tpu_device(
|
|
replica=replica_id, logical_core=logical_core, job=job_name))
|
|
if device in devices_to_vars:
|
|
replicated_vars.append(devices_to_vars[device])
|
|
break
|
|
else:
|
|
raise ValueError(
|
|
"Failed to find a variable on any device in replica {} for "
|
|
"current device assignment".format(replica_id)
|
|
)
|
|
else:
|
|
replicated_vars = vars_
|
|
|
|
# Builds a TPUReplicatedInput node for the variable, if one does not already
|
|
# exist. The TPUReplicatedInput node must belong to the enclosing
|
|
# control-flow scope of the TPUReplicateContext.
|
|
# TODO(phawkins): consider changing the contract of the TPU encapsulation
|
|
# so the TPUReplicatedInput nodes go inside the TPUReplicateContext scope
|
|
# instead.
|
|
|
|
_, graph = _enclosing_tpu_context_and_graph()
|
|
with graph.as_default():
|
|
# If replicated_vars are variables, get the handles. Note that this can be
|
|
# done inside TPUReplicateContext because replicated_vars.handle may
|
|
# create new ops.
|
|
if isinstance(replicated_vars[0], variables.Variable):
|
|
replicated_vars = [v.handle for v in replicated_vars]
|
|
# pylint: disable=protected-access
|
|
saved_context = graph._get_control_flow_context()
|
|
graph._set_control_flow_context(self.outer_context)
|
|
handle = tpu_ops.tpu_replicated_input(
|
|
replicated_vars,
|
|
name=name + "/handle",
|
|
is_mirrored_variable=is_mirrored,
|
|
is_packed=is_packed)
|
|
graph._set_control_flow_context(saved_context)
|
|
# pylint: enable=protected-access
|
|
self._replicated_vars[handle_id] = handle
|
|
return handle
|
|
|
|
def report_unsupported_operations(self) -> None:
|
|
if self._unsupported_ops:
|
|
op_str = "\n".join(
|
|
" %s (%s)" % (op.type, op.name) for op in
|
|
self._unsupported_ops[:_MAX_WARNING_LINES])
|
|
logging.warning("%d unsupported operations found: \n%s",
|
|
len(self._unsupported_ops), op_str)
|
|
if len(self._unsupported_ops
|
|
) > _MAX_WARNING_LINES:
|
|
logging.warning("... and %d more",
|
|
(len(self._unsupported_ops) - _MAX_WARNING_LINES))
|
|
|
|
def EnterGradientColocation(self, op: ops.Operation, gradient_uid: Text):
|
|
if op is not None:
|
|
if ops.get_default_graph()._control_flow_context is None: # pylint: disable=protected-access
|
|
# If we are in TF 2 functions (control flow V2 functions, or
|
|
# tf.function()), we need to attach _xla_outside_compilation attribute
|
|
# directly because we are not in TPUReplicateContext.
|
|
try:
|
|
outside_attr = op.get_attr(_OUTSIDE_COMPILATION_ATTR).decode("ascii")
|
|
except ValueError:
|
|
# The attr was not present: do nothing.
|
|
return
|
|
parts = outside_attr.split(".")
|
|
cluster = parts[0] + "." + gradient_uid
|
|
self._outside_compilation_v2_context = OutsideCompilationV2Context(
|
|
cluster)
|
|
self._outside_compilation_v2_context.Enter()
|
|
return
|
|
self._gradient_colocation_stack.append(op)
|
|
if not self._outside_compilation_cluster:
|
|
try:
|
|
outside_attr = op.get_attr(_OUTSIDE_COMPILATION_ATTR).decode("ascii")
|
|
if self._in_gradient_colocation:
|
|
raise NotImplementedError(
|
|
"Cannot nest gradient colocation operations outside compilation"
|
|
)
|
|
if gradient_uid == "__unsupported__":
|
|
raise NotImplementedError(
|
|
"No gradient_uid calling gradient within outside_compilation")
|
|
# When we take the gradient of an op X in an outside_compilation
|
|
# cluster C in a forward computation we would like to put the ops
|
|
# corresponding to the gradient of X into a new outside_compilation
|
|
# cluster C'. However, if we take the gradient of X twice, the second
|
|
# one should get yet another new outside_compilation cluster C''.
|
|
#
|
|
# The mechanism we adopt is to use a 'root_cluster' which is the
|
|
# cluster that X was in before we took gradients, and a 'gradient_uid'
|
|
# which is different for every invocation of gradients, and put the
|
|
# gradient of X in cluster 'root_cluster.gradient_uid'.
|
|
#
|
|
# When taking a gradient of a gradient, some ops will be colocated
|
|
# with Op in the forward pass (e.g., cluster root_cluster) and some in
|
|
# the backward pass (e.g., cluster root_cluster.initial_gradient_uid).
|
|
# We need all of the grad-of-grad ops to be in the same cluster to
|
|
# avoid cyclic dependencies between clusters. We adopt a heuristic
|
|
# that puts any op clustered with root_cluster.<xxx> in
|
|
# root_cluster.gradient_uid, even if xxx was initial_gradient_uid.
|
|
self._in_gradient_colocation = op
|
|
parts = outside_attr.split(".")
|
|
cluster = parts[0] + "." + gradient_uid
|
|
self._EnterOutsideCompilationScope(cluster=cluster)
|
|
except ValueError:
|
|
# The attr was not present: do nothing.
|
|
pass
|
|
|
|
def ExitGradientColocation(self, op: ops.Operation, gradient_uid: Text):
|
|
if op is not None:
|
|
if ops.get_default_graph()._control_flow_context is None: # pylint: disable=protected-access
|
|
# Inside a TF2 tf.function or control flow graph and `op` was not
|
|
# marked to be outside compiled.
|
|
assert self._outside_compilation_v2_context is None
|
|
return
|
|
if self._outside_compilation_v2_context is not None:
|
|
# Inside a TF2 tf.function or control flow graph and `op` was
|
|
# marked to be outside compiled.
|
|
self._outside_compilation_v2_context.Exit()
|
|
self._outside_compilation_v2_context = None
|
|
return
|
|
if not self._gradient_colocation_stack:
|
|
raise errors.InternalError(
|
|
op.node_def, op,
|
|
("Badly nested gradient colocation: "
|
|
+ f"empty stack when popping Op {op.name}")
|
|
)
|
|
last_op = self._gradient_colocation_stack.pop()
|
|
if op is last_op:
|
|
if op is self._in_gradient_colocation:
|
|
self._in_gradient_colocation = None
|
|
self._ExitOutsideCompilationScope()
|
|
else:
|
|
raise errors.InternalError(
|
|
op.node_def, op,
|
|
("Badly nested gradient colocation, " +
|
|
f"expected {last_op}, got {op.name}")
|
|
)
|
|
|
|
def _EnterOutsideCompilationScope(
|
|
self, cluster: Optional[Text] = None, is_map_outside_compilation=False
|
|
):
|
|
class FakeOp(object):
|
|
"""A helper class to determine the current device.
|
|
|
|
Supports only the type and device set/get methods needed to run the
|
|
graph's _apply_device_function method.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._device = ""
|
|
|
|
@property
|
|
def type(self):
|
|
return "FakeOp"
|
|
|
|
@property
|
|
def device(self):
|
|
return self._device
|
|
|
|
def _set_device(self, device):
|
|
if isinstance(device, pydev.DeviceSpec):
|
|
self._device = device.to_string()
|
|
else:
|
|
self._device = device
|
|
|
|
def _set_device_from_string(self, device_str):
|
|
self._device = device_str
|
|
|
|
if self._outside_compilation_cluster:
|
|
raise NotImplementedError("Cannot nest outside_compilation clusters")
|
|
if cluster:
|
|
self._outside_compilation_cluster = cluster
|
|
else:
|
|
self._outside_compilation_cluster = str(self._outside_compilation_counter)
|
|
self._outside_compilation_counter += 1
|
|
if is_map_outside_compilation:
|
|
self._is_map_outside_compilation = True
|
|
graph = ops.get_default_graph()
|
|
fake_op = FakeOp()
|
|
graph._apply_device_functions(fake_op) # pylint: disable=protected-access
|
|
device = pydev.DeviceSpec.from_string(fake_op.device)
|
|
if (device.device_type == "TPU_REPLICATED_CORE" and
|
|
device.device_index is not None):
|
|
self._host_compute_core.append(self._outside_compilation_cluster + ":" +
|
|
str(device.device_index))
|
|
self._oc_dev_fn_stack = graph._device_function_stack # pylint: disable=protected-access
|
|
graph._device_function_stack = self._outer_device_function_stack # pylint: disable=protected-access
|
|
|
|
def _ExitOutsideCompilationScope(self):
|
|
if not self._outside_compilation_cluster:
|
|
raise ValueError(
|
|
"Attempted to exit outside_compilation scope when not in scope")
|
|
self._outside_compilation_cluster = None
|
|
self._is_map_outside_compilation = False
|
|
graph = ops.get_default_graph()
|
|
graph._device_function_stack = self._oc_dev_fn_stack # pylint: disable=protected-access
|
|
|
|
def Enter(self) -> None:
|
|
if not self._outer_device_function_stack:
|
|
# Capture the device function stack at the time of first entry
|
|
# since that is the stack that will be used outside_compilation.
|
|
graph = ops.get_default_graph()
|
|
# pylint: disable=protected-access
|
|
self._outer_device_function_stack = graph._device_function_stack.copy()
|
|
# pylint: enable=protected-access
|
|
super(TPUReplicateContext, self).Enter()
|
|
|
|
def HostComputeCore(self) -> List[Text]:
|
|
return self._host_compute_core
|
|
|
|
def _RemoveExternalControlEdges(
|
|
self,
|
|
op: ops.Operation) -> Tuple[List[ops.Operation], List[ops.Operation]]:
|
|
"""Remove any external control dependency on this op."""
|
|
internal_control_inputs = []
|
|
external_control_inputs = []
|
|
for x in op.control_inputs:
|
|
# pylint: disable=protected-access
|
|
is_internal_op = False
|
|
ctxt = x._get_control_flow_context()
|
|
while ctxt is not None:
|
|
if ctxt == self:
|
|
is_internal_op = True
|
|
break
|
|
ctxt = ctxt._outer_context
|
|
if is_internal_op:
|
|
internal_control_inputs.append(x)
|
|
else:
|
|
external_control_inputs.append(x)
|
|
# pylint: enable=protected-access
|
|
# pylint: disable=protected-access
|
|
op._remove_all_control_inputs()
|
|
op._add_control_inputs(internal_control_inputs)
|
|
# pylint: enable=protected-access
|
|
return internal_control_inputs, external_control_inputs
|
|
|
|
def AddOp(self, op: ops.Operation) -> None:
|
|
# pylint: disable=protected-access
|
|
if op.type in _DENYLISTED_OPS:
|
|
logging.error(
|
|
"Operation of type %s (%s) is not supported on the TPU. "
|
|
"Execution will fail if this op is used in the graph. ", op.type,
|
|
op.name)
|
|
|
|
if op.type in _UNSUPPORTED_OPS:
|
|
self._unsupported_ops.append(op)
|
|
|
|
if any(x.dtype._is_ref_dtype for x in op.inputs):
|
|
raise NotImplementedError(
|
|
f"Non-resource Variables are not supported inside TPU computations "
|
|
f"(operator name: {op.name})")
|
|
|
|
# TensorFlowOpLayer may clone nodes that are in tpu.rewrite()s. It'll add
|
|
# the "_cloned" attribute and we should continue in that case.
|
|
if (_TPU_REPLICATE_ATTR in op.node_def.attr and
|
|
"_cloned" not in op.node_def.attr):
|
|
raise ValueError(f"TPU computations cannot be nested on op ({op})")
|
|
op._set_attr(_TPU_REPLICATE_ATTR, self._tpu_replicate_attr)
|
|
if self._outside_compilation_cluster:
|
|
op._set_attr(
|
|
_OUTSIDE_COMPILATION_ATTR,
|
|
attr_value_pb2.AttrValue(
|
|
s=compat.as_bytes(self._outside_compilation_cluster)))
|
|
if self._is_map_outside_compilation:
|
|
op._set_attr(
|
|
_MAP_OUTSIDE_COMPILATION_ATTR,
|
|
attr_value_pb2.AttrValue(b=True),
|
|
)
|
|
if self._num_replicas > 1 or not self._outside_compilation_cluster:
|
|
# Prevent feeding or fetching anything that is being compiled,
|
|
# and any replicated outside_compilation Op.
|
|
op.graph.prevent_feeding(op)
|
|
op.graph.prevent_fetching(op)
|
|
|
|
# Remove any control edges from outer control flow contexts. These may cause
|
|
# mismatched frame errors.
|
|
(internal_control_inputs,
|
|
external_control_inputs) = self._RemoveExternalControlEdges(op)
|
|
|
|
if not op.inputs:
|
|
# Add a control edge from the control pivot to this op.
|
|
if not internal_control_inputs:
|
|
# pylint: disable=protected-access
|
|
op._add_control_input(self.GetControlPivot())
|
|
# pylint: enable=protected-access
|
|
else:
|
|
for index in range(len(op.inputs)):
|
|
x = op.inputs[index]
|
|
real_x = self.AddValue(x)
|
|
if real_x is not x:
|
|
op._update_input(index, real_x) # pylint: disable=protected-access
|
|
|
|
if external_control_inputs:
|
|
# Use an identity to pull control inputs as data inputs. Note that we
|
|
# ignore ops which don't have outputs. TODO(phawkins): fix that.
|
|
with ops.control_dependencies(None):
|
|
self.Enter()
|
|
external_control_inputs = [
|
|
array_ops.identity(x.outputs[0]).op
|
|
for x in external_control_inputs
|
|
if x.outputs
|
|
]
|
|
self.Exit()
|
|
# pylint: disable=protected-access
|
|
op._add_control_inputs(external_control_inputs)
|
|
# pylint: enable=protected-access
|
|
|
|
# Mark op's outputs as seen by this context and any outer contexts.
|
|
output_names = [x.name for x in op.outputs]
|
|
context = self
|
|
while context is not None:
|
|
# pylint: disable=protected-access
|
|
context._values.update(output_names)
|
|
context = context._outer_context
|
|
# pylint: enable=protected-access
|
|
|
|
if self._outer_context:
|
|
self._outer_context.AddInnerOp(op)
|
|
|
|
def AddValue(self, val: core_types.Tensor) -> core_types.Tensor:
|
|
"""Add `val` to the current context and its outer context recursively."""
|
|
if not self._outer_context:
|
|
return val
|
|
|
|
if val.name in self._values:
|
|
# Use the real value if it comes from outer context.
|
|
result = self._external_values.get(val.name)
|
|
return val if result is None else result
|
|
|
|
result = val
|
|
self._values.add(val.name)
|
|
if self._outer_context:
|
|
result = self._outer_context.AddValue(val)
|
|
self._values.add(result.name)
|
|
|
|
self._external_values[val.name] = result
|
|
|
|
return result
|
|
|
|
def AddInnerOp(self, op: ops.Operation):
|
|
self.AddOp(op)
|
|
if self._outer_context:
|
|
self._outer_context.AddInnerOp(op)
|
|
|
|
@property
|
|
def grad_state(self):
|
|
# Define the gradient loop state associated with the TPUReplicateContext to
|
|
# be None as the TPUReplicateContext does not get nested nor does the
|
|
# grad_state outside the TPUReplicateContext affect the graph inside so the
|
|
# grad_state should be as if this is the top-level gradient state.
|
|
return None
|
|
|
|
@property
|
|
def back_prop(self):
|
|
"""Forwards to the enclosing while context, if any."""
|
|
if self.GetWhileContext():
|
|
return self.GetWhileContext().back_prop
|
|
return False
|
|
|
|
def GetControlPivot(self) -> ops.Operation:
|
|
return self._pivot
|
|
|
|
def RequiresUniqueFunctionRetracing(self):
|
|
# More context: b/158152827. TPU stack uses the TPUReplicateContext to
|
|
# create replicated variable handles and cluster TPU computations, thus we
|
|
# always retrace a tf.function when the wrapped TPUReplicateContext changes.
|
|
return True
|
|
|
|
|
|
def _enclosing_tpu_context_and_graph() -> Tuple[Any, Any]:
|
|
"""Returns the TPUReplicateContext and its associated graph."""
|
|
graph = ops.get_default_graph()
|
|
while graph is not None:
|
|
# pylint: disable=protected-access
|
|
context_ = graph._get_control_flow_context()
|
|
# pylint: enable=protected-access
|
|
while context_ is not None:
|
|
if isinstance(context_, TPUReplicateContext):
|
|
return context_, graph
|
|
context_ = context_.outer_context
|
|
graph = getattr(graph, "outer_graph", None)
|
|
raise ValueError("get_replicated_var_handle() called without "
|
|
"TPUReplicateContext. This shouldn't happen. Please file "
|
|
"a bug.")
|
|
|
|
|
|
class OutsideCompilationV2Context(control_flow_ops.ControlFlowContext):
|
|
"""The context for outside compilation in Tensorflow 2.0.
|
|
|
|
Every op added in this context will be assigned an _xla_outside_compilation
|
|
attribute.
|
|
"""
|
|
|
|
def __init__(self, name: Text, is_map_outside_compilation=False):
|
|
control_flow_ops.ControlFlowContext.__init__(self)
|
|
self._name = name
|
|
self._is_map_outside_compilation = is_map_outside_compilation
|
|
|
|
def AddOp(self, op: ops.Operation) -> None:
|
|
if self._outer_context:
|
|
self._outer_context.AddOp(op)
|
|
self._set_outside_compilation_attributes(op)
|
|
|
|
def AddInnerOp(self, op: ops.Operation) -> None:
|
|
if self._outer_context:
|
|
self._outer_context.AddInnerOp(op)
|
|
self._set_outside_compilation_attributes(op)
|
|
|
|
def to_control_flow_context_def(self, context_def, export_scope=None):
|
|
raise NotImplementedError
|
|
|
|
def _set_outside_compilation_attributes(self, op: ops.Operation) -> None:
|
|
# pylint: disable=protected-access
|
|
op._set_attr(
|
|
_OUTSIDE_COMPILATION_ATTR,
|
|
attr_value_pb2.AttrValue(s=compat.as_bytes(self._name)),
|
|
)
|
|
if self._is_map_outside_compilation:
|
|
op._set_attr(
|
|
_MAP_OUTSIDE_COMPILATION_ATTR, attr_value_pb2.AttrValue(b=True)
|
|
)
|
|
# pylint: enable=protected-access
|
|
|
|
|
|
def outside_compilation_impl(
|
|
is_map, computation: Callable[..., Any], *args, **kwargs
|
|
) -> Any:
|
|
"""Tags ops in `computation` with outside compilation attributes for ordinary `outside_compilation` or `map_outside_compilation`."""
|
|
args = [] if args is None else args
|
|
graph = ops.get_default_graph()
|
|
|
|
# If we are in TF 2 functions (control flow V2 functions, or tf.function()),
|
|
# we need to attach _xla_outside_compilation attribute directly because we are
|
|
# not in TPUReplicateContext.
|
|
if isinstance(graph, func_graph.FuncGraph):
|
|
try:
|
|
tpu_context, _ = _enclosing_tpu_context_and_graph()
|
|
except ValueError:
|
|
logging.warning(
|
|
"Outside compilation attempted outside TPUReplicateContext "
|
|
"scope. As no enclosing TPUReplicateContext can be found, "
|
|
"returning the result of `computation` as is."
|
|
)
|
|
return computation(*args, **kwargs)
|
|
|
|
# pylint: disable=protected-access
|
|
outside_compilation_name = str(tpu_context._outside_compilation_counter)
|
|
tpu_context._outside_compilation_counter = (
|
|
tpu_context._outside_compilation_counter + 1
|
|
)
|
|
# pylint: enable=protected-access
|
|
|
|
outside_compilation_context = OutsideCompilationV2Context(
|
|
outside_compilation_name, is_map_outside_compilation=is_map
|
|
)
|
|
outside_compilation_context.Enter()
|
|
args = [] if args is None else args
|
|
retval = computation(*args, **kwargs)
|
|
outside_compilation_context.Exit()
|
|
return retval
|
|
|
|
# If we are in a TPUReplicateContext, signal that we are now
|
|
# outside_compilation
|
|
initial_context = graph._get_control_flow_context() # pylint: disable=protected-access
|
|
context = initial_context
|
|
while context:
|
|
if isinstance(context, TPUReplicateContext):
|
|
context._EnterOutsideCompilationScope(is_map_outside_compilation=is_map) # pylint: disable=protected-access
|
|
context = context.outer_context
|
|
|
|
retval = computation(*args, **kwargs)
|
|
|
|
# If we are in a TPUReplicateContext, signal that we are no longer
|
|
# outside_compilation
|
|
final_context = graph._get_control_flow_context() # pylint: disable=protected-access
|
|
if initial_context is not final_context:
|
|
raise NotImplementedError(
|
|
"Control-flow context cannot be different at start and end of an "
|
|
"outside_compilation scope"
|
|
)
|
|
context = initial_context
|
|
while context:
|
|
if isinstance(context, TPUReplicateContext):
|
|
context._ExitOutsideCompilationScope() # pylint: disable=protected-access
|
|
context = context.outer_context
|
|
|
|
return retval
|
|
|
|
|
|
@tf_export(v1=["tpu.outside_compilation"])
|
|
def outside_compilation(
|
|
computation: Callable[..., Any], *args, **kwargs
|
|
) -> Any:
|
|
"""Builds part of a computation outside any current TPU replicate scope.
|
|
|
|
`tf.tpu.outside_compilation()` is used to run ops in `computation` on CPU
|
|
instead of running on TPU. For example, users can run ops that are not
|
|
supported on TPU's (e.g. tf.summary.write()) by explicitly placing those
|
|
ops on CPU's. Below usage of outside compilation will place ops in
|
|
`computation_with_string_ops` on CPU.
|
|
|
|
Example usage:
|
|
|
|
```python
|
|
def computation_with_string_ops(x):
|
|
# strings types are not supported on TPU's and below ops must
|
|
# run on CPU instead.
|
|
output = tf.strings.format('1{}', x)
|
|
return tf.strings.to_number(output)
|
|
|
|
def tpu_computation():
|
|
# Expected output is 11.
|
|
output = tf.tpu.outside_compilation(computation_with_string_ops, 1)
|
|
```
|
|
|
|
Outside compilation should be called inside TPUReplicateContext. That is,
|
|
`tf.tpu.outside_compilation()` should be called inside a function that is
|
|
passed to `tpu.split_compile_and_replicate()` -- this is implied when
|
|
outside compilation is invoked inside a function passed to TPUStrategy
|
|
`run()`. If invoked outside of TPUReplicateContext,
|
|
then this simply returns the result of `computation`, and therefore,
|
|
would be a no-op. Note that outside compilation is different from
|
|
`tf.distribute.experimental.TPUStrategy.merge_call()` as logic in
|
|
outside compilation is replicated and executed separately for each
|
|
replica. On the other hand, `merge_call()` requires a `merge_fn`
|
|
to aggregate the inputs from different replicas and is executed only
|
|
once.
|
|
|
|
For variables placed in TPU device, which includes variables created inside
|
|
TPUStrategy scope, outside compilation logic must not include variable
|
|
read/write. For variables placed on host, variable read/write is only allowed
|
|
if the variable is not accessed by any other ops in the TPU computation.
|
|
Variable read/write from outside compilation cluster is not visible from TPU
|
|
computation and vice versa. Therefore, if outside compilation logic contains
|
|
such host variables read/write ops and if the variables are accessed by TPU
|
|
computation as well, then this may lead to deadlock.
|
|
|
|
Internally, `tf.tpu.outside_compilation()` adds outside compilation
|
|
attributes to all ops in `computation`. During a later passes ops with outside
|
|
compilation attributes are moved to a host-side graph. Inputs to this extract
|
|
host-side graph are sent from TPU computation graph to host graph via a pair
|
|
of XlaSendToHost and XlaRecvFromHost ops. Note that using
|
|
`tf.tpu.outside_compilation()` may result in tensor transfer between TPU and
|
|
CPU, leading to non-trivial performance impact.
|
|
|
|
Args:
|
|
computation: A Python function that builds the computation to place on the
|
|
host.
|
|
*args: the positional arguments for the computation.
|
|
**kwargs: the keyword arguments for the computation.
|
|
|
|
Returns:
|
|
The Tensors returned by computation.
|
|
"""
|
|
return outside_compilation_impl(False, computation, *args, **kwargs)
|
|
|
|
|
|
def experimental_map_outside_compilation(
|
|
computation: Callable[..., Any], *args, **kwargs
|
|
) -> Any:
|
|
"""Maps `computation` onto shards and puts it outside any current TPU replicate scope.
|
|
|
|
`experimental_map_outside_compilation(f, x)` maps `f` onto the shards
|
|
of `x`, where `x` is split-sharded. Each invocation of `f` on a split occurs
|
|
on the CPU that's associated with the TPU that owns the split.
|
|
|
|
Example usage:
|
|
|
|
```python
|
|
def normalize_each_split(split):
|
|
return split - tf.math.reduce_mean(split)
|
|
|
|
def tpu_computation(x):
|
|
x_split = strategy.experimental_split_to_logical_devices(
|
|
x, [num_cores_per_replica, 1])
|
|
y = experimental_map_outside_compilation(
|
|
normalize_each_split, x_split)
|
|
y_split = strategy.experimental_split_to_logical_devices(
|
|
x, [num_cores_per_replica, 1])
|
|
return y_split
|
|
```
|
|
|
|
`experimental_map_outside_compilation` should be called inside
|
|
TPUReplicateContext. That is, `outside_compilation()` should be called
|
|
inside a function that is passed to `tpu.split_compile_and_replicate()` --
|
|
this is implied when outside compilation is invoked inside a function passed
|
|
to TPUStrategy `run()`. It is invalid to invoke outside of
|
|
TPUReplicateContext.
|
|
|
|
`experimental_map_outside_compilation` should input and output tensors that
|
|
are located on the TPU.
|
|
|
|
Internally, `experimental_map_outside_compilation()` adds outside
|
|
compilation attributes to all ops in `computation` and moves outside-compiled
|
|
ops to a host-side graph. This is similar to `tf.tpu.outside_compilation()`.
|
|
Send/recv ops from/to the TPU send each split directly to the TPU's host.
|
|
|
|
Args:
|
|
computation: A Python function that builds the computation to place on the
|
|
host.
|
|
*args: the positional arguments for the computation.
|
|
**kwargs: the keyword arguments for the computation.
|
|
|
|
Returns:
|
|
The Tensors returned by computation.
|
|
"""
|
|
return outside_compilation_impl(True, computation, *args, **kwargs)
|