/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

==============================================================================*/

#include "mlir-hlo/Dialect/mhlo/IR/chlo_ops.h"
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir-hlo/Dialect/mhlo/transforms/map_chlo_to_hlo_op.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"

namespace mlir {
namespace {

// TODO(herhut): Generate these out of op definitions.
#define MAP_XLA_OPERATION_CWISE_UNARY(fn, sep)                                \
  fn(AbsOp) sep fn(CeilOp) sep fn(ClzOp) sep fn(ConvertOp) sep fn(CosOp)      \
      sep fn(ExpOp) sep fn(Expm1Op) sep fn(FloorOp) sep fn(ImagOp)            \
          sep fn(IsFiniteOp) sep fn(LogOp) sep fn(Log1pOp) sep fn(LogisticOp) \
              sep fn(NotOp) sep fn(NegOp) sep fn(PopulationCountOp)           \
                  sep fn(RealOp) sep fn(RoundOp) sep fn(RsqrtOp)              \
                      sep fn(SignOp) sep fn(SinOp) sep fn(SqrtOp)             \
                          sep fn(TanhOp)

// TODO(herhut): Generate these out of op definitions.
#define MAP_XLA_OPERATION_CWISE_BINARY(fn, sep)                            \
  fn(AddOp) sep fn(AndOp) sep fn(Atan2Op) sep fn(ComplexOp) sep fn(DivOp)  \
      sep fn(MaxOp) sep fn(MinOp) sep fn(MulOp) sep fn(OrOp) sep fn(PowOp) \
          sep fn(RemOp) sep fn(ShiftLeftOp) sep fn(ShiftRightArithmeticOp) \
              sep fn(ShiftRightLogicalOp) sep fn(SubOp) sep fn(XorOp)

// TODO(herhut): Generate these out of op definitions.
#define MAP_CHLO_OPERATION_CWISE_UNARY(fn, sep)                            \
  fn(AcosOp) sep fn(AcoshOp) sep fn(AsinOp) sep fn(AsinhOp) sep fn(AtanOp) \
      sep fn(AtanhOp) sep fn(ConjOp) sep fn(CoshOp) sep fn(DigammaOp)      \
          sep fn(ErfOp) sep fn(ErfcOp) sep fn(IsInfOp) sep fn(LgammaOp)    \
              sep fn(SinhOp) sep fn(TanOp)

// TODO(herhut): Generate these out of op definitions.
#define MAP_CHLO_OPERATION_CWISE_BINARY(fn, sep) fn(ZetaOp)

template <typename OpTy>
inline void AddLegalOpOnRankedTensor(ConversionTarget *target) {
  target->addDynamicallyLegalOp<OpTy>([](OpTy op) {
    return llvm::all_of(op.getOperation()->getOperandTypes(),
                        [&](Type t) { return t.isa<RankedTensorType>(); });
  });
}

/// Element-wise operations on unranked tensors can be applied to the flattened
/// tensor operands with the same effect.  This pattern rewrites every such
/// operation to
///   (i)   flatten the input tensor,
///   (ii)  apply the operation, and
///   (iii) restore the original shape.
template <typename OpTy>
struct ElementwiseOpConversion : public OpRewritePattern<OpTy> {
  explicit ElementwiseOpConversion(MLIRContext *context)
      : OpRewritePattern<OpTy>(context) {}

  LogicalResult matchAndRewrite(OpTy op,
                                PatternRewriter &rewriter) const override {
    // Don't apply conversion unless all operands are unranked.
    if (!llvm::all_of(op.getOperation()->getOperands(), [&](Value operand) {
          return operand.getType().isa<UnrankedTensorType>();
        })) {
      return failure();
    }

    // Get operands' shape.
    auto loc = op.getLoc();
    Type extentTensorTy = shape::getExtentTensorType(rewriter.getContext());
    SmallVector<Value, 3> operandShapes;
    for (Value operand : op.getOperation()->getOperands()) {
      Value shape =
          rewriter.create<shape::ShapeOfOp>(loc, extentTensorTy, operand);
      operandShapes.push_back(shape);
    }
    Value shape =
        operandShapes.size() == 1
            ? operandShapes.front()
            : rewriter.create<shape::AnyOp>(loc, extentTensorTy, operandShapes);

    // Derive flat shape.
    Type indexTy = rewriter.getIndexType();
    Value numElements =
        rewriter.create<shape::NumElementsOp>(loc, indexTy, shape);
    Value flatShape = rewriter.create<tensor::FromElementsOp>(loc, numElements);

    // Flatten operands.
    SmallVector<Value, 3> flatOperands;
    for (Value operand : op.getOperation()->getOperands()) {
      Type operandElementTy =
          operand.getType().template cast<ShapedType>().getElementType();
      Type flatTy =
          RankedTensorType::get({ShapedType::kDynamicSize}, operandElementTy);
      Value flat = rewriter.create<mhlo::DynamicReshapeOp>(loc, flatTy, operand,
                                                           flatShape);
      flatOperands.push_back(flat);
    }

    // Apply operation to flattened operands.
    Type resultElementTy =
        op.getType().template cast<ShapedType>().getElementType();
    Type flatResultTy =
        RankedTensorType::get({ShapedType::kDynamicSize}, resultElementTy);
    Value flatResult =
        rewriter.create<OpTy>(loc, flatResultTy, flatOperands, op->getAttrs());

    // Restore original shape.
    rewriter.replaceOpWithNewOp<mhlo::DynamicReshapeOp>(op, op.getType(),
                                                        flatResult, shape);

    return success();
  }
};

// Converts a broadcasting binary operation with a scalar operand and an
// unranked operand to a ranked broadcasting operation by dynamically reshaping
// the unranked operand to a 1D tensor. This will always be safe because
// broadcasting from a scalar to another shape always works.
template <typename ChloOpTy, typename HloOpTy, typename Adaptor>
struct ConvertUnrankedScalarDynamicBroadcastBinaryOp
    : public OpConversionPattern<ChloOpTy> {
  using OpConversionPattern<ChloOpTy>::OpConversionPattern;
  LogicalResult matchAndRewrite(
      ChloOpTy op, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    typename ChloOpTy::Adaptor transformed(operands);
    Value lhs = transformed.lhs();
    Value rhs = transformed.rhs();

    auto lhs_ranked_type = lhs.getType().dyn_cast<RankedTensorType>();
    auto lhs_unranked_type = lhs.getType().dyn_cast<UnrankedTensorType>();

    auto rhs_ranked_type = rhs.getType().dyn_cast<RankedTensorType>();
    auto rhs_unranked_type = rhs.getType().dyn_cast<UnrankedTensorType>();

    bool lhs_is_scalar = lhs_ranked_type &&
                         lhs_ranked_type.getShape().empty() &&
                         rhs_unranked_type;
    bool rhs_is_scalar = rhs_ranked_type &&
                         rhs_ranked_type.getShape().empty() &&
                         lhs_unranked_type;

    // Only support the case where exactly one operand is scalar and the other
    // is unranked. Other patterns in chlo-to-hlo legalization will create more
    // efficient lowerings for cases where both ranks are known or will handle
    // the more generic case of both inputs being unranked.
    if (!(lhs_is_scalar ^ rhs_is_scalar)) return failure();

    auto scalar_element_type = lhs_is_scalar ? lhs_ranked_type.getElementType()
                                             : rhs_ranked_type.getElementType();
    auto result_type = op.getResult().getType().template dyn_cast<TensorType>();
    auto result_element_type = result_type.getElementType();

    // Reshape the non-scalar value into a dynamically sized, rank-1 tensor
    Value shape =
        rewriter.create<shape::ShapeOfOp>(loc, lhs_is_scalar ? rhs : lhs);
    Value num_elements = rewriter.create<shape::NumElementsOp>(loc, shape);
    Value size_tensor =
        rewriter.create<tensor::FromElementsOp>(loc, num_elements);
    Value reshaped = rewriter.create<mhlo::DynamicReshapeOp>(
        loc, RankedTensorType::get({-1}, scalar_element_type),
        lhs_is_scalar ? rhs : lhs, size_tensor);

    // Create a new ranked Chlo op that will be further lowered by other
    // patterns into Mhlo.
    SmallVector<Value, 2> new_operands{lhs_is_scalar ? lhs : reshaped,
                                       rhs_is_scalar ? rhs : reshaped};
    Value computed = rewriter.create<ChloOpTy>(
        loc, TypeRange{RankedTensorType::get({-1}, result_element_type)},
        new_operands, op->getAttrs());

    // Reshape the result back into an unranked tensor.
    rewriter.replaceOpWithNewOp<mhlo::DynamicReshapeOp>(op, result_type,
                                                        computed, shape);

    return success();
  }
};

template <typename ChloOpTy, typename HloOpTy>
struct ConvertUnrankedDynamicBroadcastOpHelper {
  // Returns the dynamic result of checking the given value is effectively a
  // scalar shape (i.e. the number of elements is 1).
  static Value GreaterRankIsN(OpBuilder &builder, Location loc,
                              Value actual_rank, int targeted_rank) {
    return builder.create<CmpIOp>(
        loc, CmpIPredicate::eq, actual_rank,
        builder.create<ConstantIndexOp>(loc, targeted_rank));
  }

  static scf::IfOp createIfOpForRankSpecializedBroadcastAndOp(
      OpBuilder &builder, ChloOpTy op, Value actual_rank, int targeted_rank) {
    // Create the if block to place the current specialized logic in.
    Value greater_rank_is_n =
        GreaterRankIsN(builder, op.getLoc(), actual_rank, targeted_rank);
    return builder.create<scf::IfOp>(op.getLoc(), op.getResult().getType(),
                                     greater_rank_is_n, true);
  }

  static Value createBroadcastToKnownRank(OpBuilder &builder, ChloOpTy op,
                                          Value shape, int targeted_rank) {
    auto loc = op.getLoc();
    SmallVector<int64_t, 6> ranked_shape(targeted_rank, 1);
    auto unknown_rank_extent_tensor_type = RankedTensorType::get(
        {RankedTensorType::kDynamicSize}, builder.getIndexType());
    auto known_rank_extent_tensor_type =
        RankedTensorType::get({targeted_rank}, builder.getIndexType());
    Value ranked_shape_val = builder.create<shape::ConstShapeOp>(
        loc, known_rank_extent_tensor_type,
        mlir::DenseIntElementsAttr::get(known_rank_extent_tensor_type,
                                        ranked_shape));
    Value extended_value = builder.create<shape::BroadcastOp>(
        loc, unknown_rank_extent_tensor_type, shape, ranked_shape_val, nullptr);
    return builder.create<tensor::CastOp>(loc, known_rank_extent_tensor_type,
                                          extended_value);
  }

  // Create the if statement and code for a broadcasting op with a result of a
  // given rank.
  static void createRankSpecializedBroadcastAndOp(OpBuilder &if_builder,
                                                  ChloOpTy op,
                                                  ValueRange operands,
                                                  ValueRange operand_shapes,
                                                  int targeted_rank) {
    auto loc = op.getLoc();
    SmallVector<Value, 2> reshaped_operands;

    auto dynamic_dimensions = llvm::SmallVector<int64_t, 6>(
        targeted_rank, RankedTensorType::kDynamicSize);

    for (auto it : llvm::zip(operands, operand_shapes)) {
      Value operand, shape;
      std::tie(operand, shape) = it;
      // Handle shape broadcasting and inference.
      Value extended_operand_casted =
          createBroadcastToKnownRank(if_builder, op, shape, targeted_rank);

      // 1. Reshape operands to the given rank (with the same number of
      // elements)
      // 2. Compute the ranked-broadcasted ChloOp (which will assert that the
      // ops
      //    can be broadcasted and do the actual broadcasting)
      // 3. Type erase the output back to unranked
      auto reshaped_type = RankedTensorType::get(
          dynamic_dimensions,
          operand.getType().template dyn_cast<TensorType>().getElementType());
      Value reshaped_operand = if_builder.create<mhlo::DynamicReshapeOp>(
          loc, reshaped_type, operand, extended_operand_casted);
      reshaped_operands.push_back(reshaped_operand);
    }
    auto result_element_type = op.getResult()
                                   .getType()
                                   .template dyn_cast<TensorType>()
                                   .getElementType();
    auto result_type =
        RankedTensorType::get(dynamic_dimensions, result_element_type);
    Value result = if_builder.create<ChloOpTy>(
        loc, ArrayRef<Type>{result_type}, reshaped_operands, op->getAttrs());
    Value reshaped_result = if_builder.create<tensor::CastOp>(
        loc, UnrankedTensorType::get(result_element_type), result);
    if_builder.create<scf::YieldOp>(loc, reshaped_result);
  }

  // Iterates over the desired ranks to be specialized and generates the code
  // snippet for each case.
  static Value HandleBroadcastAndOp(OpBuilder &rewriter, ChloOpTy op,
                                    ValueRange operands) {
    auto loc = op.getLoc();

    // Get the minimum broadcast shapes of the operands.
    SmallVector<Value> shapes;
    shapes.reserve(operands.size());
    auto extent_tensor_type = RankedTensorType::get({ShapedType::kDynamicSize},
                                                    rewriter.getIndexType());
    for (Value operand : operands) {
      Value shape =
          rewriter.create<shape::ShapeOfOp>(loc, extent_tensor_type, operand);
      shapes.push_back(shape);
    }
    auto broadcast_shape = rewriter.create<shape::BroadcastOp>(
        loc, extent_tensor_type, shapes, nullptr);
    SmallVector<Type> result_types(shapes.size(), extent_tensor_type);
    auto reduced_shapes =
        rewriter
            .create<chlo::MinimumBroadcastShapesOp>(loc, result_types, shapes)
            .results();
    SmallVector<Value> reshaped_operands;
    reshaped_operands.reserve(operands.size());
    for (auto it : llvm::zip(operands, reduced_shapes)) {
      Value operand;
      Value reduced_shape;
      std::tie(operand, reduced_shape) = it;
      auto reshaped_operand = rewriter.create<mhlo::DynamicReshapeOp>(
          loc, operand.getType(), operand, reduced_shape);
      reshaped_operands.push_back(reshaped_operand);
    }

    // Find the largest rank of the operands.
    Value greater_rank;
    for (Value shape : reduced_shapes) {
      Value rank =
          rewriter.create<shape::RankOp>(loc, rewriter.getIndexType(), shape);
      if (!greater_rank) {
        greater_rank = rank;
      } else {
        Value greater_rank_compare = rewriter.create<CmpIOp>(
            loc, CmpIPredicate::sgt, greater_rank, rank);
        greater_rank = rewriter.create<SelectOp>(loc, greater_rank_compare,
                                                 greater_rank, rank);
      }
    }

    // Generate a list of nested if/else statements to handle rank
    // specializations from 1 to `kMaxRankSpecialization`.
    scf::IfOp if_op = createIfOpForRankSpecializedBroadcastAndOp(
        rewriter, op, greater_rank, 1);
    OpBuilder if_builder = if_op.getThenBodyBuilder(rewriter.getListener());
    createRankSpecializedBroadcastAndOp(if_builder, op, reshaped_operands,
                                        reduced_shapes, 1);

    // Put each subsequent rank specialization inside the else statement of the
    // previous one.
    OpBuilder else_builder = if_op.getElseBodyBuilder(rewriter.getListener());
    constexpr int kMaxRankSpecialization = 5;
    for (int i = 2; i < kMaxRankSpecialization; i++) {
      auto inner_if = createIfOpForRankSpecializedBroadcastAndOp(
          else_builder, op, greater_rank, i);
      if_builder = inner_if.getThenBodyBuilder(rewriter.getListener());
      createRankSpecializedBroadcastAndOp(if_builder, op, reshaped_operands,
                                          reduced_shapes, i);
      else_builder.create<scf::YieldOp>(loc, inner_if.getResult(0));
      else_builder = inner_if.getElseBodyBuilder(rewriter.getListener());
    }
    // Fire an assertion if none of the rank specializations applied (one of
    // the ranks was greater than `kMaxRankSpecialization`).
    else_builder.create<AssertOp>(
        loc,
        GreaterRankIsN(else_builder, op.getLoc(), greater_rank,
                       kMaxRankSpecialization),
        "Input for dynamic binary op lowering was of a rank greater than " +
            std::to_string(kMaxRankSpecialization));
    // Add the rank 5 specialization to the innermost else block.
    createRankSpecializedBroadcastAndOp(else_builder, op, reshaped_operands,
                                        reduced_shapes, kMaxRankSpecialization);

    // Return the reshaped result of the outermost if statement.
    auto result = if_op.getResult(0);
    auto reshaped_result = rewriter.create<mhlo::DynamicReshapeOp>(
        loc, result.getType(), result, broadcast_shape);
    return reshaped_result;
  }
};

// Handles lowering of the following pattern to patterns that will be further
// matched by other patterns until they result in LHLO:
//   %result = "chlo.op"(%lhs, %rhs) : (<*xTy>, <*xTy>) -> <*xTy>
//
// The sequence of specializations this handles is:
//   - Either operand being scalar
//   - Operands having equal shapes
//   - The resulting value being any of ranks [2,6]
template <typename ChloOpTy, typename HloOpTy, typename Adaptor>
struct ConvertUnrankedDynamicBroadcastBinaryOp
    : public OpConversionPattern<ChloOpTy> {
  using OpConversionPattern<ChloOpTy>::OpConversionPattern;

  LogicalResult matchAndRewrite(
      ChloOpTy op, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    typename ChloOpTy::Adaptor transformed(operands);
    Value lhs = transformed.lhs();
    Value rhs = transformed.rhs();
    auto lhs_type = lhs.getType().dyn_cast<UnrankedTensorType>();
    auto rhs_type = rhs.getType().dyn_cast<UnrankedTensorType>();
    auto result_type = op.getResult().getType().template dyn_cast<TensorType>();

    // Only support unranked operands. If either operand is ranked, another
    // pattern will handle the lowering.
    if (!lhs_type || !rhs_type) return failure();

    Value shape_of_lhs = rewriter.create<shape::ShapeOfOp>(loc, lhs);
    Value shape_of_rhs = rewriter.create<shape::ShapeOfOp>(loc, rhs);

    // If lhs has exactly one element
    auto if_op = rewriter.create<scf::IfOp>(
        loc, result_type, IsSingleElementShape(rewriter, op, shape_of_lhs),
        true);
    OpBuilder if_lhs_scalar_builder =
        if_op.getThenBodyBuilder(rewriter.getListener());
    Value reshaped_lhs = if_lhs_scalar_builder.create<mhlo::ReshapeOp>(
        loc, RankedTensorType::get({}, lhs_type.getElementType()), lhs);
    Value if_lhs_scalar_result = if_lhs_scalar_builder.create<ChloOpTy>(
        loc, ArrayRef<Type>{result_type}, ArrayRef<Value>{reshaped_lhs, rhs},
        op->getAttrs());
    Value extended_if_lhs_scalar_result =
        extendToBroadcastShape(if_lhs_scalar_builder, loc, if_lhs_scalar_result,
                               shape_of_lhs, shape_of_rhs);
    if_lhs_scalar_builder.create<scf::YieldOp>(loc,
                                               extended_if_lhs_scalar_result);

    // If lhs does not have exactly one element
    //
    // See if rhs has exactly one element
    OpBuilder else_lhs_scalar_builder =
        if_op.getElseBodyBuilder(rewriter.getListener());
    auto if_rhs_scalar_op = else_lhs_scalar_builder.create<scf::IfOp>(
        loc, result_type,
        IsSingleElementShape(else_lhs_scalar_builder, op, shape_of_rhs), true);
    else_lhs_scalar_builder.create<scf::YieldOp>(loc,
                                                 if_rhs_scalar_op.getResult(0));
    OpBuilder if_rhs_scalar_builder =
        if_rhs_scalar_op.getThenBodyBuilder(rewriter.getListener());
    Value reshaped_rhs = if_rhs_scalar_builder.create<mhlo::ReshapeOp>(
        loc, RankedTensorType::get({}, rhs_type.getElementType()), rhs);
    Value if_rhs_scalar_result = if_rhs_scalar_builder.create<ChloOpTy>(
        loc, ArrayRef<Type>{result_type}, ArrayRef<Value>{lhs, reshaped_rhs},
        op->getAttrs());
    Value extended_if_rhs_scalar_result =
        extendToBroadcastShape(if_rhs_scalar_builder, loc, if_rhs_scalar_result,
                               shape_of_lhs, shape_of_rhs);
    if_rhs_scalar_builder.create<scf::YieldOp>(loc,
                                               extended_if_rhs_scalar_result);

    // If NEITHER shape has exactly one element
    //
    // See if shapes are equal.
    OpBuilder else_no_scalars_builder =
        if_rhs_scalar_op.getElseBodyBuilder(rewriter.getListener());
    Value equal_shapes = else_no_scalars_builder.create<shape::ShapeEqOp>(
        loc, shape_of_lhs, shape_of_rhs);

    auto if_eq_shapes_op = else_no_scalars_builder.create<scf::IfOp>(
        loc, result_type, equal_shapes, true);
    else_no_scalars_builder.create<scf::YieldOp>(loc,
                                                 if_eq_shapes_op.getResult(0));

    OpBuilder if_eq_shapes_builder =
        if_eq_shapes_op.getThenBodyBuilder(rewriter.getListener());
    Value non_broadcast_op =
        Adaptor::CreateOp(op, result_type, lhs, rhs, if_eq_shapes_builder);
    if_eq_shapes_builder.create<scf::YieldOp>(loc, non_broadcast_op);

    // If shapes do not have exactly one element, nor are equal
    //
    // See if values are of a rank that we support.
    OpBuilder if_neq_shapes_builder =
        if_eq_shapes_op.getElseBodyBuilder(rewriter.getListener());
    if_neq_shapes_builder.create<scf::YieldOp>(
        loc, ConvertUnrankedDynamicBroadcastOpHelper<
                 ChloOpTy, HloOpTy>::HandleBroadcastAndOp(if_neq_shapes_builder,
                                                          op, {lhs, rhs}));

    rewriter.replaceOp(op, {if_op.getResult(0)});
    return success();
  }

 private:
  // Returns the dynamic result of checking the given value is effectively a
  // scalar shape (i.e. the number of elements is 1).
  Value IsSingleElementShape(OpBuilder &rewriter, ChloOpTy op,
                             Value shape_of_tensor) const {
    auto loc = op.getLoc();

    Value num_elements =
        rewriter.create<shape::NumElementsOp>(loc, shape_of_tensor);
    return rewriter.create<CmpIOp>(loc, rewriter.getI1Type(), CmpIPredicate::eq,
                                   num_elements,
                                   rewriter.create<ConstantIndexOp>(loc, 1));
  }

  Value extendToBroadcastShape(OpBuilder &builder, Location loc, Value value,
                               Value shape_of_lhs, Value shape_of_rhs) const {
    auto unknown_rank_extent_tensor_type = RankedTensorType::get(
        {RankedTensorType::kDynamicSize}, builder.getIndexType());
    Value broadcast_shape =
        builder.create<shape::BroadcastOp>(loc, unknown_rank_extent_tensor_type,
                                           shape_of_lhs, shape_of_rhs, nullptr);
    return builder.create<mhlo::DynamicReshapeOp>(loc, value.getType(), value,
                                                  broadcast_shape);
  }
};

// Rank-specialize chlo.broadcast_select ops.
struct ConvertUnrankedDynamicBroadcastSelectOp
    : public OpConversionPattern<chlo::BroadcastSelectOp> {
  using OpConversionPattern<chlo::BroadcastSelectOp>::OpConversionPattern;

  LogicalResult matchAndRewrite(
      chlo::BroadcastSelectOp op, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    // For now only do the bare minimum and specialize for every rank. There is
    // more potential for optimization here. This also is missing the
    // specialization for rank 0.
    rewriter.replaceOp(
        op, {ConvertUnrankedDynamicBroadcastOpHelper<
                chlo::BroadcastSelectOp,
                mhlo::SelectOp>::HandleBroadcastAndOp(rewriter, op, operands)});
    return success();
  }
};

struct TransformUnrankedHloPass
    : public PassWrapper<TransformUnrankedHloPass, FunctionPass> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<chlo::HloClientDialect, mhlo::MhloDialect,
                    shape::ShapeDialect>();
  }

  void runOnFunction() override {
    // Setup conversion target.
    MLIRContext &ctx = getContext();
    ConversionTarget target(ctx);
    target.addLegalDialect<chlo::HloClientDialect, mhlo::MhloDialect,
                           StandardOpsDialect, shape::ShapeDialect,
                           scf::SCFDialect, tensor::TensorDialect>();
    target.addLegalOp<FuncOp>();
#define ADD_LEGAL_MHLO(op) AddLegalOpOnRankedTensor<mhlo::op>(&target)
#define ADD_LEGAL_CHLO(op) AddLegalOpOnRankedTensor<chlo::op>(&target)
    MAP_XLA_OPERATION_CWISE_UNARY(ADD_LEGAL_MHLO, ;);
    MAP_XLA_OPERATION_CWISE_BINARY(ADD_LEGAL_MHLO, ;);
    MAP_CHLO_OPERATION_CWISE_UNARY(ADD_LEGAL_CHLO, ;);
#undef ADD_LEGAL_MHLO
#undef ADD_LEGAL_CHLO
    AddLegalOpOnRankedTensor<mhlo::CompareOp>(&target);
    AddLegalOpOnRankedTensor<mhlo::SelectOp>(&target);
    target.addDynamicallyLegalDialect<chlo::HloClientDialect>(
        [](Operation *op) {
          return !llvm::any_of(op->getOperandTypes(), [](Type type) {
            return type.isa<UnrankedTensorType>();
          });
        });

    // Populate rewrite patterns.
    OwningRewritePatternList patterns;
    mhlo::PopulateTransformUnrankedHloPatterns(&ctx, &patterns);

    // Apply transformation.
    if (failed(
            applyPartialConversion(getFunction(), target, std::move(patterns))))
      return signalPassFailure();
  }
};

}  // namespace

namespace mhlo {

void PopulateTransformUnrankedHloPatterns(MLIRContext *context,
                                          OwningRewritePatternList *patterns) {
#define MAP_HLO(op) ElementwiseOpConversion<mhlo::op>
#define MAP_CHLO(op) ElementwiseOpConversion<chlo::op>
#define COMMA ,
  // clang-format off
  patterns->insert<
      MAP_XLA_OPERATION_CWISE_UNARY(MAP_HLO, COMMA),
      MAP_XLA_OPERATION_CWISE_BINARY(MAP_HLO, COMMA),
      MAP_CHLO_OPERATION_CWISE_UNARY(MAP_CHLO, COMMA),
      MAP_CHLO_OPERATION_CWISE_BINARY(MAP_CHLO, COMMA),
      ElementwiseOpConversion<mhlo::CompareOp>,
      ElementwiseOpConversion<mhlo::SelectOp>>(context);
  // clang-format on
#undef MAP_HLO
#undef MAP_CHLO
#undef COMMA
  chlo::PopulateForBroadcastingBinaryOp<
      ConvertUnrankedDynamicBroadcastBinaryOp>(context, patterns);
  chlo::PopulateForBroadcastingBinaryOp<
      ConvertUnrankedScalarDynamicBroadcastBinaryOp>(context, patterns);
  patterns->insert<ConvertUnrankedDynamicBroadcastSelectOp>(context);
}

std::unique_ptr<FunctionPass> createTransformUnrankedHloPass() {
  return std::make_unique<TransformUnrankedHloPass>();
}

}  // namespace mhlo
}  // namespace mlir