android-12.0.0_r34/s

/* Copyright 2020 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/all_reduce_combiner.h"

#include <algorithm>
#include <list>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
#include "tensorflow/compiler/xla/service/hlo_domain_map.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_instructions.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/compiler/xla/service/hlo_query.h"
#include "tensorflow/compiler/xla/service/hlo_reachability.h"
#include "tensorflow/compiler/xla/service/shape_inference.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/types.h"

namespace xla {
namespace {

// Combines the elements of to_combine into a single AllReduce op. All
// entries in to_combine must be AllReduce ops with exactly one operand
// and the same reduction operation.
Status CombineAllReduces(absl::Span<HloInstruction* const> to_combine) {
  if (to_combine.size() < 2) {
    return Status::OK();
  }
  VLOG(1) << "Combined " << to_combine.size() << " CRS ops";

  HloComputation& computation = *to_combine.back()->parent();
  HloComputation* reduction = to_combine[0]->to_apply();
  const HloOpcode type = reduction->root_instruction()->opcode();

  // Create a single bigger AllReduce of the operands of the smaller
  // AllReduces.
  std::vector<HloInstruction*> operands;
  std::vector<Shape> operand_shapes;
  VLOG(1) << "Combining set";
  for (HloInstruction* hlo : to_combine) {
    VLOG(1) << "Set element: " << hlo->ToString();
    TF_RET_CHECK(hlo->opcode() == HloOpcode::kAllReduce);
    TF_RET_CHECK(hlo->operands().size() == 1);
    TF_RET_CHECK(hlo->to_apply() == reduction ||
                 (hlo->to_apply()->instruction_count() == 3 &&
                  hlo->to_apply()->num_parameters() == 2 &&
                  hlo->to_apply()->root_instruction()->opcode() == type));
    TF_RET_CHECK(hlo->shape().IsArray());
    for (HloInstruction* operand : hlo->operands()) {
      operands.push_back(operand);
      operand_shapes.push_back(operand->shape());
    }
  }

  HloInstruction* combined;
  // AllReduce ops with more than one operand produce a tuple.
  TF_RET_CHECK(operands.size() >= 2);
  combined = computation.AddInstruction(HloInstruction::CreateAllReduce(
      ShapeUtil::MakeTupleShape(operand_shapes), operands, reduction,
      to_combine.front()->replica_groups(),
      /*constrain_layout=*/false, to_combine.front()->channel_id(),
      Cast<HloAllReduceInstruction>(to_combine.front())
          ->use_global_device_ids()));

  // We have to propagate the sharding manually because Domain instructions are
  // not guaranteed to preserve it for side effecting instructions.
  if (to_combine.front()->has_sharding()) {
    combined->set_sharding(to_combine.front()->sharding());
  }
  VLOG(1) << "Replacing with : " << combined->ToString();

  // Replace all the smaller AllReduces with elements of the tuple output
  // of the single bigger AllReduce.
  for (int64 i = 0; i < to_combine.size(); ++i) {
    auto replace_with = HloInstruction::CreateGetTupleElement(
        to_combine[i]->shape(), combined, i);
    TF_RETURN_IF_ERROR(computation.ReplaceWithNewInstruction(
        to_combine[i], std::move(replace_with)));
  }
  return Status::OK();
}

struct GroupKey {
  GroupKey(const HloInstruction* hlo, const HloDomainMap& domain_map)
      : opcode(hlo->to_apply()->root_instruction()->opcode()),
        accum_type(hlo->to_apply()->root_instruction()->shape().element_type()),
        domain_id(domain_map.GetDomainMetadataId(hlo)),
        is_cross_shard(hlo->channel_id().has_value()),
        use_global_device_ids(
            Cast<HloAllReduceInstruction>(hlo)->use_global_device_ids()),
        replica_groups(hlo->replica_groups()) {}

  bool operator<(const GroupKey& other) const {
    if (opcode != other.opcode) {
      return opcode < other.opcode;
    }
    if (accum_type != other.accum_type) {
      return accum_type < other.accum_type;
    }
    if (domain_id != other.domain_id) {
      return domain_id < other.domain_id;
    }
    if (is_cross_shard != other.is_cross_shard) {
      return is_cross_shard < other.is_cross_shard;
    }
    if (use_global_device_ids != other.use_global_device_ids) {
      return use_global_device_ids < other.use_global_device_ids;
    }
    if (replica_groups.size() != other.replica_groups.size()) {
      return replica_groups.size() < other.replica_groups.size();
    }
    for (int64 i = 0; i < replica_groups.size(); ++i) {
      const auto& rg = replica_groups[i];
      const auto& org = other.replica_groups[i];
      if (rg.replica_ids_size() != org.replica_ids_size()) {
        return rg.replica_ids_size() < org.replica_ids_size();
      }
      for (int64 j = 0; j < rg.replica_ids_size(); ++j) {
        if (rg.replica_ids(j) != org.replica_ids(j)) {
          return rg.replica_ids(j) < org.replica_ids(j);
        }
      }
    }
    return false;
  }

  HloOpcode opcode;
  PrimitiveType accum_type;
  int64 domain_id;
  bool is_cross_shard;
  bool use_global_device_ids;
  std::vector<ReplicaGroup> replica_groups;
};

// Group AllReduce instructions by the reduction types, e.g., add, min,
// max, replica groups and domain. For cross-module all reduce instructions
// we group them by the set of domains they are reducing across.
//
// Note that the shape of the reduction computation is not included in the
// reduction types, e.g.: "f32[] add" and "bf16[] add" will be the same type. We
// need to disallow combining CRS instructions with different domain metadata as
// well as that could end up short-cutting two or more different domains.
//
// In each group, the instructions should be in post order. We will then iterate
// each group and try to combine them, so to prevent non-determinism, we use
// std::map here.
//
// The return value is a list of groups where every group contains a list of
// all-reduce instruction sets in topological order and with a deterministic
// order within the set. Additionally due to the above constraints every all
// reduce set within a group will contain the same number of elements
// and every instruction within an all reduce set will have the same
// all-reduce-id (if specified) and thus shape (all reduce sets without an
// all-reduce-id will have a single instruction).
using InstructionGroups =
    std::vector<std::vector<std::vector<HloInstruction*>>>;
StatusOr<InstructionGroups> CreateComputationGroups(
    HloComputation* computation) {
  TF_ASSIGN_OR_RETURN(auto domain_map, HloDomainMap::Create(computation, ""));

  // Group instructions by opcode, domain id and replica group.
  std::map<GroupKey, std::vector<HloInstruction*>> opcode_groups;
  for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) {
    if (instruction->opcode() != HloOpcode::kAllReduce) {
      continue;
    }
    if (instruction->to_apply()->instruction_count() != 3 ||
        instruction->to_apply()->num_parameters() != 2) {
      VLOG(1) << "Skipping due to non-trivial reduction function.";
      continue;
    }
    opcode_groups[GroupKey(instruction, *domain_map)].push_back(instruction);
  }

  // Generate a unique all-reduce-id for instructions without one by negating
  // the unique id of the hlo. This way we can treat cross module and normal CRS
  // instructions uniformly.
  auto channel_id = [](const HloInstruction* all_reduce) {
    return all_reduce->IsCrossModuleAllReduce()
               ? all_reduce->channel_id().value()
               : -1 * all_reduce->unique_id();
  };

  // Group instructions by all-reduce id with instructions for an all-reduce id
  // is listed along their group id and the (group id, instruction) pairs are
  // sorted by group id in the vector.
  std::map<int64, std::vector<std::pair<int64, HloInstruction*>>>
      all_reduce_sets;
  int64 group_id = 0;
  for (auto& domain_groups : opcode_groups) {
    for (HloInstruction* hlo : domain_groups.second) {
      all_reduce_sets[channel_id(hlo)].emplace_back(group_id, hlo);
    }
    ++group_id;
  }

  // Group instructions by participating group ids. Instructions within a group
  // are sorted by topological order and instructions within an all reduce group
  // is still sorted by group id.
  std::map<std::vector<int64>, std::vector<std::vector<HloInstruction*>>>
      all_reduce_group_map;
  for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) {
    if (instruction->opcode() != HloOpcode::kAllReduce) {
      continue;
    }
    if (instruction->to_apply()->instruction_count() != 3 ||
        instruction->to_apply()->num_parameters() != 2) {
      VLOG(1) << "Skipping due to non-trivial reduction function.";
      continue;
    }

    int64 arid = channel_id(instruction);
    if (all_reduce_sets.count(arid) == 0) {
      // Already processed.
      continue;
    }

    std::vector<int64> group_ids;
    std::vector<HloInstruction*> instructions;
    for (const auto& hlo : all_reduce_sets[arid]) {
      group_ids.push_back(hlo.first);
      instructions.push_back(hlo.second);
    }
    all_reduce_group_map[group_ids].push_back(std::move(instructions));
    all_reduce_sets.erase(arid);
  }
  CHECK(all_reduce_sets.empty());

  InstructionGroups groups;
  for (const auto& all_reduce_group : all_reduce_group_map) {
    groups.push_back(all_reduce_group.second);
  }
  return std::move(groups);
}

}  // namespace

AllReduceCombiner::AllReduceCombiner(int64 combine_threshold_in_bytes,
                                     int64 combine_threshold_count)
    : combine_threshold_in_bytes_(combine_threshold_in_bytes),
      combine_threshold_count_(combine_threshold_count) {}

StatusOr<bool> AllReduceCombiner::Run(HloModule* module) {
  VLOG(1) << "Running AllReduceCombiner with threshold of "
          << combine_threshold_in_bytes_ << " bytes";

  if (combine_threshold_in_bytes_ <= 0 || combine_threshold_count_ <= 0) {
    VLOG(1) << "Skip AllReduceCombiner because the threshold is zero";
    return false;
  }

  if (hlo_query::ContainsLayoutConstrainedAllReduce(*module)) {
    VLOG(1) << "Skip AllReduceCombiner because the module contains all-reduce "
               "with constrained layouts";
    return false;
  }

  bool changed = false;
  for (HloComputation* computation : module->MakeNonfusionComputations()) {
    TF_ASSIGN_OR_RETURN(auto groups, CreateComputationGroups(computation));
    for (auto group : groups) {
      // Recompute reachability after every combine group because we can't
      // maintain a cross group topolgical order to be able to rely on the
      // transitive dependencies to detect cycles.
      auto reachability = HloReachabilityMap::Build(computation);

      // Create a map to be able to find an instruction group based on the first
      // instruction in the group. It will be used during the post order
      // iteration to be able to process full groups at a time. Doing it only
      // for one instruction in every group will be sufficient because all
      // instruction have to schedule at the same time due to cross core
      // dependencies.
      absl::flat_hash_map<HloInstruction*, std::vector<HloInstruction*>*>
          group_map;
      for (auto& instruction : group) {
        group_map[instruction.front()] = &instruction;
      }

      // Collect sets of AllReduce instructions to combine.
      std::vector<std::vector<std::vector<HloInstruction*>>> combine_sets(1);
      int64 current_size_in_bytes = 0;
      int64 current_operand_count = 0;

      // Iterate all instructions in post order and skip the ones not in the
      // current group. We have to create a new post order iteration for every
      // group because merging instructions in the previous group can made the
      // original post order no longer hold.
      // This will make it likely that we won't increase memory pressure much
      // above combine_threshold_in_bytes, since two AllReduces that are
      // near in post order are most likely, but not for sure, also near in
      // scheduled order.
      //
      // TODO(b/70235266): This should usually be fine, but it's probably
      // possible to construct some case where the memory usage increases beyond
      // the threshold due to reordering of the instructions in scheduling. If
      // this ever comes up as a real problem, it would be nice to implement
      // safeguards so that that cannot possibly happen.
      for (const HloInstruction* inst :
           computation->MakeInstructionPostOrder()) {
        auto it = group_map.find(inst);
        if (it == group_map.end()) {
          // Instruction belongs to a different group.
          continue;
        }
        const auto& instructions = *it->second;

        VLOG(1) << "Considering HLO " << instructions.front()->ToString()
                << " with current set size of " << current_size_in_bytes
                << " and current operand count of " << current_operand_count;

        // We do not handle AllReduce ops that do not have exactly 1
        // operand since that is simpler and this pass is the only way to
        // generate such ops and it should rarely be important to consider the
        // same ops again.
        if (instructions.front()->operands().size() != 1) {
          VLOG(1) << "Skipping due to "
                  << instructions.front()->operands().size() << " operands";
          continue;
        }

        int64 size_in_bytes;
        TF_RET_CHECK(instructions.front()->shape().IsArray());
        size_in_bytes = ShapeUtil::ByteSizeOf(instructions.front()->shape());

        if (size_in_bytes > combine_threshold_in_bytes_) {
          VLOG(1) << "Skipping due to size " << size_in_bytes
                  << " above threshold";
          // If the instruction is greather than the threshold, then we can
          // never combine it with anything.
          continue;
        }

        // If the current set is dependent on the instruction, then create a new
        // one to avoid the dependency. We move on from the current set instead
        // of ignoring the instruction since otherwise a single AllReduce
        // instruction that all the other ones depend on (such as one on the
        // forward pass of a model) could disable this optimization entirely.
        TF_RET_CHECK(!combine_sets.empty());
        for (const auto& previous : combine_sets.back()) {
          // The reachability information does not reflect the planned
          // combination from combine_sets. We cannot just bring it up to date
          // cheaply since HloReachabilityMap does not track reachability
          // updates transitively and doing it directly is expensive. However,
          // leaving it stale has no effect on the reachability queries that we
          // are doing here because we are considering the ops in a topological
          // order, so we can just leave it stale.
          //
          // Proof: Suppose A is the instruction we are looking to combine and B
          // is an element of the current combine set that we are looking to
          // combine A into.
          //
          // First of all, we check that all elements in each set do not depend
          // on each other, so combining the *current* combine set cannot create
          // new dependencies between A and B. It remains to prove that
          // combining the prior combine sets also cannot create a dependency
          // between A and B.
          //
          // Assume to get a contradiction that there are two AllReduce
          // ops C and D in combine_sets that will be combined and that A and B
          // are not connected now but that they will be after combining C and
          // D. Then there exist paths in the dependency graph such that one of
          // these cases is true:
          //
          //   A -> ... -> C and D -> ... -> B
          //   A -> ... -> D and C -> ... -> B
          //   B -> ... -> C and D -> ... -> A
          //   B -> ... -> D and C -> ... -> A
          //
          // None of these cases are possible because we are visiting the nodes
          // in a topological order, so C and D cannot be in-between A and B.
          // That is a contradiction, so combining the prior combine sets also
          // cannot create a dependency between A and B.
          bool new_set = false;
          for (int64 i = 0; i < instructions.size(); ++i) {
            if (reachability->IsReachable(previous[i], instructions[i])) {
              VLOG(1) << "Starting new set due to dependency between "
                      << previous[i]->ToString() << " AND "
                      << instructions[i]->ToString();
              new_set = true;
              break;
            }
          }
          if (new_set) {
            combine_sets.emplace_back();
            current_size_in_bytes = 0;
            current_operand_count = 0;
            break;
          }
        }

        if (current_size_in_bytes + size_in_bytes >
                combine_threshold_in_bytes_ ||
            current_operand_count + 1 > combine_threshold_count_) {
          VLOG(1) << "The instruction cannot be entered into the set due "
                     "to the combined size being too large.";
          // In this case we cannot include the instruction into the current set
          // since then it would grow beyond the threshold. The set of
          // instructions to carry forward will either be the current set or the
          // instruction by itself, whichever is smaller, since that maximizes
          // the chance of being able to combine with the next instruction.
          if (size_in_bytes > current_size_in_bytes) {
            VLOG(1) << "Skipping as the instruction is larger than the set.";
            continue;  // keep the current set
          }
          VLOG(1)
              << "Resetting the set as the set is larger than the instruction.";
          combine_sets.emplace_back();
          current_size_in_bytes = 0;
          current_operand_count = 0;
        }

        VLOG(1) << "Adding instruction to set.";
        combine_sets.back().push_back(instructions);
        current_size_in_bytes += size_in_bytes;
        current_operand_count += 1;
        TF_RET_CHECK(current_size_in_bytes <= combine_threshold_in_bytes_);
        TF_RET_CHECK(current_operand_count <= combine_threshold_count_);
      }
      VLOG(1) << "Done constructing sets. Final set size is "
              << current_size_in_bytes << " bytes and " << current_operand_count
              << " operands";

      // Combine the collected sets of AllReduce instructions.
      for (const auto& combine_set : combine_sets) {
        if (combine_set.size() >= 2) {
          changed = true;
          for (int64 i = 0; i < combine_set.front().size(); ++i) {
            std::vector<HloInstruction*> to_combine;
            to_combine.reserve(combine_set.size());
            for (const auto& c : combine_set) {
              to_combine.push_back(c[i]);
            }
            TF_RETURN_IF_ERROR(CombineAllReduces(to_combine));
          }
        }
      }
    }
  }

  return changed;
}

}  // namespace xla