# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# =============================================================================
"""
Utility functions for sharding graphs
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
"""

from tensorflow.compiler.xla import xla_data_pb2
from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager import execute
from tensorflow.python.framework import ops

try:
  from tensorflow.python import pywrap_tfe as pywrap
except ImportError:
  # pywrap_tfe is under a different name in TF1
  from tensorflow.python import pywrap_tensorflow as pywrap

_XLA_SHARDING = '_XlaSharding'


def get_sharding(op):  # pylint: disable=missing-type-doc,missing-return-type-doc
  """Get the sharding for the given op.

  Args:
    op: An operation.

  Returns:
    None if the operation has no sharding, otherwise
    the shard number.
  """
  if has_attr(op, _XLA_SHARDING):
    attr = op.get_attr(_XLA_SHARDING)

    sharding_proto = xla_data_pb2.OpSharding()
    sharding_proto.ParseFromString(attr)

    # Don't return a protobuf type.
    return list(sharding_proto.tile_assignment_devices)

  return None


def has_attr(o, attr_name):
  """Test for the presence of a specific attribute.

  Args:
    o: An operation.
    attr_name: The name of an attribute to test for.

  Returns:
    True if the operation has the given attribute.
  """
  for i in o.node_def.attr.items():
    if i[0] == attr_name:
      return True
  return False


def dependencies(roots):
  """Find a list of ancestor operations for a given set of root operations

  Args:
    roots: The root operations from which to start.

  """
  working = roots
  op_set = set()
  while len(working) > 0:
    o = working[0]
    if type(o) is ops.Tensor:
      o = o.op
    working = working[1:]
    op_set.add(o)
    working += [t.op for t in o.inputs if t.op not in op_set]
    working += [c for c in o.control_inputs if c not in op_set]
  return op_set


def get_shard_from_colocation(op):
  """Find the shard number from an op which shares co-location information with
  the given operation.

  Args:
    op: The operation to apply sharding to.
  """
  g = op.graph
  for c in op.colocation_groups():
    try:
      coloc_op = g.get_operation_by_name(c.decode('utf-8')[5:])
      if has_attr(coloc_op, _XLA_SHARDING):
        attr = coloc_op.get_attr(_XLA_SHARDING)
        return attr
    except KeyError:
      continue
  for c in op.inputs:
    coloc_op = c.op
    if has_attr(coloc_op, _XLA_SHARDING):
      attr = coloc_op.get_attr(_XLA_SHARDING)
      return attr
  return None


def propagate_sharding(g):
  """Move the sharding from the forward pass operations onto their co-located
  backward pass operations.

  Args:
    g: The graph.
  """
  changed = True
  while changed:
    changed = False

    op_list = g.get_operations()
    op_list = filter(lambda o: has_attr(o, '_class'), op_list)
    op_list = filter(lambda o: not has_attr(o, '_XlaSharding'), op_list)
    for o in op_list:
      attr = get_shard_from_colocation(o)
      if attr is not None:
        o._set_attr(_XLA_SHARDING, attr_value_pb2.AttrValue(s=attr))
        changed = True
        break


def enable_sharded_gradient_tape():
  """Enable backward ops generated by `tf.GradientTape` to inherit the
  sharding of their forward op.
  """
  # Gradients are automatically recorded during construction of the forward
  # Op. We register our own registration callback to extend the captured
  # attributes which allows us to keep track of the sharding information for
  # the backward Op.
  # We don't check whether sharding is supported here as the eager state can
  # change after calling this.
  pywrap.TFE_Py_RegisterGradientFunction(_sharded_backprop_gradient_function)

  execute.record_gradient = _record_gradient_with_sharding


def _sharded_backprop_gradient_function(op_name, attr_tuple, *args, **kwargs):
  attr_scope = {}
  # pylint: disable=protected-access
  if _sharding_supported() and attr_tuple and len(attr_tuple) >= 2:

    attr_name, value = attr_tuple[-2:]
    if attr_name == _XLA_SHARDING:
      attr_scope[attr_name] = value
      attr_tuple = attr_tuple[:-2]

  with ops.get_default_graph()._attr_scope(attr_scope):
    return backprop._gradient_function(op_name, attr_tuple, *args, **kwargs)
  # pylint: enable=protected-access


def _record_gradient_with_sharding(op_name, inputs, attrs, *args, **kwargs):
  # Pipeline backward stage cant be explicitly sharded, its implicitly
  # sharded based on the corresponding forward stage.
  if _sharding_supported() and op_name != "PipelineStage":
    graph = ops.get_default_graph()

    # pylint: disable=protected-access
    scoped_attrs = graph._attr_scope_map
    shard = scoped_attrs.get(_XLA_SHARDING)
    if shard:
      shard_attrs = (_XLA_SHARDING, shard)
      attrs = attrs + shard_attrs if attrs is not None else shard_attrs

  return backprop.record_gradient(op_name, inputs, attrs, *args, **kwargs)


def _sharding_supported():
  # sharding is for compiled graphs only and so not supported
  # in eager mode.
  return not context.executing_eagerly()
