path: root/tensorflow/compiler/xla/service/hlo_sharding_metadata.cc

/* Copyright 2018 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.
==============================================================================*/

#include "tensorflow/compiler/xla/service/hlo_sharding_metadata.h"

#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/shape_tree.h"
#include "tensorflow/compiler/xla/shape_util.h"

namespace xla {

namespace {

struct PassThrough {
  PassThrough(HloInstruction* user, HloInstruction* operand)
      : user(user), operand(operand) {}

  HloInstruction* user = nullptr;
  HloInstruction* operand = nullptr;
};

void SetSingleSharding(HloInstruction* instruction,
                       const HloSharding& sharding) {
  VLOG(4) << "  " << instruction->name() << " to " << sharding;
  instruction->set_single_sharding(sharding);
}

bool ShardingMatches(const HloSharding& sharding1,
                     const HloSharding& sharding2) {
  auto single_sharding1 = sharding1.ExtractSingleSharding();
  if (single_sharding1) {
    auto single_sharding2 = sharding2.ExtractSingleSharding();
    if (single_sharding2) {
      return *single_sharding1 == single_sharding2;
    }
  }
  // Anything which is not unique across all elements, gets a full sharding
  // compare.
  return sharding1 == sharding2;
}

// When we create domains, they are never "empty", where with empty we mean
// that a kDomain instruction has as operand another kDomain instruction of the
// same kind.
// But when the HLO optimizations are run, empty domains can be created.
// For example:
//
//  Domain(device=None, device=0) ->
//    Tuple(device=0) ->
//      GTE(device=0) ->
//        Domain(device=0, device=None)
//
// In that case the tuple simplifier could create something like:
//
//  Domain(device=None, device=0) -> Domain(device=0, device=None)
//
// Which is a so called empty domain.
// In the case above, crossing an empty domain which was transiting through
// device 0, requires the normalization phase to fixup the empty domain by
// adding back a Tuple+GTE pair with the proper device.
// One particular case where this can create problems is the result of the
// entry computation, where the GTE assignments are used by TF to tell the
// XLA where the results should be sent.
std::vector<PassThrough> LocatePassThroughDomainLinks(
    const DomainMetadata::Domain& domain) {
  std::vector<PassThrough> pass_through;
  for (HloInstruction* instruction : domain.enter_domains) {
    CHECK(instruction->opcode() == HloOpcode::kDomain)
        << "Instruction is not a kDomain: " << instruction->ToString();
    for (HloInstruction* user : instruction->users()) {
      if (user->opcode() == HloOpcode::kDomain &&
          domain.exit_domains.count(user) != 0) {
        pass_through.emplace_back(user, instruction);
        VLOG(2) << "Found passthrough domain link:";
        VLOG(2) << "  " << user->ToString();
        VLOG(2) << "  " << instruction->ToString();
      }
    }
  }
  return pass_through;
}

Status FixupPassThroughDomainLinks(const DomainMetadata::Domain& domain,
                                   const HloSharding& sharding) {
  for (auto& pass_through : LocatePassThroughDomainLinks(domain)) {
    HloInstruction* tuple = pass_through.operand->parent()->AddInstruction(
        HloInstruction::CreateTuple({pass_through.operand}));
    HloInstruction* gte = pass_through.operand->parent()->AddInstruction(
        HloInstruction::CreateGetTupleElement(pass_through.operand->shape(),
                                              tuple, 0));
    gte->set_sharding(sharding);
    TF_RETURN_IF_ERROR(
        pass_through.operand->ReplaceUseWith(pass_through.user, gte));
  }
  return Status::OK();
}

std::unique_ptr<HloSharding> CloneShardingForDomain(
    const HloSharding& sharding) {
  auto single_sharding = sharding.ExtractSingleSharding();
  if (!single_sharding) {
    return MakeUnique<HloSharding>(sharding);
  }
  return MakeUnique<HloSharding>(*single_sharding);
}

Status ApplyDomainSingleSharding(const DomainMetadata::Domain& domain,
                                 const HloSharding& sharding) {
  VLOG(4) << "Applying " << sharding << " sharding";
  for (HloInstruction* instruction : domain.instructions) {
    // We only change instructions without sharding, since otherwise we might
    // mess up with eventual HLO passes which has knowledge of it.
    if (!instruction->has_sharding()) {
      SetSingleSharding(instruction, sharding);
    } else {
      VLOG(4) << "  " << instruction->name() << " already has sharding "
              << instruction->sharding();
    }
  }
  return Status::OK();
}

// Retrieves the sharding of a tuple shaped instruction in form of a ShapeTree.
// If the instruction has no sharding, a ShapeTree with HloSharding::Replicate()
// sharding will be returned.
ShapeTree<HloSharding> GetTupleSharding(HloInstruction* tuple) {
  if (tuple->has_sharding()) {
    return tuple->sharding().GetAsShapeTree(tuple->shape());
  }
  return ShapeTree<HloSharding>(tuple->shape(), HloSharding::Replicate());
}

// Retrieves the sharding of operand, asked from a user instruction which is
// within domain. If operand is a kDomain, it means that sharding argument is
// the operand sharding, otherwise the operand's own sharding will be returned.
const HloSharding* GetOperandSharding(const HloInstruction* operand,
                                      const DomainMetadata::Domain& domain,
                                      const HloSharding& sharding) {
  DCHECK_EQ(domain.reach_set.count(const_cast<HloInstruction*>(operand)), 1);
  // Here the user of operand is within the domain instruction set, and since it
  // is user of operand, we need to look into the enter_domains set. If this is
  // not a kDomain within the user domains set, then return the operand
  // sharding, if any.
  if (operand->opcode() != HloOpcode::kDomain ||
      domain.enter_domains.count(const_cast<HloInstruction*>(operand)) == 0) {
    return operand->has_sharding() ? &operand->sharding() : nullptr;
  }
  // At this point operand is a kDomain of the currently processed domain, so we
  // can refer to sharding as the domain sharding.
  return &sharding;
}

// Tries to propagate the sharding information into the instructions that are
// part of the domain, in a post order manner (operand propagate to user).
StatusOr<int64> ApplyDomainShardingPass(const DomainMetadata::Domain& domain,
                                        const HloSharding& sharding) {
  int64 assigned = 0;
  for (HloInstruction* instruction : domain.instructions) {
    if (instruction->has_sharding()) {
      continue;
    }
    if (instruction->opcode() == HloOpcode::kGetTupleElement) {
      HloInstruction* tuple = instruction->mutable_operand(0);
      const HloSharding* tuple_sharding =
          GetOperandSharding(tuple, domain, sharding);
      if (tuple_sharding != nullptr) {
        if (tuple_sharding->IsTuple()) {
          HloSharding sub_sharding = tuple_sharding->GetSubSharding(
              tuple->shape(), {instruction->tuple_index()});
          VLOG(4) << "  " << instruction->name() << " to sharding "
                  << sub_sharding;
          instruction->set_sharding(sub_sharding);
        } else {
          SetSingleSharding(instruction, *tuple_sharding);
        }
        ++assigned;
      }
    } else if (instruction->opcode() == HloOpcode::kTuple) {
      int64 tuple_assigned = 0;
      ShapeTree<HloSharding> shape_tree = GetTupleSharding(instruction);
      for (int64 i = 0; i < instruction->operand_count(); ++i) {
        const HloSharding* operand_sharding =
            GetOperandSharding(instruction->operand(i), domain, sharding);
        if (operand_sharding != nullptr &&
            shape_tree.element({i}) != *operand_sharding) {
          *shape_tree.mutable_element({i}) = *operand_sharding;
          ++tuple_assigned;
        }
      }
      if (tuple_assigned > 0) {
        HloSharding tuple_sharding = HloSharding::Tuple(shape_tree);
        VLOG(4) << "  " << instruction->name() << " to sharding "
                << tuple_sharding;
        instruction->set_sharding(tuple_sharding);
        ++assigned;
      }
    } else {
      // If all the operand of the given instruction has the same single device
      // assignment, assign that device to this instruction as well.
      const HloSharding* common_sharding = nullptr;
      for (const HloInstruction* operand : instruction->operands()) {
        const HloSharding* operand_sharding =
            GetOperandSharding(operand, domain, sharding);
        if (operand_sharding != nullptr) {
          if (common_sharding != nullptr &&
              *common_sharding != *operand_sharding) {
            common_sharding = nullptr;
            break;
          }
          common_sharding = operand_sharding;
        }
      }
      if (common_sharding != nullptr) {
        VLOG(4) << "  " << instruction->name() << " to sharding "
                << *common_sharding;
        instruction->set_sharding(*common_sharding);
        ++assigned;
      }
    }
  }
  return assigned;
}

Status ApplyDomainSharding(const DomainMetadata::Domain& domain,
                           const HloSharding& sharding) {
  // Here is the place to call external sharding normalizers, which are
  // implemented in other modules (ie, spatial partitioning).
  // The signature of the external normalizer function should be something
  // like:
  //
  //   StatusOr<bool> Normalizer(const DomainMetadata::Domain&,
  //                             const HloSharding& sharding);
  //
  // The function should return true if it has processed the domain
  // normalization, false if domain was not one recognized by it, or an error.
  // We will call the functions in order below, and fall back to local code if
  // none of the external normalizers acted on the domain.
  // External normalizers should not handle the cases that are already handled
  // locally.

  // None of the external normalizers handled the domain sharding, try to see
  // whether this is a single sharding first.
  auto single_sharding = sharding.ExtractSingleSharding();
  if (single_sharding) {
    // Shortcut the simple case. We have a unique sharding, so we call
    // the ApplyDomainSingleSharding() API which will apply array or tuple
    // shaped sharding to the domain instructions.
    return ApplyDomainSingleSharding(domain, *single_sharding);
  }
  VLOG(1) << "Assigning non-trivial sharding " << sharding;
  for (;;) {
    TF_ASSIGN_OR_RETURN(int64 assigned,
                        ApplyDomainShardingPass(domain, sharding));
    if (assigned == 0) {
      break;
    }
  }
  int64 unassigned = 0;
  for (HloInstruction* instruction : domain.instructions) {
    if (!instruction->has_sharding()) {
      LOG(WARNING) << "Unassigned instruction: " << instruction->ToString();
      ++unassigned;
    }
  }
  // Should we error out if unassigned > 0?
  return Status::OK();
}

// Creates a kDomain instruction to be placed between instruction and operand.
// The kDomain instruction will be created only if the sharding differ between
// the instruction and the operand.
std::unique_ptr<HloInstruction> CreateDomain(HloInstruction* instruction,
                                             HloInstruction* operand) {
  const HloSharding* instruction_sharding =
      instruction->has_sharding() ? &instruction->sharding() : nullptr;
  const HloSharding* operand_sharding =
      operand->has_sharding() ? &operand->sharding() : nullptr;
  // No need for domain if they both have no sharding.
  if (instruction_sharding == nullptr && operand_sharding == nullptr) {
    return nullptr;
  }
  // No need for domain if they match.
  if (instruction_sharding != nullptr && operand_sharding != nullptr &&
      ShardingMatches(*instruction_sharding, *operand_sharding)) {
    return nullptr;
  }
  std::unique_ptr<HloSharding> real_instruction_sharding;
  std::unique_ptr<HloSharding> real_operand_sharding;
  if (instruction_sharding != nullptr) {
    real_instruction_sharding = CloneShardingForDomain(*instruction_sharding);
  }
  if (operand_sharding != nullptr) {
    real_operand_sharding = CloneShardingForDomain(*operand_sharding);
  }
  VLOG(3) << "Creating domain:";
  VLOG(3) << "  Instruction: " << instruction->name();
  VLOG(3) << "  Operand: " << operand->name();
  VLOG(3) << "    User side sharding: "
          << (real_instruction_sharding != nullptr
                  ? real_instruction_sharding->ToString()
                  : "None");
  VLOG(3) << "    Operand side sharding: "
          << (real_operand_sharding != nullptr
                  ? real_operand_sharding->ToString()
                  : "None");

  std::unique_ptr<DomainMetadata> operand_side_metadata =
      MakeUnique<ShardingMetadata>(std::move(real_operand_sharding));
  std::unique_ptr<DomainMetadata> user_side_metadata =
      MakeUnique<ShardingMetadata>(std::move(real_instruction_sharding));
  return HloInstruction::CreateDomain(operand->shape(), operand,
                                      std::move(operand_side_metadata),
                                      std::move(user_side_metadata));
}

StatusOr<std::unique_ptr<HloSharding>> ExtractOriginalCommonSharding(
    tensorflow::gtl::ArraySlice<HloInstruction*> instructions) {
  // If we are here, all the instructions being passed had the same sharding
  // (or no sharding), by the means of the ShardingMatches() API.
  // As such, no kDomain was inserted, and here we are asked to extract the
  // original common sharding.
  // All the instructions passed to this API are part of the same computation.
  const HloSharding* sharding = nullptr;
  for (HloInstruction* instruction : instructions) {
    if (instruction->has_sharding()) {
      if (sharding == nullptr) {
        sharding = &instruction->sharding();
      } else {
        TF_RET_CHECK(ShardingMatches(*sharding, instruction->sharding()))
            << "Sharding " << *sharding << " does not match the one in "
            << instruction->ToString();
      }
    }
  }
  if (sharding == nullptr) {
    return std::unique_ptr<HloSharding>();
  }
  VLOG(4) << "Extracted sharding is " << *sharding;
  return CloneShardingForDomain(*sharding);
}

}  // namespace

std::unique_ptr<DomainMetadata> ShardingMetadata::Clone() const {
  std::unique_ptr<HloSharding> sharding;
  if (sharding_ != nullptr) {
    sharding = MakeUnique<HloSharding>(*sharding_);
  }
  return MakeUnique<ShardingMetadata>(std::move(sharding));
}

bool ShardingMetadata::Matches(const DomainMetadata& other) const {
  const ShardingMetadata* other_ptr =
      dynamic_cast<const ShardingMetadata*>(&other);
  if (other_ptr == nullptr) {
    // If other is not a ShardingMetadata, then it is clearly a no match.
    return false;
  }
  if (sharding_ == nullptr) {
    return other_ptr->sharding_ == nullptr;
  }
  return other_ptr->sharding_ != nullptr
             ? ShardingMatches(*sharding_, *other_ptr->sharding_)
             : false;
}

string ShardingMetadata::ToString() const {
  return sharding_ != nullptr ? sharding_->ToString() : "{}";
}

Status ShardingMetadata::NormalizeInstructions(
    const DomainMetadata::Domain& domain) const {
  if (sharding_ != nullptr) {
    VLOG(4) << "Normalizing sharding to " << sharding_->ToString() << ":";
    TF_RETURN_IF_ERROR(ApplyDomainSharding(domain, *sharding_));
    TF_RETURN_IF_ERROR(FixupPassThroughDomainLinks(domain, *sharding_));
  }
  return Status::OK();
}

Status NormalizeShardingDomain(const DomainMetadata::Domain& domain) {
  TF_ASSIGN_OR_RETURN(std::unique_ptr<HloSharding> sharding,
                      ExtractOriginalCommonSharding(domain.instructions));
  if (sharding != nullptr) {
    VLOG(4) << "Normalizing sharding-less domain to " << sharding->ToString()
            << ":";
    TF_RETURN_IF_ERROR(ApplyDomainSharding(domain, *sharding));
  } else {
    VLOG(1) << "Unable to find common sharding";
  }
  return Status::OK();
}

std::unique_ptr<HloInstruction> CreateShardingDomain(
    HloInstruction* instruction, HloInstruction* operand) {
  return CreateDomain(instruction, operand);
}

}  // namespace xla