xref: /external/tensorflow/tensorflow/compiler/xla/service/conditional_code_motion.cc

/* Copyright 2017 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/conditional_code_motion.h"

#include <iterator>
#include <stack>
#include <string>
#include <utility>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/xla/debug_options_flags.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/service/call_graph.h"
#include "tensorflow/compiler/xla/service/call_inliner.h"
#include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_cse.h"
#include "tensorflow/compiler/xla/service/hlo_dce.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_pass_pipeline.h"
#include "tensorflow/compiler/xla/service/hlo_verifier.h"
#include "tensorflow/compiler/xla/service/tuple_simplifier.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/errors.h"

namespace xla {

namespace conditional_opt {

class BoundaryVisitor {
 public:
  // start with an existing conditional computation.
  explicit BoundaryVisitor(HloInstruction* conditional) {
    Boundary b(Boundary::Position::kInsideBranch);
    b.mutable_operands().push_back(conditional);
    worklist_.push_back(b);
  }
  // Start with an empty work list.
  BoundaryVisitor() {}
  // Get next boundary to visit.
  Boundary PopNextBoundary() {
    CHECK(!worklist_.empty());
    Boundary b = worklist_.front();
    worklist_.pop_front();
    // if b is already visited, it must have multiple users and is already in
    // new boundaries. Skip it. Only checking the first operand of b because b
    // is expected to have at least one operand, and all the operands in b
    // must be identical instructions from different branches for b to be moved.
    while (!worklist_.empty() && ContainsKey(visited_, b.operands()[0])) {
      b = worklist_.front();
      worklist_.pop_front();
    }
    visited_.insert(b.operands()[0]);
    return b;
  }
  void AddToWorkList(const Boundary& b) {
    CHECK(!b.operands().empty());
    worklist_.push_back(b);
  }

  bool HasNextBoundary() {
    while (!worklist_.empty()) {
      Boundary b = worklist_.front();
      if (!ContainsKey(visited_, b.operands()[0])) {
        break;
      }
      worklist_.pop_front();
    }
    return !worklist_.empty();
  }

 private:
  // worklist is the deque that contains instructions to be visited.
  std::deque<Boundary> worklist_;
  absl::flat_hash_set<HloInstruction*> visited_;
};

template <class OpCollection>
int64 CountNonLeafOps(const OpCollection& ops) {
  absl::flat_hash_set<HloInstruction*> op_set;
  for (auto op : ops) {
    if (!op_set.contains(op) && op->opcode() != HloOpcode::kConstant) {
      op_set.insert(op);
    }
  }
  return op_set.size();
}

// Returns estimation of potential reuses carried by a given pair of
// instructions.  Use different integers to classify different levels
// of reuses This is used as a placeholder only, assuming all
// instructions can be fused to enable data reuses
int64 ReusesCarriedBy(HloOpcode op, HloOpcode user) {
  // Reuses in some way work like forces that pull instructions
  // towards each other. We use a number 0-10 to classify how strong the force
  // is between a pair of operations. Given a group of instructions that can be
  // moved together, if the forces inside a conditional are stronger, the group
  // will be moved incide or remain inside the conditional; otherwise, it will
  // be moved outside to or remain outside of the conditional.
  switch (user) {
    case HloOpcode::kGetTupleElement:
      return 0;
    case HloOpcode::kConvert:
      // Because convert is treated not moveable when following Dot or
      // convolution, here if op is dot or convolution, they must be separated
      // by a conditional boundary. Here we do not try to pull convert inside
      // conditionals to be together with the dot or convolution.
      switch (op) {
        case HloOpcode::kConvolution:
        case HloOpcode::kDot:
          return 0;
        default:
          break;
      }
      break;
    default:
      break;
  }
  switch (op) {
      // These instructions do not carry weight of reuse themselves.
    case HloOpcode::kParameter:
    case HloOpcode::kConstant:
    case HloOpcode::kGetTupleElement:
      return 0;
    case HloOpcode::kConditional:
      return 10;
    default:
      return -10;
  }
}

// Returns true if `op` is worth hoisting.
bool WorthHoisting(HloOpcode op, HloOpcode child_op) {
  // TOOD[b/169182921] The following cost model is rather incomplete. Will
  // need to extend to cover most of element-wise ops.
  switch (op) {
    case HloOpcode::kConvert:
      // If Convert is after AllReduce, it is worth moving out AllReduce
      // out of conditional for AR/CRS combine. If Convert is after other
      // ops such as Dot or Convolutional, it is better to keep convert
      // within conditional so that convert can be fused with Dot or
      // Convolutional.
      switch (child_op) {
        case HloOpcode::kAllReduce:
        case HloOpcode::kReshape:
        case HloOpcode::kGetTupleElement:
          return true;
        default:
          return false;
      }
    case HloOpcode::kGetTupleElement:
      switch (child_op) {
        // do not move GTE if its operand is a parameter
        case HloOpcode::kParameter:
          return false;
        default:
          return true;
      }
    case HloOpcode::kAllReduce:
    case HloOpcode::kAbs:
    case HloOpcode::kReduce:
    case HloOpcode::kAdd:
    case HloOpcode::kPower:
    case HloOpcode::kCopy:
    case HloOpcode::kConstant:
    case HloOpcode::kSubtract:
    case HloOpcode::kMultiply:
    case HloOpcode::kDivide:
    case HloOpcode::kTuple:
    case HloOpcode::kSqrt:
    case HloOpcode::kRsqrt:
    case HloOpcode::kReshape:
    case HloOpcode::kMinimum:
    case HloOpcode::kMaximum:
      return true;
    default:
      return false;
  }
}

// Compare if the instructions to be visited at each branches are identical.
bool InstructionWithinBranchIdentical(
    const std::vector<HloInstruction*>& instructions,
    bool is_layout_sensitive) {
  // Identical includes the shape of each operands are equal.
  auto eq_operand = [&](const HloInstruction* a, const HloInstruction* b) {
    bool eq_operands = is_layout_sensitive
                           ? ShapeUtil::Equal(a->shape(), b->shape())
                           : ShapeUtil::Compatible(a->shape(), b->shape());
    return eq_operands;
  };

  auto eq_computations = [](const HloComputation* a, const HloComputation* b) {
    return *a == *b;
  };

  if (instructions.empty()) {
    return false;
  }

  if (instructions[0]->IsCrossModuleAllReduce()) {
    return std::all_of(
        instructions.begin(), instructions.end(),
        [&](HloInstruction* instruction) {
          if (!instruction->IsCrossModuleAllReduce()) {
            return false;
          }
          auto old_channel_id = instruction->channel_id();
          instruction->set_channel_id(instructions[0]->channel_id());
          bool eq_instructions = instructions[0]->Identical(
              *instruction, eq_operand, eq_computations, is_layout_sensitive);
          instruction->set_channel_id(old_channel_id);
          return eq_instructions;
        });
  }

  return std::all_of(instructions.begin(), instructions.end(),
                     [&](HloInstruction* instruction) {
                       return instructions[0]->Identical(
                           *instruction, eq_operand, eq_computations,
                           is_layout_sensitive);
                     });
}

// Copy the ith instruction in boundary to outside of conditional, or do the
// opposite (for moving in).
Status CopyInOrOutOfConditional(
    Boundary& boundary, int64 dest_index, HloComputation* parent,
    absl::flat_hash_map<HloInstruction*, Boundary>& hoisted_instructions) {
  CHECK(dest_index == 0 || boundary.IsOutsideBranch());
  HloInstruction* op = boundary.operands()[0];
  absl::InlinedVector<HloInstruction*, 4> new_operands;
  for (int i = 0; i < op->operands().size(); ++i) {
    auto op_i = op->operands()[i];
    VLOG(2) << "Looking for " << op_i->ToString() << "\n";
    if (ContainsKey(hoisted_instructions, op_i)) {
      auto new_op_i =
          FindOrDie(hoisted_instructions, op_i).operands()[dest_index];
      VLOG(2) << "new instruction:" << new_op_i->ToString() << "\n";
      new_operands.push_back(new_op_i);
    } else {
      switch (op_i->opcode()) {
        case HloOpcode::kConstant: {
          auto new_op_i = parent->AddInstruction(op_i->Clone());
          VLOG(2) << "new instruction:" << new_op_i->ToString() << "\n";
          new_operands.push_back(new_op_i);
          break;
        }
        case HloOpcode::kGetTupleElement: {
          auto gte = Cast<HloGetTupleElementInstruction>(op_i);
          int64 index = gte->tuple_index();
          HloInstruction* root = parent->root_instruction();
          CHECK(root->opcode() == HloOpcode::kTuple &&
                index < root->operand_count());
          auto new_op_i = root->mutable_operand(index);
          VLOG(2) << "new instruction:" << new_op_i->ToString() << "\n";
          new_operands.push_back(new_op_i);
          break;
        }
        default:
          LOG(FATAL) << "Unexpected out-of-boundary instruction:"
                     << op_i->ToString() << "\n";
      }
    }
  }
  HloInstruction* new_instruction = parent->AddInstruction(
      op->CloneWithNewOperands(op->shape(), new_operands));
  VLOG(2) << "new instruction:" << new_instruction->ToString() << "\n";
  // Maps the instruction outside of conditional to the instruction
  // inside of the conditional.
  for (HloInstruction* op : boundary.operands()) {
    Boundary b2 = ContainsKey(hoisted_instructions, op)
                      ? hoisted_instructions[op]
                      : Boundary(boundary.IsOutsideBranch()
                                     ? Boundary::Position::kInsideBranch
                                     : Boundary::Position::kOutsideBranch);
    b2.mutable_operands().push_back(new_instruction);
    hoisted_instructions[op] = b2;
  }
  return Status::OK();
}

// Identify converts to be hoisted/rematerialized out of the branch
// computations.
absl::flat_hash_set<int64> FindSpecialConverts(HloInstruction* old_root,
                                               int branch_count,
                                               HloInstruction* conditional,
                                               bool is_layout_sensitive) {
  absl::flat_hash_set<int64> kspecial_convert;
  for (int64 operand_num = 0; operand_num < old_root->operand_count();
       ++operand_num) {
    if (old_root->operand(operand_num)->opcode() != HloOpcode::kConvert) {
      continue;
    }
    bool replica = true;
    HloInstruction* kspecial_convert_candidate =
        old_root->mutable_operand(operand_num);
    // Check whether an identical candidate appears in other branches
    for (int others = 1; others < branch_count; ++others) {
      HloInstruction* others_root =
          conditional->branch_computation(others)->root_instruction();
      bool eq_shape =
          is_layout_sensitive
              ? ShapeUtil::Equal(others_root->operand(operand_num)->shape(),
                                 kspecial_convert_candidate->shape())
              : ShapeUtil::Compatible(
                    others_root->operand(operand_num)->shape(),
                    kspecial_convert_candidate->shape());
      if ((others_root->operand(operand_num)->opcode() ==
           HloOpcode::kConvert) &&
          eq_shape) {
        // Nothing to be done.
      } else {
        replica = false;
        break;
      }
    }
    if (replica) {
      kspecial_convert.insert(operand_num);
    }
  }
  return kspecial_convert;
}

// Restructuring the conditional instruction as follows:
// i.e., %result = conditional() becomes
// x = conditional()
// y.{0..n} = gte(x, {0..n})
// z = tuple(y.0, y.1, ...y.n)
// Doing so ensures that we can accommodate the possible shape-change of the
// conditional when the instructions are hoisted.
Status RestructureConditionalInstruction(HloComputation* computation,
                                         HloInstruction* conditional) {
  HloInstruction* old_root = computation->root_instruction();
  std::vector<HloInstruction*> new_operands;
  int cur_index = 0;
  for (; cur_index < ShapeUtil::TupleElementCount(conditional->shape());
       ++cur_index) {
    new_operands.push_back(
        computation->AddInstruction(HloInstruction::CreateGetTupleElement(
            ShapeUtil::GetTupleElementShape(conditional->shape(), cur_index),
            conditional, cur_index)));
  }
  HloInstruction* new_tuple =
      computation->AddInstruction(HloInstruction::CreateTuple(new_operands));
  if (old_root == conditional) {
    computation->set_root_instruction(new_tuple);
  } else {
    std::vector<HloInstruction*> new_tuple_users;
    for (auto conditional_user : conditional->users()) {
      auto is_new_gte = absl::c_find_if(
          new_operands,
          [&](HloInstruction* instr) { return instr == conditional_user; });
      if (is_new_gte == new_operands.end()) {
        new_tuple_users.push_back(conditional_user);
      }
    }
    for (auto new_tuple_user : new_tuple_users) {
      TF_RETURN_IF_ERROR(
          conditional->ReplaceUseWith(new_tuple_user, new_tuple));
    }
  }
  VLOG(2) << "computation after root restructure:\n" << computation->ToString();
  return Status::OK();
}

StatusOr<bool> ConvertSpecialMove(HloInstruction* conditional,
                                  bool is_layout_sensitive) {
  int branch_count = conditional->branch_count();
  if (branch_count <= 0) {
    return false;
  }

  // Determining whether all branch roots are tuples
  for (int branch_num = 0; branch_num < branch_count; ++branch_num) {
    HloInstruction* branch_root =
        conditional->branch_computation(branch_num)->root_instruction();
    if (branch_root->opcode() != HloOpcode::kTuple) {
      return false;
    }
  }

  HloInstruction* old_root =
      conditional->branch_computation(0)->root_instruction();
  VLOG(2) << "BEFORE :" << conditional->parent()->parent()->ToString();
  // Identify the gte using `index'.
  auto find_gte = [](const HloInstruction* conditional_result,
                     int64 index) -> HloInstruction* {
    for (HloInstruction* instr : conditional_result->users()) {
      if (instr->opcode() != HloOpcode::kGetTupleElement) {
        return nullptr;
      }
      if (instr->tuple_index() == index) {
        return instr;
      }
    }
    return nullptr;
  };

  // Captures tuple indices refering to converts to be rematerialized/hoisted.
  absl::flat_hash_set<int64> kspecial_convert = FindSpecialConverts(
      old_root, branch_count, conditional, is_layout_sensitive);

  // Exit if we cannot find any converts to be hoisted.
  if (kspecial_convert.empty()) {
    return false;
  }

  TF_RETURN_IF_ERROR(
      RestructureConditionalInstruction(conditional->parent(), conditional));

  for (int branch = 0; branch < branch_count; branch++) {
    old_root = conditional->branch_computation(branch)->root_instruction();
    absl::flat_hash_map<HloInstruction*, int64> map_inst_to_tuple_index;
    std::vector<HloInstruction*> new_operands(old_root->operand_count());
    absl::flat_hash_set<HloInstruction*> to_hoist_set;

    for (int64 operand_num = 0; operand_num < old_root->operand_count();
         ++operand_num) {
      map_inst_to_tuple_index[old_root->mutable_operand(operand_num)] =
          operand_num;
    }
    for (int64 operand_num = 0; operand_num < old_root->operand_count();
         ++operand_num) {
      HloInstruction* hoist = old_root->mutable_operand(operand_num);
      if (!kspecial_convert.contains(operand_num)) {
        new_operands[operand_num] = old_root->mutable_operand(operand_num);
        continue;
      }

      to_hoist_set.insert(hoist);
      int64 new_tuple_count = old_root->operand_count();

      // Replace the hoisted instr in the tuple with the operand/operands.
      // We will replace at least one of the operands of the hoist at the
      // tuple place; the rest will be added at the end.
      bool inplace = true;
      CHECK(!hoist->operands().empty());
      for (HloInstruction* prod : hoist->operands()) {
        if (inplace) {
          map_inst_to_tuple_index[prod] = map_inst_to_tuple_index[hoist];
          new_operands[map_inst_to_tuple_index[hoist]] = prod;
          inplace = false;
        } else {
          map_inst_to_tuple_index[prod] = new_tuple_count++;
          new_operands.push_back(prod);
        }
      }
    }

    // Create the new root instruction.
    HloComputation* cur_branch = conditional->branch_computation(branch);
    HloInstruction* new_branch_root =
        cur_branch->AddInstruction(HloInstruction::CreateTuple(new_operands));
    // The shape can vary since the operands to convert are now
    // being returned through the branches' root.
    cur_branch->set_root_instruction(new_branch_root, true /*new shape*/);
    TF_CHECK_OK(cur_branch->RemoveInstruction(old_root));

    // Only one of the branches needs to change the conditional->parent().
    if (branch != 0) {
      continue;
    }
    HloComputation* conditional_parent = conditional->parent();
    HloInstruction* newconditional =
        conditional_parent->AddInstruction(HloInstruction::CreateConditional(
            cur_branch->root_instruction()->shape(),
            conditional->mutable_operand(0),
            absl::MakeSpan(conditional->branch_computations()),
            absl::MakeSpan(conditional->operands()).subspan(1)));
    // Ensure that all the users of conditional refer to the new one.
    TF_RETURN_IF_ERROR(
        conditional->ReplaceAllUsesWithDifferentShape(newconditional));
    TF_CHECK_OK(conditional_parent->RemoveInstruction(conditional));
    conditional = newconditional;
    // Add the hoisted instructions in the parent.
    for (HloInstruction* hoist : to_hoist_set) {
      VLOG(2) << "Hoisting instruction:" << hoist->ToString();
      int64 hoist_index = map_inst_to_tuple_index[hoist];
      // Find out the gte that captured the hoisted instr result.
      HloInstruction* gte_hoist = find_gte(conditional, hoist_index);
      CHECK(gte_hoist != nullptr);
      std::vector<HloInstruction*> new_operands;
      for (HloInstruction* op : hoist->operands()) {
        HloInstruction* gte = conditional_parent->AddInstruction(
            HloInstruction::CreateGetTupleElement(op->shape(), conditional,
                                                  map_inst_to_tuple_index[op]));
        new_operands.push_back(gte);
      }
      HloInstruction* hoisted = conditional_parent->AddInstruction(
          hoist->CloneWithNewOperands(hoist->shape(), new_operands));
      VLOG(2) << "Hoisted instruction in parent:" << hoisted->ToString();
      TF_RETURN_IF_ERROR(gte_hoist->ReplaceAllUsesWith(hoisted));
      TF_CHECK_OK(conditional_parent->RemoveInstruction(gte_hoist));
    }
    // No need to explicitly delete a hoisted instruction since if its dead
    // then the subsequent DCE will remove it.
  }
  VLOG(2) << "AFTER :" << conditional->parent()->parent()->ToString();
  return true;
}

// Hoist identical ops out of the conditional. The definition of identical
// are the shape of the operands are identical and their properties are
// identical. Will start from the root instruction of each branch and get
// the identical ops to hoist.
StatusOr<bool> ConditionalCodeMotion::MoveInstructionOut(
    HloInstruction* conditional, std::vector<Boundary>& to_move_out,
    std::vector<Boundary>& new_boundaries) {
  if (to_move_out.empty()) {
    return false;
  }
  VLOG(1) << "Modifying code--number of boundaries to move out:"
          << to_move_out.size() << "\n";
  HloComputation* conditional_parent = conditional->parent();
  // save the old users before add new conditional user instructions
  std::vector<HloInstruction*> old_conditional_users = conditional->users();
  // Maps instructions in the conditional body to instructions hoisted outside
  // the conditional that compute the same value.
  absl::flat_hash_map<HloInstruction*, Boundary> hoisted_instructions;
  // Insert GetTupleElement before the instructions whose operands might still
  // be within the conditional.
  VLOG(1) << "before opt:"
          << conditional_parent->ToString(HloPrintOptions::Fingerprint())
          << "\n";
  int64 op_index = 0;
  for (const Boundary& b : new_boundaries) {
    HloInstruction* op = b.operands()[0];
    CHECK(op != nullptr);
    VLOG(2) << "Mapping new boundary instr: " << op->ToString() << "\n";
    HloInstruction* gtr = conditional_parent->AddInstruction(
        HloInstruction::CreateGetTupleElement(op->shape(), conditional,
                                              op_index++));
    Boundary b2(Boundary::Position::kOutsideBranch);
    b2.mutable_operands().push_back(gtr);
    hoisted_instructions[op] = b2;
  }
  // Copy boundary instructions out of the conditional.
  // Visit the operands before its users and copy it, so that the copied
  // user will point to the correct operand.
  for (int64 i = to_move_out.size() - 1; i >= 0; i--) {
    TF_RETURN_IF_ERROR(CopyInOrOutOfConditional(
        to_move_out[i], 0, conditional_parent, hoisted_instructions));
  }
  VLOG(2) << "Done copy branch instructions out\n"
          << conditional_parent->ToString(HloPrintOptions::Fingerprint())
          << "\n";
  // Change original users of the conditional to use the correct operands.
  HloInstruction* old_root =
      conditional->branch_computation(0)->root_instruction();
  for (auto user_instr : old_conditional_users) {
    VLOG(2) << "Checking conditional user: " << user_instr->ToString() << "\n";
    CHECK(user_instr->opcode() == HloOpcode::kGetTupleElement);
    auto tuple_opd = static_cast<HloGetTupleElementInstruction*>(user_instr);
    int64 index = tuple_opd->tuple_index();
    CHECK(old_root->operands().size() > index);
    HloInstruction* old_opd = old_root->operands()[index];
    VLOG(2) << "old opd = " << old_opd << "\n";
    CHECK(ContainsKey(hoisted_instructions, old_opd));
    HloInstruction* new_opd = hoisted_instructions[old_opd].operands()[0];
    CHECK(old_opd != nullptr);
    CHECK(new_opd != nullptr);
    VLOG(2) << "Try replace all uses of :" << old_opd->ToString() << "\n";
    TF_RETURN_IF_ERROR(user_instr->ReplaceAllUsesWith(new_opd));
    TF_RETURN_IF_ERROR(conditional_parent->RemoveInstruction(user_instr));
  }
  VLOG(2) << "Done changing conditional users\n"
          << conditional_parent->ToString() << "\n";
  // Create tuple element within each branch and set it as root.
  int64 branch_count = conditional->branch_count();
  for (int i = 0; i < branch_count; i++) {
    auto computation = conditional->branch_computation(i);
    std::vector<HloInstruction*> elements;
    for (const auto& b1 : new_boundaries) {
      HloInstruction* op = b1.operands()[i];
      CHECK(op != nullptr);
      VLOG(2) << "Adding to root " << i << " with " << op->ToString() << "\n";
      elements.push_back(op);
    }
    HloInstruction* tuple =
        computation->AddInstruction(HloInstruction::CreateTuple(elements));
    computation->set_root_instruction(tuple, true);
    VLOG(2) << "computation is :" << computation->ToString() << "\n";
    // Remove hoisted instructions from the branches.
    for (const auto& b2 : to_move_out) {
      auto instr_to_remove = b2.operands()[i];
      // Double check to make sure it is safe to delete the instruction.
      // Complications may arise due to some operations in the alternative
      // branches (branches 1..n) being placed into the boundaries multiple
      // times.
      if (!computation->IsMarkedAsDead(instr_to_remove) &&
          instr_to_remove->user_count() == 0) {
        VLOG(2) << "Removing boundary:" << b2.ToString() << "\n";
        TF_RETURN_IF_ERROR(computation->RemoveInstruction(instr_to_remove));
      }
    }
  }
  // Change conditional instruction shape to the shape of the new root.
  HloInstruction* new_root =
      conditional->branch_computation(0)->root_instruction();
  *conditional->mutable_shape() = new_root->shape();
  VLOG(1) << "done moving instructions out of branches\n"
          << conditional_parent->ToString(HloPrintOptions::Fingerprint())
          << "\n";
  return true;
}

// Hoist ops from outside of the conditional to inside the branches.
StatusOr<bool> ConditionalCodeMotion::MoveInstructionIn(
    HloInstruction* conditional, std::vector<Boundary>& to_move_in,
    std::vector<Boundary>& new_boundaries) {
  if (to_move_in.empty()) {
    return false;
  }
  VLOG(1) << "Modifying code---number of boundaries to move in:"
          << to_move_in.size() << "\n";
  VLOG(1) << "before opt:"
          << conditional->parent()->ToString(HloPrintOptions::Fingerprint())
          << "\n";
  // Mapping instructions to be moved to their new representations.
  absl::flat_hash_map<HloInstruction*, Boundary> hoisted_instructions;
  int64 to_move_in_size = to_move_in.size();
  int64 branch_count = conditional->branch_count();
  HloGetTupleElementInstruction* tuple_use =
      DynCast<HloGetTupleElementInstruction>(to_move_in[0].operands()[0]);
  // If use_index is -1, the old conditional root entry used by to_move_in
  // instructions still need to be included as an entry of the modified
  // conditional root, and the new result of the to_move_in instructions
  // need to be added as an extra entry of the modified root; otherwise, the
  // old root entry will be replaced with the new result in the modified root.
  // The entry replacement should be allowed only if tuple_use has <=1 users.
  int64 use_index = (tuple_use != nullptr && tuple_use->user_count() == 1)
                        ? tuple_use->tuple_index()
                        : -1;
  VLOG(2) << "Tuple use index = " << use_index << "\n";
  // Number of old conditional entries still to be used outside.
  // If conditional shape is not tuple, will create a tuple and use subscript
  // 0 to save the old operand being used.
  int64 op_index =
      conditional->shape().IsTuple()
          ? ((use_index >= 0) ? conditional->shape().tuple_shapes_size() - 1
                              : conditional->shape().tuple_shapes_size())
          : 0;
  // Use to map the tuple_use instruction to its operand;
  Boundary b_opd_use(Boundary::Position::kInsideBranch);
  Boundary b_old_root(Boundary::Position::kInsideBranch);
  // Create a new root instruction in each branch.
  for (int i = 0; i < branch_count; i++) {
    auto computation = conditional->branch_computation(i);
    auto old_root = computation->root_instruction();
    b_old_root.mutable_operands().push_back(old_root);
    std::vector<HloInstruction*> operands;
    if (old_root->opcode() == HloOpcode::kTuple) {
      // Use operands of old_root directly, so old_root can be removed later.
      for (int i = 0; i < old_root->operand_count(); ++i) {
        if (i != use_index) {
          operands.push_back(old_root->operands()[i]);
        } else {  // Map conditional use to the tuple operand.
          b_opd_use.mutable_operands().push_back(old_root->operands()[i]);
        }
      }
    } else if (old_root->shape().IsTuple()) {
      // If old_root is not a kTuple but has tuple shape, elements within the
      // tuple must be extracted first to be used by the new instructions.
      const Shape& old_shape = old_root->shape();
      for (int64 i = 0; i < old_shape.tuple_shapes_size(); ++i) {
        auto element =
            computation->AddInstruction(HloInstruction::CreateGetTupleElement(
                old_shape.tuple_shapes(i), old_root, i));
        if (i != use_index) {
          operands.push_back(element);
        } else {
          b_opd_use.mutable_operands().push_back(element);
        }
      }
    } else {
      // If old_root is not a tuple and does not have tuple shape, use it
      // to replace the conditional directly in the new computation.
      b_opd_use.mutable_operands().push_back(conditional);
    }

    HloInstruction* new_root =
        computation->AddInstruction(HloInstruction::CreateTuple(operands));
    VLOG(2) << "setting new root: " << new_root->ToString() << "\n";
    computation->set_root_instruction(new_root,
                                      /*accept_different_shape*/ true);
    if (old_root->opcode() == HloOpcode::kTuple) {
      TF_RETURN_IF_ERROR(computation->RemoveInstruction(old_root));
    }
    VLOG(2) << "new branch computation: " << computation->ToString() << "\n";
  }
  // Update get tuple element index of the conditional.
  if (use_index != -1) {
    for (auto* user : conditional->users()) {
      if (user->opcode() == HloOpcode::kGetTupleElement &&
          user->tuple_index() > use_index) {
        user->set_tuple_index(user->tuple_index() - 1);
      }
    }
  }
  hoisted_instructions[conditional] = b_old_root;
  int64 cp_start = 0;
  if (use_index >= 0) {
    VLOG(2) << "Mapping GTE: " << tuple_use->ToString() << "\n";
    hoisted_instructions[tuple_use] = b_opd_use;
  }
  cp_start = (tuple_use != nullptr) ? 1 : 0;
  for (int64 to_move_index = cp_start; to_move_index < to_move_in_size;
       to_move_index++) {
    Boundary b_to_move = to_move_in[to_move_index];
    HloInstruction* op = b_to_move.operands()[0];
    CHECK(op != nullptr);
    bool to_be_used_outside = true;
    VLOG(2) << "Mapping new boundary instr: " << op->ToString() << "\n";
    if (to_move_index < to_move_in_size - 1 && op->user_count() == 1 &&
        op->users()[0] == to_move_in[to_move_index + 1].operands()[0]) {
      to_be_used_outside = false;
      VLOG(2) << "Instruction is not to be used outside the branch\n";
    }
    Boundary b(Boundary::Position::kInsideBranch);
    for (int i = 0; i < branch_count; i++) {
      auto computation = conditional->branch_computation(i);
      VLOG(2) << "Copying to branch: " << i << "\n";
      TF_RETURN_IF_ERROR(CopyInOrOutOfConditional(b_to_move, i, computation,
                                                  hoisted_instructions));
      VLOG(2) << "Done:" << computation->ToString() << "\n";
      if (to_be_used_outside) {
        auto new_op = hoisted_instructions[op].operands()[i];
        auto new_root = computation->root_instruction();
        new_root->AppendOperand(new_op);
        *new_root->mutable_shape()->add_tuple_shapes() = new_op->shape();
        VLOG(2) << "Extending conditional root " << i << " : "
                << new_root->ToString() << "\n";
      }
      VLOG(2) << "After extending branch root: " << computation->ToString()
              << "\n";
    }
    if (to_be_used_outside) {
      // Modify uses of instructions outside of the conditionals
      HloInstruction* gtr = conditional->parent()->AddInstruction(
          HloInstruction::CreateGetTupleElement(op->shape(), conditional,
                                                op_index++));
      TF_RETURN_IF_ERROR(op->ReplaceAllUsesWith(gtr));
      if (conditional->parent()->root_instruction() == op) {
        conditional->parent()->set_root_instruction(gtr);
      }
    }
  }
  VLOG(2) << "Done copying instructions inside branch: "
          << conditional->ToString(HloPrintOptions::Fingerprint()) << "\n";
  // Change conditional instruction shape to the shape of the new root.
  HloInstruction* new_root =
      conditional->branch_computation(0)->root_instruction();
  *conditional->mutable_shape() = new_root->shape();
  VLOG(2) << "Before removing instructions:"
          << conditional->parent()->ToString() << "\n";
  // Remove hoisted instructions from the branches.
  for (int64 i = to_move_in_size - 1; i >= 0; i--) {
    Boundary boundary_to_move_in = to_move_in[i];
    HloInstruction* op = boundary_to_move_in.operands()[0];
    if (op->user_count() == 0) {
      VLOG(2) << "Removing boundary:" << boundary_to_move_in.ToString() << "\n";
      TF_RETURN_IF_ERROR(conditional->parent()->RemoveInstruction(op));
      VLOG(2) << "Done removing boundary.\n";
    }
  }

  // Reset shapes of user gtes to the new shape.
  if (use_index != -1) {
    for (auto* user : conditional->users()) {
      if (user->opcode() == HloOpcode::kGetTupleElement) {
        VLOG(2) << "Resetting shape of user: " << user->ToString() << "\n";
        *user->mutable_shape() =
            conditional->shape().tuple_shapes(user->tuple_index());
      }
    }
  }
  VLOG(1) << "Done moving instructions inside branches\n"
          << conditional->parent()->ToString(HloPrintOptions::Fingerprint())
          << "\n";
  return true;
}

// Group single chains of operands or uses of boundaries into new boundaries
class GroupConnectedBoundaries {
 private:
  std::vector<Boundary> connected_boundaries_, new_boundaries_;
  HloInstruction* conditional_;
  HloComputation* conditional_parent_;
  bool is_layout_sensitive_;
  // Instructions that have been visited but are not going to be moved.
  absl::flat_hash_map<HloInstruction*, int>& visited_count_;
  // The following four lines are configurations of the cost model, which will
  // be used to determine whether to move an instruction (move_config_) and how
  // strongly preferred it is to keep a pair of ops together (reuse_config_).
  // The search_config_ is used to control how to navigate the search space of
  // the cost model in the context of auto/manual tuning. The flipped array is
  // used to save which entries in the configuration have been changed in the
  // search/tuning process.
  std::vector<std::vector<int64>>& move_config_;
  std::vector<std::vector<int64>>& reuse_config_;
  int& search_config_;
  absl::flat_hash_map<const int64*, int64> flipped_;

  // The FlipMutation function serves to implement the search of alternative
  // cost models by deciding whether to flip a given configuration, saved in
  // the loc parameter. The non_zero parameter provides the new value to use
  // to flip a zero. The msg parameter is only used for debugging purpposes.
  int64 FlipMutation(int64* loc, const int64 non_zero, const std::string& msg) {
    if (search_config_ == 0 || ContainsKey(flipped_, loc)) {
      VLOG(2) << "Configured not to search or loc is already flipped.";
      return *loc;
    }

    // The 8-16 digits control the maximum number of times to flip a config.
    int flip_count = (search_config_ >> 8) & 255;
    if (flip_count == 0) {
      VLOG(2) << "Maximum flip count has reached. ";
      return *loc;
    }

    // The last 8 digits control when to start the first flip.
    int c = search_config_ & 255;
    VLOG(2) << "flip start index = " << c << "\n";
    // Only flip the decision if c reaches 0.
    if (c > 0) {
      search_config_--;
      return *loc;
    }

    // Decrement flip count so we can stop if it reaches 0.
    search_config_ -= 256;
    // Reload the 16-23 digits of the configuration, which controls how
    // frequently a configuration should be flipped.
    search_config_ += (search_config_ >> 16) & 255;
    VLOG(2) << "Updating Flipping configuration = " << search_config_ << "\n";

    flipped_[loc] = *loc;
    // Copy the last 8 bits back to the first 8 bits of configuration.
    switch (*loc) {
      case 0:
        *loc = non_zero;
        break;
      default:
        *loc = 0;
        break;
    }
    VLOG(2) << "Flipping decision for: " << msg << ": from " << flipped_[loc]
            << " to " << *loc << "\n";
    return *loc;
  }

 public:
  explicit GroupConnectedBoundaries(
      HloInstruction* conditional, bool is_layout_sensitive,
      absl::flat_hash_map<HloInstruction*, int>& visited_count,
      std::vector<std::vector<int64>>* move_config,
      std::vector<std::vector<int64>>* reuse_config, int* search_config)
      : conditional_(conditional),
        conditional_parent_(conditional->parent()),
        is_layout_sensitive_(is_layout_sensitive),
        visited_count_(visited_count),
        move_config_(*move_config),
        reuse_config_(*reuse_config),
        search_config_(*search_config) {}
  // Returns estimation of potential reuses carried by a given pair of
  // instructions. Use different integers to classify different levels
  // of reuses. Assume all instructions can be fused to enable data reuses.
  int64 ReusesCarriedBy(HloInstruction* op, HloInstruction* user) {
    std::vector<int64>& curconfig =
        reuse_config_[static_cast<uint32>(op->opcode())];
    // Flip the reuse configuration if tuning the cost model.
    // When flipping, use -10 if flipping to the default reuse model. Other
    // values can be specified if needed to fine-control the decision making.
    int64 config =
        (search_config_ < 0)
            ? FlipMutation(&curconfig[static_cast<uint32>(user->opcode())], -10,
                           HloOpcodeString(op->opcode()) + "->" +
                               HloOpcodeString(user->opcode()))
            : curconfig[static_cast<uint32>(user->opcode())];
    VLOG(2) << "ConditionalCodeMotion: Add reuses carried by instr: "
            << op->ToString() << "=>" << user->ToString() << " : " << config
            << "\n";
    if (config < 0) {
      // Assume the reuse decreases with increasing user count.
      int count1 = CountNonLeafOps(op->users());
      int count2 = CountNonLeafOps(user->operands());
      return (-config) / count1 / count2;
    }
    return config;
  }
  void clear_recently_visited() {
    for (const auto& boundary : new_boundaries_) {
      visited_count_.erase(boundary.operands()[0]);
    }
  }
  // Returns true if `instruction` is worth hoisting.
  bool WorthHoisting(HloInstruction* instruction, bool is_inside_branch) {
    // This is needed for the "moving-in" transformation, to prevent the root
    // of the parent computation (which contains the conditional) to be moved
    // inside the conditional.
    HloOpcode opcode = instruction->opcode();
    if (opcode == HloOpcode::kTuple &&
        instruction == conditional_parent_->root_instruction()) {
      return false;
    }
    // It is not safe to move collective ops from outside to inside
    // conditional branches, as it may cause synchronization problems,
    // when different layouts are assigned to different branches.
    if (opcode == HloOpcode::kAllReduce && !is_inside_branch) {
      return false;
    }

    // It is not legal to move the parameter instructions.
    if (opcode == HloOpcode::kParameter) {
      return false;
    }

    // Use configuration given from outside (e.g., by autotuner).
    std::vector<int64>& curconfig = move_config_[static_cast<uint32>(opcode)];
    auto col = (curconfig.size() == 1) ? 0
               : (instruction->operand_count() > 0)
                   ? static_cast<uint32>(instruction->operand(0)->opcode())
                   : 0;
    VLOG(2) << "column = " << col << "\n";
    VLOG(2) << "config size = " << curconfig.size() << "\n";
    VLOG(2) << "search_config = " << search_config_ << "\n";
    CHECK(col < curconfig.size());
    uint32 config = (search_config_ > 0)
                        ? FlipMutation(&curconfig[col], 1,
                                       "Move-" + HloOpcodeString(opcode))
                        : curconfig[col];
    VLOG(2) << "Checking instruction is worth moving: " << config << "\n";
    return (config != 0);
  }

  int64 ReusesBeforeBoundary(HloInstruction* user) {
    int64 reuses = 0;
    for (auto op : user->operands()) {
      // The operand must be an instruction that is not going to be moved (if
      // user is inside the conditional); otherwise it must be the conditional
      // itself and its user must be outside of the conditional.
      if (!ContainsKey(visited_count_, op) && op != conditional_) {
        continue;
      }
      if (auto tuple_gte = DynCast<HloGetTupleElementInstruction>(user)) {
        if (op->opcode() == HloOpcode::kConditional) {
          auto tuple = op->branch_computation(0)->root_instruction();
          if (tuple->opcode() == HloOpcode::kTuple) {
            auto index = tuple_gte->tuple_index();
            CHECK(index < tuple->operand_count());
            op = tuple->mutable_operand(index);
          }
        }
        reuses += ReusesCarriedBy(op, user->users()[0]);
      } else {
        reuses += ReusesCarriedBy(op, user);
      }
    }
    VLOG(2) << "Reuses before instruction " << user->ToString() << ":" << reuses
            << "\n";
    return reuses;
  }

  int64 ReusesAfterBoundary(HloInstruction* user) {
    CHECK(user != nullptr);
    auto all_users = user->users();
    // For now, assume that if an instruction has multiple-consumers, it
    // will not be reused, as the reuse may require duplication in
    // fusion and so is expensive. If the situation changes in the future,
    // some aspects of the overall algorithm need to be redesigned to
    // accommandate the change.
    if (all_users.size() > 1) {
      VLOG(2) << "Having multiple users from: " << user->ToString() << "\n";
      return 0;
    }
    if (!all_users.empty()) {
      auto op = all_users[0];
      int64 reuses = 0;
      // Only count reuses that run through the conditional root.
      if (op == conditional_->branch_computation(0)->root_instruction()) {
        int64 index = op->operand_index(user);
        for (auto op2 : conditional_->users()) {
          // If the use is not get tuple, right now do not consider it.
          if (op2->opcode() == HloOpcode::kGetTupleElement) {
            auto tuple_opd = static_cast<HloGetTupleElementInstruction*>(op2);
            if (index == tuple_opd->tuple_index()) {
              all_users = op2->users();
              if (!all_users.empty()) {
                reuses += ReusesCarriedBy(user, all_users[0]);
                break;
              }
            }
          }
        }
      } else if (ContainsKey(visited_count_, op)) {
        reuses += ReusesCarriedBy(user, op);
      }
      VLOG(2) << "reuses after instruction " << user->ToString() << ":"
              << reuses << "\n";
      return reuses;
    }
    return 0;
  }

  int64 BenefitForMovingBoundaries(const std::vector<Boundary>& boundaries) {
    int64 reuses_before = 0, reuses_after = 0;
    if (boundaries.size() == 1) {
      if (boundaries[0].IsOutsideBranch() &&
          boundaries[0].operands()[0]->opcode() ==
              HloOpcode::kGetTupleElement) {
        // The only boundary of moving-in is the get_tuple_element op.
        return -1;
      }
      if (boundaries[0].IsInsideBranch() &&
          boundaries[0].operands()[0]->opcode() == HloOpcode::kTuple) {
        // The only boundary of moving-out is the tuple op inside branches.
        return -1;
      }
    }
    // If trying alternative moving configurations, turn off reuse analysis.
    if (search_config_ > 0) {
      return 1;
    }
    // For cases like :
    // branch0 {
    //   ROOT copy
    // }
    // branch1 {
    //   ...
    // }
    // cond = conditional(branch0, branch1)
    // copy = copy(cond)
    //
    // We can fold the two copies thus reducing computation.
    auto get_copy_folding_benefit = [&](HloInstruction* hlo) -> int64 {
      if (hlo->opcode() != HloOpcode::kCopy) {
        return 0;
      }
      const HloGetTupleElementInstruction* gte =
          DynCast<HloGetTupleElementInstruction>(hlo->operand(0));
      if (gte == nullptr) {
        return 0;
      }
      const HloInstruction* conditional = gte->operand(0);
      if (conditional != conditional_) {
        return 0;
      }
      int64 benefit = 0;
      for (auto* branch : conditional->called_computations()) {
        HloInstruction* root = branch->root_instruction();
        if (root->opcode() == HloOpcode::kTuple) {
          const auto* tuple_operand = root->operand(gte->tuple_index());
          if (tuple_operand->opcode() == HloOpcode::kCopy) {
            if (Shape::Equal()(tuple_operand->operand(0)->shape(),
                               hlo->shape())) {
              benefit += 10;
            }
          }
        }
      }
      return benefit;
    };
    for (const Boundary& b : boundaries) {
      auto op = b.operands()[0];
      if (op == conditional_->branch_computation(0)->root_instruction()) {
        continue;
      }
      VLOG(2) << "Benefit for " << op->ToString();
      reuses_before += ReusesBeforeBoundary(op);
      VLOG(2) << "Reuses before boundary so far: " << reuses_before << "\n";
      reuses_after += ReusesAfterBoundary(op);
      VLOG(2) << "Reuese after boundary so far : " << reuses_after << "\n";
    }

    int64 copy_folding_benefit = 0;
    if (boundaries[0].IsOutsideBranch()) {
      for (const Boundary& b : boundaries) {
        auto op = b.operands()[0];
        copy_folding_benefit += get_copy_folding_benefit(op);
      }
    }
    VLOG(2) << "Copy folding benefit: " << copy_folding_benefit;

    if (reuses_after == 0 && reuses_before == 0 && copy_folding_benefit == 0) {
      return -1;
    } else if (boundaries[0].IsInsideBranch()) {
      return reuses_after - reuses_before;
    } else {
      return reuses_before - reuses_after - 1 + copy_folding_benefit;
    }
  }

  Boundary GetNextBoundary(const Boundary& b, int64 op_index) {
    Boundary b2(b.GetPosition());
    for (int j = 0; j < b.operands().size(); ++j) {
      HloInstruction* inst = b.operands()[j];
      CHECK(inst != nullptr);
      HloInstruction* op = (b.IsInsideBranch()) ? inst->operands()[op_index]
                                                : inst->users()[op_index];
      CHECK(op != nullptr);
      b2.mutable_operands().push_back(op);
    }
    return b2;
  }

  // Checking whether it is safe to move a boundary when visited through a
  // dependent already considered for moving.
  bool IsSafeToMoveBoundary(const Boundary& next_boundary) {
    int64 next_boundary_count =
        (next_boundary.IsInsideBranch())
            ? next_boundary.operands()[0]->user_count()
            : CountNonLeafOps(next_boundary.operands()[0]->operands());
    if (next_boundary_count <= 1) {
      // If boundary has only a single or no dependent, safe to move.
      return true;
    } else {
      if (!ContainsKey(visited_count_, next_boundary.operands()[0])) {
        VLOG(2) << "Skip next boundary " << next_boundary.ToString() << "\n"
                << " because it has multiple dependents: "
                << next_boundary_count << "\n";
        visited_count_[next_boundary.operands()[0]] = 1;
        new_boundaries_.push_back(next_boundary);
      } else {
        auto pos = std::find(new_boundaries_.begin(), new_boundaries_.end(),
                             next_boundary);
        if (pos != new_boundaries_.end() ||
            next_boundary.operands().size() == 1) {
          int count = ++visited_count_[next_boundary.operands()[0]];
          if (count == next_boundary_count) {
            VLOG(2) << "Recovering next boundary " << next_boundary.ToString()
                    << "\n"
                    << " because all of its dependents have been visited: "
                    << next_boundary_count << "\n";
            visited_count_.erase(next_boundary.operands()[0]);
            if (pos != new_boundaries_.end()) {
              new_boundaries_.erase(pos);
            }
            return true;
          }
        } else {
          VLOG(2) << "Skip incompatible multi-dependent boundary: "
                  << next_boundary.ToString() << ":" << next_boundary_count
                  << "\n";
        }
      }
    }
    return false;
  }
  // This function is reused both for moving the boundary outside or into a
  // conditional. As the result, the readability is somewhat compromised.
  // It might be nice to refactor this function to factor the outside-inside
  // considerations into separate function pointer parameters to improve
  // readability.
  void AddBoundaries(const Boundary& boundary) {
    BoundaryVisitor visitor;
    visitor.AddToWorkList(boundary);
    while (visitor.HasNextBoundary()) {
      Boundary b = visitor.PopNextBoundary();
      VLOG(2) << "visiting boundary " << b.ToString() << "\n";
      if ((b.IsOutsideBranch() || InstructionWithinBranchIdentical(
                                      b.operands(), is_layout_sensitive_)) &&
          IsSafeToMoveBoundary(b) &&
          WorthHoisting(b.operands()[0], b.IsInsideBranch())) {
        connected_boundaries_.push_back(b);
        VLOG(2) << "boundary can be moved\n";
        int64 operand_count = (b.IsInsideBranch())
                                  ? b.operands()[0]->operand_count()
                                  : b.operands()[0]->users().size();
        for (int i = 0; i < operand_count; i++) {
          Boundary next_boundary = GetNextBoundary(b, i);
          VLOG(2) << "Add operand/user " << i << " to visit later\n";
          visitor.AddToWorkList(next_boundary);
        }
      } else {
        VLOG(2) << "boundary cannot be moved\n";
        visited_count_[b.operands()[0]] = 1;
        new_boundaries_.push_back(b);
      }
    }
  }
  std::vector<Boundary> BoundariesToMoveInOrOut(HloInstruction* conditional,
                                                const Boundary& b) {
    // At the beginning of optimization, a conditional itself is added to a
    // worklist. Here the conditional is expanded into two sets of boundaries:
    // the first set contains the boundary that is inside branches and
    // contains the root of all branches; the second set of boundaries
    // contains all the users of the conditional.
    HloInstruction* inst = b.operands()[0];
    if (inst == conditional) {
      int branch_count = inst->branch_count();
      // Add conditional roots as a new boundary to visit.
      Boundary boundary_in(Boundary::Position::kInsideBranch);
      for (int i = 0; i < branch_count; i++) {
        HloComputation* branch_computation = inst->branch_computation(i);
        HloInstruction* root_inst = branch_computation->root_instruction();
        CHECK(root_inst != nullptr);
        boundary_in.mutable_operands().push_back(root_inst);
      }
      new_boundaries_.push_back(boundary_in);
      // Add conditional users as new boundaries to visit.
      for (auto u : inst->users()) {
        Boundary boundary_in(Boundary::Position::kOutsideBranch);
        boundary_in.mutable_operands().push_back(u);
        new_boundaries_.push_back(boundary_in);
      }
    } else {
      AddBoundaries(b);
    }
    return connected_boundaries_;
  }
  void AddNewBoundaries(std::vector<Boundary>& b) {
    b.insert(b.end(), new_boundaries_.begin(), new_boundaries_.end());
  }
};

ConditionalCodeMotion::Decision ConditionalCodeMotion::ConsiderCodeMotion(
    HloInstruction* conditional, const Boundary& cur_boundary,
    std::vector<Boundary>& to_move, std::vector<Boundary>& new_boundaries,
    absl::flat_hash_map<HloInstruction*, int>& visited_count) {
  GroupConnectedBoundaries connect(conditional, is_layout_sensitive_,
                                   visited_count, &move_config_, &reuse_config_,
                                   &search_config_);
  auto move_in_or_out =
      connect.BoundariesToMoveInOrOut(conditional, cur_boundary);
  if (!move_in_or_out.empty()) {
    auto benefit = connect.BenefitForMovingBoundaries(move_in_or_out);
    VLOG(2) << "benefit of moving in or out "
            << cur_boundary.operands()[0]->ToString() << ":" << benefit << "\n";
    if (benefit >= 0) {
      new_boundaries.clear();
      connect.AddNewBoundaries(new_boundaries);
      // The whole sequence in move_in_or_out is either all moving into a
      // conditional, or all moving out of a conditional. So looking only
      // at the first entry of the sequence is sufficient to know which
      // direction the move is intended.
      to_move = move_in_or_out;
      return Decision(to_move[0].IsInsideBranch()
                          ? Decision::Direction::kMoveOutOfBranch
                          : Decision::Direction::kMoveIntoBranch,
                      benefit);
    } else {
      connect.clear_recently_visited();
    }
  } else {
    connect.AddNewBoundaries(new_boundaries);
  }
  return Decision(Decision::Direction::kNoChange, 0);
}

StatusOr<bool> ConditionalCodeMotion::Run(HloModule* module) {
  VLOG(2) << "Begin a new pass of conditional code motion optimization.\n";
  // Use to support debugging of optimization, by disabling the opt after it has
  // been applied a pre-determined times (to isolate impact of transformations).
  if (!ConsumeFuel("conditional_code_motion", [&] {
        return "Skipping conditional opt after allowed limit reaching 0.\n";
      })) {
    return false;
  }
  bool changed = false;
  bool cleanup_changed = false;
  {
    HloPassPipeline subpipeline("before_conditional_code_motion");
    subpipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/is_layout_sensitive_);
    subpipeline.AddPass<HloDCE>();
    TF_ASSIGN_OR_RETURN(auto cleanup_changed_now, subpipeline.Run(module));
    cleanup_changed |= cleanup_changed_now;
  }
  // set the default configuration
  VLOG(2) << "Obtaining default configuration\n";
  SetDefaultMoveConfig();
  VLOG(2) << "Done obtaining default configuration\n";
  // Gather all the conditional ops in the module ahead of time, to avoid
  // potential complications of modifying the code that affecting traversal.
  std::vector<HloInstruction*> conditional_ops;
  // Track how many times each branch computation is shared.
  absl::flat_hash_map<HloComputation*, int> conditional_computations;
  for (auto* comp : module->MakeComputationPostOrder()) {
    for (auto* instr : comp->MakeInstructionPostOrder()) {
      if (instr->opcode() == HloOpcode::kConditional) {
        int branch_count = instr->branch_count();
        for (int i = 0; i < branch_count; ++i) {
          HloComputation* branch_i = instr->branch_computation(i);
          if (ContainsKey(conditional_computations, branch_i)) {
            conditional_computations[branch_i]++;
          } else {
            conditional_computations[branch_i] = 0;
          }
        }
        if (instr->shape().IsTuple()) {
          bool can_change_tuple_shape = true;
          for (auto user : instr->users()) {
            VLOG(2) << "user is : " << user->ToString() << "\n";
            if (user->opcode() != HloOpcode::kGetTupleElement) {
              can_change_tuple_shape = false;
            }
          }
          if (can_change_tuple_shape) {
            conditional_ops.push_back(instr);
          }
        } else {
          conditional_ops.push_back(instr);
        }
      }
    }
  }

  // Use to collect mappings between cloned instructions.
  HloCloneContext clone_context(module);
  for (HloInstruction* conditional : conditional_ops) {
    int branch_count = conditional->branch_count();
    // check for shared conditional computations
    bool conditional_is_shared = false;
    for (int i = 0; i < branch_count; ++i) {
      HloComputation* branch_i = conditional->branch_computation(i);
      if (conditional_computations[branch_i] > 0) {
        conditional_is_shared = true;
        break;
      }
    }

    // Boundaries to move out or to move into the branches.
    std::vector<std::vector<Boundary> > to_move_out, to_move_in;
    std::vector<std::vector<Boundary> > new_boundaries_for_moveout;
    std::vector<std::vector<Boundary> > new_boundaries_for_movein;
    // Number of times each instruction has been visited for moving.
    absl::flat_hash_map<HloInstruction*, int> visited_count;
    int benefit_move_out = 0, benefit_move_in = 0;
    Decision::Direction final_d = Decision::Direction::kNoChange;
    // The conditional is moved into a worklist as the seed (starting point).
    // The conditional will be expanded into multiple seeds (starting points),
    // its roots and its users, when it is visited by GroupConnectedBoundaries.
    // A NO_CHANGE decision will always be returned for the conditional itself,
    // so that the other seeding boundaries can be visited in turn.
    BoundaryVisitor visitor(conditional);
    VLOG(2) << "Analyzing conditional:" << conditional->ToString() << "\n";
    // Try visit all the boundaries, collect the analysis results, and save
    // all the benefitical non-conflicting decisions. If two decisions conflict
    // with each other, save the more benefitical one.
    while (visitor.HasNextBoundary()) {
      std::vector<Boundary> to_move, next_boundary;
      Boundary boundary = visitor.PopNextBoundary();
      VLOG(2) << "Analyzing boundary:" << boundary.ToString() << "\n";
      auto d = ConsiderCodeMotion(conditional, boundary, to_move, next_boundary,
                                  visited_count);
      switch (d.GetDirection()) {
        case Decision::Direction::kMoveOutOfBranch:
          VLOG(2) << "Local Decision is move out of branch\n";
          to_move_out.push_back(to_move);
          new_boundaries_for_moveout.push_back(next_boundary);
          benefit_move_out += d.GetBenefit();
          if (benefit_move_out >= benefit_move_in) {
            final_d = Decision::Direction::kMoveOutOfBranch;
            VLOG(2) << "Current Decision is move out of branch ("
                    << to_move_out.size() << ")\n";
          } else {
            VLOG(2) << "Current Decision remains move into branch\n";
          }
          break;
        case Decision::Direction::kMoveIntoBranch:
          VLOG(2) << "Decision is move into branch\n";
          to_move_in.push_back(to_move);
          new_boundaries_for_movein.push_back(next_boundary);
          benefit_move_in += d.GetBenefit();
          if (benefit_move_out >= benefit_move_in) {
            VLOG(2) << "Current Decision remains move out of branch\n";
          } else {
            final_d = Decision::Direction::kMoveIntoBranch;
            VLOG(2) << "Current Decision is move into branch ("
                    << to_move_in.size() << ")\n";
          }
          break;
        case Decision::Direction::kNoChange:
          VLOG(2) << "Decision is no change\n";
          for (const Boundary& b : next_boundary) {
            visitor.AddToWorkList(b);
            VLOG(2) << "Adding new boundary to worklist:" << b.ToString()
                    << "\n";
          }
          break;
      }
    }
    // If modification is to be made, need to clone the shared branches.
    if (final_d != Decision::Direction::kNoChange && conditional_is_shared) {
      for (int i = 0; i < branch_count; ++i) {
        HloComputation* branch_i = conditional->branch_computation(i);
        if (conditional_computations[branch_i] > 0) {
          // Cloning is absolutely needed if the computation is shared by
          // different branches, but the cloning can be potentially avoided
          // if the sharing is only among branches of the same conditional.
          // If cloning these branches causes a problem due to space issues,
          // a fix can pass a vector of unique branches to the actual
          // transformations, as an alternative representation of the
          // conditional branches to be modified. Right now we assume the
          // overhead of cloning is minimal since later stages of the compiler
          // inline all the computations anyway.
          HloComputation* clone_i =
              conditional->parent()->parent()->AddEmbeddedComputation(
                  branch_i->Clone("clone", &clone_context));
          conditional->set_branch_computation(i, clone_i);
          conditional_computations[branch_i]--;
          // Need to translate the analysis result to generate correct result.
          auto update_boundary = [&](Boundary& boundary) {
            auto cloned_instr =
                clone_context.FindInstruction(boundary.operands()[i]);
            CHECK(cloned_instr != nullptr);
            VLOG(2) << "boundary before cloning:" << boundary.operands()[i]
                    << "\n";
            boundary.mutable_operands()[i] = cloned_instr;
            VLOG(2) << "boundary after cloning:" << boundary.operands()[i]
                    << "\n";
          };
          // Only boundaries to move out need to be updated.
          if (final_d == Decision::Direction::kMoveOutOfBranch) {
            for (int i = 0; i < to_move_out.size(); ++i) {
              std::vector<Boundary>& m = to_move_out[i];
              std::for_each(m.begin(), m.end(), update_boundary);
            }
            for (int i = 0; i < new_boundaries_for_moveout.size(); ++i) {
              std::vector<Boundary>& m = new_boundaries_for_moveout[i];
              std::for_each(m.begin(), m.end(), update_boundary);
            }
          }
        }
      }
      VLOG(2) << "Cloned branches as needed: " << conditional->ToString()
              << "\n";
    }
    // At most one of to_move_out or to_move_in can be non-empty, since there is
    // only one optimization decision.
    if (final_d == Decision::Direction::kMoveOutOfBranch) {
      CHECK(to_move_out.size() == new_boundaries_for_moveout.size());
      for (int i = 0; i < to_move_out.size(); ++i) {
        TF_ASSIGN_OR_RETURN(bool result,
                            MoveInstructionOut(conditional, to_move_out[i],
                                               new_boundaries_for_moveout[i]));
        changed |= result;
      }
      VLOG(2) << "Done moving out of branches " << to_move_out.size()
              << " times. \n";
      if (!ConsumeFuel("conditional_code_motion", [&] {
            return "Skipping conditional opt after allowed limit reaching 0.\n";
          })) {
        break;
      }
    } else if (final_d == Decision::Direction::kMoveIntoBranch) {
      CHECK(to_move_in.size() == new_boundaries_for_movein.size());
      for (int i = 0; i < to_move_in.size(); ++i) {
        TF_ASSIGN_OR_RETURN(bool result,
                            MoveInstructionIn(conditional, to_move_in[i],
                                              new_boundaries_for_movein[i]));
        changed |= result;
      }
      VLOG(2) << "Done moving into branches " << to_move_in.size()
              << " times. \n";
      if (!ConsumeFuel("conditional_code_motion", [&] {
            return "Skipping conditional opt after allowed limit reaching 0.\n";
          })) {
        break;
      }
    } else if (pursue_full_conditional_code_motion_ && !conditional_is_shared) {
      // Invoke special handling for convert rematerialization/hoisting
      // We need to make sure no sharing is present in the branches because no
      // cloning has been done by the earlier analysis.
      // TOOD[b/165848866]: extend solution to handle cloning for special move.
      TF_ASSIGN_OR_RETURN(
          bool convert_result,
          ConvertSpecialMove(conditional, is_layout_sensitive_));
      if (convert_result) {
        VLOG(2) << "Done special moving of convert\n";
        if (!ConsumeFuel("conditional_code_motion", [&] {
              return "Skipping conditional opt after allowed limit reaching "
                     "0.\n";
            })) {
          break;
        }
      }
      changed |= convert_result;
    }
  }
  if (changed) {
    HloPassPipeline subpipeline(
        "after_conditional_code_motion_after_convert_hoisting");
    VLOG(2) << "starting after motion passes: DCE\n";
    subpipeline.AddPass<HloDCE>();
    subpipeline.AddPass<TupleSimplifier>();
    subpipeline.AddPass<HloDCE>();
    TF_ASSIGN_OR_RETURN(auto cleanup_changed_now, subpipeline.Run(module));
    cleanup_changed |= cleanup_changed_now;
  }
  if (cleanup_changed) {
    VLOG(2) << "subpipeline cleanup have modified code\n";
  }
  return changed;
}

void ConditionalCodeMotion::SetDefaultMoveConfig() {
  int tuning_option = (search_config_ == 0) ? 0 : (search_config_ > 0) ? 1 : 2;

  auto row = HloOpcodeCount();
  auto col = row;
  VLOG(2) << "Start setting default configuration\n";
  reuse_config_.reserve(row);
  move_config_.reserve(row);
  for (int64 opcode = 0; opcode < row; ++opcode) {
    // To save whether an instruction is preferred to be moved.
    std::vector<int64> reuse_vec(col, 0);
    for (uint32 j = 0; j < col; ++j) {
      reuse_vec[j] = ReusesCarriedBy(static_cast<HloOpcode>(opcode),
                                     static_cast<HloOpcode>(j));
    }
    reuse_config_.push_back(reuse_vec);
    std::vector<int64> move_vec;
    switch (tuning_option) {
      case 1:
        // Tuning transformation decision --- start with all yes.
        // Only a single entry is needed if we don't consider operands of an op
        // when searching/tuning transformation decisions.
        move_vec.push_back(1);
        break;
      case 2:  // Tune the ReusesCarriedBy results only.
      case 0:
        // No tuning --- use the default configuration.
        // Use the opcode of first operand to configure default.
        move_vec.reserve(col);
        for (uint32 j = 0; j < col; ++j) {
          move_vec.push_back(WorthHoisting(static_cast<HloOpcode>(opcode),
                                           static_cast<HloOpcode>(j)));
        }
        break;
    }
    move_config_.push_back(move_vec);
  }
}

}  // namespace conditional_opt

}  // namespace xla