// Copyright 2020 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// Implements StableHLO-to-StableHLO IREE-specific preprocessing.

#include <numeric>
#include <random>

#include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Passes.h"
#include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Rewriters.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"

namespace mlir::iree_compiler::stablehlo {

#define GEN_PASS_DEF_STABLEHLOTOSTABLEHLOPREPROCESSING
#include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Passes.h.inc"

namespace {

bool isIota(ArrayRef<int64_t> array) {
  for (auto [idx, value] : llvm::enumerate(array)) {
    if (static_cast<int64_t>(idx) != value)
      return false;
  }
  return true;
}

DenseI64ArrayAttr make1DElementsAttr(OpBuilder &b, ArrayRef<int64_t> integers) {
  return b.getDenseI64ArrayAttr(integers);
}

Value getF32Const(ImplicitLocOpBuilder b, ArrayRef<int64_t> shapes,
                  ArrayRef<float> values) {
  RankedTensorType ty = RankedTensorType::get(shapes, b.getF32Type());
  return b
      .create<mlir::stablehlo::ConstantOp>(DenseFPElementsAttr::get(ty, values))
      .getResult();
}

// Guarantee that the input dimensions are ordered batch, spatial_dims, feature
// dim.
struct ReorderConvOpInputDimensions final
    : OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
                                PatternRewriter &rewriter) const override {
    auto lhsType = cast<ShapedType>(op.getLhs().getType());
    auto lhsShape = lhsType.getShape();
    if (!lhsType.hasRank()) {
      return failure();
    }

    auto dimensionNumbers = op.getDimensionNumbers();
    auto spatialDims = dimensionNumbers.getInputSpatialDimensions();

    // Compute the permutation required to create a standard order.
    llvm::SmallVector<int64_t> permutations;
    permutations.push_back(dimensionNumbers.getInputBatchDimension());
    llvm::append_range(permutations, spatialDims);
    permutations.push_back(dimensionNumbers.getInputFeatureDimension());

    // If the permutation is iota then no reordering is required.
    if (isIota(permutations)) {
      return failure();
    }

    llvm::SmallVector<int64_t> transposeShape;
    for (int64_t p : permutations) {
      transposeShape.push_back(lhsShape[p]);
    }

    auto transposed = rewriter.create<mlir::stablehlo::TransposeOp>(
        op.getLoc(),
        RankedTensorType::get(transposeShape, lhsType.getElementType()),
        op.getLhs(), rewriter.getDenseI64ArrayAttr(permutations));

    llvm::SmallVector<int64_t> newSpatialDimensions(spatialDims.size());
    std::iota(newSpatialDimensions.begin(), newSpatialDimensions.end(), 1);

    auto newDimensionNumbers = mlir::stablehlo::ConvDimensionNumbersAttr::get(
        op.getContext(),
        /*input_batch_dimension=*/0,
        /*input_feature_dimension=*/newSpatialDimensions.size() + 1,
        /*input_spatial_dimensions=*/newSpatialDimensions,
        dimensionNumbers.getKernelInputFeatureDimension(),
        dimensionNumbers.getKernelOutputFeatureDimension(),
        dimensionNumbers.getKernelSpatialDimensions(),
        dimensionNumbers.getOutputBatchDimension(),
        dimensionNumbers.getOutputFeatureDimension(),
        dimensionNumbers.getOutputSpatialDimensions());

    SmallVector<Value, 2> operands = {transposed, op.getRhs()};
    auto newConv = rewriter.create<mlir::stablehlo::ConvolutionOp>(
        op.getLoc(), op.getType(), operands, op->getAttrs());
    newConv.setDimensionNumbersAttr(newDimensionNumbers);
    rewriter.replaceOp(op, newConv.getResult());

    return success();
  }
};

struct ReorderConvOpKernelDimensions final
    : OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
  using OpRewritePattern::OpRewritePattern;
  LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
                                PatternRewriter &rewriter) const override {
    auto kernel = op.getRhs();
    auto kernelType = cast<ShapedType>(kernel.getType());
    if (!kernelType.hasRank())
      return failure();
    auto kernelShape = kernelType.getShape();

    auto dimensionNumbers = op.getDimensionNumbers();

    auto spatialDims = dimensionNumbers.getKernelSpatialDimensions();

    auto inputFeatureDimension =
        dimensionNumbers.getKernelInputFeatureDimension();
    auto outputFeatureDimension =
        dimensionNumbers.getKernelOutputFeatureDimension();

    // Compute the permutation for the transpose.
    llvm::SmallVector<int64_t> permutation(spatialDims.begin(),
                                           spatialDims.end());
    permutation.push_back(inputFeatureDimension);
    permutation.push_back(outputFeatureDimension);

    // If the permutation is iota, then no transpose is required.
    if (isIota(permutation))
      return failure();

    llvm::SmallVector<int64_t> transposeShape;
    for (int64_t perm : permutation) {
      transposeShape.push_back(kernelShape[perm]);
    }

    llvm::SmallVector<int64_t> newSpatialDimensions(spatialDims.size());
    std::iota(newSpatialDimensions.begin(), newSpatialDimensions.end(), 0);

    auto transposeKernel = rewriter.create<mlir::stablehlo::TransposeOp>(
        op.getLoc(),
        RankedTensorType::get(transposeShape, kernelType.getElementType()),
        kernel, rewriter.getDenseI64ArrayAttr(permutation));

    auto newDimensionNumbers = mlir::stablehlo::ConvDimensionNumbersAttr::get(
        op.getContext(), dimensionNumbers.getInputBatchDimension(),
        dimensionNumbers.getInputFeatureDimension(),
        dimensionNumbers.getInputSpatialDimensions(),
        /*kernel_input_feature_dimension=*/
        newSpatialDimensions.size(),
        /*kernel_output_feature_dimension=*/
        newSpatialDimensions.size() + 1, newSpatialDimensions,
        dimensionNumbers.getOutputBatchDimension(),
        dimensionNumbers.getOutputFeatureDimension(),
        dimensionNumbers.getOutputSpatialDimensions());

    SmallVector<Value, 2> operands = {op.getLhs(), transposeKernel};
    mlir::stablehlo::ConvolutionOp newConv =
        rewriter.create<mlir::stablehlo::ConvolutionOp>(
            op.getLoc(), op.getType(), operands, op->getAttrs());
    newConv.setDimensionNumbersAttr(newDimensionNumbers);

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

// Guarantee that the output dimensions are ordered batch, spatial_dims, feature
// dim.
struct ReorderConvOpOutputDimensions final
    : OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
                                PatternRewriter &rewriter) const override {
    auto resultType = llvm::cast<ShapedType>(op.getType());
    auto resultShape = resultType.getShape();
    if (!resultType.hasRank()) {
      return failure();
    }

    auto dimensionNumbers = op.getDimensionNumbers();
    auto spatialDims = dimensionNumbers.getOutputSpatialDimensions();

    // Compute the permutation to transpose to an ordered output.
    llvm::SmallVector<int64_t> permutation;
    permutation.push_back(dimensionNumbers.getOutputBatchDimension());
    permutation.append(spatialDims.begin(), spatialDims.end());
    permutation.push_back(dimensionNumbers.getOutputFeatureDimension());

    // If the permutation is iota then no reordering is required.
    if (isIota(permutation)) {
      return failure();
    }

    // Compute what the new conv shape should be.
    llvm::SmallVector<int64_t> convShape;
    for (auto p : permutation) {
      convShape.push_back(resultShape[p]);
    }

    // Compute the inverse transpose to unordered and ordered output.
    llvm::SmallVector<int64_t> invertPermutation(permutation.size());
    for (auto it : llvm::enumerate(permutation)) {
      invertPermutation[it.value()] = it.index();
    }

    llvm::SmallVector<int64_t> newSpatialDimensions(spatialDims.size());
    std::iota(newSpatialDimensions.begin(), newSpatialDimensions.end(), 1);

    auto newDimensionNumbers = mlir::stablehlo::ConvDimensionNumbersAttr::get(
        op.getContext(), dimensionNumbers.getInputBatchDimension(),
        dimensionNumbers.getInputFeatureDimension(),
        dimensionNumbers.getInputSpatialDimensions(),
        dimensionNumbers.getKernelInputFeatureDimension(),
        dimensionNumbers.getKernelOutputFeatureDimension(),
        dimensionNumbers.getKernelSpatialDimensions(),
        /*output_batch_dimension=*/0,
        /*output_feature_dimension=*/newSpatialDimensions.size() + 1,
        /*output_spatial_dimensions=*/newSpatialDimensions);

    SmallVector<Value, 2> operands = {op.getLhs(), op.getRhs()};
    auto newConv = rewriter.create<mlir::stablehlo::ConvolutionOp>(
        op.getLoc(),
        RankedTensorType::get(convShape, resultType.getElementType()), operands,
        op->getAttrs());
    newConv.setDimensionNumbersAttr(newDimensionNumbers);

    auto transposed = rewriter.create<mlir::stablehlo::TransposeOp>(
        op.getLoc(), resultType, newConv,
        rewriter.getDenseI64ArrayAttr(invertPermutation));

    rewriter.replaceOp(op, transposed.getResult());
    return success();
  }
};

bool isConsecutive(ArrayRef<int64_t> array) {
  for (size_t i = 1, e = array.size(); i < e; ++i) {
    if (array[i] - array[i - 1] != 1)
      return false;
  }
  return true;
}

/// Rewrites stablehlo.dot_general so lhs contraction dimensions are innermost
/// and rhs contraction dimensions are dims right after batch dimension. The
/// pattern inserts transposes so the dot_general always has the form:
/// {batch_dims, parallel_dims, contraction_dims}.
///   {batch_dims, contraction_dims, parallel_dims}
/// After that, batch_dims, contraction_dims, parallel_dims are
/// in consecutive order and not spliting the domain. This pattern inserts
/// reshapes to collapse consecutive reduction and parallel dims to always
/// generate a rank-3 dot_general op.
struct TransposeReshapeGenericDotGeneral final
    : OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
  using OpRewritePattern::OpRewritePattern;

  Value TransposeIfNonConsecutive(OpBuilder &b, Location loc, Value src,
                                  ArrayRef<int64_t> targetOrder) const {
    if (isConsecutive(targetOrder))
      return src;

    auto type = cast<RankedTensorType>(src.getType());
    SmallVector<int64_t> transposeShape;
    for (int64_t i : targetOrder) {
      transposeShape.push_back(type.getDimSize(i));
    }
    return b.create<mlir::stablehlo::TransposeOp>(
        loc, RankedTensorType::get(transposeShape, type.getElementType()), src,
        b.getDenseI64ArrayAttr(targetOrder));
  }

  Value ReshapeIfNonStandard(OpBuilder &b, Location loc, Value src,
                             size_t dimsBorder0, size_t dimsBorder1) const {
    auto type = cast<RankedTensorType>(src.getType());
    ArrayRef<int64_t> shape = type.getShape();
    if (dimsBorder0 <= 1 && dimsBorder1 - dimsBorder0 <= 1 &&
        shape.size() - dimsBorder1 <= 1)
      return src;

    SmallVector<int64_t> result_shape = {
        std::accumulate(shape.begin(), shape.begin() + dimsBorder0, 1,
                        std::multiplies<int64_t>()),
        std::accumulate(shape.begin() + dimsBorder0,
                        shape.begin() + dimsBorder1, 1,
                        std::multiplies<int64_t>()),
        std::accumulate(shape.begin() + dimsBorder1, shape.end(), 1,
                        std::multiplies<int64_t>())};
    return b.create<mlir::stablehlo::ReshapeOp>(
        loc, RankedTensorType::get(result_shape, type.getElementType()), src);
  }

  LogicalResult matchAndRewrite(mlir::stablehlo::DotGeneralOp op,
                                PatternRewriter &rewriter) const override {
    auto lhsShapeType = dyn_cast<RankedTensorType>(op.getLhs().getType());
    auto rhsShapeType = dyn_cast<RankedTensorType>(op.getRhs().getType());
    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
    if (!lhsShapeType || !rhsShapeType || !resultType)
      return failure();

    // TODO(jpienaar): This pattern is not safe for dynamic shapes and seems to
    // be (now) redundant with later pass that does handle them. To decouple
    // fixing and verifying redundant, this just limits to static shapes and
    // then will remove this in follow up.
    if (!lhsShapeType.hasStaticShape() || !rhsShapeType.hasStaticShape())
      return failure();

    SmallVector<int64_t> lhsTargetOrder, rhsTargetOrder;
    mlir::stablehlo::DotDimensionNumbersAttr dimNumbers =
        op.getDotDimensionNumbers();
    ArrayRef<int64_t> lhsBatchingDims = dimNumbers.getLhsBatchingDimensions();
    ArrayRef<int64_t> lhsContractingDims =
        dimNumbers.getLhsContractingDimensions();
    ArrayRef<int64_t> rhsBatchingDims = dimNumbers.getRhsBatchingDimensions();
    ArrayRef<int64_t> rhsContractingDims =
        dimNumbers.getRhsContractingDimensions();

    // No contraction dims means this can be represented as a mul.
    if (lhsContractingDims.empty() || rhsContractingDims.empty()) {
      return rewriter.notifyMatchFailure(op,
                                         "can be represented as stablehlo.mul");
    }

    // No batching dimensions means this can be represented a dot.
    if (lhsBatchingDims.empty() || rhsBatchingDims.empty()) {
      return rewriter.notifyMatchFailure(op,
                                         "can be represented as stablehlo.dot");
    }

    SmallVector<bool> isLhsParallel(lhsShapeType.getRank(), true);
    for (int64_t i : lhsBatchingDims) {
      lhsTargetOrder.push_back(i);
      isLhsParallel[i] = false;
    }
    for (int64_t i : lhsContractingDims) {
      isLhsParallel[i] = false;
    }
    for (int64_t i = 0, e = lhsShapeType.getRank(); i < e; ++i) {
      if (isLhsParallel[i]) {
        lhsTargetOrder.push_back(i);
      }
    }
    for (int64_t i : lhsContractingDims) {
      lhsTargetOrder.push_back(i);
    }

    SmallVector<bool> isRhsParallel(rhsShapeType.getRank(), true);

    for (int64_t i : rhsBatchingDims) {
      rhsTargetOrder.push_back(i);
      isRhsParallel[i] = false;
    }
    for (int64_t i : rhsContractingDims) {
      rhsTargetOrder.push_back(i);
      isRhsParallel[i] = false;
    }
    for (int64_t i = 0, e = rhsShapeType.getRank(); i < e; ++i) {
      if (isRhsParallel[i]) {
        rhsTargetOrder.push_back(i);
      }
    }

    Value lhs = TransposeIfNonConsecutive(rewriter, op.getLoc(), op.getLhs(),
                                          lhsTargetOrder);
    Value rhs = TransposeIfNonConsecutive(rewriter, op.getLoc(), op.getRhs(),
                                          rhsTargetOrder);

    // The dimensions of this will always be transposed into {batch_dims,
    // parallel_dims, contraction_dims}, and the
    // following logic is based on this assumption.
    // TODO(#7443): If we consider transpose performance, the above assumptions
    // may not be true.
    int64_t numLhsContractionDims = lhsContractingDims.size();
    int64_t lhsContractionBase = lhsShapeType.getRank() - numLhsContractionDims;
    int64_t rhsContractionBase = rhsBatchingDims.size();
    int64_t numRhsContractionDims =
        rhsContractionBase + rhsContractingDims.size();

    lhs = ReshapeIfNonStandard(rewriter, op.getLoc(), lhs,
                               lhsBatchingDims.size(), lhsContractionBase);
    rhs = ReshapeIfNonStandard(rewriter, op.getLoc(), rhs,
                               rhsBatchingDims.size(), numRhsContractionDims);

    if (lhs == op.getLhs() && rhs == op.getRhs())
      return rewriter.notifyMatchFailure(op, "already in canonical form");

    auto dimensionNumbers = mlir::stablehlo::DotDimensionNumbersAttr::get(
        rewriter.getContext(), /*lhsBatchingDimensions=*/0,
        /*rhsBatchingDimensions=*/0,
        /*lhsContractingDimensions=*/
        cast<ShapedType>(lhs.getType()).getRank() - 1,
        /*rhsContractingDimensions=*/1);
    auto lhsNewType = cast<RankedTensorType>(lhs.getType());
    auto rhsNewType = cast<RankedTensorType>(rhs.getType());

    // batching、lhs parallel、rhs parallel this order is a conversion
    SmallVector<int64_t, 3> newShape = {lhsNewType.getShape()[0]};

    if (lhsNewType.getRank() > 2)
      newShape.push_back(lhsNewType.getDimSize(1));

    if (rhsNewType.getRank() > 2)
      newShape.push_back(rhsNewType.getDimSize(2));

    TensorType newResultType =
        RankedTensorType::get(newShape, resultType.getElementType());

    auto newOp = rewriter.create<mlir::stablehlo::DotGeneralOp>(
        op.getLoc(), newResultType, lhs, rhs, dimensionNumbers,
        op.getPrecisionConfigAttr(), op.getAlgorithmAttr());

    // Copy over unknown attributes as we currently rely on it to let user tune
    // lowering parameters.
    ArrayRef<StringRef> odsAttrs = op.getAttributeNames();
    for (NamedAttribute kv : op->getAttrs()) {
      if (!llvm::is_contained(odsAttrs, kv.getName().getValue())) {
        newOp->setAttr(kv.getName(), kv.getValue());
      }
    }

    Value result = newOp.getResult();
    if (op.getType() != newResultType) {
      result = rewriter.create<mlir::stablehlo::ReshapeOp>(
          op.getLoc(), op.getType(), newOp.getResult());
    }

    rewriter.replaceOp(op, result);
    return success();
  }
};

struct ScatterInt64Indices final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    auto indices = op.getScatterIndices();
    auto indicesTy = indices.getType();
    auto indicesETy = indicesTy.getElementType();
    if (indicesETy.isInteger(32)) {
      return rewriter.notifyMatchFailure(op, "already has i32 index type");
    }

    if (!indicesTy.hasStaticShape()) {
      return rewriter.notifyMatchFailure(op, "cannot validate legal size");
    }

    uint64_t maxSize = std::numeric_limits<int32_t>::max();
    if (indicesETy.getIntOrFloatBitWidth() > 32) {
      for (int i = 0, s = indicesTy.getRank(); i < s; ++i) {
        if (indicesTy.getDimSize(i) > maxSize) {
          return rewriter.notifyMatchFailure(op, "index may exceed i32 max");
        }
      }
    }

    indices = rewriter.create<mlir::stablehlo::ConvertOp>(
        op.getLoc(), indicesTy.clone(rewriter.getI32Type()), indices);

    auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
        op.getLoc(), op.getResultTypes(), op.getInputs(), indices,
        op.getUpdates(), op.getScatterDimensionNumbers(),
        op.getIndicesAreSorted(), op.getUniqueIndices());

    Region &region = newScatter.getUpdateComputation();
    rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
    rewriter.replaceOp(op, newScatter.getResults());

    return success();
  }
};

// If the indices tensor has an implicit index vector dim we expand and make it
// an explicit dim.
struct ScatterImplicitIndex final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    auto dimNumbers = op.getScatterDimensionNumbers();
    auto indexVectorDim = dimNumbers.getIndexVectorDim();
    Value indices = op.getScatterIndices();
    auto indicesTy = llvm::cast<ShapedType>(indices.getType());

    // Check indices vector has an implicit dim.
    if (indexVectorDim != indicesTy.getRank()) {
      return rewriter.notifyMatchFailure(op, "no implicit index dim");
    }

    // Materialize the implicit indices dim.
    SmallVector<ReassociationExprs> reassociationMap;
    reassociationMap.resize(indicesTy.getRank());
    SmallVector<int64_t> newShape;
    for (int i = 0, s = indicesTy.getRank(); i < s; ++i) {
      reassociationMap[i].push_back(rewriter.getAffineDimExpr(i));
      newShape.push_back(indicesTy.getDimSize(i));
    }
    if (!reassociationMap.empty()) {
      reassociationMap.back().push_back(
          rewriter.getAffineDimExpr(indicesTy.getRank()));
    }
    newShape.push_back(1);
    indicesTy = RankedTensorType::get(newShape, indicesTy.getElementType());
    indices = rewriter.create<tensor::ExpandShapeOp>(op.getLoc(), indicesTy,
                                                     indices, reassociationMap);

    auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
        op.getLoc(), op.getResultTypes(), op.getInputs(), indices,
        op.getUpdates(), dimNumbers, op.getIndicesAreSorted(),
        op.getUniqueIndices());
    Region &region = newScatter.getUpdateComputation();
    rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
    rewriter.replaceOp(op, newScatter.getResults());
    return success();
  }
};

struct ScatterImplicitBatch final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using OpRewritePattern::OpRewritePattern;

  static Value addUnitBatchDim(Location loc, Value value,
                               PatternRewriter &rewriter) {
    auto valueTy = cast<ShapedType>(value.getType());
    if (!valueTy.hasRank())
      return nullptr;

    // Materialize the implicit indices dim.
    SmallVector<ReassociationExprs> reassociationMap(valueTy.getRank());
    if (!reassociationMap.empty()) {
      reassociationMap.front().push_back(rewriter.getAffineDimExpr(0));
    }

    SmallVector<int64_t> newShape = {1};
    for (int i = 0, s = valueTy.getRank(); i < s; i++) {
      reassociationMap[i].push_back(rewriter.getAffineDimExpr(i + 1));
      newShape.push_back(valueTy.getDimSize(i));
    }

    valueTy = RankedTensorType::get(newShape, valueTy.getElementType());
    return rewriter.create<tensor::ExpandShapeOp>(loc, valueTy, value,
                                                  reassociationMap);
  }

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    auto dimNumbers = op.getScatterDimensionNumbers();
    auto indexVectorDim = dimNumbers.getIndexVectorDim();
    auto indices = llvm::cast<Value>(op.getScatterIndices());
    auto indicesTy = llvm::dyn_cast<RankedTensorType>(indices.getType());

    // Check whether indices has no batch dimension.
    if (!indicesTy)
      return failure();
    if (indicesTy.getRank() != 1 || indexVectorDim != 0) {
      return rewriter.notifyMatchFailure(op,
                                         "no implicit batch dimension to add.");
    }

    indices = addUnitBatchDim(op.getLoc(), indices, rewriter);
    if (!indices) {
      return rewriter.notifyMatchFailure(
          op, "Unable to add implicit batch dim to indices.");
    }

    llvm::SmallVector<int64_t> newUpdateWindowDims;
    for (auto dim : dimNumbers.getUpdateWindowDims()) {
      // Batch dimension is inserted at the start so window dimensions are shift
      // forwards.
      newUpdateWindowDims.push_back(dim + 1);
    }

    llvm::SmallVector<Value> updates;
    for (Value update : op.getUpdates()) {
      update = addUnitBatchDim(op.getLoc(), update, rewriter);
      if (!update) {
        return rewriter.notifyMatchFailure(
            op, "Unable to add implicit batch dim to update.");
      }
      updates.push_back(update);
    }

    auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
        op.getContext(), newUpdateWindowDims,
        dimNumbers.getInsertedWindowDims(), dimNumbers.getInputBatchingDims(),
        dimNumbers.getScatterIndicesBatchingDims(),
        dimNumbers.getScatterDimsToOperandDims(),
        dimNumbers.getIndexVectorDim() + 1);

    auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
        op.getLoc(), op.getResultTypes(), op.getInputs(), indices, updates,
        newDimNumbers, op.getIndicesAreSorted(), op.getUniqueIndices());
    Region &region = newScatter.getUpdateComputation();
    rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
    rewriter.replaceOp(op, newScatter.getResults());
    return success();
  }
};

struct ScatterCollapseBatch final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using OpRewritePattern::OpRewritePattern;

  static Value collapseBatchDims(Location loc, Value value, int64_t batchCount,
                                 PatternRewriter &rewriter) {
    auto valueTy = dyn_cast<ShapedType>(value.getType());
    if (!valueTy)
      return nullptr;

    SmallVector<ReassociationExprs> reassociationMap(1);
    reassociationMap.reserve(valueTy.getRank() - batchCount + 1);
    int64_t batchSize = 1;
    for (int i = 0, s = batchCount; i < s; i++) {
      reassociationMap.front().push_back(rewriter.getAffineDimExpr(i));
      bool isDynamic =
          valueTy.isDynamicDim(i) || ShapedType::isDynamic(batchSize);
      batchSize =
          isDynamic ? ShapedType::kDynamic : valueTy.getDimSize(i) * batchSize;
    }

    SmallVector<int64_t> newShape = {batchSize};
    for (int i = batchCount, s = valueTy.getRank(); i < s; i++) {
      reassociationMap.push_back({rewriter.getAffineDimExpr(i)});
      newShape.push_back(valueTy.getDimSize(i));
    }

    valueTy = RankedTensorType::get(newShape, valueTy.getElementType());
    return rewriter.create<tensor::CollapseShapeOp>(loc, valueTy, value,
                                                    reassociationMap);
  }

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    auto dimNumbers = op.getScatterDimensionNumbers();
    auto indexVectorDim = dimNumbers.getIndexVectorDim();
    auto indices = llvm::cast<Value>(op.getScatterIndices());
    auto indicesTy = llvm::cast<ShapedType>(indices.getType());
    auto updatedWindowDims = dimNumbers.getUpdateWindowDims();

    if (!indicesTy.hasRank()) {
      return rewriter.notifyMatchFailure(op, "indices has unknown rank");
    }

    // Check for an explicit index dimension.
    if (indexVectorDim != indicesTy.getRank() - 1) {
      return rewriter.notifyMatchFailure(op, "no explicit indices dimension");
    }

    // Check that there are multiple batch dimensions.
    if (indicesTy.getRank() < 3) {
      return rewriter.notifyMatchFailure(op, "no multiple batch dimensions");
    }

    const int64_t batchCount = indicesTy.getRank() - 1;
    for (auto it : llvm::enumerate(updatedWindowDims)) {
      if (it.index() != it.value() - batchCount) {
        return rewriter.notifyMatchFailure(
            op, "update windows should be at the end.");
      }
    }

    indices = collapseBatchDims(op.getLoc(), indices, batchCount, rewriter);
    if (!indices) {
      return rewriter.notifyMatchFailure(op,
                                         "cannot collapse indices batch dims");
    }

    llvm::SmallVector<Value> updates;
    for (Value update : op.getUpdates()) {
      update = collapseBatchDims(op.getLoc(), update, batchCount, rewriter);
      if (!update) {
        return rewriter.notifyMatchFailure(op,
                                           "cannot collapse update batch dims");
      }
      updates.push_back(update);
    }

    llvm::SmallVector<int64_t> newUpdatedWindowDims;
    for (auto dim : updatedWindowDims) {
      newUpdatedWindowDims.push_back(dim - batchCount + 1);
    }

    auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
        op.getContext(), newUpdatedWindowDims,
        dimNumbers.getInsertedWindowDims(), dimNumbers.getInputBatchingDims(),
        dimNumbers.getScatterIndicesBatchingDims(),
        dimNumbers.getScatterDimsToOperandDims(),
        /*indexVectorDim=*/1);

    auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
        op.getLoc(), op.getResultTypes(), op.getInputs(), indices, updates,
        newDimNumbers, op.getIndicesAreSorted(), op.getUniqueIndices());
    Region &region = newScatter.getUpdateComputation();
    rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
    rewriter.replaceOp(op, newScatter.getResults());
    return success();
  }
};

// Ensure the batch dimensions of both the indices and updates are the first
// dimensions. If they are not, transpose them to the start.
struct ScatterBatchFirst final : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
    auto dimNumbers = op.getScatterDimensionNumbers();

    // If the index vector dim is not implicitly or explicitly at the end
    // we need to transpose the batch dimensions to the start.
    Value indices = op.getScatterIndices();
    auto indicesTy = llvm::cast<ShapedType>(indices.getType());
    auto indexVectorDim = dimNumbers.getIndexVectorDim();
    if (indexVectorDim < indicesTy.getRank() - 1) {
      llvm::SmallVector<int64_t> perm;
      perm.reserve(indicesTy.getRank());
      for (int i = 0, s = indicesTy.getRank(); i < s; ++i) {
        if (i != indexVectorDim)
          perm.push_back(i);
      }

      if (perm.size() < indicesTy.getRank())
        perm.push_back(indexVectorDim);

      llvm::SmallVector<int64_t> newShape;
      for (int i = 0, s = perm.size(); i < s; ++i) {
        newShape.push_back(indicesTy.getDimSize(perm[i]));
      }

      indices = builder.create<mlir::stablehlo::TransposeOp>(
          indicesTy.clone(newShape), indices,
          builder.getDenseI64ArrayAttr(perm));
      indicesTy = llvm::cast<RankedTensorType>(indices.getType());
      indexVectorDim = indicesTy.getRank() - 1;
    }

    // Compute the permutation require to transpose the batch dimensions to
    // the beginning.
    auto updates = op.getUpdates();
    auto updates0 = updates.front();
    auto updates0Ty = llvm::cast<ShapedType>(updates0.getType());
    auto updatedWindowDims = dimNumbers.getUpdateWindowDims();

    // Determine which dimensions are batch dimensions.
    llvm::SmallVector<bool> isBatch(updates0Ty.getRank(), true);
    for (int i = 0, s = updatedWindowDims.size(); i < s; ++i)
      isBatch[updatedWindowDims[i]] = false;

    // Permute batch dimensions to the start of the update tensor.
    llvm::SmallVector<int64_t> updatePerm;
    updatePerm.reserve(updates0Ty.getRank());
    for (int i = 0, s = isBatch.size(); i < s; ++i)
      if (isBatch[i])
        updatePerm.push_back(i);
    updatePerm.append(updatedWindowDims.begin(), updatedWindowDims.end());

    llvm::SmallVector<int64_t> newUpdatedWindowDims;
    int64_t batchCount = updates0Ty.getRank() - updatedWindowDims.size();
    for (int i = batchCount, s = updates0Ty.getRank(); i < s; i++)
      newUpdatedWindowDims.push_back(i);

    bool indicesChanged = indices != op.getScatterIndices();
    bool updatesChanged =
        llvm::any_of(llvm::enumerate(updatePerm),
                     [](auto it) { return it.index() != it.value(); });
    llvm::SmallVector<Value> newUpdates(updates.begin(), updates.end());
    if (updatesChanged) {
      for (Value &update : newUpdates) {
        auto updateTy = llvm::cast<ShapedType>(update.getType());
        llvm::SmallVector<int64_t> newShape;
        newShape.reserve(updateTy.getRank());
        for (int i = 0, s = updatePerm.size(); i < s; i++)
          newShape.push_back(updateTy.getDimSize(updatePerm[i]));
        update = builder.create<mlir::stablehlo::TransposeOp>(
            updateTy.clone(newShape), update,
            builder.getDenseI64ArrayAttr(updatePerm));
      }
    }

    if (!indicesChanged && !updatesChanged)
      return rewriter.notifyMatchFailure(
          op, "batch dimensions are already leading");

    auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
        op.getContext(), newUpdatedWindowDims,
        dimNumbers.getInsertedWindowDims(), dimNumbers.getInputBatchingDims(),
        dimNumbers.getScatterIndicesBatchingDims(),
        dimNumbers.getScatterDimsToOperandDims(),
        /*indexVectorDim=*/indexVectorDim);

    auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
        op.getLoc(), op.getResultTypes(), op.getInputs(), indices, newUpdates,
        newDimNumbers, op.getIndicesAreSorted(), op.getUniqueIndices());
    Region &region = newScatter.getUpdateComputation();
    rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
    rewriter.replaceOp(op, newScatter.getResults());
    return success();
  }
};

// stablehlo.scatter can materialize a unit dimension at both indexed dimensions
// or at unary dimensions in the destination matrix. linalg_ext.scatter only
// allows unit dimensions at indexed dimensions. This pattern inserts all
// unary dimensions that are not index dimensions to be compatible with
// linalg_ext.scatter.
//
// If converts an stablehlo.scatter as below:
//  %result = "stablehlo.scatter"(...) ({
//    indices_are_sorted = true,
//    scatter_dimension_numbers = #stablehlo.scatter<
//            update_window_dims = [1],
//            inserted_window_dims = [0, 2],
//            scatter_dims_to_operand_dims = [0],
//            index_vector_dim = 1>,
//    unique_indices = true} :
//        (tensor<5x4x1xi32>, tensor<1x1xi32>, tensor<1x4xi32>)
//
// To:
//  %result = "stablehlo.scatter"(...) ({
//    indices_are_sorted = true,
//    scatter_dimension_numbers = #stablehlo.scatter<
//            update_window_dims = [1, 2],
//            inserted_window_dims = [0],
//            scatter_dims_to_operand_dims = [0],
//            index_vector_dim = 1>,
//     unique_indices = true} :
//        (tensor<5x4x1xi32>, tensor<1x1xi32>, tensor<1x4x1xi32>)
//  return %0 : tensor<5x4x1xi32>
struct ScatterMaterializeInsertedDim final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    auto indices = op.getScatterIndices();
    auto operand = op.getInputs().front();
    auto indicesTy = llvm::cast<ShapedType>(indices.getType());
    auto operandTy = llvm::cast<ShapedType>(operand.getType());

    if (!operandTy.hasRank() || !indicesTy.hasRank()) {
      return rewriter.notifyMatchFailure(op, "operand/indices have no rank");
    }

    auto dimNumbers = op.getScatterDimensionNumbers();
    auto updateDims = dimNumbers.getUpdateWindowDims();

    if (indicesTy.getRank() != 2 || dimNumbers.getIndexVectorDim() != 1) {
      return rewriter.notifyMatchFailure(
          op, "indices is not of shape [batch, indices]");
    }

    if (!updateDims.empty() && updateDims.front() == 0) {
      return rewriter.notifyMatchFailure(
          op, "updates is not of shape [batch, ...]");
    }

    auto scatterDimsToOperandDims = dimNumbers.getScatterDimsToOperandDims();
    llvm::SmallVector<bool> isIndexDim(operandTy.getRank(), false);
    for (auto val : scatterDimsToOperandDims) {
      isIndexDim[val] = true;
    }

    int64_t firstNonIndex = 0;
    for (int64_t s = scatterDimsToOperandDims.size(); firstNonIndex < s;
         ++firstNonIndex) {
      if (!isIndexDim[firstNonIndex])
        break;
    }

    llvm::SmallVector<bool> isInsertDims(operandTy.getRank(), false);
    for (auto val : dimNumbers.getInsertedWindowDims()) {
      isInsertDims[val] = true;
    }

    int64_t frontInsertedDims = 0;
    for (; frontInsertedDims < firstNonIndex; ++frontInsertedDims) {
      if (!isInsertDims[frontInsertedDims]) {
        break;
      }
    }

    llvm::ArrayRef<bool> toInsertDims =
        llvm::ArrayRef<bool>(isInsertDims).drop_front(frontInsertedDims);
    if (!llvm::any_of(toInsertDims, [](auto d) { return d; })) {
      return rewriter.notifyMatchFailure(op, "no dimensions to insert");
    }

    // Create a reassociation map that starts with the batch dims.
    SmallVector<ReassociationExprs> reassociationMap;
    reassociationMap.push_back({rewriter.getAffineDimExpr(0)});

    for (auto it : llvm::enumerate(llvm::ArrayRef<bool>(toInsertDims))) {
      if (!it.value())
        reassociationMap.push_back({});
      reassociationMap.back().push_back(
          rewriter.getAffineDimExpr(it.index() + 1));
    }

    llvm::SmallVector<Value> expandedUpdates;
    for (auto update : op.getUpdates()) {
      auto updatesTy = llvm::cast<ShapedType>(update.getType());

      llvm::SmallVector<int64_t> newShape;
      for (int i = 0, s = reassociationMap.size(); i < s; ++i) {
        newShape.push_back(updatesTy.getDimSize(i));
        for (int j = 1, s = reassociationMap[i].size(); j < s; ++j) {
          newShape.push_back(1);
        }
      }

      Value expandUpdate = rewriter.create<tensor::ExpandShapeOp>(
          op.getLoc(),
          RankedTensorType::get(newShape, updatesTy.getElementType()), update,
          reassociationMap);
      expandedUpdates.push_back(expandUpdate);
    }

    llvm::SmallVector<int64_t> newUpdatedWindowDims(toInsertDims.size());
    llvm::SmallVector<int64_t> newInsertedWindowDims(frontInsertedDims);
    std::iota(newUpdatedWindowDims.begin(), newUpdatedWindowDims.end(), 1);
    std::iota(newInsertedWindowDims.begin(), newInsertedWindowDims.end(), 0);

    auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
        op.getContext(), newUpdatedWindowDims, newInsertedWindowDims,
        dimNumbers.getInputBatchingDims(),
        dimNumbers.getScatterIndicesBatchingDims(),
        dimNumbers.getScatterDimsToOperandDims(),
        /*indexVectorDim=*/1);

    auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
        op.getLoc(), op.getResultTypes(), op.getInputs(),
        op.getScatterIndices(), expandedUpdates, newDimNumbers,
        op.getIndicesAreSorted(), op.getUniqueIndices());
    Region &region = newScatter.getUpdateComputation();
    rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
    rewriter.replaceOp(op, newScatter.getResults());
    return success();
  }
};

// Traverse upward past common operations to see if the value came from a
// boolean tensor.
bool isFromBool(Value val) {
  while (true) {
    Operation *op = val.getDefiningOp();
    if (!op)
      return false;

    if (auto convertOp = dyn_cast<mlir::stablehlo::ConvertOp>(op)) {
      auto inTy = llvm::cast<ShapedType>(convertOp.getOperand().getType());
      if (inTy.getElementType().isInteger(1)) {
        return true;
      }
      val = convertOp.getOperand();
      continue;
    }

    if (isa<mlir::stablehlo::DynamicBroadcastInDimOp>(op) ||
        isa<mlir::stablehlo::BroadcastInDimOp>(op) ||
        isa<mlir::stablehlo::BroadcastOp>(op)) {
      val = op->getOperand(0);
      continue;
    }

    return false;
  }
}

// Mul of non-finite values (e.g. NaN, inf) and 0.0 produce 0.0 in StableHLO.
// For linalg we need to convert these to select operations.
struct MulCastOfBool final : OpRewritePattern<mlir::stablehlo::MulOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::MulOp op,
                                PatternRewriter &rewriter) const override {
    auto resultTy = cast<ShapedType>(op.getType());
    if (!isa<FloatType>(resultTy.getElementType()))
      return failure();
    Value lhs = op.getLhs();
    Value rhs = op.getRhs();
    bool lhsIsBool = isFromBool(lhs);
    bool rhsIsBool = isFromBool(rhs);

    if (lhsIsBool == rhsIsBool)
      return failure();
    if (rhsIsBool)
      std::swap(lhs, rhs);

    Type eType = resultTy.getElementType();
    auto lhsTy = cast<ShapedType>(lhs.getType());
    Value lhsBool = rewriter.create<mlir::stablehlo::ConvertOp>(
        op.getLoc(), lhsTy.clone(rewriter.getIntegerType(1)), lhs);
    Value zero = rewriter.create<mlir::stablehlo::ConstantOp>(
        op.getLoc(), DenseElementsAttr::get(RankedTensorType::get({}, eType),
                                            rewriter.getZeroAttr(eType)));

    auto lhsShape = rewriter.create<shape::ShapeOfOp>(
        op.getLoc(),
        RankedTensorType::get({lhsTy.getRank()}, rewriter.getIndexType()), lhs);

    int64_t resultRank = resultTy.getRank();
    auto broadcast = [&](Value value) -> Value {
      auto valueTy = cast<ShapedType>(value.getType());
      auto newTy =
          RankedTensorType::get(resultTy.getShape(), valueTy.getElementType());
      if (valueTy == newTy)
        return value;
      auto dimensions = llvm::to_vector(
          llvm::seq<int64_t>(resultRank - valueTy.getRank(), resultRank));
      return rewriter.create<mlir::stablehlo::DynamicBroadcastInDimOp>(
          op.getLoc(), newTy, value, lhsShape,
          rewriter.getDenseI64ArrayAttr(dimensions));
    };

    zero = broadcast(zero);

    rewriter.replaceOpWithNewOp<mlir::stablehlo::SelectOp>(op, resultTy,
                                                           lhsBool, rhs, zero);
    return success();
  }
};

// Generates Gaussian noise with uniform random generator based on Box-Muller
// transform.
struct ExpandRngNormal final : OpRewritePattern<mlir::stablehlo::RngOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::RngOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getRngDistribution() != mlir::stablehlo::RngDistribution::NORMAL)
      return failure();

    auto resTy = dyn_cast<RankedTensorType>(op.getType());
    // We can support static shapes, but it's easier to implement Box-Muller
    // transform if we know the number of elements.
    if (!resTy || !resTy.hasStaticShape())
      return failure();

    // The algorithm requires even numbers and will generate pairs.
    auto numElems = resTy.getNumElements();
    if (numElems & 1)
      numElems++;
    auto halfNumElems = numElems / 2;

    ImplicitLocOpBuilder b(op.getLoc(), rewriter);

    // Explicitly set the seed to 0, so we have stateless generator. This is not
    // a hard limit. Random generator is still a new topic, and we start with
    // stateless random generator.
    std::mt19937 rng{0};
    std::uniform_real_distribution<> runif(0.0, 1.0);
    SmallVector<float> sqrtValues(halfNumElems), cosValues(halfNumElems),
        sinValues(halfNumElems);
    for (auto i : llvm::seq<unsigned>(0, numElems / 2)) {
      static constexpr float kEpsilon = std::numeric_limits<float>::epsilon();
      static constexpr float kTwoPi = static_cast<float>(2.0 * M_PI);
      float u1, u2;
      do {
        u1 = runif(rng);
        u2 = runif(rng);
      } while (u1 <= kEpsilon);
      sqrtValues[i] = -2.0 * log(u1);
      cosValues[i] = cos(kTwoPi * u2);
      sinValues[i] = sin(kTwoPi * u2);
    }

    // mag = sigma * sqrt(-2.0 * log(u1));
    Value mag = getF32Const(b, /*shapes=*/{halfNumElems}, sqrtValues);
    Value sigma = b.create<mlir::stablehlo::BroadcastOp>(
        mag.getType(), op.getB(), make1DElementsAttr(b, halfNumElems));
    mag = b.create<mlir::stablehlo::MulOp>(
        sigma, b.create<mlir::stablehlo::SqrtOp>(mag));

    // z0 = mag * cos(two_pi * u2) + mu;
    // z1 = mag * sin(two_pi * u2) + mu;
    Value mu = b.create<mlir::stablehlo::BroadcastOp>(
        mag.getType(), op.getA(), make1DElementsAttr(b, halfNumElems));
    Value z0 = getF32Const(b, /*shapes=*/{halfNumElems}, cosValues);
    z0 = b.create<mlir::stablehlo::MulOp>(mag, z0);
    z0 = b.create<mlir::stablehlo::AddOp>(z0, mu);
    Value z1 = getF32Const(b, /*shapes=*/{halfNumElems}, sinValues);
    z1 = b.create<mlir::stablehlo::MulOp>(mag, z1);
    z1 = b.create<mlir::stablehlo::AddOp>(z1, mu);

    Value res = b.create<mlir::stablehlo::ConcatenateOp>(
        ValueRange{z0, z1}, b.getI64IntegerAttr(0));
    if (numElems != resTy.getNumElements()) {
      OpFoldResult zero = b.getIndexAttr(0);
      OpFoldResult one = b.getIndexAttr(1);
      OpFoldResult size = b.getIndexAttr(resTy.getNumElements());
      res = b.create<tensor::ExtractSliceOp>(res, zero, size, one);
    }
    if (resTy.getRank() != 1) {
      res = b.create<mlir::stablehlo::ReshapeOp>(resTy, res);
    }
    rewriter.replaceOp(op, res);
    return success();
  }
};

// clang-format off
//
// Reorder BroadcastInDimOp and N-ary elementwise op.
//
// Rewrites the following pattern (take binary elementwise op as example)
//
// %bcastx = "stablehlo.broadcast_in_dim"(%x) {broadcast_dimensions = %[[BCAST_DIMS]]} : (%[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
// %bcasty = "stablehlo.broadcast_in_dim"(%y) {broadcast_dimensions = %[[BCAST_DIMS]]} : (%[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
// %result = "BinaryElementwiseOpT"(%bcastx, %bcasty) : (%[[SHAPE_AFTER_BCAST]], %[[SHAPE_AFTER_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
//
// into
//
// %z = "BinaryElementwiseOpT"(%x, %y) : (%[[SHAPE_BEFORE_BCAST]], %[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_BEFORE_BCAST]]
// %result = "stablehlo.broadcast_in_dim"(%z) {broadcast_dimensions = %[[BCAST_DIMS]]} : (%[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
//
// clang-format on
template <typename ElementwiseOpT>
struct ReorderBroadcastInDimOpAndElementwiseOp final
    : OpRewritePattern<ElementwiseOpT> {
  using OpRewritePattern<ElementwiseOpT>::OpRewritePattern;

  LogicalResult matchAndRewrite(ElementwiseOpT op,
                                PatternRewriter &rewriter) const override {
    Operation *operation = op.getOperation();
    assert(operation->getNumOperands() >= 1 && operation->getNumResults() == 1);

    // Verify if all operands are from BroadcastInDimOp and its
    // broadcast_dimensions is the same.
    llvm::SmallVector<mlir::stablehlo::BroadcastInDimOp, 2> bcastOps;
    for (auto operand : operation->getOperands()) {
      if (auto bcastOp =
              operand.getDefiningOp<mlir::stablehlo::BroadcastInDimOp>()) {
        bcastOps.push_back(bcastOp);
      } else {
        return failure();
      }
    }

    if (llvm::any_of(bcastOps,
                     [&bcastOps](mlir::stablehlo::BroadcastInDimOp bcastOp) {
                       return bcastOp.getBroadcastDimensions() !=
                              bcastOps[0].getBroadcastDimensions();
                     })) {
      return failure();
    }

    // Verify if all operands of BroadcastInDimOp are of same type and have
    // static shape.
    auto bcastOperandType =
        dyn_cast<ShapedType>(bcastOps[0].getOperand().getType());
    llvm::SmallVector<Value, 2> bcastOperands;
    for (auto bcastOp : bcastOps) {
      auto bcastOperand = bcastOp.getOperand();
      auto type = dyn_cast<ShapedType>(bcastOperand.getType());
      if (!type || !type.hasStaticShape() || type != bcastOperandType) {
        return failure();
      }
      bcastOperands.push_back(bcastOperand);
    }

    // Some elementwise ops, mlir::stablehlo::RealOp for example, do not have
    // SameOperandsAndResultType trait, so resultType might be different
    // from bcastOperandType.
    auto elementType = getElementTypeOrSelf(op.getResult());
    auto resultShape = bcastOperandType.getShape();
    auto resultType = RankedTensorType::get(resultShape, elementType);

    Value result =
        rewriter.create<ElementwiseOpT>(op.getLoc(), resultType, bcastOperands);
    rewriter.replaceOpWithNewOp<mlir::stablehlo::BroadcastInDimOp>(
        op, op.getType(), result, bcastOps[0].getBroadcastDimensionsAttr());

    // NOTE: bcastOps may contain duplicates.
    SetVector<Operation *> deadOps;
    for (auto bcastOp : bcastOps) {
      if (bcastOp.getOperation()->use_empty())
        deadOps.insert(bcastOp);
    }
    for (auto *deadOp : deadOps)
      rewriter.eraseOp(deadOp);

    return success();
  }
};

// Identifies cases where a dense operation has inputs that come from widening
// operations. For instance, a dot product widening from FP16 to FP32 is better
// to have the casting operation fused into the dot operation. This decreases
// the loading required during a dense computation.
template <class Op>
struct FuseWidenOperands final : OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;

  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    llvm::SmallVector<Value> operands;
    for (Value operand : op->getOperands()) {
      auto convertOp =
          dyn_cast_or_null<mlir::stablehlo::ConvertOp>(operand.getDefiningOp());
      if (convertOp) {
        auto inputType = getElementTypeOrSelf(convertOp.getOperand().getType());
        auto castedType = getElementTypeOrSelf(convertOp.getResult().getType());
        if (!isa<FloatType, IntegerType>(inputType) ||
            !isa<FloatType, IntegerType>(castedType)) {
          return rewriter.notifyMatchFailure(
              op, "non-integer or floating point type");
          ;
        }

        if (inputType.getIntOrFloatBitWidth() <
            castedType.getIntOrFloatBitWidth()) {
          operands.push_back(convertOp.getOperand());
          continue;
        }
      }
      operands.push_back(operand);
    }

    if (llvm::all_of(
            llvm::zip_equal(operands, op->getOperands()),
            [](auto pair) { return std::get<0>(pair) == std::get<1>(pair); }))
      return failure();

    rewriter.replaceOpWithNewOp<Op>(op, op->getResultTypes(), operands,
                                    op->getAttrs());
    return success();
  }
};

struct DotToMul final : OpRewritePattern<mlir::stablehlo::DotOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::DotOp op,
                                PatternRewriter &rewriter) const override {
    Value lhs = op.getLhs();
    Value rhs = op.getRhs();
    auto lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
    auto rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
    auto resultTy = dyn_cast<RankedTensorType>(op.getType());

    if (!lhsTy || !rhsTy) {
      return rewriter.notifyMatchFailure(op, "lhs and rhs must be ranked");
    }

    if (lhsTy.getRank() != 2 || rhsTy.getRank() != 2) {
      return rewriter.notifyMatchFailure(op, "lhs and rhs must be rank-2");
    }

    if (lhsTy.getDimSize(1) != 1)
      return failure();

    // Dynamically compute the shape of the result of the DotOp by querying
    // the 0-th dimensions, of the left, and the 1st dimension of the right.
    // Concatenating them togething to make the final shape.
    Value batchSize = rewriter.create<mlir::stablehlo::GetDimensionSizeOp>(
        op.getLoc(), lhs, rewriter.getI64IntegerAttr(0));
    Value batchSize1 = rewriter.create<mlir::stablehlo::ReshapeOp>(
        op.getLoc(), RankedTensorType::get({1}, rewriter.getI32Type()),
        batchSize);

    Value featureSize = rewriter.create<mlir::stablehlo::GetDimensionSizeOp>(
        op.getLoc(), rhs, rewriter.getI64IntegerAttr(1));
    Value featureSize1 = rewriter.create<mlir::stablehlo::ReshapeOp>(
        op.getLoc(), RankedTensorType::get({1}, rewriter.getI32Type()),
        featureSize);

    Value outSize = rewriter.create<mlir::stablehlo::ConcatenateOp>(
        op.getLoc(), RankedTensorType::get({2}, rewriter.getI32Type()),
        ValueRange{batchSize1, featureSize1}, rewriter.getI64IntegerAttr(0));

    lhs = rewriter.create<mlir::stablehlo::DynamicBroadcastInDimOp>(
        op.getLoc(), resultTy.clone(lhsTy.getElementType()), lhs, outSize,
        rewriter.getDenseI64ArrayAttr({0, 1}));

    rhs = rewriter.create<mlir::stablehlo::DynamicBroadcastInDimOp>(
        op.getLoc(), resultTy.clone(rhsTy.getElementType()), rhs, outSize,
        rewriter.getDenseI64ArrayAttr({0, 1}));

    auto computeETy = lhsTy.getElementType();
    if (computeETy.getIntOrFloatBitWidth() < rhsTy.getElementTypeBitWidth())
      computeETy = rhsTy.getElementType();
    if (computeETy.getIntOrFloatBitWidth() < resultTy.getElementTypeBitWidth())
      computeETy = resultTy.getElementType();

    auto computeTy = resultTy.clone(computeETy);

    rhs = rewriter.create<mlir::stablehlo::ConvertOp>(op.getLoc(), computeTy,
                                                      rhs);
    lhs = rewriter.create<mlir::stablehlo::ConvertOp>(op.getLoc(), computeTy,
                                                      lhs);

    auto result = rewriter.create<mlir::stablehlo::MulOp>(
        op.getLoc(), resultTy.clone(computeETy), lhs, rhs);
    rewriter.replaceOpWithNewOp<mlir::stablehlo::ConvertOp>(op, resultTy,
                                                            result);
    return success();
  }
};

// Rewrite RngBitGenerator with f32 return type to instead generate the same
// number of i32 outputs, then BitcastConvert to return f32.
struct RngBitcastFloat final
    : OpRewritePattern<mlir::stablehlo::RngBitGeneratorOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::RngBitGeneratorOp op,
                                PatternRewriter &rewriter) const override {
    ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
    auto resultTy = dyn_cast<RankedTensorType>(op.getType(1));
    auto stateTy = dyn_cast<RankedTensorType>(op.getType(0));

    if (!isa<FloatType>(resultTy.getElementType())) {
      return failure();
    }

    llvm::SmallVector<Type> castedTypes;
    castedTypes.push_back(stateTy);
    castedTypes.push_back(resultTy.clone(rewriter.getI32Type()));

    TypeRange castedTypeRange = TypeRange{castedTypes};

    auto resultOp = rewriter.create<mlir::stablehlo::RngBitGeneratorOp>(
        op.getLoc(), castedTypeRange, op.getRngAlgorithm(), op.getOperand());

    auto casted = rewriter.create<mlir::stablehlo::BitcastConvertOp>(
        resultOp.getLoc(), resultTy, resultOp.getResult(1));

    llvm::SmallVector<Value> results;
    results.push_back(resultOp.getResult(0));
    results.push_back(casted);

    rewriter.replaceOp(op, results);
    return success();
  }
};

struct ZeroConcat final : OpRewritePattern<mlir::stablehlo::ConcatenateOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::ConcatenateOp op,
                                PatternRewriter &rewriter) const override {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type || !type.hasStaticShape())
      return failure();

    uint64_t axis = op.getDimension();
    OperandRange origInputs = op.getInputs();
    SmallVector<Value> nonzeroInputs;
    for (auto input : origInputs) {
      auto type = dyn_cast<RankedTensorType>(input.getType());
      ArrayRef<int64_t> shape = type.getShape();
      if (axis > shape.size())
        return failure();

      if (shape[axis] != 0)
        nonzeroInputs.push_back(input);
    }

    if (nonzeroInputs.size() == origInputs.size())
      return failure();

    rewriter.replaceOpWithNewOp<mlir::stablehlo::ConcatenateOp>(
        op, nonzeroInputs, /*dimension=*/axis);
    return success();
  }
};

// Similar to DotIsMul, this finds the case where a dot general
// can be represented using a mul operation. This includes possibly making
// an implicit cast explicit prior the mul.
struct DotGeneralIsMul final : OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::DotGeneralOp op,
                                PatternRewriter &rewriter) const override {
    auto lhs = cast<Value>(op.getLhs());
    auto rhs = cast<Value>(op.getRhs());
    auto lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
    auto rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
    auto resultTy = dyn_cast<RankedTensorType>(op.getType());
    ImplicitLocOpBuilder builder(op.getLoc(), rewriter);

    if (!lhsTy || !rhsTy || !resultTy)
      return failure();

    auto dNums = op.getDotDimensionNumbers();
    auto batchDimsL = dNums.getLhsBatchingDimensions();
    auto batchDimsR = dNums.getRhsBatchingDimensions();
    auto contractDimsL = dNums.getLhsContractingDimensions();
    auto contractDimsR = dNums.getRhsContractingDimensions();

    llvm::SmallVector<bool> isLhsParallelDim(lhsTy.getRank(), true);
    llvm::SmallVector<bool> isRhsParallelDim(rhsTy.getRank(), true);

    for (auto dim : batchDimsL)
      isLhsParallelDim[dim] = false;
    for (auto dim : batchDimsR)
      isRhsParallelDim[dim] = false;
    for (auto dim : contractDimsL)
      isLhsParallelDim[dim] = false;
    for (auto dim : contractDimsR)
      isRhsParallelDim[dim] = false;

    for (auto dim : contractDimsL) {
      if (lhsTy.getDimSize(dim) != 1) {
        return rewriter.notifyMatchFailure(op, "Non unit contract dimensions");
      }
    }

    // Generate the permutation matrix to order BatchDims, ParallelDims,
    // ContractDims.
    llvm::SmallVector<int64_t> permLhs;
    llvm::SmallVector<int64_t> permRhs;
    permLhs.append(batchDimsL.begin(), batchDimsL.end());
    permRhs.append(batchDimsR.begin(), batchDimsR.end());

    for (auto [idx, value] : llvm::enumerate(isLhsParallelDim)) {
      if (value)
        permLhs.push_back(idx);
    }

    for (auto [idx, value] : llvm::enumerate(isRhsParallelDim)) {
      if (value)
        permRhs.push_back(idx);
    }

    llvm::append_range(permLhs, contractDimsL);
    llvm::append_range(permRhs, contractDimsR);

    // Determine the transpose shape based on the generate permutations.
    llvm::SmallVector<int64_t> lhsTransposeShape;
    llvm::SmallVector<int64_t> rhsTransposeShape;
    for (auto dim : permLhs)
      lhsTransposeShape.push_back(lhsTy.getDimSize(dim));
    for (auto dim : permRhs)
      rhsTransposeShape.push_back(rhsTy.getDimSize(dim));

    // Transpose the left hand side and the right hand side.
    lhs = builder.create<mlir::stablehlo::TransposeOp>(
        RankedTensorType::get(lhsTransposeShape, lhsTy.getElementType()), lhs,
        builder.getDenseI64ArrayAttr(permLhs));
    lhsTy = llvm::cast<RankedTensorType>(lhs.getType());

    rhs = builder.create<mlir::stablehlo::TransposeOp>(
        RankedTensorType::get(rhsTransposeShape, rhsTy.getElementType()), rhs,
        builder.getDenseI64ArrayAttr(permRhs));
    rhsTy = llvm::cast<RankedTensorType>(rhs.getType());

    auto dimI32Ty = RankedTensorType::get({1}, builder.getI32Type());

    // Drop all of the non-concat dimensions from the lhs.
    llvm::SmallVector<Value> lhsReshapeDims;
    for (int i = 0, s = lhsTy.getRank() - contractDimsL.size(); i < s; i++) {
      Value dim = builder.create<mlir::stablehlo::GetDimensionSizeOp>(lhs, i);
      lhsReshapeDims.push_back(
          builder.create<mlir::stablehlo::ReshapeOp>(dimI32Ty, dim));
    }
    Value lhsDynShape = builder.create<mlir::stablehlo::ConcatenateOp>(
        RankedTensorType::get({static_cast<int64_t>(lhsReshapeDims.size())},
                              builder.getI32Type()),
        lhsReshapeDims, 0);
    lhsTy =
        RankedTensorType::get(lhsTy.getShape().drop_back(contractDimsL.size()),
                              lhsTy.getElementType());
    lhs = builder.create<mlir::stablehlo::DynamicReshapeOp>(lhsTy, lhs,
                                                            lhsDynShape);

    // Drop all of the non concat dimensions from the rhs.
    llvm::SmallVector<Value> rhsReshapeDims;
    for (int i = 0, s = rhsTy.getRank() - contractDimsR.size(); i < s; i++) {
      Value dim = builder.create<mlir::stablehlo::GetDimensionSizeOp>(rhs, i);
      rhsReshapeDims.push_back(
          builder.create<mlir::stablehlo::ReshapeOp>(dimI32Ty, dim));
    }
    Value rhsDynShape = builder.create<mlir::stablehlo::ConcatenateOp>(
        RankedTensorType::get({static_cast<int64_t>(rhsReshapeDims.size())},
                              builder.getI32Type()),
        rhsReshapeDims, 0);
    rhsTy =
        RankedTensorType::get(rhsTy.getShape().drop_back(contractDimsR.size()),
                              rhsTy.getElementType());
    rhs = builder.create<mlir::stablehlo::DynamicReshapeOp>(rhsTy, rhs,
                                                            rhsDynShape);

    // Compute the size of the output shape with dynamic shape support using the
    // lhs and rhs dimensions.
    llvm::SmallVector<Value> outputDims;
    outputDims.append(lhsReshapeDims);
    outputDims.append(rhsReshapeDims.begin() + batchDimsR.size(),
                      rhsReshapeDims.end());
    Value outputShape = builder.create<mlir::stablehlo::ConcatenateOp>(
        RankedTensorType::get({resultTy.getRank()}, builder.getI32Type()),
        outputDims, 0);

    // Broadcast the left hand side to match the expect output shape.
    llvm::SmallVector<int64_t> lhsDimMapping(lhsTy.getRank());
    std::iota(lhsDimMapping.begin(), lhsDimMapping.end(), 0);
    auto lhsBroadcastTy =
        RankedTensorType::get(resultTy.getShape(), lhsTy.getElementType());
    lhs = builder.createOrFold<mlir::stablehlo::DynamicBroadcastInDimOp>(
        lhsBroadcastTy, lhs, outputShape,
        rewriter.getDenseI64ArrayAttr(lhsDimMapping));

    // Broadcast the right hand side to match the expected output shape.
    llvm::SmallVector<int64_t> rhsDimMapping(rhsTy.getRank());
    std::iota(rhsDimMapping.begin(), rhsDimMapping.begin() + batchDimsR.size(),
              0);
    std::iota(rhsDimMapping.begin() + batchDimsR.size(), rhsDimMapping.end(),
              lhsTy.getRank());
    auto rhsBroadcastTy =
        RankedTensorType::get(resultTy.getShape(), rhsTy.getElementType());
    rhs = builder.createOrFold<mlir::stablehlo::DynamicBroadcastInDimOp>(
        rhsBroadcastTy, rhs, outputShape,
        rewriter.getDenseI64ArrayAttr(rhsDimMapping));

    lhs = builder.createOrFold<mlir::stablehlo::ConvertOp>(resultTy, lhs);
    rhs = builder.createOrFold<mlir::stablehlo::ConvertOp>(resultTy, rhs);
    rewriter.replaceOpWithNewOp<mlir::stablehlo::MulOp>(op, resultTy, lhs, rhs);
    return success();
  }
};

// TODO(suderman): Move this to solve code.
struct CustomCallIsTopK final
    : OpRewritePattern<mlir::stablehlo::CustomCallOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getCallTargetName() != "TopK") {
      return rewriter.notifyMatchFailure(op, "not a TopK custom call");
    }

    if (op.getNumOperands() != 1 ||
        !(op.getNumResults() == 1 || op.getNumResults() == 2)) {
      return rewriter.notifyMatchFailure(
          op, "incorrect number of operands / results");
    }

    ArrayAttr computations = op.getCalledComputations();
    if (computations.size() != 1) {
      return rewriter.notifyMatchFailure(op,
                                         "incorrect number of computations");
    }

    SymbolRefAttr computation = dyn_cast<SymbolRefAttr>(computations[0]);
    if (!computation) {
      return rewriter.notifyMatchFailure(op, "not a ref attr");
    }

    auto operand = op.getOperand(0);
    auto operandTy = cast<ShapedType>(operand.getType());
    if (!operandTy.hasRank() || operandTy.getRank() != 2) {
      return rewriter.notifyMatchFailure(op, "rank-2 input not found");
    }

    ShapedType topVTy;
    ShapedType topITy;
    if (op.getNumResults() == 1) {
      if (auto tupleTy = dyn_cast<TupleType>(op.getType(0))) {
        if (tupleTy.size() != 2) {
          return rewriter.notifyMatchFailure(
              op, "tuple return does not tuple two values");
        }
        topVTy = dyn_cast<ShapedType>(tupleTy.getType(0));
        topITy = dyn_cast<ShapedType>(tupleTy.getType(1));
      }
    }

    if (op.getNumResults() == 2) {
      topVTy = dyn_cast<ShapedType>(op.getType(0));
      topITy = dyn_cast<ShapedType>(op.getType(1));
    }

    if (!topVTy || !topITy) {
      return rewriter.notifyMatchFailure(op, "unknown return type behavior");
    }

    int64_t k = topVTy.getDimSize(1);
    if (ShapedType::isDynamic(k)) {
      return rewriter.notifyMatchFailure(op, "dynamic top-k k value");
    }

    auto moduleOp = op->getParentOfType<ModuleOp>();
    auto funcOp = dyn_cast<func::FuncOp>(moduleOp.lookupSymbol(computation));

    Block &block = funcOp.getRegion().front();
    auto stablehloCompareOp =
        dyn_cast<mlir::stablehlo::CompareOp>(block.front());
    if (!stablehloCompareOp) {
      return rewriter.notifyMatchFailure(op, "not stablehlo compare op");
    }

    auto returnOp = dyn_cast<func::ReturnOp>(block.getTerminator());
    if (!returnOp) {
      return rewriter.notifyMatchFailure(op, "could not find ReturnOp");
    }

    if (returnOp.getNumOperands() != 1 ||
        returnOp.getOperand(0) != stablehloCompareOp.getResult()) {
      return rewriter.notifyMatchFailure(op, "ReturnOp operand not compare op");
    }

    auto direction = stablehloCompareOp.getComparisonDirection();
    bool getTop = direction == mlir::stablehlo::ComparisonDirection::GT ||
                  direction == mlir::stablehlo::ComparisonDirection::GE;

    if (!getTop) {
      return rewriter.notifyMatchFailure(op,
                                         "Unsupported comparison direction");
    }

    auto newTopK = rewriter.create<chlo::TopKOp>(
        op.getLoc(), TypeRange{topVTy, topITy}, operand, k);

    if (op.getNumResults() == 2) {
      rewriter.replaceOp(op, newTopK.getResults());
      return success();
    }

    if (auto tupleTy = dyn_cast<TupleType>(op.getType(0))) {
      rewriter.replaceOpWithNewOp<mlir::stablehlo::TupleOp>(
          op, op.getType(0), newTopK.getResults());
      return success();
    }

    return failure();
  }
};

// Recursive helper function that identifies an Iota followed by a set of
// broadcasts where the last dimension of the iota is preserved throughout.
bool isIotaOrIotaBroadcast(PatternRewriter &rewriter, Value input) {
  if (auto iotaOp =
          dyn_cast_or_null<mlir::stablehlo::IotaOp>(input.getDefiningOp())) {
    int64_t iotaDim = iotaOp.getIotaDimension();
    auto iotaLastDim = cast<ShapedType>(iotaOp.getType()).getRank() - 1;
    if (iotaDim == iotaLastDim) {
      return true;
    }

    (void)rewriter.notifyMatchFailure(iotaOp, "iota must be on last dimension");
    return false;
  }

  if (auto broadcastOp = dyn_cast_or_null<mlir::stablehlo::BroadcastInDimOp>(
          input.getDefiningOp())) {
    auto broadcastLastDim =
        cast<ShapedType>(broadcastOp.getType()).getRank() - 1;
    if (broadcastOp.getBroadcastDimensions().back() != broadcastLastDim) {
      (void)rewriter.notifyMatchFailure(
          broadcastOp, "last dimension must be maintained in broadcast");
      return false;
    }
    return isIotaOrIotaBroadcast(rewriter, broadcastOp.getOperand());
  }

  return false;
}

struct IotaSortSliceIsTopK final : OpRewritePattern<mlir::stablehlo::SortOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::SortOp op,
                                PatternRewriter &rewriter) const override {
    auto opOperands = op.getOperands();
    auto opResults = op.getResults();
    Value topKInput;
    if (opOperands.size() != 2 || opResults.size() != 2) {
      return rewriter.notifyMatchFailure(
          op, "slice that maps to TopK must have exactly two inputs/outputs");
    }

    Value inputIota;
    // Check that one of the inputs is iota, assume that the other one is the
    // input.
    for (Value operand : opOperands) {
      if (isIotaOrIotaBroadcast(rewriter, operand)) {
        inputIota = operand;
      } else {
        topKInput = operand;
      }
    }

    if (!inputIota) {
      return rewriter.notifyMatchFailure(op, "sort isn't called from Iota");
    }

    Block &block = op.getRegion().front();
    auto stablehloCompareOp =
        dyn_cast<mlir::stablehlo::CompareOp>(block.front());
    if (!stablehloCompareOp) {
      return rewriter.notifyMatchFailure(op, "not stablehlo compare op");
    }

    auto direction = stablehloCompareOp.getComparisonDirection();
    bool getTop = direction == mlir::stablehlo::ComparisonDirection::GT ||
                  direction == mlir::stablehlo::ComparisonDirection::GE;

    if (!getTop) {
      return rewriter.notifyMatchFailure(op,
                                         "unsupported comparison direction");
    }

    Value topV, topI;
    int64_t k;
    // Check that the output of the sort op gets fed into a slice.
    for (auto [idx, result] : llvm::enumerate(opResults)) {
      if (result.getUsers().empty())
        return rewriter.notifyMatchFailure(
            op, "sort isn't calling into a slice op");
      auto sliceOp =
          dyn_cast<mlir::stablehlo::SliceOp>(*result.getUsers().begin());
      if (!sliceOp) {
        return rewriter.notifyMatchFailure(
            op, "sort isn't calling into a slice op");
      }

      for (auto stride : sliceOp.getStrides()) {
        if (stride != 1) {
          return rewriter.notifyMatchFailure(
              op, "all slice strides must be 1 in order to match to TopK");
        }
      }

      // Treat the first slice as inputs, the second as indices.
      if (idx == 0) {
        topV = sliceOp.getResult();
        k = sliceOp.getLimitIndices().back();
      } else {
        topI = sliceOp.getResult();
      }
    }

    auto topK = rewriter.create<chlo::TopKOp>(
        op.getLoc(), TypeRange{topV.getType(), topI.getType()}, topKInput, k);
    topV.replaceAllUsesWith(topK.getResults()[0]);
    topI.replaceAllUsesWith(topK.getResults()[1]);
    return success();
  }
};

// TODO(suderman): Move this to solve code.
struct ApproxTopK final : OpRewritePattern<mlir::stablehlo::CustomCallOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getCallTargetName() != "ApproxTopK")
      return rewriter.notifyMatchFailure(op, "not ApproxTopK operation.");

    auto computationName =
        dyn_cast<SymbolRefAttr>(op.getCalledComputationsAttr()[0]);
    Operation *funcOp;
    for (auto parent = op->getParentOp(); parent;
         parent = parent->getParentOp()) {
      funcOp = SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
          parent, computationName);
      if (funcOp)
        break;
    }
    if (!funcOp)
      return rewriter.notifyMatchFailure(op, "computation function not found.");

    int64_t k = cast<ShapedType>(op.getType(0)).getShape().back();
    auto input = op.getOperand(0);
    auto iota = op.getOperand(1);

    if (auto iotaOp =
            dyn_cast_or_null<mlir::stablehlo::IotaOp>(iota.getDefiningOp())) {
      int64_t iotaDim = iotaOp.getIotaDimension();
      auto iotaLastDim = cast<ShapedType>(iotaOp.getType()).getRank() - 1;
      if (iotaDim != iotaLastDim) {
        return rewriter.notifyMatchFailure(op, "Iota of last dim not found.");
      }
    }

    Block &block = funcOp->getRegion(0).front();
    auto stablehloCompareOp =
        dyn_cast<mlir::stablehlo::CompareOp>(block.front());
    if (!stablehloCompareOp) {
      return rewriter.notifyMatchFailure(op, "not stablehlo compare op");
    }

    auto returnOp = block.getTerminator();
    auto freturnOp = dyn_cast<func::ReturnOp>(returnOp);
    auto sreturnOp = dyn_cast<mlir::stablehlo::ReturnOp>(returnOp);
    if (!freturnOp && !sreturnOp) {
      return rewriter.notifyMatchFailure(op, "could not find ReturnOp");
    }

    if (returnOp->getNumOperands() != 1 ||
        returnOp->getOperand(0) != stablehloCompareOp.getResult()) {
      return rewriter.notifyMatchFailure(op, "ReturnOp operand not compare op");
    }

    auto direction = stablehloCompareOp.getComparisonDirection();
    bool getTop = direction == mlir::stablehlo::ComparisonDirection::GT ||
                  direction == mlir::stablehlo::ComparisonDirection::GE;
    if (getTop) {
      auto topK =
          rewriter.create<chlo::TopKOp>(op.getLoc(), op.getResultTypes(), input,
                                        rewriter.getI64IntegerAttr(k));
      rewriter.replaceOp(op, topK);
      return success();
    }

    bool getBottom = direction == mlir::stablehlo::ComparisonDirection::LT ||
                     direction == mlir::stablehlo::ComparisonDirection::LE;
    if (getBottom) {
      input = rewriter.create<mlir::stablehlo::NegOp>(op.getLoc(), input);
      auto topK =
          rewriter.create<chlo::TopKOp>(op.getLoc(), op.getResultTypes(), input,
                                        rewriter.getI64IntegerAttr(k));
      rewriter.replaceOp(op, topK);
      return success();
    }

    return failure();
  }
};

struct StableHLOToStableHLOPreprocessing final
    : impl::StableHLOToStableHLOPreprocessingBase<
          StableHLOToStableHLOPreprocessing> {
  using StableHLOToStableHLOPreprocessingBase::
      StableHLOToStableHLOPreprocessingBase;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<shape::ShapeDialect, mlir::stablehlo::StablehloDialect,
                    chlo::ChloDialect, tensor::TensorDialect>();
  }

  void runOnOperation() override {
    MLIRContext *context = &getContext();
    ConversionTarget conversionTarget(*context);
    RewritePatternSet conversionPatterns(context);
    conversionTarget
        .addLegalDialect<shape::ShapeDialect, chlo::ChloDialect,
                         mlir::stablehlo::StablehloDialect, math::MathDialect,
                         mlir::func::FuncDialect, mlir::arith::ArithDialect,
                         mlir::tensor::TensorDialect>();
    if (failed(applyPartialConversion(getOperation(), conversionTarget,
                                      std::move(conversionPatterns)))) {
      return signalPassFailure();
    }

    RewritePatternSet patterns(context);
    // General StableHLO canonicalization patterns. Run these with a high
    // benefit to enable more rewrites and avoid needless expansions that could
    // be more difficult to fold away. Note that we need to manually add these
    // because StableHLO does not provide constant folders or canonicalization
    // patterns for its ops.
    populateCanonicalizationPatterns(context, &patterns, PatternBenefit{1024});

    // TODO: Remove once we have a general contraction to matmul pass.
    populatePreprocessingEinsumToDotGeneralPatterns(context, &patterns);
    populatePreprocessingUnfuseBatchNormPatterns(context, &patterns);
    populatePreprocessingComplexPatterns(context, &patterns);
    populatePreprocessingGatherToTorchIndexSelectPatterns(context, &patterns);
    patterns.insert<ExpandRngNormal, MulCastOfBool>(context);

    // rng float conversion pattern
    patterns.insert<RngBitcastFloat>(context);

    // scatter canonicalization patterns
    patterns.insert<ScatterInt64Indices, ScatterImplicitIndex,
                    ScatterImplicitBatch, ScatterMaterializeInsertedDim,
                    ScatterCollapseBatch, ScatterBatchFirst>(context);

    // dot_general canonicalization patterns.
    populatePreprocessingDotGeneralToDotPatterns(context, &patterns);

    // TODO(jpienaar): This may be redundant with lower_general_dot. Remove if
    // so.
    patterns.insert<TransposeReshapeGenericDotGeneral>(context,
                                                       /*benefit=*/200);
    patterns.insert<DotGeneralIsMul>(context, /*benefit=*/300);

    // Fusion operations.
    // TODO: Reconsider performing this optimization in preprocessing.
    patterns.insert<FuseWidenOperands<mlir::stablehlo::DotOp>,
                    FuseWidenOperands<mlir::stablehlo::DotGeneralOp>,
                    FuseWidenOperands<mlir::stablehlo::ConvolutionOp>>(
        context,
        /*benefit=*/400);

    // Identify known custom calls and convert them to equivalent StableHLO.
    patterns.insert<CustomCallIsTopK>(context);

    // Identify an iota->sort->slice pattern that maps to TopK.
    patterns.insert<IotaSortSliceIsTopK, ApproxTopK>(context);

    // Additional canonicalizers that simplify to computationally
    // less-complex operations.
    patterns.insert<DotToMul, ZeroConcat>(context);

    // Unary elementwise op.
    patterns.insert<
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::AbsOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::CeilOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ConvertOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ClzOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::CosineOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ExpOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::Expm1Op>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::FloorOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ImagOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::IsFiniteOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::LogOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::Log1pOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::LogisticOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::NotOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::NegOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<
            mlir::stablehlo::PopulationCountOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RealOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RoundOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RsqrtOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SignOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SineOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SqrtOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::TanhOp>>(
        context);
    // Binary elementwise op.
    patterns.insert<
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::AddOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::Atan2Op>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ComplexOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::DivOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::MaxOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::MinOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::MulOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::PowOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RemOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ShiftLeftOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<
            mlir::stablehlo::ShiftRightArithmeticOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<
            mlir::stablehlo::ShiftRightLogicalOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SubtractOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::AndOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::OrOp>,
        ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::XorOp>>(
        context);
    if (orderConvFeatures) {
      patterns.insert<ReorderConvOpInputDimensions>(context);
      patterns.insert<ReorderConvOpKernelDimensions>(context);
      patterns.insert<ReorderConvOpOutputDimensions>(context);
    }
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

} // namespace
} // namespace mlir::iree_compiler::stablehlo
