path: root/tensorflow/compiler/xla/service/algebraic_simplifier.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/algebraic_simplifier.h"

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

#include "tensorflow/compiler/xla/layout_util.h"
#include "tensorflow/compiler/xla/literal_util.h"
#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/compiler/xla/service/hlo_query.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/compiler/xla/window_util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"

namespace xla {
namespace {

// Returns whether operand is a literal with the given value.
bool IsLiteralWithValue(const HloInstruction* operand, int value) {
  return operand->opcode() == HloOpcode::kConstant &&
         LiteralUtil::IsAll(operand->literal(), value);
}

// Returns whether the given transpose produces a result which is bit-wise
// identical to its operand and thus may be replaced with a bitcast.
bool TransposeIsBitcast(
    const HloInstruction* transpose,
    const AlgebraicSimplifier::ValidBitcastCallback& valid_bitcast_callback) {
  CHECK_EQ(HloOpcode::kTranspose, transpose->opcode());
  const HloInstruction* operand = transpose->operand(0);

  // Can't insert bitcasts if the compiler used a memory layout which isn't
  // compatible.
  if (!valid_bitcast_callback(operand->shape(), transpose->shape())) {
    return false;
  }

  return ShapeUtil::TransposeIsBitcast(operand->shape(), transpose->shape(),
                                       transpose->dimensions());
}

// Returns true if the given reshape produces a result which is bit-wise
// identical to its operand and thus may be replaced with a bitcast.
//
// This function is conservative -- even if this function returns false, the
// reshape may still be a bitcast. For example, a reshape from [28x28] to [784].
bool ReshapeIsBitcast(
    const HloInstruction* reshape,
    const AlgebraicSimplifier::ValidBitcastCallback& valid_bitcast_callback) {
  CHECK_EQ(HloOpcode::kReshape, reshape->opcode());

  const HloInstruction* operand = reshape->operand(0);
  // Can't insert bitcasts if the compiler used a memory layout which isn't
  // compatible.
  if (!valid_bitcast_callback(operand->shape(), reshape->shape())) {
    return false;
  }

  return ShapeUtil::ReshapeIsBitcast(operand->shape(), reshape->shape());
}
}  // namespace

// AlgebraicSimplifierVisitor traverses the HLO computation and reduces certain
// algebraic expressions to simplified forms. Note: This only supports
// simplifications that simply look at the operands of an instruction. For the
// more general case a worklist based approach would be needed.
class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault {
 public:
  // Default visitor action is to do nothing and return OK.
  Status DefaultAction(HloInstruction* /*hlo_instruction*/) override {
    return Status::OK();
  }

  Status HandleAdd(HloInstruction* add, HloInstruction* lhs,
                   HloInstruction* rhs) override;

  Status HandleBroadcast(HloInstruction* broadcast) override;

  Status HandleCopy(HloInstruction* copy, HloInstruction* operand) override;

  Status HandleConvert(HloInstruction* convert,
                       HloInstruction* operand) override;

  Status HandleConvolution(HloInstruction* convolution, HloInstruction* lhs,
                           HloInstruction* rhs, const Window& window) override;

  Status HandleDivide(HloInstruction* divide, HloInstruction* lhs,
                      HloInstruction* rhs) override;

  Status HandleGetTupleElement(HloInstruction* get_tuple_element,
                               HloInstruction* operand) override;

  Status HandleLog(HloInstruction* log, HloInstruction* operand) override;

  Status HandleMultiply(HloInstruction* multiply, HloInstruction* lhs,
                        HloInstruction* rhs) override;

  Status HandlePad(HloInstruction* pad) override;

  Status HandlePower(HloInstruction* power, HloInstruction* lhs,
                     HloInstruction* rhs) override;

  Status HandleReshape(HloInstruction* reshape) override;

  Status HandleReduce(HloInstruction* reduce, HloInstruction* arg,
                      HloInstruction* init_value,
                      tensorflow::gtl::ArraySlice<int64> dimensions,
                      HloComputation* function) override;

  Status HandleSlice(HloInstruction* slice, HloInstruction* operand) override;

  Status HandleTranspose(HloInstruction* transpose) override;

  Status HandleSubtract(HloInstruction* sub, HloInstruction* lhs,
                        HloInstruction* rhs) override;

  Status HandleMaximum(HloInstruction* maximum, HloInstruction* lhs,
                       HloInstruction* rhs) override;

  Status HandleMinimum(HloInstruction* minimum, HloInstruction* lhs,
                       HloInstruction* rhs) override;

  // Returns whether algebraic simplification has occurred.
  const bool changed() const { return changed_; }

  // Runs the visitor on a computation.
  static bool Run(
      HloComputation* computation, bool is_layout_sensitive,
      AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback);

 private:
  explicit AlgebraicSimplifierVisitor(
      HloComputation* computation, bool is_layout_sensitive,
      AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback)
      : computation_(computation),
        is_layout_sensitive_(is_layout_sensitive),
        valid_bitcast_callback_(std::move(valid_bitcast_callback)) {}

  // Convenience method for replacing an instruction with a bitcast.
  void ReplaceWithBitcast(HloInstruction* instruction);

  // Replace old instruction with new instruction if old and new instructions
  // have the same shape. Updates uses and root instruction. Returns whether a
  // replacement was made.
  bool ReplaceInstructionIfSameShape(HloInstruction* old_instruction,
                                     HloInstruction* new_instruction);

  // Returns whether the shape of the output of the given instructions are the
  // same for the purposes of simplification. If is_layout_sensitive_ is true,
  // then this tests shape equality including layout (ShapeUtil::Equal). If
  // is_layout_sensitive_ is false, then the tests shape compatibility
  // (ShapeUtil::Compatible).
  bool SameShape(const HloInstruction* lhs, const HloInstruction* rhs) const;

  // Returns whether it was possible to transform `root` to a clamp instruction.
  // With min a minimum instruction, max a maximum instruction, min_operand a
  // operand of min and max_operand a operand of max.
  // Precondition: root is either a minimum or a maximum.
  bool TransformToClampIfSameShape(HloInstruction* root, HloInstruction* min,
                                   HloInstruction* min_operand,
                                   HloInstruction* operand, HloInstruction* max,
                                   HloInstruction* max_operand);

  // Current HloComputation instance the AlgebraicSimplifierVisitor is
  // traversing.
  HloComputation* computation_;

  // Whether algebraic simplification has occurred.
  bool changed_ = false;

  // Whether layout is considered during transformation.
  bool is_layout_sensitive_;

  // Callback used to determine if a bitcast is valid.
  AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback_;
};

bool AlgebraicSimplifierVisitor::Run(
    HloComputation* computation, bool is_layout_sensitive,
    AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback) {
  AlgebraicSimplifierVisitor visitor(computation, is_layout_sensitive,
                                     std::move(valid_bitcast_callback));
  TF_CHECK_OK(computation->root_instruction()->Accept(&visitor));
  return visitor.changed_;
}

bool AlgebraicSimplifierVisitor::SameShape(const HloInstruction* lhs,
                                           const HloInstruction* rhs) const {
  if (is_layout_sensitive_) {
    return ShapeUtil::Equal(lhs->shape(), rhs->shape());
  } else {
    return ShapeUtil::Compatible(lhs->shape(), rhs->shape());
  }
}

void AlgebraicSimplifierVisitor::ReplaceWithBitcast(
    HloInstruction* instruction) {
  CHECK_EQ(1, instruction->operand_count());
  CHECK_EQ(ShapeUtil::ElementsIn(instruction->shape()),
           ShapeUtil::ElementsIn(instruction->operand(0)->shape()));
  CHECK_EQ(ShapeUtil::ByteSizeOf(instruction->shape()),
           ShapeUtil::ByteSizeOf(instruction->operand(0)->shape()));

  auto bitcast = computation_->AddInstruction(
      HloInstruction::CreateUnary(instruction->shape(), HloOpcode::kBitcast,
                                  instruction->mutable_operand(0)));
  computation_->ReplaceInstruction(instruction, bitcast);
  changed_ = true;
}

bool AlgebraicSimplifierVisitor::ReplaceInstructionIfSameShape(
    HloInstruction* old_instruction, HloInstruction* new_instruction) {
  if (!SameShape(old_instruction, new_instruction)) {
    return false;
  }
  computation_->ReplaceInstruction(old_instruction, new_instruction);
  changed_ = true;
  return true;
}

Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add,
                                             HloInstruction* lhs,
                                             HloInstruction* rhs) {
  // A + 0 => A
  VLOG(10) << "trying transform [A + 0 => A]: " << add->ToString();
  if (IsLiteralWithValue(rhs, 0) && ReplaceInstructionIfSameShape(add, lhs)) {
    return Status::OK();
  }
  // 0 + A => A
  VLOG(10) << "trying transform [0 + A => A]: " << add->ToString();
  if (IsLiteralWithValue(lhs, 0) && ReplaceInstructionIfSameShape(add, rhs)) {
    return Status::OK();
  }

  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy,
                                              HloInstruction* operand) {
  // All copies can be eliminated (assuming layout constraints are satisified).
  ReplaceInstructionIfSameShape(copy, operand);
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub,
                                                  HloInstruction* lhs,
                                                  HloInstruction* rhs) {
  // A - 0 => A
  VLOG(10) << "trying transform [A - 0 => A]: " << sub->ToString();
  if (IsLiteralWithValue(rhs, 0) && ReplaceInstructionIfSameShape(sub, lhs)) {
    return Status::OK();
  }

  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide,
                                                HloInstruction* lhs,
                                                HloInstruction* rhs) {
  // A/1 => A
  VLOG(10) << "trying transform [A/1 => A]: " << divide->ToString();
  if (IsLiteralWithValue(rhs, 1) && ReplaceInstructionIfSameShape(divide, lhs)) {
    return Status::OK();
  }

  // exp(A)/exp(B) => exp(A-B)
  if (lhs->opcode() == HloOpcode::kExp && rhs->opcode() == HloOpcode::kExp) {
    VLOG(10) << "transform [exp(A)/exp(B) => exp(A-B)]: " << divide->ToString();
    HloInstruction* subtract =
        computation_->AddInstruction(HloInstruction::CreateBinary(
            divide->shape(), HloOpcode::kSubtract, lhs->mutable_operand(0),
            rhs->mutable_operand(0)));
    computation_->ReplaceWithNewInstruction(
        divide, HloInstruction::CreateUnary(divide->shape(), HloOpcode::kExp,
                                            subtract));
    changed_ = true;
  }

  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply,
                                                  HloInstruction* lhs,
                                                  HloInstruction* rhs) {
  // A*1 => A
  VLOG(10) << "trying transform [A*1 => A]: " << multiply->ToString();
  if (IsLiteralWithValue(rhs, 1) &&
      ReplaceInstructionIfSameShape(multiply, lhs)) {
    return Status::OK();
  }
  // 1*A => A
  VLOG(10) << "trying transform [1*A => A]: " << multiply->ToString();
  if (IsLiteralWithValue(lhs, 1) &&
      ReplaceInstructionIfSameShape(multiply, rhs)) {
    return Status::OK();
  }
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log,
                                             HloInstruction* operand) {
  // ln(exp(A)) => A
  VLOG(10) << "trying transform [ln(exp(A)) => A]: " << log->ToString();
  if (operand->opcode() == HloOpcode::kExp &&
      ReplaceInstructionIfSameShape(log, operand->mutable_operand(0))) {
    return Status::OK();
  }
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleGetTupleElement(
    HloInstruction* get_tuple_element, HloInstruction* operand) {
  if (operand->opcode() == HloOpcode::kTuple) {
    // get_tuple_element(make_tuple({A_0, A_1, ..., A_n}), i) => A_i
    VLOG(10) << "trying transform "
             << "[get_tuple_element(make_tuple({...,A_i,...}), i)] => A_i: "
             << get_tuple_element->ToString();
    if (ReplaceInstructionIfSameShape(
            get_tuple_element,
            operand->mutable_operand(get_tuple_element->tuple_index()))) {
      return Status::OK();
    }
  }
  return Status::OK();
}

namespace {

// Return whether the given reshape instruction leaves the dimensions at the
// given input indices unmodified, and returns their output indices.
//
// Example:
//   input_dim_indices = {2, 3}
//   input  shape = T[a, b, x, y, cd]
//   output shape = T[ab, x, 1, y, c, d]
//   return value = {1, 3}
//
// Precondition: input_dim_indices is sorted.
std::pair<bool, std::vector<int64>> ReshapeLeavesDimensionsUnmodified(
    const HloInstruction* hlo,
    tensorflow::gtl::ArraySlice<int64> input_dim_indices) {
  CHECK_EQ(HloOpcode::kReshape, hlo->opcode());
  CHECK(std::is_sorted(input_dim_indices.begin(), input_dim_indices.end()));

  std::vector<int64> output_dim_indices;
  std::vector<std::pair<int64, int64>> unmodified_dims =
      ShapeUtil::DimensionsUnmodifiedByReshape(hlo->operand(0)->shape(),
                                               hlo->shape());
  size_t i = 0;  // index to unmodified_dims
  for (int64 input_dim_index : input_dim_indices) {
    // Search unmodified_dims for input_dim_index. We can search from the last
    // matching position because input_dim_indices is guaranteed to be sorted.
    while (i < unmodified_dims.size() &&
           unmodified_dims[i].first < input_dim_index) {
      ++i;
    }
    if (i >= unmodified_dims.size() ||
        unmodified_dims[i].first != input_dim_index) {
      return std::make_pair(false, std::vector<int64>());
    }
    output_dim_indices.push_back(unmodified_dims[i].second);
  }
  return std::make_pair(true, output_dim_indices);
}

// Returns true if the output of "instruction" is a permutation of the elements
// of "operand". Precondition: "operand" is an operand of "instruction".
bool OutputIsPermutationOfOperandElements(HloInstruction* instruction,
                                          HloInstruction* operand) {
  DCHECK(!instruction->OperandIndices(operand).empty());
  switch (instruction->opcode()) {
    case HloOpcode::kReshape:
    case HloOpcode::kReverse:
    case HloOpcode::kSort:
    case HloOpcode::kTranspose:
      return true;
    default:
      return false;
  }
}

// Returns true if the output of "instruction" is a subset of the elements of
// "operand". Precondition: "operand" is an operand of "instruction".
bool OutputIsSubsetOfOperandElements(HloInstruction* instruction,
                                     HloInstruction* operand) {
  std::vector<int64> operand_indices = instruction->OperandIndices(operand);
  CHECK(!operand_indices.empty());
  if (operand_indices.size() != 1) {
    return false;
  }
  int64 operand_index = operand_indices[0];
  switch (instruction->opcode()) {
    case HloOpcode::kSlice:
      CHECK_EQ(0, operand_index);
      return true;
    case HloOpcode::kDynamicSlice:
      return operand_index == 0;
    default:
      return false;
  }
}

}  // namespace

Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) {
  auto operand = broadcast->mutable_operand(0);
  // A degenerate broadcast of a reshape that does not change the number of
  // elements can be replaced by a reshape.
  if (std::is_sorted(broadcast->dimensions().begin(),
                     broadcast->dimensions().end()) &&
      ShapeUtil::ElementsIn(broadcast->shape()) ==
          ShapeUtil::ElementsIn(operand->shape())) {
    VLOG(10) << "transform broadcast(X) -> reshape(X) where "
                "n(broadcast(X)) == n(X)";
    computation_->ReplaceWithNewInstruction(
        broadcast, HloInstruction::CreateReshape(broadcast->shape(), operand));
    changed_ = true;
    return Status::OK();
  }

  // A broadcast of a reshape which merely inserts 1-sized dimensions can elide
  // its operand.
  {
    bool merely_inserts_or_deletes_1_sized_dimensions;
    std::vector<int64> inserted_indices, deleted_indices;
    std::tie(merely_inserts_or_deletes_1_sized_dimensions, deleted_indices,
             inserted_indices) =
        operand->ReshapeMerelyInsertsOrDeletes1SizedDimensions();
    if (merely_inserts_or_deletes_1_sized_dimensions &&
        deleted_indices.empty()) {
      std::reverse(inserted_indices.begin(), inserted_indices.end());
      auto dims = broadcast->dimensions();
      for (auto inserted_index : inserted_indices) {
        dims.erase(dims.begin() + inserted_index);
      }
      computation_->ReplaceWithNewInstruction(
          broadcast,
          HloInstruction::CreateBroadcast(broadcast->shape(),
                                          operand->mutable_operand(0), dims));
      changed_ = true;
      return Status::OK();
    }
  }

  // A scalar broadcast feeding an instruction which only permutes (reshape,
  // transpose, sort, reverse) or selects a subset of operand elements (slice,
  // dynamic slice) can be replaced with a broadcast directly to the output
  // shape of the instruction.
  if (ShapeUtil::IsScalar(operand->shape())) {
    for (HloInstruction* user : broadcast->users()) {
      if (OutputIsPermutationOfOperandElements(user, broadcast) ||
          OutputIsSubsetOfOperandElements(user, broadcast)) {
        HloInstruction* new_broadcast = computation_->AddInstruction(
            HloInstruction::CreateBroadcast(user->shape(), operand, {}));
        // Use ReplaceUsesOfInstruction instead of ReplaceWithNewInstruction
        // because we are replacing an instruction other than the visited
        // instruction.
        computation_->ReplaceUsesOfInstruction(user, new_broadcast);
        changed_ = true;
        return Status::OK();
      }
    }
  }
  return Status::OK();
}

template <PrimitiveType primitive_src_type, PrimitiveType primitive_dest_type>
static std::unique_ptr<HloInstruction> ConvertIfTypesMatch(
    const Literal& src_literal) {
  CHECK_EQ(primitive_src_type, src_literal.shape().element_type());

  return HloInstruction::CreateConstant(
      LiteralUtil::Convert<typename primitive_util::PrimitiveTypeToNative<
                               primitive_src_type>::type,
                           typename primitive_util::PrimitiveTypeToNative<
                               primitive_dest_type>::type>(src_literal));
}

template <PrimitiveType primitive_src_type>
static std::unique_ptr<HloInstruction> ConvertIfDestTypeMatches(
    const Literal& src_literal, PrimitiveType primitive_dest_type) {
  switch (primitive_dest_type) {
#define CONVERT_IF_TYPES_MATCH(type) \
  case (type):                       \
    return ConvertIfTypesMatch<primitive_src_type, (type)>(src_literal);
    CONVERT_IF_TYPES_MATCH(PRED)
    CONVERT_IF_TYPES_MATCH(S8)
    CONVERT_IF_TYPES_MATCH(S32)
    CONVERT_IF_TYPES_MATCH(S64)
    CONVERT_IF_TYPES_MATCH(U8)
    CONVERT_IF_TYPES_MATCH(U32)
    CONVERT_IF_TYPES_MATCH(U64)
    CONVERT_IF_TYPES_MATCH(F32)
    CONVERT_IF_TYPES_MATCH(F64)
#undef CONVERT_IF_TYPES_MATCH
    // Other types are not yet supported.
    default:
      LOG(FATAL) << "Unimplemented: ConvertIfDestTypeMatches for type "
                 << PrimitiveType_Name(src_literal.shape().element_type());
  }
}

static std::unique_ptr<HloInstruction> ConvertIfSrcTypeMatches(
    const Literal& src_literal, PrimitiveType primitive_dest_type) {
  switch (src_literal.shape().element_type()) {
#define CONVERT_IF_DEST_TYPE_MATCHES(type) \
  case (type):                             \
    return ConvertIfDestTypeMatches<(type)>(src_literal, primitive_dest_type);
    CONVERT_IF_DEST_TYPE_MATCHES(PRED)
    CONVERT_IF_DEST_TYPE_MATCHES(S8)
    CONVERT_IF_DEST_TYPE_MATCHES(S32)
    CONVERT_IF_DEST_TYPE_MATCHES(S64)
    CONVERT_IF_DEST_TYPE_MATCHES(U8)
    CONVERT_IF_DEST_TYPE_MATCHES(U32)
    CONVERT_IF_DEST_TYPE_MATCHES(U64)
    CONVERT_IF_DEST_TYPE_MATCHES(F32)
    CONVERT_IF_DEST_TYPE_MATCHES(F64)
#undef CONVERT_IF_DEST_TYPE_MATCHES
    // Other types are not yet supported.
    default:
      LOG(FATAL) << "Unimplemented: ConvertIfSrcTypeMatches for type "
                 << PrimitiveType_Name(src_literal.shape().element_type());
  }
}

// A conversion to the same element type as the operand is a nop and can be
// removed.  A conversion of a constant can be simplified by making a new
// constant.
Status AlgebraicSimplifierVisitor::HandleConvert(HloInstruction* convert,
                                                 HloInstruction* operand) {
  PrimitiveType src_type = operand->shape().element_type();
  PrimitiveType dest_type = convert->shape().element_type();
  if (src_type == dest_type) {
    computation_->ReplaceInstruction(convert, operand);
    changed_ = true;
    return Status::OK();
  }
  if (operand->opcode() == HloOpcode::kConstant) {
    const Literal& src_literal = operand->literal();
    std::unique_ptr<HloInstruction> new_constant =
        ConvertIfSrcTypeMatches(src_literal, dest_type);
    computation_->ReplaceWithNewInstruction(convert, std::move(new_constant));
    changed_ = true;
    return Status::OK();
  }
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) {
  // The pad instruction does nothing if the output shape is the same as the
  // input shape, i.e, all paddings are zero.
  ReplaceInstructionIfSameShape(pad, pad->mutable_operand(0));
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power,
                                               HloInstruction* lhs,
                                               HloInstruction* rhs) {
  VLOG(10) << "trying transform [pow(A, 0) => 1]: " << power->ToString();
  if (IsLiteralWithValue(rhs, 0)) {
    auto one = HloInstruction::CreateConstant(LiteralUtil::CloneToUnique(
        LiteralUtil::One(power->shape().element_type())));
    std::unique_ptr<HloInstruction> ones;
    if (ShapeUtil::IsScalar(power->shape())) {
      ones = std::move(one);
    } else {
      ones = HloInstruction::CreateBroadcast(
          power->shape(), computation_->AddInstruction(std::move(one)), {});
    }
    computation_->ReplaceWithNewInstruction(power, std::move(ones));
    changed_ = true;
    return Status::OK();
  }

  VLOG(10) << "trying transform [pow(A, 1) => A]: " << power->ToString();
  if (IsLiteralWithValue(rhs, 1) && ReplaceInstructionIfSameShape(power, lhs)) {
    return Status::OK();
  }

  VLOG(10) << "trying transform [pow(A, 2) => A*A]: " << power->ToString();
  if (IsLiteralWithValue(rhs, 2)) {
    computation_->ReplaceWithNewInstruction(
        power, HloInstruction::CreateBinary(power->shape(),
                                            HloOpcode::kMultiply, lhs, lhs));
    changed_ = true;
    return Status::OK();
  }

  VLOG(10) << "trying transform [pow(A, -1) => 1/A]: " << power->ToString();
  if (IsLiteralWithValue(rhs, -1)) {
    auto* one = computation_->AddInstruction(
        HloInstruction::CreateConstant(LiteralUtil::CloneToUnique(
            LiteralUtil::One(rhs->shape().element_type()))));
    computation_->ReplaceWithNewInstruction(
        power, HloInstruction::CreateBinary(power->shape(), HloOpcode::kDivide,
                                            one, lhs));
    changed_ = true;
    return Status::OK();
  }
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) {
  auto operand = reshape->mutable_operand(0);

  // Delete no-op reshapes, i.e. where shape = operand shape.
  if (SameShape(reshape, operand)) {
    VLOG(10) << "deleting no-op reshape";
    computation_->ReplaceInstruction(reshape, operand);
    changed_ = true;
    return Status::OK();
  }

  // Merge reshapes.
  if (HloOpcode::kReshape == operand->opcode()) {
    computation_->ReplaceWithNewInstruction(
        reshape, HloInstruction::CreateReshape(reshape->shape(),
                                               operand->mutable_operand(0)));
    changed_ = true;
    return Status::OK();
  }

  if (HloOpcode::kBroadcast == reshape->operand(0)->opcode()) {
    auto opt_dims = ReshapeLeavesDimensionsUnmodified(
        reshape, reshape->operand(0)->dimensions());
    if (opt_dims.first) {
      computation_->ReplaceWithNewInstruction(
          reshape,
          HloInstruction::CreateBroadcast(
              reshape->shape(), reshape->mutable_operand(0)->mutable_operand(0),
              opt_dims.second));
      changed_ = true;
      return Status::OK();
    }
  }

  // Make this a bitcast if possible.
  if (is_layout_sensitive_ &&
      ReshapeIsBitcast(reshape, valid_bitcast_callback_)) {
    ReplaceWithBitcast(reshape);
    return Status::OK();
  }

  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice,
                                               HloInstruction* operand) {
  // Delete no-op slices, i.e. where shape = operand shape.
  if (ReplaceInstructionIfSameShape(slice, operand)) {
    return Status::OK();
  }
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleReduce(
    HloInstruction* reduce, HloInstruction* arg, HloInstruction* init_value,
    tensorflow::gtl::ArraySlice<int64> dimensions, HloComputation* function) {
  if (ShapeUtil::ElementsIn(reduce->shape()) ==
      ShapeUtil::ElementsIn(arg->shape())) {
    auto reshape = computation_->AddInstruction(
        HloInstruction::CreateReshape(reduce->shape(), arg));
    computation_->ReplaceWithNewInstruction(
        reduce, HloInstruction::CreateMap(reduce->shape(),
                                          {reshape, init_value}, function));
    return Status::OK();
  }
  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleTranspose(HloInstruction* transpose) {
  auto operand = transpose->mutable_operand(0);

  if (std::is_sorted(transpose->dimensions().begin(),
                     transpose->dimensions().end())) {
    VLOG(10) << "deleting no-op transpose";
    computation_->ReplaceInstruction(transpose, operand);
    changed_ = true;
    return Status::OK();
  }

  if (HloOpcode::kTranspose == operand->opcode()) {
    computation_->ReplaceWithNewInstruction(
        transpose, HloInstruction::CreateTranspose(
                       transpose->shape(), operand->mutable_operand(0),
                       ComposePermutations(operand->dimensions(),
                                           transpose->dimensions())));
    changed_ = true;
    return Status::OK();
  }

  if (is_layout_sensitive_ &&
      TransposeIsBitcast(transpose, valid_bitcast_callback_)) {
    ReplaceWithBitcast(transpose);
    return Status::OK();
  }

  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleConvolution(
    HloInstruction* convolution, HloInstruction* lhs, HloInstruction* rhs,
    const Window& window) {
  // HandleConvolution tries to replace a convolution with a DOT instruction.
  //
  // Only add when bitcasts can be used:
  // - if bitcasts are not supported, then reshapes could be used but will
  //   end up with another copy.
  // - if bitcasts are supported, the simplifier will be called again with
  //   bitcasts_ == true.

  // TODO(cwhipkey): b/31337498, make this layout insensitive.
  if (!is_layout_sensitive_) return Status::OK();

  const ConvolutionDimensionNumbers& dnums =
      convolution->convolution_dimension_numbers();
  const Shape& input_shape = lhs->shape();
  const Shape& filter_shape = rhs->shape();
  const Shape& convolution_shape = convolution->shape();
  TF_RET_CHECK(LayoutUtil::HasLayout(input_shape));
  TF_RET_CHECK(LayoutUtil::HasLayout(filter_shape));
  TF_RET_CHECK(LayoutUtil::HasLayout(convolution_shape));

  // Require 1x1 filter in the spatial dimensions (so no need to extract image
  // patches).
  if (filter_shape.dimensions(dnums.kernel_spatial_dimensions(0)) != 1 ||
      filter_shape.dimensions(dnums.kernel_spatial_dimensions(1)) != 1) {
    return Status::OK();
  }

  // Stride ignores part of the output, which matrix multiplication does not do,
  // so require no stride. Padding and base (lhs) dilation both implicitly
  // extend the data, which matrix multiplication also does not do, so require
  // no padding and no base (lhs) dilation. Window (rhs) dilation has no effect
  // for a 1x1 window, so window dilation is no problem.
  if (window_util::HasStride(window) || window_util::HasPadding(window) ||
      window_util::HasBaseDilation(window)) {
    return Status::OK();
  }

  // Also, the shapes must align for a rowmajor matmul:
  // - the input and output have the same layout.
  // - for input/output, the channel dimension must be the most minor. Other
  //   spatial dims can be in any order.
  // - for filters, the input channel dimension must be more major than the
  //   output channel dimension. The width+height don't matter because
  //   they are 1.
  //
  // These constraints are harsh. If the channel dimension is the most major
  // and/or the layout of input/output feature dimensions are reversed, we can
  // still convert Conv into more efficient Matmul with operand transposition
  // (such as the transposition flags in cuBLAS SGEMM).
  if (!LayoutUtil::Equal(input_shape.layout(), convolution_shape.layout()) ||
      input_shape.layout().minor_to_major(0) != dnums.feature_dimension() ||
      // The input feature dimension should come later in the minor-to-major
      // order.
      (PositionInContainer(AsInt64Slice(filter_shape.layout().minor_to_major()),
                           dnums.kernel_input_feature_dimension()) <
       PositionInContainer(AsInt64Slice(filter_shape.layout().minor_to_major()),
                           dnums.kernel_output_feature_dimension()))) {
    return Status::OK();
  }

  auto add_bitcast = [&](Shape shape, HloInstruction* operand) {
    std::vector<int64> dims(operand->shape().dimensions_size());
    std::iota(dims.begin(), dims.end(), 0);
    return computation_->AddInstruction(
        HloInstruction::CreateUnary(shape, HloOpcode::kBitcast, operand));
  };

  // Replace it with a dot, with bitcasts around it to get the right shape.
  const int64 input_channels =
      input_shape.dimensions(dnums.feature_dimension());
  const int64 output_channels =
      filter_shape.dimensions(dnums.kernel_output_feature_dimension());

  // Computes the product of the non-feature dimensions.
  int64 conv_width = 1;
  for (int i = 0; i < input_shape.dimensions_size(); ++i) {
    if (i != dnums.feature_dimension()) {
      conv_width *= input_shape.dimensions(i);
    }
  }

  // We already checked feature_dimension is most minor, so data in input_shape
  // and row-major {conv_width,input_channels} are bitwise identical.
  const Shape new_input_shape =
      ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout(
          input_shape.element_type(), {conv_width, input_channels});
  // We already checked input_feature_dimension is more major than
  // output_feature_dimension, so data in filter_shape and row-major
  // {input_channels,output_channels} are bitwise identical.
  const Shape new_filter_shape =
      ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout(
          filter_shape.element_type(), {input_channels, output_channels});
  const Shape dot_output_shape =
      ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout(
          convolution_shape.element_type(), {conv_width, output_channels});

  // We cannot insert bitcasts if the layouts will not be compatible.
  // TODO(b/33178038): Consider inserting a transpose if a bitcast would be
  // invalid.
  if (!valid_bitcast_callback_(lhs->shape(), input_shape) ||
      !valid_bitcast_callback_(rhs->shape(), new_filter_shape) ||
      !valid_bitcast_callback_(dot_output_shape, convolution_shape)) {
    return Status::OK();
  }

  auto new_lhs = add_bitcast(new_input_shape, lhs);
  auto new_rhs = add_bitcast(new_filter_shape, rhs);
  auto dot = computation_->AddInstruction(HloInstruction::CreateBinary(
      dot_output_shape, HloOpcode::kDot, new_lhs, new_rhs));
  computation_->ReplaceInstruction(convolution,
                                   add_bitcast(convolution_shape, dot));
  changed_ = true;
  return Status::OK();
}

bool AlgebraicSimplifierVisitor::TransformToClampIfSameShape(
    HloInstruction* root, HloInstruction* min, HloInstruction* min_operand,
    HloInstruction* operand, HloInstruction* max, HloInstruction* max_operand) {
  // Ensure shapes of min and max operand are equal to match current shape
  // inference.
  if (!SameShape(min_operand, max_operand)) {
    return false;
  }

  auto clamp = HloInstruction::CreateTernary(root->shape(), HloOpcode::kClamp,
                                             max_operand, operand, min_operand);
  computation_->ReplaceWithNewInstruction(root, std::move(clamp));
  changed_ = true;
  return true;
}

Status AlgebraicSimplifierVisitor::HandleMaximum(HloInstruction* maximum,
                                                 HloInstruction* lhs,
                                                 HloInstruction* rhs) {
  // Match the following tree:
  //          min_operand     operand
  //                     \   /
  //      max_operand     min
  //                 \   /
  //                  max
  // where max_operand and min_operand are scalar constants.
  {
    HloInstruction* min;
    HloInstruction* max_operand;
    HloInstruction* min_operand;
    HloInstruction* operand;

    if (hlo_query::MatchBinaryInstructionOperandOpcode(
            HloOpcode::kMinimum, maximum,
            /*matching_operand=*/&min,
            /*other_operand=*/&max_operand) &&
        hlo_query::MatchBinaryInstructionOperand(
            hlo_query::IsScalarConstant, min,
            /*matching_operand=*/&min_operand,
            /*other_operand=*/&operand) &&
        TransformToClampIfSameShape(maximum, min, min_operand, operand, maximum,
                                    max_operand)) {
      return Status::OK();
    }
  }

  return Status::OK();
}

Status AlgebraicSimplifierVisitor::HandleMinimum(HloInstruction* minimum,
                                                 HloInstruction* lhs,
                                                 HloInstruction* rhs) {
  // Match the following tree:
  //          max_operand     operand
  //                     \   /
  //      min_operand     max
  //                 \   /
  //                  min
  // where max_operand and min_operand are scalar constants.
  {
    HloInstruction* max;
    HloInstruction* max_operand;
    HloInstruction* min_operand;
    HloInstruction* operand;

    if (hlo_query::MatchBinaryInstructionOperandOpcode(
            HloOpcode::kMaximum, minimum,
            /*matching_operand=*/&max,
            /*other_operand=*/&min_operand) &&
        hlo_query::MatchBinaryInstructionOperand(
            hlo_query::IsScalarConstant, max,
            /*matching_operand=*/&max_operand,
            /*other_operand=*/&operand) &&
        TransformToClampIfSameShape(minimum, minimum, min_operand, operand, max,
                                    max_operand)) {
      return Status::OK();
    }
  }

  return Status::OK();
}

StatusOr<bool> AlgebraicSimplifier::Run(HloModule* module) {
  return std::any_of(
      module->computations().begin(), module->computations().end(),
      [=](const std::unique_ptr<HloComputation>& computation) {
        return AlgebraicSimplifierVisitor::Run(
            computation.get(), is_layout_sensitive_, valid_bitcast_callback_);
      });
}

}  // namespace xla