// 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/Base.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 mlir::stablehlo::ConstantOp::create(
             b, 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 Base::Base;

  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 = mlir::stablehlo::TransposeOp::create(
        rewriter, 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 = mlir::stablehlo::ConvolutionOp::create(
        rewriter, 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 Base::Base;
  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 = mlir::stablehlo::TransposeOp::create(
        rewriter, 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 =
        mlir::stablehlo::ConvolutionOp::create(
            rewriter, 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 Base::Base;

  LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
                                PatternRewriter &rewriter) const override {
    auto resultType = 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 = mlir::stablehlo::ConvolutionOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(convShape, resultType.getElementType()), operands,
        op->getAttrs());
    newConv.setDimensionNumbersAttr(newDimensionNumbers);

    auto transposed = mlir::stablehlo::TransposeOp::create(
        rewriter, 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 splitting 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 Base::Base;

  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 mlir::stablehlo::TransposeOp::create(
        b, 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;
    }

    int64_t resultShape[] = {
        llvm::product_of(shape.take_front(dimsBorder0)),
        llvm::product_of(shape.slice(dimsBorder0, dimsBorder1 - dimsBorder0)),
        llvm::product_of(shape.drop_front(dimsBorder1))};
    return mlir::stablehlo::ReshapeOp::create(
        b, loc, RankedTensorType::get(resultShape, 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 = mlir::stablehlo::DotGeneralOp::create(
        rewriter, 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 = mlir::stablehlo::ReshapeOp::create(
          rewriter, op.getLoc(), op.getType(), newOp.getResult());
    }

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

// The following helpers and the ScatterBatchingDimsExpander pattern below are
// ported from upstream StableHLO's `ScatterWithBatchingDimsExpander` and its
// supporting helpers in
// third_party/stablehlo/stablehlo/transforms/StablehloCompatibilityExpander.cpp
// (mergeSortedDims / fitsInIntegralType / promoteTypeForSize /
// getUpdatedIndicesAreSorted / createConcatIndices). They are copied rather
// than reused because the upstream pattern lives in an anonymous namespace, and
// its only public entry point (populateStablehloCompatibilityExpanderPatterns)
// is a version-gated bundle that would also pull in unrelated compatibility
// expanders.
//
// Merges two sorted lists of dimensions into a single sorted list.
SmallVector<int64_t> mergeSortedDims(ArrayRef<int64_t> dims1,
                                     ArrayRef<int64_t> dims2) {
  SmallVector<int64_t> result;
  result.reserve(dims1.size() + dims2.size());
  std::merge(dims1.begin(), dims1.end(), dims2.begin(), dims2.end(),
             std::back_inserter(result));
  return result;
}

bool fitsInIntegralType(int64_t size, IntegerType type) {
  if (type.isUnsigned()) {
    return llvm::isUIntN(type.getWidth(), size);
  }
  return llvm::isIntN(type.getWidth(), size);
}

// If `type` is an integer type in which `size` doesn't fit, promote it to i32
// or i64 (depending on `size`).
Type promoteTypeForSize(Type type, int64_t size, OpBuilder &builder) {
  // Gather/Scatter should have an integer type, but we check just in case.
  auto intType = dyn_cast<IntegerType>(type);
  if (!intType || fitsInIntegralType(size, intType)) {
    return type;
  }
  if (fitsInIntegralType(size, builder.getI32Type())) {
    return builder.getI32Type();
  }
  return builder.getI64Type();
}

// If `indicesBatchingDims` and `updatedIndexMap` are both sorted, then the
// `indices_are_sorted` property is preserved: each concatenated iota is
// monotonically increasing.
bool getUpdatedIndicesAreSorted(bool indicesAreSorted,
                                ArrayRef<int64_t> indicesBatchingDims,
                                ArrayRef<int64_t> updatedIndexMap) {
  return indicesAreSorted && llvm::is_sorted(indicesBatchingDims) &&
         llvm::is_sorted(updatedIndexMap);
}

// Returns an updated indices tensor such that an `IotaOp` is prepended for each
// dim in `indicesBatchingDims` with a `ConcatenateOp`.
//
// If `indexVectorDim` is equal to the rank of `indices`, it is reshaped to have
// a trailing dimension of size 1 so it can be concatenated with the `IotaOp`s.
Value createConcatIndices(Value indices, int64_t indexVectorDim,
                          ArrayRef<int64_t> indicesBatchingDims,
                          PatternRewriter &rewriter) {
  Location loc = indices.getLoc();
  auto indicesType = cast<RankedTensorType>(indices.getType());
  Type elementType = indicesType.getElementType();

  // The batching dim sizes might not fit in the existing element type, in which
  // case we need to promote it.
  for (int64_t batchingDim : indicesBatchingDims) {
    elementType = promoteTypeForSize(
        elementType, indicesType.getDimSize(batchingDim), rewriter);
  }
  if (elementType != indicesType.getElementType()) {
    indicesType = RankedTensorType::get(indicesType.getShape(), elementType);
    indices =
        mlir::stablehlo::ConvertOp::create(rewriter, loc, indicesType, indices);
  }

  bool indexVectorDimOnLastDim = indexVectorDim == indicesType.getRank();
  SmallVector<int64_t> iotaShape(indicesType.getShape());
  if (indexVectorDimOnLastDim) {
    iotaShape.push_back(1);
  } else {
    iotaShape[indexVectorDim] = 1;
  }
  auto iotaType = RankedTensorType::get(iotaShape, elementType);

  if (indexVectorDimOnLastDim) {
    indices =
        mlir::stablehlo::ReshapeOp::create(rewriter, loc, iotaType, indices);
  }

  SmallVector<Value> indicesToConcat;
  indicesToConcat.reserve(indicesBatchingDims.size() + 1);
  for (int64_t batchingDim : indicesBatchingDims) {
    indicesToConcat.push_back(
        mlir::stablehlo::IotaOp::create(rewriter, loc, iotaType, batchingDim));
  }
  indicesToConcat.push_back(indices);
  return mlir::stablehlo::ConcatenateOp::create(rewriter, loc, indicesToConcat,
                                                indexVectorDim);
}

// Converts a `stablehlo.scatter` with batching dims to one without batching
// dims, such that each batching dim becomes an inserted window dim with a
// corresponding `IotaOp` concatenated to the scatter indices. This mirrors
// upstream StableHLO's `ScatterWithBatchingDimsExpander`. The other scatter
// canonicalization patterns do not understand batching dims and bail while any
// remain (see `failIfScatterHasBatchingDims`).
struct ScatterBatchingDimsExpander final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    auto dimNumbers = op.getScatterDimensionNumbers();
    ArrayRef<int64_t> inputBatchingDims = dimNumbers.getInputBatchingDims();
    ArrayRef<int64_t> scatterIndicesBatchingDims =
        dimNumbers.getScatterIndicesBatchingDims();
    if (inputBatchingDims.empty()) {
      return rewriter.notifyMatchFailure(op, "scatter op has no batching dims");
    }

    if (!cast<ShapedType>(op.getScatterIndices().getType()).hasStaticShape()) {
      return rewriter.notifyMatchFailure(
          op, "scatter indices have dynamic shape, can't expand");
    }

    SmallVector<int64_t> newInsertedWindowDims =
        mergeSortedDims(inputBatchingDims, dimNumbers.getInsertedWindowDims());
    SmallVector<int64_t> newScatterDimsToOperandDims =
        llvm::to_vector(llvm::concat<const int64_t>(
            inputBatchingDims, dimNumbers.getScatterDimsToOperandDims()));
    Value newIndices = createConcatIndices(
        op.getScatterIndices(), dimNumbers.getIndexVectorDim(),
        scatterIndicesBatchingDims, rewriter);

    auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
        op.getContext(), dimNumbers.getUpdateWindowDims(),
        newInsertedWindowDims,
        /*inputBatchingDims=*/{}, /*scatterIndicesBatchingDims=*/{},
        newScatterDimsToOperandDims, dimNumbers.getIndexVectorDim());

    auto newScatter = mlir::stablehlo::ScatterOp::create(
        rewriter, op.getLoc(), op->getResultTypes(), op.getInputs(), newIndices,
        op.getUpdates(), newDimNumbers,
        getUpdatedIndicesAreSorted(op.getIndicesAreSorted(),
                                   scatterIndicesBatchingDims,
                                   newScatterDimsToOperandDims),
        op.getUniqueIndices());
    newScatter.getUpdateComputation().takeBody(op.getUpdateComputation());
    rewriter.replaceOp(op, newScatter.getResults());
    return success();
  }
};

// The scatter canonicalization patterns below assume a scatter with no batching
// dims. This helper checks if there are any batching dims remaining so the
// patterns can bail and let them be expanded by `ScatterBatchingDimsExpander`.
static LogicalResult failIfScatterHasBatchingDims(mlir::stablehlo::ScatterOp op,
                                                  PatternRewriter &rewriter) {
  auto dimNumbers = op.getScatterDimensionNumbers();
  if (!dimNumbers.getInputBatchingDims().empty() ||
      !dimNumbers.getScatterIndicesBatchingDims().empty()) {
    return rewriter.notifyMatchFailure(
        op,
        "scatter has batching dims; ScatterBatchingDimsExpander runs first");
  }
  return success();
}

// Converts a `stablehlo.scatter` that writes a single, full-rank contiguous
// slice with an overwrite computation into a `stablehlo.dynamic_update_slice`.
//
// Such scatters (a scalar index selecting an offset into an operand dim that is
// *not* collapsed, i.e. a partial slice) are dynamic-update-slice semantics and
// cannot be represented as an `iree_linalg_ext.scatter`, which requires the
// index depth to map to fully-inserted leading operand dims.
struct ScatterToDynamicUpdateSlice final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getInputs().size() != 1 || op.getUpdates().size() != 1) {
      return rewriter.notifyMatchFailure(op, "variadic scatter");
    }
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    auto dimNumbers = op.getScatterDimensionNumbers();
    if (!dimNumbers.getInsertedWindowDims().empty()) {
      return rewriter.notifyMatchFailure(op, "has inserted dims");
    }

    Value operand = op.getInputs().front();
    Value update = op.getUpdates().front();
    Value indices = op.getScatterIndices();
    auto operandTy = dyn_cast<RankedTensorType>(operand.getType());
    auto updateTy = dyn_cast<RankedTensorType>(update.getType());
    auto indicesTy = dyn_cast<RankedTensorType>(indices.getType());
    if (!operandTy || !updateTy || !indicesTy || !operandTy.hasStaticShape() ||
        !updateTy.hasStaticShape() || !indicesTy.hasStaticShape()) {
      return rewriter.notifyMatchFailure(op, "dynamic/unranked shapes");
    }

    int64_t rank = operandTy.getRank();

    // The update must be a single full-rank contiguous slice: no batch dims and
    // an identity window mapping over all operand dims.
    ArrayRef<int64_t> updateWindowDims = dimNumbers.getUpdateWindowDims();
    if (updateTy.getRank() != rank ||
        static_cast<int64_t>(updateWindowDims.size()) != rank) {
      return rewriter.notifyMatchFailure(op, "update is not a full-rank slice");
    }
    for (auto [idx, dim] : llvm::enumerate(updateWindowDims)) {
      if (static_cast<int64_t>(idx) != dim) {
        return rewriter.notifyMatchFailure(op, "non-identity window dims");
      }
    }

    // There must be exactly one scatter index vector.
    int64_t indexVectorDim = dimNumbers.getIndexVectorDim();
    int64_t numUpdates = 1;
    for (int64_t d = 0, s = indicesTy.getRank(); d < s; ++d) {
      if (d != indexVectorDim) {
        numUpdates *= indicesTy.getDimSize(d);
      }
    }
    if (numUpdates != 1) {
      return rewriter.notifyMatchFailure(op, "more than one scatter index");
    }

    // The update computation must be a plain overwrite: return the update value
    // (block argument #1) unchanged.
    Region &region = op.getUpdateComputation();
    if (!region.hasOneBlock() || region.front().getNumArguments() != 2) {
      return rewriter.notifyMatchFailure(op, "unexpected update computation");
    }
    Block &block = region.front();
    auto retOp = dyn_cast<mlir::stablehlo::ReturnOp>(block.getTerminator());
    if (!retOp || retOp.getNumOperands() != 1 ||
        retOp.getOperand(0) != block.getArgument(1)) {
      return rewriter.notifyMatchFailure(op, "not an overwrite computation");
    }

    ArrayRef<int64_t> scatterDims = dimNumbers.getScatterDimsToOperandDims();
    int64_t indexDepth = scatterDims.size();

    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    Type indexElemTy = indicesTy.getElementType();
    auto scalarTy = RankedTensorType::get({}, indexElemTy);

    // Flatten indices to `indexDepth` scalar components (numUpdates == 1).
    auto flatTy = RankedTensorType::get({indexDepth}, indexElemTy);
    Value flatIndices = mlir::stablehlo::ReshapeOp::create(b, flatTy, indices);

    Value zero =
        mlir::stablehlo::ConstantOp::create(b, b.getZeroAttr(indexElemTy));

    // Map each operand dim to a start offset: the scatter index component if
    // the dim is indexed, otherwise zero. dynamic_update_slice clamps offsets.
    auto sliceTy = RankedTensorType::get({1}, indexElemTy);
    llvm::SmallVector<Value> startIndices(rank, zero);
    for (auto [j, operandDim] : llvm::enumerate(scatterDims)) {
      Value comp = mlir::stablehlo::SliceOp::create(
          b, sliceTy, flatIndices,
          b.getDenseI64ArrayAttr({static_cast<int64_t>(j)}),
          b.getDenseI64ArrayAttr({static_cast<int64_t>(j) + 1}),
          b.getDenseI64ArrayAttr({1}));
      startIndices[operandDim] =
          mlir::stablehlo::ReshapeOp::create(b, scalarTy, comp);
    }

    rewriter.replaceOpWithNewOp<mlir::stablehlo::DynamicUpdateSliceOp>(
        op, operandTy, operand, update, startIndices);
    return success();
  }
};

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

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    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 = mlir::stablehlo::ConvertOp::create(
        rewriter, op.getLoc(), indicesTy.clone(rewriter.getI32Type()), indices);

    auto newScatter = mlir::stablehlo::ScatterOp::create(
        rewriter, 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 Base::Base;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    auto dimNumbers = op.getScatterDimensionNumbers();
    auto indexVectorDim = dimNumbers.getIndexVectorDim();
    Value indices = op.getScatterIndices();
    auto indicesTy = 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 = tensor::ExpandShapeOp::create(rewriter, op.getLoc(), indicesTy,
                                            indices, reassociationMap);

    auto newScatter = mlir::stablehlo::ScatterOp::create(
        rewriter, 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 Base::Base;

  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 tensor::ExpandShapeOp::create(rewriter, loc, valueTy, value,
                                         reassociationMap);
  }

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    auto dimNumbers = op.getScatterDimensionNumbers();
    auto indexVectorDim = dimNumbers.getIndexVectorDim();
    auto indices = cast<Value>(op.getScatterIndices());
    auto indicesTy = 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 = mlir::stablehlo::ScatterOp::create(
        rewriter, 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 Base::Base;

  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 tensor::CollapseShapeOp::create(rewriter, loc, valueTy, value,
                                           reassociationMap);
  }

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    auto dimNumbers = op.getScatterDimensionNumbers();
    auto indexVectorDim = dimNumbers.getIndexVectorDim();
    auto indices = cast<Value>(op.getScatterIndices());
    auto indicesTy = 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 = mlir::stablehlo::ScatterOp::create(
        rewriter, 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 Base::Base;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    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 = 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 = mlir::stablehlo::TransposeOp::create(
          builder, indicesTy.clone(newShape), indices,
          builder.getDenseI64ArrayAttr(perm));
      indicesTy = cast<RankedTensorType>(indices.getType());
      indexVectorDim = indicesTy.getRank() - 1;
    }

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

    // Determine which dimensions are batch dimensions.
    llvm::SmallVector<bool> isBatch(updates0Ty.getRank(), true);
    for (int64_t updatedWindowDim : updatedWindowDims) {
      isBatch[updatedWindowDim] = 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 = cast<ShapedType>(update.getType());
        llvm::SmallVector<int64_t> newShape;
        newShape.reserve(updateTy.getRank());
        for (int64_t i : updatePerm) {
          newShape.push_back(updateTy.getDimSize(i));
        }
        update = mlir::stablehlo::TransposeOp::create(
            builder, 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 = mlir::stablehlo::ScatterOp::create(
        rewriter, 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();
  }
};

// Move indexed operand dimensions to the front so scatter_dims_to_operand_dims
// becomes [0, 1, ..., index_depth - 1]. This canonicalizes StableHLO scatter
// to the form expected by downstream IREE scatter lowering.
//
// This rewrite only permutes operand, update and result dimensions.
struct ScatterIndexedDimsFirst final
    : OpRewritePattern<mlir::stablehlo::ScatterOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    Location loc = op.getLoc();
    auto dimNumbers = op.getScatterDimensionNumbers();
    ArrayRef<int64_t> scatterDimsToOperandDims =
        dimNumbers.getScatterDimsToOperandDims();
    ArrayRef<int64_t> insertedWindowDims = dimNumbers.getInsertedWindowDims();
    ArrayRef<int64_t> updateWindowDims = dimNumbers.getUpdateWindowDims();
    ArrayRef<int64_t> inputBatchingDims = dimNumbers.getInputBatchingDims();
    ArrayRef<int64_t> scatterIndicesBatchingDims =
        dimNumbers.getScatterIndicesBatchingDims();

    // Keep batching cases out of this rewrite. This pattern only reorders
    // indexed and window dimensions within a single scatter slice layout.
    if (!inputBatchingDims.empty() || !scatterIndicesBatchingDims.empty()) {
      return rewriter.notifyMatchFailure(op, "batching dims are not handled");
    }

    auto isIdentityPermutation = [](ArrayRef<int64_t> perm) {
      return llvm::all_of(llvm::enumerate(perm), [](auto it) {
        return static_cast<int64_t>(it.index()) == it.value();
      });
    };

    if (isIdentityPermutation(scatterDimsToOperandDims)) {
      return rewriter.notifyMatchFailure(
          op, "indexed operand dims are already leading");
    }

    ValueRange inputs = op.getInputs();
    ValueRange updates = op.getUpdates();
    Value operand = inputs.front();
    Value update = updates.front();
    auto operandTy = dyn_cast<RankedTensorType>(operand.getType());
    if (!operandTy) {
      return rewriter.notifyMatchFailure(op, "operand has no ranked type");
    }
    auto updateTy = dyn_cast<RankedTensorType>(update.getType());
    if (!updateTy) {
      return rewriter.notifyMatchFailure(op, "update has no ranked type");
    }

    int64_t operandRank = operandTy.getRank();
    llvm::SmallVector<bool> isIndexedDim(operandRank, false);
    for (int64_t dim : scatterDimsToOperandDims) {
      isIndexedDim[dim] = true;
    }

    // Move indexed operand dims to the front while preserving the relative
    // order of all remaining dims.
    llvm::SmallVector<int64_t> operandPerm;
    operandPerm.reserve(operandRank);
    operandPerm.append(scatterDimsToOperandDims.begin(),
                       scatterDimsToOperandDims.end());
    for (int64_t dim = 0; dim < operandRank; ++dim) {
      if (!isIndexedDim[dim]) {
        operandPerm.push_back(dim);
      }
    }

    llvm::SmallVector<int64_t> inverseOperandPerm(operandRank);
    for (auto [newDim, oldDim] : llvm::enumerate(operandPerm)) {
      inverseOperandPerm[oldDim] = newDim;
    }

    // Recover the operand dims that correspond to update window dims so the
    // updates tensor can be permuted consistently with the operand.
    llvm::SmallVector<bool> isInsertedDim(operandRank, false);
    for (int64_t dim : insertedWindowDims) {
      isInsertedDim[dim] = true;
    }

    llvm::SmallVector<int64_t> windowOperandDims;
    windowOperandDims.reserve(updateWindowDims.size());
    for (int64_t dim = 0; dim < operandRank; ++dim) {
      if (!isInsertedDim[dim]) {
        windowOperandDims.push_back(dim);
      }
    }
    if (windowOperandDims.size() != updateWindowDims.size()) {
      return rewriter.notifyMatchFailure(
          op, "window dims do not match non-inserted operand dims");
    }

    // Map each operand window dim back to the update dim that materializes it.
    llvm::SmallVector<int64_t> windowOperandDimToUpdateDim(operandRank, -1);
    for (auto [operandDim, updateDim] :
         llvm::zip_equal(windowOperandDims, updateWindowDims)) {
      windowOperandDimToUpdateDim[operandDim] = updateDim;
    }

    int64_t updateRank = updateTy.getRank();
    llvm::SmallVector<bool> isUpdateWindowDim(updateRank, false);
    for (int64_t dim : updateWindowDims) {
      isUpdateWindowDim[dim] = true;
    }

    // Keep non-window update dims fixed. Permute window dims to match the
    // operand window dims after operandPerm is applied.
    llvm::SmallVector<int64_t> updatePerm;
    updatePerm.reserve(updateRank);
    auto reorderedWindowDims = llvm::make_filter_range(
        operandPerm, [&](int64_t dim) { return !isInsertedDim[dim]; });
    auto nextWindowDim = reorderedWindowDims.begin();
    for (int64_t dim = 0; dim < updateRank; ++dim) {
      if (isUpdateWindowDim[dim]) {
        int64_t operandDim = *nextWindowDim;
        ++nextWindowDim;
        updatePerm.push_back(windowOperandDimToUpdateDim[operandDim]);
      } else {
        updatePerm.push_back(dim);
      }
    }

    auto transposeType = [&](RankedTensorType type,
                             ArrayRef<int64_t> perm) -> RankedTensorType {
      llvm::SmallVector<int64_t> shape;
      shape.reserve(type.getRank());
      for (int64_t dim : perm) {
        shape.push_back(type.getDimSize(dim));
      }
      return RankedTensorType::get(shape, type.getElementType(),
                                   type.getEncoding());
    };

    auto createTranspose = [&](Value value, RankedTensorType type,
                               ArrayRef<int64_t> perm) -> Value {
      if (isIdentityPermutation(perm)) {
        return value;
      }
      return mlir::stablehlo::TransposeOp::create(
          rewriter, loc, transposeType(type, perm), value,
          rewriter.getDenseI64ArrayAttr(perm));
    };

    llvm::SmallVector<Value> transposedInputs;
    llvm::SmallVector<Type> transposedResultTypes;
    transposedInputs.reserve(inputs.size());
    transposedResultTypes.reserve(inputs.size());
    for (Value input : inputs) {
      auto inputTy = cast<RankedTensorType>(input.getType());
      transposedInputs.push_back(createTranspose(input, inputTy, operandPerm));
      transposedResultTypes.push_back(transposeType(inputTy, operandPerm));
    }

    llvm::SmallVector<Value> transposedUpdates;
    transposedUpdates.reserve(updates.size());
    for (Value updateValue : updates) {
      auto updateValueTy = cast<RankedTensorType>(updateValue.getType());
      transposedUpdates.push_back(
          createTranspose(updateValue, updateValueTy, updatePerm));
    }

    llvm::SmallVector<int64_t> newInsertedWindowDims;
    newInsertedWindowDims.reserve(insertedWindowDims.size());
    for (int64_t dim : insertedWindowDims) {
      newInsertedWindowDims.push_back(inverseOperandPerm[dim]);
    }
    llvm::sort(newInsertedWindowDims);

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

    auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
        op.getContext(), updateWindowDims, newInsertedWindowDims,
        inputBatchingDims, scatterIndicesBatchingDims,
        newScatterDimsToOperandDims, dimNumbers.getIndexVectorDim());

    auto newScatter = mlir::stablehlo::ScatterOp::create(
        rewriter, loc, transposedResultTypes, transposedInputs,
        op.getScatterIndices(), transposedUpdates, newDimNumbers,
        op.getIndicesAreSorted(), op.getUniqueIndices());
    Region &region = newScatter.getUpdateComputation();
    rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());

    llvm::SmallVector<Value> results;
    results.reserve(newScatter.getNumResults());
    for (auto [result, originalResultTy] :
         llvm::zip_equal(newScatter.getResults(), op.getResultTypes())) {
      results.push_back(mlir::stablehlo::TransposeOp::create(
          rewriter, loc, cast<RankedTensorType>(originalResultTy), result,
          rewriter.getDenseI64ArrayAttr(inverseOperandPerm)));
    }

    rewriter.replaceOp(op, results);
    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 Base::Base;

  LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
                                PatternRewriter &rewriter) const override {
    if (failed(failIfScatterHasBatchingDims(op, rewriter))) {
      return failure();
    }
    auto indices = op.getScatterIndices();
    Value operand = op.getInputs().front();
    auto indicesTy = cast<ShapedType>(indices.getType());
    auto operandTy = 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;
      }
    }

    auto toInsertDims =
        llvm::ArrayRef<bool>(isInsertDims).drop_front(frontInsertedDims);
    if (llvm::none_of(toInsertDims, [](bool 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(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 = 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 = tensor::ExpandShapeOp::create(
          rewriter, 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 = mlir::stablehlo::ScatterOp::create(
        rewriter, 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 = 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 Base::Base;

  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 = mlir::stablehlo::ConvertOp::create(
        rewriter, op.getLoc(), lhsTy.clone(rewriter.getIntegerType(1)), lhs);
    Value zero = mlir::stablehlo::ConstantOp::create(
        rewriter, op.getLoc(),
        DenseElementsAttr::get(RankedTensorType::get({}, eType),
                               rewriter.getZeroAttr(eType)));

    auto lhsShape = shape::ShapeOfOp::create(
        rewriter, 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 mlir::stablehlo::DynamicBroadcastInDimOp::create(
          rewriter, 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 Base::Base;

  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 = mlir::stablehlo::BroadcastOp::create(
        b, mag.getType(), op.getB(), make1DElementsAttr(b, halfNumElems));
    mag = mlir::stablehlo::MulOp::create(
        b, sigma, mlir::stablehlo::SqrtOp::create(b, mag));

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

    Value res = mlir::stablehlo::ConcatenateOp::create(b, 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 = tensor::ExtractSliceOp::create(b, res, zero, size, one);
    }
    if (resTy.getRank() != 1) {
      res = mlir::stablehlo::ReshapeOp::create(b, 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 = ElementwiseOpT::create(rewriter, 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();
  }
};

// Accepts fused widening operands for ops that have no additional type
// constraints.
template <typename Op>
LogicalResult verifyFusedWidenOperands(Op op, ArrayRef<Value> operands) {
  return success();
}

// Rejects convolution fusion when the resulting operand types violate
// StableHLO.
LogicalResult verifyFusedWidenOperands(mlir::stablehlo::ConvolutionOp op,
                                       ArrayRef<Value> operands) {
  Type lhsType = getElementTypeOrSelf(operands[0].getType());
  Type rhsType = getElementTypeOrSelf(operands[1].getType());
  if (!mlir::hlo::isCompatibleForHloTypeInference(lhsType, rhsType)) {
    return failure();
  }
  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. Fusion is skipped when it
// would produce invalid operand types.
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_if_present<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();
    }

    // Recheck type constraints after replacing widened operands with their
    // inputs.
    if (failed(verifyFusedWidenOperands(op, operands))) {
      return rewriter.notifyMatchFailure(
          op, "fused operands do not satisfy op type constraints");
    }

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

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

  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 = mlir::stablehlo::GetDimensionSizeOp::create(
        rewriter, op.getLoc(), lhs, rewriter.getI64IntegerAttr(0));
    Value batchSize1 = mlir::stablehlo::ReshapeOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get({1}, rewriter.getI32Type()), batchSize);

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

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

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

    rhs = mlir::stablehlo::DynamicBroadcastInDimOp::create(
        rewriter, 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 = mlir::stablehlo::ConvertOp::create(rewriter, op.getLoc(), computeTy,
                                             rhs);
    lhs = mlir::stablehlo::ConvertOp::create(rewriter, op.getLoc(), computeTy,
                                             lhs);

    auto result = mlir::stablehlo::MulOp::create(
        rewriter, 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 Base::Base;

  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 = mlir::stablehlo::RngBitGeneratorOp::create(
        rewriter, op.getLoc(), castedTypeRange, op.getRngAlgorithm(),
        op.getOperand());

    auto casted = mlir::stablehlo::BitcastConvertOp::create(
        rewriter, 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 Base::Base;

  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 Base::Base;

  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 = mlir::stablehlo::TransposeOp::create(
        builder,
        RankedTensorType::get(lhsTransposeShape, lhsTy.getElementType()), lhs,
        builder.getDenseI64ArrayAttr(permLhs));
    lhsTy = cast<RankedTensorType>(lhs.getType());

    rhs = mlir::stablehlo::TransposeOp::create(
        builder,
        RankedTensorType::get(rhsTransposeShape, rhsTy.getElementType()), rhs,
        builder.getDenseI64ArrayAttr(permRhs));
    rhsTy = 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 = mlir::stablehlo::GetDimensionSizeOp::create(builder, lhs, i);
      lhsReshapeDims.push_back(
          mlir::stablehlo::ReshapeOp::create(builder, dimI32Ty, dim));
    }
    Value lhsDynShape = mlir::stablehlo::ConcatenateOp::create(
        builder,
        RankedTensorType::get({static_cast<int64_t>(lhsReshapeDims.size())},
                              builder.getI32Type()),
        lhsReshapeDims, 0);
    lhsTy =
        RankedTensorType::get(lhsTy.getShape().drop_back(contractDimsL.size()),
                              lhsTy.getElementType());
    lhs = mlir::stablehlo::DynamicReshapeOp::create(builder, 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 = mlir::stablehlo::GetDimensionSizeOp::create(builder, rhs, i);
      rhsReshapeDims.push_back(
          mlir::stablehlo::ReshapeOp::create(builder, dimI32Ty, dim));
    }
    Value rhsDynShape = mlir::stablehlo::ConcatenateOp::create(
        builder,
        RankedTensorType::get({static_cast<int64_t>(rhsReshapeDims.size())},
                              builder.getI32Type()),
        rhsReshapeDims, 0);
    rhsTy =
        RankedTensorType::get(rhsTy.getShape().drop_back(contractDimsR.size()),
                              rhsTy.getElementType());
    rhs = mlir::stablehlo::DynamicReshapeOp::create(builder, 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 = mlir::stablehlo::ConcatenateOp::create(
        builder,
        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 Base::Base;

  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 = chlo::TopKOp::create(rewriter, 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_if_present<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_if_present<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 Base::Base;

  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 = chlo::TopKOp::create(rewriter, 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 Base::Base;

  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_if_present<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 =
          chlo::TopKOp::create(rewriter, 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 = mlir::stablehlo::NegOp::create(rewriter, op.getLoc(), input);
      auto topK =
          chlo::TopKOp::create(rewriter, 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<ScatterToDynamicUpdateSlice>(context, /*benefit=*/2);
    patterns.insert<ScatterBatchingDimsExpander, ScatterInt64Indices,
                    ScatterImplicitIndex, ScatterImplicitBatch,
                    ScatterMaterializeInsertedDim, ScatterCollapseBatch,
                    ScatterBatchFirst, ScatterIndexedDimsFirst>(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
