//===- EnzymeWrapPass.cpp - Replace calls with their derivatives ------------ //
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to create wrapper functions which differentiate
// ops.
//===----------------------------------------------------------------------===//

#include "Enzyme/MLIR/Dialect/Dialect.h"
#include "Enzyme/MLIR/Dialect/Ops.h"
#include "Enzyme/MLIR/Passes/EnzymeBatchPass.h"
#include "mlir/Analysis/Presburger/PresburgerRelation.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
#include "mlir/Dialect/CommonFolders.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Visitors.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"
#include "shardy/dialect/sdy/ir/utils.h"
#include "src/enzyme_ad/jax/Analysis/PartialSymmetryAnalysis.h"
#include "src/enzyme_ad/jax/CheckedRewrite.h"
#include "src/enzyme_ad/jax/Dialect/Dialect.h"
#include "src/enzyme_ad/jax/Dialect/Ops.h"
#include "src/enzyme_ad/jax/Implementations/ReferenceOps.h"
#include "src/enzyme_ad/jax/Implementations/WhileLoopInfo.h"
#include "src/enzyme_ad/jax/Passes/EnzymeHLOPatterns.h"
#include "src/enzyme_ad/jax/Passes/Passes.h"
#include "src/enzyme_ad/jax/Utils.h"
#include "stablehlo/dialect/Base.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"
#include "stablehlo/dialect/TypeInference.h"
#include "stablehlo/reference/Ops.h"
#include "stablehlo/reference/Scope.h"
#include "stablehlo/reference/Types.h"
#include "stablehlo/transforms/ChloDecompositionUtils.h"
#include "stablehlo/transforms/PassUtils.h"
#include "stablehlo/transforms/Passes.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

#include "Interfaces/AutoDiffTypeInterface.h"

#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallSet.h"

#include "llvm/ADT/MapVector.h"
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <llvm/ADT/MapVector.h>
#include <llvm/ADT/STLExtras.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/Support/Casting.h>
#include <llvm/Support/LogicalResult.h>
#include <mlir/IR/BuiltinAttributes.h>
#include <mlir/IR/Value.h>
#include <numeric>
#include <optional>
#include <ostream>
#define DEBUG_TYPE "enzymehloopt"

namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_ENZYMEHLOOPTPASS
#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

using namespace mlir;
using namespace mlir::enzyme;
using namespace mlir::stablehlo;

namespace mlir {
// Implementation of helper function to lower MultiRotateOp into individual
// RotateOps
LogicalResult lowerMultiRotateToRotates(enzymexla::MultiRotateOp op,
                                        PatternRewriter &rewriter) {
  int32_t leftAmount = op.getLeftAmount();
  int32_t rightAmount = op.getRightAmount();
  int32_t totalResults = leftAmount + rightAmount + 1;
  int32_t centerIdx = leftAmount;

  Value input = op.getOperand();
  auto inputType = cast<RankedTensorType>(input.getType());
  int64_t dimSize = inputType.getShape()[op.getDimension()];

  // Get sharding info if present
  auto shard = sdy::getShardingPerValue(op);

  SmallVector<Value> replacements(totalResults);

  for (int i = 0; i < totalResults; i++) {
    // Calculate rotation amount for this result
    // Result at centerIdx corresponds to amount 0 (identity)
    // Results before centerIdx have positive amounts (rotate left)
    // Results after centerIdx have negative amounts (rotate right)
    int32_t amount = centerIdx - i;

    // Normalize negative amounts to positive equivalent
    if (amount < 0) {
      amount += dimSize;
    }

    if (amount == 0) {
      replacements[i] = input;
      continue;
    }

    auto rotateOp = enzymexla::RotateOp::create(
        rewriter, op.getLoc(), inputType, input, amount, op.getDimension());

    // Propagate sharding if present
    if (shard) {
      sdy::TensorShardingAttr shards[1] = {shard.getShardings()[i]};
      auto shard2 =
          sdy::TensorShardingPerValueAttr::get(rotateOp.getContext(), shards);
      sdy::setShardings(rotateOp, shard2);
    }

    replacements[i] = rotateOp.getResult();
  }

  rewriter.replaceOp(op, replacements);
  return success();
}
} // namespace mlir

// Check if any of the pad sizes are negative
bool anyPadSizesNegative(stablehlo::PadOp pad) {
  for (auto &&[low, high, inner] :
       llvm::zip(pad.getEdgePaddingLow(), pad.getEdgePaddingHigh(),
                 pad.getInteriorPadding())) {
    if (low < 0 || high < 0 || inner < 0)
      return true;
  }
  return false;
}

namespace {

class ReshapeDimMapping {
public:
  void addMapping(int64_t left, int64_t right) {
    mapping.push_back(std::make_pair(left, right));
  }

  SmallVector<int64_t> getMappingFromResultDim(int64_t dim) const {
    SmallVector<int64_t> result;
    for (auto &[left, right] : mapping) {
      if (left == dim)
        result.push_back(right);
    }
    return result;
  }

  SmallVector<int64_t> getMappingFromOperandDim(int64_t dim) const {
    SmallVector<int64_t> result;
    for (auto &[left, right] : mapping) {
      if (right == dim)
        result.push_back(left);
    }
    return result;
  }

  bool isOnlySplitting() const {
    llvm::SmallDenseSet<int64_t> keys;
    for (auto &[left, right] : mapping) {
      if (!std::get<1>(keys.insert(left)))
        return false;
    }
    return true;
  }

  void dump() const {
    for (auto &[left, right] : mapping) {
      llvm::outs() << left << " -> " << right << "\n";
    }
  }

private:
  // Left is result dim, right is operand dim.
  SmallVector<std::pair<int64_t, int64_t>> mapping;
};

// Analyze if a reshape is clearly merging or splitting dimensions.
std::optional<ReshapeDimMapping>
tryFindReshapeDimMapping(stablehlo::ReshapeOp op) {
  ReshapeDimMapping mapping;
  int64_t lhsPos = 0;
  int64_t rhsPos = 0;
  auto lhsShape = cast<TensorType>(op.getOperand().getType()).getShape();
  auto rhsShape = cast<TensorType>(op.getResult().getType()).getShape();
  while (lhsPos < lhsShape.size() && rhsPos < rhsShape.size()) {
    if (lhsShape[lhsPos] == rhsShape[rhsPos]) {
      // Nice 1-to-1 mapping.
      mapping.addMapping(lhsPos, rhsPos);
    } else if (lhsShape[lhsPos] < rhsShape[rhsPos]) {
      // Potential many-to-one mapping.
      int64_t product = lhsShape[lhsPos];
      mapping.addMapping(lhsPos, rhsPos);
      while (product < rhsShape[rhsPos]) {
        if (++lhsPos >= lhsShape.size())
          break;
        product *= lhsShape[lhsPos];
        mapping.addMapping(lhsPos, rhsPos);
      }
      if (product != rhsShape[rhsPos])
        return std::nullopt;
    } else {
      // Potential one-to-many mapping.
      assert(lhsShape[lhsPos] > rhsShape[rhsPos]);
      int64_t product = rhsShape[rhsPos];
      mapping.addMapping(lhsPos, rhsPos);
      while (product < lhsShape[lhsPos]) {
        if (++rhsPos >= rhsShape.size())
          break;
        product *= rhsShape[rhsPos];
        mapping.addMapping(lhsPos, rhsPos);
      }
      if (product != lhsShape[lhsPos])
        return std::nullopt;
    }
    ++lhsPos;
    ++rhsPos;
  };
  return mapping;
}

class StaticSlice {
private:
  using VecTy = SmallVector<int64_t>;
  using TensorValue = TypedValue<RankedTensorType>;
  VecTy starts;
  VecTy limits;
  VecTy inputShape;
  VecTy outputShape;
  VecTy strides;
  unsigned rank;
  TensorValue input, output;
  RankedTensorType inputTy, outputTy;

public:
  TensorValue getOutput() { return output; }
  TensorValue getInput() { return input; }
  int64_t getBeginOffset(unsigned dim) { return starts[dim]; }
  int64_t getEndOffset(unsigned dim) { return inputShape[dim] - limits[dim]; }

  int64_t getOutputShape(unsigned dim) const {
    assert(dim < rank);
    return outputShape[dim];
  }

  bool isFullInDim(unsigned dim) const {
    assert(dim < rank);
    return starts[dim] == 0 && limits[dim] == inputShape[dim];
  }

  bool isSliceInDim(unsigned dim) const {
    assert(dim < rank);
    return starts[dim] != 0 || limits[dim] != inputShape[dim];
  }

  bool isFromStartInDim(unsigned dim) const {
    assert(dim < rank);
    return starts[dim] == 0;
  }

  bool isToEndInDim(unsigned dim) const {
    assert(dim < rank);
    return limits[dim] == inputShape[dim];
  }

  std::optional<unsigned> isOneDimSlice() const {
    std::optional<unsigned> found = std::nullopt;
    for (unsigned i = 0; i < rank; i++) {
      if (isSliceInDim(i)) {
        if (!found)
          found = i;
        else
          return std::nullopt;
      }
    }
    return found;
  }

  bool isFullSlice() const {
    for (unsigned i = 0; i < rank; i++)
      if (isSliceInDim(i))
        return false;
    return true;
  }

  bool isStrideOneAtDim(unsigned dim) const { return strides[dim] == 1; }

  bool isStrideOne() const {
    return llvm::all_of(strides, [](int64_t stride) { return stride == 1; });
  }

  static bool isEquivalentInDim(const StaticSlice &a, const StaticSlice &b,
                                unsigned dim) {
    if (a.rank != b.rank)
      return false;
    if (a.input != b.input)
      return false;

    return a.starts[dim] == b.starts[dim] && a.limits[dim] == b.limits[dim];
  }

  static bool isPrefixInDim(const StaticSlice &a, const StaticSlice &b,
                            unsigned dim) {
    if (!isEquivalentExceptDim(a, b, dim))
      return false;
    return a.starts[dim] == b.starts[dim] && a.limits[dim] <= b.limits[dim];
  }

  static bool isSuffixInDim(const StaticSlice &a, const StaticSlice &b,
                            unsigned dim) {
    if (!isEquivalentExceptDim(a, b, dim))
      return false;
    return a.starts[dim] >= b.starts[dim] && a.limits[dim] == b.limits[dim];
  }

  static bool isEquivalentExceptDim(const StaticSlice &a, const StaticSlice &b,
                                    unsigned dim) {
    return llvm::all_of(llvm::seq(a.rank), [&](unsigned i) {
      return dim == i || isEquivalentInDim(a, b, i);
    });
  }

  static std::optional<StaticSlice> get(Value v) {
    if (!v)
      return std::nullopt;

    StaticSlice res;
    RankedTensorType ty = dyn_cast<RankedTensorType>(v.getType());
    if (!ty)
      return std::nullopt;

    unsigned rank = ty.getRank();
    res.rank = rank;
    res.output = cast<TypedValue<RankedTensorType>>(v);
    res.outputTy = ty;

    if (stablehlo::SliceOp slice = v.getDefiningOp<stablehlo::SliceOp>()) {
      res.inputTy = slice.getOperand().getType();
      res.starts = VecTy(slice.getStartIndices());
      res.limits = VecTy(slice.getLimitIndices());
      res.strides = VecTy(slice.getStrides());
      res.input = slice.getOperand();
    } else {
      res.inputTy = ty;
      res.starts = VecTy(rank, 0);
      res.limits = VecTy(ty.getShape());
      res.strides = VecTy(rank, 1);
      res.input = res.output;
    }
    res.inputShape = VecTy(res.inputTy.getShape());
    res.outputShape = VecTy(res.outputTy.getShape());

    return res;
  }
};

template <typename OpTy, typename Child>
struct NoNanCheckedOpRewritePattern
    : public CheckedOpRewritePattern<OpTy, Child> {
  using Base = CheckedOpRewritePattern<OpTy, Child>;
  using Base::Base;

  NoNanCheckedOpRewritePattern(bool allowOnFloatingPointMath, MLIRContext *ctx,
                               PatternBenefit benefit = 1)
      : Base(ctx, benefit), allowOnFloatingPointMath(allowOnFloatingPointMath) {
  }

protected:
  bool allowOnFloatingPointMath;
};

struct NoopSlice final
    : CheckedOpRewritePattern<stablehlo::SliceOp, NoopSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();
    for (auto sz : type.getShape())
      if (sz < 0)
        return failure();
    if (op.getOperand().getType() == type) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }
    return failure();
  }
};

SmallVector<Value> promoteToLargestBitWidth(SmallVector<Value> values,
                                            PatternRewriter &rewriter) {
  if (values.empty())
    return {};

  int64_t maxBitWidth = 0;
  for (auto v : values) {
    auto tensorTy = cast<RankedTensorType>(v.getType());
    assert(tensorTy && tensorTy.getRank() == 0 && "Expected rank 0 tensor");
    auto intTy = cast<IntegerType>(tensorTy.getElementType());
    assert(intTy && "Expected integer type");
    maxBitWidth = std::max(maxBitWidth, static_cast<int64_t>(intTy.getWidth()));
  }

  SmallVector<Value> promotedValues(values.size());
  for (size_t i = 0; i < values.size(); i++) {
    auto tensorTy = cast<RankedTensorType>(values[i].getType());
    auto intTy = cast<IntegerType>(tensorTy.getElementType());

    if (intTy.getWidth() == maxBitWidth) {
      promotedValues[i] = values[i];
      continue;
    }

    promotedValues[i] = stablehlo::ConvertOp::create(
        rewriter, values[i].getLoc(),
        RankedTensorType::get({}, rewriter.getIntegerType(maxBitWidth)),
        values[i]);
  }
  return promotedValues;
}

void sliceSliceHelper(stablehlo::DynamicSliceOp prev,
                      SmallVector<Value> &startIndices,
                      SmallVector<int64_t> &sliceSizes,
                      PatternRewriter &rewriter) {
  assert(startIndices.size() == prev.getType().getShape().size());
  assert(sliceSizes.size() == prev.getType().getShape().size());

  auto prevStartIndices = prev.getStartIndices();
  for (size_t i = 0; i < startIndices.size(); ++i) {
    if (startIndices[i]) {
      SmallVector<Value> addInputs = promoteToLargestBitWidth(
          {startIndices[i], prevStartIndices[i]}, rewriter);
      startIndices[i] = stablehlo::AddOp::create(rewriter, prev.getLoc(),
                                                 addInputs[0], addInputs[1]);
    } else {
      startIndices[i] = prevStartIndices[i];
    }
  }
}

void sliceSliceHelper(stablehlo::SliceOp prev, SmallVector<int64_t> &starts,
                      SmallVector<int64_t> &limits,
                      SmallVector<int64_t> &strides) {
  assert(starts.size() == prev.getType().getShape().size());
  assert(limits.size() == prev.getType().getShape().size());
  assert(strides.size() == prev.getType().getShape().size());

  for (auto &&[pstart, pend, pstep, nstart, nend, nstep, size] : llvm::zip(
           prev.getStartIndices(), prev.getLimitIndices(), prev.getStrides(),
           starts, limits, strides, prev.getOperand().getType().getShape())) {

    auto start2 = pstart + pstep * nstart;
    auto step2 = pstep * nstep;
    auto end2 = pstart + pstep * nstart + pstep * (nend - nstart);
    if (start2 > size)
      start2 = size;
    if (end2 > size)
      end2 = size;
    nstart = start2;
    nstep = step2;
    nend = end2;
  }
}

struct SliceSlice final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto prev = op.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (!prev)
      return failure();

    SmallVector<int64_t> start(op.getStartIndices().begin(),
                               op.getStartIndices().end());
    SmallVector<int64_t> end(op.getLimitIndices().begin(),
                             op.getLimitIndices().end());
    SmallVector<int64_t> step(op.getStrides().begin(), op.getStrides().end());

    sliceSliceHelper(prev, start, end, step);
    auto resTy = op.getType();
    auto res = rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
        op, prev.getOperand(), start, end, step);
    assert(res.getType() == resTy);
    (void)res;
    (void)resTy;
    return success();
  }
};

void rewriteDynamicSliceDynamicSlice(stablehlo::DynamicSliceOp op,
                                     stablehlo::DynamicSliceOp prev,
                                     PatternRewriter &rewriter) {
  SmallVector<Value> startIndices(op.getStartIndices().begin(),
                                  op.getStartIndices().end());
  SmallVector<int64_t> sliceSizes(op.getSliceSizes().begin(),
                                  op.getSliceSizes().end());

  rewriter.setInsertionPoint(op);
  sliceSliceHelper(prev, startIndices, sliceSizes, rewriter);

  auto resTy = op.getType();
  auto res = rewriter.replaceOpWithNewOp<stablehlo::DynamicSliceOp>(
      op, prev.getOperand(), startIndices, sliceSizes);
  (void)res;
  (void)resTy;
  assert(res.getType() == resTy);
}

stablehlo::DynamicSliceOp SliceToDynamicSlice(stablehlo::SliceOp op,
                                              PatternRewriter &rewriter) {
  for (auto stride : op.getStrides()) {
    if (stride != 1) // dynamic slice only supports stride 1
      return nullptr;
  }

  rewriter.setInsertionPoint(op);

  SmallVector<Value> startIndices;
  SmallVector<int64_t> sliceSizes;
  auto elemTy = RankedTensorType::get({}, rewriter.getI32Type());
  for (auto &&[start, end] :
       llvm::zip(op.getStartIndices(), op.getLimitIndices())) {
    startIndices.push_back(stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), elemTy,
        cast<ElementsAttr>(makeAttr(elemTy, start))));
    sliceSizes.push_back(end - start);
  }

  auto res = rewriter.replaceOpWithNewOp<stablehlo::DynamicSliceOp>(
      op, op.getOperand(), startIndices, sliceSizes);
  return res;
}

struct DynamicSliceDynamicSlice final
    : CheckedOpRewritePattern<stablehlo::DynamicSliceOp,
                              DynamicSliceDynamicSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto prev = op.getOperand().getDefiningOp<stablehlo::DynamicSliceOp>();
    if (!prev || !llvm::hasSingleElement(prev->getUsers()))
      return failure();

    rewriteDynamicSliceDynamicSlice(op, prev, rewriter);
    return success();
  }
};

struct SliceDynamicSlice final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceDynamicSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto prev = op.getOperand().getDefiningOp<stablehlo::DynamicSliceOp>();
    if (!prev || !llvm::hasSingleElement(prev->getUsers()))
      return failure();

    auto curDS = SliceToDynamicSlice(op, rewriter);
    if (!curDS)
      return failure();

    rewriteDynamicSliceDynamicSlice(curDS, prev, rewriter);
    return success();
  }
};

struct DynamicSliceSlice final
    : CheckedOpRewritePattern<stablehlo::DynamicSliceOp, DynamicSliceSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto prev = op.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (!prev || !llvm::hasSingleElement(prev->getUsers()))
      return failure();

    auto prevDS = SliceToDynamicSlice(prev, rewriter);
    if (!prevDS)
      return failure();

    rewriteDynamicSliceDynamicSlice(op, prevDS, rewriter);
    return success();
  }
};

struct DynamicSliceToStatic final
    : CheckedOpRewritePattern<stablehlo::DynamicSliceOp, DynamicSliceToStatic> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    SmallVector<int64_t> starts;
    SmallVector<int64_t> ends;
    SmallVector<int64_t> steps;
    for (auto &&[start, size, shape] :
         llvm::zip(op.getStartIndices(), op.getSliceSizes(),
                   op.getOperand().getType().getShape())) {

      DenseIntElementsAttr startattr;
      if (!matchPattern(start, m_Constant(&startattr))) {
        return failure();
      }
      int64_t startv = (*startattr.begin()).getSExtValue();
      if (startv < 0)
        return failure();
      if (startv + size > shape)
        return failure();
      starts.push_back(startv);
      ends.push_back(startv + size);
      steps.push_back(1);
    }
    rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
        op, op.getType(), op.getOperand(), starts, ends, steps);
    return success();
  }
};

struct DynamicUpdateSliceElim final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                              DynamicUpdateSliceElim> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    if (op.getUpdate().getType() != type)
      return failure();

    for (auto start : op.getStartIndices()) {
      DenseIntElementsAttr startattr;
      if (!matchPattern(start, m_Constant(&startattr))) {
        return failure();
      }
      int64_t startv = (*startattr.begin()).getSExtValue();
      if (startv != 0)
        return failure();
    }
    rewriter.replaceOp(op, op.getUpdate());
    return success();
  }
};

// Given a reshape fromType to toType, remove the eliminated indices from start,
// and fill any new dimensions with toFill. Check that any removed indices have
// value checkRemoved, if set
template <typename T>
bool transformReshapeSlice(
    RankedTensorType fromType, RankedTensorType toType,
    SmallVectorImpl<T> &start, std::function<T()> toFillFn,
    std::function<bool(const T &)> checkRemovedFn = nullptr) {
  auto fromShape = fromType.getShape();
  auto toShape = toType.getShape();

  assert(start.size() == fromShape.size());

  int i = 0, j = 0;
  int startidx = 0;
  std::optional<T> toFillVal;
  while (i < fromShape.size() && j < toShape.size()) {
    if (fromShape[i] == toShape[j]) {
      // Leave everything as is and carry on
      i++;
      j++;
      startidx++;
      continue;
    }

    if (fromShape[i] == 1) {
      if (checkRemovedFn) {
        if (!checkRemovedFn(start[startidx])) {
          if (i > 0 && j > 0 && fromShape[i - 1] == fromShape[j - 1] &&
              fromShape[i - 1] == 1 && checkRemovedFn(start[startidx - 1])) {
            i--;
            j--;
            startidx--;
          } else {
            return false;
          }
        }
      }
      start.erase(start.begin() + startidx);
      i++;
      continue;
    }

    if (toShape[j] == 1) {
      T toFill = toFillVal ? *toFillVal : toFillFn();
      if (!toFillVal) {
        toFillVal = toFill;
      }
      start.insert(start.begin() + startidx, toFill);
      startidx++;
      j++;
      continue;
    }
    return false;
  }

  while (i < fromShape.size()) {
    if (fromShape[i] == 1) {
      if (checkRemovedFn) {
        if (!checkRemovedFn(start[startidx])) {
          return false;
        }
      }
      start.erase(start.begin() + startidx);
      i++;
      continue;
    }
    return false;
  }

  while (j < toShape.size()) {
    if (toShape[j] == 1) {
      T toFill = toFillVal ? *toFillVal : toFillFn();
      if (!toFillVal) {
        toFillVal = toFill;
      }
      start.insert(start.begin() + startidx, toFill);
      startidx++;
      j++;
      continue;
    }
  }
  assert(start.size() == toShape.size());
  return true;
}

template <typename T>
bool transformReshapeSlice(RankedTensorType fromType, RankedTensorType toType,
                           SmallVectorImpl<T> &start,
                           std::function<T()> toFillFn, T *checkRemoved) {
  std::function<bool(const T &)> checkRemovedFn;
  if (checkRemoved) {
    checkRemovedFn = [checkRemoved](const T &val) {
      return val == *checkRemoved;
    };
  }
  return transformReshapeSlice<T>(fromType, toType, start, toFillFn,
                                  checkRemovedFn);
}

template <typename T>
bool transformReshapeSlice(RankedTensorType fromType, RankedTensorType toType,
                           SmallVectorImpl<T> &start, T toFill,
                           T *checkRemoved = nullptr) {
  return transformReshapeSlice<T>(
      fromType, toType, start, [=]() { return toFill; }, checkRemoved);
}

template <typename T>
bool transformReshapeSlice(RankedTensorType fromType, RankedTensorType toType,
                           SmallVectorImpl<T> &start, T toFill,
                           std::function<bool(const T &)> checkRemovedFn) {
  return transformReshapeSlice<T>(
      fromType, toType, start, [=]() { return toFill; }, checkRemovedFn);
}

template <typename T>
bool transformReshapeSlice(stablehlo::ReshapeOp op, SmallVectorImpl<T> &start,
                           T toFill, T *checkRemoved = nullptr) {
  return transformReshapeSlice<T>(op.getOperand().getType(), op.getType(),
                                  start, toFill, checkRemoved);
}

template <typename T>
bool transformReshapeSlice(stablehlo::ReshapeOp op, SmallVectorImpl<T> &start,
                           T toFill,
                           std::function<bool(const T &)> checkRemovedFn) {
  return transformReshapeSlice<T>(op.getOperand().getType(), op.getType(),
                                  start, toFill, checkRemovedFn);
}

template <typename T>
bool transformReshapeSlice(stablehlo::ReshapeOp op, SmallVectorImpl<T> &start,
                           std::function<T()> toFill,
                           T *checkRemoved = nullptr) {
  return transformReshapeSlice<T>(op.getOperand().getType(), op.getType(),
                                  start, toFill, checkRemoved);
}

template <typename T>
bool transformReshapeSlice(stablehlo::ReshapeOp op, SmallVectorImpl<T> &start,
                           std::function<T()> toFill,
                           std::function<bool(const T &)> checkRemovedFn) {
  return transformReshapeSlice<T>(op.getOperand().getType(), op.getType(),
                                  start, toFill, checkRemovedFn);
}

stablehlo::Element conj(const stablehlo::Element &orig) {
  if (stablehlo::isSupportedComplexType(orig.getType())) {
    mlir::Complex<APFloat> val = orig.getComplexValue();
    return stablehlo::Element(orig.getType(),
                              mlir::Complex<APFloat>(val.real(), -val.imag()));
  }

  llvm_unreachable("Unsupported type");
}

stablehlo::Tensor conjOp(const stablehlo::Tensor &orig, ShapedType resultType) {
  stablehlo::Tensor result(resultType);
  for (auto it = result.index_begin(); it != result.index_end(); ++it)
    result.set(*it, conj(orig.get(*it)));
  return result;
}

struct ReshapeDUS final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Reshape is a DynamicUpdateSlice
    auto dus = op.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
    if (!dus)
      return failure();

    if (!llvm::hasSingleElement(dus->getUsers()))
      return failure();

    // %0 = dynamic_update_slice %arg0, %update
    // %1 = reshape %0
    //
    // ->
    //
    // %arg0_reshaped = reshape %arg0
    // %update_resahepd = reshape %update
    // dynamic_update_slice %arg0_reshaped, %update_reshaped

    SmallVector<Value> startIndices(dus.getStartIndices().begin(),
                                    dus.getStartIndices().end());
    auto itype = startIndices.size() > 0
                     ? cast<RankedTensorType>(startIndices[0].getType())
                     : RankedTensorType::get({}, rewriter.getI64Type());

    if (!transformReshapeSlice<mlir::Value>(
            op, startIndices, /*toFill*/
            [&]() -> mlir::Value {
              return stablehlo::ConstantOp::create(
                  rewriter, dus.getLoc(), itype,
                  cast<ElementsAttr>(makeAttr(itype, 0)));
            },
            [](mlir::Value v) -> bool { return matchPattern(v, m_Zero()); }))
      return failure();

    SmallVector<int64_t> updateShape(
        dus.getUpdate().getType().getShape().begin(),
        dus.getUpdate().getType().getShape().end());

    int64_t one = 1;
    if (!transformReshapeSlice<int64_t>(op, updateShape, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    auto newOperand = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), dus.getOperand(),
        cast<RankedTensorType>(op.getType()).getShape());

    auto newUpdate = stablehlo::ReshapeOpCreate(rewriter, op.getLoc(),
                                                dus.getUpdate(), updateShape);

    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
        op, newOperand, newUpdate, startIndices);

    return success();
  }
};

struct ReshapeSlice final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeSlice> {
  bool onlySingleUser;

  ReshapeSlice(bool onlySingleUser, MLIRContext *context,
               PatternBenefit benefit = 1,
               ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        onlySingleUser(onlySingleUser) {}

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Reshape is a DynamicUpdateSlice
    auto slice = op.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (!slice)
      return failure();

    if (onlySingleUser && !llvm::hasSingleElement(slice->getUsers()))
      return failure();

    SmallVector<int64_t> startIndices(slice.getStartIndices().begin(),
                                      slice.getStartIndices().end());

    int64_t zero = 0;
    int64_t one = 1;
    if (!transformReshapeSlice<int64_t>(op, startIndices, /*toFill*/ 0,
                                        /*checkRemoved*/ &zero))
      return failure();

    SmallVector<int64_t> limitIndices(slice.getLimitIndices().begin(),
                                      slice.getLimitIndices().end());

    if (!transformReshapeSlice<int64_t>(op, limitIndices, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    SmallVector<int64_t> stepIndices(slice.getStrides().begin(),
                                     slice.getStrides().end());

    if (!transformReshapeSlice<int64_t>(op, stepIndices, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    SmallVector<int64_t> operandShape(
        slice.getOperand().getType().getShape().begin(),
        slice.getOperand().getType().getShape().end());

    if (!transformReshapeSlice<int64_t>(op, operandShape, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    auto newOperand = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), slice.getOperand(), operandShape);

    rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
        op, newOperand, startIndices, limitIndices, stepIndices);

    return success();
  }
};

struct ReshapeDynamicSlice final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeDynamicSlice> {
  bool onlySingleUser;

  ReshapeDynamicSlice(bool onlySingleUser, MLIRContext *ctx,
                      PatternBenefit benefit = 1)
      : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeDynamicSlice>(
            ctx, benefit),
        onlySingleUser(onlySingleUser) {}

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Reshape is a DynamicSlice
    auto slice = op.getOperand().getDefiningOp<stablehlo::DynamicSliceOp>();
    if (!slice)
      return failure();

    if (onlySingleUser && !llvm::hasSingleElement(slice->getUsers()))
      return failure();

    // %0 = dynamic_slice %arg0
    // %1 = reshape %0
    //
    // ->
    //
    // %arg0_reshaped = reshape %arg0
    // dynamic_slice %arg0_reshaped
    SmallVector<Value> startIndices(slice.getStartIndices().begin(),
                                    slice.getStartIndices().end());

    auto itype = startIndices.size() > 0
                     ? cast<RankedTensorType>(startIndices[0].getType())
                     : RankedTensorType::get({}, rewriter.getI64Type());

    if (!transformReshapeSlice<mlir::Value>(
            op, startIndices, /*toFill*/
            [&]() -> mlir::Value {
              return stablehlo::ConstantOp::create(
                  rewriter, slice.getLoc(), itype,
                  cast<ElementsAttr>(makeAttr(itype, 0)));
            },
            [](mlir::Value v) -> bool { return matchPattern(v, m_Zero()); }))
      return failure();

    SmallVector<int64_t> sliceSizes = llvm::to_vector(slice.getSliceSizes());

    int64_t one = 1;
    if (!transformReshapeSlice<int64_t>(op, sliceSizes, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    SmallVector<int64_t> operandShape(
        slice.getOperand().getType().getShape().begin(),
        slice.getOperand().getType().getShape().end());

    if (!transformReshapeSlice<int64_t>(op, operandShape, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    auto newOperand = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), slice.getOperand(), operandShape);

    rewriter.replaceOpWithNewOp<stablehlo::DynamicSliceOp>(
        op, newOperand, startIndices,
        rewriter.getDenseI64ArrayAttr(sliceSizes));
    return success();
  }
};

// reshape(extend) -> extend(reshape)
struct ReshapeExtend final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeExtend> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Reshape is an ExtendOp
    auto extendOp = op.getOperand().getDefiningOp<enzymexla::ExtendOp>();
    if (!extendOp)
      return failure();

    if (!llvm::hasSingleElement(extendOp->getUsers()))
      return failure();

    // Get extend operation parameters
    int64_t extendDim = extendOp.getDimension();
    int64_t lhs = extendOp.getLhs();
    int64_t rhs = extendOp.getRhs();

    // Get the shape of the operand
    SmallVector<int64_t> operandShape(
        extendOp.getOperand().getType().getShape().begin(),
        extendOp.getOperand().getType().getShape().end());

    // Calculate the new extended dimension after the reshape
    int64_t one = 1;
    if (!transformReshapeSlice<int64_t>(op, operandShape, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    int64_t newExtendDim = -1;
    SmallVector<int64_t> oneHotDim(extendOp.getType().getShape().size(), 0);
    oneHotDim[extendDim] = 1;

    int64_t zero = 0;
    if (!transformReshapeSlice<int64_t>(op, oneHotDim, /*toFill*/ 0,
                                        /*checkRemoved*/ &zero))
      return failure();

    for (size_t i = 0; i < oneHotDim.size(); i++) {
      if (oneHotDim[i]) {
        if (newExtendDim != -1) {
          return failure();
        }
        newExtendDim = i;
      }
    }
    if (newExtendDim == -1)
      return failure();

    // First reshape the extend's operand
    auto newReshapeOp = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), extendOp.getOperand(), operandShape);

    // Then create a new extend operation on the reshaped data
    auto newExtendOp = enzymexla::ExtendOp::create(
        rewriter, op.getLoc(), newReshapeOp, lhs, rhs, newExtendDim);

    // Replace the original reshape op with the new extend operation
    rewriter.replaceOp(op, newExtendOp);

    return success();
  }
};

// reshape(wrap) -> wrap(reshape)
struct ReshapeWrap final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeWrap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Reshape is a WrapOp
    auto wrapOp = op.getOperand().getDefiningOp<enzymexla::WrapOp>();
    if (!wrapOp)
      return failure();

    if (!llvm::hasSingleElement(wrapOp->getUsers()))
      return failure();

    // Get wrap operation parameters
    int64_t wrapDim = wrapOp.getDimension();
    int64_t lhs = wrapOp.getLhs();
    int64_t rhs = wrapOp.getRhs();

    // Get the shape of the operand
    SmallVector<int64_t> operandShape(
        wrapOp.getOperand().getType().getShape().begin(),
        wrapOp.getOperand().getType().getShape().end());

    // Calculate the new wrapped dimension after the reshape
    int64_t one = 1;
    if (!transformReshapeSlice<int64_t>(op, operandShape, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    SmallVector<int64_t> oneHotDim(wrapOp.getType().getShape().size(), 0);
    oneHotDim[wrapDim] = 1;

    int64_t zero = 0;
    if (!transformReshapeSlice<int64_t>(op, oneHotDim, /*toFill*/ 0,
                                        /*checkRemoved*/ &zero))
      return failure();

    int64_t newWrapDim = -1;
    for (size_t i = 0; i < oneHotDim.size(); i++) {
      if (oneHotDim[i]) {
        if (newWrapDim != -1) {
          return failure();
        }
        newWrapDim = i;
      }
    }
    if (newWrapDim == -1)
      return failure();

    // First reshape the wrap's operand
    auto newReshapeOp = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), wrapOp.getOperand(), operandShape);

    // Then create a new wrap operation on the reshaped data
    auto newWrapOp =
        enzymexla::WrapOp::create(rewriter, op.getLoc(), op.getType(),
                                  newReshapeOp, lhs, rhs, newWrapDim);

    // Replace the original reshape op with the new wrap operation
    rewriter.replaceOp(op, newWrapOp);

    return success();
  }
};

// reshape(rotate) -> rotate(reshape)
struct ReshapeRotate final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Reshape is a RotateOp
    auto rotateOp = op.getOperand().getDefiningOp<enzymexla::RotateOp>();
    if (!rotateOp)
      return failure();

    if (!llvm::hasSingleElement(rotateOp->getUsers()))
      return failure();

    // Get rotate operation parameters
    int64_t rotateDim = rotateOp.getDimension();
    int64_t rotateAmount = rotateOp.getAmount();

    // Get the shape of the operand
    SmallVector<int64_t> operandShape(
        rotateOp.getOperand().getType().getShape().begin(),
        rotateOp.getOperand().getType().getShape().end());

    // Calculate the new rotated dimension after the reshape
    int64_t one = 1;
    if (!transformReshapeSlice<int64_t>(op, operandShape, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    SmallVector<int64_t> oneHotDim(rotateOp.getType().getShape().size(), 0);
    oneHotDim[rotateDim] = 1;

    int64_t zero = 0;
    if (!transformReshapeSlice<int64_t>(op, oneHotDim, /*toFill*/ 0,
                                        /*checkRemoved*/ &zero))
      return failure();

    int64_t newRotateDim = -1;
    for (size_t i = 0; i < oneHotDim.size(); i++) {
      if (oneHotDim[i]) {
        if (newRotateDim != -1) {
          return failure();
        }
        newRotateDim = i;
      }
    }
    if (newRotateDim == -1)
      return failure();

    // First reshape the rotate's operand
    auto newReshapeOp = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), rotateOp.getOperand(), operandShape);

    // Then create a new rotate operation on the reshaped data
    auto newRotateOp = enzymexla::RotateOp::create(
        rewriter, op.getLoc(), newReshapeOp, rotateAmount, newRotateDim);

    // Replace the original reshape op with the new rotate operation
    rewriter.replaceOp(op, newRotateOp);

    return success();
  }
};

struct ReshapePad final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Reshape is a DynamicUpdateSlice
    auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    if (!llvm::hasSingleElement(pad->getUsers()))
      return failure();

    SmallVector<int64_t> interior(pad.getInteriorPadding().begin(),
                                  pad.getInteriorPadding().end());

    int64_t zero = 0;
    if (!transformReshapeSlice<int64_t>(op, interior, /*toFill*/ 0,
                                        /*checkRemoved*/ &zero))
      return failure();

    SmallVector<int64_t> low(pad.getEdgePaddingLow().begin(),
                             pad.getEdgePaddingLow().end());

    if (!transformReshapeSlice<int64_t>(op, low, /*toFill*/ 0,
                                        /*checkRemoved*/ &zero))
      return failure();

    SmallVector<int64_t> high(pad.getEdgePaddingHigh().begin(),
                              pad.getEdgePaddingHigh().end());

    if (!transformReshapeSlice<int64_t>(op, high, /*toFill*/ 0,
                                        /*checkRemoved*/ &zero))
      return failure();

    SmallVector<int64_t> operandShape(
        pad.getOperand().getType().getShape().begin(),
        pad.getOperand().getType().getShape().end());

    int64_t one = 1;
    if (!transformReshapeSlice<int64_t>(op, operandShape, /*toFill*/ 1,
                                        /*checkRemoved*/ &one))
      return failure();

    auto newOperand = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), pad.getOperand(), operandShape);

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, newOperand, pad.getPaddingValue(), low, high, interior);

    return success();
  }
};

struct TransposeDUS final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    // Check if the input to Transpose is a DynamicUpdateSlice
    auto dus = op.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
    if (!dus)
      return failure();

    if (!llvm::hasSingleElement(dus->getUsers()))
      return failure();

    SmallVector<int64_t> permutation;
    for (auto perm : op.getPermutation()) {
      permutation.push_back(perm);
    }

    auto loc = op.getLoc();
    auto transposedOperand = stablehlo::TransposeOp::create(
        rewriter, loc, dus.getOperand(), op.getPermutation());
    auto transposedUpdate = stablehlo::TransposeOp::create(
        rewriter, loc, dus.getUpdate(), op.getPermutation());

    SmallVector<Value> permutedStartIndices;
    permutedStartIndices.resize(dus.getStartIndices().size());
    for (auto [i, permIndex] : llvm::enumerate(permutation)) {
      permutedStartIndices[i] = dus.getStartIndices()[permIndex];
    }

    auto newDus = stablehlo::DynamicUpdateSliceOp::create(
        rewriter, loc, op.getType(), transposedOperand, transposedUpdate,
        permutedStartIndices);

    rewriter.replaceOp(op, newDus);
    rewriter.eraseOp(dus);
    return success();
  }
};

struct LowerWrap
    : public CheckedOpRewritePattern<enzymexla::WrapOp, LowerWrap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::WrapOp wrap,
                                    PatternRewriter &rewriter) const {
    auto concat = lowerWrap(wrap, rewriter, /*replace*/ true);
    if (concat.getInputs().size() == 1) {
      rewriter.replaceOp(concat, concat.getInputs()[0]);
    }
    return success();
  }
};

stablehlo::ConcatenateOp lowerExtend(enzymexla::ExtendOp extend,
                                     PatternRewriter &rewriter, bool replace) {
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(extend);

  auto loc = extend.getLoc();
  auto operand = extend.getOperand();

  SmallVector<Value> args;

  SmallVector<int64_t> strides(operand.getType().getRank(), 1);
  auto shard = sdy::getShardingPerValue(extend);

  if (extend.getLhs() != 0) {
    SmallVector<int64_t> lhsStarts(operand.getType().getRank(), 0);
    SmallVector<int64_t> lhsLimits(operand.getType().getShape());
    lhsLimits[extend.getDimension()] = extend.getLhs();
    auto lhs = stablehlo::SliceOp::create(rewriter, loc, operand, lhsStarts,
                                          lhsLimits, strides);
    if (shard)
      sdy::setShardings(lhs, shard);

    args.push_back(lhs);
  }

  args.push_back(operand);

  if (extend.getRhs() != 0) {
    SmallVector<int64_t> rhsStarts(operand.getType().getRank(), 0);
    SmallVector<int64_t> rhsLimits(operand.getType().getShape());
    rhsStarts[extend.getDimension()] =
        rhsLimits[extend.getDimension()] - extend.getRhs();
    auto rhs = stablehlo::SliceOp::create(rewriter, loc, operand, rhsStarts,
                                          rhsLimits, strides);
    if (shard)
      sdy::setShardings(rhs, shard);

    args.push_back(rhs);
  }

  auto newConcat = stablehlo::ConcatenateOp::create(
      rewriter, extend.getLoc(), args, extend.getDimension());
  if (replace)
    rewriter.replaceOp(extend, newConcat);
  if (shard)
    sdy::setShardings(newConcat, shard);
  return newConcat;
}

struct LowerExtend
    : public CheckedOpRewritePattern<enzymexla::ExtendOp, LowerExtend> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::ExtendOp extend,
                                    PatternRewriter &rewriter) const {
    auto concat = lowerExtend(extend, rewriter, /*replace*/ true);
    if (concat.getInputs().size() == 1) {
      rewriter.replaceOp(concat, concat.getInputs()[0]);
    }
    return success();
  }
};

stablehlo::ConcatenateOp lowerRotate(enzymexla::RotateOp rotate,
                                     PatternRewriter &rewriter, bool replace) {
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(rotate);

  // sl0[A:end], sl1[0:A]
  auto shard = sdy::getShardingPerValue(rotate);
  SmallVector<int64_t> strides(rotate.getType().getShape().size(), 1);
  SmallVector<int64_t> sl0_starts(rotate.getType().getShape().size(), 0);
  SmallVector<int64_t> sl0_ends(rotate.getType().getShape());
  SmallVector<int64_t> sl1_starts(rotate.getType().getShape().size(), 0);
  SmallVector<int64_t> sl1_ends(rotate.getType().getShape());
  sl0_starts[rotate.getDimension()] = rotate.getAmount();
  sl1_ends[rotate.getDimension()] = rotate.getAmount();
  auto sl0 =
      stablehlo::SliceOp::create(rewriter, rotate.getLoc(), rotate.getOperand(),
                                 sl0_starts, sl0_ends, strides);
  if (shard) {
    sdy::setShardings(sl0, shard);
  }
  auto sl1 =
      stablehlo::SliceOp::create(rewriter, rotate.getLoc(), rotate.getOperand(),
                                 sl1_starts, sl1_ends, strides);
  if (shard) {
    sdy::setShardings(sl1, shard);
  }
  Value args[] = {sl0, sl1};
  auto newConcat = stablehlo::ConcatenateOp::create(
      rewriter, rotate.getLoc(), args, rotate.getDimension());
  if (replace)
    rewriter.replaceOp(rotate, newConcat);
  if (shard) {
    sdy::setShardings(newConcat, shard);
  }
  return newConcat;
}

struct LowerRotate
    : public CheckedOpRewritePattern<enzymexla::RotateOp, LowerRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::RotateOp rotate,
                                    PatternRewriter &rewriter) const {
    lowerRotate(rotate, rewriter, /*replace*/ true);
    return success();
  }
};

struct LowerUpdateWithoutCorners
    : public CheckedOpRewritePattern<enzymexla::UpdateWithoutCornersOp,
                                     LowerUpdateWithoutCorners> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::UpdateWithoutCornersOp up,
                                    PatternRewriter &rewriter) const {
    mlir::enzyme::commonLowerUpdateWithoutCorners(up, rewriter);
    return success();
  }
};

static Value createInBoxCondition(OpBuilder &rewriter, Location loc,
                                  ArrayRef<int64_t> shape,
                                  ArrayRef<int64_t> bounds) {
  Value condition;
  for (int64_t dim = 0; dim < shape.size(); dim++) {
    auto i32TensorType = RankedTensorType::get(shape, rewriter.getI32Type());
    Value iota = stablehlo::IotaOp::create(rewriter, loc, i32TensorType, dim);

    Value lb = stablehlo::ConstantOp::create(
        rewriter, loc,
        SplatElementsAttr::get(i32TensorType,
                               rewriter.getI32IntegerAttr(bounds[2 * dim])));
    Value ub = stablehlo::ConstantOp::create(
        rewriter, loc,
        SplatElementsAttr::get(
            i32TensorType, rewriter.getI32IntegerAttr(bounds[2 * dim + 1])));

    Value cmp1 = stablehlo::CompareOp::create(
        rewriter, loc, iota, lb, stablehlo::ComparisonDirection::GE);
    Value cmp2 = stablehlo::CompareOp::create(
        rewriter, loc, iota, ub, stablehlo::ComparisonDirection::LT);
    Value cmp = stablehlo::AndOp::create(rewriter, loc, cmp1, cmp2);
    if (!condition) {
      condition = cmp;
      continue;
    }
    condition = stablehlo::AndOp::create(rewriter, loc, condition, cmp);
  }
  return condition;
}

struct LowerPiecewiseSelect
    : public CheckedOpRewritePattern<enzymexla::PiecewiseSelectOp,
                                     LowerPiecewiseSelect> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::PiecewiseSelectOp op,
                                    PatternRewriter &rewriter) const {
    Value condition;
    for (Attribute box : op.getBoxes()) {
      SmallVector<int64_t> boxCoords = llvm::map_to_vector(
          cast<ArrayAttr>(box).getAsValueRange<IntegerAttr>(),
          [](APInt apInt) { return apInt.getSExtValue(); });
      Value boxCondition = createInBoxCondition(
          rewriter, op.getLoc(), op.getType().getShape(), boxCoords);
      if (!condition) {
        condition = boxCondition;
        continue;
      }
      condition = stablehlo::OrOp::create(rewriter, op.getLoc(), condition,
                                          boxCondition);
    }

    auto select =
        stablehlo::SelectOp::create(rewriter, op.getLoc(), condition,
                                    op.getWhenInBox(), op.getWhenOutOfBox());
    rewriter.replaceOp(op, select);

    return success();
  }
};

// Optimization: DUSConcat
// Pattern:
//   %concat = stablehlo.concatenate %A, %B, %C, dimension = D
//   %dus = stablehlo.dynamic_update_slice %concat, %update, %idx0, %idxD, %idxN
// Constraint: %concat has only %dus as its user.
// Constraint: The slice defined by %idx0..N and %update shape falls entirely
//             within one of the inputs (e.g., %B) along dimension D, and
//             spans fully along other dimensions.
// Rewrite to:
//   %idxD_new = %idxD - offset_of_B
//   %new_dus = stablehlo.dynamic_update_slice %B, %update, %idx0_0, %idxD_new,
//   %idxN_0 %new_concat = stablehlo.concatenate %A, %new_dus, %C, dimension = D
//   replaceAllUsesWith(%dus, %new_concat)
struct DUSConcat final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp, DUSConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dus,
                                    PatternRewriter &rewriter) const {
    auto updateType = cast<RankedTensorType>(dus.getUpdate().getType());

    ArrayRef<int64_t> updateShape = updateType.getShape();
    SmallVector<Value> startIndices = dus.getStartIndices();

    int64_t currentOffset = 0;
    int targetInputIdx = -1;

    int64_t concatDim;
    SmallVector<int64_t> inputSizes;
    bool legal = false;

    if (auto concatOp =
            dus.getOperand().getDefiningOp<stablehlo::ConcatenateOp>()) {
      concatDim = concatOp.getDimension();
      for (auto input : concatOp.getInputs()) {
        inputSizes.push_back(
            cast<RankedTensorType>(input.getType()).getShape()[concatDim]);
      }
      legal = true;
    }
    if (auto wrap = dus.getOperand().getDefiningOp<enzymexla::WrapOp>()) {
      concatDim = wrap.getDimension();
      inputSizes.push_back(wrap.getLhs());
      inputSizes.push_back(cast<RankedTensorType>(wrap.getOperand().getType())
                               .getShape()[concatDim]);
      inputSizes.push_back(wrap.getRhs());
      legal = true;
    }
    if (auto extend = dus.getOperand().getDefiningOp<enzymexla::ExtendOp>()) {
      concatDim = extend.getDimension();
      inputSizes.push_back(extend.getLhs());
      inputSizes.push_back(cast<RankedTensorType>(extend.getOperand().getType())
                               .getShape()[concatDim]);
      inputSizes.push_back(extend.getRhs());
      legal = true;
    }
    if (auto rotate = dus.getOperand().getDefiningOp<enzymexla::RotateOp>()) {
      concatDim = rotate.getDimension();
      inputSizes.push_back(cast<RankedTensorType>(rotate.getOperand().getType())
                               .getShape()[concatDim] -
                           rotate.getAmount());
      inputSizes.push_back(rotate.getAmount());
      legal = true;
    }

    if (!legal) {
      return failure();
    }

    DenseIntElementsAttr concatStartAttr;
    if (!matchPattern(startIndices[concatDim], m_Constant(&concatStartAttr))) {
      return rewriter.notifyMatchFailure(
          dus, "Requires constant start index for concat dimension");
    }

    int64_t concatStartVal = (*concatStartAttr.begin()).getSExtValue();
    int64_t concatUpdateSize = updateShape[concatDim];
    int64_t concatEndVal = concatStartVal + concatUpdateSize; // Exclusive end

    for (const auto &indexedInput : llvm::enumerate(inputSizes)) {
      int64_t inputSize = indexedInput.value();

      // Check if the DUS update region falls entirely within this input's
      // region
      if (concatStartVal >= currentOffset &&
          concatEndVal <= (currentOffset + inputSize)) {
        targetInputIdx = indexedInput.index();
        break; // Found the target input
      }
      currentOffset += inputSize;
    }

    // If no suitable input was found
    if (targetInputIdx == -1) {
      return rewriter.notifyMatchFailure(
          dus,
          "DUS update region does not fall entirely within one concat input");
    }

    auto concatOp = dus.getOperand().getDefiningOp<stablehlo::ConcatenateOp>();

    if (auto wrap = dus.getOperand().getDefiningOp<enzymexla::WrapOp>()) {
      concatOp = lowerWrap(wrap, rewriter, /*replace*/ false);
    } else if (auto extend =
                   dus.getOperand().getDefiningOp<enzymexla::ExtendOp>()) {
      concatOp = lowerExtend(extend, rewriter, /*replace*/ false);
    } else if (auto rotate =
                   dus.getOperand().getDefiningOp<enzymexla::RotateOp>()) {
      concatOp = lowerRotate(rotate, rewriter, /*replace*/ false);
    }
    assert(concatOp);
    assert(concatOp.getType() == dus.getOperand().getType());

    SmallVector<Value> newDusStartIndices =
        llvm::to_vector(dus.getStartIndices());
    Location loc = dus.getLoc();

    newDusStartIndices[concatDim] = stablehlo::ConstantOp::create(
        rewriter, dus.getLoc(), newDusStartIndices[concatDim].getType(),
        cast<ElementsAttr>(makeAttr(newDusStartIndices[concatDim].getType(),
                                    concatStartVal - currentOffset)));

    Value targetInputVal = concatOp.getInputs()[targetInputIdx];

    // Create the new, smaller DUS.
    auto newDus = stablehlo::DynamicUpdateSliceOp::create(
        rewriter, loc,
        targetInputVal
            .getType(), // Result type is the same as the input being updated
        targetInputVal, dus.getUpdate(), newDusStartIndices);

    // Create the list of operands for the new concatenate op.
    SmallVector<Value> newConcatOperands =
        llvm::to_vector(concatOp.getInputs());
    newConcatOperands[targetInputIdx] = newDus.getResult();

    // Create the new concatenate op.
    auto newConcat = stablehlo::ConcatenateOp::create(
        rewriter,
        concatOp
            .getLoc(), // Use concatOp's location, maybe dus.getLoc() is better?
        newConcatOperands, concatDim);
    assert(
        newConcat.getType() ==
        dus.getType()); // The result type should match the original DUS result

    // Replace the original DUS with the new concatenate op.
    rewriter.replaceOp(dus, newConcat.getResult());

    return success();
  }
};

template <typename T>
struct SimplifyBoundary final
    : CheckedOpRewritePattern<T, SimplifyBoundary<T>> {
  using CheckedOpRewritePattern<T,
                                SimplifyBoundary<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {

    SplatElementsAttr elems;
    if (!matchPattern(op.getOperand(), m_Constant(&elems))) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, op.getType(), elems.resizeSplat(op.getType()));
    return success();
  }
};

struct SliceInternal final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceInternal> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp slice,
                                    PatternRewriter &rewriter) const {
    int64_t currentOffset = 0;
    int targetInputIdx = -1;

    int64_t concatDim;
    SmallVector<int64_t> inputSizes;
    SmallVector<int64_t> actualStartSizes;
    bool legal = false;

    Value operand;

    if (auto wrap = slice.getOperand().getDefiningOp<enzymexla::WrapOp>()) {
      concatDim = wrap.getDimension();
      inputSizes.push_back(wrap.getLhs());
      inputSizes.push_back(cast<RankedTensorType>(wrap.getOperand().getType())
                               .getShape()[concatDim]);
      inputSizes.push_back(wrap.getRhs());
      actualStartSizes.push_back(
          cast<RankedTensorType>(wrap.getOperand().getType())
              .getShape()[concatDim] -
          wrap.getLhs());
      actualStartSizes.push_back(0);
      actualStartSizes.push_back(0);
      legal = true;
      operand = wrap.getOperand();
    }
    if (auto extend = slice.getOperand().getDefiningOp<enzymexla::ExtendOp>()) {
      concatDim = extend.getDimension();
      inputSizes.push_back(extend.getLhs());
      inputSizes.push_back(cast<RankedTensorType>(extend.getOperand().getType())
                               .getShape()[concatDim]);
      inputSizes.push_back(extend.getRhs());
      actualStartSizes.push_back(0);
      actualStartSizes.push_back(0);
      actualStartSizes.push_back(
          cast<RankedTensorType>(extend.getOperand().getType())
              .getShape()[concatDim]);
      legal = true;
      operand = extend.getOperand();
    }
    if (auto rotate = slice.getOperand().getDefiningOp<enzymexla::RotateOp>()) {
      concatDim = rotate.getDimension();
      inputSizes.push_back(cast<RankedTensorType>(rotate.getOperand().getType())
                               .getShape()[concatDim] -
                           rotate.getAmount());
      inputSizes.push_back(rotate.getAmount());
      actualStartSizes.push_back(rotate.getAmount());
      actualStartSizes.push_back(0);
      operand = rotate.getOperand();
      legal = true;
    }

    if (!legal) {
      return failure();
    }

    for (const auto &indexedInput : llvm::enumerate(inputSizes)) {
      int64_t inputSize = indexedInput.value();

      // Check if the DUS update region falls entirely within this input's
      // region
      if (slice.getStartIndices()[concatDim] >= currentOffset &&
          slice.getLimitIndices()[concatDim] <= (currentOffset + inputSize)) {
        targetInputIdx = indexedInput.index();
        break; // Found the target input
      }
      currentOffset += inputSize;
    }

    // If no suitable input was found
    if (targetInputIdx == -1) {
      return rewriter.notifyMatchFailure(
          slice, "Slice region does not fall entirely within one concat input");
    }

    SmallVector<int64_t> newStart = llvm::to_vector(slice.getStartIndices());
    newStart[concatDim] -= currentOffset;
    newStart[concatDim] += actualStartSizes[targetInputIdx];
    SmallVector<int64_t> newLimit = llvm::to_vector(slice.getLimitIndices());
    newLimit[concatDim] -= currentOffset;
    newLimit[concatDim] += actualStartSizes[targetInputIdx];

    rewriter.modifyOpInPlace(slice, [&]() {
      slice.setStartIndices(newStart);
      slice.setLimitIndices(newLimit);
      slice.getOperandMutable().assign(operand);
    });
    return success();
  }
};

struct ConcatConcatToDUS final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatConcatToDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp outer,
                                    PatternRewriter &rewriter) const {
    if (outer.getOperands().size() < 2)
      return failure();
    SmallVector<stablehlo::ConcatenateOp> inners;

    for (auto v : outer.getOperands()) {
      auto concatOp = v.getDefiningOp<stablehlo::ConcatenateOp>();
      if (!concatOp)
        return failure();
      if (concatOp.getDimension() == outer.getDimension())
        return failure();
      if (concatOp.getOperands().size() < 2)
        return failure();
      if (inners.size()) {
        if (inners[0].getOperands().size() != concatOp.getOperands().size()) {
          return failure();
        }
        if (inners[0].getDimension() != concatOp.getDimension()) {
          return failure();
        }
      }
      inners.push_back(concatOp);
    }

    stablehlo::SliceOp lhs = nullptr;
    if (auto lhsSlice =
            inners[0].getOperands()[0].getDefiningOp<stablehlo::SliceOp>()) {
      bool legal = true;
      size_t nextStart = lhsSlice.getLimitIndices()[outer.getDimension()];
      for (int i = 1; i < inners.size(); i++) {
        auto slice2 =
            inners[i].getOperands()[0].getDefiningOp<stablehlo::SliceOp>();
        if (!slice2) {
          legal = false;
          break;
        }
        if (slice2.getOperand() != lhsSlice.getOperand()) {
          legal = false;
          break;
        }
        for (int j = 0; j < lhsSlice.getType().getShape().size(); j++) {
          if (j == outer.getDimension()) {
            if (slice2.getStartIndices()[j] != nextStart) {
              legal = false;
              break;
            }
          } else {
            if (slice2.getStartIndices()[j] != lhsSlice.getStartIndices()[j]) {
              legal = false;
              break;
            }
            if (slice2.getLimitIndices()[j] != lhsSlice.getLimitIndices()[j]) {
              legal = false;
              break;
            }
            if (slice2.getStrides()[j] != 1 || lhsSlice.getStrides()[j] != 1) {
              legal = false;
              break;
            }
          }
        }
        nextStart = slice2.getLimitIndices()[outer.getDimension()];
      }
      if (legal) {
        lhs = lhsSlice;
      }
    }

    stablehlo::SliceOp rhs = nullptr;
    if (auto rhsSlice = inners[0]
                            .getOperands()
                            .back()
                            .getDefiningOp<stablehlo::SliceOp>()) {
      bool legal = true;
      size_t nextStart = rhsSlice.getLimitIndices()[outer.getDimension()];
      for (int i = 1; i < inners.size(); i++) {
        auto slice2 =
            inners[i].getOperands().back().getDefiningOp<stablehlo::SliceOp>();
        if (!slice2) {
          legal = false;
          break;
        }
        if (slice2.getOperand() != rhsSlice.getOperand()) {
          legal = false;
          break;
        }
        for (int j = 0; j < rhsSlice.getType().getShape().size(); j++) {
          if (j == outer.getDimension()) {
            if (slice2.getStartIndices()[j] != nextStart) {
              legal = false;
              break;
            }
          } else {
            if (slice2.getStartIndices()[j] != rhsSlice.getStartIndices()[j]) {
              legal = false;
              break;
            }
            if (slice2.getLimitIndices()[j] != rhsSlice.getLimitIndices()[j]) {
              legal = false;
              break;
            }
            if (slice2.getStrides()[j] != 1 || rhsSlice.getStrides()[j] != 1) {
              legal = false;
              break;
            }
          }
        }
        nextStart = slice2.getLimitIndices()[outer.getDimension()];
      }
      if (legal) {
        rhs = rhsSlice;
      }
    }

    if (rhs && lhs && rhs.getOperand() == lhs.getOperand() &&
        rhs.getLimitIndices()[inners[0].getDimension()] -
                lhs.getStartIndices()[inners[0].getDimension()] ==
            outer.getType().getShape()[inners[0].getDimension()]) {
      bool legal = true;
      for (int j = 0; j < lhs.getType().getShape().size(); j++) {
        if (j != inners[0].getDimension()) {
          if (lhs.getStartIndices()[j] != rhs.getStartIndices()[j]) {
            legal = false;
            break;
          }
          if (lhs.getLimitIndices()[j] != rhs.getLimitIndices()[j]) {
            legal = false;
            break;
          }
        }
      }
      if (legal) {
        SmallVector<Value> newInner;
        for (auto concat : inners) {
          SmallVector<Value> innerOps;
          for (int i = 1; i < concat.getOperands().size() - 1; i++) {
            innerOps.push_back(concat.getOperands()[i]);
          }
          newInner.push_back(stablehlo::ConcatenateOp::create(
              rewriter, concat.getLoc(), innerOps, concat.getDimension()));
        }
        auto newOuter = stablehlo::ConcatenateOp::create(
            rewriter, outer.getLoc(), newInner, outer.getDimension());
        auto newSlice = stablehlo::SliceOp::create(
            rewriter, outer.getLoc(), lhs.getOperand(), lhs.getStartIndices(),
            inners.back()
                .getOperands()
                .back()
                .getDefiningOp<stablehlo::SliceOp>()
                .getLimitIndices(),
            lhs.getStrides());
        auto iTy = RankedTensorType::get({}, rewriter.getI64Type());
        SmallVector<Value> starts(lhs.getType().getShape().size(),
                                  stablehlo::ConstantOp::create(
                                      rewriter, outer.getLoc(), iTy,
                                      cast<ElementsAttr>(makeAttr(iTy, 0))));
        starts[outer.getDimension()] = stablehlo::ConstantOp::create(
            rewriter, outer.getLoc(),
            cast<ElementsAttr>(
                makeAttr(iTy, lhs.getType().getShape()[outer.getDimension()])));
        rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
            outer, newSlice, newOuter, starts);
        return success();
      }
    }

    if (lhs) {
      llvm::errs() << "TODO: implement LHS side concat optimization\n";
    }

    if (rhs) {
      llvm::errs() << "TODO: implement RHS side concat optimization\n";
    }
    return failure();
  }
};

enum class SliceRequirement {
  AnyOperand = 0,
  NeedsSlice = 1,
  NeedsFull = 2,
};

stablehlo::DynamicUpdateSliceOp
concat_to_dus_slice_common(PatternRewriter &rewriter, Location loc,
                           RankedTensorType concatType, int dimension,
                           ArrayRef<Value> operands, SliceRequirement sliceReq,
                           mlir::sdy::TensorShardingPerValueAttr shard) {
  if (operands.size() < 3)
    return nullptr;

  stablehlo::SliceOp lhs = nullptr;
  bool hasSlice = false;
  if (auto lhsSlice = operands[0].getDefiningOp<stablehlo::SliceOp>()) {
    bool legal = true;
    for (int i = 0; i < lhsSlice.getType().getShape().size(); i++) {
      if (lhsSlice.getStartIndices()[i] != 0) {
        hasSlice = true;
      }
      if (lhsSlice.getStrides()[i] != 1) {
        legal = false;
        break;
      }
    }
    if (legal)
      lhs = lhsSlice;
  }

  if (!lhs)
    return nullptr;

  stablehlo::SliceOp rhs = nullptr;
  if (auto rhsSlice = operands.back().getDefiningOp<stablehlo::SliceOp>()) {
    bool legal = true;
    for (int i = 0; i < rhsSlice.getType().getShape().size(); i++) {
      if (rhsSlice.getStrides()[i] != 1) {
        legal = false;
        break;
      }
      if (i == dimension) {
        if (lhs.getStartIndices()[i] + concatType.getShape()[i] !=
            rhsSlice.getLimitIndices()[i]) {
          legal = false;
          break;
        }
      } else {
        if (lhs.getStartIndices()[i] != rhsSlice.getStartIndices()[i]) {
          legal = false;
          break;
        }
        if (lhs.getLimitIndices()[i] != rhsSlice.getLimitIndices()[i]) {
          legal = false;
          break;
        }
      }
      if (rhsSlice.getOperand().getType().getShape()[i] !=
          concatType.getShape()[i]) {
        hasSlice = true;
      }
    }
    if (legal)
      rhs = rhsSlice;
  }

  if (!rhs)
    return nullptr;
  if (sliceReq == SliceRequirement::NeedsSlice && !hasSlice)
    return nullptr;
  if (sliceReq == SliceRequirement::NeedsFull && hasSlice)
    return nullptr;

  if (rhs.getOperand() != lhs.getOperand())
    return nullptr;

  SmallVector<Value> newOps;
  int start = lhs ? 1 : 0;
  int end = operands.size() - (rhs ? 1 : 0);
  for (int i = start; i < end; i++) {
    newOps.push_back(operands[i]);
  }
  Value innerConcat = newOps[0];
  if (newOps.size() != 1) {
    auto nConcat =
        stablehlo::ConcatenateOp::create(rewriter, loc, newOps, dimension);
    innerConcat = nConcat;
    if (shard) {
      sdy::setShardings(nConcat, shard);
    }
  }

  auto iTy = RankedTensorType::get({}, rewriter.getI64Type());

  Value operand = lhs.getOperand();

  if (!shard) {
    if (auto opSharding = mlir::sdy::getSharding(operand)) {
      sdy::TensorShardingAttr opShardings[] = {opSharding};
      shard = sdy::TensorShardingPerValueAttr::get(operand.getContext(),
                                                   opShardings);
    }
  }

  if (hasSlice) {
    auto sloperand = stablehlo::SliceOp::create(
        rewriter, lhs.getLoc(), lhs.getOperand(), lhs.getStartIndices(),
        rhs.getLimitIndices(), lhs.getStrides());
    if (shard) {
      sdy::setShardings(sloperand, shard);
    }
    operand = sloperand;
  }

  SmallVector<Value> starts(
      concatType.getShape().size(),
      stablehlo::ConstantOp::create(rewriter, loc, iTy,
                                    cast<ElementsAttr>(makeAttr(iTy, 0))));

  if (lhs) {
    starts[dimension] = stablehlo::ConstantOp::create(
        rewriter, loc, iTy,
        cast<ElementsAttr>(makeAttr(iTy, lhs.getType().getShape()[dimension])));
  }

  auto dus = stablehlo::DynamicUpdateSliceOp::create(rewriter, loc, operand,
                                                     innerConcat, starts);
  if (shard) {
    sdy::setShardings(dus, shard);
  }
  return dus;
}

struct ConcatToOneDimDUSSlice final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                              ConcatToOneDimDUSSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp outer,
                                    PatternRewriter &rewriter) const {

    stablehlo::DynamicUpdateSliceOp replacement = concat_to_dus_slice_common(
        rewriter, outer.getLoc(), outer.getType(), outer.getDimension(),
        llvm::to_vector(outer.getOperands()), SliceRequirement::NeedsSlice,
        sdy::getShardingPerValue(outer));
    if (!replacement)
      return failure();
    rewriter.replaceOp(outer, replacement);

    return success();
  }
};

struct ConcatReshapeToOneDimDUS final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                              ConcatReshapeToOneDimDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp outer,
                                    PatternRewriter &rewriter) const {

    SmallVector<Value> pre_reshape;
    for (auto operand : outer.getOperands()) {
      auto re = operand.getDefiningOp<stablehlo::ReshapeOp>();
      if (!re)
        return failure();
      auto pre_shape =
          cast<RankedTensorType>(re.getOperand().getType()).getShape();
      if (re.getType().getShape().size() + 1 != pre_shape.size()) {
        return failure();
      }
      for (auto &&[lhs, rhs] :
           llvm::zip_equal(re.getType().getShape(), pre_shape.slice(1))) {
        if (lhs != rhs)
          return failure();
      }
      if (pre_shape[0] != 1) {
        return failure();
      }
      pre_reshape.push_back(re.getOperand());
    }
    SmallVector<int64_t> subShape = llvm::to_vector(outer.getType().getShape());
    subShape.insert(subShape.begin(), 1);
    RankedTensorType subType =
        RankedTensorType::get(subShape, outer.getType().getElementType());
    stablehlo::DynamicUpdateSliceOp replacement = concat_to_dus_slice_common(
        rewriter, outer.getLoc(), subType, outer.getDimension() + 1,
        pre_reshape, SliceRequirement::AnyOperand, nullptr);
    if (!replacement)
      return failure();

    auto shard = sdy::getShardingPerValue(outer);
    auto reshaped = rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(
        outer, outer.getType(), replacement);
    if (shard) {
      sdy::setShardings(reshaped, shard);
    }
    return success();
  }
};

// Optimization: DUSPad
// Pattern:
//   %padded_val = stablehlo.pad %original_data, %pad_val, low=[L...],
//   high=[H...], interior=[0...] %dus = stablehlo.dynamic_update_slice
//   %padded_val, %update_data, %idx...
// Constraint: The update region defined by %update_data and %idx falls entirely
//             within the %original_data region of %padded_val.
// Constraint: %padded_val has only %dus as its user.
// Constraint: Interior padding is zero.
// Rewrite to:
//   %idx_new = %idx - L // Adjust indices relative to original_data
//   %new_dus = stablehlo.dynamic_update_slice %original_data, %update_data,
//   %idx_new... %new_pad = stablehlo.pad %new_dus, %pad_val, low=[L...],
//   high=[H...], interior=[0...] replaceAllUsesWith(%dus, %new_pad)
struct DUSPad final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp, DUSPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dus,
                                    PatternRewriter &rewriter) const {

    // 1. Check if the operand being updated comes from a pad op.
    auto padOp = dus.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!padOp) {
      return rewriter.notifyMatchFailure(dus, "Operand is not a PadOp");
    }

    // 3. Constraint: Check for zero interior padding and non-negative edge
    // padding.
    if (llvm::any_of(padOp.getInteriorPadding(),
                     [](int64_t i) { return i != 0; })) {
      return rewriter.notifyMatchFailure(dus, "Requires zero interior padding");
    }
    if (llvm::any_of(padOp.getEdgePaddingLow(),
                     [](int64_t i) { return i < 0; }) ||
        llvm::any_of(padOp.getEdgePaddingHigh(),
                     [](int64_t i) { return i < 0; })) {
      return rewriter.notifyMatchFailure(dus,
                                         "Requires non-negative edge padding");
    }

    // 4. Get shapes and indices, require static shapes.
    auto dusOperandType = dyn_cast<RankedTensorType>(
        dus.getOperand().getType()); // = padOp result type
    auto updateType = dyn_cast<RankedTensorType>(dus.getUpdate().getType());
    auto originalDataType =
        dyn_cast<RankedTensorType>(padOp.getOperand().getType());

    if (!dusOperandType || !updateType || !originalDataType ||
        !dusOperandType.hasStaticShape() || !updateType.hasStaticShape() ||
        !originalDataType.hasStaticShape()) {
      return rewriter.notifyMatchFailure(
          dus, "Requires static shapes for involved tensors");
    }

    ArrayRef<int64_t> dusOperandShape = dusOperandType.getShape();
    ArrayRef<int64_t> updateShape = updateType.getShape();
    ArrayRef<int64_t> originalDataShape = originalDataType.getShape();
    SmallVector<Value> startIndices = dus.getStartIndices();
    ArrayRef<int64_t> lowPadding = padOp.getEdgePaddingLow();

    SmallVector<int64_t> newDusStartIndexValues;
    SmallVector<Value> newDusStartIndices;
    Location loc = dus.getLoc();
    auto indexElementType =
        cast<ShapedType>(startIndices[0].getType()).getElementType();
    auto indexScalarType = RankedTensorType::get({}, indexElementType);

    // 5. Check if update region is within original data bounds & calculate new
    // indices.
    for (int64_t d = 0; d < dusOperandShape.size(); ++d) {
      DenseIntElementsAttr startAttr;
      if (!matchPattern(startIndices[d], m_Constant(&startAttr)) ||
          startAttr.getNumElements() != 1) {
        return rewriter.notifyMatchFailure(
            dus, "Requires constant scalar start indices");
      }
      int64_t startVal = (*startAttr.begin()).getSExtValue();
      int64_t updateSize = updateShape[d];
      int64_t endVal = startVal + updateSize; // Exclusive end
      int64_t lowPad = lowPadding[d];
      int64_t originalSize = originalDataShape[d];

      // Check bounds: update [startVal, endVal) must be within original data
      // region [lowPad, lowPad + originalSize)
      if (startVal < lowPad || endVal > (lowPad + originalSize)) {
        return rewriter.notifyMatchFailure(
            dus, "Update region extends into padded area for dim " +
                     std::to_string(d));
      }

      // Calculate new start index relative to original data
      newDusStartIndexValues.push_back(startVal - lowPad);
    }

    Value valueToPad = nullptr;

    if (padOp.getOperand().getType() == dus.getUpdate().getType()) {
      valueToPad = dus.getUpdate();
    } else {
      // 2. Constraint: Check if the pad op has only the DUS as its user.
      if (!llvm::hasSingleElement(padOp->getUsers())) {
        return rewriter.notifyMatchFailure(dus, "PadOp has multiple users");
      }

      for (auto val : newDusStartIndexValues) {
        auto newStartConst = stablehlo::ConstantOp::create(
            rewriter, loc, indexScalarType,
            cast<ElementsAttr>(makeAttr(indexScalarType, val)));
        newDusStartIndices.push_back(newStartConst);
      }

      // 6. Perform the rewrite.
      // Create the new DUS on the original data.
      valueToPad = stablehlo::DynamicUpdateSliceOp::create(
          rewriter, loc,
          originalDataType,    // Result type matches original data
          padOp.getOperand(),  // Update the original data
          dus.getUpdate(),     // Use the same update value
          newDusStartIndices); // Use the adjusted indices
    }

    // Create the new Pad operation using the result of the new DUS.
    auto newPad = stablehlo::PadOp::create(
        rewriter, padOp.getLoc(),
        dus.getType(), // Result type matches original DUS result
        valueToPad, padOp.getPaddingValue(), padOp.getEdgePaddingLowAttr(),
        padOp.getEdgePaddingHighAttr(),
        padOp.getInteriorPaddingAttr()); // Assumed to be zeros by check above

    // Replace the original DUS with the new Pad operation.
    rewriter.replaceOp(dus, newPad.getResult());

    return success();
  }
};

struct DUSDUSConcat final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp, DUSDUSConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dus,
                                    PatternRewriter &rewriter) const {
    auto dus2 =
        dus.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();

    if (!dus2)
      return failure();

    RankedTensorType tys[2];
    stablehlo::DynamicUpdateSliceOp duses[2] = {dus, dus2};
    for (auto en : llvm::enumerate(duses)) {
      auto ty = dyn_cast<RankedTensorType>(en.value().getUpdate().getType());
      if (!ty)
        return failure();
      tys[en.index()] = ty;
    }

    if (dus.getOperand().getType() != dus2.getOperand().getType())
      return failure();

    ssize_t diffidx = -1;
    bool allStatic = true;
    SmallVector<int64_t> starts;
    for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
      DenseIntElementsAttr startattr;
      if (!matchPattern(dus.getStartIndices()[i], m_Constant(&startattr))) {
        allStatic = false;
      } else {
        int64_t ival = (*startattr.begin()).getSExtValue();
        starts.push_back(ival);
      }

      if (dus.getStartIndices()[i] == dus2.getStartIndices()[i])
        continue;
      if (diffidx != -1) {
        return failure();
      }
      diffidx = i;
    }

    if (diffidx == -1) {
      for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
        if (tys[0].getShape()[i] == tys[1].getShape()[i])
          continue;
        if (diffidx != -1) {
          return failure();
        }
        diffidx = i;
      }
    }

    if (diffidx == -1) {
      return failure();
    }

    // Sizes must be the same except for the differing index
    for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
      if (i == diffidx)
        continue;
      if (tys[0].getShape()[i] != tys[1].getShape()[i])
        return failure();
    }

    int64_t idxs[2];

    for (auto en : llvm::enumerate(duses)) {
      auto val = en.value().getStartIndices()[diffidx];
      DenseIntElementsAttr startattr;
      if (!matchPattern(val, m_Constant(&startattr))) {
        return failure();
      }
      int64_t ival = (*startattr.begin()).getSExtValue();
      idxs[en.index()] = ival;
    }

    // We have dus(dus2(operand, update2, start=idxs[1]), update1,
    // start=idxs[0])
    //
    // only one index differs, which may differ in start, size, or both
    //
    if (idxs[1] == idxs[0] + tys[0].getShape()[diffidx]) {
      // Case 1:     [idxs[0]   ... update1 ...  idxs[0] + tys[0].size][idxs[1]
      // .... update2 ...   ]
      Value operands[2] = {dus.getUpdate(), dus2.getUpdate()};
      auto concat = stablehlo::ConcatenateOp::create(rewriter, dus.getLoc(),
                                                     operands, diffidx);
      rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
          dus, dus2.getOperand(), concat, dus.getStartIndices());
      return success();
    } else if (idxs[0] == idxs[1] + tys[1].getShape()[diffidx]) {
      // Case 2:     [idxs[1]   ... update2 ...  idxs[1] + tys[1].size][idxs[0]
      // .... update1 ...   ]
      Value operands[2] = {dus2.getUpdate(), dus.getUpdate()};
      auto concat = stablehlo::ConcatenateOp::create(rewriter, dus.getLoc(),
                                                     operands, diffidx);
      rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
          dus, dus2.getOperand(), concat, dus2.getStartIndices());
      return success();
    } else if (idxs[1] >= idxs[0] && idxs[1] + tys[1].getShape()[diffidx] <=
                                         idxs[0] + tys[0].getShape()[diffidx]) {

      // Case 3:     [idxs[1]   ... update2 ...  idxs[1] + tys[1].size]
      //          [idxs[0]     .... update1 ...              idxs[0] +
      //          tys[0].size  ]

      // the previous update (in dus1) was completely overwritten [e.g. dus0
      // starts before and end later]
      rewriter.modifyOpInPlace(
          dus, [&]() { dus.getOperandMutable().set(dus2.getOperand()); });
      return success();
    } else if (idxs[0] >= idxs[1] && idxs[0] + tys[0].getShape()[diffidx] <=
                                         idxs[1] + tys[1].getShape()[diffidx]) {
      // the new update is entirely within the space of the old update

      // Case 4:
      //         [idxs[1]      ... update2 ...         idxs[1] + tys[1].size]
      //              [idxs[0] ... update1 ...    idxs[0] + tys[0].size  ]

      auto itype = dus.getStartIndices()[diffidx].getType();
      auto c0 =
          stablehlo::ConstantOp::create(rewriter, dus.getLoc(), itype,
                                        cast<ElementsAttr>(makeAttr(itype, 0)));
      auto cidx = stablehlo::ConstantOp::create(
          rewriter, dus.getLoc(), itype,
          cast<ElementsAttr>(makeAttr(itype, idxs[0] - idxs[1])));

      SmallVector<Value> idxs(dus.getStartIndices().size());
      for (size_t i = 0; i < dus.getStartIndices().size(); i++)
        idxs[i] = c0;
      idxs[diffidx] = cidx;

      auto within_dus = stablehlo::DynamicUpdateSliceOp::create(
          rewriter, dus2.getLoc(), dus2.getUpdate(), dus.getUpdate(), idxs);
      rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
          dus, dus2.getOperand(), within_dus, dus2.getStartIndices());
      return success();
    } else if (idxs[0] >= idxs[1] &&
               idxs[0] < idxs[1] + tys[1].getShape()[diffidx] &&
               idxs[0] + tys[0].getShape()[diffidx] >
                   idxs[1] + tys[1].getShape()[diffidx] &&
               allStatic &&
               idxs[0] + tys[0].getShape()[diffidx] <=
                   dus.getOperand().getType().getShape()[diffidx]) {
      // the new update overlaps, following the old update

      // Case 5:
      //         [idxs[1]      ... update2 ...    idxs[1] + tys[1].size]
      //              [idxs[0] ... update1 ...         idxs[0] + tys[0].size  ]

      SmallVector<int64_t> limits =
          llvm::to_vector(dus2.getUpdate().getType().getShape());
      limits[diffidx] = idxs[0] - idxs[1];
      SmallVector<int64_t> zeros(starts.size(), 0);
      SmallVector<int64_t> ones(starts.size(), 1);
      auto begin = stablehlo::SliceOp::create(
          rewriter, dus.getLoc(), dus2.getUpdate(), zeros, limits, ones);
      Value operands[2] = {begin, dus.getUpdate()};
      auto concat = stablehlo::ConcatenateOp::create(rewriter, dus.getLoc(),
                                                     operands, diffidx);

      rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
          dus, dus2.getOperand(), concat, dus2.getStartIndices());
      return success();
    } else if (idxs[0] < idxs[1] &&
               idxs[1] < idxs[0] + tys[0].getShape()[diffidx] &&
               idxs[0] + tys[0].getShape()[diffidx] <
                   idxs[1] + tys[1].getShape()[diffidx] &&
               allStatic &&
               idxs[1] + tys[1].getShape()[diffidx] <=
                   dus.getOperand().getType().getShape()[diffidx]) {
      // the new update overlaps, following the old update

      // Case 5:
      //             [idxs[1]      ... update2 ...    idxs[1] + tys[1].size]
      //     [idxs[0] ... update1 ...         idxs[0] + tys[0].size  ]

      SmallVector<int64_t> limits =
          llvm::to_vector(dus2.getUpdate().getType().getShape());
      limits[diffidx] = idxs[1] + tys[1].getShape()[diffidx] -
                        (idxs[0] + tys[0].getShape()[diffidx]);
      SmallVector<int64_t> zeros(starts.size(), 0);
      SmallVector<int64_t> ones(starts.size(), 1);
      auto end = stablehlo::SliceOp::create(
          rewriter, dus.getLoc(), dus2.getUpdate(), zeros, limits, ones);
      Value operands[2] = {dus.getUpdate(), end};
      auto concat = stablehlo::ConcatenateOp::create(rewriter, dus.getLoc(),
                                                     operands, diffidx);

      rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
          dus, dus2.getOperand(), concat, dus.getStartIndices());
      return success();
    }

    return failure();
  }
};

struct DynamicUpdateToConcat final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                              DynamicUpdateToConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    SmallVector<size_t> mismatches;
    size_t idx = 0;
    for (auto &&[start, update_size, res_size] :
         llvm::zip(op.getStartIndices(), op.getUpdate().getType().getShape(),
                   op.getType().getShape())) {
      DenseIntElementsAttr startattr;
      if (!matchPattern(start, m_Constant(&startattr))) {
        return failure();
      }
      int64_t startv = (*startattr.begin()).getSExtValue();
      if (startv < 0)
        return failure();

      if (startv + update_size > res_size)
        return failure();

      if (startv == 0 && update_size == res_size) {
        idx++;
        continue;
      }
      mismatches.push_back(idx);
      idx++;
    }

    if (mismatches.size() != 1)
      return failure();
    auto dim = mismatches[0];

    DenseIntElementsAttr startattr;
    if (!matchPattern(op.getStartIndices()[dim], m_Constant(&startattr))) {
      return failure();
    }
    int64_t startv = (*startattr.begin()).getSExtValue();

    SmallVector<Value> toConcat;

    if (startv != 0) {
      SmallVector<int64_t> starts(op.getType().getShape().size(), 0);
      SmallVector<int64_t> ends(op.getType().getShape().begin(),
                                op.getType().getShape().end());
      SmallVector<int64_t> steps(op.getType().getShape().size(), 1);
      ends[dim] = startv;
      toConcat.push_back(stablehlo::SliceOp::create(
          rewriter, op.getLoc(), op.getOperand(), starts, ends, steps));
    }
    toConcat.push_back(op.getUpdate());
    auto update_size = op.getUpdate().getType().getShape()[dim];
    auto res_size = op.getType().getShape()[dim];
    if (startv + update_size != res_size) {
      SmallVector<int64_t> starts(op.getType().getShape().size(), 0);
      SmallVector<int64_t> ends(op.getType().getShape().begin(),
                                op.getType().getShape().end());
      SmallVector<int64_t> steps(op.getType().getShape().size(), 1);
      starts[dim] = startv + update_size;
      toConcat.push_back(stablehlo::SliceOp::create(
          rewriter, op.getLoc(), op.getOperand(), starts, ends, steps));
    }

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, op.getType(),
                                                          toConcat, dim);
    return success();
  }
};

struct SliceOfDynamicUpdate final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceOfDynamicUpdate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto dyn = op.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
    if (!dyn)
      return failure();

    // Try to use the updated value
    {
      SmallVector<int64_t> start;
      SmallVector<int64_t> end;
      SmallVector<int64_t> step;

      bool legal = true;
      for (auto &&[nstart, nend, nstep, update_start, update_size] : llvm::zip(
               op.getStartIndices(), op.getLimitIndices(), op.getStrides(),
               dyn.getStartIndices(), dyn.getUpdate().getType().getShape())) {
        DenseIntElementsAttr startattr;
        if (!matchPattern(update_start, m_Constant(&startattr))) {
          legal = false;
          break;
        }
        int64_t startv = (*startattr.begin()).getSExtValue();
        if (startv < 0) {
          legal = false;
          break;
        }

        // see if the slice is exclusively inside the update.

        // slice starts below insertion
        if (nstart < startv) {
          legal = false;
          break;
        }

        // slice ends after insertion
        if (nend > startv + update_size) {
          legal = false;
          break;
        }

        start.push_back(nstart - startv);
        end.push_back(nend - startv);
        step.push_back(nstep);
      }

      if (legal) {
        rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(op, dyn.getUpdate(),
                                                        start, end, step);
        return success();
      }
    }

    // Try proving that there can be no overlap
    {
      bool no_overlap = false;

      for (auto &&[nstart, nend, nstep, update_start, update_size] : llvm::zip(
               op.getStartIndices(), op.getLimitIndices(), op.getStrides(),
               dyn.getStartIndices(), dyn.getUpdate().getType().getShape())) {
        DenseIntElementsAttr startattr;
        if (!matchPattern(update_start, m_Constant(&startattr))) {
          continue;
        }

        int64_t startv = (*startattr.begin()).getSExtValue();
        // slice ends below insertion
        if (nend <= startv) {
          no_overlap = true;
          break;
        }

        // slice starts after insertion end
        if (nstart >= startv + update_size) {
          no_overlap = true;
          break;
        }
      }

      if (no_overlap) {
        rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
            op, dyn.getOperand(), op.getStartIndices(), op.getLimitIndices(),
            op.getStrides());
        return success();
      }
    }

    return failure();
  }
};

struct SliceOfUpdateWithoutCorners final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceOfUpdateWithoutCorners> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {

    auto dyn =
        op.getOperand().getDefiningOp<enzymexla::UpdateWithoutCornersOp>();
    if (!dyn)
      return failure();

    // Try to use the updated value, currently extremely conservative and checks
    // everything is within the innermost box This can be extended to just check
    // that nothing overlaps with the corners.
    if (op.getStrides()[dyn.getDimensionX()] == 1 &&
        op.getStrides()[dyn.getDimensionY()] == 1 &&
        op.getStartIndices()[dyn.getDimensionX()] >= dyn.getX1() &&
        op.getStartIndices()[dyn.getDimensionY()] >= dyn.getY1() &&
        op.getLimitIndices()[dyn.getDimensionX()] <= dyn.getX2() &&
        op.getLimitIndices()[dyn.getDimensionY()] <= dyn.getY2()) {
      rewriter.modifyOpInPlace(
          op, [&]() { op.getOperandMutable().assign(dyn.getUpdate()); });
      return success();
    }

    // Try to use the old value, currently extremely conservative and checks
    // everything is within the innermost box This can be extended to just check
    // that nothing overlaps with the corners.
    if (op.getStrides()[dyn.getDimensionX()] == 1 &&
        op.getStrides()[dyn.getDimensionY()] == 1 &&
        (op.getStartIndices()[dyn.getDimensionX()] >= dyn.getX2() ||
         op.getLimitIndices()[dyn.getDimensionX()] <= dyn.getX1()) &&
        (op.getStartIndices()[dyn.getDimensionY()] >= dyn.getY2() ||
         op.getLimitIndices()[dyn.getDimensionY()] <= dyn.getY1())) {
      rewriter.modifyOpInPlace(
          op, [&]() { op.getOperandMutable().assign(dyn.getOperand()); });
      return success();
    }
    return failure();
  }
};

struct SliceDUSToConcat final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceDUSToConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp sliceOp,
                                    PatternRewriter &rewriter) const {

    auto dusOp =
        sliceOp.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
    if (!dusOp) {
      return rewriter.notifyMatchFailure(sliceOp, "Operand is not a DUS");
    }

    // Constraint: DUS has only the slice as user
    if (!dusOp.getResult().hasOneUse()) {
      return rewriter.notifyMatchFailure(sliceOp, "DUS has multiple users");
    }

    // Get DUS info
    Value targetOperand = dusOp.getOperand();
    Value updateVal = dusOp.getUpdate();
    SmallVector<Value> dusIndexVals = dusOp.getStartIndices();

    auto targetType = dyn_cast<RankedTensorType>(targetOperand.getType());
    auto updateType = dyn_cast<RankedTensorType>(updateVal.getType());
    auto sliceResultType =
        dyn_cast<RankedTensorType>(sliceOp.getType()); // Type of %283

    if (!targetType || !updateType || !sliceResultType ||
        !targetType.hasStaticShape() || !updateType.hasStaticShape() ||
        !sliceResultType.hasStaticShape()) {
      return rewriter.notifyMatchFailure(sliceOp, "Requires static shapes");
    }

    ArrayRef<int64_t> updateShape = updateType.getShape();
    int rank = targetType.getRank();

    // Constraint: DUS indices must be constant
    SmallVector<int64_t> dusStartIndices;
    dusStartIndices.reserve(rank);
    for (Value idxVal : dusIndexVals) {
      DenseIntElementsAttr idxAttr;
      if (!matchPattern(idxVal, m_Constant(&idxAttr)) ||
          idxAttr.getNumElements() != 1) {
        return rewriter.notifyMatchFailure(
            sliceOp, "DUS indices must be constant scalars");
      }
      dusStartIndices.push_back((*idxAttr.begin()).getSExtValue());
    }

    if (llvm::any_of(sliceOp.getStrides(), [](int64_t s) { return s != 1; })) {
      return rewriter.notifyMatchFailure(sliceOp,
                                         "Requires slice strides of 1");
    }

    int concatDim = -1; // Dimension along which to concatenate results

    SmallVector<int64_t> dusLimitIndices;
    for (int d = 0; d < rank; ++d) {
      int64_t dusStart = dusStartIndices[d];
      int64_t dusEnd = dusStart + updateShape[d];
      dusLimitIndices.push_back(dusEnd);
      int64_t sliceStart = sliceOp.getStartIndices()[d];
      int64_t sliceLimit = sliceOp.getLimitIndices()[d];

      if (dusStart == sliceStart && dusEnd == sliceLimit) {
        continue;
      }

      if (concatDim != -1) {
        return failure();
      }
      concatDim = d;

      // if entirely within one or the other, use the other pattern
      // 1) if we are outside the update
      if (sliceLimit <= dusStart || sliceStart >= dusEnd) {
        return failure();
      }

      // 2) we are fully within the update
      if (sliceStart >= dusStart & sliceStart <= dusEnd) {
        return failure();
      }
    }

    if (concatDim == -1)
      return failure();

    SmallVector<Value> toConcat;

    if (sliceOp.getStartIndices()[concatDim] < dusStartIndices[concatDim]) {
      assert(sliceOp.getLimitIndices()[concatDim] > dusStartIndices[concatDim]);
      SmallVector<int64_t> newLimit =
          llvm::to_vector(sliceOp.getLimitIndices());
      newLimit[concatDim] = dusStartIndices[concatDim];
      toConcat.push_back(stablehlo::SliceOp::create(
          rewriter, sliceOp.getLoc(), dusOp.getOperand(),
          sliceOp.getStartIndices(), newLimit, sliceOp.getStrides()));
    }

    if (sliceOp.getLimitIndices()[concatDim] >= dusLimitIndices[concatDim] &&
        sliceOp.getStartIndices()[concatDim] <= dusStartIndices[concatDim]) {
      toConcat.push_back(dusOp.getUpdate());
    } else if (sliceOp.getLimitIndices()[concatDim] <=
               dusLimitIndices[concatDim]) {
      assert(sliceOp.getStartIndices()[concatDim] < dusStartIndices[concatDim]);
      SmallVector<int64_t> newStart(dusStartIndices);
      SmallVector<int64_t> newLimit =
          llvm::to_vector(sliceOp.getLimitIndices());
      for (int i = 0; i < rank; i++) {
        newStart[i] -= dusStartIndices[i];
        newLimit[i] -= dusStartIndices[i];
      }
      toConcat.push_back(stablehlo::SliceOp::create(
          rewriter, sliceOp.getLoc(), dusOp.getUpdate(), newStart, newLimit,
          sliceOp.getStrides()));
    } else if (sliceOp.getStartIndices()[concatDim] >=
               dusStartIndices[concatDim]) {
      assert(sliceOp.getLimitIndices()[concatDim] > dusLimitIndices[concatDim]);
      SmallVector<int64_t> newStart =
          llvm::to_vector(sliceOp.getStartIndices());
      SmallVector<int64_t> newLimit(dusLimitIndices);
      for (int i = 0; i < rank; i++) {
        newStart[i] -= dusStartIndices[i];
        newLimit[i] -= dusStartIndices[i];
      }
      toConcat.push_back(stablehlo::SliceOp::create(
          rewriter, sliceOp.getLoc(), dusOp.getUpdate(), newStart, newLimit,
          sliceOp.getStrides()));
    }

    if (sliceOp.getLimitIndices()[concatDim] > dusLimitIndices[concatDim]) {
      assert(sliceOp.getStartIndices()[concatDim] < dusLimitIndices[concatDim]);
      SmallVector<int64_t> newStart =
          llvm::to_vector(sliceOp.getStartIndices());
      newStart[concatDim] = dusLimitIndices[concatDim];
      toConcat.push_back(stablehlo::SliceOp::create(
          rewriter, sliceOp.getLoc(), dusOp.getOperand(), newStart,
          sliceOp.getLimitIndices(), sliceOp.getStrides()));
    }

    auto preType = sliceOp.getType();

    auto newConcat = rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
        sliceOp, toConcat, concatDim);
    assert(preType == newConcat.getType());
    (void)preType;
    (void)newConcat;

    return success();
  }
};

static bool definedOutside(Value v, Operation *op) {
  return !op->isAncestor(v.getParentBlock()->getParentOp());
}

// Replace while op iteration variables which are not updated with their
// upcoming value
template <typename T> struct LICM : public CheckedOpRewritePattern<T, LICM<T>> {
  using CheckedOpRewritePattern<T, LICM<T>>::CheckedOpRewritePattern;

  bool single_user;
  LICM(bool single_user, MLIRContext *context, PatternBenefit benefit = 1,
       ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern<T, LICM<T>>(context, benefit, generatedNames),
        single_user(single_user) {}

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    auto whileOp = op->template getParentOfType<stablehlo::WhileOp>();
    if (!whileOp)
      return failure();
    for (auto operand : op->getOperands()) {
      if (!definedOutside(operand, whileOp))
        return failure();
      if (single_user) {
        for (auto U : operand.getUsers()) {
          if (U == op)
            continue;
          if (whileOp->isAncestor(U))
            return failure();
        }
      }
    }
    rewriter.modifyOpInPlace(op, [&]() { op->moveBefore(whileOp); });
    return success();
  }
};

struct LICMElementwise
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          LICMElementwise> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise, LICMElementwise>::CheckedOpTraitRewritePattern;

  bool single_user;
  LICMElementwise(bool single_user, MLIRContext *context,
                  PatternBenefit benefit = 1)
      : CheckedOpTraitRewritePattern<OpTrait::Elementwise, LICMElementwise>(
            context, benefit),
        single_user(single_user) {}

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (!isa<stablehlo::WhileOp>(op->getParentOp()))
      return failure();
    for (auto operand : op->getOperands()) {
      if (!definedOutside(operand, op->getParentOp()))
        return failure();
      if (single_user) {
        for (auto U : operand.getUsers()) {
          if (U == op)
            continue;
          if (op->getParentOp()->isAncestor(U))
            return failure();
        }
      }
    }
    rewriter.modifyOpInPlace(op, [&]() { op->moveBefore(op->getParentOp()); });
    return success();
  }
};

static bool isIotaRange(ArrayRef<int64_t> dims) {
  return llvm::all_of(llvm::enumerate(dims), [](const auto &it) {
    return static_cast<int64_t>(it.index()) == it.value();
  });
}

/// Matches when either of the submatchers match.
template <typename MatcherA, typename MatcherB> struct m_AnyOf {
  m_AnyOf(MatcherA a, MatcherB b) : matcherA(a), matcherB(b) {}

  bool match(Operation *op) { return matcherA.match(op) || matcherB.match(op); }

  MatcherA matcherA;
  MatcherB matcherB;
};

template <typename MatcherA, typename MatcherB>
m_AnyOf(MatcherA, MatcherB) -> m_AnyOf<MatcherA, MatcherB>;

/// Binary constant folder that used a generic folder function to handle both
/// ints and floats.
template <typename Fn>
static TypedAttr foldBinaryOpIntOrFloat(TypedAttr lhs, TypedAttr rhs,
                                        Fn &&folder) {
  Attribute operands[2] = {lhs, rhs};
  Type elemTy = getElementTypeOrSelf(lhs);

  Attribute res;
  if (isa<IntegerType>(elemTy))
    res = constFoldBinaryOp<IntegerAttr, IntegerAttr::ValueType, void>(operands,
                                                                       folder);
  if (isa<FloatType>(elemTy))
    res = constFoldBinaryOp<FloatAttr, FloatAttr::ValueType, void>(operands,
                                                                   folder);
  if (res)
    return cast<TypedAttr>(res);

  return nullptr;
}

static bool isTransposeReshapeLikeBroadcast(stablehlo::BroadcastInDimOp op) {
  TypedValue<RankedTensorType> operand = op.getOperand();
  RankedTensorType operandTy = operand.getType();
  RankedTensorType type = op.getType();

  // Fold when broadcast is a noop.
  auto dims = op.getBroadcastDimensions();
  bool isDimsIota = isIotaRange(dims);
  if (type == operandTy && isDimsIota) {
    return false;
  }

  // Handle splat broadcasts.
  if (SplatElementsAttr cstAttr; matchPattern(operand, m_Constant(&cstAttr))) {
    return false;
  }

  if (operandTy.hasStaticShape() && type.hasStaticShape() &&
      type.getNumElements() == operandTy.getNumElements()) {
    // BroadcastInDim equivalent to reshape.
    if (isDimsIota) {
      return false;
    }
    // BroadcastInDim equivalent to transpose.
    if (type.getRank() == operandTy.getRank()) {
      return false;
    }

    return true;
  }
  return false;
}

// slice(broadcast x) -> broadcast(slice x)
struct SliceBroadcast final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceBroadcast> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto bcast = op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcast)
      return failure();

    if (isTransposeReshapeLikeBroadcast(bcast))
      return failure();

    SmallVector<int64_t> nbcast_idx;

    auto preShape = cast<RankedTensorType>(bcast.getOperand().getType());
    SmallVector<int64_t> in_start(preShape.getShape().size(), 0);
    SmallVector<int64_t> in_end(preShape.getShape().begin(),
                                preShape.getShape().end());
    SmallVector<int64_t> in_stride(preShape.getShape().size(), 1);

    bool innerSlice = false;

    size_t outidx = 0;
    for (auto &&[start, end, step, indim, outdim] :
         llvm::zip(op.getStartIndices(), op.getLimitIndices(), op.getStrides(),
                   bcast.getType().getShape(), op.getType().getShape())) {
      ssize_t idx = -1;
      for (auto en : llvm::enumerate(bcast.getBroadcastDimensions())) {
        if (en.value() == outidx) {
          idx = en.index();
          break;
        }
      }

      nbcast_idx.push_back(outdim);
      if (idx == -1) {
        // being broadcast just resize the outshape
      } else {
        auto preShapeIdx = preShape.getShape()[idx];

        // slice the inner shape
        if (preShapeIdx == indim) {
          in_start[idx] = start;
          in_end[idx] = end;
          in_stride[idx] = step;
          if (start != 0 || step != 1 || end != preShape.getShape()[idx])
            innerSlice = true;
        } else if (preShapeIdx != 1) {
          return failure();
        }
      }

      outidx++;
    }

    if (innerSlice && !llvm::hasSingleElement(bcast->getUsers()))
      return failure();

    Value tobcast = bcast.getOperand();
    if (innerSlice)
      tobcast = stablehlo::SliceOp::create(rewriter, op.getLoc(), tobcast,
                                           in_start, in_end, in_stride);

    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), tobcast, bcast.getBroadcastDimensions());
    return success();
  }
};

// (dyn)slice(transpose x) -> transpose((dyn)slice x)
template <typename OpTy>
struct SliceTransposeBase final
    : CheckedOpRewritePattern<OpTy, SliceTransposeBase<OpTy>> {
  using CheckedOpRewritePattern<
      OpTy, SliceTransposeBase<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto transpose =
        op.getOperand().template getDefiningOp<stablehlo::TransposeOp>();
    if (!transpose || !llvm::hasSingleElement(transpose->getUsers())) {
      return failure();
    }

    if (transpose->getBlock() != op->getBlock()) {
      return failure();
    }

    auto sliceOp = sliceTransposeHelper(transpose, rewriter, op);
    rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(
        op, sliceOp, transpose.getPermutation());
    return success();
  }
};

// 1. slice(reduce_window x, last_idx) -> reduce x
// 2. slice(reduce_window x) -> reduce_window(slice x)
struct SliceReduceWindow
    : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceReduceWindow> {
  using CheckedOpRewritePattern<stablehlo::SliceOp,
                                SliceReduceWindow>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reduceWindow =
        op.getOperand().getDefiningOp<stablehlo::ReduceWindowOp>();
    if (!reduceWindow)
      return failure();

    if (!llvm::hasSingleElement(reduceWindow->getUsers()))
      return failure();

    if (reduceWindow.getInputs().size() != 1 ||
        reduceWindow.getInitValues().size() != 1)
      return failure();

    // Check window parameters indicate full reduction along one dimension
    auto windowDims = reduceWindow.getWindowDimensions();
    auto windowStrides = reduceWindow.getWindowStrides();
    auto windowDilations = reduceWindow.getWindowDilations();
    auto baseDilations = reduceWindow.getBaseDilations();

    // Check if the window strides are all 1 or unspecified
    bool stridesAllOne =
        !windowStrides || llvm::all_of(*windowStrides, [](int64_t stride) {
          return stride == 1;
        });

    // Check if the window dilations are all 1 or unspecified
    bool windowDilationsAllOne =
        !windowDilations ||
        llvm::all_of(*windowDilations,
                     [](int64_t dilation) { return dilation == 1; });

    // Check if the base dilations are all 1 or unspecified
    bool baseDilationsAllOne =
        !baseDilations || llvm::all_of(*baseDilations, [](int64_t dilation) {
          return dilation == 1;
        });

    auto padding = reduceWindow.getPadding();
    bool zeroPadding =
        !padding || llvm::all_of(padding->getValues<int64_t>(),
                                 [](int64_t val) { return val == 0; });

    auto sliceStarts = op.getStartIndices();
    auto sliceLimits = op.getLimitIndices();
    auto sliceStrides = op.getStrides();

    bool sliceStridesAllOne =
        sliceStrides.empty() ||
        llvm::all_of(sliceStrides, [](int64_t stride) { return stride == 1; });

    // Case 2: slice(reduce_window x) -> reduce_window(slice x)
    if (stridesAllOne && windowDilationsAllOne && baseDilationsAllOne &&
        zeroPadding && sliceStridesAllOne) {
      int64_t rank = sliceStarts.size();
      SmallVector<int64_t> newInputStarts(rank);
      SmallVector<int64_t> newInputLimits(rank);
      SmallVector<int64_t> newSliceStrides(rank, 1);

      for (int64_t i = 0; i < rank; ++i) {
        newInputStarts[i] = sliceStarts[i];
        newInputLimits[i] = sliceLimits[i] + windowDims[i] - 1;
      }

      auto newSlice = stablehlo::SliceOp::create(
          rewriter, op.getLoc(), reduceWindow.getInputs()[0], newInputStarts,
          newInputLimits, newSliceStrides);

      auto newRw = rewriter.replaceOpWithNewOp<stablehlo::ReduceWindowOp>(
          op, op.getType(), ValueRange{newSlice.getResult()},
          reduceWindow.getInitValues(), reduceWindow.getWindowDimensionsAttr(),
          reduceWindow.getWindowStridesAttr(),
          reduceWindow.getBaseDilationsAttr(),
          reduceWindow.getWindowDilationsAttr(), reduceWindow.getPaddingAttr());

      newRw.getBody().takeBody(reduceWindow.getBody());
      return success();
    }

    if (!padding)
      return failure();
    auto paddingVals = padding->getValues<int64_t>();

    if (!stridesAllOne || !windowDilationsAllOne || !baseDilationsAllOne)
      return failure();

    // Find which dimension has window size > 1 (the reduction dimension)
    int64_t reductionDim = -1;
    auto inputType =
        dyn_cast<ShapedType>(reduceWindow.getInputs()[0].getType());
    if (!inputType || !inputType.hasStaticShape())
      return failure();
    ArrayRef<int64_t> inputShape = inputType.getShape();
    for (unsigned i = 0; i < windowDims.size(); ++i) {
      if (windowDims[i] == inputShape[i]) {
        if (reductionDim != -1)
          return failure();
        reductionDim = i;
      }
    }
    if (reductionDim == -1)
      return failure();

    // Check that padding covers exactly (window_size-1) elements before in
    // reduction dim
    // FIXME: do other padding values have to be 0?
    for (unsigned i = 0; i < windowDims.size(); ++i) {
      if (i == reductionDim) {
        if (paddingVals[2 * i] != windowDims[i] - 1)
          return failure();
      }
    }

    // Check this is a slice taking the last element in reduction dim
    for (int64_t i = 0; i < (int64_t)sliceStarts.size(); ++i) {
      if (!sliceStrides.empty() && sliceStrides[i] != 1)
        return failure();

      if (i == reductionDim) {
        if (sliceStarts[i] != windowDims[i] - 1 ||
            sliceLimits[i] != windowDims[i])
          return failure();
      } else {
        if (sliceStarts[i] != 0 || sliceLimits[i] != inputShape[i])
          return failure();
      }
    }

    // Replace with direct reduce
    auto input = reduceWindow.getInputs()[0];
    auto initVal = reduceWindow.getInitValues()[0];
    // Compute the result type for the new ReduceOp by removing the reduction
    // dimension
    SmallVector<int64_t> resultShape;
    for (size_t i = 0; i < inputShape.size(); ++i) {
      if (static_cast<int64_t>(i) != reductionDim) {
        resultShape.push_back(inputShape[i]);
      }
    }
    Type resultType =
        RankedTensorType::get(resultShape, inputType.getElementType());

    rewriter.setInsertionPoint(reduceWindow);

    auto reduceOp = stablehlo::ReduceOp::create(
        rewriter, reduceWindow.getLoc(), TypeRange(resultType), input, initVal,
        rewriter.getDenseI64ArrayAttr({reductionDim}));

    // Clone the reduction body
    rewriter.inlineRegionBefore(reduceWindow.getBody(), reduceOp.getBody(),
                                reduceOp.getBody().end());

    // Create a reshape to match the slice output shape
    Value result = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), reduceOp.getResult(0),
        cast<RankedTensorType>(op.getResult().getType()).getShape());

    // Replace the slice with the reduce result
    rewriter.replaceOp(op, result);
    rewriter.eraseOp(reduceWindow);

    return success();
  }
};

// transpose(dynamic_slice x) -> dynamic_slice(transpose x)
// transpose(slice x) -> slice(transpose x)
template <typename OpTy>
struct TransposeSliceBase final
    : CheckedOpRewritePattern<stablehlo::TransposeOp,
                              TransposeSliceBase<OpTy>> {
  using CheckedOpRewritePattern<
      stablehlo::TransposeOp,
      TransposeSliceBase<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto sliceOp = op.getOperand().template getDefiningOp<OpTy>();
    if (!sliceOp) {
      return failure();
    }

    auto resShape = cast<RankedTensorType>(op.getType()).getShape();
    if (OpIsReshapeLike(op, resShape)) {
      return failure(); // transpose_is_reshape has higher priority
    }

    // If we can fuse the transpose into all of its users then we shouldn't push
    // it up (or atleast give higher priority to other passes before trying to
    // move it up)
    if (llvm::all_of(op->getUsers(),
                     [&](auto user) { return isFusible(op, user); })) {
      return failure();
    }

    bool singleUser = sliceOp->getResult(0).hasOneUse();

    auto newTranspose = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), sliceOp.getOperand(), op.getPermutation());
    auto newSlice = transposeSliceHelper(newTranspose, rewriter, sliceOp);
    rewriter.replaceOp(op, newSlice);
    if (singleUser) {
      rewriter.eraseOp(sliceOp);
    }
    return success();
  }
};

// transpose(dynamic_slice x) -> dynamic_slice(transpose x)
// transpose(slice x) -> slice(transpose x)
template <typename OpTy>
struct TransposeLikeBroadcastSliceBase final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                              TransposeLikeBroadcastSliceBase<OpTy>> {
  using CheckedOpRewritePattern<
      stablehlo::BroadcastInDimOp,
      TransposeLikeBroadcastSliceBase<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    if (!isTransposeReshapeLikeBroadcast(op))
      return failure();
    auto sliceOp = op.getOperand().template getDefiningOp<OpTy>();
    if (!sliceOp) {
      return failure();
    }

    // If we can fuse the transpose into all of its users then we shouldn't push
    // it up (or atleast give higher priority to other passes before trying to
    // move it up)
    if (llvm::all_of(op->getUsers(),
                     [&](auto user) { return isFusible(op, user); })) {
      return failure();
    }

    bool singleUser = sliceOp->getResult(0).hasOneUse();

    SmallVector<int64_t> shape = llvm::to_vector(op.getType().getShape());
    for (auto &&[i, v] : llvm::enumerate(op.getBroadcastDimensions())) {
      shape[v] = sliceOp.getOperand().getType().getShape()[i];
    }
    auto newTranspose = stablehlo::BroadcastInDimOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(shape, op.getType().getElementType()),
        sliceOp.getOperand(), op.getBroadcastDimensions());
    auto newSlice = transposeLikeSliceHelper(newTranspose, rewriter, sliceOp);
    rewriter.replaceOp(op, newSlice);
    if (singleUser) {
      rewriter.eraseOp(sliceOp);
    }
    return success();
  }
};

struct TransposeAllUsersSlice final
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeAllUsersSlice> {
  using CheckedOpRewritePattern<
      stablehlo::TransposeOp, TransposeAllUsersSlice>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    SetVector<stablehlo::SliceOp> sliceOps;
    SetVector<stablehlo::DynamicSliceOp> dsOps;
    SetVector<stablehlo::DynamicUpdateSliceOp> dusOps;
    for (auto user : op->getUsers()) {
      if (user->getBlock() != op->getBlock()) {
        continue;
      }

      TypeSwitch<Operation *>(user)
          .Case<stablehlo::SliceOp>([&](auto sliceOp) {
            if (shouldPropagateTransposeDown(op, sliceOp)) {
              sliceOps.insert(sliceOp);
            }
          })
          .Case<stablehlo::DynamicSliceOp>([&](auto sliceOp) {
            if (sliceOp.getOperand() != op) {
              return;
            }

            if (shouldPropagateTransposeDown(op, sliceOp)) {
              dsOps.insert(sliceOp);
            }
          })
          .Case<stablehlo::DynamicUpdateSliceOp>([&](auto dusOp) {
            if (dusOp.getOperand() != op) {
              return;
            }

            // check that if DUS is consumed by a DS, can we absorb the
            // transpose due to singleton dims.
            // Should we have some constraint based on the update? though
            // transposing an update will likely lead to lower cost than
            // transposing the operand
            bool valid = true;
            for (auto user : dusOp->getUsers()) {
              if (auto sOp = dyn_cast<stablehlo::SliceOp>(user)) {
                valid &= shouldPropagateTransposeDown(op, sOp);
              } else if (auto dsOp =
                             dyn_cast<stablehlo::DynamicSliceOp>(user)) {
                valid &= shouldPropagateTransposeDown(op, dsOp);
              } else {
                valid &= isFusible(op, user);
              }

              if (!valid) {
                break;
              }
            }

            if (valid) {
              dusOps.insert(dusOp);
            }
          });
    }

    if (sliceOps.empty() && dsOps.empty() && dusOps.empty()) {
      return failure();
    }

    auto replaceSliceOp = [&](auto sliceOp) {
      return this->replaceOp(sliceOp, rewriter, op);
    };

    llvm::for_each(dusOps, replaceSliceOp);
    llvm::for_each(dsOps, replaceSliceOp);
    llvm::for_each(sliceOps, replaceSliceOp);
    return success();
  }

private:
  template <typename OpTy>
  void replaceOp(OpTy sliceOp, PatternRewriter &rewriter,
                 stablehlo::TransposeOp op) const {
    rewriter.setInsertionPoint(sliceOp);
    auto newSlice = sliceTransposeHelper(op, rewriter, sliceOp);
    rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(sliceOp, newSlice,
                                                        op.getPermutation());
  }

  template <typename OpTy>
  bool shouldPropagateTransposeDown(stablehlo::TransposeOp tOp,
                                    OpTy sliceOp) const {
    if (OpIsReshapeLike(tOp,
                        cast<RankedTensorType>(sliceOp.getType()).getShape())) {
      return true;
    }

    return llvm::all_of(sliceOp->getUsers(), [&](auto downstreamUser) {
      return isFusible(tOp, downstreamUser);
    });
  }
};

struct DynamicSliceElementwise final
    : CheckedOpRewritePattern<stablehlo::DynamicSliceOp,
                              DynamicSliceElementwise> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto elem = op.getOperand().getDefiningOp();
    if (!elem) {
      return failure();
    }
    if (!stablehlo::hasTraitElementwise(elem)) {
      return failure();
    }
    if (elem->getBlock() !=
        op->getBlock()) { // very common to use DS inside loops
      return failure();
    }

    // simpler case where we trivially move the slice up
    if (llvm::hasSingleElement(elem->getUsers())) {
      SmallVector<Value> newOperands;

      for (auto v : elem->getOperands()) {
        auto newDS = stablehlo::DynamicSliceOp::create(
            rewriter, op.getLoc(), v, op.getStartIndices(), op.getSliceSizes());
        newOperands.push_back(newDS.getResult());
      }

      auto nex = rewriter.create(
          elem->getLoc(), elem->getName().getIdentifier(),
          ValueRange(newOperands), TypeRange(op->getResult(0).getType()),
          elem->getAttrs(), {}, {});
      rewriter.replaceOp(op, nex);
      return success();
    }

    // For all start indices that are not constants we require those to be
    // equal.
    // we can do this more aggressively inside loops where we have better
    // bound / offset information
    SmallVector<Value> newStartIndices(op.getStartIndices().begin(),
                                       op.getStartIndices().end());
    SmallVector<int64_t> newSliceSizes(op.getSliceSizes().begin(),
                                       op.getSliceSizes().end());

    SmallVector<int64_t> constantStarts(op.getStartIndices().size());
    for (auto [i, startIndex] : llvm::enumerate(newStartIndices)) {
      APInt idx;
      if (matchPattern(startIndex, m_ConstantInt(&idx))) {
        constantStarts[i] = idx.getSExtValue();
      } else {
        constantStarts[i] = -1;
      }
    }

    llvm::SetVector<stablehlo::DynamicSliceOp> allDSUsers;
    stablehlo::DynamicSliceOp firstDSOp;

    for (auto user : elem->getUsers()) {
      auto dsOp = dyn_cast<stablehlo::DynamicSliceOp>(user);
      if (!dsOp) { // we can support slices too by doing a full slice along the
                   // non-constant dimensions, but that is too much headache for
                   // potentially very little gain
        return failure();
      }

      if (elem->getBlock() != dsOp->getBlock()) {
        return failure();
      }

      if (allDSUsers.contains(dsOp)) {
        continue;
      }
      allDSUsers.insert(dsOp);
      if (!firstDSOp) {
        firstDSOp = dsOp;
      } else if (dsOp->isBeforeInBlock(firstDSOp)) {
        firstDSOp = dsOp;
      }

      auto startIndices = dsOp.getStartIndices();
      auto sliceSizes = dsOp.getSliceSizes();

      for (size_t i = 0; i < startIndices.size(); ++i) {
        APInt constIdx;
        if (matchPattern(startIndices[i], m_ConstantInt(&constIdx))) {
          if (constantStarts[i] == -1) {
            return failure();
          }
          int64_t startIdx =
              std::min(constIdx.getSExtValue(), constantStarts[i]);
          int64_t curLimit = constIdx.getSExtValue() + sliceSizes[i];
          int64_t prevLimit = constantStarts[i] + newSliceSizes[i];
          int64_t newLimit = std::max(curLimit, prevLimit);

          constantStarts[i] = startIdx;
          newSliceSizes[i] = newLimit - startIdx;
        } else {
          if (startIndices[i] != newStartIndices[i]) {
            return failure(); // require all non-constant starts to be the same
                              // value
          }
          // take the biggest slice
          newSliceSizes[i] = std::max(newSliceSizes[i], sliceSizes[i]);
        }
      }
    }

    auto indexElemType = newStartIndices[0].getType();
    for (size_t i = 0; i < newStartIndices.size(); ++i) {
      if (constantStarts[i] != -1) { // construct new constant indices
        newStartIndices[i] = stablehlo::ConstantOp::create(
            rewriter, op->getLoc(), indexElemType,
            cast<ElementsAttr>(makeAttr(indexElemType, constantStarts[i])));
      }
    }

    rewriter.setInsertionPoint(firstDSOp);

    SmallVector<Value> newOperands;
    for (auto v : elem->getOperands()) {
      newOperands.push_back(stablehlo::DynamicSliceOpCreate(
          rewriter, op.getLoc(), v, newStartIndices, newSliceSizes));
    }
    auto nex = rewriter.create(
        elem->getLoc(), elem->getName().getIdentifier(),
        ValueRange(newOperands),
        TypeRange(RankedTensorType::get(
            newSliceSizes, cast<RankedTensorType>(op->getResult(0).getType())
                               .getElementType())),
        elem->getAttrs(), {}, {});

    for (auto dsOp : allDSUsers) {
      // we need to replace with slices
      SmallVector<int64_t> sliceStarts(dsOp.getStartIndices().size());
      SmallVector<int64_t> sliceLimits(dsOp.getStartIndices().size());
      SmallVector<int64_t> sliceStrides(dsOp.getStartIndices().size(), 1);

      for (size_t i = 0; i < sliceStarts.size(); i++) {
        if (constantStarts[i] == -1) {
          sliceStarts[i] = 0;
          sliceLimits[i] = dsOp.getSliceSizes()[i];
        } else {
          APInt origStart;
          auto matched = matchPattern(dsOp.getStartIndices()[i],
                                      m_ConstantInt(&origStart));
          assert(matched);
          (void)matched;
          sliceStarts[i] = origStart.getSExtValue() - constantStarts[i];
          sliceLimits[i] = sliceStarts[i] + dsOp.getSliceSizes()[i];
        }
      }

      auto newSliceOp =
          stablehlo::SliceOpCreate(rewriter, dsOp.getLoc(), nex->getResult(0),
                                   sliceStarts, sliceLimits, sliceStrides);
      rewriter.replaceOp(dsOp, newSliceOp);
    }

    return success();
  }
};

struct SliceElementwise final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceElementwise> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto elem = op.getOperand().getDefiningOp();
    if (!elem)
      return failure();
    if (!stablehlo::hasTraitElementwise(elem))
      return failure();
    if (llvm::hasSingleElement(elem->getUsers())) {
      SmallVector<Value> ops;
      for (auto v : elem->getOperands()) {
        ops.push_back(stablehlo::SliceOp::create(
            rewriter, op.getLoc(), v, op.getStartIndices(),
            op.getLimitIndices(), op.getStrides()));
      }
      auto nex = rewriter.create(
          elem->getLoc(), elem->getName().getIdentifier(), ValueRange(ops),
          TypeRange(op->getResult(0).getType()), elem->getAttrs(), {}, {});
      rewriter.replaceOp(op, nex);
      return success();
    }

    SmallVector<int64_t> starts(op.getStartIndices().begin(),
                                op.getStartIndices().end());
    SmallVector<int64_t> stops(op.getLimitIndices().begin(),
                               op.getLimitIndices().end());
    SmallVector<int64_t> ints(op.getStrides().begin(), op.getStrides().end());
    SmallVector<stablehlo::SliceOp> todo;
    SmallVector<int64_t> sizes;
    for (auto u : elem->getUsers()) {
      auto sop = dyn_cast<stablehlo::SliceOp>(u);
      if (!sop)
        return failure();
      for (auto en : llvm::enumerate(sop.getType().getShape())) {
        auto start = sop.getStartIndices()[en.index()];
        auto stop = sop.getLimitIndices()[en.index()];
        auto stride = sop.getStrides()[en.index()];

        if (stride != ints[en.index()])
          ints[en.index()] = 1;
        // If the difference between the starts is not an even multiple of the
        // stride change the stride to 1
        if (start != starts[en.index()]) {
          if ((start - starts[en.index()]) % ints[en.index()] != 0) {
            ints[en.index()] = 1;
          }
        }
        if (start < starts[en.index()])
          starts[en.index()] = start;
        if (stop > stops[en.index()])
          stops[en.index()] = stop;
      }
      todo.push_back(sop);
    }
    bool changed = false;
    for (auto en : llvm::enumerate(op.getOperand().getType().getShape())) {
      if (starts[en.index()] != 0) {
        changed = true;
      }
      if (stops[en.index()] < en.value()) {
        changed = true;
      }
      if (ints[en.index()] != 1) {
        changed = true;
      }
      sizes.push_back((stops[en.index()] - starts[en.index()]) /
                      ints[en.index()]);
    }
    if (!changed)
      return failure();
    rewriter.setInsertionPoint(elem);
    SmallVector<Value> ops;
    for (auto v : elem->getOperands()) {
      ops.push_back(stablehlo::SliceOp::create(rewriter, op.getLoc(), v, starts,
                                               stops, ints));
    }
    auto nex = rewriter.create(
        elem->getLoc(), elem->getName().getIdentifier(), ValueRange(ops),
        TypeRange(RankedTensorType::get(
            sizes, op.getOperand().getType().getElementType())),
        elem->getAttrs(), {}, {});

    for (auto sl : todo) {
      SmallVector<int64_t> sstarts;
      SmallVector<int64_t> sstops;
      SmallVector<int64_t> sints;

      for (auto &&[start, stop, stride, ostart, ostop, ostride] :
           llvm::zip(sl.getStartIndices(), sl.getLimitIndices(),
                     sl.getStrides(), starts, stops, ints)) {
        if (stride == ostride) {
          sstarts.push_back(start - ostart);
          sstops.push_back((stop - ostart) / stride);
          sints.push_back(1);
        } else {
          assert(ostride == 1);
          sstarts.push_back(start - ostart);
          sstops.push_back(stop - ostart);
          sints.push_back(stride);
        }
      }
      rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(sl, nex->getResult(0),
                                                      sstarts, sstops, sints);
    }
    return success();
  }
};

LogicalResult slicePadHelper(
    stablehlo::PadOp pad, ArrayRef<int64_t> starts, ArrayRef<int64_t> limits,
    ArrayRef<int64_t> strides, SmallVectorImpl<int64_t> &start,
    SmallVectorImpl<int64_t> &end, SmallVectorImpl<int64_t> &step,
    SmallVectorImpl<int64_t> &lpads, SmallVectorImpl<int64_t> &hpads,
    SmallVectorImpl<int64_t> &interiors, bool &broadcastres, bool &needspad) {
  assert(start.size() == 0);
  assert(end.size() == 0);
  assert(step.size() == 0);
  assert(lpads.size() == 0);
  assert(hpads.size() == 0);
  assert(interiors.size() == 0);
  assert(!broadcastres);
  assert(!needspad);
  assert(starts.size() == pad.getOperand().getType().getShape().size());
  assert(limits.size() == pad.getOperand().getType().getShape().size());
  assert(strides.size() == pad.getOperand().getType().getShape().size());

  if (anyPadSizesNegative(pad))
    return failure();

  for (auto &&[nstart, nend, nstep, lpad, hpad, interior, inshape, outshape] :
       llvm::zip(starts, limits, strides, pad.getEdgePaddingLow(),
                 pad.getEdgePaddingHigh(), pad.getInteriorPadding(),
                 pad.getOperand().getType().getShape(),
                 pad.getType().getShape())) {
    if (nstep != 1)
      return failure();
    if (interior != 0)
      return failure();

    // slice goes from [nstart, nend]
    // pad result is [0..lpad][lpad... outshape - hpad][outshape - hpad...
    // outshape]

    // start of slice starts after end of value being padded
    if (nstart >= outshape - hpad) {
      broadcastres = true;
      return success();
    }
    // slice ends before the start of value being padded
    if (nend <= lpad) {
      broadcastres = true;
      return success();
    }
    if (nstart - lpad < 0) {
      start.push_back(0);
      lpads.push_back(lpad - nstart);
      needspad = true;
    } else {
      start.push_back(nstart - lpad);
      lpads.push_back(0);
    }
    if (nend - lpad > inshape) {
      end.push_back(inshape);
      hpads.push_back(nend - lpad - inshape);
      needspad = true;
    } else {
      end.push_back(nend - lpad);
      hpads.push_back(0);
    }

    step.push_back(1);
    interiors.push_back(0);
  }
  return success();
}

// slice(pad x) -> pad(slice x)
struct SlicePad final : CheckedOpRewritePattern<stablehlo::SliceOp, SlicePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    SmallVector<int64_t> start;
    SmallVector<int64_t> end;
    SmallVector<int64_t> step;

    SmallVector<int64_t> lpads;
    SmallVector<int64_t> hpads;
    SmallVector<int64_t> interiors;

    bool needspad = false;
    bool broadcastres = false;
    if (!slicePadHelper(pad, op.getStartIndices(), op.getLimitIndices(),
                        op.getStrides(), start, end, step, lpads, hpads,
                        interiors, broadcastres, needspad)
             .succeeded())
      return failure();

    if (broadcastres) {
      rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
          op, op.getType(), pad.getPaddingValue(),
          rewriter.getDenseI64ArrayAttr({}));
      return success();
    }

    if (needspad) {
      auto nslice = stablehlo::SliceOp::create(
          rewriter, op.getLoc(), pad.getOperand(), start, end, step);
      rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
          op, nslice, pad.getPaddingValue(), lpads, hpads, interiors);
    } else {
      rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(op, pad.getOperand(),
                                                      start, end, step);
    }
    return success();
  }
};

// From
// https:github.com/openxla/stablehlo/blob/5d1a9c892500c2e9fecbfedfa66ffe84ff1caf7b/stablehlo/dialect/StablehloOps.cpp#L1498C1-L1532C1
bool hasSameOperandAndResultTypes(Operation &op) {
  Type expected;
  if (op.getNumResults() != 0)
    expected = op.getResult(0).getType();
  if (op.getNumOperands() != 0)
    expected = op.getOperand(0).getType();
  if (!expected)
    return false;

  auto typeMatch = [&](Type actual) { return actual == expected; };
  return llvm::all_of(op.getOperandTypes(), typeMatch) &&
         llvm::all_of(op.getResultTypes(), typeMatch);
}

static bool isEligibleForCompactPrint(stablehlo::ReduceOp op) {
  // Check E1.
  auto &block = op.getBody().front();
  if (!hasSingleElement(block.without_terminator()))
    return false;

  Operation &innerOp = *block.begin();

  // Check E2.
  if (innerOp.getDialect() != op->getDialect())
    return false;

  if (innerOp.getNumOperands() != 2 ||
      !innerOp.hasTrait<mlir::OpTrait::OneResult>() ||
      !hasSameOperandAndResultTypes(innerOp) ||
      !innerOp.hasTrait<mlir::hlo::OpTrait::IsCommutative>() ||
      !innerOp.hasTrait<mlir::OpTrait::ZeroRegions>())
    return false;

  // Check E3.
  if (op.getInputs().empty())
    return false;

  auto elemType =
      cast<ShapedType>(op.getInputs()[0].getType()).getElementType();
  auto expectedInnerOpType = RankedTensorType::get(/*shape=*/{}, elemType);
  if (innerOp.getOperands()[0].getType() != expectedInnerOpType)
    return false;

  // Check E4.
  if (!llvm::equal(block.getArguments(), innerOp.getOperands()))
    return false;

  // Check E5.
  auto retOp = dyn_cast<stablehlo::ReturnOp>(block.getTerminator());
  if (!retOp)
    return false;

  return llvm::equal(innerOp.getResults(), retOp.getOperands());
}

struct ReduceToReshape final
    : CheckedOpRewritePattern<stablehlo::ReduceOp, ReduceToReshape> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1)
      return failure();
    if (!isEligibleForCompactPrint(op))
      return failure();
    auto inpTy = cast<RankedTensorType>(op.getInputs()[0].getType());
    for (auto idx : op.getDimensions()) {
      if (inpTy.getShape()[idx] != 1)
        return failure();
    }

    auto reshaped = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), op.getInputs()[0],
        cast<RankedTensorType>(op.getResult(0).getType()).getShape());

    Operation &innerOp = op.getBody().front().front();

    auto bcast = stablehlo::BroadcastInDimOp::create(
        rewriter, op.getLoc(), reshaped.getType(), op.getInitValues()[0],
        rewriter.getDenseI64ArrayAttr({}));
    Value vals[2] = {bcast, reshaped};
    auto res = rewriter.create(
        op.getLoc(), innerOp.getName().getIdentifier(), ValueRange(vals),
        TypeRange(op->getResult(0).getType()), innerOp.getAttrs(), {}, {});

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

struct ReducePad
    : public CheckedOpRewritePattern<stablehlo::ReduceOp, ReducePad> {
  using CheckedOpRewritePattern<stablehlo::ReduceOp,
                                ReducePad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1 || op.getInitValues().size() != 1) {
      return rewriter.notifyMatchFailure(
          op, "only single-operand single-init reduce is supported");
    }

    Value input = op.getInputs()[0];
    auto pad = input.getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return rewriter.notifyMatchFailure(op, "input source is not a pad op");
    if (anyPadSizesNegative(pad))
      return failure();

    auto checkCommonReduce = mlir::stablehlo::CheckCommonReduceOp(op);

    switch (checkCommonReduce.kind) {
    case stablehlo::ReduceOpKind::Add:
      return matchAndRewriteReduceAdd(op, rewriter, pad, input);
    case stablehlo::ReduceOpKind::Min:
      return matchAndRewriteReduceMin(op, rewriter, pad, input);
    case stablehlo::ReduceOpKind::Max:
      return matchAndRewriteReduceMax(op, rewriter, pad, input);
    case stablehlo::ReduceOpKind::Mul:
      return matchAndRewriteReduceMul(op, rewriter, pad, input);
    default:
      return failure();
    }
  }

private:
  struct PostReduceInfo {
    Value res;
    SmallVector<int64_t> low;
    SmallVector<int64_t> high;
    SmallVector<int64_t> inner;
    bool needsPostPad;
    int64_t numElementsReduced;
  };

  LogicalResult matchAndRewriteReduceMin(stablehlo::ReduceOp op,
                                         PatternRewriter &rewriter,
                                         stablehlo::PadOp pad,
                                         Value input) const {
    auto reduceInfo = matchAndRewriteReduce(op, rewriter, pad, input);
    Value res = stablehlo::MinOpCreate(rewriter, op.getLoc(), reduceInfo.res,
                                       pad.getPaddingValue());

    if (reduceInfo.needsPostPad) {
      res = stablehlo::PadOp::create(rewriter, op.getLoc(), res,
                                     pad.getPaddingValue(), reduceInfo.low,
                                     reduceInfo.high, reduceInfo.inner);
    }
    rewriter.replaceOp(op, res);
    return success();
  }

  LogicalResult matchAndRewriteReduceMax(stablehlo::ReduceOp op,
                                         PatternRewriter &rewriter,
                                         stablehlo::PadOp pad,
                                         Value input) const {
    auto reduceInfo = matchAndRewriteReduce(op, rewriter, pad, input);
    Value res = stablehlo::MaxOpCreate(rewriter, op.getLoc(), reduceInfo.res,
                                       pad.getPaddingValue());

    if (reduceInfo.needsPostPad) {
      res = stablehlo::PadOp::create(rewriter, op.getLoc(), res,
                                     pad.getPaddingValue(), reduceInfo.low,
                                     reduceInfo.high, reduceInfo.inner);
    }
    rewriter.replaceOp(op, res);
    return success();
  }

  LogicalResult matchAndRewriteReduceMul(stablehlo::ReduceOp op,
                                         PatternRewriter &rewriter,
                                         stablehlo::PadOp pad,
                                         Value input) const {
    auto reduceInfo = matchAndRewriteReduce(op, rewriter, pad, input);
    Value res = reduceInfo.res;

    // compute padValue ^ numElementsReduced and multiply this to the result
    auto cType = RankedTensorType::get(
        {}, cast<RankedTensorType>(pad.getPaddingValue().getType())
                .getElementType());
    auto mulConst = stablehlo::PowOp::create(
        rewriter, op.getLoc(), pad.getPaddingValue(),
        stablehlo::ConstantOp::create(
            rewriter, op.getLoc(), cType,
            cast<ElementsAttr>(
                makeAttr(cType, reduceInfo.numElementsReduced))));
    res = stablehlo::MulOpCreate(rewriter, op.getLoc(), res, mulConst);

    if (reduceInfo.needsPostPad) {
      res = stablehlo::PadOp::create(rewriter, op.getLoc(), res,
                                     pad.getPaddingValue(), reduceInfo.low,
                                     reduceInfo.high, reduceInfo.inner);
    }
    rewriter.replaceOp(op, res);
    return success();
  }

  LogicalResult matchAndRewriteReduceAdd(stablehlo::ReduceOp op,
                                         PatternRewriter &rewriter,
                                         stablehlo::PadOp pad,
                                         Value input) const {
    auto reduceInfo = matchAndRewriteReduce(op, rewriter, pad, input);
    Value res = reduceInfo.res;

    Attribute padValAttr;
    bool padValConst =
        matchPattern(pad.getPaddingValue(), m_Constant(&padValAttr));
    if (!padValConst || (!matchPattern(padValAttr, m_AnyZeroFloat()) &&
                         !matchPattern(padValAttr, m_AnyZeroComplex()) &&
                         !matchPattern(padValAttr, m_Zero()))) {
      // compute padValue * numElementsReduced and add this to the result
      auto cType = RankedTensorType::get(
          {}, cast<RankedTensorType>(pad.getPaddingValue().getType())
                  .getElementType());
      auto mulConst = stablehlo::MulOp::create(
          rewriter, op.getLoc(), pad.getPaddingValue(),
          stablehlo::ConstantOp::create(
              rewriter, op.getLoc(), cType,
              cast<ElementsAttr>(
                  makeAttr(cType, reduceInfo.numElementsReduced))));
      res = stablehlo::AddOpCreate(rewriter, op.getLoc(), res, mulConst);
    }

    if (reduceInfo.needsPostPad) {
      res = stablehlo::PadOp::create(rewriter, op.getLoc(), res,
                                     pad.getPaddingValue(), reduceInfo.low,
                                     reduceInfo.high, reduceInfo.inner);
    }
    rewriter.replaceOp(op, res);
    return success();
  }

  PostReduceInfo matchAndRewriteReduce(stablehlo::ReduceOp op,
                                       PatternRewriter &rewriter,
                                       stablehlo::PadOp pad,
                                       Value input) const {
    SmallVector<int64_t> low, high, inner;
    int64_t numElementsReduced = 0;
    bool needsPostPad = false;
    for (auto en : llvm::enumerate(
             cast<RankedTensorType>(pad.getOperand().getType()).getShape())) {
      int64_t lowPad = pad.getEdgePaddingLow()[en.index()];
      int64_t highPad = pad.getEdgePaddingHigh()[en.index()];
      int64_t interiorPad = pad.getInteriorPadding()[en.index()];
      if (llvm::is_contained(op.getDimensions(), en.index())) {
        numElementsReduced +=
            lowPad + highPad +
            std::max(static_cast<int32_t>(en.value() - 1), 0) * interiorPad;
        continue;
      }
      low.push_back(lowPad);
      high.push_back(highPad);
      inner.push_back(interiorPad);
      needsPostPad = true;
    }

    auto newReduction = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), ValueRange(pad.getOperand()), op.getInitValues(),
        op.getDimensions());
    newReduction.getRegion().takeBody(op.getRegion());

    Value res = newReduction->getResult(0);
    return {res, low, high, inner, needsPostPad, numElementsReduced};
  }
};

struct ConvertConcat final
    : CheckedOpRewritePattern<stablehlo::ConvertOp, ConvertConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    auto concat = op.getOperand().getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concat)
      return failure();

    SmallVector<Value> newvals;
    for (auto v : concat.getOperands()) {
      newvals.push_back(stablehlo::ConvertOp::create(
          rewriter, op.getLoc(),
          RankedTensorType::get(cast<RankedTensorType>(v.getType()).getShape(),
                                op.getType().getElementType()),
          v));
    }
    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
        op, newvals, concat.getDimension());
    return success();
  }
};

struct ConvertConvertFloat final
    : CheckedOpRewritePattern<stablehlo::ConvertOp, ConvertConvertFloat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    auto conv0 = op.getOperand().getDefiningOp<stablehlo::ConvertOp>();
    if (!conv0)
      return failure();

    auto prev = conv0.getOperand();
    if (isa<FloatType>(prev.getType().getElementType()) &&
        isa<FloatType>(op.getType().getElementType()) &&
        isa<FloatType>(conv0.getType().getElementType())) {
      if (prev.getType() == op.getType()) {
        rewriter.replaceOp(op, prev);
        return success();
      }
      rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(op, op.getType(), prev);
      return success();
    }
    return failure();
  }
};

struct ConvertConvertInt final
    : CheckedOpRewritePattern<stablehlo::ConvertOp, ConvertConvertInt> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    auto conv0 = op.getOperand().getDefiningOp<stablehlo::ConvertOp>();
    if (!conv0)
      return failure();

    auto prev = conv0.getOperand();
    if (isa<IntegerType>(prev.getType().getElementType()) &&
        isa<IntegerType>(op.getType().getElementType()) &&
        isa<IntegerType>(conv0.getType().getElementType())) {
      // we only do the elimination if we go from low bitwidth to high bitwidth
      auto prevwidth = prev.getType().getElementType().getIntOrFloatBitWidth();
      auto midwidth = conv0.getType().getElementType().getIntOrFloatBitWidth();
      if (prevwidth > midwidth)
        return rewriter.notifyMatchFailure(op, "prevwidth > midwidth");

      if (prev.getType() == op.getType()) {
        rewriter.replaceOp(op, prev);
        return success();
      }
      rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(op, op.getType(), prev);
      return success();
    }
    return failure();
  }
};

struct ReduceConcat final
    : CheckedOpRewritePattern<stablehlo::ReduceOp, ReduceConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1)
      return failure();

    auto concat = op.getInputs()[0].getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concat)
      return failure();

    auto dim = concat.getDimension();

    if (!llvm::is_contained(op.getDimensions(), dim))
      return failure();

    Value prev = op.getInitValues()[0];

    auto checkCommonReduce = mlir::stablehlo::CheckCommonReduceOp(op);

    Value identity = nullptr;
    switch (checkCommonReduce.kind) {
    case stablehlo::ReduceOpKind::Add:
      identity = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), prev.getType(),
          cast<ElementsAttr>(makeAttr(prev.getType(), 0)));
      break;
    case stablehlo::ReduceOpKind::Min:
    case stablehlo::ReduceOpKind::Max:
      identity = prev;
      break;
    default:
      return failure();
    }

    Operation &innerOp = op.getBody().front().front();

    if (prev.getType() != op.getResultTypes()[0]) {
      prev = stablehlo::BroadcastInDimOp::create(rewriter, op.getLoc(),
                                                 op.getResultTypes()[0], prev,
                                                 ArrayRef<int64_t>());
    }

    for (auto v : concat.getOperands()) {
      IRMapping map;
      map.map(op.getInitValues()[0], identity);
      map.map(op.getInputs()[0], v);
      auto next = rewriter.clone(*op, map)->getResult(0);
      map.map(innerOp.getOperand(0), prev);
      map.map(innerOp.getOperand(1), next);
      Value vals[] = {prev, next};
      prev =
          rewriter
              .create(innerOp.getLoc(), innerOp.getName().getIdentifier(), vals,
                      TypeRange(prev.getType()), innerOp.getAttrs(), {}, {})
              ->getResult(0);
    }

    assert(op.getResultTypes()[0] == prev.getType());
    rewriter.replaceOp(op, prev);
    return success();
  }
};

struct FullReduceReshapeOrTranspose final
    : CheckedOpRewritePattern<stablehlo::ReduceOp,
                              FullReduceReshapeOrTranspose> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1)
      return failure();

    auto inpTy = cast<RankedTensorType>(op.getInputs()[0].getType());
    if (op.getDimensions().size() != inpTy.getShape().size())
      return failure();

    RankedTensorType changeType = nullptr;
    SmallVector<Operation *> reshapeOrTransposes;
    llvm::MapVector<Operation *, int> toclone;
    {
      SmallVector<Value> todo = {op.getInputs()[0]};
      while (todo.size()) {
        auto cur = todo.pop_back_val();
        auto curOp = cur.getDefiningOp();
        if (!curOp)
          return failure();
        if (auto rs = dyn_cast<stablehlo::ReshapeOp>(curOp)) {
          if (changeType != nullptr) {
            if (rs.getOperand().getType() != changeType)
              return failure();
          } else {
            changeType = rs.getOperand().getType();
          }
          reshapeOrTransposes.push_back(rs);
          continue;
        }
        if (auto rs = dyn_cast<stablehlo::TransposeOp>(curOp)) {
          if (changeType != nullptr) {
            if (rs.getOperand().getType() != changeType)
              return failure();
          } else {
            changeType = rs.getOperand().getType();
          }
          reshapeOrTransposes.push_back(rs);
          continue;
        }
        if (!hasTraitElementwise(curOp))
          return failure();
        if (!isMemoryEffectFree(curOp))
          return failure();
        for (auto op : curOp->getOperands())
          todo.push_back(op);
        toclone[curOp] = curOp->getNumOperands();
      }
    }

    IRMapping map;
    SmallVector<Operation *> todo;
    for (auto reshape : reshapeOrTransposes) {
      map.map(reshape->getResult(0), reshape->getOperand(0));
      for (auto u : reshape->getResult(0).getUsers()) {
        if (toclone.contains(u)) {
          toclone[u]--;
          if (toclone[u] == 0) {
            todo.push_back(u);
            toclone.erase(u);
          }
        }
      }
    }
    for (auto pair : toclone) {
      for (auto u : pair.first->getResult(0).getUsers()) {
        if (u == op)
          continue;
        if (llvm::is_contained(reshapeOrTransposes, u))
          continue;
        if (toclone.contains(u))
          continue;
        return failure();
      }
    }

    for (auto cur : todo) {
      for (auto curOp : cur->getOperands()) {
        if (!map.contains(curOp))
          return failure();
      }
    }

    while (todo.size()) {
      auto cur = todo.pop_back_val();

      SmallVector<Value> vals;
      for (auto op : cur->getOperands())
        vals.push_back(map.lookup(op));

      auto changeType2 = RankedTensorType::get(
          changeType.getShape(),
          cast<RankedTensorType>(cur->getResult(0).getType()).getElementType());
      auto res =
          rewriter.create(cur->getLoc(), cur->getName().getIdentifier(), vals,
                          TypeRange(changeType2), cur->getAttrs(), {}, {});

      map.map(cur->getResult(0), res->getResult(0));

      for (auto u : cur->getResult(0).getUsers()) {
        if (toclone.contains(u)) {
          toclone[u]--;
          if (toclone[u] == 0) {
            todo.push_back(u);
            toclone.erase(u);
          }
        }
      }
    }

    SmallVector<int64_t> newReduceDimensions;
    for (size_t i = 0, end = changeType.getShape().size(); i < end; i++)
      newReduceDimensions.push_back(i);

    auto newReduction = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), op->getResultTypes(),
        map.lookup(op.getInputs()[0]), op.getInitValues(), newReduceDimensions);
    newReduction.getRegion().takeBody(op.getRegion());
    rewriter.replaceOp(op, newReduction);
    return success();
  }
};

LogicalResult sliceConcatHelper(stablehlo::ConcatenateOp concat,
                                PatternRewriter &rewriter,
                                ArrayRef<int64_t> starts,
                                ArrayRef<int64_t> limits,
                                ArrayRef<int64_t> strides,
                                SmallVectorImpl<Value> &postConcat) {
  auto dim = concat.getDimension();

  if (strides[dim] != 1)
    return failure();

  assert(postConcat.size() == 0);
  size_t curdim = 0;
  for (auto v : concat.getInputs()) {
    auto ty = cast<RankedTensorType>(v.getType());
    auto nextdim = ty.getShape()[dim];
    if (starts[dim] >= curdim + nextdim) {
      curdim += nextdim;
      continue;
    }
    if (limits[dim] <= curdim) {
      curdim += nextdim;
      continue;
    }
    SmallVector<int64_t> nstart(starts.begin(), starts.end());
    SmallVector<int64_t> nend(limits.begin(), limits.end());
    nstart[dim] -= curdim;
    if (nstart[dim] < 0)
      nstart[dim] = 0;
    nend[dim] -= curdim;
    if (nend[dim] > nextdim)
      nend[dim] = nextdim;
    auto subslice = stablehlo::SliceOp::create(rewriter, concat.getLoc(), v,
                                               nstart, nend, strides);
    postConcat.push_back(subslice);
    curdim += nextdim;
  }
  return success();
}

struct ConcatSlice final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto dim = op.getDimension();
    SmallVector<Value> newOperands;

    auto oldOperands = llvm::to_vector(op.getOperands());
    auto res = concatSliceSimplify(rewriter, oldOperands, dim, newOperands);
    if (!res.succeeded())
      return res;

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, newOperands, dim);
    return success();
  }
};

bool canMergePadsAlongAxis(int dimension, stablehlo::PadOp pad,
                           stablehlo::PadOp otherPad) {
  if (otherPad.getPaddingValue() != pad.getPaddingValue())
    return false;

  for (int d = 0, ndims = pad.getType().getShape().size(); d < ndims; ++d) {
    if (d == dimension) {
      if (pad.getInteriorPadding()[d] != 0) {
        return false;
      }
      if (otherPad.getInteriorPadding()[d] != 0) {
        return false;
      }
      if (pad.getEdgePaddingHigh()[d] != 0) {
        return false;
      }
      if (otherPad.getEdgePaddingLow()[d] != 0) {
        return false;
      }
    } else {
      if (pad.getInteriorPadding()[d] != otherPad.getInteriorPadding()[d]) {
        return false;
      }
      if (pad.getEdgePaddingHigh()[d] != otherPad.getEdgePaddingHigh()[d]) {
        return false;
      }
      if (pad.getEdgePaddingLow()[d] != otherPad.getEdgePaddingLow()[d]) {
        return false;
      }
    }
  }
  return true;
}

struct ConcatMultiPad final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatMultiPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto dim = op.getDimension();

    SmallVector<Value> newOperands;

    for (int i = 0, e = op->getNumOperands(); i < e; ++i) {
      auto operand = op->getOperand(i);
      auto pad = operand.getDefiningOp<stablehlo::PadOp>();

      if (!pad) {
        newOperands.push_back(operand);
        continue;
      }

      stablehlo::PadOp otherPadOp;
      while (i + 1 < e &&
             (otherPadOp =
                  op->getOperand(i + 1).getDefiningOp<stablehlo::PadOp>())) {
        if (canMergePadsAlongAxis(op.getDimension(), pad, otherPadOp)) {
          Value padops[] = {pad.getOperand(), otherPadOp.getOperand()};
          auto subConcat = stablehlo::ConcatenateOp::create(
              rewriter, op.getLoc(), padops, op.getDimension());
          pad = stablehlo::PadOp::create(
              rewriter, pad->getLoc(), subConcat, pad.getPaddingValue(),
              pad.getEdgePaddingLow(), pad.getEdgePaddingHigh(),
              pad.getInteriorPadding());
          i++;
        } else
          break;
      }

      newOperands.push_back(pad.getResult());
    }

    if (newOperands.size() == op->getNumOperands())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, newOperands, dim);
    return success();
  }
};

bool canMergeWrapsAlongAxis(int dimension, enzymexla::WrapOp wrap,
                            enzymexla::WrapOp otherWrap) {
  if (wrap.getDimension() != otherWrap.getDimension())
    return false;

  if (dimension == wrap.getDimension())
    return false;

  if (wrap.getLhs() != otherWrap.getLhs())
    return false;

  if (wrap.getRhs() != otherWrap.getRhs())
    return false;

  return true;
}

struct ConcatWrap final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatWrap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto dim = op.getDimension();

    SmallVector<Value> newOperands;

    for (int i = 0, e = op->getNumOperands(); i < e; ++i) {
      auto operand = op->getOperand(i);
      auto wrap = operand.getDefiningOp<enzymexla::WrapOp>();

      if (!wrap) {
        newOperands.push_back(operand);
        continue;
      }

      enzymexla::WrapOp otherWrap;
      while (i + 1 < e &&
             (otherWrap =
                  op->getOperand(i + 1).getDefiningOp<enzymexla::WrapOp>())) {
        if (canMergeWrapsAlongAxis(op.getDimension(), wrap, otherWrap)) {
          Value padops[] = {wrap.getOperand(), otherWrap.getOperand()};
          auto subConcat = stablehlo::ConcatenateOp::create(
              rewriter, op.getLoc(), padops, op.getDimension());
          wrap = enzymexla::WrapOp::create(rewriter, wrap->getLoc(), subConcat,
                                           wrap.getLhs(), wrap.getRhs(),
                                           wrap.getDimension());
          i++;
        } else
          break;
      }

      newOperands.push_back(wrap.getResult());
    }

    if (newOperands.size() == op->getNumOperands())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, newOperands, dim);
    return success();
  }
};

bool isOuterReducingReshape(stablehlo::ReshapeOp op) {
  auto prevT = cast<RankedTensorType>(op.getOperand().getType());
  if (prevT.getShape().size() != op.getType().getShape().size() + 1)
    return false;
  if (prevT.getShape()[0] != 1)
    return false;
  for (int i = 1; i < prevT.getShape().size(); i++) {
    if (prevT.getShape()[i] != op.getType().getShape()[i - 1])
      return false;
  }
  return true;
}

bool isSliceOf(Value wrapOperand, Value widenOperand, int dim,
               bool widenOperandOnLeft, int existingOffset) {
  auto wrapType = cast<RankedTensorType>(wrapOperand.getType());

  if (auto sliceWiden = widenOperand.getDefiningOp<stablehlo::SliceOp>()) {

    SmallVector<int64_t> wrapStarts(wrapType.getShape().size(), 0);
    SmallVector<int64_t> wrapLimits = llvm::to_vector(wrapType.getShape());
    SmallVector<int64_t> wrapStrides(wrapType.getShape().size(), 1);

    if (sliceWiden.getOperand() != wrapOperand) {
      if (auto sliceWrap = wrapOperand.getDefiningOp<stablehlo::SliceOp>()) {
        wrapStarts = llvm::to_vector(sliceWrap.getStartIndices());
        wrapLimits = llvm::to_vector(sliceWrap.getLimitIndices());
        wrapStrides = llvm::to_vector(sliceWrap.getStrides());
        if (sliceWiden.getOperand() != sliceWrap.getOperand()) {
          return false;
        }
      }
    }

    for (int j = 0; j < wrapType.getShape().size(); j++) {
      if (j != dim) {
        if (sliceWiden.getStartIndices()[j] != wrapStarts[j]) {
          return false;
        }
        if (sliceWiden.getLimitIndices()[j] != wrapLimits[j]) {
          return false;
        }
        if (sliceWiden.getStrides()[j] != wrapStrides[j]) {
          return false;
        }
        continue;
      }

      if (sliceWiden.getStrides()[j] != 1 || wrapStrides[j] != 1) {
        return false;
      }

      if (widenOperandOnLeft) {
        // widenOperand[B-lhs-existing:B-existing], wrap(wrapOperand[A:B],
        // lhs=existingOffset)
        if (sliceWiden.getLimitIndices()[j] + existingOffset != wrapLimits[j]) {
          return false;
        }
      } else {
        // wrap(wrapOperand[A:B], rhs=existing),
        // widenOperand[A+existing:A+rhs+existing]
        if (sliceWiden.getStartIndices()[j] != wrapStarts[j] + existingOffset) {
          return false;
        }
      }
    }

    return true;
  }

  if (auto concatWiden =
          widenOperand.getDefiningOp<stablehlo::ConcatenateOp>()) {
    if (auto concatWrap =
            wrapOperand.getDefiningOp<stablehlo::ConcatenateOp>()) {
      if (concatWrap.getDimension() == concatWiden.getDimension() &&
          concatWrap.getDimension() != dim &&
          concatWiden.getOperands().size() == concatWrap.getOperands().size()) {
        for (auto &&[newWrap, newWiden] : llvm::zip_equal(
                 concatWrap.getOperands(), concatWiden.getOperands())) {
          if (!isSliceOf(newWrap, newWiden, dim, widenOperandOnLeft,
                         existingOffset)) {
            return false;
          }
        }
        return true;
      }
    }
  }

  return false;
}

bool isExtendOf(Value extendOperand, Value widenOperand, int dim,
                bool widenOperandOnLeft) {
  auto extendType = cast<RankedTensorType>(extendOperand.getType());

  if (auto sliceWiden = widenOperand.getDefiningOp<stablehlo::SliceOp>()) {

    SmallVector<int64_t> extendStarts(extendType.getShape().size(), 0);
    SmallVector<int64_t> extendLimits = llvm::to_vector(extendType.getShape());
    SmallVector<int64_t> extendStrides(extendType.getShape().size(), 1);

    if (sliceWiden.getOperand() != extendOperand) {
      if (auto sliceExtend =
              extendOperand.getDefiningOp<stablehlo::SliceOp>()) {
        extendStarts = llvm::to_vector(sliceExtend.getStartIndices());
        extendLimits = llvm::to_vector(sliceExtend.getLimitIndices());
        extendStrides = llvm::to_vector(sliceExtend.getStrides());
        if (sliceWiden.getOperand() != sliceExtend.getOperand()) {
          return false;
        }
      }
    }

    for (int j = 0; j < extendType.getShape().size(); j++) {
      if (j != dim) {
        if (sliceWiden.getStartIndices()[j] != extendStarts[j]) {
          return false;
        }
        if (sliceWiden.getLimitIndices()[j] != extendLimits[j]) {
          return false;
        }
        if (sliceWiden.getStrides()[j] != extendStrides[j]) {
          return false;
        }
        continue;
      }

      if (sliceWiden.getStrides()[j] != 1 || extendStrides[j] != 1) {
        return false;
      }

      if (cast<RankedTensorType>(widenOperand.getType()).getShape()[j] != 1) {
        return false;
      }

      if (widenOperandOnLeft) {
        // widenOperand[A:A+1], extend(extendOperand[A:B])
        if (sliceWiden.getStartIndices()[j] != extendStarts[j]) {
          return false;
        }
      } else {
        // extend(extendOperand[A:B]), widenOperand[B-1:B]
        if (sliceWiden.getLimitIndices()[j] != extendLimits[j]) {
          return false;
        }
      }
    }

    return true;
  }

  if (auto concatWiden =
          widenOperand.getDefiningOp<stablehlo::ConcatenateOp>()) {
    if (auto concatExtend =
            extendOperand.getDefiningOp<stablehlo::ConcatenateOp>()) {
      if (concatExtend.getDimension() == concatWiden.getDimension() &&
          concatExtend.getDimension() != dim &&
          concatWiden.getOperands().size() ==
              concatExtend.getOperands().size()) {
        for (auto &&[newExtend, newWiden] : llvm::zip_equal(
                 concatExtend.getOperands(), concatWiden.getOperands())) {
          if (!isExtendOf(newExtend, newWiden, dim, widenOperandOnLeft)) {
            return false;
          }
        }
        return true;
      }
    }
  }

  return false;
}
struct WidenWrap final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, WidenWrap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto dim = op.getDimension();

    SmallVector<Value> newOperands;
    bool changed = false;

    for (int i = 0, e = op->getNumOperands(); i < e; ++i) {
      auto operand = op->getOperand(i);
      auto wrap = operand.getDefiningOp<enzymexla::WrapOp>();

      if (!wrap || wrap.getDimension() != dim) {
        newOperands.push_back(operand);
        continue;
      }

      if (newOperands.size()) {
        Value prev = newOperands.back();
        Value wrapOperand = wrap.getOperand();

        int withReshape = 0;
        auto rsWrap = wrapOperand.getDefiningOp<stablehlo::ReshapeOp>();
        auto rsWiden = prev.getDefiningOp<stablehlo::ReshapeOp>();
        if (rsWrap && rsWiden && isOuterReducingReshape(rsWrap) &&
            isOuterReducingReshape(rsWiden)) {

          auto postRsTy = cast<ShapedType>(rsWiden.getResult().getType());

          wrapOperand = rsWrap.getOperand();
          prev = rsWiden.getOperand();

          auto prevTy = cast<ShapedType>(prev.getType());
          withReshape = prevTy.getRank() - postRsTy.getRank();
        }

        if (isSliceOf(wrapOperand, prev, dim + withReshape,
                      /*widenOperandOnLeft*/ true, wrap.getLhs())) {
          Value newWrap = enzymexla::WrapOp::create(
              rewriter, wrap.getLoc(), wrap.getOperand(),
              wrap.getLhs() + cast<RankedTensorType>(prev.getType())
                                  .getShape()[dim + withReshape],
              wrap.getRhs(), wrap.getDimension());
          newOperands.pop_back();
          newOperands.push_back(newWrap);
          changed = true;
          continue;
        }
      }

      if (i + 1 < e) {
        Value prev = op->getOperand(i + 1);
        Value wrapOperand = wrap.getOperand();

        int withReshape = 0;
        auto rsWrap = wrapOperand.getDefiningOp<stablehlo::ReshapeOp>();
        auto rsWiden = prev.getDefiningOp<stablehlo::ReshapeOp>();
        if (rsWrap && rsWiden && isOuterReducingReshape(rsWrap) &&
            isOuterReducingReshape(rsWiden)) {

          auto postRsTy = cast<ShapedType>(rsWiden.getResult().getType());

          wrapOperand = rsWrap.getOperand();
          prev = rsWiden.getOperand();

          auto prevTy = cast<ShapedType>(prev.getType());
          withReshape = prevTy.getRank() - postRsTy.getRank();
        }

        if (isSliceOf(wrapOperand, prev, dim + withReshape,
                      /*widenOperandOnLeft*/ false, wrap.getRhs())) {
          auto newWrap = enzymexla::WrapOp::create(
              rewriter, wrap.getLoc(), wrap.getOperand(), wrap.getLhs(),
              wrap.getRhs() + cast<RankedTensorType>(prev.getType())
                                  .getShape()[dim + withReshape],
              wrap.getDimension());
          newOperands.push_back(newWrap);
          i++;
          changed = true;
          continue;
        }
      }

      newOperands.push_back(operand);
      continue;
    }

    if (!changed)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, newOperands, dim);
    return success();
  }
};

struct WidenExtend final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, WidenExtend> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto dim = op.getDimension();

    SmallVector<Value> newOperands;
    bool changed = false;

    for (int i = 0, e = op->getNumOperands(); i < e; ++i) {
      auto operand = op->getOperand(i);
      auto extend = operand.getDefiningOp<enzymexla::ExtendOp>();

      if (!extend || extend.getDimension() != dim) {
        newOperands.push_back(operand);
        continue;
      }

      if (newOperands.size()) {
        auto prev = newOperands.back();
        if (isExtendOf(extend.getOperand(), prev, dim,
                       /*widenOperandOnLeft*/ true)) {
          auto newExtend = enzymexla::ExtendOp::create(
              rewriter, extend.getLoc(), extend.getOperand(),
              extend.getLhs() +
                  cast<RankedTensorType>(prev.getType()).getShape()[dim],
              extend.getRhs(), extend.getDimension());
          newOperands.pop_back();
          newOperands.push_back(newExtend);
          changed = true;
          continue;
        }
      }

      if (i + 1 < e) {
        auto prev = op->getOperand(i + 1);
        if (isExtendOf(extend.getOperand(), prev, dim,
                       /*widenOperandOnLeft*/ false)) {
          auto newExtend = enzymexla::ExtendOp::create(
              rewriter, extend.getLoc(), extend.getOperand(), extend.getLhs(),
              extend.getRhs() +
                  cast<RankedTensorType>(prev.getType()).getShape()[dim],
              extend.getDimension());
          newOperands.push_back(newExtend);
          i++;
          changed = true;
          continue;
        }
      }

      newOperands.push_back(operand);
      continue;
    }

    if (!changed)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, newOperands, dim);
    return success();
  }
};

struct SliceConcat final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto concat = op.getOperand().getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concat)
      return failure();

    auto dim = concat.getDimension();

    SmallVector<Value> postConcat;
    if (!sliceConcatHelper(concat, rewriter, op.getStartIndices(),
                           op.getLimitIndices(), op.getStrides(), postConcat)
             .succeeded())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, postConcat, dim);
    return success();
  }
};

DenseElementsAttr fromTensor(stablehlo::Tensor tensor) {
  return stablehlo::makeDenseElementsAttr(tensor);
}

/*
%22 = stablehlo.dot_general %21, %16, contracting_dims = [1] x [0], precision
= [DEFAULT, DEFAULT] : (tensor<288x288xf32>, tensor<288xf32>) ->
tensor<288xf32>
%27 = stablehlo.reshape %22 : (tensor<288xf32>) ->
tensor<144x2xf32>
%28 = stablehlo.dot_general %6, %27, batching_dims = [0] x
[0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] :
(tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>

should become

%a21 = stablehlo.reshape %21 : (tensor<288xf32>) -> tensor<144x2xf32>

%22 = stablehlo.dot_general %a21, %16, batching_dims = [1] x [],
contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] :
(tensor<144x2x288xf32>, tensor<288xf32>) -> tensor<2x144xf32>

%28 = stablehlo.dot_general %6, %22, batching_dims = [0] x [1],
contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] :
(tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>

TODO
*/

struct DotReshapeDot final
    : CheckedOpRewritePattern<stablehlo::DotGeneralOp, DotReshapeDot> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    return failure();
  }
};

struct PadSimplify final
    : CheckedOpRewritePattern<stablehlo::PadOp, PadSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;
  size_t max_constant_expansion;

  PadSimplify(size_t max_constant_expansion, MLIRContext *context,
              PatternBenefit benefit = 1,
              ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

  LogicalResult matchAndRewriteImpl(stablehlo::PadOp op,
                                    PatternRewriter &rewriter) const {

    Attribute operandAttr, pvAttr;
    matchPattern(op.getOperand(), m_Constant(&operandAttr));
    matchPattern(op.getPaddingValue(), m_Constant(&pvAttr));

    if (operandAttr && matchPattern(operandAttr, m_AnyZeroFloat())) {
      if (pvAttr && matchPattern(pvAttr, m_AnyZeroFloat())) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            op, op.getType(), cast<ElementsAttr>(makeAttr(op.getType(), 0)));
        return success();
      }
    }

    {
      auto inp = dyn_cast_or_null<DenseElementsAttr>(operandAttr);
      auto pv = dyn_cast_or_null<DenseElementsAttr>(pvAttr);
      if (inp && pv) {

        if (inp.isSplat() && pv.isSplat() &&
            inp.getSplatValue<Attribute>() == pv.getSplatValue<Attribute>()) {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, op.getType(), inp.resizeSplat(op.getType()));
          return success();
        }

        size_t size = 1;
        for (auto sz : op.getType().getShape())
          size *= sz;
        if (size < max_constant_expansion) {
          auto ten = stablehlo::constantOp(inp);
          auto out = fromTensor(stablehlo::padOp(
              stablehlo::constantOp(inp), stablehlo::constantOp(pv),
              stablehlo::Sizes(op.getEdgePaddingLow()),
              stablehlo::Sizes(op.getInteriorPadding()), op.getType()));

          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(),
                                                             out);
          return success();
        }
      }
    }

    for (auto &&[low, high, inner] :
         llvm::zip(op.getEdgePaddingLow(), op.getEdgePaddingHigh(),
                   op.getInteriorPadding())) {
      if (low != 0)
        return failure();
      if (high != 0)
        return failure();
      if (inner != 0)
        return failure();
    }
    rewriter.replaceOp(op, op.getOperand());
    return success();
  }
};

struct RotatePad final
    : CheckedOpRewritePattern<enzymexla::RotateOp, RotatePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::RotateOp rotate,
                                    PatternRewriter &rewriter) const {

    auto pad = rotate.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    if (pad.getEdgePaddingLow()[rotate.getDimension()] != 0)
      return failure();
    if (pad.getEdgePaddingHigh()[rotate.getDimension()] != 0)
      return failure();
    if (pad.getInteriorPadding()[rotate.getDimension()] != 0)
      return failure();

    auto newRotate =
        enzymexla::RotateOp::create(rewriter, rotate.getLoc(), pad.getOperand(),
                                    rotate.getAmount(), rotate.getDimension());

    auto newPad = stablehlo::PadOp::create(
        rewriter, pad.getLoc(), newRotate, pad.getPaddingValue(),
        pad.getEdgePaddingLow(), pad.getEdgePaddingHigh(),
        pad.getInteriorPadding());

    rewriter.replaceOp(rotate, newPad);
    return success();
  }
};

struct ShiftRightLogicalSimplify final
    : CheckedOpRewritePattern<stablehlo::ShiftRightLogicalOp,
                              ShiftRightLogicalSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ShiftRightLogicalOp op,
                                    PatternRewriter &rewriter) const {

    DenseElementsAttr lhs;
    matchPattern(op.getLhs(), m_Constant(&lhs));
    DenseElementsAttr rhs;
    matchPattern(op.getRhs(), m_Constant(&rhs));
    if (lhs && rhs) {
      auto out = fromTensor(stablehlo::shiftRightLogicalOp(
          stablehlo::constantOp(lhs), stablehlo::constantOp(rhs),
          op.getType()));

      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(), out);
      return success();
    }
    return failure();
  }
};

struct WhileDeadResults final
    : CheckedOpRewritePattern<stablehlo::WhileOp, WhileDeadResults> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  void replaceTerminator(PatternRewriter &rewriter, Region &region,
                         ArrayRef<int64_t> deadResults) const {
    Operation *terminator = region.front().getTerminator();
    SmallVector<Value> terminatorOperands;
    for (auto &&[i, operand] : llvm::enumerate(terminator->getOperands())) {
      if (!llvm::is_contained(deadResults, i))
        terminatorOperands.push_back(operand);
    }
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPoint(terminator);
    rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(
        terminator, TypeRange(), terminatorOperands, terminator->getAttrs());
  }

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp op,
                                    PatternRewriter &rewriter) const {
    bool isTotalPure = isPure(op);
    std::map<int64_t, int> forwardIsPure;
    std::map<int64_t, std::set<int64_t>> backwardsUsesOfArguments;
    SmallVector<size_t> emptyResults;
    size_t NumResults = op.getNumResults();
    for (OpResult result : op.getResults()) {
      if (result.use_empty())
        emptyResults.push_back(result.getResultNumber());
    }
    if (emptyResults.size() == 0)
      return failure();

    DenseMap<Value, llvm::SmallPtrSet<size_t, 3> *> users(
        emptyResults.size() + op.getCond().front().getOperations().size());
    llvm::SmallVector<std::unique_ptr<llvm::SmallPtrSet<size_t, 3>>, 3>
        condCache;
    condCache.reserve(emptyResults.size());
    for (auto ores : emptyResults) {
      auto ba = op.getCond().front().getArgument(ores);
      auto up = std::make_unique<llvm::SmallPtrSet<size_t, 3>>();
      up->insert(ores);
      users[ba] = up.get();
      condCache.emplace_back(std::move(up));
    }

    llvm::SmallPtrSet<size_t, 3> nonPure;

    SmallVector<
        std::pair<Operation *, std::unique_ptr<llvm::SmallPtrSet<size_t, 3>>>>
        opUsers;
    for (auto &sop : op.getCond().front().without_terminator()) {
      auto setP = std::make_unique<llvm::SmallPtrSet<size_t, 3>>();
      auto &set = *setP;
      for (auto operand : sop.getOperands()) {
        auto found = users.find(operand);
        if (found != users.end()) {
          set.insert(found->second->begin(), found->second->end());
        }
      }
      if (!(sop.getNumRegions() == 0 ||
            sop.hasTrait<OpTrait::IsIsolatedFromAbove>())) {
        SmallVector<Operation *> todo;
        for (auto &reg : sop.getRegions()) {
          for (auto &blk : reg) {
            for (auto &mop : blk) {
              todo.push_back(&mop);
            }
          }
        }
        while (todo.size()) {
          auto cur = todo.pop_back_val();
          for (auto operand : cur->getOperands()) {
            auto found = users.find(operand);
            if (found != users.end()) {
              set.insert(found->second->begin(), found->second->end());
            }
          }
          if (!(cur->getNumRegions() == 0 ||
                cur->hasTrait<OpTrait::IsIsolatedFromAbove>())) {
            for (auto &reg : cur->getRegions()) {
              for (auto &blk : reg) {
                for (auto &mop : blk) {
                  todo.push_back(&mop);
                }
              }
            }
          }
        }
      }
      if (set.size() == 0)
        continue;

      if (!isTotalPure && !mlir::isPure(&sop)) {
        for (auto arg : set)
          nonPure.insert(arg);
      }
      for (auto res : sop.getResults())
        if (!res.use_empty())
          users.try_emplace(res, &set);
      opUsers.emplace_back(&sop, std::move(setP));
    }
    llvm::SmallPtrSet<size_t, 3> terminatorUsers;
    for (auto v : op.getCond().front().getTerminator()->getOperands()) {
      auto found = users.find(v);
      if (found != users.end()) {
        terminatorUsers.insert(found->second->begin(), found->second->end());
      }
    }

    SmallVector<size_t> emptyNonPure;
    for (auto ores : emptyResults) {
      if (nonPure.count(ores)) {
        continue;
      }
      if (terminatorUsers.count(ores)) {
        continue;
      }
      emptyNonPure.push_back(ores);
    }
    if (emptyNonPure.size() == 0) {
      return failure();
    }

    DenseMap<Value, llvm::SmallPtrSet<size_t, 3> *> usersBody(
        emptyNonPure.size() + op.getBody().front().getOperations().size());
    SmallVector<
        std::pair<Operation *, std::unique_ptr<llvm::SmallPtrSet<size_t, 3>>>>
        opUsersBody;
    terminatorUsers.clear();
    nonPure.clear();

    llvm::SmallVector<std::unique_ptr<llvm::SmallPtrSet<size_t, 3>>, 3>
        bodyCache;
    bodyCache.reserve(emptyNonPure.size());
    for (auto ores : emptyNonPure) {
      auto ba = op.getBody().front().getArgument(ores);
      auto up = std::make_unique<llvm::SmallPtrSet<size_t, 3>>();
      up->insert(ores);
      usersBody[ba] = up.get();
      bodyCache.emplace_back(std::move(up));
    }

    llvm::SmallSet<size_t, 3> nonPure2;

    for (auto &sop : op.getBody().front().without_terminator()) {
      auto setP = std::make_unique<llvm::SmallPtrSet<size_t, 3>>();
      auto &set = *setP;
      for (auto operand : sop.getOperands()) {
        auto found = usersBody.find(operand);
        if (found != usersBody.end()) {
          set.insert(found->second->begin(), found->second->end());
        }
      }
      if (!(sop.getNumRegions() == 0 ||
            sop.hasTrait<OpTrait::IsIsolatedFromAbove>())) {
        SmallVector<Operation *> todo;
        for (auto &reg : sop.getRegions()) {
          for (auto &blk : reg) {
            for (auto &mop : blk) {
              todo.push_back(&mop);
            }
          }
        }
        while (todo.size()) {
          auto cur = todo.pop_back_val();
          for (auto operand : cur->getOperands()) {
            auto found = usersBody.find(operand);
            if (found != usersBody.end()) {
              set.insert(found->second->begin(), found->second->end());
            }
          }
          if (!(cur->getNumRegions() == 0 ||
                cur->hasTrait<OpTrait::IsIsolatedFromAbove>())) {
            for (auto &reg : cur->getRegions()) {
              for (auto &blk : reg) {
                for (auto &mop : blk) {
                  todo.push_back(&mop);
                }
              }
            }
          }
        }
      }
      if (set.size() == 0)
        continue;

      if (!isTotalPure && !mlir::isPure(&sop)) {
        for (auto arg : set)
          nonPure2.insert(arg);
      }
      for (auto res : sop.getResults())
        if (!res.use_empty())
          usersBody.try_emplace(res, &set);
      opUsersBody.emplace_back(&sop, std::move(setP));
    }

    llvm::SmallPtrSet<size_t, 3> emptyNonPure2;
    if (isTotalPure)
      emptyNonPure2.insert(emptyNonPure.begin(), emptyNonPure.end());
    else {
      for (auto ores : emptyNonPure) {
        if (nonPure2.count(ores)) {
          continue;
        }
        emptyNonPure2.insert(ores);
      }
      if (emptyNonPure2.size() == 0) {
        return failure();
      }
    }

    llvm::BitVector removedResults(NumResults, true);
    llvm::BitVector seen(NumResults, false);
    SmallVector<size_t> todo;
    todo.reserve(NumResults);
    for (size_t residx = 0; residx < NumResults; residx++) {
      if (!emptyNonPure2.count(residx)) {
        todo.push_back(residx);
        seen.set(residx);
      }
    }

    while (todo.size()) {
      auto cur = todo.pop_back_val();
      if (!removedResults[cur])
        continue;
      removedResults.reset(cur);
      auto v = op.getBody().front().getTerminator()->getOperands()[cur];
      auto found = usersBody.find(v);
      if (found != usersBody.end()) {
        if (isTotalPure) {
          for (auto arg : *found->second) {
            if (!seen.test(arg)) {
              todo.push_back(arg);
              seen.set(arg);
            }
          }
        } else {
          for (auto arg : *found->second) {
            if (emptyNonPure2.contains(arg)) {
              if (!seen.test(arg)) {
                todo.push_back(arg);
                seen.set(arg);
              }
            }
          }
        }
      }
    }
    /*
    llvm::BitVector todo(NumResults, false);
    int start = NumResults;
    for (size_t residx = 0; residx < NumResults; residx++) {
      if (!emptyNonPure2.count(residx)) {
        todo.set(residx);
        if (residx < start) start = residx;
      }
    }

    while (true) {
      auto cur = todo.find_first_in(start, NumResults);
      if (cur == -1) break;
      start = cur+1;
      todo.reset(cur);
      if (!removedResults.test(cur))
        continue;
      removedResults.reset(cur);
      auto v = op.getBody().front().getTerminator()->getOperands()[cur];
      auto found = usersBody.find(v);
      if (found != usersBody.end()) {
        if (isTotalPure) {
          for (auto arg : found->second) {
            if (!removedResults.test(arg))
              continue;
            todo.set(arg);
            if (arg < start) start = arg;
          }
        } else {
          for (auto arg : found->second) {
            if (emptyNonPure2.contains(arg)) {
              if (!removedResults.test(arg))
                continue;
              todo.set(arg);
              if (arg < start) start = arg;
            }
          }
        }
      }
    }
    */
    /*
    SmallVector<size_t> todo;
    for (size_t residx = 0; residx < NumResults; residx++) {
      if (!emptyNonPure2.count(residx))
        todo.push_back(residx);
    }

    while (todo.size()) {
      auto cur = todo.pop_back_val();
      if (!removedResults[cur])
        continue;
      removedResults.reset(cur);
      auto v = op.getBody().front().getTerminator()->getOperands()[cur];
      auto found = usersBody.find(v);
      if (found != usersBody.end()) {
        if (isTotalPure)
          todo.append(found->second.begin(), found->second.end());
        else
          for (auto arg : found->second) {
            if (emptyNonPure2.contains(arg))
              todo.push_back(arg);
        }
      }
    }
    */

    if (!removedResults.any())
      return failure();

    SmallVector<int64_t> deadResults;
    for (auto cur = removedResults.find_first(); cur != -1;
         cur = removedResults.find_next(cur)) {
      deadResults.push_back(cur);
    }

    replaceTerminator(rewriter, op.getBody(), deadResults);

    for (auto &cop : llvm::reverse(opUsers)) {
      bool hasDead = false;
      for (int64_t i : deadResults) {
        if (cop.second->count(i)) {
          hasDead = true;
          break;
        }
      }
      if (hasDead)
        rewriter.eraseOp(cop.first);
    }
    for (auto &cop : llvm::reverse(opUsersBody)) {
      bool hasDead = false;
      for (int64_t i : deadResults) {
        if (cop.second->count(i)) {
          hasDead = true;
          break;
        }
      }
      if (hasDead)
        rewriter.eraseOp(cop.first);
    }

    SmallVector<Value> operands;
    SmallVector<Type> resultTypes;
    SmallVector<Location> condBlockArgLocs, bodyBlockArgsLocs;
    for (auto &&[i, operand, resultType] :
         llvm::enumerate(op->getOperands(), op.getResultTypes())) {
      if (llvm::is_contained(deadResults, i))
        continue;

      operands.push_back(operand);
      resultTypes.push_back(resultType);
      condBlockArgLocs.push_back(op.getCond().getArgument(i).getLoc());
      bodyBlockArgsLocs.push_back(op.getBody().getArgument(i).getLoc());
    }

    auto updated = stablehlo::WhileOp::create(
        rewriter, op->getLoc(), resultTypes, operands, op->getAttrs());
    SmallVector<Value> resultReplacements;
    for (int64_t old = 0, upd = 0, end = op->getNumResults(); old < end;
         ++old) {
      if (llvm::is_contained(deadResults, old)) {
        resultReplacements.push_back(nullptr);
        continue;
      }
      resultReplacements.push_back(updated->getResult(upd));
      ++upd;
    }

    for (int64_t i : llvm::reverse(deadResults))
      op.getCond().eraseArgument(i);
    rewriter.inlineRegionBefore(op.getCond(), updated.getCond(),
                                updated.getCond().begin());

    for (int64_t i : llvm::reverse(deadResults)) {
      op.getBody().eraseArgument(i);
    }
    rewriter.inlineRegionBefore(op.getBody(), updated.getBody(),
                                updated.getBody().begin());

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

struct NegativePadToSlice final
    : CheckedOpRewritePattern<stablehlo::PadOp, NegativePadToSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::PadOp op,
                                    PatternRewriter &rewriter) const {
    SmallVector<int64_t> starts;
    SmallVector<int64_t> limits;
    SmallVector<int64_t> strides;

    bool negative = false;
    for (auto &&[low, high, inner, dim] : llvm::zip(
             op.getEdgePaddingLow(), op.getEdgePaddingHigh(),
             op.getInteriorPadding(), op.getOperand().getType().getShape())) {
      if (low > 0)
        return failure();
      if (high > 0)
        return failure();
      if (inner != 0)
        return failure();
      if (low < 0 || high < 0)
        negative = true;

      starts.push_back(-low);
      limits.push_back(dim + high);
      strides.push_back(1);
    }
    if (!negative)
      return failure();
    rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(op, op.getOperand(), starts,
                                                    limits, strides);
    return success();
  }
};

/*
%1192 = stablehlo.pad %1189, %cst_0, low = [0], high = [1], interior = [0] :
  (tensor<1xf32>, tensor<f32>) -> tensor<2xf32>

%1193 = arith.addf %1191, %1192 : tensor<2xf32>
*/
template <typename T>
struct BinopPadToConcat final
    : CheckedOpRewritePattern<T, BinopPadToConcat<T>> {
  using CheckedOpRewritePattern<T,
                                BinopPadToConcat<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    for (int i = 0; i < 2; i++) {
      if (auto lhs =
              op->getOperand(i).template getDefiningOp<stablehlo::PadOp>()) {
        if (anyPadSizesNegative(lhs))
          continue;
        auto rhs = op->getOperand(1 - i);

        Attribute padValAttr;
        matchPattern(lhs.getPaddingValue(), m_Constant(&padValAttr));
        bool padIsZero =
            padValAttr && matchPattern(padValAttr, m_AnyZeroFloat());
        bool padIsOne = padValAttr && matchPattern(padValAttr, m_OneFloat());

        bool match = false;
        if (isa<stablehlo::AddOp>(op)) {
          match = padIsZero;
        } else if (isa<stablehlo::MulOp>(op)) {
          match = padIsOne || padIsZero;
        }
        if (!match) {
          SmallVector<Operation *> ops = {op};
          bool legal = true;
          while (!ops.empty()) {
            auto cur = ops.pop_back_val();
            if (isa<stablehlo::SliceOp>(cur))
              continue;
            if (isa<stablehlo::AddOp, stablehlo::MulOp>(cur)) {
              for (auto u : cur->getResult(0).getUsers()) {
                ops.push_back(u);
              }
              continue;
            }
            legal = false;
            break;
          }
          if (!legal)
            return failure();
        }

        bool legal = true;
        for (auto step : lhs.getInteriorPadding()) {
          if (step != 0) {
            legal = false;
            break;
          }
        }
        if (!legal)
          continue;

        ssize_t padidx = -1;

        SmallVector<size_t> idxs;
        for (auto &&[low, high, dim] :
             llvm::zip(lhs.getEdgePaddingLow(), lhs.getEdgePaddingHigh(),
                       type.getShape())) {
          padidx++;
          if (low == 0 && high == 0)
            continue;
          if (low < 0 || high < 0)
            return failure();
          idxs.push_back(padidx);
        }

        if (idxs.size() == 1 &&
            lhs.getOperand().getType().getShape()[idxs[0]] * 2 <=
                type.getShape()[idxs[0]]) {
          auto idx = idxs[0];

          SmallVector<int64_t> strides(type.getShape().size(), 1);
          SmallVector<int64_t> starts(type.getShape().size(), 0);
          SmallVector<int64_t> limits(type.getShape().begin(),
                                      type.getShape().end());

          SmallVector<Value, 1> vals;

          if (lhs.getEdgePaddingLow()[idx] != 0) {
            starts[idx] = 0;
            limits[idx] = lhs.getEdgePaddingLow()[idx];
            Value prevSlice = stablehlo::SliceOp::create(
                rewriter, op.getLoc(), rhs, starts, limits, strides);

            if (isa<stablehlo::AddOp>(op) && padIsZero) {
              // If adding we're adding 0, no need to do extra work
            } else if (isa<stablehlo::MulOp>(op) && padIsZero) {
              // If multiplying by 0, broadcast the zero
              prevSlice = stablehlo::BroadcastInDimOp::create(
                  rewriter, op.getLoc(), prevSlice.getType(),
                  lhs.getPaddingValue(), ArrayRef<int64_t>());
            } else if (isa<stablehlo::MulOp>(op) && padIsOne) {
              // If multiplying by 1, no need to do extra work
            } else
              prevSlice =
                  T::create(rewriter, op.getLoc(), prevSlice,
                            stablehlo::BroadcastInDimOp::create(
                                rewriter, op.getLoc(), prevSlice.getType(),
                                lhs.getPaddingValue(), ArrayRef<int64_t>()));
            vals.push_back(prevSlice);
          }

          starts[idx] = lhs.getEdgePaddingLow()[idx];
          limits[idx] = type.getShape()[idx] - lhs.getEdgePaddingHigh()[idx];

          auto midSlice = stablehlo::SliceOp::create(rewriter, op.getLoc(), rhs,
                                                     starts, limits, strides);
          auto mid =
              T::create(rewriter, op.getLoc(), midSlice, lhs.getOperand());
          vals.push_back(mid);

          if (lhs.getEdgePaddingHigh()[idx] != 0) {
            starts[idx] = type.getShape()[idx] - lhs.getEdgePaddingHigh()[idx];
            limits[idx] = type.getShape()[idx];
            Value postSlice = stablehlo::SliceOp::create(
                rewriter, op.getLoc(), rhs, starts, limits, strides);

            if (isa<stablehlo::AddOp>(op) && padIsZero) {
              // If adding we're adding 0, no need to do extra work
            } else if (isa<stablehlo::MulOp>(op) && padIsZero) {
              // If multiplying by 0, broadcast the zero
              postSlice = stablehlo::BroadcastInDimOp::create(
                  rewriter, op.getLoc(), postSlice.getType(),
                  lhs.getPaddingValue(), ArrayRef<int64_t>());
            } else if (isa<stablehlo::MulOp>(op) && padIsOne) {
              // If multiplying by 1, no need to do extra work
            } else
              postSlice =
                  T::create(rewriter, op.getLoc(), postSlice,
                            stablehlo::BroadcastInDimOp::create(
                                rewriter, op.getLoc(), postSlice.getType(),
                                lhs.getPaddingValue(), ArrayRef<int64_t>()));
            vals.push_back(postSlice);
          }

          rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, vals, idx);
          return success();
        }
      }
    }

    return failure();
  }
};

struct ReshapeIota final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeIota> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto iota = op.getOperand().getDefiningOp<stablehlo::IotaOp>();
    if (!iota)
      return failure();

    // A scalar (rank-0) IotaOp is invalid in stableHLO. Bail out and let
    // other passes (e.g. constant folding) handle this case.
    if (op.getType().getRank() == 0)
      return failure();

    size_t curiotaidx = 0;
    size_t iotadim = 0;
    for (auto en : llvm::enumerate(op.getType().getShape())) {
      if (en.value() == 1)
        continue;

      if (curiotaidx == iota.getType().getShape().size())
        return failure();
      auto ival = iota.getType().getShape()[curiotaidx];
      while (ival == 1 && curiotaidx < iota.getType().getShape().size()) {
        if (curiotaidx == iota.getIotaDimension()) {
          return failure();
        }
        curiotaidx++;
        ival = iota.getType().getShape()[curiotaidx];
      }
      if (en.value() == ival) {
        if (curiotaidx == iota.getIotaDimension()) {
          iotadim = en.index();
        }
        curiotaidx++;
        continue;
      }
      return failure();
    }
    rewriter.replaceOpWithNewOp<stablehlo::IotaOp>(op, op.getType(), iotadim);
    return success();
  }
};

struct BroadcastIota final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp, BroadcastIota> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    auto iota = op.getOperand().getDefiningOp<stablehlo::IotaOp>();
    if (!iota)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::IotaOp>(
        op, op.getType(), op.getBroadcastDimensions()[iota.getIotaDimension()]);
    return success();
  }
};

LogicalResult reshapePadHelper(stablehlo::ReshapeOp op,
                               PatternRewriter &rewriter) {
  auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
  if (!pad)
    return failure();
  if (anyPadSizesNegative(pad))
    return failure();
  size_t curiotaidx = 0;
  SmallVector<int64_t> lows;
  SmallVector<int64_t> highs;
  SmallVector<int64_t> interiors;

  SmallVector<int64_t> inner_shape;
  for (auto en : llvm::enumerate(op.getType().getShape())) {
    if (en.value() == 1) {
      lows.push_back(0);
      highs.push_back(0);
      interiors.push_back(0);
      inner_shape.push_back(1);
      continue;
    }

    if (curiotaidx == pad.getType().getShape().size())
      return failure();

    auto ival = pad.getType().getShape()[curiotaidx];
    while (ival == 1 && curiotaidx < pad.getType().getShape().size()) {
      assert(pad.getEdgePaddingLow()[curiotaidx] == 0);
      assert(pad.getEdgePaddingHigh()[curiotaidx] == 0);
      assert(pad.getInteriorPadding()[curiotaidx] == 0);
      curiotaidx++;
      ival = pad.getType().getShape()[curiotaidx];
    }
    if (en.value() == ival) {
      lows.push_back(pad.getEdgePaddingLow()[curiotaidx]);
      highs.push_back(pad.getEdgePaddingHigh()[curiotaidx]);
      interiors.push_back(pad.getInteriorPadding()[curiotaidx]);
      inner_shape.push_back(pad.getOperand().getType().getShape()[curiotaidx]);
      curiotaidx++;
      continue;
    }
    return failure();
  }
  auto inner = stablehlo::ReshapeOpCreate(rewriter, op.getLoc(),
                                          pad.getOperand(), inner_shape);
  rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
      op, inner, pad.getPaddingValue(), lows, highs, interiors);
  return success();
}

struct DotReshapePad final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, DotReshapePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    if (!llvm::hasSingleElement(pad->getUsers()))
      return failure();

    for (auto u : op->getUsers())
      if (!isa<stablehlo::DotGeneralOp>(u))
        return failure();

    if (!reshapePadHelper(op, rewriter).succeeded())
      return failure();
    return success();
  }
};

struct ZeroProductReshapePad final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ZeroProductReshapePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    if (!llvm::hasSingleElement(pad->getUsers()))
      return failure();

    if (!matchPattern(pad.getPaddingValue(), m_AnyZeroFloat()))
      return failure();

    for (auto u : op->getUsers()) {
      if (!isa<stablehlo::MulOp>(u) && !isa<stablehlo::DivOp>(u))
        return failure();
    }
    if (!reshapePadHelper(op, rewriter).succeeded())
      return failure();
    return success();
  }
};

struct PadReshapePad final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, PadReshapePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    if (!llvm::hasSingleElement(pad->getUsers()))
      return failure();

    for (auto u : op->getUsers()) {
      auto pad2 = dyn_cast<stablehlo::PadOp>(u);
      if (!pad2)
        return failure();
      if (pad2.getPaddingValue() != pad.getPaddingValue())
        return failure();
    }
    if (!reshapePadHelper(op, rewriter).succeeded())
      return failure();
    return success();
  }
};

struct BinopConstReshapePad final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, BinopConstReshapePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    if (!llvm::hasSingleElement(pad->getUsers()))
      return failure();

    for (auto u : op->getUsers()) {
      if (isa<stablehlo::AddOp>(u) || isa<stablehlo::SubtractOp>(u) ||
          isa<stablehlo::MulOp>(u) || isa<stablehlo::DivOp>(u)) {
        bool hasConst = false;
        for (auto op : u->getOperands())
          hasConst |= op.getDefiningOp<stablehlo::ConstantOp>() != nullptr;
        if (hasConst)
          continue;
      }
      return failure();
    }
    if (!reshapePadHelper(op, rewriter).succeeded())
      return failure();
    return success();
  }
};

struct ConcatAppendingReshape final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                              ConcatAppendingReshape> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() != 2)
      return failure();

    SmallVector<Value> lhs;

    SmallVector<Type> converts;

    size_t frontSize = 0;
    for (auto v : op.getOperands()) {
      if (auto t = v.getDefiningOp<stablehlo::ConvertOp>()) {
        v = t.getOperand();
        converts.push_back(
            cast<RankedTensorType>(v.getType()).getElementType());
      } else
        converts.push_back(nullptr);
      if (auto t = v.getDefiningOp<stablehlo::ReshapeOp>()) {
        lhs.push_back(t->getOperand(0));

        auto prevshape = t.getOperand().getType().getShape();
        auto postshape = t.getType().getShape();
        if (prevshape.size() == 0)
          return failure();
        if (prevshape.size() + 1 != postshape.size())
          return failure();
        if (postshape[0] != 1)
          return failure();

        frontSize += prevshape[0];

        for (auto en : llvm::enumerate(prevshape)) {
          if (en.value() != postshape[1 + en.index()])
            return failure();
        }

      } else
        return failure();
    }

    Type typeconvert = converts[0];
    for (auto c : converts)
      if (c != typeconvert)
        return failure();

    RankedTensorType nextType = op.getType();
    auto nextDim = op.getDimension();
    if (nextDim == 0) {
      SmallVector<int64_t> nextShape(nextType.getShape().begin() + 1,
                                     nextType.getShape().end());

      nextShape[0] = frontSize;
      nextType = RankedTensorType::get(
          nextShape, typeconvert ? typeconvert : nextType.getElementType());
      nextDim = 0;
    } else {
      nextType = RankedTensorType::get(nextType.getShape().drop_front(),
                                       typeconvert ? typeconvert
                                                   : nextType.getElementType());
      nextDim--;
    }
    auto lhs2 = stablehlo::ConcatenateOp::create(rewriter, op.getLoc(),
                                                 nextType, lhs, nextDim);

    Value res2 = stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), lhs2,
                                            op.getType().getShape());

    if (typeconvert)
      res2 = stablehlo::ConvertOp::create(rewriter, op.getLoc(), op.getType(),
                                          res2);

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

template <typename T>
struct UnaryPadPush final : CheckedOpRewritePattern<T, UnaryPadPush<T>> {
  using CheckedOpRewritePattern<T, UnaryPadPush<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    auto pad = op->getOperand(0).template getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();
    if (anyPadSizesNegative(pad))
      return failure();

    auto padval = pad.getPaddingValue();
    auto padval2 = T::create(
        rewriter, op.getLoc(),
        RankedTensorType::get({}, op.getType().getElementType()), padval);

    auto val = pad.getOperand();
    auto val2 = T::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(cast<RankedTensorType>(val.getType()).getShape(),
                              op.getType().getElementType()),
        val);

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, val2, padval2, pad.getEdgePaddingLow(), pad.getEdgePaddingHigh(),
        pad.getInteriorPadding());
    return success();
  }
};

struct TransposePad final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto pad = op->getOperand(0).template getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();
    if (anyPadSizesNegative(pad))
      return failure();

    if (!llvm::hasSingleElement(pad->getUsers()))
      return failure();

    auto padval = pad.getPaddingValue();

    auto val = pad.getOperand();
    auto val2 = stablehlo::TransposeOp::create(rewriter, op.getLoc(), val,
                                               op.getPermutation());

    SmallVector<int64_t> low;
    SmallVector<int64_t> high;
    SmallVector<int64_t> inner;
    for (auto idx : op.getPermutation()) {
      low.push_back(pad.getEdgePaddingLow()[idx]);
      high.push_back(pad.getEdgePaddingHigh()[idx]);
      inner.push_back(pad.getInteriorPadding()[idx]);
    }

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(op, val2, padval, low, high,
                                                  inner);
    rewriter.eraseOp(pad);
    return success();
  }
};

template <typename T>
struct ConcatPushBinop final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatPushBinop<T>> {
  using CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                ConcatPushBinop<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() != 2)
      return failure();

    SmallVector<Value> lhs;
    SmallVector<Value> rhs;

    SmallVector<Type> converts;

    for (auto v : op.getOperands()) {
      if (auto t = v.getDefiningOp<stablehlo::ConvertOp>()) {
        converts.push_back(
            cast<RankedTensorType>(t.getType()).getElementType());
        v = t.getOperand();
      } else
        converts.push_back(nullptr);
      if (auto t = v.getDefiningOp<T>()) {
        lhs.push_back(t->getOperand(0));
        rhs.push_back(t->getOperand(1));
      } else
        return failure();
    }

    Type typeconvert = converts[0];
    for (auto c : converts)
      if (c != typeconvert)
        return failure();

    auto lhs2 = stablehlo::ConcatenateOp::create(rewriter, op.getLoc(), lhs,
                                                 op.getDimension());
    auto rhs2 = stablehlo::ConcatenateOp::create(rewriter, op.getLoc(), rhs,
                                                 op.getDimension());

    Value res2 = T::create(rewriter, op.getLoc(), lhs2, rhs2);

    if (typeconvert)
      res2 = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(),
          RankedTensorType::get(
              cast<RankedTensorType>(res2.getType()).getShape(), typeconvert),
          res2);

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

struct ConcatFuse final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatFuse> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() == 1 &&
        op->getOperand(0).getType() == op.getType()) {
      rewriter.replaceOp(op, op->getOperand(0));
      return success();
    }
    SmallVector<Value> vals;
    bool changed = false;
    for (auto v : op->getOperands()) {
      if (auto c2 = v.getDefiningOp<stablehlo::ConcatenateOp>()) {
        if (c2.getDimension() == op.getDimension()) {
          for (auto v2 : c2->getOperands())
            vals.push_back(v2);
          changed = true;
          continue;
        }
      }
      if (cast<RankedTensorType>(v.getType()).getShape()[op.getDimension()] ==
          0) {
        changed = true;
        continue;
      }
      vals.push_back(v);
    }
    if (!changed)
      return failure();
    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
        op, op.getType(), vals, op.getDimensionAttr());
    return success();
  }
};

// concat(slice(src,[N-1:N,...]), slice(src,[N-2:N-1,...]), ...,
// slice(src,[0:1,...]), dim=D)
// -> reverse(src, dims=[D])
struct ConcatSlicesToReverse final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatSlicesToReverse> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto concatDim = op.getDimension();
    auto operands = op.getOperands();
    int64_t N = (int64_t)operands.size();

    if (N == 0)
      return failure();

    auto firstSlice = operands[0].getDefiningOp<stablehlo::SliceOp>();
    if (!firstSlice)
      return failure();

    Value src = firstSlice.getOperand();
    auto srcType = cast<RankedTensorType>(src.getType());
    int64_t rank = srcType.getRank();

    if (srcType.getShape()[concatDim] != N)
      return failure();

    for (int64_t i = 0; i < N; ++i) {
      auto slice = operands[i].getDefiningOp<stablehlo::SliceOp>();
      if (!slice || slice.getOperand() != src)
        return failure();
      auto starts = slice.getStartIndices();
      auto limits = slice.getLimitIndices();
      auto strides = slice.getStrides();
      for (int64_t d = 0; d < rank; ++d) {
        if (strides[d] != 1)
          return failure();
        if (d == (int64_t)concatDim) {
          if (starts[d] != N - 1 - i || limits[d] != N - i)
            return failure();
        } else {
          if (starts[d] != 0 || limits[d] != srcType.getShape()[d])
            return failure();
        }
      }
    }

    rewriter.replaceOpWithNewOp<stablehlo::ReverseOp>(
        op, src, SmallVector<int64_t>{(int64_t)concatDim});
    return success();
  }
};

struct ConcatToBroadcast final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatToBroadcast> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() <= 1)
      return failure();
    for (auto opv : op->getOperands())
      if (opv != op->getOperand(0))
        return failure();
    auto InTy = cast<RankedTensorType>(op->getOperand(0).getType());
    if (InTy.getShape()[op.getDimension()] == 1) {
      SmallVector<int64_t> bcast;
      for (auto en : llvm::enumerate(op.getType().getShape())) {
        bcast.push_back(en.index());
      }
      auto bcast2 = rewriter.getDenseI64ArrayAttr(bcast);
      rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
          op, op.getType(), op->getOperand(0), bcast2);
      return success();
    } else {
      SmallVector<int64_t> reshaped = llvm::to_vector(InTy.getShape());
      reshaped.insert(reshaped.begin() + op.getDimension(), 1);
      auto reshapeVal = stablehlo::ReshapeOpCreate(rewriter, op.getLoc(),
                                                   op->getOperand(0), reshaped);

      SmallVector<int64_t> bcast;
      for (auto en : llvm::enumerate(
               cast<RankedTensorType>(reshapeVal.getType()).getShape())) {
        bcast.push_back(en.index());
      }
      reshaped[op.getDimension()] = op->getNumOperands();
      auto bcast2 = rewriter.getDenseI64ArrayAttr(bcast);
      auto bcastVal = stablehlo::BroadcastInDimOp::create(
          rewriter, op.getLoc(),
          RankedTensorType::get(reshaped, op.getType().getElementType()),
          reshapeVal, bcast2);

      auto newReshapeVal =
          stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), bcastVal,
                                     cast<ShapedType>(op.getType()).getShape());
      rewriter.replaceOp(op, newReshapeVal);
      return success();
    }
  }
};

struct GammaConstProp final
    : CheckedOpRewritePattern<chlo::LgammaOp, GammaConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::LgammaOp op,
                                    PatternRewriter &rewriter) const {
    // return if not constant
    DenseElementsAttr inputAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&inputAttr)))
      return failure();
    Value result =
        stablehlo::materializeLgamma(rewriter, op.getLoc(), op->getOperands());
    rewriter.replaceOp(op, result);

    return success();
  }
};

struct TGammaConstProp final
    : CheckedOpRewritePattern<enzymexla::TGammaOp, TGammaConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::TGammaOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr inputAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&inputAttr)))
      return failure();

    auto resultType = cast<ShapedType>(op.getType());
    auto floatTy = dyn_cast<FloatType>(resultType.getElementType());
    if (!floatTy)
      return failure();

    const auto &sem = floatTy.getFloatSemantics();

    auto computeTGamma = [&](APFloat val) -> APFloat {
      double x = val.convertToDouble();
      double res =
          (x < 0.0) ? std::numeric_limits<double>::quiet_NaN() : std::tgamma(x);
      bool losesInfo;
      APFloat apRes(res);
      apRes.convert(sem, APFloat::rmNearestTiesToEven, &losesInfo);
      return apRes;
    };

    if (inputAttr.isSplat()) {
      APFloat result = computeTGamma(inputAttr.getSplatValue<APFloat>());
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, DenseElementsAttr::get(resultType, result));
      return success();
    }

    SmallVector<APFloat> results;
    for (auto val : inputAttr.getValues<APFloat>())
      results.push_back(computeTGamma(val));

    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, DenseElementsAttr::get(resultType, results));
    return success();
  }
};

struct LGammaConstProp final
    : CheckedOpRewritePattern<enzymexla::LGammaOp, LGammaConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::LGammaOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr inputAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&inputAttr)))
      return failure();

    auto resultType = cast<ShapedType>(op.getType());
    auto floatTy = dyn_cast<FloatType>(resultType.getElementType());
    if (!floatTy)
      return failure();

    const auto &sem = floatTy.getFloatSemantics();

    auto computeLGamma = [&](APFloat val) -> APFloat {
      double x = val.convertToDouble();
      double res = std::lgamma(x);
      bool losesInfo;
      APFloat apRes(res);
      apRes.convert(sem, APFloat::rmNearestTiesToEven, &losesInfo);
      return apRes;
    };

    if (inputAttr.isSplat()) {
      APFloat result = computeLGamma(inputAttr.getSplatValue<APFloat>());
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, DenseElementsAttr::get(resultType, result));
      return success();
    }

    SmallVector<APFloat> results;
    for (auto val : inputAttr.getValues<APFloat>())
      results.push_back(computeLGamma(val));

    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, DenseElementsAttr::get(resultType, results));
    return success();
  }
};

struct BinomialProgressConstProp final
    : CheckedOpRewritePattern<enzyme::BinomialProgressOp,
                              BinomialProgressConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  // The Revolve advance *distance*: with beta(s, t) = C(s + t, t) and t minimal
  // such that beta(s, t) >= n, every advance in
  // [n - beta(s-1, t), beta(s, t-1)] attains that optimal t. Clamp that window
  // to [1, n-1] and take its midpoint. Must stay in sync with
  // enzyme/Enzyme/MLIR/Dialect/Ops.cpp's binomialProgress and with
  // LowerBinomialProgressOpToStableHLO.
  //
  // With a single checkpoint left there is nothing better than replaying the
  // whole remaining stretch from it, so advance all of it. Together with every
  // other advance landing in [1, n-1], that is what makes the per-slot advances
  // sum to exactly `n` across `budget` slots -- which is what the callers'
  // outer loop relies on to reach the end of the primal.
  static int64_t binomialProgress(int64_t n, int64_t s) {
    if (n <= 0)
      return 0;
    if (n == 1)
      return 1;
    if (s <= 1)
      return n;

    int64_t t = 0, beta = 1; // beta == C(s + t, t)
    while (beta < n) {
      ++t;
      beta = beta * (s + t) / t; // C(s+t,t) from C(s+t-1,t-1); exact
    }

    // beta(s-1, t) == beta * s / (s + t) and beta(s, t-1) == beta * t / (s +
    // t), both exact in integers.
    int64_t lo = std::max<int64_t>(n - beta * s / (s + t), 1);
    int64_t hi = std::min<int64_t>(beta * t / (s + t), n - 1);
    int64_t m = (lo + hi) / 2; // lo <= hi, so this is already in [1, n-1]

    // Leave at least one step for each of the s-1 checkpoints still to be
    // placed. Without this the advances can exhaust the interval before the
    // slots run out, and a caller that walks one slot per iteration then
    // records slots at a step past the end -- holding the final state rather
    // than a checkpoint.
    m = std::min<int64_t>(m, n - (s - 1));
    return std::max<int64_t>(m, 1);
  }

  LogicalResult matchAndRewriteImpl(enzyme::BinomialProgressOp op,
                                    PatternRewriter &rewriter) const {
    APInt numSteps, budget;

    // The op is also used on plain scalars by the scf path, where a
    // stablehlo.constant replacement would be ill-typed; Enzyme's own folder
    // handles those.
    if (!isa<TensorType>(op.getType()))
      return failure();

    if (!matchPattern(op.getNumSteps(), m_ConstantInt(&numSteps)) ||
        !matchPattern(op.getBudget(), m_ConstantInt(&budget)))
      return failure();

    if (budget.isNonPositive() || numSteps.isNonPositive())
      return failure();

    int64_t progress =
        binomialProgress(numSteps.getSExtValue(), budget.getSExtValue());

    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, cast<ElementsAttr>(makeAttr(op.getResult().getType(), progress)));

    return success();
  }
};

struct CHLOLGammaConstProp final
    : CheckedOpRewritePattern<chlo::LgammaOp, CHLOLGammaConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::LgammaOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr inputAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&inputAttr)))
      return failure();

    auto resultType = cast<ShapedType>(op.getType());
    auto floatTy = dyn_cast<FloatType>(resultType.getElementType());
    if (!floatTy)
      return failure();

    const auto &sem = floatTy.getFloatSemantics();

    auto computeLGamma = [&](APFloat val) -> APFloat {
      double x = val.convertToDouble();
      double res = std::lgamma(x);
      bool losesInfo;
      APFloat apRes(res);
      apRes.convert(sem, APFloat::rmNearestTiesToEven, &losesInfo);
      return apRes;
    };

    if (inputAttr.isSplat()) {
      APFloat result = computeLGamma(inputAttr.getSplatValue<APFloat>());
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, DenseElementsAttr::get(resultType, result));
      return success();
    }

    SmallVector<APFloat> results;
    for (auto val : inputAttr.getValues<APFloat>())
      results.push_back(computeLGamma(val));

    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, DenseElementsAttr::get(resultType, results));
    return success();
  }
};

struct DynamicUpdateSliceConstProp final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                              DynamicUpdateSliceConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;
  size_t max_constant_expansion;

  DynamicUpdateSliceConstProp(size_t max_constant_expansion,
                              MLIRContext *context, PatternBenefit benefit = 1,
                              ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp op,
                                    PatternRewriter &rewriter) const {

    size_t size = 1;
    for (auto sz : op.getType().getShape())
      size *= sz;
    if (size >= max_constant_expansion)
      return failure();

    auto startIndices = op.getStartIndices();

    bool legal = true;

    DenseElementsAttr operandConstant;
    DenseElementsAttr updateConstant;

    SmallVector<DenseElementsAttr> constants(startIndices.size(),
                                             DenseElementsAttr());
    for (auto &operand : op->getOpOperands()) {
      if (operand.getOperandNumber() == 0)
        legal &= matchPattern(operand.get(), m_Constant(&operandConstant));
      else if (operand.getOperandNumber() == 1)
        legal &= matchPattern(operand.get(), m_Constant(&updateConstant));
      else
        legal &= matchPattern(
            operand.get(),
            m_Constant(&constants[operand.getOperandNumber() - 2]));
    }

    if (operandConstant && updateConstant && operandConstant.isSplat() &&
        updateConstant.isSplat() &&
        operandConstant.getSplatValue<Attribute>() ==
            updateConstant.getSplatValue<Attribute>()) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }

    if (!legal)
      return failure();

    stablehlo::Tensor operandTen = stablehlo::constantOp(operandConstant);
    stablehlo::Tensor updateTen = stablehlo::constantOp(updateConstant);
    SmallVector<stablehlo::Tensor> inps;
    for (auto &c : constants)
      inps.push_back(stablehlo::constantOp(c));

    auto out = stablehlo::dynamicUpdateSliceOp(operandTen, updateTen, inps,
                                               op.getType());
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(),
                                                       fromTensor(out));

    return success();
  }
};

template <auto f>
LogicalResult binaryConstProp(Operation *op, PatternRewriter &rewriter) {
  // return if not constant
  DenseElementsAttr lhsAttr;
  DenseElementsAttr rhsAttr;
  if (!matchPattern(op->getOperand(0), m_Constant(&lhsAttr)) ||
      !matchPattern(op->getOperand(1), m_Constant(&rhsAttr)))
    return failure();

  stablehlo::Tensor lhsTen;
  stablehlo::Tensor rhsTen;
  RankedTensorType ty = cast<RankedTensorType>(op->getResultTypes()[0]);

  // only const prop if the constant has a single user to prevent create many
  // constants
  if ((!lhsAttr.isSplat() || !rhsAttr.isSplat()) &&
      !llvm::hasSingleElement(op->getResult(0).getUsers()))
    return failure();

  if (lhsAttr.isSplat() && rhsAttr.isSplat()) {
    ty = RankedTensorType::get(
        {}, cast<ShapedType>(op->getResultTypes()[0]).getElementType());

    lhsTen = stablehlo::makeTensor(lhsAttr.resizeSplat(RankedTensorType::get(
        {}, cast<ShapedType>(op->getOperand(0).getType()).getElementType())));

    rhsTen = stablehlo::makeTensor(rhsAttr.resizeSplat(RankedTensorType::get(
        {}, cast<ShapedType>(op->getOperand(1).getType()).getElementType())));
  } else {
    lhsTen = stablehlo::constantOp(lhsAttr);
    rhsTen = stablehlo::constantOp(rhsAttr);
  }

  // get the resultType
  auto resultType = cast<ShapedType>(ty);

  auto out = fromTensor(f(lhsTen, rhsTen, resultType));

  if (lhsAttr.isSplat() && rhsAttr.isSplat()) {
    out = out.resizeSplat(cast<ShapedType>(op->getResultTypes()[0]));
  }

  // Replace with new constant op containing the computed result
  rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
      op, op->getResultTypes()[0], out);
  return success();
}

template <typename OpTy, auto constPropFn>
struct BinaryConstProp final
    : CheckedOpRewritePattern<OpTy, BinaryConstProp<OpTy, constPropFn>> {
  using CheckedOpRewritePattern<
      OpTy, BinaryConstProp<OpTy, constPropFn>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    return binaryConstProp<constPropFn>(op, rewriter);
  }
};

template <auto f>
LogicalResult unaryConstProp(Operation *op, PatternRewriter &rewriter) {
  // return if not constant
  DenseElementsAttr inputAttr;
  if (!matchPattern(op->getOperand(0), m_Constant(&inputAttr)))
    return failure();

  stablehlo::Tensor inputTen;
  RankedTensorType ty = cast<RankedTensorType>(op->getResultTypes()[0]);

  // only const prop if the constant has a single user to prevent create many
  // constants
  if (!inputAttr.isSplat() &&
      !llvm::hasSingleElement(op->getResult(0).getUsers()))
    return failure();

  if (inputAttr.isSplat()) {
    ty = RankedTensorType::get(
        {}, cast<ShapedType>(op->getResultTypes()[0]).getElementType());
    auto inputTy = RankedTensorType::get(
        {}, cast<ShapedType>(op->getOperand(0).getType()).getElementType());
    inputTen = stablehlo::makeTensor(inputAttr.resizeSplat(inputTy));
  } else {
    inputTen = stablehlo::constantOp(inputAttr);
  }
  // get the resultType
  auto resultType = cast<ShapedType>(ty);

  // Convert constant to tensor, compute log, then convert back to attribute
  auto out = fromTensor(f(inputTen, resultType));

  if (inputAttr.isSplat()) {
    out = out.resizeSplat(cast<ShapedType>(op->getResultTypes()[0]));
  }
  // Replace with new constant op containing the computed result
  rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
      op, op->getResultTypes()[0], out);

  return success();
}

template <typename OpTy, auto constPropFn>
struct UnaryConstProp final
    : CheckedOpRewritePattern<OpTy, UnaryConstProp<OpTy, constPropFn>> {
  using CheckedOpRewritePattern<
      OpTy, UnaryConstProp<OpTy, constPropFn>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    return unaryConstProp<constPropFn>(op, rewriter);
  }
};

struct GeluConstProp final
    : CheckedOpRewritePattern<enzymexla::GeluOp, GeluConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::GeluOp op,
                                    PatternRewriter &rewriter) const {
    switch (op.getGeluApproximation()) {
    case enzymexla::GeluApproximation::TANH:
      return unaryConstProp<geluTanhOp>(op, rewriter);
    case enzymexla::GeluApproximation::SIGMOID:
      return unaryConstProp<geluSigmoidOp>(op, rewriter);
    case enzymexla::GeluApproximation::NONE:
      return failure();
    }

    return failure();
  }
};

struct ClampConstProp final
    : CheckedOpRewritePattern<stablehlo::ClampOp, ClampConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ClampOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr minAttr, inputAttr, maxAttr;
    if (!matchPattern(op.getMin(), m_Constant(&minAttr)) ||
        !matchPattern(op.getOperand(), m_Constant(&inputAttr)) ||
        !matchPattern(op.getMax(), m_Constant(&maxAttr)))
      return failure();

    // TODO: for only min or max with input being constant we can convert this
    // to a min/max op
    stablehlo::Tensor minTen, maxTen, inputTen;
    bool splattedVersion = false;
    RankedTensorType ty = cast<RankedTensorType>(op->getResultTypes()[0]);
    if (minAttr.isSplat() && maxAttr.isSplat() && inputAttr.isSplat()) {
      splattedVersion = true;
      ty = RankedTensorType::get(
          {}, cast<ShapedType>(op->getResultTypes()[0]).getElementType());
      auto inputTy = RankedTensorType::get(
          {}, cast<ShapedType>(op->getOperand(0).getType()).getElementType());
      minTen = stablehlo::makeTensor(minAttr.resizeSplat(inputTy));
      maxTen = stablehlo::makeTensor(maxAttr.resizeSplat(inputTy));
      inputTen = stablehlo::makeTensor(inputAttr.resizeSplat(inputTy));
    } else {
      minTen = stablehlo::constantOp(minAttr);
      maxTen = stablehlo::constantOp(maxAttr);
      inputTen = stablehlo::constantOp(inputAttr);
    }

    auto out =
        fromTensor(clampOp(inputTen, minTen, maxTen, cast<ShapedType>(ty)));

    if (splattedVersion) {
      out = out.resizeSplat(cast<ShapedType>(op->getResultTypes()[0]));
    }
    // Replace with new constant op containing the computed result
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, op->getResultTypes()[0], out);

    return success();
  }
};

struct ChloInfConstProp final
    : CheckedOpRewritePattern<chlo::IsInfOp, ChloInfConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::IsInfOp op,
                                    PatternRewriter &rewriter) const {
    // return if not constant
    DenseElementsAttr inputAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&inputAttr)))
      return failure();

    DenseElementsAttr outAttr;
    auto resultTy = cast<ShapedType>(op->getResultTypes()[0]);

    // handle splat separately
    if (inputAttr.isSplat()) {
      llvm::APInt resVals;
      if (matchPattern(op.getOperand(), m_PosInfFloat()) ||
          matchPattern(op.getOperand(), m_NegInfFloat())) {
        // true
        resVals = llvm::APInt(1, 1);
      } else {
        // false
        resVals = llvm::APInt(1, 0);
      }

      outAttr = DenseElementsAttr::get(resultTy, resVals);
    } else {
      SmallVector<APInt> resVals;
      resVals.reserve(inputAttr.getNumElements());

      // iterate over every element in inputAttr and run check.
      for (APFloat val : inputAttr.getValues<APFloat>()) {
        bool isInf = val.isInfinity();
        resVals.push_back(APInt(1, isInf ? 1 : 0));
      }

      outAttr = DenseElementsAttr::get(resultTy, resVals);
    }

    // replace op with the bool const op
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, op->getResultTypes()[0], outAttr);

    return success();
  }
};

struct ConcatConstProp final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatConstProp> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;
  size_t max_constant_expansion;
  ConcatConstProp(size_t max_constant_expansion, MLIRContext *context,
                  PatternBenefit benefit = 1,
                  ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

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

    if (op->getNumOperands() == 1) {
      rewriter.replaceOp(op, op->getOperand(0));
      return success();
    }

    {
      SmallVector<Value> subconcat;
      bool changed = false;
      for (auto v : op->getOperands()) {
        if (auto c2 = v.getDefiningOp<stablehlo::ConcatenateOp>())
          if (c2.getDimension() == op.getDimension()) {
            for (auto v2 : c2->getOperands())
              subconcat.push_back(v2);
            changed = true;
            continue;
          }
        subconcat.push_back(v);
      }
      if (changed) {
        rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
            op, subconcat, op.getDimension());
        return success();
      }
    }

    SmallVector<DenseElementsAttr> constants;
    constants.assign(op->getNumOperands(), DenseElementsAttr());
    bool legal = true;
    for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i) {
      matchPattern(op->getOperand(i), m_Constant(&constants[i]));
      if (!constants[i])
        legal = false;
    }

    if (legal) {

      if (constants[0].isSplat()) {
        bool allSplat = true;
        for (int i = 1; i < op->getNumOperands(); i++) {
          if (!constants[i].isSplat()) {
            allSplat = false;
            break;
          }
          if (constants[0].getSplatValue<Attribute>() !=
              constants[i].getSplatValue<Attribute>()) {
            allSplat = false;
            break;
          }
        }
        if (allSplat) {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, op.getType(), constants[0].resizeSplat(op.getType()));
          return success();
        }
      }
      size_t size = 1;
      for (auto sz : op.getType().getShape())
        size *= sz;
      if (size >= max_constant_expansion)
        return failure();

      SmallVector<stablehlo::Tensor> inps;
      for (auto &c : constants)
        inps.push_back(stablehlo::constantOp(c));
      auto out =
          stablehlo::concatenateOp(inps, op.getDimension(), op.getType());
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(),
                                                         fromTensor(out));
      return success();
    }
    return failure();
  }
};

struct ScatterConstFold final
    : CheckedOpRewritePattern<stablehlo::ScatterOp, ScatterConstFold> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;
  size_t max_constant_expansion;

  ScatterConstFold(size_t max_constant_expansion, MLIRContext *context,
                   PatternBenefit benefit = 1,
                   ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp op,
                                    PatternRewriter &rewriter) const {
    // First we check that all inputs/updates are constants and then we
    // materialize them using constantOp
    SmallVector<DenseElementsAttr> inputConstants, updateConstants;
    DenseElementsAttr scatterIndicesConstant;
    for (auto input : op.getInputs()) {
      DenseElementsAttr attr;
      if (matchPattern(input, m_Constant(&attr))) {
        inputConstants.push_back(attr);
      } else {
        return failure();
      }
    }
    for (auto update : op.getUpdates()) {
      DenseElementsAttr attr;
      if (matchPattern(update, m_Constant(&attr))) {
        updateConstants.push_back(attr);
      } else {
        return failure();
      }
    }
    if (!matchPattern(op.getScatterIndices(),
                      m_Constant(&scatterIndicesConstant)))
      return failure();

    SmallVector<ShapedType> resultTypes(op->getResultTypes());
    int totalResultSize = 1;
    for (auto resultType : resultTypes) {
      int resultSize = 1;
      for (auto sidx : resultType.getShape()) {
        resultSize *= sidx;
      }
      totalResultSize += resultSize;
    }
    if (totalResultSize > max_constant_expansion)
      return rewriter.notifyMatchFailure(op, "result size too large");

    SmallVector<stablehlo::Tensor> inputs, updates;
    for (auto input : inputConstants)
      inputs.push_back(stablehlo::constantOp(input));
    for (auto update : updateConstants)
      updates.push_back(stablehlo::constantOp(update));
    stablehlo::Tensor scatterIndices =
        stablehlo::constantOp(scatterIndicesConstant);

    // construct the scope and add all live variables (accessible from update
    // computation) to it. If any of the variables are defined outside the
    // region but isn't constant, we cant constant fold.
    Scope scope(nullptr);
    auto &updateComputation = op.getUpdateComputation();

    if (updateComputation.empty())
      return failure();

    llvm::SetVector<Value> liveValues;
    Block &block = updateComputation.front();
    for (auto &nestedOp : block.getOperations()) {
      for (Value operand : nestedOp.getOperands()) {
        if (operand.getParentRegion() != &updateComputation) {
          liveValues.insert(operand);
        }
      }
    }
    SmallVector<DenseElementsAttr> liveValuesAttr;
    for (Value liveValue : liveValues) {
      DenseElementsAttr attr;
      if (matchPattern(liveValue, m_Constant(&attr))) {
        liveValuesAttr.push_back(attr);
      } else {
        return failure();
      }
    }
    for (auto [liveValue, attr] : llvm::zip(liveValues, liveValuesAttr)) {
      auto tensor = stablehlo::constantOp(attr);
      scope.add(liveValue, tensor);
    }

    auto scatterDimensionNumbers = op.getScatterDimensionNumbers();
    stablehlo::Axes updateWindowDims(
        scatterDimensionNumbers.getUpdateWindowDims());
    stablehlo::Axes insertedWindowDims(
        scatterDimensionNumbers.getInsertedWindowDims());
    stablehlo::Axes inputBatchingDims(
        scatterDimensionNumbers.getInputBatchingDims());
    stablehlo::Axes scatterIndicesBatchingDims(
        scatterDimensionNumbers.getScatterIndicesBatchingDims());
    stablehlo::Axes scatterDimsToOperandDims(
        scatterDimensionNumbers.getScatterDimsToOperandDims());
    stablehlo::Axis indexVectorDim(scatterDimensionNumbers.getIndexVectorDim());

    auto results = stablehlo::scatterOp(
        inputs, scatterIndices, updates, updateWindowDims, insertedWindowDims,
        inputBatchingDims, scatterIndicesBatchingDims, scatterDimsToOperandDims,
        indexVectorDim, updateComputation, nullptr, scope, resultTypes);

    for (auto [i, result] : llvm::enumerate(results)) {
      auto constOp = stablehlo::ConstantOp::create(rewriter, op.getLoc(),
                                                   fromTensor(result));
      rewriter.replaceAllUsesWith(op->getResult(0), constOp.getResult());
    }

    return success();
  }
};

struct ReshapeEmptyBroadcast final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeEmptyBroadcast> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto bcast = op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcast)
      return failure();
    if (bcast.getBroadcastDimensions().size() != 0)
      return failure();
    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), bcast.getOperand(), bcast.getBroadcastDimensions());
    return success();
  }
};

// BroadcastInDim produces singleton dims and the reshape removes certain
// singleton dims, we can avoid constructing those dims in the first place
struct ReshapeBroadcast final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeBroadcast> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto bcast = op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcast || !llvm::hasSingleElement(bcast->getUsers())) {
      return failure();
    }

    auto deletionDims =
        findReshapeInsertionDims(op.getType(), op.getOperand().getType());
    if (deletionDims.empty()) {
      return failure();
    }

    // for each of the deletion dims check if that was introduced by the
    // broadcast
    auto bcastOutputShape = bcast.getType().getShape();
    auto bcastDims = bcast.getBroadcastDimensions();

    SmallVector<int64_t> dimsToRemove;
    for (auto dim : deletionDims) {
      if (!llvm::is_contained(bcastDims, dim) && bcastOutputShape[dim] == 1) {
        dimsToRemove.push_back(dim);
      }
    }

    // If no dims can be removed, bail out
    if (dimsToRemove.empty()) {
      return failure();
    }

    // Compute the new broadcast dimensions by adjusting for removed dimensions
    SmallVector<int64_t> newBroadcastDims;
    for (int64_t bcastDim : bcastDims) {
      // Count how many dimensions before this one are being removed
      int64_t offset = 0;
      for (int64_t removedDim : dimsToRemove) {
        if (removedDim < bcastDim) {
          offset++;
        }
      }
      newBroadcastDims.push_back(bcastDim - offset);
    }

    // Compute the intermediate broadcast shape (after removing dimsToRemove)
    SmallVector<int64_t> newBcastShape;
    for (size_t i = 0; i < bcastOutputShape.size(); i++) {
      if (!llvm::is_contained(dimsToRemove, (int64_t)i)) {
        newBcastShape.push_back(bcastOutputShape[i]);
      }
    }

    auto newBcast = stablehlo::BroadcastInDimOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(newBcastShape, bcast.getType().getElementType()),
        bcast.getOperand(), newBroadcastDims);

    // If the new broadcast shape matches the reshape output, we're done
    Value reshaped = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), newBcast.getResult(), op.getType().getShape());
    rewriter.replaceOp(op, reshaped);
    return success();
  }
};

struct BroadcastReshape final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp, BroadcastReshape> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    if (reshape.getType().hasStaticShape() && type.hasStaticShape() &&
        type.getShape().size() >=
            reshape.getOperand().getType().getShape().size()) {

      auto deletionDims = findReshapeInsertionDims(
          reshape.getType(), reshape.getOperand().getType());
      if (!deletionDims.empty()) {
        SmallVector<int64_t> unusedSingletons;

        for (auto &&[i, s] : llvm::enumerate(type.getShape())) {
          if (llvm::is_contained(op.getBroadcastDimensions(), i))
            continue;
          unusedSingletons.push_back(i);
        }

        SmallVector<int64_t> bcast(
            reshape.getOperand().getType().getShape().size(), 0);

        size_t bcast_idx = 0;
        size_t unused_idx = 0;

        for (size_t i = 0; i < reshape.getOperand().getType().getShape().size();
             i++) {
          if (llvm::is_contained(deletionDims, i)) {
            bcast[i] = unusedSingletons[unused_idx];
            unused_idx++;

          } else {
            bcast[i] = op.getBroadcastDimensions()[bcast_idx];
            bcast_idx++;
          }
        }

        rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
            op, op.getType(), reshape.getOperand(), bcast);

        return success();
      }
    }

    SmallVector<int64_t> dims;

    size_t pre_reshape_idx = 0;
    size_t postidx = 0;

    SmallVector<int64_t> oneOutIdxs;
    for (auto en : llvm::enumerate(op.getType().getShape()))
      if (en.value() == 1)
        oneOutIdxs.push_back(en.index());

    for (auto en : llvm::enumerate(reshape.getType().getShape())) {
      if (en.value() == 1) {
        continue;
      }

      if (pre_reshape_idx == reshape.getOperand().getType().getShape().size())
        return failure();
      auto ival = reshape.getOperand().getType().getShape()[pre_reshape_idx];
      while (ival == 1 &&
             pre_reshape_idx + 1 <
                 reshape.getOperand().getType().getShape().size()) {
        if (postidx == oneOutIdxs.size())
          return failure();
        dims.push_back(oneOutIdxs[postidx]);
        postidx++;
        pre_reshape_idx++;
        ival = reshape.getOperand().getType().getShape()[pre_reshape_idx];
      }
      if (en.value() == ival) {
        dims.push_back(op.getBroadcastDimensions()[en.index()]);
        pre_reshape_idx++;
        continue;
      }
      return failure();
    }
    while (pre_reshape_idx !=
           reshape.getOperand().getType().getShape().size()) {
      if (reshape.getOperand().getType().getShape()[pre_reshape_idx] != 1) {
        LLVM_DEBUG(llvm::dbgs() << "bcast: " << op << "\n";
                   llvm::dbgs() << "reshape: " << reshape << "\n\n");
        return failure();
      }

      size_t nextdim = 0;
      if (postidx == oneOutIdxs.size()) {
        return failure();
      } else {
        nextdim = oneOutIdxs[postidx];
        postidx++;
      }
      dims.push_back(nextdim);
      pre_reshape_idx++;
    }
    assert(dims.size() == reshape.getOperand().getType().getShape().size());
    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), reshape.getOperand(), dims);
    return success();
  }
};

struct BroadcastToReshape final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp, BroadcastToReshape> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();
    assert(op.getBroadcastDimensions().size() ==
           op.getOperand().getType().getShape().size());

    // Ensure these are sorted
    for (auto en : llvm::enumerate(op.getBroadcastDimensions())) {
      if (en.index() == 0)
        continue;
      if (op.getBroadcastDimensions()[en.index() - 1] >= en.value()) {
        return failure();
      }
    }

    // Check that no new data is added
    for (auto en : llvm::enumerate(op.getType().getShape())) {
      ssize_t idx = -1;
      for (auto en2 : llvm::enumerate(op.getBroadcastDimensions())) {
        if (en2.value() == en.index())
          idx = en2.index();
      }
      if (idx != -1) {
        if (en.value() != op.getOperand().getType().getShape()[idx]) {
          return failure();
        }
        continue;
      }
      if (en.value() != 1)
        return failure();
    }

    // replace with reshape
    if (op.getType() == op.getOperand().getType())
      rewriter.replaceOp(op, op.getOperand());
    else
      rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                        op.getOperand());
    return success();
  }
};

struct BroadcastPad final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp, BroadcastPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {

    auto pad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    SplatElementsAttr splat1;
    if (!llvm::hasSingleElement(pad->getUsers()) &&
        !(matchPattern(pad.getOperand(), m_Constant(&splat1))))
      return failure();

    SmallVector<int64_t> paddingLow;
    SmallVector<int64_t> paddingHigh;
    SmallVector<int64_t> paddingInt;

    // broadcast dim map idx from pad output -> which broadcast output

    SmallVector<int64_t> midShape;

    for (auto en : llvm::enumerate(op.getType().getShape())) {
      // which pad in/output dim
      ssize_t origIdx = -1;
      for (auto en2 : llvm::enumerate(op.getBroadcastDimensions())) {
        if (en2.value() == en.index()) {
          origIdx = en2.index();
          break;
        }
      }

      if (origIdx == -1) {
        paddingLow.push_back(0);
        paddingHigh.push_back(0);
        paddingInt.push_back(0);
        midShape.push_back(en.value());
      } else {
        if (pad.getType().getShape()[origIdx] == en.value()) {
          paddingLow.push_back(pad.getEdgePaddingLow()[origIdx]);
          paddingHigh.push_back(pad.getEdgePaddingHigh()[origIdx]);
          paddingInt.push_back(pad.getInteriorPadding()[origIdx]);
          midShape.push_back(pad.getOperand().getType().getShape()[origIdx]);
        } else {
          if (pad.getEdgePaddingLow()[origIdx] != 0)
            return failure();
          if (pad.getEdgePaddingHigh()[origIdx] != 0)
            return failure();
          if (pad.getInteriorPadding()[origIdx] != 0)
            return failure();
          paddingLow.push_back(pad.getEdgePaddingLow()[origIdx]);
          paddingHigh.push_back(pad.getEdgePaddingHigh()[origIdx]);
          paddingInt.push_back(pad.getInteriorPadding()[origIdx]);
          midShape.push_back(en.value());
        }
      }
    }

    auto bcast2 = stablehlo::BroadcastInDimOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(midShape, pad.getType().getElementType()),
        pad.getOperand(), op.getBroadcastDimensions());

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, bcast2, pad.getPaddingValue(), paddingLow, paddingHigh, paddingInt);
    return success();
  }
};

// Given a value and index idx, determine whether all values are the same along
// idx. If so, return said value
std::optional<Value> is_same_in_axis(OpBuilder &rewriter, ShapedType outTy,
                                     Value v, size_t idx) {
  mlir::SplatElementsAttr splat;
  if (matchPattern(v, m_Constant(&splat))) {
    return stablehlo::ConstantOp::create(rewriter, v.getLoc(), outTy,
                                         splat.resizeSplat(outTy));
  }

  return {};
}

struct ScatterToDynamicUpdateSlice final
    : CheckedOpRewritePattern<stablehlo::ScatterOp,
                              ScatterToDynamicUpdateSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1) {
      return failure();
    }

    CheckCommonScatterOp scatterCheck(op);
    if (scatterCheck.kind != ScatterOpKind::Setindex &&
        scatterCheck.kind != ScatterOpKind::ConstantSetindex) {
      return failure();
    }

    auto dims = op.getScatterDimensionNumbers();

    auto input = op.getInputs()[0];
    auto scatter = op.getScatterIndices();
    auto updatesTy = cast<RankedTensorType>(op.getUpdates()[0].getType());
    auto updateRank = updatesTy.getRank();

    if (dims.getInsertedWindowDims().size() != 0 ||
        dims.getUpdateWindowDims().size() != updateRank) {
      return failure();
    }

    Value update;
    if (scatterCheck.kind == ScatterOpKind::Setindex) {
      update = op.getUpdates()[0];
    } else if (scatterCheck.kind == ScatterOpKind::ConstantSetindex) {
      auto splatAttr = DenseElementsAttr::get(
          cast<ShapedType>(op.getUpdates()[0].getType()),
          scatterCheck.constant.getSplatValue<Attribute>());
      update = stablehlo::ConstantOp::create(rewriter, op.getLoc(), updatesTy,
                                             splatAttr);
    } else {
      return failure();
    }

    auto ity = RankedTensorType::get(
        {}, cast<ShapedType>(scatter.getType()).getElementType());
    SmallVector<Value> start(updateRank, 0);
    for (auto en : llvm::enumerate(dims.getScatterDimsToOperandDims())) {
      auto startval = is_same_in_axis(rewriter, ity, scatter, en.index());
      if (!startval)
        return failure();
      start[en.value()] = *startval;
    }
    for (auto &v : start) {
      if (v != nullptr)
        continue;
      v = stablehlo::ConstantOp::create(rewriter, op.getLoc(), ity,
                                        cast<ElementsAttr>(makeAttr(ity, 0)));
    }
    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
        op, op.getResult(0).getType(), input, update, start);
    return success();
  }
};

struct ElementwiseAllTransposeOperandsSimplify
    : public CheckedOpTraitRewritePattern<
          OpTrait::Elementwise, ElementwiseAllTransposeOperandsSimplify> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise,
      ElementwiseAllTransposeOperandsSimplify>::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() < 2)
      return rewriter.notifyMatchFailure(
          op, "ElementwiseAllTransposeOperandsSimplify needs >= 2 operands");

    SmallVector<Value> operands;
    SmallVector<OperandKind> kinds;
    DenseI64ArrayAttr permutation;
    bool foundTranspose = false;
    for (auto operand : op->getOperands()) {
      if (matchPattern(operand, m_Constant())) {
        kinds.push_back(OperandKind::Const);
        operands.push_back(operand);
        continue;
      } else if (auto transposeOp =
                     operand.getDefiningOp<stablehlo::TransposeOp>()) {
        if (!foundTranspose) {
          foundTranspose = true;
          permutation = transposeOp.getPermutationAttr();
        } else if (permutation != transposeOp.getPermutationAttr()) {
          return rewriter.notifyMatchFailure(
              op,
              "ElementwiseAllTransposeOperandsSimplify needs all operands to "
              "have the same permutation");
        }
        kinds.push_back(OperandKind::Transpose);
        operands.push_back(transposeOp.getOperand());
        continue;
      }

      return rewriter.notifyMatchFailure(
          op,
          "ElementwiseAllTransposeOperandsSimplify needs all operands to be "
          "either constants or transpose ops");
    }

    if (!foundTranspose)
      return rewriter.notifyMatchFailure(
          op, "ElementwiseAllTransposeOperandsSimplify needs at least one "
              "transpose op");

    auto invPerm =
        rewriter.getDenseI64ArrayAttr(getInversePermutation(permutation));
    for (size_t i = 0; i < operands.size(); i++) {
      switch (kinds[i]) {
      case OperandKind::Transpose:
        break;
      case OperandKind::Const:
        // This will be eliminated by a transpose(constant) -> constant
        // optimization
        auto constTransposeOp = stablehlo::TransposeOp::create(
            rewriter, op->getLoc(), operands[i], invPerm);
        operands[i] = constTransposeOp;
        break;
      }
    }

    auto operandTy = cast<RankedTensorType>(operands[0].getType());
    SmallVector<int64_t> elemResShape(operandTy.getShape().begin(),
                                      operandTy.getShape().end());
    auto newOp = Operation::create(
        op->getLoc(), op->getName(),
        {RankedTensorType::get(
            elemResShape,
            cast<TensorType>(op->getResult(0).getType()).getElementType())},
        operands, op->getAttrs(), mlir::PropertyRef(), op->getSuccessors(), 0);
    rewriter.insert(newOp);
    auto newTransposeOp = stablehlo::TransposeOp::create(
        rewriter, op->getLoc(), newOp->getResult(0), permutation);
    rewriter.replaceOp(op, newTransposeOp);
    return success();
  }

private:
  enum class OperandKind { Transpose, Const };
};

struct TransposeElementwiseTransposeSimplify
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeElementwiseTransposeSimplify> {
  using CheckedOpRewritePattern<
      stablehlo::TransposeOp,
      TransposeElementwiseTransposeSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto elem = op.getOperand().getDefiningOp();
    if (!elem)
      return failure();
    if (!stablehlo::hasTraitElementwise(elem))
      return failure();

    SmallVector<Value> newOperands;

    auto invPerm = rewriter.getDenseI64ArrayAttr(
        getInversePermutation(op.getPermutation()));

    for (auto operand : elem->getOperands()) {
      auto innerTransposeOp = operand.getDefiningOp<stablehlo::TransposeOp>();
      if (!innerTransposeOp)
        return failure();
      if (innerTransposeOp.getPermutationAttr() != invPerm)
        return failure();
      newOperands.push_back(innerTransposeOp.getOperand());
    }

    auto newElem = Operation::create(elem->getLoc(), elem->getName(),
                                     {op->getResult(0).getType()}, newOperands,
                                     elem->getAttrs(), mlir::PropertyRef(),
                                     elem->getSuccessors(), 0);
    rewriter.insert(newElem);
    rewriter.replaceOp(op, newElem);
    return success();
  }
};

struct AddSimplify
    : public CheckedOpRewritePattern<stablehlo::AddOp, AddSimplify> {
  using CheckedOpRewritePattern<stablehlo::AddOp,
                                AddSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AddOp op,
                                    PatternRewriter &rewriter) const {
    Attribute lhsAttr, rhsAttr;
    bool lhsConst = matchPattern(op.getLhs(), m_Constant(&lhsAttr));
    bool rhsConst = matchPattern(op.getRhs(), m_Constant(&rhsAttr));

    if (lhsConst && (matchPattern(lhsAttr, m_AnyZeroFloat()) ||
                     matchPattern(lhsAttr, m_Zero()) ||
                     matchPattern(lhsAttr, m_AnyZeroComplex()))) {
      rewriter.replaceOp(op, op.getRhs());
      return success();
    }

    if (rhsConst && (matchPattern(rhsAttr, m_AnyZeroFloat()) ||
                     matchPattern(rhsAttr, m_Zero()) ||
                     matchPattern(rhsAttr, m_AnyZeroComplex()))) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    return failure();
  }
};

struct ReplaceNegAddWithSubtract
    : public CheckedOpRewritePattern<stablehlo::AddOp,
                                     ReplaceNegAddWithSubtract> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AddOp op,
                                    PatternRewriter &rewriter) const {
    if (auto rhsNegateOp = op.getRhs().getDefiningOp<stablehlo::NegOp>()) {
      if (llvm::hasSingleElement(rhsNegateOp->getUsers())) {
        rewriter.replaceOpWithNewOp<stablehlo::SubtractOp>(
            op, op.getLhs(), rhsNegateOp.getOperand());
        return success();
      }
    }

    if (auto lhsNegateOp = op.getLhs().getDefiningOp<stablehlo::NegOp>()) {
      if (llvm::hasSingleElement(lhsNegateOp->getUsers())) {
        rewriter.replaceOpWithNewOp<stablehlo::SubtractOp>(
            op, op.getRhs(), lhsNegateOp.getOperand());
        return success();
      }
    }

    return failure();
  }
};

struct ReplaceSubtractNegWithAdd
    : CheckedOpRewritePattern<stablehlo::SubtractOp,
                              ReplaceSubtractNegWithAdd> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SubtractOp op,
                                    PatternRewriter &rewriter) const {
    if (auto rhsNegateOp = op.getRhs().getDefiningOp<stablehlo::NegOp>()) {
      if (llvm::hasSingleElement(rhsNegateOp->getUsers())) {
        rewriter.replaceOpWithNewOp<stablehlo::AddOp>(op, op.getLhs(),
                                                      rhsNegateOp.getOperand());
        return success();
      }
    }

    return failure();
  }
};

struct SubSimplify
    : public CheckedOpRewritePattern<stablehlo::SubtractOp, SubSimplify> {
  using CheckedOpRewritePattern<stablehlo::SubtractOp,
                                SubSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SubtractOp op,
                                    PatternRewriter &rewriter) const {
    Attribute lhsAttr, rhsAttr;
    bool lhsConst = matchPattern(op.getLhs(), m_Constant(&lhsAttr));
    bool rhsConst = matchPattern(op.getRhs(), m_Constant(&rhsAttr));

    if (rhsConst && (matchPattern(rhsAttr, m_AnyZeroFloat()) ||
                     matchPattern(rhsAttr, m_Zero()) ||
                     matchPattern(rhsAttr, m_AnyZeroComplex()))) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    if (lhsConst && (matchPattern(lhsAttr, m_AnyZeroFloat()) ||
                     matchPattern(lhsAttr, m_Zero()) ||
                     matchPattern(lhsAttr, m_AnyZeroComplex()))) {
      rewriter.replaceOpWithNewOp<stablehlo::NegOp>(op, op.getRhs());
      return success();
    }

    if (isa<IntegerType>(op.getType().getElementType()) &&
        op.getLhs() == op.getRhs()) {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, rewriter.getZeroAttr(op.getType()));
      return success();
    }

    return failure();
  }
};

struct ExponentialMinusOneFuse
    : public CheckedOpRewritePattern<stablehlo::SubtractOp,
                                     ExponentialMinusOneFuse> {
  using CheckedOpRewritePattern<
      stablehlo::SubtractOp, ExponentialMinusOneFuse>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SubtractOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    { // exp(x) - 1 -> expm1(x)
      auto defOp = lhs.getDefiningOp<stablehlo::ExpOp>();
      if (defOp && llvm::hasSingleElement(defOp->getUsers()) &&
          (matchPattern(rhs, m_One()) || matchPattern(rhs, m_OneFloat()))) {
        rewriter.replaceOpWithNewOp<stablehlo::Expm1Op>(op, defOp.getOperand());
        return success();
      }
    }

    { // 1 - exp(x) -> -expm1(x)
      auto defOp = rhs.getDefiningOp<stablehlo::ExpOp>();
      if (defOp && llvm::hasSingleElement(defOp->getUsers()) &&
          (matchPattern(lhs, m_One()) || matchPattern(lhs, m_OneFloat()))) {
        auto expm1 = stablehlo::Expm1Op::create(
            rewriter, op.getLoc(), op.getType(), defOp.getOperand());
        rewriter.replaceOpWithNewOp<stablehlo::NegOp>(op, expm1);
        return success();
      }
    }

    return failure();
  }
};

struct ExponentialMinusOneAddFuse
    : public CheckedOpRewritePattern<stablehlo::AddOp,
                                     ExponentialMinusOneAddFuse> {
  using CheckedOpRewritePattern<
      stablehlo::AddOp, ExponentialMinusOneAddFuse>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AddOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto isMinusOne = [](Value val) {
      SplatElementsAttr attr;
      if (!matchPattern(val, m_Constant(&attr)))
        return false;
      auto doubleVal = getDoubleFromAttr(attr.getSplatValue<mlir::TypedAttr>());
      return doubleVal && *doubleVal == -1.0;
    };

    { // exp(x) + -1 -> expm1(x)
      auto defOp = lhs.getDefiningOp<stablehlo::ExpOp>();
      if (defOp && llvm::hasSingleElement(defOp->getUsers()) &&
          isMinusOne(rhs)) {
        rewriter.replaceOpWithNewOp<stablehlo::Expm1Op>(op, defOp.getOperand());
        return success();
      }
    }

    { // -1 + exp(x) -> expm1(x)
      auto defOp = rhs.getDefiningOp<stablehlo::ExpOp>();
      if (defOp && llvm::hasSingleElement(defOp->getUsers()) &&
          isMinusOne(lhs)) {
        rewriter.replaceOpWithNewOp<stablehlo::Expm1Op>(op, defOp.getOperand());
        return success();
      }
    }

    return failure();
  }
};

struct TransposeSymmetricSimplify
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeSymmetricSimplify> {
  using CheckedOpRewritePattern<
      stablehlo::TransposeOp,
      TransposeSymmetricSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto operand = op.getOperand();
    auto perm = op.getPermutation();
    if (perm.size() != 2 || perm[0] != 1 || perm[1] != 0)
      return failure();

    auto resTy = cast<RankedTensorType>(op.getResult().getType());
    if (resTy.getRank() != 2 || resTy.getDimSize(0) != resTy.getDimSize(1))
      return failure(); // quick check and exit

    if (canApplySymmetricPattern(
            operand, rewriter)) { // tranpose(symmetric) -> symmetric
      rewriter.replaceOp(op, operand);
      return success();
    }
    return failure();
  }
};

struct TransposePartialSymmetrySimplify
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposePartialSymmetrySimplify> {
  using CheckedOpRewritePattern<
      stablehlo::TransposeOp,
      TransposePartialSymmetrySimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto operand = op.getOperand();
    auto perm = op.getPermutation();

    auto annotationOpt =
        enzyme::PartialSymmetryAnnotation::createFromIR(operand);
    if (!annotationOpt.has_value())
      return failure();

    auto annotation = annotationOpt.value();

    bool isIdentity = true;
    for (int64_t i = 0; i < (int64_t)perm.size(); ++i) {
      if (annotation.getSetId(i) != annotation.getSetId(perm[i])) {
        isIdentity = false;
        break;
      }
    }

    if (isIdentity) {
      rewriter.replaceOp(op, operand);
      return success();
    }

    return failure();
  }
};

struct NoNanSelfSubSimplify
    : public NoNanCheckedOpRewritePattern<stablehlo::SubtractOp,
                                          NoNanSelfSubSimplify> {
  using NoNanCheckedOpRewritePattern<
      stablehlo::SubtractOp,
      NoNanSelfSubSimplify>::NoNanCheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SubtractOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getLhs() == op.getRhs()) {
      if (canApplyNoNanPattern(allowOnFloatingPointMath, op.getType(),
                               op.getLhs().getType(), op, rewriter)) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            op, rewriter.getZeroAttr(op.getType()));
        return success();
      }
    }

    return failure();
  }
};

struct AndSimplify
    : public CheckedOpRewritePattern<stablehlo::AndOp, AndSimplify> {
  using CheckedOpRewritePattern<stablehlo::AndOp,
                                AndSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AndOp op,
                                    PatternRewriter &rewriter) const {

    if (op.getLhs() == op.getRhs()) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    for (int i = 0; i < 2; i++) {
      Attribute attr;
      if (!matchPattern(op.getOperand(i), m_Constant(&attr)))
        continue;
      // false & x -> false
      if (matchPattern(attr, m_Zero())) {
        rewriter.replaceOp(op, op.getOperand(i));
        return success();
      }
      // true & x -> x
      if (matchPattern(attr, m_AllOnes())) {
        rewriter.replaceOp(op, op.getOperand(1 - i));
        return success();
      }
    }

    return failure();
  }
};

struct OrSimplify
    : public CheckedOpRewritePattern<stablehlo::OrOp, OrSimplify> {
  using CheckedOpRewritePattern<stablehlo::OrOp,
                                OrSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::OrOp op,
                                    PatternRewriter &rewriter) const {

    if (op.getLhs() == op.getRhs()) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    for (int i = 0; i < 2; i++) {
      Attribute attr;
      if (!matchPattern(op.getOperand(i), m_Constant(&attr)))
        continue;
      // true | x -> true
      if (matchPattern(attr, m_AllOnes())) {
        rewriter.replaceOp(op, op.getOperand(i));
        return success();
      }
      // false | x -> x
      if (matchPattern(attr, m_Zero())) {
        rewriter.replaceOp(op, op.getOperand(1 - i));
        return success();
      }
    }

    return failure();
  }
};

struct XorSimplify
    : public CheckedOpRewritePattern<stablehlo::XorOp, XorSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::XorOp op,
                                    PatternRewriter &rewriter) const {

    for (int i = 0; i < 2; i++) {
      Attribute attr;
      if (!matchPattern(op.getOperand(i), m_Constant(&attr)))
        continue;
      // false ^ x -> x
      if (matchPattern(attr, m_Zero())) {
        rewriter.replaceOp(op, op.getOperand(1 - i));
        return success();
      }
      // true ^ x -> not x
      if (matchPattern(attr, m_AllOnes())) {
        rewriter.replaceOpWithNewOp<stablehlo::NotOp>(op, op.getOperand(1 - i));
        return success();
      }
    }

    return failure();
  }
};

struct MulSimplify
    : public CheckedOpRewritePattern<stablehlo::MulOp, MulSimplify> {
  using CheckedOpRewritePattern<stablehlo::MulOp,
                                MulSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    Attribute lhsAttr, rhsAttr;
    bool lhsConst = matchPattern(op.getLhs(), m_Constant(&lhsAttr));
    bool rhsConst = matchPattern(op.getRhs(), m_Constant(&rhsAttr));

    // 1 * x -> x
    if (lhsConst && (matchPattern(lhsAttr, m_One()) ||
                     matchPattern(lhsAttr, m_OneFloat()))) {
      rewriter.replaceOp(op, op.getRhs());
      return success();
    }

    // -1 * x -> negate x
    if (lhsConst && (matchPattern(lhsAttr, m_NegOne()) ||
                     matchPattern(lhsAttr, m_NegOneFloat()))) {
      rewriter.replaceOpWithNewOp<stablehlo::NegOp>(op, op.getRhs());
      return success();
    }

    if (!rhsConst)
      return failure();

    // x * 1 -> x
    if (rhsConst && (matchPattern(rhsAttr, m_One()) ||
                     matchPattern(rhsAttr, m_OneFloat()))) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    // x * -1 -> negate x
    if (rhsConst && (matchPattern(rhsAttr, m_NegOne()) ||
                     matchPattern(rhsAttr, m_NegOneFloat()))) {
      rewriter.replaceOpWithNewOp<stablehlo::NegOp>(op, op.getLhs());
      return success();
    }

    return failure();
  }
};

struct DivSimplify
    : public CheckedOpRewritePattern<stablehlo::DivOp, DivSimplify> {
  using CheckedOpRewritePattern<stablehlo::DivOp,
                                DivSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DivOp op,
                                    PatternRewriter &rewriter) const {
    Attribute rhsAttr;
    bool rhsConst = matchPattern(op.getRhs(), m_Constant(&rhsAttr));

    // x / 1 -> x
    if (rhsConst && (matchPattern(rhsAttr, m_OneFloat()) ||
                     matchPattern(rhsAttr, m_One()))) {
      rewriter.replaceOp(op, op.getLhs());
      return success();
    }

    // x / -1 -> negate x
    if (rhsConst && (matchPattern(rhsAttr, m_NegOneFloat()) ||
                     matchPattern(rhsAttr, m_NegOne()))) {
      rewriter.replaceOpWithNewOp<stablehlo::NegOp>(op, op.getLhs());
      return success();
    }

    // x / const -> x * (1 / const)
    if (isa<FloatType>(op.getType().getElementType())) {
      if (auto rhsDenseAttr = dyn_cast_or_null<DenseElementsAttr>(rhsAttr)) {
        {
          DenseElementsAttr lhsAttr;
          if (matchPattern(op.getLhs(), m_Constant(&lhsAttr)))
            return failure(); // const prop will evaluate this
        }

        auto ty = op.getType();
        if (rhsDenseAttr.isSplat()) {
          ty = RankedTensorType::get(
              {}, cast<ShapedType>(op->getResultTypes()[0]).getElementType());
          rhsDenseAttr = rhsDenseAttr.resizeSplat(ty);
        }

        auto rhsTen = stablehlo::constantOp(rhsDenseAttr);
        auto oneTen =
            stablehlo::constantOp(cast<ElementsAttr>(makeAttr(ty, 1)));
        auto out = fromTensor(stablehlo::divideOp(oneTen, rhsTen, ty));

        if (ty != op.getType()) {
          out = out.resizeSplat(op.getType());
        }

        rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
            op, op.getLhs(),
            stablehlo::ConstantOp::create(rewriter, op.getLoc(), op.getType(),
                                          out));
        return success();
      }
    }

    return failure();
  }
};

struct NoNanDivSimplify final
    : public NoNanCheckedOpRewritePattern<stablehlo::DivOp, NoNanDivSimplify> {
  using NoNanCheckedOpRewritePattern<
      stablehlo::DivOp, NoNanDivSimplify>::NoNanCheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DivOp op,
                                    PatternRewriter &rewriter) const {
    // 0 / x -> 0
    Attribute lhsAttr;
    if (matchPattern(op.getLhs(), m_Constant(&lhsAttr)) &&
        (matchPattern(lhsAttr, m_AnyZeroFloat()) ||
         matchPattern(lhsAttr, m_Zero()))) {
      if (canApplyNoNanPattern(allowOnFloatingPointMath, op.getType(), op,
                               rewriter)) {
        rewriter.replaceOp(op, op.getLhs());
        return success();
      }
    }

    // x / x -> 1
    if (op.getLhs() == op.getRhs()) {
      if (canApplyNoNanPattern(allowOnFloatingPointMath, op.getType(), op,
                               rewriter)) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            op, op.getType(), cast<ElementsAttr>(makeAttr(op.getType(), 1)));
        return success();
      }
    }

    return failure();
  }
};

struct RemSimplify
    : public CheckedOpRewritePattern<stablehlo::RemOp, RemSimplify> {
  using CheckedOpRewritePattern<stablehlo::RemOp,
                                RemSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::RemOp op,
                                    PatternRewriter &rewriter) const {

    if (matchPattern(op.getRhs(), m_One())) {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, cast<ElementsAttr>(makeAttr(op.getType(), 0)));
      return success();
    }

    return failure();
  }
};

struct PowSimplify
    : public CheckedOpRewritePattern<stablehlo::PowOp, PowSimplify> {
  using CheckedOpRewritePattern<stablehlo::PowOp,
                                PowSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::PowOp op,
                                    PatternRewriter &rewriter) const {
    Attribute lhsAttr, rhsAttr;
    matchPattern(op.getLhs(), m_Constant(&lhsAttr));
    matchPattern(op.getRhs(), m_Constant(&rhsAttr));

    // pow(x, 1) -> x
    if (rhsAttr && (matchPattern(rhsAttr, m_One()) ||
                    matchPattern(rhsAttr, m_OneFloat()))) {
      rewriter.replaceAllUsesWith(op, op.getLhs());
      return success();
    }

    // pow(x, 0) -> 1 || pow(1, x) -> 1
    if ((rhsAttr && (matchPattern(rhsAttr, m_Zero()) ||
                     matchPattern(rhsAttr, m_AnyZeroFloat()))) ||
        (lhsAttr && (matchPattern(lhsAttr, m_One()) ||
                     matchPattern(lhsAttr, m_OneFloat())))) {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, op.getType(), cast<ElementsAttr>(makeAttr(op.getType(), 1)));
      return success();
    }

    if (isa<FloatType>(op.getType().getElementType())) {
      auto rhs = dyn_cast_or_null<DenseFPElementsAttr>(rhsAttr);
      if (rhs && rhs.isSplat()) {
        bool allHalf = true, allNegOne = true, allNegHalf = true, allTwo = true;
        auto v = rhs.getSplatValue<APFloat>();
        allHalf &= v.isExactlyValue(0.5);
        allNegOne &= v.isExactlyValue(-1.0);
        allNegHalf &= v.isExactlyValue(-0.5);
        allTwo &= v.isExactlyValue(2.0);

        // pow(X, -1) -> 1 / X
        if (allNegOne) {
          rewriter.replaceOpWithNewOp<stablehlo::DivOp>(
              op,
              stablehlo::ConstantOp::create(
                  rewriter, op.getLoc(), op.getType(),
                  cast<ElementsAttr>(makeAttr(op.getType(), 1))),
              op.getLhs());
          return success();
        }

        // pow(X, -0.5) -> rsqrt(X)
        if (allNegHalf) {
          rewriter.replaceOpWithNewOp<stablehlo::RsqrtOp>(op, op.getLhs());
          return success();
        }

        // pow(X, 0.5) -> sqrt(X)
        if (allHalf) {
          rewriter.replaceOpWithNewOp<stablehlo::SqrtOp>(op, op.getLhs());
          return success();
        }

        // pow(X, 2) -> X * X
        if (allTwo) {
          rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, op.getLhs(),
                                                        op.getLhs());
          return success();
        }
      }
    }

    return failure();
  }
};

struct NoNanZeroBasePowSimplify final
    : public NoNanCheckedOpRewritePattern<stablehlo::PowOp,
                                          NoNanZeroBasePowSimplify> {
  using NoNanCheckedOpRewritePattern<
      stablehlo::PowOp, NoNanZeroBasePowSimplify>::NoNanCheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::PowOp op,
                                    PatternRewriter &rewriter) const {
    Attribute lhsAttr;
    if (matchPattern(op.getLhs(), m_Constant(&lhsAttr)) &&
        (matchPattern(lhsAttr, m_Zero()) ||
         matchPattern(lhsAttr, m_AnyZeroFloat()))) {

      DenseElementsAttr attr;
      if (matchPattern(op.getRhs(), m_Constant(&attr)))
        return failure(); // let constant propagation handle this

      if (!canApplyNoNanPattern(allowOnFloatingPointMath, op.getType(), op,
                                rewriter))
        return failure();

      // 0 ^ x => x == 0 ? 1 : (x > 0 ? 0 : Inf)
      auto zero = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), rewriter.getZeroAttr(op.getType()));
      auto rhsPositive =
          stablehlo::CompareOp::create(rewriter, op.getLoc(), op.getRhs(), zero,
                                       stablehlo::ComparisonDirection::GT);
      auto inf = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), op.getType(),
          cast<ElementsAttr>(
              makeAttr(op.getType(), std::numeric_limits<float>::infinity())));
      auto nonZeroCase = stablehlo::SelectOp::create(rewriter, op.getLoc(),
                                                     rhsPositive, zero, inf);
      auto rhsIsZero =
          stablehlo::CompareOp::create(rewriter, op.getLoc(), op.getRhs(), zero,
                                       stablehlo::ComparisonDirection::EQ);
      auto one = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), op.getType(),
          cast<ElementsAttr>(makeAttr(op.getType(), 1)));
      rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(op, rhsIsZero, one,
                                                       nonZeroCase);
      return success();
    }

    return failure();
  }
};

bool is_broadcastable_compare(Value operand) {
  if (auto cmp = operand.getDefiningOp<stablehlo::CompareOp>()) {

    for (int i = 0; i < 2; i++) {
      auto v = cmp->getOperand(i);
      if (v.getDefiningOp<stablehlo::IotaOp>()) {
        continue;
      }
      DenseElementsAttr attr;
      if (matchPattern(v, m_Constant(&attr))) {
        if (attr.isSplat()) {
          continue;
        }
      }
      return false;
    }
    return true;
  }
  if (auto andv = operand.getDefiningOp<stablehlo::AndOp>()) {
    return is_broadcastable_compare(andv.getLhs()) &&
           is_broadcastable_compare(andv.getRhs());
  }
  return false;
}

struct BroadcastCompare
    : public CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                                     BroadcastCompare> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    if (!is_broadcastable_compare(op.getOperand()))
      return failure();

    auto operand = op.getOperand().getDefiningOp();

    Value newops[2];

    for (int i = 0; i < 2; i++) {
      auto v = operand->getOperand(i);
      auto RT = RankedTensorType::get(
          op.getType().getShape(),
          cast<RankedTensorType>(v.getType()).getElementType());
      newops[i] = stablehlo::BroadcastInDimOp::create(
          rewriter, op.getLoc(), RT, v, op.getBroadcastDimensions());
    }

    if (auto cmp = op.getOperand().getDefiningOp<stablehlo::CompareOp>()) {
      auto cmp2 =
          stablehlo::CompareOp::create(rewriter, cmp.getLoc(), newops[0],
                                       newops[1], cmp.getComparisonDirection());
      rewriter.replaceOp(op, cmp2);
    } else {
      auto and2 = stablehlo::AndOp::create(rewriter, op.getOperand().getLoc(),
                                           newops[0], newops[1]);
      rewriter.replaceOp(op, and2);
    }
    return success();
  }
};

struct IotaSimplify
    : public CheckedOpRewritePattern<stablehlo::IotaOp, IotaSimplify> {
  using CheckedOpRewritePattern<stablehlo::IotaOp,
                                IotaSimplify>::CheckedOpRewritePattern;

  size_t max_constant_expansion;
  IotaSimplify(size_t max_constant_expansion, MLIRContext *context,
               PatternBenefit benefit = 1,
               ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

  LogicalResult matchAndRewriteImpl(stablehlo::IotaOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getType().getNumElements() >= max_constant_expansion) {
      return failure();
    }

    // Bail out on degenerate iota ops (e.g. rank-0 tensors) where the iota
    // dimension is out of bounds for the result type's rank.
    if (op.getIotaDimension() >= op.getType().getRank()) {
      return failure();
    }

    auto out = stablehlo::iotaOp(op.getIotaDimension(), op.getType());
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(),
                                                       fromTensor(out));
    return success();
  }
};

struct ConcatToPad
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatToPad> {
  using CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                ConcatToPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getNumOperands() < 2)
      return failure();

    for (unsigned ind : {(unsigned int)0, op.getNumOperands() - 1}) {
      DenseElementsAttr inp;
      if (!matchPattern(op->getOperand(ind), m_Constant(&inp)))
        continue;
      if (!inp.isSplat())
        continue;

      auto subconcat = stablehlo::ConcatenateOp::create(
          rewriter, op.getLoc(),
          (ind == 0) ? op.getOperands().drop_front()
                     : op.getOperands().drop_back(),
          op.getDimension());

      SmallVector<int64_t> low(op.getType().getShape().size(), 0);
      SmallVector<int64_t> high(op.getType().getShape().size(), 0);
      SmallVector<int64_t> interior(op.getType().getShape().size(), 0);
      if (ind == 0)
        low[op.getDimension()] = inp.getType().getShape()[op.getDimension()];
      else
        high[op.getDimension()] = inp.getType().getShape()[op.getDimension()];
      auto type0 = RankedTensorType::get({}, inp.getType().getElementType());
      rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
          op, op.getType(), subconcat,
          stablehlo::ConstantOp::create(rewriter, op.getLoc(), type0,
                                        inp.resizeSplat(type0)),
          low, high, interior);
      return success();
    }
    return failure();
  }
};

// reduce_window(pad(x, lo, hi, 0)) -> reduce_window(x, pad_lo=lo, pad_hi=hi)
struct PadReduceWindow
    : public CheckedOpRewritePattern<stablehlo::ReduceWindowOp,
                                     PadReduceWindow> {
  using CheckedOpRewritePattern<stablehlo::ReduceWindowOp,
                                PadReduceWindow>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceWindowOp op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() != 2)
      return failure();

    if (op.getPadding().has_value() &&
        !llvm::all_of(op.getPadding().value(),
                      [](auto pad) { return pad.isZero(); }))
      return failure();

    Value operand = op->getOperand(0), initValue = op->getOperand(1);

    auto padOp = operand.getDefiningOp<stablehlo::PadOp>();
    if (!padOp || !llvm::all_of(padOp.getInteriorPadding(),
                                [](int64_t pad) { return pad == 0; }))
      return failure();

    if (padOp.getPaddingValue() != initValue)
      return failure();

    auto highValues = padOp.getEdgePaddingHigh();
    auto lowValues = padOp.getEdgePaddingLow();

    int64_t N = highValues.size();

    SmallVector<int64_t> newPaddingValues(2 * N, 0);

    for (int i = 0; i < N; ++i) {
      newPaddingValues[2 * i] = lowValues[i];
      newPaddingValues[2 * i + 1] = highValues[i];
    }

    auto paddingType =
        mlir::RankedTensorType::get({N, 2}, rewriter.getI64Type());
    auto newPaddingAttr =
        mlir::DenseIntElementsAttr::get(paddingType, newPaddingValues);

    auto newOp = stablehlo::ReduceWindowOp::create(
        rewriter, op.getLoc(), op.getResult(0).getType(), padOp.getOperand(),
        initValue, op.getWindowDimensionsAttr(), op.getWindowStridesAttr(),
        op.getBaseDilationsAttr(), op.getWindowDilationsAttr(), newPaddingAttr);
    newOp.getRegion().takeBody(op.getRegion());

    rewriter.replaceOp(op, newOp);

    return success();
  }
};

// conv(pad(x, lo, hi, 0)) -> conv(x, pad_lo=lo, pad_hi=hi)
struct ConvolutionPad
    : public CheckedOpRewritePattern<stablehlo::ConvolutionOp, ConvolutionPad> {
  using CheckedOpRewritePattern<stablehlo::ConvolutionOp,
                                ConvolutionPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvolutionOp op,
                                    PatternRewriter &rewriter) const {
    auto padOp = op.getLhs().getDefiningOp<stablehlo::PadOp>();
    if (!padOp || !llvm::all_of(padOp.getInteriorPadding(),
                                [](int64_t pad) { return pad == 0; }))
      return failure();

    Attribute padValAttr;
    if (!matchPattern(padOp.getPaddingValue(), m_Constant(&padValAttr)) ||
        (!matchPattern(padValAttr, m_Zero()) &&
         !matchPattern(padValAttr, m_AnyZeroFloat())))
      return failure();

    auto paddingLow = padOp.getEdgePaddingLow();
    auto paddingHigh = padOp.getEdgePaddingHigh();

    auto convDims = op.getDimensionNumbers();
    int64_t inputBatchDim = convDims.getInputBatchDimension();
    int64_t inputFeatureDim = convDims.getInputFeatureDimension();
    if (paddingLow[inputBatchDim] != 0 || paddingHigh[inputBatchDim] != 0 ||
        paddingLow[inputFeatureDim] != 0 || paddingHigh[inputFeatureDim] != 0)
      return failure();

    auto inputSpatialDims = convDims.getInputSpatialDimensions();
    int64_t N = paddingLow.size();
    SmallVector<int64_t> newPaddingValues(2 * (N - 2), 0);
    for (auto [i, dim] : llvm::enumerate(inputSpatialDims)) {
      newPaddingValues[2 * i] = paddingLow[dim];
      newPaddingValues[2 * i + 1] = paddingHigh[dim];
    }
    auto paddingType =
        mlir::RankedTensorType::get({N - 2, 2}, rewriter.getI64Type());
    auto newPaddingAttr =
        mlir::DenseIntElementsAttr::get(paddingType, newPaddingValues);

    rewriter.replaceOpWithNewOp<stablehlo::ConvolutionOp>(
        op, op.getType(), padOp.getOperand(), op.getRhs(),
        op.getWindowStridesAttr(), newPaddingAttr, op.getLhsDilationAttr(),
        op.getRhsDilationAttr(), op.getWindowReversalAttr(), convDims,
        op.getFeatureGroupCountAttr(), op.getBatchGroupCountAttr(),
        op.getPrecisionConfigAttr());
    return success();
  }
};

struct ConcatPad
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatPad> {
  using CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                ConcatPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getNumOperands() < 2)
      return failure();

    for (unsigned ind : {(unsigned int)0, op.getNumOperands() - 1}) {

      auto pad = op->getOperand(ind).getDefiningOp<stablehlo::PadOp>();
      if (!pad)
        continue;

      if (pad.getInteriorPadding()[op.getDimension()] != 0)
        continue;

      if (ind == 0) {
        if (pad.getEdgePaddingHigh()[op.getDimension()] != 0)
          continue;
      } else {
        if (pad.getEdgePaddingLow()[op.getDimension()] != 0)
          continue;
      }

      bool legal = true;
      for (size_t i = 0; i < pad.getType().getShape().size(); i++) {
        if (i == op.getDimension())
          continue;
        if (pad.getInteriorPadding()[i] != 0) {
          legal = false;
          break;
        }
        if (pad.getEdgePaddingLow()[i] != 0) {
          legal = false;
          break;
        }
        if (pad.getEdgePaddingHigh()[i] != 0) {
          legal = false;
          break;
        }
      }

      if (!legal)
        continue;

      auto prevArgs = (ind == 0) ? op.getOperands().drop_front()
                                 : op.getOperands().drop_back();
      SmallVector<Value> subArgs(prevArgs.begin(), prevArgs.end());
      if (ind == 0)
        subArgs.insert(subArgs.begin(), pad.getOperand());
      else
        subArgs.push_back(pad.getOperand());

      auto subconcat = stablehlo::ConcatenateOp::create(
          rewriter, op.getLoc(), subArgs, op.getDimension());

      rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
          op, op.getType(), subconcat, pad.getPaddingValue(),
          pad.getEdgePaddingLow(), pad.getEdgePaddingHigh(),
          pad.getInteriorPadding());
      return success();
    }
    return failure();
  }
}; // namespace

struct ConvertSimplify
    : public CheckedOpRewritePattern<stablehlo::ConvertOp, ConvertSimplify> {
  using CheckedOpRewritePattern<stablehlo::ConvertOp,
                                ConvertSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr inp;
    matchPattern(op->getOperand(0), m_Constant(&inp));
    if (inp) {
      stablehlo::Tensor ten;
      RankedTensorType ty = op.getType();
      if (inp.isSplat()) {
        ten = stablehlo::makeTensor(inp.resizeSplat(
            RankedTensorType::get({}, inp.getType().getElementType())));
        ty = RankedTensorType::get({}, op.getType().getElementType());
      } else {
        ten = stablehlo::constantOp(inp);
      }
      auto out = fromTensor(stablehlo::convertOp(ten, ty));
      if (inp.isSplat())
        out = out.resizeSplat(op.getType());

      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(), out);
      return success();
    }
    return failure();
  }
};

struct ConvertIotaSimplify
    : public CheckedOpRewritePattern<stablehlo::ConvertOp,
                                     ConvertIotaSimplify> {
  using CheckedOpRewritePattern<stablehlo::ConvertOp,
                                ConvertIotaSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp convertOp,
                                    PatternRewriter &rewriter) const {
    auto operand = convertOp.getOperand();
    auto iota = operand.getDefiningOp<stablehlo::IotaOp>();
    if (!iota)
      return failure();

    auto targetType = convertOp.getType();
    if (!targetType.getElementType().isInteger())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::IotaOp>(convertOp, targetType,
                                                   iota.getIotaDimension());
    return success();
  }
};

struct SliceSimplify
    : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceSimplify> {
  using CheckedOpRewritePattern<stablehlo::SliceOp,
                                SliceSimplify>::CheckedOpRewritePattern;

  static size_t getDenseElementBitWidth(Type eltType) {
    // Align the width for complex to 8 to make storage and interpretation
    // easier.
    if (ComplexType comp = llvm::dyn_cast<ComplexType>(eltType))
      return llvm::alignTo<8>(getDenseElementBitWidth(comp.getElementType())) *
             2;
    if (eltType.isIndex())
      return IndexType::kInternalStorageBitWidth;
    return eltType.getIntOrFloatBitWidth();
  }

  static size_t getDenseElementStorageWidth(size_t origWidth) {
    return origWidth == 1 ? origWidth : llvm::alignTo<8>(origWidth);
  }
  static size_t getDenseElementStorageWidth(Type elementType) {
    return getDenseElementStorageWidth(getDenseElementBitWidth(elementType));
  }

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr inp;
    matchPattern(op->getOperand(0), m_Constant(&inp));
    if (inp) {
      DenseElementsAttr out;
      if (inp.isSplat()) {
        out = inp.resizeSplat(op.getType());
      } else {
        // See
        // https://github.com/EnzymeAD/Reactant.jl/pull/1740#issuecomment-3393612379
        // for why this is commented out
        // bool contiguous = true;
        // size_t offset = 0;
        // auto inshape = op.getOperand().getType().getShape();
        // auto outshape = op.getType().getShape();
        // size_t total = 1;
        // for (int i = 0; i < inshape.size(); i++) {
        //   if (op.getStrides()[i] != 1) {
        //     contiguous = false;
        //   }
        //   auto start = op.getStartIndices()[i];
        //   auto lim = op.getLimitIndices()[i];
        //   if (start != 0 || lim != inshape[i]) {
        //     if (offset != 0) {
        //       contiguous = false;
        //     }
        //   }
        //   offset *= inshape[i];
        //   offset += start;
        //   total *= outshape[i];
        // }
        // auto elementType = op.getOperand().getType().getElementType();
        // auto bw = getDenseElementStorageWidth(elementType);
        // if (contiguous && bw != 1) {
        //   const char *elementPtr = inp.getRawData().data() + (bw / 8) *
        //   offset;

        //   auto values = ArrayRef((char *)elementPtr, (bw / 8) * total);
        //   out =
        //       DenseIntOrFPElementsAttr::getFromRawBuffer(op.getType(),
        //       values);
        // } else {
        out = fromTensor(stablehlo::sliceOp(
            stablehlo::constantOp(inp), stablehlo::Sizes(op.getStartIndices()),
            stablehlo::Sizes(op.getStrides()), op.getType()));
        // }
      }
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(), out);
      return success();
    }

    return failure();
  }
};

struct BroadcastInDimSimplify
    : public CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                                     BroadcastInDimSimplify> {
  using CheckedOpRewritePattern<
      stablehlo::BroadcastInDimOp,
      BroadcastInDimSimplify>::CheckedOpRewritePattern;

  size_t max_constant_expansion;
  BroadcastInDimSimplify(size_t max_constant_expansion, MLIRContext *context,
                         PatternBenefit benefit = 1,
                         ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr inp;
    matchPattern(op->getOperand(0), m_Constant(&inp));
    if (inp) {
      DenseElementsAttr out;
      if (inp.isSplat()) {
        out = inp.resizeSplat(op.getType());
      } else {
        size_t size = 1;
        for (auto sz : op.getType().getShape())
          size *= sz;
        if (size >= max_constant_expansion)
          return failure();
        auto ten = stablehlo::constantOp(inp);
        out = fromTensor(stablehlo::broadcastInDimOp(
            ten, stablehlo::Axes(op.getBroadcastDimensions()), op.getType()));
      }

      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(), out);
      return success();
    }

    return failure();
  }
};

static bool getIotaLikeBounds(const enzyme::IotaLikeTensor &iotaLike,
                              double &min_offset, double &max_offset) {
  auto startOpt = enzyme::getDoubleFromAttr(iotaLike.start);
  auto scaleOpt = enzyme::getDoubleFromAttr(iotaLike.scale);
  if (!startOpt || !scaleOpt)
    return false;
  double startV = *startOpt;
  double scaleV = *scaleOpt;

  auto T = dyn_cast<RankedTensorType>(iotaLike.tensorType);
  if (!T)
    return false;

  double max_index = T.getShape()[iotaLike.dimension] - 1;
  double offset1 = startV;
  double offset2 = startV + max_index * scaleV;
  min_offset = std::min(offset1, offset2);
  max_offset = std::max(offset1, offset2);
  return true;
}

template <typename OpTy>
struct MinMaxIotaConstSimplify
    : public CheckedOpRewritePattern<OpTy, MinMaxIotaConstSimplify<OpTy>> {
  using CheckedOpRewritePattern<
      OpTy, MinMaxIotaConstSimplify<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsIota = enzyme::detectIotaLikeTensor(lhs);
    auto rhsIota = enzyme::detectIotaLikeTensor(rhs);

    if ((!lhsIota && !rhsIota) || (lhsIota && rhsIota))
      return failure();

    auto iotaLike = lhsIota ? *lhsIota : *rhsIota;
    Value iotaVal = lhsIota ? lhs : rhs;
    Value cst = lhsIota ? rhs : lhs;

    APInt cstAPInt;
    APFloat cstAPFloat(0.0);
    double cstF;
    if (matchPattern(cst, m_ConstantInt(&cstAPInt))) {
      cstF = cstAPInt.getSExtValue();
    } else if (matchPattern(cst, m_ConstantFloat(&cstAPFloat))) {
      cstF = cstAPFloat.convertToDouble();
    } else {
      return failure();
    }

    double min_offset, max_offset;
    if (!getIotaLikeBounds(iotaLike, min_offset, max_offset))
      return failure();

    if (std::is_same_v<OpTy, stablehlo::MaxOp>) {
      if (cstF <= min_offset) {
        rewriter.replaceOp(op, iotaVal);
        return success();
      }
      if (cstF >= max_offset) {
        rewriter.replaceOp(op, cst);
        return success();
      }
    } else {
      if (cstF <= min_offset) {
        rewriter.replaceOp(op, cst);
        return success();
      }
      if (cstF >= max_offset) {
        rewriter.replaceOp(op, iotaVal);
        return success();
      }
    }

    return failure();
  }
};

struct ClampIotaConstSimplify
    : public CheckedOpRewritePattern<stablehlo::ClampOp,
                                     ClampIotaConstSimplify> {
  using CheckedOpRewritePattern<
      stablehlo::ClampOp, ClampIotaConstSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ClampOp op,
                                    PatternRewriter &rewriter) const {
    auto min_cst = op.getMin();
    auto max_cst = op.getMax();
    auto operand = op.getOperand();

    auto iotaLike = enzyme::detectIotaLikeTensor(operand);
    if (!iotaLike)
      return failure();

    APInt minAPInt, maxAPInt;
    APFloat minAPFloat(0.0), maxAPFloat(0.0);
    double minF, maxF;

    if (matchPattern(min_cst, m_ConstantInt(&minAPInt)))
      minF = minAPInt.getSExtValue();
    else if (matchPattern(min_cst, m_ConstantFloat(&minAPFloat)))
      minF = minAPFloat.convertToDouble();
    else
      return failure();

    if (matchPattern(max_cst, m_ConstantInt(&maxAPInt)))
      maxF = maxAPInt.getSExtValue();
    else if (matchPattern(max_cst, m_ConstantFloat(&maxAPFloat)))
      maxF = maxAPFloat.convertToDouble();
    else
      return failure();

    double iota_min, iota_max;
    if (!getIotaLikeBounds(*iotaLike, iota_min, iota_max))
      return failure();

    if (minF <= iota_min && maxF >= iota_max) {
      rewriter.replaceOp(op, operand);
      return success();
    }
    if (iota_max <= minF) {
      rewriter.replaceOp(op, min_cst);
      return success();
    }
    if (iota_min >= maxF) {
      rewriter.replaceOp(op, max_cst);
      return success();
    }

    return failure();
  }
};

struct CompareIotaConstSimplify
    : public CheckedOpRewritePattern<stablehlo::CompareOp,
                                     CompareIotaConstSimplify> {
  using CheckedOpRewritePattern<
      stablehlo::CompareOp, CompareIotaConstSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp cmpOp,
                                    PatternRewriter &rewriter) const {
    auto lhs = cmpOp.getLhs();
    auto rhs = cmpOp.getRhs();

    auto getIota = [](Value v) -> stablehlo::IotaOp {
      if (auto iota = v.getDefiningOp<stablehlo::IotaOp>())
        return iota;
      if (auto conv = v.getDefiningOp<stablehlo::ConvertOp>()) {
        if (conv.getType().getElementType().isInteger() &&
            conv.getOperand().getType().getElementType().isInteger()) {
          return conv.getOperand().getDefiningOp<stablehlo::IotaOp>();
        }
      }
      return nullptr;
    };
    auto lhsIota = getIota(lhs);
    auto rhsIota = getIota(rhs);

    if ((!lhsIota && !rhsIota) || (lhsIota && rhsIota))
      return rewriter.notifyMatchFailure(cmpOp, "Requires single iota user");

    auto iota = lhsIota ? lhsIota : rhsIota;
    Value cst = lhsIota ? rhs : lhs;

    APInt cstAPInt;
    if (!matchPattern(cst, m_ConstantInt(&cstAPInt)))
      return rewriter.notifyMatchFailure(cmpOp, "Non-constant comparison");

    std::optional<stablehlo::ComparisonType> compType = cmpOp.getCompareType();

    int64_t cstI = compType == stablehlo::ComparisonType::SIGNED
                       ? cstAPInt.getSExtValue()
                       : cstAPInt.getZExtValue();
    auto dir = cmpOp.getComparisonDirection();

    auto T = cast<RankedTensorType>(iota.getType());

    ssize_t max_offset = T.getShape()[iota.getIotaDimension()] - 1;
    ssize_t min_offset = 0;

    if (lhs != iota) {
      dir = reversedComparisonDirection(dir);
    }

    auto boolType = rewriter.getI1Type();
    int64_t lb = 0, ub = T.getShape()[iota.getIotaDimension()];

    auto padInner = [&](bool valueIn, int64_t cstI, int64_t iotaDim) {
      SmallVector<int64_t> shape(T.getShape().begin(), T.getShape().end());
      SmallVector<int64_t> slow(T.getShape().size(), 0);
      SmallVector<int64_t> shigh(T.getShape().size(), 0);
      SmallVector<int64_t> sint(T.getShape().size(), 0);

      shape[iotaDim] = 1;
      slow[iotaDim] = cstI;
      shigh[iotaDim] = T.getShape()[iotaDim] - (cstI + 1);

      Value innerValue = stablehlo::ConstantOp::create(
          rewriter, cmpOp.getLoc(),
          SplatElementsAttr::get(RankedTensorType::get(shape, boolType),
                                 rewriter.getBoolAttr(valueIn)));
      Value paddingValue = stablehlo::ConstantOp::create(
          rewriter, cmpOp.getLoc(),
          SplatElementsAttr::get(RankedTensorType::get({}, boolType),
                                 rewriter.getBoolAttr(!valueIn)));

      rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
          cmpOp, innerValue, paddingValue, slow, shigh, sint);
    };

    switch (dir) {
    case stablehlo::ComparisonDirection::EQ:
      if (cstI < lb || cstI >= ub) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(false)));
        return success();
      }

      padInner(true, cstI, iota.getIotaDimension());
      return success();
    case stablehlo::ComparisonDirection::NE:
      if (cstI < lb || cstI >= ub) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(true)));
        return success();
      }

      padInner(false, cstI, iota.getIotaDimension());
      return success();
    case stablehlo::ComparisonDirection::LE:
    case stablehlo::ComparisonDirection::LT:
    case stablehlo::ComparisonDirection::GE:
    case stablehlo::ComparisonDirection::GT: {
      // iota <= cst [0, 1, 2, 3] .<= 2 -> [1, 1, 1, 0]

      // [0, 1, 2, 3] < 4   => all true if max_offset < cst
      if (dir == stablehlo::ComparisonDirection::LT && max_offset < cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(true)));
        return success();
      }
      // [0, 1, 2, 3] <= 3  => all true if max_offset <= cst
      if (dir == stablehlo::ComparisonDirection::LE && max_offset <= cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(true)));
        return success();
      }
      // [0, 1, 2, 3] > -1  => all true if min_offset > cst
      if (dir == stablehlo::ComparisonDirection::GT && min_offset > cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(true)));
        return success();
      }
      // [0, 1, 2, 3] >= 0  => all true if min_offset > cst
      if (dir == stablehlo::ComparisonDirection::GE && min_offset >= cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(true)));
        return success();
      }

      // [0, 1, 2, 3] < 0   => all false if min_offset >= cst
      if (dir == stablehlo::ComparisonDirection::LT && min_offset >= cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(false)));
        return success();
      }

      // [0, 1, 2, 3] <= -1  => all false if min_offset > cst
      if (dir == stablehlo::ComparisonDirection::LE && min_offset > cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(false)));
        return success();
      }

      // [0, 1, 2, 3] > 3   => all false if max_offset <= cst
      if (dir == stablehlo::ComparisonDirection::GT && max_offset <= cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(false)));
        return success();
      }

      // [0, 1, 2, 3] >= 4   => all false if max_offset < cst
      if (dir == stablehlo::ComparisonDirection::GE && max_offset < cstI) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, SplatElementsAttr::get(cmpOp.getType(),
                                          rewriter.getBoolAttr(false)));
        return success();
      }

      bool isLess = dir == stablehlo::ComparisonDirection::LE ||
                    dir == stablehlo::ComparisonDirection::LT;

      SmallVector<int64_t> leftShape;
      SmallVector<int64_t> rightShape;

      for (auto [i, S] : llvm::enumerate(T.getShape())) {
        if (i == iota.getIotaDimension()) {

          // [0, 1, 2, 3] ?= 2

          // is <= 2, left = [0, 1, 2]
          // is < 2, left = [0, 1]
          // is >= 2, left = [0, 1]
          // is > 2, left = [0, 1, 2]
          int left = (dir == stablehlo::ComparisonDirection::LE ||
                      dir == stablehlo::ComparisonDirection::GT)
                         ? cstI + 1
                         : cstI;
          leftShape.push_back(left);
          rightShape.push_back(S - left);
          continue;
        }

        leftShape.push_back(S);
        rightShape.push_back(S);
      }

      Value ops[] = {
          stablehlo::ConstantOp::create(
              rewriter, iota.getLoc(),
              SplatElementsAttr::get(RankedTensorType::get(leftShape, boolType),
                                     rewriter.getBoolAttr(isLess))),
          stablehlo::ConstantOp::create(
              rewriter, iota.getLoc(),
              SplatElementsAttr::get(
                  RankedTensorType::get(rightShape, boolType),
                  rewriter.getBoolAttr(!isLess)))};
      rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
          cmpOp, cmpOp.getType(), ValueRange(ops), iota.getIotaDimension());
      return success();
    }
    default:
      // TODO: other directions
      break;
    }

    return failure();
  }
};

struct CompareAbs
    : public CheckedOpRewritePattern<stablehlo::CompareOp, CompareAbs> {
  using CheckedOpRewritePattern<stablehlo::CompareOp,
                                CompareAbs>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp cmpOp,
                                    PatternRewriter &rewriter) const {
    for (int i = 0; i < 2; i++) {
      auto operand = cmpOp->getOperand(i);
      auto abs = operand.getDefiningOp<stablehlo::AbsOp>();
      if (!abs)
        continue;
      if (!enzyme::isZero(cmpOp->getOperand(1 - i)))
        continue;

      auto dir = cmpOp.getComparisonDirection();
      if (i == 1) {
        dir = reversedComparisonDirection(dir);
      }
      // now its always abs ?= 0

      // abs(x) < 0 -> false
      if (dir == stablehlo::ComparisonDirection::LT) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, cmpOp.getType(),
            SplatElementsAttr::get(cmpOp.getType(),
                                   rewriter.getBoolAttr(false)));
        return success();
      }
      // abs(x) <= 0 -> x == 0
      if (dir == stablehlo::ComparisonDirection::LE) {
        rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
            cmpOp, abs.getOperand(), cmpOp->getOperand(1 - i),
            stablehlo::ComparisonDirection::EQ);
        return success();
      }
      // abs(x) >= 0 -> true
      if (dir == stablehlo::ComparisonDirection::GE) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            cmpOp, cmpOp.getType(),
            SplatElementsAttr::get(cmpOp.getType(),
                                   rewriter.getBoolAttr(true)));
        return success();
      }
      // abs(x) > 0 -> x != 0
      if (dir == stablehlo::ComparisonDirection::GT) {
        rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
            cmpOp, abs.getOperand(), cmpOp->getOperand(1 - i),
            stablehlo::ComparisonDirection::NE);
        return success();
      }

      // abs(x) == 0 -> x == 0
      if (dir == stablehlo::ComparisonDirection::EQ) {
        rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
            cmpOp, abs.getOperand(), cmpOp->getOperand(1 - i),
            stablehlo::ComparisonDirection::EQ);
        return success();
      }

      // abs(x) != 0 -> x != 0
      if (dir == stablehlo::ComparisonDirection::NE) {
        rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
            cmpOp, abs.getOperand(), cmpOp->getOperand(1 - i),
            stablehlo::ComparisonDirection::NE);
        return success();
      }
    }
    return failure();
  }
};

struct CompareMul
    : public CheckedOpRewritePattern<stablehlo::CompareOp, CompareMul> {
  using CheckedOpRewritePattern<stablehlo::CompareOp,
                                CompareMul>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp cmpOp,
                                    PatternRewriter &rewriter) const {
    for (int i = 0; i < 2; i++) {
      auto operand = cmpOp->getOperand(i);
      auto mul = operand.getDefiningOp<stablehlo::MulOp>();
      if (!mul)
        continue;
      if (!enzyme::isZero(cmpOp->getOperand(1 - i)))
        continue;

      for (int j = 0; j < 2; j++) {
        SplatElementsAttr splat;
        if (!matchPattern(mul->getOperand(j), m_Constant(&splat))) {
          continue;
        }

        bool isNegative = false;

        auto attr = splat.getSplatValue<Attribute>();
        if (auto fp = dyn_cast<FloatAttr>(attr)) {
          if (fp.getValue().isZero())
            continue;
          isNegative = fp.getValue().isNegative();
        } else if (auto fp = dyn_cast<IntegerAttr>(attr)) {
          if (fp.getValue().isZero())
            continue;
          isNegative = fp.getValue().isNegative();
        } else {
          continue;
        }
        auto dir = cmpOp.getComparisonDirection();
        if (i == 1) {
          dir = reversedComparisonDirection(dir);
        }

        // now its always mul(x, cst) ?= 0

        // mul(x, cst) < 0 ->
        //    if cst > 0,   x < 0
        //    elif cst < 0, x > 0
        if (dir == stablehlo::ComparisonDirection::LT) {
          rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
              cmpOp, mul->getOperand(1 - j), cmpOp->getOperand(1 - i),
              isNegative ? stablehlo::ComparisonDirection::GT
                         : stablehlo::ComparisonDirection::LT);
          return success();
        }
        // mul(x, cst) <= 0 ->
        //    if cst > 0,   x <= 0
        //    elif cst < 0, x >= 0
        if (dir == stablehlo::ComparisonDirection::LE) {
          rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
              cmpOp, mul->getOperand(1 - j), cmpOp->getOperand(1 - i),
              isNegative ? stablehlo::ComparisonDirection::GE
                         : stablehlo::ComparisonDirection::LE);
          return success();
        }

        // mul(x, cst) > 0 ->
        //    if cst > 0,   x > 0
        //    elif cst < 0, x < 0
        if (dir == stablehlo::ComparisonDirection::GT) {
          rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
              cmpOp, mul->getOperand(1 - j), cmpOp->getOperand(1 - i),
              isNegative ? stablehlo::ComparisonDirection::LT
                         : stablehlo::ComparisonDirection::GT);
          return success();
        }
        // mul(x, cst) <= 0 ->
        //    if cst > 0,   x <= 0
        //    elif cst < 0, x >= 0
        if (dir == stablehlo::ComparisonDirection::GE) {
          rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
              cmpOp, mul->getOperand(1 - j), cmpOp->getOperand(1 - i),
              isNegative ? stablehlo::ComparisonDirection::LE
                         : stablehlo::ComparisonDirection::GE);
          return success();
        }

        // mul(x, cst) == 0 -> x == 0
        if (dir == stablehlo::ComparisonDirection::EQ) {
          rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
              cmpOp, mul->getOperand(1 - j), cmpOp->getOperand(1 - i),
              stablehlo::ComparisonDirection::EQ);
          return success();
        }

        // mul(x, cst) != 0 -> x != 0
        if (dir == stablehlo::ComparisonDirection::NE) {
          rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
              cmpOp, mul->getOperand(1 - j), cmpOp->getOperand(1 - i),
              stablehlo::ComparisonDirection::NE);
          return success();
        }
      }
    }
    return failure();
  }
};

struct CompareConvert
    : public CheckedOpRewritePattern<stablehlo::CompareOp, CompareConvert> {
  using CheckedOpRewritePattern<stablehlo::CompareOp,
                                CompareConvert>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp cmpOp,
                                    PatternRewriter &rewriter) const {
    for (int i = 0; i < 2; i++) {
      auto operand = cmpOp->getOperand(i);
      auto conv = operand.getDefiningOp<stablehlo::ConvertOp>();
      if (!conv)
        continue;
      if (!conv.getType().getElementType().isFloat())
        continue;
      if (!conv.getOperand().getType().getElementType().isInteger())
        continue;

      if (!enzyme::isZero(cmpOp->getOperand(1 - i)))
        continue;

      SplatElementsAttr splat;
      if (!matchPattern(cmpOp->getOperand(1 - i), m_Constant(&splat))) {
        continue;
      }

      auto attr = splat.getSplatValue<FloatAttr>().getValue();
      if (!attr.isInteger())
        continue;

      auto newCmp = stablehlo::ConvertOp::create(rewriter, cmpOp.getLoc(),
                                                 conv.getOperand().getType(),
                                                 cmpOp->getOperand(1 - i));

      if (i == 0)
        rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
            cmpOp, conv.getOperand(), newCmp, cmpOp.getComparisonDirection());
      else
        rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
            cmpOp, newCmp, conv.getOperand(), cmpOp.getComparisonDirection());
      return success();
    }
    return failure();
  }
};

bool isAddOfSelects(Value lhs, Value rhs, Value &tval, Value &fval) {
  auto sel1 = lhs.getDefiningOp<stablehlo::SelectOp>();
  if (!sel1)
    return false;
  auto sel2 = rhs.getDefiningOp<stablehlo::SelectOp>();
  if (!sel2)
    return false;

  if (sel1.getPred() != sel2.getPred())
    return false;
  if (sel1.getOnTrue() != sel2.getOnFalse() ||
      sel1.getOnFalse() != sel2.getOnTrue())
    return false;

  tval = sel1.getOnTrue();
  fval = sel2.getOnTrue();
  return true;
}

struct AddSelects
    : public CheckedOpRewritePattern<stablehlo::AddOp, AddSelects> {
  using CheckedOpRewritePattern<stablehlo::AddOp,
                                AddSelects>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AddOp addOp,
                                    PatternRewriter &rewriter) const {
    Value tval, fval;
    if (isAddOfSelects(addOp.getLhs(), addOp.getRhs(), tval, fval)) {
      rewriter.modifyOpInPlace(addOp, [&]() {
        addOp.getLhsMutable().assign(tval);
        addOp.getRhsMutable().assign(fval);
      });
      return success();
    }
    for (int i = 0; i < 2; i++) {
      if (auto sub =
              addOp->getOperand(i).getDefiningOp<stablehlo::SubtractOp>()) {
        if (isAddOfSelects(sub.getLhs(), addOp->getOperand(1 - i), tval,
                           fval)) {
          auto newAdd =
              stablehlo::AddOp::create(rewriter, addOp.getLoc(), tval, fval);
          rewriter.replaceOpWithNewOp<stablehlo::SubtractOp>(addOp, newAdd,
                                                             sub.getRhs());
          return success();
        }
      }
    }
    return failure();
  }
};

struct CompareNegateConstSimplify
    : public CheckedOpRewritePattern<stablehlo::CompareOp,
                                     CompareNegateConstSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp cmpOp,
                                    PatternRewriter &rewriter) const {
    auto lhsNegate = cmpOp.getLhs().getDefiningOp<stablehlo::NegOp>();
    auto rhsNegate = cmpOp.getRhs().getDefiningOp<stablehlo::NegOp>();

    auto shardingAttr = sdy::getShardingPerValue(cmpOp);

    if (lhsNegate && rhsNegate) {
      auto newOp = rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
          cmpOp, rhsNegate.getOperand(), lhsNegate.getOperand(),
          cmpOp.getComparisonDirection(), cmpOp.getCompareTypeAttr());
      if (shardingAttr)
        sdy::setShardings(newOp, shardingAttr);
      return success();
    }

    if (!lhsNegate && !rhsNegate)
      return failure();

    if (lhsNegate) {
      if (auto lhsNegateOp = lhsNegate.getOperand().getDefiningOp()) {
        if (!lhsNegateOp->hasTrait<mlir::OpTrait::ConstantLike>()) {
          if (auto cmpRhsOp = cmpOp.getRhs().getDefiningOp()) {
            if (cmpRhsOp->hasTrait<mlir::OpTrait::ConstantLike>()) {
              auto negConst = stablehlo::NegOp::create(rewriter, cmpOp.getLoc(),
                                                       cmpOp.getRhs());
              auto newOp = rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
                  cmpOp, lhsNegate.getOperand(), negConst,
                  reversedComparisonDirection(cmpOp.getComparisonDirection()),
                  cmpOp.getCompareTypeAttr());
              if (shardingAttr) {
                sdy::setShardings(newOp, shardingAttr);
                sdy::setShardings(negConst, shardingAttr);
              }
              return success();
            }
          }
        }
      }
    }

    if (rhsNegate) {
      if (auto rhsNegateOp = rhsNegate.getOperand().getDefiningOp()) {
        if (!rhsNegateOp->hasTrait<mlir::OpTrait::ConstantLike>()) {
          if (auto cmpLhsOp = cmpOp.getLhs().getDefiningOp()) {
            if (cmpLhsOp->hasTrait<mlir::OpTrait::ConstantLike>()) {
              auto negConst = stablehlo::NegOp::create(rewriter, cmpOp.getLoc(),
                                                       cmpOp.getLhs());
              auto newOp = rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
                  cmpOp, negConst, rhsNegate.getOperand(),
                  reversedComparisonDirection(cmpOp.getComparisonDirection()),
                  cmpOp.getCompareTypeAttr());
              if (shardingAttr) {
                sdy::setShardings(newOp, shardingAttr);
                sdy::setShardings(negConst, shardingAttr);
              }
              return success();
            }
          }
        }
      }
    }

    return failure();
  }
};

// (cst - x) OP cst2  →  x reversed_OP (cst - cst2)
// Both cst and cst2 must be constants; the new SubtractOp will be folded by
// a subsequent constant-propagation pass.
struct CompareSubtractConstSimplify
    : public CheckedOpRewritePattern<stablehlo::CompareOp,
                                     CompareSubtractConstSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp cmpOp,
                                    PatternRewriter &rewriter) const {
    auto lhsSub = cmpOp.getLhs().getDefiningOp<stablehlo::SubtractOp>();
    if (!lhsSub)
      return rewriter.notifyMatchFailure(cmpOp, "lhs is not subtract");

    if (!matchPattern(lhsSub.getLhs(), m_Constant()))
      return rewriter.notifyMatchFailure(cmpOp,
                                         "lhs of subtract is not constant");

    if (matchPattern(lhsSub.getRhs(), m_Constant()))
      return rewriter.notifyMatchFailure(cmpOp,
                                         "rhs of subtract is a constant");

    if (!matchPattern(cmpOp.getRhs(), m_Constant()))
      return rewriter.notifyMatchFailure(cmpOp, "rhs is not constant");

    auto newSub = stablehlo::SubtractOp::create(
        rewriter, cmpOp.getLoc(), lhsSub.getLhs(), cmpOp.getRhs());
    auto reversedDir =
        reversedComparisonDirection(cmpOp.getComparisonDirection());

    auto shardingAttr = sdy::getShardingPerValue(cmpOp);
    auto newOp = rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
        cmpOp, lhsSub.getRhs(), newSub, reversedDir,
        cmpOp.getCompareTypeAttr());
    if (shardingAttr)
      sdy::setShardings(newOp, shardingAttr);
    return success();
  }
};

struct BroadcastIotaSimplify
    : public CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                                     BroadcastIotaSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp broadcast,
                                    PatternRewriter &rewriter) const {
    auto operand = broadcast.getOperand();
    DenseIntElementsAttr input;
    matchPattern(operand, m_Constant(&input));

    auto RTO = cast<RankedTensorType>(operand.getType());
    if (input && !input.isSplat() && RTO.getShape().size() == 1 &&
        RTO.getShape()[0] >= 2) {
      auto elemType = input.getElementType();

      if (auto int_attr_arr = input.tryGetValues<::mlir::IntegerAttr>();
          llvm::succeeded(int_attr_arr)) {
        const auto end = int_attr_arr->end();
        auto curr = int_attr_arr->begin();
        auto next = int_attr_arr->begin();
        if (next++ == end)
          return failure();
        if (next == end)
          return failure();

        const auto start = (*curr).getValue().getSExtValue();
        const auto diff = (*next).getValue().getSExtValue() -
                          (*curr).getValue().getSExtValue();

        auto result_type = broadcast->getResultTypes();
        auto loc = broadcast.getLoc();

        // find the dimension to broadcast in
        auto result_shape =
            cast<mlir::ShapedType>(result_type.front()).getShape();

        if (broadcast.getType().getElementType().isInteger(1)) {
          // true, false, .... false.  -> iota == 0
          if (start != 0 && (*next).getValue().isZero()) {
            bool legal = true;
            for (auto idx = next; idx != end; idx++) {
              if (!(*idx).getValue().isZero()) {
                legal = false;
                break;
              }
            }
            // only 1 at the start
            if (legal) {
              auto ITy = RankedTensorType::get(
                  result_shape, rewriter.getIntegerType(32, false));
              auto iota = stablehlo::IotaOp::create(
                  rewriter, loc, ITy, broadcast.getBroadcastDimensions()[0]);
              auto cmp = stablehlo::CompareOp::create(
                  rewriter, loc, iota,
                  stablehlo::ConstantOp::create(
                      rewriter, loc, ITy, cast<ElementsAttr>(makeAttr(ITy, 0))),
                  stablehlo::ComparisonDirection::EQ);
              rewriter.replaceOp(broadcast, cmp);
              return success();
            }
          }

          // false, true, .... true.  -> iota != 0
          if (start == 0 && !(*next).getValue().isZero()) {
            bool legal = true;
            for (auto idx = next; idx != end; idx++) {
              if ((*idx).getValue().isZero()) {
                legal = false;
                break;
              }
            }
            // only 1 at the start
            if (legal) {
              auto ITy = RankedTensorType::get(
                  result_shape, rewriter.getIntegerType(32, false));
              auto iota = stablehlo::IotaOp::create(
                  rewriter, loc, ITy, broadcast.getBroadcastDimensions()[0]);
              auto cmp = stablehlo::CompareOp::create(
                  rewriter, loc, iota,
                  stablehlo::ConstantOp::create(
                      rewriter, loc, ITy, cast<ElementsAttr>(makeAttr(ITy, 0))),
                  stablehlo::ComparisonDirection::NE);
              rewriter.replaceOp(broadcast, cmp);
              return success();
            }
          }

          // false, .... false, true -> iota == end
          auto lastVal = (*(--int_attr_arr->end())).getInt();
          if (lastVal != 0) {
            bool legal = true;
            for (auto idx = int_attr_arr->begin();;) {
              if (!(*idx).getValue().isZero()) {
                legal = false;
                break;
              }
              idx++;
              auto nextv = idx;
              nextv++;
              if (nextv == end) {
                break;
              }
            }
            // only 1 at the end
            if (legal) {
              auto ITy = RankedTensorType::get(
                  result_shape, rewriter.getIntegerType(32, false));
              auto iota = stablehlo::IotaOp::create(
                  rewriter, loc, ITy, broadcast.getBroadcastDimensions()[0]);
              auto cmp = stablehlo::CompareOp::create(
                  rewriter, loc, iota,
                  stablehlo::ConstantOp::create(
                      rewriter, loc, ITy,
                      cast<ElementsAttr>(makeAttr(ITy, RTO.getShape()[0] - 1))),
                  stablehlo::ComparisonDirection::EQ);
              rewriter.replaceOp(broadcast, cmp);
              return success();
            }
          }

          // true, .... true, false  -> iota != end
          if (lastVal == 0) {
            bool legal = true;
            for (auto idx = int_attr_arr->begin();;) {
              if ((*idx).getValue().isZero()) {
                legal = false;
                break;
              }
              idx++;
              auto nextv = idx;
              nextv++;
              if (nextv == end) {
                break;
              }
            }
            // only 1 at the end
            if (legal) {
              auto ITy = RankedTensorType::get(
                  result_shape, rewriter.getIntegerType(32, false));
              auto iota = stablehlo::IotaOp::create(
                  rewriter, loc, ITy, broadcast.getBroadcastDimensions()[0]);
              auto cmp = stablehlo::CompareOp::create(
                  rewriter, loc, iota,
                  stablehlo::ConstantOp::create(
                      rewriter, loc, ITy,
                      cast<ElementsAttr>(makeAttr(ITy, RTO.getShape()[0] - 1))),
                  stablehlo::ComparisonDirection::NE);
              rewriter.replaceOp(broadcast, cmp);
              return success();
            }
          }

          // false, false, ..., true, true, true, ..., false, false, ...
          if (lastVal == 0 && start == 0) {
            bool legal = true;
            int firstTrue = -1;
            int firstFalseAgain = -1;
            int i = 0;
            for (auto idx = int_attr_arr->begin(); idx != end; i++, idx++) {
              bool isTrue = !(*idx).getValue().isZero();
              // Before the transition.
              if (firstTrue == -1) {
                if (isTrue) {
                  firstTrue = i;
                  continue;
                } else {
                  continue;
                }
              }
              // In the true region
              if (firstFalseAgain == -1) {
                if (isTrue) {
                  continue;
                } else {
                  firstFalseAgain = i;
                  continue;
                }
              }

              // Post transition
              if (isTrue) {
                legal = false;
                break;
              }
              continue;
            }
            if (legal && firstTrue != -1 && firstFalseAgain != -1) {
              auto ITy = RankedTensorType::get(
                  result_shape, rewriter.getIntegerType(32, false));
              auto iota = stablehlo::IotaOp::create(
                  rewriter, loc, ITy, broadcast.getBroadcastDimensions()[0]);
              auto cmp1 = stablehlo::CompareOp::create(
                  rewriter, loc, iota,
                  stablehlo::ConstantOp::create(
                      rewriter, loc, ITy,
                      cast<ElementsAttr>(makeAttr(ITy, firstTrue))),
                  stablehlo::ComparisonDirection::GE);
              auto cmp2 = stablehlo::CompareOp::create(
                  rewriter, loc, iota,
                  stablehlo::ConstantOp::create(
                      rewriter, loc, ITy,
                      cast<ElementsAttr>(makeAttr(ITy, firstFalseAgain))),
                  stablehlo::ComparisonDirection::LT);
              rewriter.replaceOpWithNewOp<stablehlo::AndOp>(broadcast, cmp1,
                                                            cmp2);
              return success();
            }
          }
        }

        if (diff == 0)
          return failure();

        while (next != end) {
          auto curr_diff = (*next).getValue().getSExtValue() -
                           (*curr).getValue().getSExtValue();
          if (curr_diff != diff)
            return failure();
          ++curr;
          ++next;
        }

        // The input is 1D, so getBroadcastDimensions()[0] gives the output
        // dimension that the input data maps to. The iota must vary along
        // that dimension, not the replicated one.
        int32_t broadcast_dim = broadcast.getBroadcastDimensions()[0];

        // build the replacement operations
        auto iota = stablehlo::IotaOp::create(rewriter, loc, result_type,
                                              broadcast_dim);
        auto stride_attr = mlir::DenseElementsAttr::get(
            operand.getType().cloneWith(result_shape, elemType),
            rewriter.getIntegerAttr(elemType, diff));
        auto start_attr = mlir::DenseElementsAttr::get(
            operand.getType().cloneWith(result_shape, elemType),
            rewriter.getIntegerAttr(elemType, start));
        auto stride_const = stablehlo::ConstantOp::create(
            rewriter, loc, result_type, stride_attr);
        auto start_const = stablehlo::ConstantOp::create(
            rewriter, loc, result_type, start_attr);
        auto mul = stablehlo::MulOp::create(rewriter, loc, iota, stride_const);

        rewriter.replaceOpWithNewOp<stablehlo::AddOp>(broadcast, start_const,
                                                      mul);
        return success();
      }
      return failure();
    }

    return failure();
  }
};

struct DotGeneralSimplify
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                     DotGeneralSimplify> {
  using CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                DotGeneralSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    Attribute lhsAttr, rhsAttr;
    matchPattern(op.getLhs(), m_Constant(&lhsAttr));
    matchPattern(op.getRhs(), m_Constant(&rhsAttr));
    if ((lhsAttr && (matchPattern(lhsAttr, m_AnyZeroFloat()) ||
                     matchPattern(lhsAttr, m_AnyZeroComplex()))) ||
        (rhsAttr && (matchPattern(rhsAttr, m_AnyZeroFloat()) ||
                     matchPattern(rhsAttr, m_AnyZeroComplex())))) {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, rewriter.getZeroAttr(op.getType()));
      return success();
    }

    auto dimNumbers = op.getDotDimensionNumbers();
    auto lhsContractingDims = dimNumbers.getLhsContractingDimensions();
    auto rhsContractingDims = dimNumbers.getRhsContractingDimensions();
    auto lhsBatchingDims = dimNumbers.getLhsBatchingDimensions();
    auto rhsBatchingDims = dimNumbers.getRhsBatchingDimensions();
    auto lhsType = cast<RankedTensorType>(op.getLhs().getType());
    auto rhsType = cast<RankedTensorType>(op.getRhs().getType());
    auto lhsShape = lhsType.getShape();
    auto rhsShape = rhsType.getShape();
    auto lhsNonContractingDims = getNonContractingDims(
        lhsBatchingDims, lhsContractingDims, lhsType.getRank());
    auto rhsNonContractingDims = getNonContractingDims(
        rhsBatchingDims, rhsContractingDims, rhsType.getRank());

    Value reduceSumInput;
    SmallVector<int64_t> reduceDims, broadcastDims, broadcastShape;
    SplatElementsAttr splatElementsAttr;

    if (matchPattern(op.getLhs(), m_Constant(&splatElementsAttr))) {
      reduceDims = llvm::to_vector(rhsContractingDims);
      reduceSumInput = op.getRhs();

      SmallVector<int64_t> broadcastDimsTmp(rhsType.getRank(), -1);

      int64_t idx = 0;
      for (auto dim : rhsBatchingDims) {
        broadcastDimsTmp[dim] = idx;
        broadcastShape.push_back(rhsShape[dim]);
        idx++;
      }
      for (auto dim : lhsNonContractingDims) {
        broadcastShape.push_back(lhsShape[dim]);
        idx++;
      }
      for (auto dim : rhsNonContractingDims) {
        broadcastDimsTmp[dim] = idx;
        broadcastShape.push_back(rhsShape[dim]);
        idx++;
      }

      for (auto dim : broadcastDimsTmp) {
        if (dim != -1) {
          broadcastDims.push_back(dim);
        }
      }
    } else if (matchPattern(op.getRhs(), m_Constant(&splatElementsAttr))) {
      reduceDims = llvm::to_vector(lhsContractingDims);
      reduceSumInput = op.getLhs();

      SmallVector<int64_t> broadcastDimsTmp(lhsType.getRank(), -1);

      int64_t idx = 0;
      for (auto dim : lhsBatchingDims) {
        broadcastDimsTmp[dim] = idx;
        broadcastShape.push_back(lhsShape[dim]);
        idx++;
      }
      for (auto dim : lhsNonContractingDims) {
        broadcastDimsTmp[dim] = idx;
        broadcastShape.push_back(lhsShape[dim]);
        idx++;
      }
      for (auto dim : rhsNonContractingDims) {
        broadcastShape.push_back(rhsShape[dim]);
        idx++;
      }

      for (auto dim : broadcastDimsTmp) {
        if (dim != -1) {
          broadcastDims.push_back(dim);
        }
      }
    }

    if (splatElementsAttr) {
      auto inputType = cast<RankedTensorType>(reduceSumInput.getType());
      auto elemType = inputType.getElementType();
      auto rank0Type = RankedTensorType::get({}, elemType);

      auto zero = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), rank0Type,
          cast<ElementsAttr>(makeAttr(rank0Type, 0)));
      auto reduceOp = stablehlo::ReduceOp::create(
          rewriter, op.getLoc(), ValueRange{reduceSumInput}, ValueRange{zero},
          rewriter.getDenseI64ArrayAttr(reduceDims));

      Block &block = reduceOp.getBody().emplaceBlock();
      BlockArgument arg0 = block.addArgument(rank0Type, op.getLoc());
      BlockArgument arg1 = block.addArgument(rank0Type, op.getLoc());

      {
        OpBuilder::InsertionGuard guard(rewriter);
        rewriter.setInsertionPointToStart(&block);

        Value sum = stablehlo::AddOp::create(rewriter, op.getLoc(), arg0, arg1);
        stablehlo::ReturnOp::create(rewriter, op.getLoc(), sum);
      }

      auto bcastOp = stablehlo::BroadcastInDimOp::create(
          rewriter, op.getLoc(),
          RankedTensorType::get(broadcastShape, elemType),
          reduceOp.getResult(0), broadcastDims);

      rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
          op, bcastOp,
          stablehlo::ConstantOp::create(
              rewriter, op.getLoc(), bcastOp.getType(),
              splatElementsAttr.resizeSplat(bcastOp.getType())));
      return success();
    }

    return failure();
  }

private:
  SmallVector<int64_t> getNonContractingDims(ArrayRef<int64_t> batchingDims,
                                             ArrayRef<int64_t> contractingDims,
                                             int64_t rank) const {
    SmallVector<int64_t> nonContractingDims;
    for (int i = 0; i < rank; i++) {
      if (!llvm::is_contained(batchingDims, i) &&
          !llvm::is_contained(contractingDims, i)) {
        nonContractingDims.push_back(i);
      }
    }
    return nonContractingDims;
  }
};

struct TransposeSimplify
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeSimplify> {
  using CheckedOpRewritePattern<stablehlo::TransposeOp,
                                TransposeSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr inp;
    matchPattern(op->getOperand(0), m_Constant(&inp));
    if (inp) {

      DenseElementsAttr out;
      if (inp.isSplat()) {
        out = inp.resizeSplat(op.getType());
      } else {
        out = fromTensor(stablehlo::transposeOp(
            stablehlo::constantOp(inp), stablehlo::Axes(op.getPermutation()),
            op.getType()));
      }
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(), out);
      return success();
    }
    return failure();
  }
};

struct MaxSimplify
    : public CheckedOpRewritePattern<stablehlo::MaxOp, MaxSimplify> {
  using CheckedOpRewritePattern<stablehlo::MaxOp,
                                MaxSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MaxOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getOperand(0) == op.getOperand(1)) {
      rewriter.replaceOp(op, op.getOperand(0));
      return success();
    }

    return failure();
  }
};

struct MinSimplify
    : public CheckedOpRewritePattern<stablehlo::MinOp, MinSimplify> {
  using CheckedOpRewritePattern<stablehlo::MinOp,
                                MinSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MinOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getOperand(0) == op.getOperand(1)) {
      rewriter.replaceOp(op, op.getOperand(0));
      return success();
    }

    return failure();
  }
};

template <typename T>
struct BinBroadcastSplat final
    : CheckedOpRewritePattern<T, BinBroadcastSplat<T>> {
  using CheckedOpRewritePattern<T,
                                BinBroadcastSplat<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    for (int i = 0; i < 2; i++) {
      mlir::Value opi = op->getOperand(i);
      if (auto broadcast = opi.getDefiningOp<stablehlo::BroadcastInDimOp>()) {
        if (isTransposeReshapeLikeBroadcast(broadcast))
          continue;
        SplatElementsAttr other;
        if (matchPattern(op->getOperand(1 - i), m_Constant(&other))) {
          IRMapping map;
          mlir::Value vals[2];
          vals[i] = broadcast.getOperand();
          vals[1 - i] = stablehlo::ConstantOp::create(
              rewriter, op.getLoc(), broadcast.getOperand().getType(),
              other.resizeSplat(broadcast.getOperand().getType()));
          auto pushed = T::create(rewriter, op.getLoc(), vals[0], vals[1]);
          map.map(broadcast.getOperand(), pushed->getResult(0));
          auto bc2 = rewriter.clone(*broadcast, map);
          rewriter.replaceOp(op, bc2);
          return success();
        }
      }
    }
    return failure();
  }
};

struct AllFiniteIsFinite
    : public CheckedOpRewritePattern<stablehlo::IsFiniteOp, AllFiniteIsFinite> {
  using CheckedOpRewritePattern<stablehlo::IsFiniteOp,
                                AllFiniteIsFinite>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IsFiniteOp op,
                                    PatternRewriter &rewriter) const {
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, op.getType(), cast<ElementsAttr>(makeAttr(op.getType(), 1)));
    return success();
  }
};

struct AllFiniteIsInf
    : public CheckedOpRewritePattern<chlo::IsInfOp, AllFiniteIsInf> {
  using CheckedOpRewritePattern<chlo::IsInfOp,
                                AllFiniteIsInf>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::IsInfOp op,
                                    PatternRewriter &rewriter) const {
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, rewriter.getZeroAttr(op.getType()));
    return success();
  }
};

struct AllFiniteIsPosInf
    : public CheckedOpRewritePattern<chlo::IsPosInfOp, AllFiniteIsPosInf> {
  using CheckedOpRewritePattern<chlo::IsPosInfOp,
                                AllFiniteIsPosInf>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::IsPosInfOp op,
                                    PatternRewriter &rewriter) const {
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, rewriter.getZeroAttr(op.getType()));
    return success();
  }
};

struct AllFiniteIsNegInf
    : public CheckedOpRewritePattern<chlo::IsNegInfOp, AllFiniteIsNegInf> {
  using CheckedOpRewritePattern<chlo::IsNegInfOp,
                                AllFiniteIsNegInf>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::IsNegInfOp op,
                                    PatternRewriter &rewriter) const {
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, rewriter.getZeroAttr(op.getType()));
    return success();
  }
};

struct NoNanCompareSimplify
    : public NoNanCheckedOpRewritePattern<stablehlo::CompareOp,
                                          NoNanCompareSimplify> {
  using NoNanCheckedOpRewritePattern<
      stablehlo::CompareOp, NoNanCompareSimplify>::NoNanCheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getLhs() == op.getRhs()) {
      if (op.getComparisonDirection() == stablehlo::ComparisonDirection::EQ) {
        if (canApplyNoNanPattern(allowOnFloatingPointMath, op.getType(),
                                 op.getLhs().getType(), op, rewriter)) {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, op.getType(), cast<ElementsAttr>(makeAttr(op.getType(), 1)));
          return success();
        }
      }
      if (op.getComparisonDirection() == stablehlo::ComparisonDirection::NE) {
        if (canApplyNoNanPattern(allowOnFloatingPointMath, op.getType(),
                                 op.getLhs().getType(), op, rewriter)) {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, op.getType(), cast<ElementsAttr>(makeAttr(op.getType(), 0)));
          return success();
        }
      }
    }
    return failure();
  }
};

struct TransposeTranspose
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeTranspose> {
  using CheckedOpRewritePattern<stablehlo::TransposeOp,
                                TransposeTranspose>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto operand = op.getOperand();

    auto convertOp = operand.getDefiningOp<stablehlo::ConvertOp>();
    if (convertOp) {
      operand = convertOp.getOperand();
    }

    auto definingTranspose = operand.getDefiningOp<stablehlo::TransposeOp>();
    if (!definingTranspose)
      return rewriter.notifyMatchFailure(op, "not a transpose(transpose)");

    llvm::ArrayRef<int64_t> thisPermutation = op.getPermutation();
    llvm::ArrayRef<int64_t> prevPermutation =
        definingTranspose.getPermutation();

    SmallVector<int64_t> newPermutation;
    newPermutation.resize(thisPermutation.size());
    for (unsigned i = 0, e = thisPermutation.size(); i != e; ++i) {
      newPermutation[i] = prevPermutation[thisPermutation[i]];
    }

    if (!convertOp) {
      rewriter.modifyOpInPlace(op, [&]() {
        op.setPermutation(newPermutation);
        op.setOperand(definingTranspose.getOperand());
      });
    } else {
      auto midPerm = stablehlo::TransposeOp::create(
          rewriter, op.getLoc(), definingTranspose.getOperand(),
          newPermutation);
      rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(op, op.getType(),
                                                        midPerm);
    }
    return success();
  }
};

size_t getBitWidth(mlir::Type ty) {
  if (auto CT = dyn_cast<ComplexType>(ty)) {
    return 2 * getBitWidth(CT.getElementType());
  }
  return ty.getIntOrFloatBitWidth();
}

struct TransposeConvert
    : public CheckedOpRewritePattern<stablehlo::ConvertOp, TransposeConvert> {
  using CheckedOpRewritePattern<stablehlo::ConvertOp,
                                TransposeConvert>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    auto resultType = cast<TensorType>(op.getResult().getType());
    auto operandType = cast<TensorType>(op.getOperand().getType());
    if (!resultType.hasStaticShape() || !operandType.hasStaticShape())
      return failure();
    if (resultType.getNumElements() *
            getBitWidth(resultType.getElementType()) >=
        operandType.getNumElements() *
            getBitWidth(operandType.getElementType()))
      return failure();

    auto transpose = op.getOperand().getDefiningOp<stablehlo::TransposeOp>();
    if (!transpose || !llvm::hasSingleElement(transpose->getUsers()))
      return failure();

    auto newConvert = stablehlo::ConvertOp::create(rewriter, op.getLoc(),
                                                   transpose.getOperand(),
                                                   resultType.getElementType());
    auto newTranspose = stablehlo::TransposeOp::create(
        rewriter, transpose.getLoc(), newConvert.getResult(),
        transpose.getPermutation());
    rewriter.replaceOp(op, newTranspose);
    rewriter.eraseOp(transpose);

    return success();
  }
};

struct TransposeDotReorder
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeDotReorder> {
  using CheckedOpRewritePattern<stablehlo::TransposeOp,
                                TransposeDotReorder>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {

    auto operand = op.getOperand();
    auto convert = operand.getDefiningOp<stablehlo::ConvertOp>();
    if (convert) {
      operand = convert.getOperand();
      if (!llvm::hasSingleElement(convert->getUsers()))
        return failure();
    }

    auto dot = operand.getDefiningOp<stablehlo::DotGeneralOp>();
    if (!dot || !llvm::hasSingleElement(dot->getUsers()))
      return failure();

    auto perm = op.getPermutation();
    auto dimensionNumbers = dot.getDotDimensionNumbers();

    size_t permidx = 0;

    for (size_t bidx = 0,
                end = dimensionNumbers.getLhsBatchingDimensions().size();
         bidx < end; bidx++) {
      if (perm[permidx] != bidx)
        return failure();
      permidx++;
    }

    size_t numLHSResults =
        dot.getLhs().getType().getShape().size() -
        dimensionNumbers.getLhsBatchingDimensions().size() -
        dimensionNumbers.getLhsContractingDimensions().size();

    {
      size_t residx = 0;
      for (size_t ridx = 0, end = dot.getRhs().getType().getShape().size();
           ridx < end; ridx++) {
        if (llvm::is_contained(dimensionNumbers.getRhsBatchingDimensions(),
                               ridx))
          continue;
        if (llvm::is_contained(dimensionNumbers.getRhsContractingDimensions(),
                               ridx))
          continue;
        if (perm[permidx] !=
            dimensionNumbers.getLhsBatchingDimensions().size() + numLHSResults +
                residx)
          return failure();
        permidx++;
        residx++;
      }
    }

    {
      size_t residx = 0;
      for (size_t lidx = 0, end = dot.getLhs().getType().getShape().size();
           lidx < end; lidx++) {
        if (llvm::is_contained(dimensionNumbers.getLhsBatchingDimensions(),
                               lidx))
          continue;
        if (llvm::is_contained(dimensionNumbers.getLhsContractingDimensions(),
                               lidx))
          continue;
        if (perm[permidx] !=
            dimensionNumbers.getLhsBatchingDimensions().size() + residx)
          return failure();
        permidx++;
        residx++;
      }
    }

    auto ndim = stablehlo::DotDimensionNumbersAttr::get(
        dimensionNumbers.getContext(),
        dimensionNumbers.getRhsBatchingDimensions(),
        dimensionNumbers.getLhsBatchingDimensions(),
        dimensionNumbers.getRhsContractingDimensions(),
        dimensionNumbers.getLhsContractingDimensions());
    if (!convert) {
      rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
          op, op.getType(), dot.getRhs(), dot.getLhs(), ndim,
          dot.getPrecisionConfigAttr(), dot.getAlgorithmAttr());
    } else {
      auto middot = stablehlo::DotGeneralOp::create(
          rewriter, op.getLoc(),
          RankedTensorType::get(op.getType().getShape(),
                                dot.getType().getElementType()),
          dot.getRhs(), dot.getLhs(), ndim, dot.getPrecisionConfigAttr(),
          dot.getAlgorithmAttr());
      rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(op, op.getType(),
                                                        middot);
    }
    return success();
  }
};

struct TransposeReduce
    : public CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeReduce> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp transpose,
                                    PatternRewriter &rewriter) const {
    auto operand = transpose.getOperand();
    auto reduce = operand.getDefiningOp<stablehlo::ReduceOp>();
    if (!reduce)
      return failure();

    for (auto result : reduce->getResults())
      if (!llvm::hasSingleElement(result.getUsers()))
        return failure();

    unsigned resultNum = std::distance(
        reduce.getResults().begin(), llvm::find(reduce.getResults(), operand));

    auto reduceDims = reduce.getDimensions();
    auto reduceInput = reduce.getInputs()[resultNum];
    auto reduceInputType = dyn_cast<RankedTensorType>(reduceInput.getType());
    if (!reduceInputType)
      return rewriter.notifyMatchFailure(reduce, "Reduce input not tensor");

    auto transposePermutation = transpose.getPermutation();

    SmallVector<int64_t> newTransposePermutation(transposePermutation);

    for (int64_t reduceDim : reduceDims) {
      for (auto &transposeDim : newTransposePermutation) {
        if (reduceDim <= transposeDim)
          transposeDim++;
      }
      newTransposePermutation.insert(
          std::next(newTransposePermutation.begin(), reduceDim), reduceDim);
    }

    rewriter.setInsertionPoint(reduce);
    auto newTransposeOp = stablehlo::TransposeOp::create(
        rewriter, transpose.getLoc(), reduceInput, newTransposePermutation);

    SmallVector<Type> newReduceResultTypes(reduce.getResultTypes());
    newReduceResultTypes[resultNum] = transpose.getResult().getType();
    SmallVector<Value> newReduceInputs(reduce.getInputs());
    newReduceInputs[resultNum] = newTransposeOp.getResult();

    auto newReduce = stablehlo::ReduceOp::create(
        rewriter, reduce.getLoc(), newReduceResultTypes, newReduceInputs,
        reduce.getInitValues(), reduceDims);
    rewriter.inlineRegionBefore(reduce.getRegion(), newReduce.getRegion(),
                                newReduce.getRegion().begin());
    for (auto [i, oldRes, newRes] :
         llvm::enumerate(reduce.getResults(), newReduce.getResults())) {
      if (i == resultNum) {
        rewriter.replaceOp(transpose, newRes);
      } else {
        rewriter.replaceAllUsesWith(oldRes, newRes);
      }
    }
    rewriter.eraseOp(reduce);

    return success();
  }
};

struct TransposeConvolution
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeConvolution> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp transpose,
                                    PatternRewriter &rewriter) const {
    auto operand = transpose.getOperand();
    auto conv = operand.getDefiningOp<stablehlo::ConvolutionOp>();
    if (!conv || !llvm::hasSingleElement(operand.getUsers()))
      return failure();

    auto permutation = transpose.getPermutation();
    SmallVector<int64_t> mapping(permutation.size(), -1);
    for (int64_t i = 0; i < permutation.size(); i++)
      mapping[permutation[i]] = i;

    auto dimensionNumbers = conv.getDimensionNumbers();

    int64_t outputBatchDimension =
        mapping[dimensionNumbers.getOutputBatchDimension()];
    int64_t outputFeatureDimension =
        mapping[dimensionNumbers.getOutputFeatureDimension()];

    SmallVector<int64_t> outputSpatialDimensions(
        dimensionNumbers.getOutputSpatialDimensions().begin(),
        dimensionNumbers.getOutputSpatialDimensions().end());
    for (auto &dim : outputSpatialDimensions)
      dim = mapping[dim];

    auto newDimensionNumbers = stablehlo::ConvDimensionNumbersAttr::get(
        dimensionNumbers.getContext(),
        dimensionNumbers.getInputBatchDimension(),
        dimensionNumbers.getInputFeatureDimension(),
        dimensionNumbers.getInputSpatialDimensions(),
        dimensionNumbers.getKernelInputFeatureDimension(),
        dimensionNumbers.getKernelOutputFeatureDimension(),
        dimensionNumbers.getKernelSpatialDimensions(), outputBatchDimension,
        outputFeatureDimension, outputSpatialDimensions);

    rewriter.replaceOpWithNewOp<stablehlo::ConvolutionOp>(
        transpose, transpose.getType(), conv.getLhs(), conv.getRhs(),
        conv.getWindowStridesAttr(), conv.getPaddingAttr(),
        conv.getLhsDilationAttr(), conv.getRhsDilationAttr(),
        conv.getWindowReversalAttr(), newDimensionNumbers,
        conv.getFeatureGroupCountAttr(), conv.getBatchGroupCountAttr(),
        conv.getPrecisionConfigAttr());

    return success();
  }
};

struct ConvolutionTranspose
    : public CheckedOpRewritePattern<stablehlo::ConvolutionOp,
                                     ConvolutionTranspose> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvolutionOp conv,
                                    PatternRewriter &rewriter) const {
    auto lhs_trans = conv.getLhs().getDefiningOp<stablehlo::TransposeOp>();
    auto rhs_trans = conv.getRhs().getDefiningOp<stablehlo::TransposeOp>();
    if (!lhs_trans && !rhs_trans)
      return failure();

    auto dim = conv.getDimensionNumbers();
    int64_t inputBatchDimension = dim.getInputBatchDimension();
    int64_t inputFeatureDimension = dim.getInputFeatureDimension();
    SmallVector<int64_t> inputSpatialDimensions(
        dim.getInputSpatialDimensions().begin(),
        dim.getInputSpatialDimensions().end());

    int64_t kernelInputFeatureDimension = dim.getKernelInputFeatureDimension();
    int64_t kernelOutputFeatureDimension =
        dim.getKernelOutputFeatureDimension();
    SmallVector<int64_t> kernelSpatialDimensions(
        dim.getKernelSpatialDimensions().begin(),
        dim.getKernelSpatialDimensions().end());

    if (lhs_trans) {
      inputBatchDimension = lhs_trans.getPermutation()[inputBatchDimension];
      inputFeatureDimension = lhs_trans.getPermutation()[inputFeatureDimension];
      for (auto &dim : inputSpatialDimensions)
        dim = lhs_trans.getPermutation()[dim];
    }

    if (rhs_trans) {
      kernelInputFeatureDimension =
          rhs_trans.getPermutation()[kernelInputFeatureDimension];
      kernelOutputFeatureDimension =
          rhs_trans.getPermutation()[kernelOutputFeatureDimension];
      for (auto &dim : kernelSpatialDimensions)
        dim = rhs_trans.getPermutation()[dim];
    }

    auto ndim = stablehlo::ConvDimensionNumbersAttr::get(
        dim.getContext(), inputBatchDimension, inputFeatureDimension,
        inputSpatialDimensions, kernelInputFeatureDimension,
        kernelOutputFeatureDimension, kernelSpatialDimensions,
        dim.getOutputBatchDimension(), dim.getOutputFeatureDimension(),
        dim.getOutputSpatialDimensions());

    rewriter.replaceOpWithNewOp<stablehlo::ConvolutionOp>(
        conv, conv.getType(),
        lhs_trans ? lhs_trans.getOperand() : conv.getLhs(),
        rhs_trans ? rhs_trans.getOperand() : conv.getRhs(),

        conv.getWindowStridesAttr(), conv.getPaddingAttr(),
        conv.getLhsDilationAttr(), conv.getRhsDilationAttr(),
        conv.getWindowReversalAttr(), ndim, conv.getFeatureGroupCountAttr(),
        conv.getBatchGroupCountAttr(), conv.getPrecisionConfigAttr());
    return success();
  }
};

// transpose(einsum) -> einsum
struct TransposeEinsum
    : public CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeEinsum> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp transpose,
                                    PatternRewriter &rewriter) const {
    auto operand = transpose.getOperand();
    auto einsum = operand.getDefiningOp<stablehlo::EinsumOp>();
    if (!einsum || !llvm::hasSingleElement(operand.getUsers()))
      return failure();

    auto einsumConfig = einsum.getEinsumConfig();
    auto arrowPos = einsumConfig.find("->");

    if (arrowPos == StringRef::npos)
      return failure();

    auto permutation = transpose.getPermutation();

    if (einsumConfig.size() - (arrowPos + 2) < permutation.size())
      return failure();

    auto newEinsumConfig = std::string(einsumConfig.str());
    for (int i = 0; i < permutation.size(); ++i) {
      newEinsumConfig[arrowPos + 2 + i] =
          einsumConfig[arrowPos + 2 + permutation[i]];
    }

    rewriter.modifyOpInPlace(einsum, [&einsum, &transpose, newEinsumConfig] {
      einsum.setEinsumConfig(
          StringAttr::get(einsum.getContext(), newEinsumConfig));
      einsum.getResult().setType(transpose.getType());
    });
    rewriter.replaceAllUsesWith(transpose.getResult(), einsum.getResult());
    return success();
  }
};

// einsum(transpose(x), transpose(y)) -> einsum(x, y)
struct EinsumTranspose
    : public CheckedOpRewritePattern<stablehlo::EinsumOp, EinsumTranspose> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::EinsumOp einsum,
                                    PatternRewriter &rewriter) const {
    llvm::StringRef einsumConfig = einsum.getEinsumConfig();

    auto lhs_trans = einsum.getLhs().getDefiningOp<stablehlo::TransposeOp>();
    auto rhs_trans = einsum.getRhs().getDefiningOp<stablehlo::TransposeOp>();
    if (!lhs_trans && !rhs_trans)
      return failure();

    size_t commaPos = einsumConfig.find(",");
    size_t arrowPos = einsumConfig.find("->");
    if (commaPos != einsum.getLhs().getType().getRank() ||
        einsumConfig.size() - commaPos < einsum.getRhs().getType().getRank() ||
        (arrowPos != StringRef::npos &&
         arrowPos - commaPos < einsum.getRhs().getType().getRank()))
      return failure();

    auto newEinsumConfig = std::string(einsumConfig.str());

    if (lhs_trans) {
      for (int i = 0; i < commaPos; ++i) {
        newEinsumConfig[i] = einsumConfig[lhs_trans.getPermutation()[i]];
      }
    }

    if (rhs_trans) {
      int64_t rhsRank = einsum.getRhs().getType().getRank();
      for (int i = 0; i < rhsRank; ++i) {
        newEinsumConfig[commaPos + 1 + i] =
            einsumConfig[commaPos + 1 + rhs_trans.getPermutation()[i]];
      }
    }

    rewriter.replaceOpWithNewOp<stablehlo::EinsumOp>(
        einsum, einsum.getType(),
        lhs_trans ? lhs_trans.getOperand() : einsum.getLhs(),
        rhs_trans ? rhs_trans.getOperand() : einsum.getRhs(),
        StringAttr::get(einsum.getContext(), newEinsumConfig));
    return success();
  }
};

struct DotTranspose
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp, DotTranspose> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp dot,
                                    PatternRewriter &rewriter) const {

    auto dim = dot.getDotDimensionNumbers();
    size_t numLHSResults = dot.getLhs().getType().getRank() -
                           dim.getLhsBatchingDimensions().size() -
                           dim.getLhsContractingDimensions().size();

    size_t numRHSResults = dot.getRhs().getType().getRank() -
                           dim.getRhsBatchingDimensions().size() -
                           dim.getRhsContractingDimensions().size();

    auto lhsTrans = dot.getLhs().getDefiningOp<stablehlo::TransposeOp>();
    bool fuseLhs = false;
    if (lhsTrans)
      fuseLhs = numLHSResults <= 1; // otherwise we need to transpose the result

    auto rhsTrans = dot.getRhs().getDefiningOp<stablehlo::TransposeOp>();
    bool fuseRhs = false;
    if (rhsTrans)
      fuseRhs = numRHSResults <= 1; // otherwise we need to transpose the result

    if (!fuseLhs && !fuseRhs)
      return failure();

    SmallVector<int64_t> lhsBatch(dim.getLhsBatchingDimensions().begin(),
                                  dim.getLhsBatchingDimensions().end());
    SmallVector<int64_t> rhsBatch(dim.getRhsBatchingDimensions().begin(),
                                  dim.getRhsBatchingDimensions().end());

    SmallVector<int64_t> lhsContract(dim.getLhsContractingDimensions().begin(),
                                     dim.getLhsContractingDimensions().end());
    SmallVector<int64_t> rhsContract(dim.getRhsContractingDimensions().begin(),
                                     dim.getRhsContractingDimensions().end());

    if (fuseLhs) {
      for (auto &dim : lhsBatch)
        dim = lhsTrans.getPermutation()[dim];
      for (auto &dim : lhsContract)
        dim = lhsTrans.getPermutation()[dim];
    }

    if (fuseRhs) {
      for (auto &dim : rhsBatch)
        dim = rhsTrans.getPermutation()[dim];
      for (auto &dim : rhsContract)
        dim = rhsTrans.getPermutation()[dim];
    }

    auto ndim = stablehlo::DotDimensionNumbersAttr::get(
        dim.getContext(), lhsBatch, rhsBatch, lhsContract, rhsContract);

    rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
        dot, dot.getType(), fuseLhs ? lhsTrans.getOperand() : dot.getLhs(),
        fuseRhs ? rhsTrans.getOperand() : dot.getRhs(), ndim,
        dot.getPrecisionConfigAttr(), dot.getAlgorithmAttr());
    return success();
  }
};

// Optimize reduce(max/min, a * b) where a or b is a positive scalar
// If a > 0, then max(a*b1, a*b2, ...) = a * max(b1, b2, ...)
// Same for min with positive scalars
struct ReduceMaxMinMulPositiveScalar
    : public CheckedOpRewritePattern<stablehlo::ReduceOp,
                                     ReduceMaxMinMulPositiveScalar> {
  using CheckedOpRewritePattern<
      stablehlo::ReduceOp,
      ReduceMaxMinMulPositiveScalar>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1 || op.getInitValues().size() != 1) {
      return rewriter.notifyMatchFailure(
          op, "only single-operand single-init reduce is supported");
    }

    auto checkCommonReduce = mlir::stablehlo::CheckCommonReduceOp(op);

    // Only support max and min reductions
    if (checkCommonReduce.kind != stablehlo::ReduceOpKind::Max &&
        checkCommonReduce.kind != stablehlo::ReduceOpKind::Min) {
      return rewriter.notifyMatchFailure(
          op, "only max and min reductions are supported for this pattern");
    }

    Value input = op.getInputs()[0];
    if (!llvm::hasSingleElement(input.getUsers())) {
      return rewriter.notifyMatchFailure(op, "multiple users of input");
    }

    auto mulOp = input.getDefiningOp<stablehlo::MulOp>();
    if (!mulOp) {
      return rewriter.notifyMatchFailure(op, "input is not a multiply op");
    }

    Value lhs = mulOp.getLhs();
    Value rhs = mulOp.getRhs();

    // Check which operand is the positive scalar
    Value scalarOperand = nullptr;
    Value otherOperand = nullptr;

    // Helper to check if a value originates from a positive scalar
    auto isPositiveScalar = [&](Value v) -> bool {
      // First check if it's a scalar (broadcast from rank-0 tensor or splat)
      if (!stablehlo::isScalarValue(v))
        return false;

      // Now check if it's guaranteed non-negative and positive
      // We need to check the original scalar value, not the broadcasted one
      if (guaranteedNonNegativeResult(v, rewriter)) {
        // For strict positivity, check if it's guaranteed to be > 0
        // For now, we accept non-negative which includes zero
        // This is still correct for max/min but could potentially
        // multiply by zero which doesn't break correctness
        return true;
      }
      return false;
    };

    if (isPositiveScalar(lhs)) {
      scalarOperand = lhs;
      otherOperand = rhs;
    } else if (isPositiveScalar(rhs)) {
      scalarOperand = rhs;
      otherOperand = lhs;
    } else {
      return rewriter.notifyMatchFailure(
          op, "neither operand is a positive scalar");
    }

    // Get the scalar value to multiply the result with
    Value scalar = stablehlo::getScalarValue(scalarOperand, rewriter);
    if (!scalar) {
      return rewriter.notifyMatchFailure(op, "failed to extract scalar value");
    }

    // Create new reduce op on the non-scalar operand
    // First we need to make sure the otherOperand has the right shape
    auto inputType = cast<RankedTensorType>(input.getType());
    auto otherType = cast<RankedTensorType>(otherOperand.getType());

    // If otherOperand doesn't match the input shape, we need to broadcast it
    Value reduceInput = otherOperand;
    if (otherType.getShape() != inputType.getShape()) {
      // The original multiply would have broadcasted, so we need to find
      // how to get from otherOperand to input shape
      // Since one is scalar and one is the tensor, the tensor should match
      // the input shape - if not, we need to handle broadcasting
      return rewriter.notifyMatchFailure(
          op, "shape mismatch between operand and original input");
    }

    // Create the new reduction with the non-scalar input
    auto newReduction = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), op->getResultTypes(), ValueRange{reduceInput},
        op.getInitValues(), op.getDimensions());
    newReduction.getRegion().takeBody(op.getRegion());

    auto res = stablehlo::MulOpCreate(rewriter, op.getLoc(),
                                      newReduction.getResult(0), scalar);
    rewriter.replaceOp(op, res);
    return success();
  }
};

struct BroadcastReduce
    : public CheckedOpRewritePattern<stablehlo::ReduceOp, BroadcastReduce> {
  using CheckedOpRewritePattern<stablehlo::ReduceOp,
                                BroadcastReduce>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1 || op.getInitValues().size() != 1) {
      return rewriter.notifyMatchFailure(
          op, "only single-operand single-init reduce is supported");
    }

    auto checkCommonReduce = mlir::stablehlo::CheckCommonReduceOp(op);

    switch (checkCommonReduce.kind) {
    case stablehlo::ReduceOpKind::Add:
    case stablehlo::ReduceOpKind::Min:
    case stablehlo::ReduceOpKind::Max:
    case stablehlo::ReduceOpKind::And:
    case stablehlo::ReduceOpKind::Or:
    case stablehlo::ReduceOpKind::Xor:
    case stablehlo::ReduceOpKind::Mul:
      break;
    default:
      return rewriter.notifyMatchFailure(
          op, "only common reduce ops like add, mul, max, min are currently "
              "supported");
    }

    Value input = op.getInputs()[0];
    auto inputType = cast<TensorType>(input.getType());
    auto broadcast = input.getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!broadcast) {
      return rewriter.notifyMatchFailure(op,
                                         "input source is not a broadcast op");
    }

    // If any of the dimensions that are being reduced was initially
    // broadcasted, we can multiply the result with the dimension instead.
    ArrayRef<int64_t> broadcastDims = broadcast.getBroadcastDimensions();
    SmallVector<int64_t> broadcastFromNothingDims, broadcastFromOneDims;
    auto broadcastSourceType =
        cast<TensorType>(broadcast.getOperand().getType());
    for (int64_t reductionDim : op.getDimensions()) {
      if (inputType.isDynamicDim(reductionDim))
        continue;
      auto it = llvm::find(broadcastDims, reductionDim);
      if (it == broadcastDims.end()) {
        broadcastFromNothingDims.push_back(reductionDim);
        continue;
      }
      size_t originalDim = std::distance(broadcastDims.begin(), it);
      if (broadcastSourceType.getDimSize(originalDim) == 1 &&
          inputType.getDimSize(reductionDim) != 1) {
        broadcastFromOneDims.push_back(reductionDim);
      }
    }
    if (broadcastFromNothingDims.empty() && broadcastFromOneDims.empty())
      return rewriter.notifyMatchFailure(op, "no dimensions to remove");

    int64_t size = 1;
    for (int64_t dim : broadcastFromNothingDims) {
      size *= inputType.getDimSize(dim);
    }
    for (int64_t dim : broadcastFromOneDims) {
      size *= inputType.getDimSize(dim);
    }

    int64_t numRemoved = 0;
    SmallVector<int64_t> newReduceDimensions;
    llvm::sort(broadcastFromNothingDims);
    for (int64_t reductionDim : op.getDimensions()) {
      if (llvm::is_contained(broadcastFromNothingDims, reductionDim)) {
        numRemoved++;
        continue;
      }
      newReduceDimensions.push_back(reductionDim - numRemoved);
    }

    assert(numRemoved == broadcastFromNothingDims.size());

    Value operand = broadcast.getOperand();
    // Some of the broadcasted values aren't being reduced over
    // if (broadcastFromNothingDims.size() + broadcastFromOneDims.size() +
    // cast<RankedTensorType>(op.getType(0)).getShape().size() !=
    // inputType.getShape().size()) {
    SmallVector<int64_t> newBroadcast(
        cast<RankedTensorType>(broadcast.getOperand().getType())
            .getShape()
            .size(),
        -1);
    SmallVector<int64_t> newShape;
    for (int i = 0; i < inputType.getShape().size(); i++) {
      // These dimensions are removed
      if (llvm::is_contained(broadcastFromNothingDims, i)) {
        continue;
      }
      auto it = llvm::find(broadcastDims, i);
      if (llvm::is_contained(broadcastFromOneDims, i))
        newShape.push_back(1);
      else
        newShape.push_back(inputType.getShape()[i]);
      if (it == broadcastDims.end()) {
        // This dimension does not come from an earlier input
      } else {
        size_t originalDim = std::distance(broadcastDims.begin(), it);
        newBroadcast[originalDim] = newShape.size() - 1;
      }
    }
    operand = stablehlo::BroadcastInDimOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(newShape, inputType.getElementType()), operand,
        newBroadcast);
    // }

    auto newReduction = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), op->getResultTypes(), ValueRange{operand},
        op.getInitValues(), newReduceDimensions);
    newReduction.getRegion().takeBody(op.getRegion());

    auto newResultType = cast<TensorType>(newReduction.getResult(0).getType());

    switch (checkCommonReduce.kind) {
    case stablehlo::ReduceOpKind::Add: {
      auto constantInt = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(),
          makeAttr(newResultType.clone(rewriter.getI64Type()), size));
      auto converted = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), constantInt, newResultType.getElementType());
      assert(op.getType(0) == newReduction.getResult(0).getType());
      assert(op.getType(0) == converted.getType());
      rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
          op, newReduction.getResult(0), converted.getResult());
      break;
    }
    case stablehlo::ReduceOpKind::Min:
    case stablehlo::ReduceOpKind::Max:
    case stablehlo::ReduceOpKind::And:
    case stablehlo::ReduceOpKind::Or:
    case stablehlo::ReduceOpKind::Xor:
      rewriter.replaceAllUsesWith(op.getResult(0), newReduction.getResult(0));
      break;
    case stablehlo::ReduceOpKind::Mul: {
      auto constantInt = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(),
          makeAttr(newResultType.clone(rewriter.getI64Type()), size));
      auto converted = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), constantInt, newResultType.getElementType());
      assert(op.getType(0) == newReduction.getResult(0).getType());
      assert(op.getType(0) == converted.getType());
      rewriter.replaceOpWithNewOp<stablehlo::PowOp>(
          op, newReduction.getResult(0), converted.getResult());
      break;
    }
    default:
      llvm_unreachable(
          "probably missed out on a new reduction check introduced elsewhere.");
    }

    return success();
  }
};

template <typename OpTy>
static LogicalResult getDefiningZeroPadding(OpTy op, PatternRewriter &rewriter,
                                            stablehlo::PadOp &pad,
                                            Value &otherArg,
                                            bool &isOtherArgLHS) {
  pad = op.getLhs().template getDefiningOp<stablehlo::PadOp>();
  otherArg = op.getRhs();
  isOtherArgLHS = false;
  if (!pad) {
    pad = op.getRhs().template getDefiningOp<stablehlo::PadOp>();
    otherArg = op.getLhs();
    isOtherArgLHS = true;
  }
  if (!pad)
    return rewriter.notifyMatchFailure(op, "operands not produced by pad");
  // if (!llvm::hasSingleElement(pad->getUsers()))
  //  return rewriter.notifyMatchFailure(op, "pad has multiple users");

  if (!matchPattern(pad.getPaddingValue(), m_AnyZeroFloat()))
    return rewriter.notifyMatchFailure(op, "padding value not zero");
  return success();
}

template <typename T>
struct BinopConstPad : public CheckedOpRewritePattern<T, BinopConstPad<T>> {
  using CheckedOpRewritePattern<T, BinopConstPad<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    for (int i = 0; i < 2; i++) {
      DenseElementsAttr inp;
      if (!matchPattern(op->getOperand(i), m_Constant(&inp)))
        continue;
      if (!inp.isSplat())
        continue;

      auto pad =
          op->getOperand(1 - i).template getDefiningOp<stablehlo::PadOp>();
      if (!pad)
        continue;
      if (anyPadSizesNegative(pad))
        return failure();

      auto pval = pad.getPaddingValue();
      auto pval_cst = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), pval.getType(),
          inp.resizeSplat(cast<ShapedType>(pval.getType())));
      auto pval2 = T::create(rewriter, op.getLoc(), (i == 0) ? pval_cst : pval,
                             (i == 0) ? pval : pval_cst);

      auto val = pad.getOperand();
      auto val_cst = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), val.getType(),
          inp.resizeSplat(cast<ShapedType>(val.getType())));
      auto val2 = T::create(rewriter, op.getLoc(), (i == 0) ? val_cst : val,
                            (i == 0) ? val : val_cst);

      rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
          op, val2, pval2, pad.getEdgePaddingLow(), pad.getEdgePaddingHigh(),
          pad.getInteriorPadding());
      return success();
    }

    return failure();
  }
};

template <typename T>
struct BinopPadPad : public CheckedOpRewritePattern<T, BinopPadPad<T>> {
  using CheckedOpRewritePattern<T, BinopPadPad<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    auto pad1 = op->getOperand(0).template getDefiningOp<stablehlo::PadOp>();
    if (!pad1 || anyPadSizesNegative(pad1))
      return failure();

    auto pad2 = op->getOperand(1).template getDefiningOp<stablehlo::PadOp>();
    if (!pad2 || anyPadSizesNegative(pad2))
      return failure();

    if (pad1.getEdgePaddingLow() != pad2.getEdgePaddingLow())
      return failure();

    if (pad1.getEdgePaddingHigh() != pad2.getEdgePaddingHigh())
      return failure();

    if (pad1.getInteriorPadding() != pad2.getInteriorPadding())
      return failure();

    auto pv2 = T::create(rewriter, op.getLoc(), pad1.getPaddingValue(),
                         pad2.getPaddingValue());
    auto op2 =
        T::create(rewriter, op.getLoc(), pad1.getOperand(), pad2.getOperand());

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, op2, pv2, pad1.getEdgePaddingLow(), pad1.getEdgePaddingHigh(),
        pad1.getInteriorPadding());
    return success();
  }
};

struct AddPadPadToConcat
    : public CheckedOpRewritePattern<stablehlo::AddOp, AddPadPadToConcat> {
  using CheckedOpRewritePattern<stablehlo::AddOp,
                                AddPadPadToConcat>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AddOp op,
                                    PatternRewriter &rewriter) const {
    auto pad1 = op->getOperand(0).template getDefiningOp<stablehlo::PadOp>();
    if (!pad1 || anyPadSizesNegative(pad1))
      return failure();

    auto pad2 = op->getOperand(1).template getDefiningOp<stablehlo::PadOp>();
    if (!pad2 || anyPadSizesNegative(pad2))
      return failure();

    if (!matchPattern(pad1.getPaddingValue(), m_AnyZeroFloat()))
      return failure();

    if (!matchPattern(pad2.getPaddingValue(), m_AnyZeroFloat()))
      return failure();

    for (auto en : llvm::enumerate(op.getType().getShape())) {
      auto h1 = pad1.getEdgePaddingHigh()[en.index()];
      auto h2 = pad2.getEdgePaddingHigh()[en.index()];

      auto p1 = pad1.getOperand().getType().getShape()[en.index()];
      auto p2 = pad2.getOperand().getType().getShape()[en.index()];

      if (pad1.getInteriorPadding()[en.index()])
        continue;
      if (pad2.getInteriorPadding()[en.index()])
        continue;

      //  pad1: [ 0s ][ 0s   ][ data ][ 0s ]
      //  pad2: [ 0s ][ data ][ 0s   ][ 0s ]
      if (h2 == h1 + p1) {
        bool legal = true;
        for (auto en2 : llvm::enumerate(op.getType().getShape())) {
          if (en2.index() == en.index())
            continue;
          auto sl1 = pad1.getEdgePaddingLow()[en2.index()];
          auto sl2 = pad2.getEdgePaddingLow()[en2.index()];
          auto sh1 = pad1.getEdgePaddingHigh()[en2.index()];
          auto sh2 = pad2.getEdgePaddingHigh()[en2.index()];
          if (sl1 != sl2 || sh1 != sh2) {
            legal = false;
            break;
          }
        }
        if (legal) {
          Value data[] = {pad2.getOperand(), pad1.getOperand()};
          auto concat = stablehlo::ConcatenateOp::create(rewriter, op.getLoc(),
                                                         data, en.index());

          SmallVector<int64_t> lows(pad2.getEdgePaddingLow().begin(),
                                    pad2.getEdgePaddingLow().end());
          SmallVector<int64_t> highs(pad1.getEdgePaddingHigh().begin(),
                                     pad1.getEdgePaddingHigh().end());
          SmallVector<int64_t> ints(pad1.getInteriorPadding().begin(),
                                    pad1.getInteriorPadding().end());

          rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
              op, concat, pad1.getPaddingValue(), lows, highs, ints);
          return success();
        }
      }

      //  pad2: [ 0s ][ 0s   ][ data ][ 0s ]
      //  pad1: [ 0s ][ data ][ 0s   ][ 0s ]
      if (h1 == h2 + p2) {
        bool legal = true;
        for (auto en2 : llvm::enumerate(op.getType().getShape())) {
          if (en2.index() == en.index())
            continue;
          auto sl1 = pad1.getEdgePaddingLow()[en2.index()];
          auto sl2 = pad2.getEdgePaddingLow()[en2.index()];
          auto sh1 = pad1.getEdgePaddingHigh()[en2.index()];
          auto sh2 = pad2.getEdgePaddingHigh()[en2.index()];
          if (sl1 != sl2 || sh1 != sh2) {
            legal = false;
            break;
          }
        }
        if (legal) {
          Value data[] = {pad1.getOperand(), pad2.getOperand()};

          auto concat = stablehlo::ConcatenateOp::create(rewriter, op.getLoc(),
                                                         data, en.index());

          SmallVector<int64_t> lows(pad1.getEdgePaddingLow().begin(),
                                    pad1.getEdgePaddingLow().end());
          SmallVector<int64_t> highs(pad2.getEdgePaddingHigh().begin(),
                                     pad2.getEdgePaddingHigh().end());
          SmallVector<int64_t> ints(pad1.getInteriorPadding().begin(),
                                    pad1.getInteriorPadding().end());

          rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
              op, concat, pad1.getPaddingValue(), lows, highs, ints);
          return success();
        }
      }

      //  pad1: [ 0s ][      data            ][ 0s ]
      //  pad2: [ 0s ][ 0s ][  data   ][ 0s  ][ 0s ]
      if (h2 >= h1 && h2 + p2 <= h1 + p1) {
        bool legal = true;
        for (auto en2 : llvm::enumerate(op.getType().getShape())) {
          if (en2.index() == en.index())
            continue;
          auto sl1 = pad1.getEdgePaddingLow()[en2.index()];
          auto sl2 = pad2.getEdgePaddingLow()[en2.index()];
          auto sh1 = pad1.getEdgePaddingHigh()[en2.index()];
          auto sh2 = pad2.getEdgePaddingHigh()[en2.index()];
          if (sl1 != sl2 || sh1 != sh2) {
            legal = false;
            break;
          }
        }
        if (legal) {

          SmallVector<int64_t> slow(pad2.getEdgePaddingLow().size(), 0);
          SmallVector<int64_t> shigh(pad2.getEdgePaddingLow().size(), 0);
          SmallVector<int64_t> sint(pad2.getEdgePaddingLow().size(), 0);

          slow[en.index()] = h1 + p1 - (h2 + p2);
          shigh[en.index()] = h2 - h1;

          auto inPad = stablehlo::PadOp::create(
              rewriter, op.getLoc(), pad2.getOperand(), pad2.getPaddingValue(),
              slow, shigh, sint);
          assert(inPad.getType() == pad1.getOperand().getType());

          auto add = stablehlo::AddOp::create(rewriter, op.getLoc(), inPad,
                                              pad1.getOperand());

          rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
              op, add, pad1.getPaddingValue(), pad1.getEdgePaddingLow(),
              pad1.getEdgePaddingHigh(), pad1.getInteriorPadding());
          return success();
        }
      }

      //  pad2: [ 0s ][      data            ][ 0s ]
      //  pad1: [ 0s ][ 0s ][  data   ][ 0s  ][ 0s ]
      if (h1 >= h2 && h1 + p1 <= h2 + p2) {
        bool legal = true;
        for (auto en2 : llvm::enumerate(op.getType().getShape())) {
          if (en2.index() == en.index())
            continue;
          auto sl1 = pad1.getEdgePaddingLow()[en2.index()];
          auto sl2 = pad2.getEdgePaddingLow()[en2.index()];
          auto sh1 = pad1.getEdgePaddingHigh()[en2.index()];
          auto sh2 = pad2.getEdgePaddingHigh()[en2.index()];
          if (sl1 != sl2 || sh1 != sh2) {
            legal = false;
            break;
          }
        }
        if (legal) {

          SmallVector<int64_t> slow(pad1.getEdgePaddingLow().size(), 0);
          SmallVector<int64_t> shigh(pad1.getEdgePaddingLow().size(), 0);
          SmallVector<int64_t> sint(pad1.getEdgePaddingLow().size(), 0);

          slow[en.index()] = h2 + p2 - (h1 + p1);
          shigh[en.index()] = h1 - h2;

          auto inPad = stablehlo::PadOp::create(
              rewriter, op.getLoc(), pad1.getOperand(), pad1.getPaddingValue(),
              slow, shigh, sint);
          assert(inPad.getType() == pad2.getOperand().getType());

          auto add = stablehlo::AddOp::create(rewriter, op.getLoc(), inPad,
                                              pad2.getOperand());

          rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
              op, add, pad2.getPaddingValue(), pad2.getEdgePaddingLow(),
              pad2.getEdgePaddingHigh(), pad2.getInteriorPadding());
          return success();
        }
      }
    }

    return failure();
  }
};

template <typename T>
struct BinopBinopPadPad
    : public CheckedOpRewritePattern<T, BinopBinopPadPad<T>> {
  using CheckedOpRewritePattern<T,
                                BinopBinopPadPad<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    for (int i = 0; i < 2; i++) {
      auto pad1 = op->getOperand(i).template getDefiningOp<stablehlo::PadOp>();

      DenseElementsAttr inp1;
      if (!pad1) {
        if (!matchPattern(op->getOperand(i), m_Constant(&inp1)))
          continue;
        if (!inp1.isSplat())
          continue;
      } else if (anyPadSizesNegative(pad1))
        return failure();

      auto op2 = op->getOperand(1 - i).template getDefiningOp<T>();
      if (!op2)
        continue;

      for (int j = 0; j < 2; j++) {
        auto pad2 =
            op2->getOperand(j).template getDefiningOp<stablehlo::PadOp>();

        DenseElementsAttr inp2;
        if (!pad2) {
          if (!matchPattern(op2->getOperand(j), m_Constant(&inp2)))
            continue;
          if (!inp2.isSplat())
            continue;
        } else if (anyPadSizesNegative(pad2))
          return failure();

        if (pad1 && pad2) {
          if (pad1.getEdgePaddingLow() != pad2.getEdgePaddingLow())
            continue;

          if (pad1.getEdgePaddingHigh() != pad2.getEdgePaddingHigh())
            continue;

          if (pad1.getInteriorPadding() != pad2.getInteriorPadding())
            continue;
        }
        if (inp1 && inp2)
          continue;

        auto p1val =
            pad1 ? pad1.getPaddingValue()
                 : stablehlo::ConstantOp::create(
                       rewriter, op.getLoc(), pad2.getPaddingValue().getType(),
                       inp1.resizeSplat(
                           cast<ShapedType>(pad2.getPaddingValue().getType())));
        auto p2val =
            pad2 ? pad2.getPaddingValue()
                 : stablehlo::ConstantOp::create(
                       rewriter, op.getLoc(), pad1.getPaddingValue().getType(),
                       inp2.resizeSplat(
                           cast<ShapedType>(pad1.getPaddingValue().getType())));

        auto pval = T::create(rewriter, op.getLoc(), p1val, p2val);

        auto o1val =
            pad1 ? pad1.getOperand()
                 : stablehlo::ConstantOp::create(
                       rewriter, op.getLoc(), pad2.getOperand().getType(),
                       inp1.resizeSplat(
                           cast<ShapedType>(pad2.getOperand().getType())));
        auto o2val =
            pad2 ? pad2.getOperand()
                 : stablehlo::ConstantOp::create(
                       rewriter, op.getLoc(), pad1.getOperand().getType(),
                       inp2.resizeSplat(
                           cast<ShapedType>(pad1.getOperand().getType())));

        auto val = T::create(rewriter, op.getLoc(), o1val, o2val);

        auto npad =
            stablehlo::PadOp::create(rewriter, op.getLoc(), val, pval,
                                     (pad1 ? pad1 : pad2).getEdgePaddingLow(),
                                     (pad1 ? pad1 : pad2).getEdgePaddingHigh(),
                                     (pad1 ? pad1 : pad2).getInteriorPadding());

        rewriter.replaceOpWithNewOp<T>(op, op2->getOperand(1 - j), npad);
        return success();
      }
    }
    return failure();
  }
};

template <typename T>
struct BinopBinopPadConst
    : public CheckedOpRewritePattern<T, BinopBinopPadConst<T>> {
  using CheckedOpRewritePattern<T,
                                BinopBinopPadConst<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    for (int i = 0; i < 2; i++) {
      auto pad1 = op->getOperand(i).template getDefiningOp<stablehlo::PadOp>();
      if (!pad1)
        continue;
      if (anyPadSizesNegative(pad1))
        return failure();
      auto op2 = op->getOperand(1 - i).template getDefiningOp<T>();
      if (!op2)
        continue;

      for (int j = 0; j < 2; j++) {
        auto pad2 =
            op2->getOperand(j).template getDefiningOp<stablehlo::PadOp>();
        if (!pad2)
          continue;
        if (anyPadSizesNegative(pad2))
          return failure();

        if (pad1.getEdgePaddingLow() != pad2.getEdgePaddingLow())
          continue;

        if (pad1.getEdgePaddingHigh() != pad2.getEdgePaddingHigh())
          continue;

        if (pad1.getInteriorPadding() != pad2.getInteriorPadding())
          continue;

        auto pval = T::create(rewriter, pad2.getLoc(), pad1.getPaddingValue(),
                              pad2.getPaddingValue());
        auto val = T::create(rewriter, pad2.getLoc(), pad1.getOperand(),
                             pad2.getOperand());

        auto npad = stablehlo::PadOp::create(
            rewriter, op.getLoc(), val, pval, pad1.getEdgePaddingLow(),
            pad1.getEdgePaddingHigh(), pad1.getInteriorPadding());

        rewriter.replaceOpWithNewOp<T>(op, op2->getOperand(1 - j), npad);
        return success();
      }
    }
    return failure();
  }
};

struct MulZeroPad
    : public CheckedOpRewritePattern<stablehlo::MulOp, MulZeroPad> {
  using CheckedOpRewritePattern<stablehlo::MulOp,
                                MulZeroPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    stablehlo::PadOp pad;
    Value otherArg;
    bool otherIsLHS;
    if (failed(getDefiningZeroPadding(op, rewriter, pad, otherArg, otherIsLHS)))
      return failure();
    if (anyPadSizesNegative(pad))
      return failure();

    auto otherArgType = cast<TensorType>(otherArg.getType());
    SmallVector<int64_t> limitDims = llvm::to_vector(otherArgType.getShape());
    for (auto &&[limit, pad] : llvm::zip(limitDims, pad.getEdgePaddingHigh())) {
      limit -= pad;
    }
    SmallVector<int64_t> interior = llvm::to_vector(pad.getInteriorPadding());
    for (int64_t &value : interior) {
      value += 1;
    }

    auto slice = stablehlo::SliceOp::create(rewriter, pad.getLoc(), otherArg,
                                            pad.getEdgePaddingLow(), limitDims,
                                            interior);
    auto mul = stablehlo::MulOp::create(
        rewriter, op.getLoc(),
        otherIsLHS ? slice.getResult() : pad.getOperand(),
        otherIsLHS ? pad.getOperand() : slice.getResult());
    auto newPad = stablehlo::PadOp::create(
        rewriter, pad.getLoc(), mul.getResult(), pad.getPaddingValue(),
        pad.getEdgePaddingLowAttr(), pad.getEdgePaddingHighAttr(),
        pad.getInteriorPaddingAttr());
    assert(op.getType() == newPad.getType());
    rewriter.replaceOp(op, newPad);

    return success();
  }
};

struct DivZeroPad
    : public CheckedOpRewritePattern<stablehlo::DivOp, DivZeroPad> {
  using CheckedOpRewritePattern<stablehlo::DivOp,
                                DivZeroPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DivOp op,
                                    PatternRewriter &rewriter) const {
    stablehlo::PadOp pad;
    Value otherArg;
    bool otherIsLHS;
    if (failed(getDefiningZeroPadding(op, rewriter, pad, otherArg, otherIsLHS)))
      return failure();
    if (anyPadSizesNegative(pad))
      return failure();

    if (otherIsLHS)
      return failure();

    auto otherArgType = cast<TensorType>(otherArg.getType());
    SmallVector<int64_t> limitDims = llvm::to_vector(otherArgType.getShape());
    for (auto &&[limit, pad] : llvm::zip(limitDims, pad.getEdgePaddingHigh())) {
      limit -= pad;
    }
    SmallVector<int64_t> interior = llvm::to_vector(pad.getInteriorPadding());
    for (int64_t &value : interior) {
      value += 1;
    }

    auto slice = stablehlo::SliceOp::create(rewriter, pad.getLoc(), otherArg,
                                            pad.getEdgePaddingLow(), limitDims,
                                            interior);
    auto mul = stablehlo::DivOp::create(
        rewriter, op.getLoc(),
        otherIsLHS ? slice.getResult() : pad.getOperand(),
        otherIsLHS ? pad.getOperand() : slice.getResult());
    auto newPad = stablehlo::PadOp::create(
        rewriter, pad.getLoc(), mul.getResult(), pad.getPaddingValue(),
        pad.getEdgePaddingLowAttr(), pad.getEdgePaddingHighAttr(),
        pad.getInteriorPaddingAttr());
    rewriter.replaceOp(op, newPad);

    return success();
  }
};

template <typename T> DenseI64ArrayAttr addLists(T lhs, T rhs) {
  MLIRContext *context = lhs.getContext();
  SmallVector<int64_t> sum;
  for (auto &&[lhsv, rhsv] :
       llvm::zip(cast<DenseI64ArrayAttr>(lhs).asArrayRef(),
                 cast<DenseI64ArrayAttr>(rhs).asArrayRef())) {
    sum.push_back(lhsv + rhsv);
  }
  return DenseI64ArrayAttr::get(context, sum);
}

struct PadPad : public CheckedOpRewritePattern<stablehlo::PadOp, PadPad> {
  using CheckedOpRewritePattern<stablehlo::PadOp,
                                PadPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::PadOp op,
                                    PatternRewriter &rewriter) const {
    auto definingPad = op.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!definingPad || definingPad.getPaddingValue() != op.getPaddingValue()) {
      return rewriter.notifyMatchFailure(op, "no compatible defining pad");
    }

    auto allZero = [](ArrayRef<int64_t> values) {
      return llvm::all_of(values, [](int64_t v) { return v == 0; });
    };

    if (!allZero(op.getInteriorPadding()) ||
        !allZero(definingPad.getInteriorPadding())) {
      return rewriter.notifyMatchFailure(op, "cannot combine interior padding");
    }

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, definingPad.getOperand(), definingPad.getPaddingValue(),
        addLists(op.getEdgePaddingLowAttr(),
                 definingPad.getEdgePaddingLowAttr()),
        addLists(op.getEdgePaddingHighAttr(),
                 definingPad.getEdgePaddingHighAttr()),
        addLists(op.getInteriorPaddingAttr(),
                 definingPad.getInteriorPaddingAttr()));
    return success();
  }
};

std::tuple<Value, Value, RankedTensorType>
sliceDotGeneralHelper(stablehlo::DotGeneralOp dot, ArrayRef<int64_t> starts,
                      ArrayRef<int64_t> limits, ArrayRef<int64_t> strides,
                      PatternRewriter &rewriter) {
  SmallVector<int64_t> lhs_lb(dot.getLhs().getType().getShape().size(), 0);
  SmallVector<int64_t> lhs_ub(dot.getLhs().getType().getShape().begin(),
                              dot.getLhs().getType().getShape().end());
  SmallVector<int64_t> lhs_step(dot.getLhs().getType().getShape().size(), 1);

  SmallVector<int64_t> rhs_lb(dot.getRhs().getType().getShape().size(), 0);
  SmallVector<int64_t> rhs_ub(dot.getRhs().getType().getShape().begin(),
                              dot.getRhs().getType().getShape().end());
  SmallVector<int64_t> rhs_step(dot.getRhs().getType().getShape().size(), 1);

  auto dimensionNumbers = dot.getDotDimensionNumbers();

  size_t residx = 0;
  SmallVector<int64_t> resShape;
  for (auto &&[lhs, rhs] :
       llvm::zip(dimensionNumbers.getLhsBatchingDimensions(),
                 dimensionNumbers.getRhsBatchingDimensions())) {
    lhs_lb[lhs] = starts[residx];
    lhs_ub[lhs] = limits[residx];
    lhs_step[lhs] = strides[residx];

    rhs_lb[rhs] = starts[residx];
    rhs_ub[rhs] = limits[residx];
    rhs_step[rhs] = strides[residx];
    resShape.push_back((limits[residx] - starts[residx]) / strides[residx]);
    residx++;
  }

  for (size_t i = 0, end = dot.getLhs().getType().getShape().size(); i < end;
       i++) {
    if (llvm::is_contained(dimensionNumbers.getLhsContractingDimensions(), i))
      continue;
    if (llvm::is_contained(dimensionNumbers.getLhsBatchingDimensions(), i))
      continue;

    lhs_lb[i] = starts[residx];
    lhs_ub[i] = limits[residx];
    lhs_step[i] = strides[residx];
    resShape.push_back((limits[residx] - starts[residx]) / strides[residx]);
    residx++;
  }

  for (size_t i = 0, end = dot.getRhs().getType().getShape().size(); i < end;
       i++) {
    if (llvm::is_contained(dimensionNumbers.getRhsContractingDimensions(), i))
      continue;
    if (llvm::is_contained(dimensionNumbers.getRhsBatchingDimensions(), i))
      continue;

    rhs_lb[i] = starts[residx];
    rhs_ub[i] = limits[residx];
    rhs_step[i] = strides[residx];
    resShape.push_back((limits[residx] - starts[residx]) / strides[residx]);
    residx++;
  }

  assert(residx == dot.getType().getShape().size());

  auto lhs2 = stablehlo::SliceOp::create(rewriter, dot.getLoc(), dot.getLhs(),
                                         lhs_lb, lhs_ub, lhs_step);
  auto rhs2 = stablehlo::SliceOp::create(rewriter, dot.getLoc(), dot.getRhs(),
                                         rhs_lb, rhs_ub, rhs_step);

  return std::tuple<Value, Value, RankedTensorType>(
      lhs2, rhs2,
      RankedTensorType::get(resShape, dot.getType().getElementType()));
}

struct SliceDotGeneral
    : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceDotGeneral> {
  using CheckedOpRewritePattern<stablehlo::SliceOp,
                                SliceDotGeneral>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto dot = op.getOperand().getDefiningOp<stablehlo::DotGeneralOp>();
    if (!dot) {
      return rewriter.notifyMatchFailure(op, "defining op is not a reshape");
    }
    if (!llvm::hasSingleElement(dot->getUsers()))
      return failure();

    auto &&[lhs2, rhs2, resTy] =
        sliceDotGeneralHelper(dot, op.getStartIndices(), op.getLimitIndices(),
                              op.getStrides(), rewriter);

    Value operands[2] = {lhs2, rhs2};
    rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
        op, TypeRange(resTy), operands, dot->getAttrs());

    return success();
  }
};

struct PadDotGeneral
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp, PadDotGeneral> {
  using CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                PadDotGeneral>::CheckedOpRewritePattern;

  bool postPad;
  PadDotGeneral(size_t postPad, MLIRContext *context,
                PatternBenefit benefit = 1,
                ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        postPad(postPad) {}

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    stablehlo::PadOp pad;
    Value otherArg;
    bool otherIsLHS;
    if (failed(getDefiningZeroPadding(op, rewriter, pad, otherArg, otherIsLHS)))
      return failure();
    if (anyPadSizesNegative(pad))
      return failure();

    for (auto u : pad->getUsers())
      if (!isa<stablehlo::DotGeneralOp>(u))
        return failure();

    auto dimensionNumbers = op.getDotDimensionNumbers();
    auto padContractingDimensions =
        dimensionNumbers.getLhsContractingDimensions();
    auto otherContractingDimensions =
        dimensionNumbers.getRhsContractingDimensions();

    auto padBatchDimensions = dimensionNumbers.getLhsBatchingDimensions();
    auto otherBatchDimensions = dimensionNumbers.getRhsBatchingDimensions();

    SmallVector<int64_t> padResultDimensions;
    for (size_t i = 0, end = op.getLhs().getType().getShape().size(); i < end;
         i++) {
      if (llvm::is_contained(dimensionNumbers.getLhsContractingDimensions(), i))
        continue;
      if (llvm::is_contained(dimensionNumbers.getLhsBatchingDimensions(), i))
        continue;
      padResultDimensions.push_back(i);
    }

    SmallVector<int64_t> otherResultDimensions;
    for (size_t i = 0, end = op.getRhs().getType().getShape().size(); i < end;
         i++) {
      if (llvm::is_contained(dimensionNumbers.getRhsContractingDimensions(), i))
        continue;
      if (llvm::is_contained(dimensionNumbers.getRhsBatchingDimensions(), i))
        continue;
      otherResultDimensions.push_back(i);
    }

    if (otherIsLHS) {
      std::swap(padContractingDimensions, otherContractingDimensions);
      std::swap(padBatchDimensions, otherBatchDimensions);
      std::swap(padResultDimensions, otherResultDimensions);
    }

    // Need to figure out which dimension(s) to slice. For this purpose,
    // look the pairs of contracting dimensions.
    SmallVector<std::tuple<int64_t, int64_t, int64_t, int64_t>>
        otherDimsToSlice;
    for (auto &&[padDim, otherDim] :
         llvm::zip(padContractingDimensions, otherContractingDimensions)) {
      // If padding along the dim, mark the corresponding other dim for
      // slicing.
      int64_t low = pad.getEdgePaddingLow()[padDim];
      int64_t high = pad.getEdgePaddingHigh()[padDim];
      int64_t interior = pad.getInteriorPadding()[padDim];
      if (low == 0 && high == 0 && interior == 0)
        continue;
      if (!postPad)
        otherDimsToSlice.emplace_back(otherDim, low, high, interior);
    }

    SmallVector<std::tuple<int64_t, int64_t, int64_t, int64_t>> resultDimsToPad;
    SmallVector<int64_t> resultShape;

    {
      size_t resultidx = 0;
      for (auto &&[padDim, otherDim] :
           llvm::zip(padBatchDimensions, otherBatchDimensions)) {
        // If padding along the dim, mark the corresponding other dim for
        // slicing.
        int64_t low = pad.getEdgePaddingLow()[padDim];
        int64_t high = pad.getEdgePaddingHigh()[padDim];
        int64_t interior = pad.getInteriorPadding()[padDim];
        auto padSize = cast<RankedTensorType>(pad.getOperand().getType())
                           .getShape()[padDim];
        resultShape.push_back(padSize);
        if (low == 0 && high == 0 && interior == 0) {
          resultidx++;
          continue;
        }
        if (!postPad)
          otherDimsToSlice.emplace_back(otherDim, low, high, interior);
        resultDimsToPad.emplace_back(resultidx, low, high, interior);
        resultidx++;
      }

      if (otherIsLHS) {
        for (auto dim : otherResultDimensions) {
          resultidx++;
          resultShape.push_back(
              cast<RankedTensorType>(otherArg.getType()).getShape()[dim]);
        }
      }

      for (auto padDim : padResultDimensions) {
        int64_t low = pad.getEdgePaddingLow()[padDim];
        int64_t high = pad.getEdgePaddingHigh()[padDim];
        int64_t interior = pad.getInteriorPadding()[padDim];
        auto padSize = cast<RankedTensorType>(pad.getOperand().getType())
                           .getShape()[padDim];
        resultShape.push_back(padSize);
        if (low == 0 && high == 0 && interior == 0) {
          resultidx++;
          continue;
        }

        resultDimsToPad.emplace_back(resultidx, low, high, interior);
        resultidx++;
      }

      if (!otherIsLHS) {
        for (auto dim : otherResultDimensions) {
          resultidx++;
          resultShape.push_back(
              cast<RankedTensorType>(otherArg.getType()).getShape()[dim]);
        }
      }
    }

    if (!resultDimsToPad.empty() && !postPad)
      return failure();

    if (otherDimsToSlice.empty() && resultDimsToPad.empty()) {
      return rewriter.notifyMatchFailure(op,
                                         "contracting dimensions not padded");
    }

    Value nextOtherArg = otherArg;
    if (!otherDimsToSlice.empty()) {
      SmallVector<int64_t> sliceLow, sliceHigh, sliceStride;
      for (auto &&[pos, size] :
           llvm::enumerate(cast<TensorType>(otherArg.getType()).getShape())) {
        std::optional<std::tuple<int64_t, int64_t, int64_t, int64_t>> it;
        for (auto tup : otherDimsToSlice)
          if (std::get<0>(tup) == pos)
            it = tup;
        if (!it) {
          sliceLow.push_back(0);
          sliceHigh.push_back(size);
          sliceStride.push_back(1);
          continue;
        }

        sliceLow.push_back(std::get<1>(*it));
        sliceHigh.push_back(size - std::get<2>(*it));
        sliceStride.push_back(std::get<3>(*it) + 1);
      }

      auto slice = stablehlo::SliceOp::create(rewriter, op.getLoc(), otherArg,
                                              sliceLow, sliceHigh, sliceStride);
      nextOtherArg = slice.getResult();
    }

    Value res = stablehlo::DotGeneralOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(resultShape, op.getType().getElementType()),
        otherIsLHS ? nextOtherArg : pad.getOperand(),
        otherIsLHS ? pad.getOperand() : nextOtherArg,
        op.getDotDimensionNumbersAttr(), op.getPrecisionConfigAttr(),
        op.getAlgorithmAttr());

    if (!resultDimsToPad.empty()) {
      SmallVector<int64_t> low(op.getType().getShape().size(), 0);
      SmallVector<int64_t> high(op.getType().getShape().size(), 0);
      SmallVector<int64_t> interior(op.getType().getShape().size(), 0);
      for (auto &&[idx, lval, hval, ival] : resultDimsToPad) {
        low[idx] = lval;
        high[idx] = hval;
        interior[idx] = ival;
      }
      res =
          stablehlo::PadOp::create(rewriter, op.getLoc(), op.getType(), res,
                                   pad.getPaddingValue(), low, high, interior);
    }
    rewriter.replaceOp(op, res);
    return success();
  }
};

// sliceIndices which are null are essentially 0 (we don't construct here since
// the pass might still fail)
LogicalResult sliceReshapeHelper(stablehlo::DynamicSliceOp op,
                                 SmallVectorImpl<Value> &startIndices,
                                 SmallVectorImpl<int64_t> &sliceSizes) {
  auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
  if (!reshape) {
    return failure();
  }

  assert(startIndices.size() == 0);
  assert(sliceSizes.size() == 0);

  auto reshapeOperandType = cast<TensorType>(reshape.getOperand().getType());
  auto reshapeType = cast<TensorType>(reshape.getType());
  size_t indim = 0;
  size_t outdim = 0;

  while (indim < reshapeOperandType.getShape().size() &&
         outdim < reshapeType.getShape().size()) {
    if (reshapeOperandType.getShape()[indim] ==
        reshapeType.getShape()[outdim]) {
      startIndices.push_back(op.getStartIndices()[outdim]);
      sliceSizes.push_back(op.getSliceSizes()[outdim]);
      indim++;
      outdim++;
      continue;
    }

    if (reshapeOperandType.getShape()[indim] == 1) {
      startIndices.push_back(nullptr);
      sliceSizes.push_back(1);
      indim++;
      continue;
    }

    if (reshapeType.getShape()[outdim] == 1) {
      if (!matchPattern(op.getStartIndices()[outdim], m_Zero()))
        return failure();
      if (op.getSliceSizes()[outdim] != 1)
        return failure();
      outdim++;
      continue;
    }
    return failure();
  }
  while (indim < reshapeOperandType.getShape().size()) {
    if (reshapeOperandType.getShape()[indim] != 1)
      return failure();
    // It's a full slice of the original dimension.
    startIndices.push_back(nullptr);
    sliceSizes.push_back(1);
    indim++;
  }
  while (outdim < reshapeType.getShape().size()) {
    if (reshapeType.getShape()[outdim] != 1)
      return failure();
    if (!matchPattern(op.getStartIndices()[outdim], m_Zero()))
      return failure();
    if (op.getSliceSizes()[outdim] != 1)
      return failure();
    outdim++;
  }
  return success();
}

LogicalResult sliceReshapeHelper(stablehlo::SliceOp op,
                                 SmallVectorImpl<int64_t> &starts,
                                 SmallVectorImpl<int64_t> &limits,
                                 SmallVectorImpl<int64_t> &strides) {
  auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
  if (!reshape) {
    return failure();
  }

  assert(starts.size() == 0);
  assert(limits.size() == 0);
  assert(strides.size() == 0);

  auto reshapeOperandType = cast<TensorType>(reshape.getOperand().getType());
  auto reshapeType = cast<TensorType>(reshape.getType());
  size_t indim = 0;
  size_t outdim = 0;
  while (indim < reshapeOperandType.getShape().size() &&
         outdim < reshapeType.getShape().size()) {
    if (reshapeOperandType.getShape()[indim] ==
        reshapeType.getShape()[outdim]) {
      starts.push_back(op.getStartIndices()[outdim]);
      limits.push_back(op.getLimitIndices()[outdim]);
      strides.push_back(op.getStrides()[outdim]);
      indim++;
      outdim++;
      continue;
    }
    if (reshapeOperandType.getShape()[indim] == 1) {
      starts.push_back(0);
      limits.push_back(1);
      strides.push_back(1);
      indim++;
      continue;
    }
    if (reshapeType.getShape()[outdim] == 1) {
      if (op.getStartIndices()[outdim] != 0)
        return failure();
      if (op.getLimitIndices()[outdim] != 1)
        return failure();
      if (op.getStrides()[outdim] != 1)
        return failure();
      outdim++;
      continue;
    }
    return failure();
  }
  while (indim < reshapeOperandType.getShape().size()) {
    if (reshapeOperandType.getShape()[indim] != 1)
      return failure();
    // It's a full slice of the original dimension.
    starts.push_back(0);
    limits.push_back(1);
    strides.push_back(1);
    indim++;
  }
  while (outdim < reshapeType.getShape().size()) {
    if (reshapeType.getShape()[outdim] != 1)
      return failure();
    if (op.getStartIndices()[outdim] != 0)
      return failure();
    if (op.getLimitIndices()[outdim] != 1)
      return failure();
    if (op.getStrides()[outdim] != 1)
      return failure();
    outdim++;
  }
  return success();
}

struct SliceReshape
    : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceReshape> {
  using CheckedOpRewritePattern<stablehlo::SliceOp,
                                SliceReshape>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape) {
      return rewriter.notifyMatchFailure(op, "defining op is not a reshape");
    }
    if (!llvm::hasSingleElement(reshape->getUsers()))
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    auto newSlice = stablehlo::SliceOp::create(
        rewriter, op->getLoc(), reshape.getOperand(), starts, limits, strides);
    auto newReshape = stablehlo::ReshapeOpCreate(
        rewriter, reshape->getLoc(), newSlice.getResult(),
        cast<RankedTensorType>(op.getResult().getType()).getShape());
    rewriter.replaceOp(op, newReshape);

    return success();
  }
};

// slice(reshape(pad x)) -> pad(slice x)
struct SliceReshapePad final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceReshapePad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    auto pad = reshape.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    SmallVector<int64_t> start;
    SmallVector<int64_t> end;
    SmallVector<int64_t> step;

    SmallVector<int64_t> lpads;
    SmallVector<int64_t> hpads;
    SmallVector<int64_t> interiors;

    bool needspad = false;
    bool broadcastres = false;
    if (!slicePadHelper(pad, starts, limits, strides, start, end, step, lpads,
                        hpads, interiors, broadcastres, needspad)
             .succeeded())
      return failure();

    if (broadcastres) {
      rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
          op, op.getType(), pad.getPaddingValue(),
          rewriter.getDenseI64ArrayAttr({}));
      return success();
    }

    if (needspad && !llvm::hasSingleElement(pad->getUsers()) &&
        !llvm::hasSingleElement(reshape->getUsers()))
      return failure();

    mlir::Value nslice = stablehlo::SliceOp::create(
        rewriter, op.getLoc(), pad.getOperand(), start, end, step);
    if (needspad) {
      nslice = stablehlo::PadOp::create(rewriter, op.getLoc(), nslice,
                                        pad.getPaddingValue(), lpads, hpads,
                                        interiors);
    }
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(), nslice);
    return success();
  }
};

struct SliceIf : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceIf> {
  using CheckedOpRewritePattern<stablehlo::SliceOp,
                                SliceIf>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto ifop = op.getOperand().getDefiningOp<stablehlo::IfOp>();
    if (!ifop)
      return failure();
    if (!llvm::hasSingleElement(op.getOperand().getUsers()))
      return failure();

    auto opres = cast<OpResult>(op.getOperand());
    ssize_t opIdx = opres.getResultNumber();

    SmallVector<Type> ifResultTypes = llvm::to_vector(ifop->getResultTypes());
    ifResultTypes[opIdx] = op.getType();

    auto newIf = stablehlo::IfOp::create(rewriter, ifop.getLoc(), ifResultTypes,
                                         ifop.getPred());

    Operation *trueTerm = ifop.getTrueBranch().front().getTerminator();
    Operation *falseTerm = ifop.getFalseBranch().front().getTerminator();

    rewriter.setInsertionPoint(trueTerm);
    auto newTrue = stablehlo::SliceOp::create(
        rewriter, op.getLoc(), trueTerm->getOperands()[opIdx],
        op.getStartIndices(), op.getLimitIndices(), op.getStrides());
    rewriter.modifyOpInPlace(trueTerm,
                             [&] { trueTerm->setOperand(opIdx, newTrue); });

    rewriter.setInsertionPoint(falseTerm);
    auto newFalse = stablehlo::SliceOp::create(
        rewriter, op.getLoc(), falseTerm->getOperands()[opIdx],
        op.getStartIndices(), op.getLimitIndices(), op.getStrides());
    rewriter.modifyOpInPlace(falseTerm,
                             [&] { falseTerm->setOperand(opIdx, newFalse); });

    newIf.getTrueBranch().takeBody(ifop.getTrueBranch());
    newIf.getFalseBranch().takeBody(ifop.getFalseBranch());

    for (int i = 0; i < ifop.getNumResults(); i++) {
      if (i == opIdx) {
        rewriter.replaceOp(op, newIf.getResult(i));
      } else {
        rewriter.replaceAllUsesWith(ifop.getResult(i), newIf.getResult(i));
      }
    }

    rewriter.eraseOp(ifop);
    return success();
  }
};

struct SliceReshapeConcat
    : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceReshapeConcat> {
  using CheckedOpRewritePattern<stablehlo::SliceOp,
                                SliceReshapeConcat>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()))
      return failure();

    auto concat =
        reshape.getOperand().getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concat)
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    auto dim = concat.getDimension();

    SmallVector<Value> postConcat;
    if (!sliceConcatHelper(concat, rewriter, starts, limits, strides,
                           postConcat)
             .succeeded())
      return failure();

    auto c2 = stablehlo::ConcatenateOp::create(rewriter, concat.getLoc(),
                                               postConcat, dim);
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(), c2);
    return success();
  }
};

struct SliceReshapeElementwise final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceReshapeElementwise> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()))
      return failure();
    auto elem = reshape.getOperand().getDefiningOp();
    if (!elem)
      return failure();
    if (!stablehlo::hasTraitElementwise(elem))
      return failure();
    if (!llvm::hasSingleElement(elem->getUsers()))
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    SmallVector<Value> ops;
    for (auto v : elem->getOperands()) {
      ops.push_back(stablehlo::SliceOp::create(rewriter, op.getLoc(), v, starts,
                                               limits, strides));
    }
    SmallVector<int64_t> sizes;
    for (auto &&[start, stop, stride] : llvm::zip(starts, limits, strides))
      sizes.push_back((stop - start) / stride);
    auto nex = rewriter.create(
        elem->getLoc(), elem->getName().getIdentifier(), ValueRange(ops),
        TypeRange(RankedTensorType::get(
            sizes, reshape.getOperand().getType().getElementType())),
        elem->getAttrs(), {}, {});
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      nex->getResult(0));
    return success();
  }
};

struct TransposeLikeBroadcastElementwise final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                              TransposeLikeBroadcastElementwise> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  bool onlySingleUser;

  TransposeLikeBroadcastElementwise(bool onlySingleUser, MLIRContext *context,
                                    PatternBenefit benefit = 1,
                                    ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        onlySingleUser(onlySingleUser) {}

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    if (!isTransposeReshapeLikeBroadcast(op))
      return failure();
    auto elem = op.getOperand().getDefiningOp();
    if (!elem)
      return failure();

    if (!stablehlo::hasTraitElementwise(elem))
      return failure();

    bool singleUser = llvm::hasSingleElement(elem->getUsers());
    if (onlySingleUser && !singleUser)
      return failure();

    SmallVector<Value> ops;
    for (auto v : elem->getOperands()) {
      if (auto rop = v.getDefiningOp()) {
        rewriter.setInsertionPointAfter(rop);
      } else if (auto ba = dyn_cast<BlockArgument>(v)) {
        rewriter.setInsertionPointToStart(ba.getOwner());
      }
      auto nt = RankedTensorType::get(
          op.getType().getShape(),
          cast<RankedTensorType>(v.getType()).getElementType());
      ops.push_back(stablehlo::BroadcastInDimOp::create(
          rewriter, op.getLoc(), nt, v, op.getBroadcastDimensions()));
    }
    if (singleUser) {
      rewriter.modifyOpInPlace(elem, [&]() {
        elem->setOperands(ops);
        elem->getResult(0).setType(op.getType());
      });
      rewriter.replaceOp(op, elem);
    } else {
      rewriter.setInsertionPointAfter(elem);
      auto newOp = rewriter.create(
          elem->getLoc(), elem->getName().getIdentifier(), ValueRange(ops),
          TypeRange(op.getType()), elem->getAttrs(), {}, {});
      rewriter.replaceOp(op, newOp);
    }
    return success();
  }
};

struct TransposeElementwise final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeElementwise> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  bool onlySingleUser;

  TransposeElementwise(bool onlySingleUser, MLIRContext *context,
                       PatternBenefit benefit = 1,
                       ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        onlySingleUser(onlySingleUser) {}

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto elem = op.getOperand().getDefiningOp();
    if (!elem)
      return failure();

    if (!stablehlo::hasTraitElementwise(elem))
      return failure();

    bool singleUser = llvm::hasSingleElement(elem->getUsers());
    if (onlySingleUser && !singleUser)
      return failure();

    SmallVector<Value> ops;
    for (auto v : elem->getOperands()) {
      if (auto rop = v.getDefiningOp()) {
        rewriter.setInsertionPointAfter(rop);
      } else if (auto ba = dyn_cast<BlockArgument>(v)) {
        rewriter.setInsertionPointToStart(ba.getOwner());
      }
      ops.push_back(stablehlo::TransposeOp::create(rewriter, op.getLoc(), v,
                                                   op.getPermutation()));
    }
    if (singleUser) {
      rewriter.modifyOpInPlace(elem, [&]() {
        elem->setOperands(ops);
        elem->getResult(0).setType(op.getType());
      });
      rewriter.replaceOp(op, elem);
    } else {
      rewriter.setInsertionPointAfter(elem);
      auto newOp = rewriter.create(
          elem->getLoc(), elem->getName().getIdentifier(), ValueRange(ops),
          TypeRange(op.getType()), elem->getAttrs(), {}, {});
      rewriter.replaceOp(op, newOp);
    }
    return success();
  }
};

struct TransposeConcat final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto concat = op.getOperand().getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concat)
      return failure();

    if (!llvm::hasSingleElement(concat->getUsers()))
      return failure();

    SmallVector<Value> ops;
    for (auto v : concat->getOperands()) {
      ops.push_back(stablehlo::TransposeOp::create(rewriter, op.getLoc(), v,
                                                   op.getPermutation()));
    }

    auto dim = concat.getDimension();
    auto dim2 = getInversePermutation(op.getPermutation())[dim];

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, ops, dim2);
    rewriter.eraseOp(concat);
    return success();
  }
};

struct TransposeIota final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeIota> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto iota = op.getOperand().getDefiningOp<stablehlo::IotaOp>();
    if (!iota)
      return failure();

    if (!llvm::hasSingleElement(iota->getUsers()))
      return failure();

    auto dim = iota.getIotaDimension();
    auto dim2 = getInversePermutation(op.getPermutation())[dim];

    rewriter.replaceOpWithNewOp<stablehlo::IotaOp>(op, op.getType(), dim2);
    rewriter.eraseOp(iota);
    return success();
  }
};

struct TransposeReduceWindow final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeReduceWindow> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto reduce = op.getOperand().getDefiningOp<stablehlo::ReduceWindowOp>();
    if (!reduce)
      return failure();

    if (reduce->getNumResults() != 1)
      return failure();

    if (!llvm::hasSingleElement(reduce->getResult(0).getUsers()))
      return failure();

    Type restys[] = {op.getType()};

    Value operands[] = {stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), reduce.getOperands()[0], op.getPermutation())};

    int64_t padding_shape[2] = {(int64_t)op.getType().getShape().size(), 2};

    SmallVector<int64_t> win_dim(reduce.getWindowDimensions().begin(),
                                 reduce.getWindowDimensions().end());
    SmallVector<int64_t> win_strides;
    if (reduce.getWindowStrides())
      win_strides.append(reduce.getWindowStrides()->begin(),
                         reduce.getWindowStrides()->end());
    SmallVector<int64_t> base_dialations;
    if (reduce.getBaseDilations())
      base_dialations.append(reduce.getBaseDilations()->begin(),
                             reduce.getBaseDilations()->end());
    SmallVector<int64_t> win_dialations;
    if (reduce.getWindowDilations())
      win_dialations.append(reduce.getWindowDilations()->begin(),
                            reduce.getWindowDilations()->end());
    SmallVector<int64_t> padding_dialations(2 * padding_shape[0]);

    auto perm = op.getPermutation();
    for (int64_t i = 0; i < perm.size(); ++i) {
      win_dim[perm[i]] = reduce.getWindowDimensions()[i];
      if (reduce.getWindowStrides())
        win_strides[perm[i]] = (*reduce.getWindowStrides())[i];
      if (reduce.getBaseDilations())
        base_dialations[perm[i]] = (*reduce.getBaseDilations())[i];
      if (reduce.getWindowDilations())
        win_dialations[perm[i]] = (*reduce.getWindowDilations())[i];
      if (reduce.getPadding()) {
        padding_dialations[2 * perm[i]] =
            (*(reduce.getPadding()->begin() + (2 * i))).getSExtValue();
        padding_dialations[2 * perm[i] + 1] =
            (*(reduce.getPadding()->begin() + (2 * i + 1))).getSExtValue();
      }
    }

    auto red2 = rewriter.replaceOpWithNewOp<stablehlo::ReduceWindowOp>(
        op, restys, operands, reduce.getInitValues(),
        rewriter.getDenseI64ArrayAttr(win_dim),
        (reduce.getWindowStrides() ? rewriter.getDenseI64ArrayAttr(win_strides)
                                   : DenseI64ArrayAttr()),
        (reduce.getBaseDilations()
             ? rewriter.getDenseI64ArrayAttr(base_dialations)
             : DenseI64ArrayAttr()),
        (reduce.getWindowDilations()
             ? rewriter.getDenseI64ArrayAttr(win_dialations)
             : DenseI64ArrayAttr()),
        (reduce.getPadding()
             ? DenseIntElementsAttr::get(
                   RankedTensorType::get(padding_shape,
                                         rewriter.getIntegerType(64)),
                   padding_dialations)
             : DenseIntElementsAttr()));

    red2.getBody().takeBody(reduce.getBody());
    rewriter.eraseOp(reduce);
    return success();
  }
};

// reshape(concat(...)) -> concat(reshape(...))
struct ReshapeOfConcatToConcatOfReshape final
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                     ReshapeOfConcatToConcatOfReshape> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  int64_t computeNewConcatDim(ArrayRef<int64_t> originalShape,
                              ArrayRef<int64_t> reshapedShape,
                              int64_t originalConcatDim) const {
    int64_t flattenedIndex = 0;
    int64_t newConcatDim = 0;

    for (int64_t i = 0; i < originalConcatDim; ++i) {
      flattenedIndex += originalShape[i];
    }

    int64_t cumulativeSize = 0;
    for (int64_t i = 0; i < reshapedShape.size(); ++i) {
      cumulativeSize += reshapedShape[i];
      if (flattenedIndex < cumulativeSize) {
        newConcatDim = i;
        break;
      }
    }
    return newConcatDim;
  }

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp reshapeOp,
                                    PatternRewriter &rewriter) const {
    // Check if the operand of the reshape is a concatenate operation
    auto concatOp =
        reshapeOp.getOperand().getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concatOp)
      return failure();

    SmallVector<int64_t> oneHot(concatOp.getType().getShape().size(), 0);
    oneHot[concatOp.getDimension()] = 1;

    int64_t zero = 0;
    if (!transformReshapeSlice<int64_t>(reshapeOp, oneHot, /*toFill*/ 0, &zero))
      return failure();

    int64_t newDim = -1;
    for (auto &&[i, val] : llvm::enumerate(oneHot)) {
      if (val == 1) {
        assert(newDim == -1);
        newDim = i;
      }
    }
    assert(newDim != -1);

    // Create reshaped operands for the concat operation
    SmallVector<Value> concatOperands;
    for (auto operand : concatOp.getOperands()) {
      auto operandType = cast<RankedTensorType>(operand.getType());
      if (!operandType)
        return failure();

      SmallVector<int64_t> shape(operandType.getShape().begin(),
                                 operandType.getShape().end());

      int64_t one = 1;
      if (!transformReshapeSlice<int64_t>(reshapeOp, shape, /*toFill*/ 1, &one))
        return failure();

      auto newReshapeOp = stablehlo::ReshapeOpCreate(
          rewriter, reshapeOp.getLoc(), operand, shape);
      concatOperands.push_back(newReshapeOp);
    }

    // Create a new concat operation with the reshaped operands
    auto newConcatOp = stablehlo::ConcatenateOp::create(
        rewriter, reshapeOp.getLoc(), concatOperands, newDim);
    rewriter.replaceOp(reshapeOp, newConcatOp);
    return success();
  }
};

// reshape(reduce_window(...)) -> reduce_window(reshape(...))
struct ReshapeReduceWindow final
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                     ReshapeReduceWindow> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp reshapeOp,
                                    PatternRewriter &rewriter) const {
    // Check if the operand of the reshape is a reduce_window operation
    auto reduceWindow =
        reshapeOp.getOperand().getDefiningOp<stablehlo::ReduceWindowOp>();
    if (!reduceWindow)
      return failure();

    // Check if there is any non-reshape user of this reduce_window operation
    if (llvm::any_of(reduceWindow->getUsers(), [&](Operation *user) {
          return !isa<stablehlo::ReshapeOp>(user);
        }))
      return failure();

    auto reduceWindowType = dyn_cast<ShapedType>(reduceWindow.getType(0));
    if (!reduceWindowType || !reduceWindowType.hasStaticShape())
      return failure();
    ArrayRef<int64_t> reduceWindowShape = reduceWindowType.getShape();

    auto reshapeType = dyn_cast<ShapedType>(reshapeOp.getType());
    if (!reshapeType || !reshapeType.hasStaticShape())
      return failure();
    ArrayRef<int64_t> reshapeShape = reshapeType.getShape();

    // Create reshaped operands for the reduce_window operation
    SmallVector<int64_t> windowDims;
    SmallVector<int64_t> windowStrides;
    SmallVector<int64_t> windowDilations;
    SmallVector<int64_t> baseDilations;
    size_t paddingSize =
        reduceWindow.getPadding() ? 2 * reshapeShape.size() : 0;
    SmallVector<int64_t> padding(paddingSize);

    size_t reduceWindowDim = 0;
    size_t reshapeDim = 0;

    while (reduceWindowDim < reduceWindowShape.size() &&
           reshapeDim < reshapeShape.size()) {
      // Original dimension, not being reshaped
      if (reduceWindowShape[reduceWindowDim] == reshapeShape[reshapeDim]) {
        windowDims.push_back(
            reduceWindow.getWindowDimensions()[reduceWindowDim]);
        if (reduceWindow.getWindowStrides())
          windowStrides.push_back(
              (*reduceWindow.getWindowStrides())[reduceWindowDim]);
        if (reduceWindow.getWindowDilations())
          windowDilations.push_back(
              (*reduceWindow.getWindowDilations())[reduceWindowDim]);
        if (reduceWindow.getBaseDilations())
          baseDilations.push_back(
              (*reduceWindow.getBaseDilations())[reduceWindowDim]);
        if (reduceWindow.getPadding()) {
          padding[2 * reshapeDim] =
              (*(reduceWindow.getPadding()->begin() + (2 * reduceWindowDim)))
                  .getSExtValue();
          padding[2 * reshapeDim + 1] = (*(reduceWindow.getPadding()->begin() +
                                           (2 * reduceWindowDim + 1)))
                                            .getSExtValue();
        }
        reduceWindowDim++;
        reshapeDim++;
        continue;
      }
      // Unit dimension dropped by reshape
      if (reduceWindowShape[reduceWindowDim] == 1) {
        reduceWindowDim++;
        continue;
      }
      // Unit dimension added by reshape
      if (reshapeShape[reshapeDim] == 1) {
        windowDims.push_back(1);
        if (reduceWindow.getWindowStrides())
          windowStrides.push_back(1);
        if (reduceWindow.getWindowDilations())
          windowDilations.push_back(1);
        if (reduceWindow.getBaseDilations())
          baseDilations.push_back(1);
        if (reduceWindow.getPadding()) {
          padding[2 * reshapeDim] = 0;
          padding[2 * reshapeDim + 1] = 0;
        }
        reshapeDim++;
        continue;
      }
      // Adding/dropping non-unit dimension is not supported
      return failure();
    }

    while (reduceWindowDim < reduceWindowShape.size()) {
      if (reduceWindowShape[reduceWindowDim] != 1)
        return failure();
      reduceWindowDim++;
    }

    while (reshapeDim < reshapeShape.size()) {
      windowDims.push_back(1);
      if (reduceWindow.getWindowStrides())
        windowStrides.push_back(1);
      if (reduceWindow.getWindowDilations())
        windowDilations.push_back(1);
      if (reduceWindow.getBaseDilations())
        baseDilations.push_back(1);
      if (reduceWindow.getPadding()) {
        padding[2 * reshapeDim] = 0;
        padding[2 * reshapeDim + 1] = 0;
      }
      reshapeDim++;
    }

    auto newReshapeOp = stablehlo::ReshapeOpCreate(rewriter, reshapeOp.getLoc(),
                                                   reduceWindow.getInputs()[0],
                                                   reshapeType.getShape());
    auto newReduceWindowOp =
        rewriter.replaceOpWithNewOp<stablehlo::ReduceWindowOp>(
            reshapeOp, TypeRange(reshapeType), ValueRange(newReshapeOp),
            reduceWindow.getInitValues(),
            rewriter.getDenseI64ArrayAttr(windowDims),
            rewriter.getDenseI64ArrayAttr(windowStrides),
            rewriter.getDenseI64ArrayAttr(baseDilations),
            rewriter.getDenseI64ArrayAttr(windowDilations),
            DenseIntElementsAttr::get(
                RankedTensorType::get({(int64_t)reshapeShape.size(), 2},
                                      rewriter.getIntegerType(64)),
                padding));

    newReduceWindowOp.getBody().takeBody(reduceWindow.getBody());
    rewriter.eraseOp(reduceWindow);
    return success();
  }
};

struct ReshapeElementwise final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeElementwise> {
  bool onlySingleUser;
  bool onlyFusible;

  ReshapeElementwise(bool onlySingleUser, bool onlyFusible,
                     MLIRContext *context, PatternBenefit benefit = 1,
                     ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        onlySingleUser(onlySingleUser), onlyFusible(onlyFusible) {}

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto elem = op.getOperand().getDefiningOp();
    if (!elem || !stablehlo::hasTraitElementwise(elem) ||
        (onlySingleUser && !llvm::hasSingleElement(elem->getUsers()))) {
      return failure();
    }

    // Collect all elementwise operations in the chain using recursive DFS.
    // By inserting ops after processing their operands, we get topological
    // order without needing a separate sort pass.
    SmallVector<Operation *> chainOps;
    DenseSet<Operation *> visited;
    SmallVector<Value> leafOperands;

    std::function<void(Operation *)> collectChain = [&](Operation *op) {
      if (!visited.insert(op).second)
        return;

      // Process operands first to ensure dependencies come before this op
      for (Value operand : op->getOperands()) {
        auto defOp = operand.getDefiningOp();
        if (!defOp || !stablehlo::hasTraitElementwise(defOp)) {
          // Not elementwise - this is a leaf operand's defining op
          leafOperands.push_back(operand);
          continue;
        }

        // Check single user constraint for intermediate ops
        if (onlySingleUser && defOp != elem &&
            !llvm::hasSingleElement(defOp->getUsers())) {
          // Treat as a leaf if it has multiple users
          leafOperands.push_back(operand);
          continue;
        }

        // Recursively collect the defining op (which adds it before this op)
        collectChain(defOp);
      }

      // Add this op after its operands have been processed
      chainOps.push_back(op);
    };

    collectChain(elem);

    if (chainOps.empty()) {
      return failure();
    }

    // Check fusibility constraint on all leaf operands after collection
    if (onlyFusible) {
      for (Value leaf : leafOperands) {
        if (auto defOp = leaf.getDefiningOp()) {
          if (!isFusible(defOp, op)) {
            return failure();
          }
        } else { // Block argument - not fusible
          return failure();
        }
      }
    }

    // Map from old values to reshaped values
    DenseMap<Value, Value> reshapedValues;

    // Helper to get or create reshaped value
    auto getOrCreateReshapedValue = [&](Value v) -> Value {
      if (reshapedValues.contains(v)) {
        return reshapedValues[v];
      }

      Value newOperand;
      DenseElementsAttr cstAttr;
      if (matchPattern(v, m_Constant(&cstAttr))) {
        newOperand = stablehlo::ConstantOp::create(
            rewriter, op->getLoc(),
            cstAttr.reshape(RankedTensorType::get(
                op.getType().getShape(),
                cast<RankedTensorType>(v.getType()).getElementType())));
      } else {
        Value reshapeOperand = v;
        if (auto reshapeOp = v.getDefiningOp<stablehlo::ReshapeOp>()) {
          reshapeOperand = reshapeOp.getOperand();
        }
        newOperand = stablehlo::ReshapeOpCreate(
            rewriter, op->getLoc(), reshapeOperand, op.getType().getShape());
      }
      reshapedValues[v] = newOperand;
      return newOperand;
    };

    // Rewrite each operation in the chain in topological order
    auto sz = cast<RankedTensorType>(op.getType()).getShape();

    for (auto leaf : leafOperands) {
      // process all the leaf operands first
      rewriter.setInsertionPointAfterValue(leaf);
      getOrCreateReshapedValue(leaf);
    }

    for (Operation *elemOp : chainOps) {
      SmallVector<Value> newOperands;
      rewriter.setInsertionPoint(elemOp);
      for (Value v : elemOp->getOperands()) {
        assert(reshapedValues.contains(v));
        newOperands.push_back(reshapedValues[v]);
      }

      auto newOp = rewriter.create(
          elemOp->getLoc(), elemOp->getName().getIdentifier(),
          ValueRange(newOperands),
          TypeRange(RankedTensorType::get(
              sz, cast<ShapedType>(elemOp->getResult(0).getType())
                      .getElementType())),
          elemOp->getAttrs(), {}, {});

      // Map the result of the original op to the result of the new op
      reshapedValues[elemOp->getResult(0)] = newOp->getResult(0);
    }

    // Replace the reshape op with the reshaped version of the top-level
    // elementwise op
    rewriter.replaceOp(op, reshapedValues[elem->getResult(0)]);
    return success();
  }
};

// slice(transpose x) -> transpose(slice x)
struct SliceReshapeTranspose final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceReshapeTranspose> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()))
      return failure();

    auto transpose =
        reshape.getOperand().getDefiningOp<stablehlo::TransposeOp>();
    if (!transpose || !llvm::hasSingleElement(transpose->getUsers()))
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    auto newslice =
        sliceTransposeHelper(transpose, rewriter, starts, limits, strides);
    auto newtransp = stablehlo::TransposeOp::create(
        rewriter, transpose.getLoc(), newslice, transpose.getPermutation());
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      newtransp);
    return success();
  }
};

struct SliceReshapeDotGeneral
    : public CheckedOpRewritePattern<stablehlo::SliceOp,
                                     SliceReshapeDotGeneral> {
  using CheckedOpRewritePattern<
      stablehlo::SliceOp, SliceReshapeDotGeneral>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()))
      return failure();

    auto dot = reshape.getOperand().getDefiningOp<stablehlo::DotGeneralOp>();
    if (!dot) {
      return rewriter.notifyMatchFailure(op, "defining op is not a reshape");
    }
    if (!llvm::hasSingleElement(dot->getUsers()))
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    auto &&[lhs2, rhs2, resTy] =
        sliceDotGeneralHelper(dot, starts, limits, strides, rewriter);

    Value operands[2] = {lhs2, rhs2};
    auto newdot = stablehlo::DotGeneralOp::create(
        rewriter, op.getLoc(), TypeRange(resTy), operands, dot->getAttrs());

    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(), newdot);
    return success();
  }
};

struct ReshuffleAndsCompares final
    : CheckedOpRewritePattern<stablehlo::AndOp, ReshuffleAndsCompares> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AndOp andOp,
                                    PatternRewriter &rewriter) const {
    SmallVector<Value> conjuncts;
    SmallVector<Operation *> worklist;
    worklist.push_back(andOp);
    while (!worklist.empty()) {
      auto andOp = cast<stablehlo::AndOp>(worklist.pop_back_val());
      if (auto lhsAndOp = andOp.getLhs().getDefiningOp<stablehlo::AndOp>()) {
        worklist.push_back(lhsAndOp);
      } else {
        conjuncts.push_back(andOp.getLhs());
      }
      if (auto lhsAndOp = andOp.getRhs().getDefiningOp<stablehlo::AndOp>()) {
        worklist.push_back(lhsAndOp);
      } else {
        conjuncts.push_back(andOp.getRhs());
      }
    }
    if (conjuncts.size() <= 2)
      return failure();

    Value compareLhs = nullptr;
    auto compares = llvm::filter_to_vector(conjuncts, [&](Value v) {
      auto cmpOp = v.getDefiningOp<stablehlo::CompareOp>();
      if (!cmpOp)
        return false;
      if (compareLhs == nullptr)
        compareLhs = cmpOp.getLhs();
      else if (compareLhs != cmpOp.getLhs())
        return false;
      return cmpOp.getComparisonDirection() ==
             stablehlo::ComparisonDirection::LE;
    });

    if (compares.size() <= 1)
      return failure();

    Value running = compares[0].getDefiningOp<stablehlo::CompareOp>().getRhs();
    for (unsigned i = 1, e = compares.size(); i < e; ++i) {
      Value minRhs = compares[i].getDefiningOp<stablehlo::CompareOp>().getRhs();
      running =
          stablehlo::MinOp::create(rewriter, andOp.getLoc(), running, minRhs);
    }
    Value replacement = stablehlo::CompareOp::create(
        rewriter, compares[0].getLoc(), compareLhs, running,
        stablehlo::ComparisonDirection::LE);

    for (Value conjunct : conjuncts) {
      if (llvm::is_contained(compares, conjunct))
        continue;
      replacement = stablehlo::AndOp::create(rewriter, andOp.getLoc(),
                                             replacement, conjunct);
    }
    rewriter.replaceOp(andOp, replacement);
    return success();
  }
};

struct SliceReshapeSlice final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceReshapeSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()))
      return failure();

    auto prev = reshape.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (!prev)
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    sliceSliceHelper(prev, starts, limits, strides);
    auto newslice = stablehlo::SliceOp::create(
        rewriter, op.getLoc(), prev.getOperand(), starts, limits, strides);
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      newslice);
    return success();
  }
};

// if the slice chain is defined inside a loop, and atleast one of the indices
// of the slice is a loop induction variable, then it makes sense to simplify
// the chain to enable loop optimizations (even if there are multiple users
// of the intermediate)
bool DSDSSimplificationSingleUserCheckException(SmallVector<Operation *> ops) {
  if (ops.empty()) {
    return false;
  }

  auto getParentWhileOp = [](auto op) {
    // this is a more strict check than getParentOfType
    auto parentOp = op->getParentOp();
    return dyn_cast_or_null<stablehlo::WhileOp>(parentOp);
  };

  auto whileOp = getParentWhileOp(ops.front());
  if (!whileOp) {
    return false;
  }

  SmallVector<stablehlo::DynamicSliceOp> dynSliceOps;
  dynSliceOps.reserve(ops.size());
  if (auto dslice = dyn_cast<stablehlo::DynamicSliceOp>(ops.front())) {
    dynSliceOps.push_back(dslice);
  }

  for (auto op : llvm::drop_begin(ops)) {
    auto parent = getParentWhileOp(op);
    if (!parent || whileOp != parent) {
      return false;
    }
    if (auto dslice = dyn_cast<stablehlo::DynamicSliceOp>(op)) {
      dynSliceOps.push_back(dslice);
    }
  }

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

  // check that atleast one of the dynamic slices have a loop iteration argument
  // as their start index.
  WhileLoopInfo loopInfo(whileOp);
  if (loopInfo.computeInfo().succeeded() && !loopInfo.isValid()) {
    return false;
  }

  auto affineIndexInfo = loopInfo.getAffineIndexInfo();
  bool anyIndVarDependent = false;
  for (auto slice : dynSliceOps) {
    for (auto start : slice.getStartIndices()) {
      if (affineIndexInfo.contains(start)) {
        anyIndVarDependent = true;
        goto loopexit;
      }
    }
  }
loopexit:
  return anyIndVarDependent;
}

struct DynamicSliceReshapeDynamicSlice final
    : CheckedOpRewritePattern<stablehlo::DynamicSliceOp,
                              DynamicSliceReshapeDynamicSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    auto prev = reshape.getOperand().getDefiningOp<stablehlo::DynamicSliceOp>();
    if (!prev)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()) &&
        !DSDSSimplificationSingleUserCheckException({reshape, op, prev}))
      return failure();

    SmallVector<Value> startIndices;
    SmallVector<int64_t> sliceSizes;
    if (!sliceReshapeHelper(op, startIndices, sliceSizes).succeeded())
      return failure();

    rewriter.setInsertionPoint(op);
    sliceSliceHelper(prev, startIndices, sliceSizes, rewriter);
    auto newSlice = stablehlo::DynamicSliceOp::create(
        rewriter, op.getLoc(), prev.getOperand(),
        promoteToLargestBitWidth(startIndices, rewriter),
        rewriter.getDenseI64ArrayAttr(sliceSizes));
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      newSlice);
    return success();
  }
};

struct SliceReshapeDynamicSlice final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceReshapeDynamicSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    auto prev = reshape.getOperand().getDefiningOp<stablehlo::DynamicSliceOp>();
    if (!prev)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()) &&
        !DSDSSimplificationSingleUserCheckException({reshape, op, prev}))
      return failure();

    SmallVector<int64_t> starts, limits, strides;
    if (!sliceReshapeHelper(op, starts, limits, strides).succeeded())
      return failure();

    for (auto stride : strides) {
      if (stride != 1) // dynamic slice only supports stride 1
        return failure();
    }

    SmallVector<Value> startIndices;
    SmallVector<int64_t> sliceSizes;
    auto elemTy = RankedTensorType::get({}, rewriter.getI32Type());
    for (auto &&[start, end] : llvm::zip(starts, limits)) {
      startIndices.push_back(stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), elemTy,
          cast<ElementsAttr>(makeAttr(elemTy, start))));
      sliceSizes.push_back(end - start);
    }

    rewriter.setInsertionPoint(op);
    sliceSliceHelper(prev, startIndices, sliceSizes, rewriter);
    auto newSlice = stablehlo::DynamicSliceOp::create(
        rewriter, op.getLoc(), prev.getOperand(),
        promoteToLargestBitWidth(startIndices, rewriter),
        rewriter.getDenseI64ArrayAttr(sliceSizes));
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      newSlice);
    return success();
  }
};

struct DynamicSliceReshapeSlice final
    : CheckedOpRewritePattern<stablehlo::DynamicSliceOp,
                              DynamicSliceReshapeSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reshape = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshape)
      return failure();

    auto prev = reshape.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (!prev)
      return failure();

    if (!llvm::hasSingleElement(reshape->getUsers()) &&
        !DSDSSimplificationSingleUserCheckException({op, reshape, prev}))
      return failure();

    SmallVector<Value> startIndices;
    SmallVector<int64_t> sliceSizes;
    if (!sliceReshapeHelper(op, startIndices, sliceSizes).succeeded())
      return failure();

    auto prevDS = SliceToDynamicSlice(prev, rewriter);
    if (!prevDS)
      return failure();

    rewriter.setInsertionPoint(op);
    sliceSliceHelper(prevDS, startIndices, sliceSizes, rewriter);

    auto newSlice = stablehlo::DynamicSliceOpCreate(
        rewriter, op.getLoc(), prevDS.getOperand(),
        promoteToLargestBitWidth(startIndices, rewriter), sliceSizes);
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      newSlice);
    return success();
  }
};

} // namespace

// Rewritten from
// https:github.com/openxla/stablehlo/blob/4f180d3c2236a15f82f29aad1b47f6ea2c14fc52/stablehlo/reference/Ops.cpp#L1381
// using https://openxla.org/xla/operation_semantics#gather
// becuase xla/openxla differ in semantics
stablehlo::Index index(stablehlo::Tensor tensor) {
  stablehlo::Index result;
  for (auto it = tensor.index_begin(); it != tensor.index_end(); ++it)
    result.push_back(tensor.get(*it).getIntegerValue().getSExtValue());
  return result;
}

stablehlo::Tensor sliceOp(const stablehlo::Tensor &operand,
                          const stablehlo::Index &index) {
  using namespace stablehlo;
  Sizes start, limit;
  for (auto i = 0; i < operand.getRank(); ++i) {
    if (index[i] == -1) {
      start.push_back(0);
      limit.push_back(operand.getShape()[i]);
    } else {
      start.push_back(index[i]);
      limit.push_back(index[i] + 1);
    }
  }
  Sizes strides(operand.getRank(), 1);

  SmallVector<Type> inferredTypes;
  Builder builder(operand.getType().getContext());
  auto inferStatus = hlo::inferSliceOp({}, operand.getType(), start, limit,
                                       strides, inferredTypes);
  (void)inferStatus;
  assert(!failed(inferStatus));

  return stablehlo::sliceOp(operand, start, strides,
                            cast<mlir::ShapedType>(inferredTypes[0]));
}

bool is_iota(ArrayRef<int64_t> idx) {
  for (auto en : llvm::enumerate(idx))
    if (en.index() != en.value())
      return false;
  return true;
}

struct CSEIota : CheckedOpRewritePattern<stablehlo::IotaOp, CSEIota> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IotaOp op,
                                    PatternRewriter &rewriter) const {
    bool anyCsed = false;
    Operation *next = op->getNextNode();
    while (next) {
      if (auto iota = dyn_cast<stablehlo::IotaOp>(next)) {
        next = next->getNextNode();
        if (iota.getIotaDimension() == op.getIotaDimension() &&
            iota.getType() == op.getType()) {
          rewriter.replaceOp(iota, op);
          anyCsed = true;
        }
      } else
        break;
    }
    return success(anyCsed);
  }
};

bool opaque_cmp(Value a, Value b) {
  return a.getAsOpaquePointer() < b.getAsOpaquePointer();
}

bool isCommutativeEquivalent(ValueRange lhs, ValueRange rhs) {
  if (lhs.size() != rhs.size())
    return false;

  if (lhs.size() == 1) {
    return lhs[0] == rhs[0];
  } else if (lhs.size() == 2) {
    return (lhs[0] == rhs[0] && lhs[1] == rhs[1]) ||
           (lhs[1] == rhs[0] && lhs[0] == rhs[1]);
  } else {
    auto lhsv = llvm::to_vector(lhs);
    auto rhsv = llvm::to_vector(rhs);
    llvm::sort(lhsv, opaque_cmp);
    llvm::sort(rhsv, opaque_cmp);
    return lhsv == rhsv;
  }
}

template <typename T> struct CSE final : CheckedOpRewritePattern<T, CSE<T>> {
  using CheckedOpRewritePattern<T, CSE<T>>::CheckedOpRewritePattern;

  bool supportsDynamicShapes() { return true; }

  LogicalResult matchAndRewriteImpl(T op, PatternRewriter &rewriter) const {
    if (op->getNumOperands() > 0)
      for (auto nop : op->getOperand(0).getUsers()) {
        if (nop == op)
          continue;
        if (!isa<T>(nop))
          continue;
        if (nop->getBlock() != op->getBlock())
          continue;

        if (op->getName() != nop->getName())
          continue;

        OperationEquivalence::Flags flags =
            OperationEquivalence::IgnoreLocations |
            OperationEquivalence::IgnoreDiscardableAttrs;

        // OperationEquivalence only checks for mlir::OpTrait::IsCommutative
        if constexpr (!std::is_base_of_v<::mlir::OpTrait::IsCommutative<T>,
                                         T>) {
          flags |= OperationEquivalence::IgnoreCommutativity;
        }

        if (!OperationEquivalence::isEquivalentTo(op, nop, flags)) {
          // stablehlo defines a special trait for commutative operations.
          // check for that here.
          if constexpr (std::is_base_of_v<
                            ::mlir::hlo::OpTrait::IsCommutative<T>, T>) {
            auto opRange = op->getOperands();
            auto nopRange = nop->getOperands();
            if (!isCommutativeEquivalent(opRange, nopRange))
              continue;
          } else {
            continue;
          }
        }

        if (nop->isBeforeInBlock(op)) {
          rewriter.replaceOp(op, nop);
          return success();
        } else {
          rewriter.replaceOp(nop, op);
          return success();
        }
      }
    return failure();
  }
};

struct ConstPropThroughBarrier final
    : CheckedOpRewritePattern<stablehlo::OptimizationBarrierOp,
                              ConstPropThroughBarrier> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::OptimizationBarrierOp op,
                                    PatternRewriter &rewriter) const {

    SmallVector<Value> replacements;
    bool changed = false;
    for (auto &&[res, inp] : llvm::zip(op.getResults(), op.getOperands())) {
      if (inp.getDefiningOp<stablehlo::ConstantOp>()) {
        changed = true;
      } else if (res.use_empty()) {
        changed = true;
      } else {
        replacements.push_back(inp);
      }
    }
    if (!changed) {
      return failure();
    }

    auto nop = stablehlo::OptimizationBarrierOp::create(rewriter, op.getLoc(),
                                                        replacements);

    size_t idx = 0;
    SmallVector<Value> results;
    for (auto &&[res, inp] : llvm::zip(op.getResults(), op.getOperands())) {
      if (res.use_empty() || inp.getDefiningOp<stablehlo::ConstantOp>()) {
        results.push_back(inp);
      } else {
        results.push_back(nop->getResult(idx));
        idx++;
      }
    }
    rewriter.replaceOp(op, results);
    return success();
  }
};

// If there is a dus where part of the updated region is not used by later slice
// operations, pre-slice the operand and update to the original dus
struct DUSSliceSimplify final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                              DUSSliceSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dusOp,
                                    PatternRewriter &rewriter) const {
    auto res = dusOp.getResult();
    auto update = dusOp.getUpdate();
    auto updateShape = update.getType().getShape();
    auto operand = dusOp.getOperand();

    auto resShape = res.getType().getShape();
    auto resRank = res.getType().getRank();

    SmallVector<int64_t> ignoredStart(resShape);
    SmallVector<int64_t> ignoredEnd(resRank, 0);
    SmallVector<stablehlo::SliceOp> slices;
    slices.reserve(resRank);
    for (auto &use : res.getUses()) {
      if (auto slice = dyn_cast<stablehlo::SliceOp>(use.getOwner())) {
        ignoredStart = llvm::map_to_vector(
            llvm::zip(slice.getStartIndices(), ignoredStart), [&](auto p) {
              auto &[start, oldIgnoreStart] = p;
              if (start < oldIgnoreStart)
                return start;
              return oldIgnoreStart;
            });
        ignoredEnd = llvm::map_to_vector(
            llvm::zip(slice.getLimitIndices(), ignoredEnd), [&](auto p) {
              auto &[thisIgnoreEnd, oldIgnoreEnd] = p;
              if (thisIgnoreEnd > oldIgnoreEnd)
                return thisIgnoreEnd;
              return oldIgnoreEnd;
            });
        slices.push_back(slice);
      } else {
        return rewriter.notifyMatchFailure(use.getOwner(),
                                           "Found non slice user");
      }
    }

    if (slices.empty())
      return rewriter.notifyMatchFailure(dusOp, "No slice users found");

    if (llvm::all_of(ignoredStart,
                     [](int64_t ignored) { return ignored == 0; }) &&
        llvm::all_of(llvm::zip(ignoredEnd, resShape), [](auto p) {
          auto &[i, s] = p;
          return i == s;
        }))
      return rewriter.notifyMatchFailure(dusOp, "No ignored regions");

    assert(llvm::all_of(llvm::zip(ignoredStart, ignoredEnd), [](auto p) {
      auto &[s, e] = p;
      return s <= e;
    }));

    SmallVector<Value> dusIndexVals = dusOp.getStartIndices();
    SmallVector<int64_t> dusStartIndices;
    dusStartIndices.reserve(resRank);
    for (Value idxVal : dusIndexVals) {
      DenseIntElementsAttr idxAttr;
      if (!matchPattern(idxVal, m_Constant(&idxAttr)) ||
          idxAttr.getNumElements() != 1)
        return rewriter.notifyMatchFailure(
            dusOp, "DUS indices must be constant scalars");
      dusStartIndices.push_back((*idxAttr.begin()).getSExtValue());
    }
    SmallVector<int64_t> duslimitIndices = llvm::map_to_vector(
        llvm::zip(dusStartIndices, updateShape),
        [](auto p) { return std::get<0>(p) + std::get<1>(p); });

    SmallVector<int64_t> strideOne(resRank, 1);

    auto loc = dusOp->getLoc();
    auto preSliceOperand = stablehlo::SliceOp::create(
        rewriter, loc, operand, ignoredStart, ignoredEnd, strideOne);

    auto ignoredUpdateStart = llvm::map_to_vector(
        llvm::zip(ignoredStart, dusStartIndices), [](auto p) -> int64_t {
          auto &[ignoredStart, dusStart] = p;
          if (ignoredStart > dusStart)
            return ignoredStart - dusStart;
          return 0;
        });
    auto ignoredUpdateEnd =
        llvm::map_to_vector(llvm::zip(ignoredEnd, duslimitIndices, updateShape),
                            [](auto p) -> int64_t {
                              auto &[ignoredEnd, dusEnd, updateShape] = p;
                              if (ignoredEnd < dusEnd)
                                return updateShape - (dusEnd - ignoredEnd);
                              return updateShape;
                            });
    Value preSliceUpdate = stablehlo::SliceOp::create(
        rewriter, loc, update, ignoredUpdateStart, ignoredUpdateEnd, strideOne);

    Type itype = dusIndexVals[0].getType();
    SmallVector<Value> newDusIndices = llvm::map_to_vector(
        llvm::zip(ignoredStart, dusStartIndices), [&](auto p) -> Value {
          auto &[ignored, dusStart] = p;
          int64_t start = 0;
          if (ignored < dusStart)
            start = dusStart - ignored;

          return stablehlo::ConstantOp::create(
              rewriter, loc, itype, cast<ElementsAttr>(makeAttr(itype, start)));
        });

    LLVM_DEBUG(
        for (auto [idx, operandSize, updateSize]
             : llvm::zip_equal(
                 newDusIndices,
                 cast<RankedTensorType>(preSliceOperand.getType()).getShape(),
                 cast<RankedTensorType>(preSliceUpdate.getType()).getShape())) {
          APInt start;
          assert(matchPattern(idx, m_ConstantInt(&start)));
          int64_t vali = start.getSExtValue();
          assert(operandSize >= vali + updateSize);
        });

    auto newDus = stablehlo::DynamicUpdateSliceOp::create(
        rewriter, loc, preSliceOperand, preSliceUpdate, newDusIndices);

    for (auto slice : slices) {
      SmallVector<int64_t> starts(slice.getStartIndices());
      SmallVector<int64_t> sizes(slice.getResult().getType().getShape());
      SmallVector<int64_t> limits;
      starts = llvm::map_to_vector(llvm::zip(starts, ignoredStart),
                                   [](auto p) -> int64_t {
                                     auto &[s, i] = p;
                                     assert(s >= i);
                                     return s - i;
                                   });
      limits =
          llvm::map_to_vector(llvm::zip(starts, sizes), [](auto p) -> int64_t {
            auto &[start, size] = p;
            return start + size;
          });
      rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
          slice, newDus, starts, limits, slice.getStrides());
    }
    rewriter.eraseOp(dusOp);

    return success();
  }
};

struct DUSToI32 final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp, DUSToI32> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dusOp,
                                    PatternRewriter &rewriter) const {
    auto i32 = rewriter.getI32Type();

    auto unrankedI32 = RankedTensorType::get({}, i32);

    auto startIndices = dusOp.getStartIndices();

    SmallVector<int64_t> newStartIndicesConst;
    for (auto idx : startIndices) {
      if (cast<RankedTensorType>(idx.getType()).getElementType() == i32)
        return failure();

      llvm::APInt val;
      if (!matchPattern(idx, m_ConstantInt(&val)))
        return failure();

      if (val.getSignificantBits() > 32)
        return failure();

      newStartIndicesConst.push_back(val.getZExtValue());
    }

    SmallVector<Value> newStartIndices;
    for (auto [val, idx] :
         llvm::zip_equal(newStartIndicesConst, startIndices)) {
      newStartIndices.push_back(stablehlo::ConstantOp::create(
          rewriter, idx.getLoc(), unrankedI32,
          cast<ElementsAttr>(makeAttr(unrankedI32, val))));
    }

    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
        dusOp, dusOp.getOperand(), dusOp.getUpdate(), newStartIndices);

    return success();
  }
};

// Replaces DUS with a combination of slices and concats.
// Each run of the pattern handles one dimension at a time.
struct DUSToConcat final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp, DUSToConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dusOp,
                                    PatternRewriter &rewriter) const {

    // --- Get Info & Check Static Requirements ---
    Value targetOperand = dusOp.getOperand();
    Value updateVal = dusOp.getUpdate();
    SmallVector<Value> dusIndexVals = dusOp.getStartIndices();

    auto targetType = dyn_cast<RankedTensorType>(targetOperand.getType());
    auto updateType = dyn_cast<RankedTensorType>(updateVal.getType());

    if (!targetType || !updateType || !targetType.hasStaticShape() ||
        !updateType.hasStaticShape()) {
      return rewriter.notifyMatchFailure(dusOp, "Requires static shapes");
    }

    ArrayRef<int64_t> targetShape = targetType.getShape();
    ArrayRef<int64_t> updateShape = updateType.getShape();
    int rank = targetType.getRank();

    if (rank == 0) {
      // Replace scalar DUS with the update value directly
      rewriter.replaceOp(dusOp, updateVal);
      return success();
    }

    SmallVector<int64_t> dusStartIndices;
    dusStartIndices.reserve(rank);
    for (Value idxVal : dusIndexVals) {
      DenseIntElementsAttr idxAttr;
      if (!matchPattern(idxVal, m_Constant(&idxAttr)) ||
          idxAttr.getNumElements() != 1) {
        return rewriter.notifyMatchFailure(
            dusOp, "Requires constant scalar start indices");
      }
      dusStartIndices.push_back((*idxAttr.begin()).getSExtValue());
    }

    // --- Check Constraints ---
    int differingDim = -1;
    bool singleDifferingDim = false;
    for (int d = 0; d < rank; ++d) {
      // Check if update spans the full dimension
      if (dusStartIndices[d] == 0 && updateShape[d] == targetShape[d]) {
        continue; // Spans fully, this dimension is fine
      }
      if (differingDim == -1)
        singleDifferingDim = true;
      else
        singleDifferingDim = false;
      differingDim = d;
    }

    if (differingDim == -1) {
      // DUS covers the entire tensor, replace with update value
      if (targetType == updateType) {
        rewriter.replaceOp(dusOp, updateVal);
        return success();
      } else {
        // Types mismatch (shouldn't happen with valid DUS?), fail for safety
        return rewriter.notifyMatchFailure(
            dusOp, "Full tensor update but types mismatch");
      }
    }

    // Check if the differing dimension is a prefix or suffix update
    bool dusFromBeginning = (dusStartIndices[differingDim] == 0);
    bool dusToEnd =
        (dusStartIndices[differingDim] + updateShape[differingDim] ==
         targetShape[differingDim]);

    Location loc = dusOp.getLoc();
    SmallVector<Value> newDusIndices = dusIndexVals;

    SmallVector<int64_t> sliceStrides(rank, 1);

    Value slicePre;
    Value slicePost;

    auto getPrePost = [&](int64_t sliceAt) {
      SmallVector<int64_t> slicePreStarts(rank, 0);
      SmallVector<int64_t> slicePreLimits = llvm::to_vector(targetShape);

      SmallVector<int64_t> slicePostStarts(rank, 0);
      SmallVector<int64_t> slicePostLimits = llvm::to_vector(targetShape);

      slicePreStarts[differingDim] = 0;
      slicePreLimits[differingDim] = sliceAt;

      slicePostStarts[differingDim] = sliceAt;
      slicePostLimits[differingDim] = targetShape[differingDim];

      slicePre = stablehlo::SliceOp::create(rewriter, loc, targetOperand,
                                            slicePreStarts, slicePreLimits,
                                            sliceStrides);

      slicePost = stablehlo::SliceOp::create(rewriter, loc, targetOperand,
                                             slicePostStarts, slicePostLimits,
                                             sliceStrides);
    };

    auto itype = dusIndexVals[differingDim].getType();
    auto c0 = stablehlo::ConstantOp::create(
        rewriter, loc, itype, cast<ElementsAttr>(makeAttr(itype, 0)));
    newDusIndices[differingDim] = c0;

    if (!dusFromBeginning) {
      int64_t sliceAt = dusStartIndices[differingDim];
      getPrePost(sliceAt);

      auto c0 = stablehlo::ConstantOp::create(
          rewriter, loc, itype, cast<ElementsAttr>(makeAttr(itype, 0)));
      newDusIndices[differingDim] = c0;

      if (!(singleDifferingDim && dusToEnd)) {
        slicePost = stablehlo::DynamicUpdateSliceOp::create(
            rewriter, loc, slicePost, updateVal, newDusIndices);
      }

    } else {
      assert(!dusToEnd);
      int64_t sliceAt =
          dusStartIndices[differingDim] + updateShape[differingDim];
      getPrePost(sliceAt);

      if (!(singleDifferingDim && dusFromBeginning)) {
        slicePre = stablehlo::DynamicUpdateSliceOp::create(
            rewriter, loc, slicePre, updateVal, newDusIndices);
      }
    }

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
        dusOp, dusOp.getType(), ValueRange{slicePre, slicePost}, differingDim);

    return success();
  }
};

// I.   (div a (div b c)) -> (div (mul a c) b)
// II.  (div (div a b) c) -> (div a (mul b c))
// III. (div (div a b) (div c d)) -> (div (mul a d) (mul b c))
struct DivideDivideSimplify
    : public CheckedOpRewritePattern<stablehlo::DivOp, DivideDivideSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DivOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    if (auto lhsDivOp = lhs.getDefiningOp<stablehlo::DivOp>()) {
      if (!lhsDivOp->hasOneUse())
        return failure();

      if (auto rhsDivOp = rhs.getDefiningOp<stablehlo::DivOp>()) {
        if (!rhsDivOp->hasOneUse())
          return failure();

        // Case III.
        auto numerator = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                                  lhsDivOp->getOperand(0),
                                                  rhsDivOp->getOperand(1));
        auto denominator = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                                    lhsDivOp->getOperand(1),
                                                    rhsDivOp->getOperand(0));
        rewriter.replaceOpWithNewOp<stablehlo::DivOp>(op, numerator,
                                                      denominator);
        return success();
      } else {
        // Case II.
        rewriter.replaceOpWithNewOp<stablehlo::DivOp>(
            op, lhsDivOp->getOperand(0),
            stablehlo::MulOp::create(rewriter, op.getLoc(),
                                     lhsDivOp->getOperand(1), rhs));
        return success();
      }
    } else {
      if (auto rhsDivOp = rhs.getDefiningOp<stablehlo::DivOp>()) {
        if (!rhsDivOp->hasOneUse())
          return failure();

        // Case I.
        rewriter.replaceOpWithNewOp<stablehlo::DivOp>(
            op,
            stablehlo::MulOp::create(rewriter, op.getLoc(), lhs,
                                     rhsDivOp->getOperand(1)),
            rhsDivOp->getOperand(0));
        return success();
      }
    }

    return failure();
  }
};

static stablehlo::ComparisonDirection
invertDirection(stablehlo::ComparisonDirection direction) {
  using stablehlo::ComparisonDirection;

  switch (direction) {
  case ComparisonDirection::EQ:
  case ComparisonDirection::NE:
    return direction;
  case ComparisonDirection::GE:
    return ComparisonDirection::LE;
  case ComparisonDirection::GT:
    return ComparisonDirection::LT;
  case ComparisonDirection::LE:
    return ComparisonDirection::GE;
  case ComparisonDirection::LT:
    return ComparisonDirection::GT;
  }

  llvm::report_fatal_error("Unhandled case");
}

static APInt calculateComp(stablehlo::ComparisonType kind,
                           stablehlo::ComparisonDirection direction,
                           const APInt &lhs, const APInt &rhs) {
  using stablehlo::ComparisonDirection;
  using stablehlo::ComparisonType;
  assert(llvm::is_contained({ComparisonType::SIGNED, ComparisonType::UNSIGNED},
                            kind) &&
         "Not an integer comparison");

  auto asBit = [](bool value) {
    return value ? APInt::getAllOnes(1) : APInt::getZero(1);
  };

  switch (direction) {
  case ComparisonDirection::EQ:
    return asBit(lhs == rhs);
  case ComparisonDirection::NE:
    return asBit(lhs != rhs);
  case ComparisonDirection::GE:
    return asBit(kind == ComparisonType::SIGNED ? lhs.sge(rhs) : lhs.uge(rhs));
  case ComparisonDirection::GT:
    return asBit(kind == ComparisonType::SIGNED ? lhs.sgt(rhs) : lhs.ugt(rhs));
  case ComparisonDirection::LE:
    return asBit(kind == ComparisonType::SIGNED ? lhs.sle(rhs) : lhs.ule(rhs));
  case ComparisonDirection::LT:
    return asBit(kind == ComparisonType::SIGNED ? lhs.slt(rhs) : lhs.ult(rhs));
  }

  llvm_unreachable("Unhandled case");
}

struct CompareOpCanon final
    : CheckedOpRewritePattern<stablehlo::CompareOp, CompareOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp op,
                                    PatternRewriter &rewriter) const {
    RankedTensorType type = op.getType();

    {
      DenseElementsAttr lhs;
      matchPattern(op.getLhs(), m_Constant(&lhs));
      DenseElementsAttr rhs;
      matchPattern(op.getRhs(), m_Constant(&rhs));
      if (lhs && rhs) {
        bool isSplat = lhs.isSplat() && rhs.isSplat();
        auto ty = isSplat
                      ? RankedTensorType::get({}, op.getType().getElementType())
                      : op.getType();
        auto out = fromTensor(stablehlo::compareOp(
            isSplat
                ? stablehlo::makeTensor(lhs.resizeSplat(RankedTensorType::get(
                      {}, lhs.getType().getElementType())))
                : stablehlo::constantOp(lhs),
            isSplat
                ? stablehlo::makeTensor(rhs.resizeSplat(RankedTensorType::get(
                      {}, rhs.getType().getElementType())))
                : stablehlo::constantOp(rhs),
            op.getComparisonDirection(), ty));
        if (isSplat)
          out = out.resizeSplat(op.getType());

        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(),
                                                           out);
        return success();
      }
    }

    // Bail out on non-integer comparison.
    // TODO: Support more comparison types.
    using stablehlo::ComparisonDirection;
    using stablehlo::ComparisonType;
    ComparisonDirection direction = op.getComparisonDirection();
    Value lhs = op.getLhs();
    Value rhs = op.getRhs();
    std::optional<ComparisonType> compType = op.getCompareType();

    TypedAttr lhsAttr;
    matchPattern(lhs, m_Constant(&lhsAttr));

    TypedAttr rhsAttr;
    matchPattern(rhs, m_Constant(&rhsAttr));

    if (compType &&
        llvm::is_contained({ComparisonType::SIGNED, ComparisonType::UNSIGNED},
                           *compType)) {

      if (lhs == rhs) {
        switch (direction) {
        case ComparisonDirection::EQ:
        case ComparisonDirection::GE:
        case ComparisonDirection::LE: {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, SplatElementsAttr::get(type, rewriter.getBoolAttr(true)));
          return success();
        }
        case ComparisonDirection::GT:
        case ComparisonDirection::LT:
        case ComparisonDirection::NE: {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, rewriter.getZeroAttr(type));
          return success();
        }
        }
        llvm_unreachable("Unhandled case");
      }

      // The canonical form has the constant operand as the RHS.
      if (lhsAttr && !rhsAttr) {
        rewriter.modifyOpInPlace(op, [&op, direction, lhs, rhs] {
          op.setComparisonDirection(invertDirection(direction));
          op->setOperands(ValueRange{rhs, lhs});
        });
        return success();
      }

      if (Attribute res;
          lhsAttr && rhsAttr &&
          (res = constFoldBinaryOp<IntegerAttr>(
               ArrayRef<Attribute>({lhsAttr, rhsAttr}), op.getType(),
               [direction, kind = *compType](const APInt &a, const APInt &b) {
                 return calculateComp(kind, direction, a, b);
               }))) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, res);
        return success();
      }
    }

    auto simplifyNonNegative = [&](Attribute attr,
                                   ComparisonDirection direction,
                                   Value otherOperand) -> LogicalResult {
      APInt val;
      APFloat valFloat(APFloat::IEEEdouble());
      bool isSplat = false, isFloat = false;

      if (auto dense = dyn_cast<DenseElementsAttr>(attr)) {
        if (dense.isSplat()) {
          auto splatAttr = dense.getSplatValue<Attribute>();
          if (auto intAttr = dyn_cast<IntegerAttr>(splatAttr)) {
            val = intAttr.getValue();
            isSplat = true;
          } else if (auto floatAttr = dyn_cast<FloatAttr>(splatAttr)) {
            valFloat = floatAttr.getValue();
            isFloat = true;
            isSplat = true;
          }
        }
      }

      if (isSplat) {
        bool alwaysTrue = false;
        bool alwaysFalse = false;

        // Check if the constant is negative or zero
        bool isNegative = isFloat ? valFloat.isNegative() : val.isNegative();
        bool isZeroVal = isFloat ? valFloat.isZero() : val.isZero();

        if ((compType && *compType == ComparisonType::SIGNED) ||
            (isFloat && (!compType || *compType == ComparisonType::FLOAT) &&
             guaranteedNoNanResult(otherOperand, rewriter))) {
          if (isNegative) {
            switch (direction) {
            case ComparisonDirection::EQ:
            case ComparisonDirection::LE:
            case ComparisonDirection::LT:
              alwaysFalse = true;
              break;
            case ComparisonDirection::NE:
            case ComparisonDirection::GE:
            case ComparisonDirection::GT:
              alwaysTrue = true;
              break;
            }
          } else if (isZeroVal) {
            if (direction == ComparisonDirection::LT)
              alwaysFalse = true;
            if (direction == ComparisonDirection::GE)
              alwaysTrue = true;
          }
        } else if (compType && *compType == ComparisonType::UNSIGNED) {
          if (isZeroVal) {
            if (direction == ComparisonDirection::LT)
              alwaysFalse = true;
            if (direction == ComparisonDirection::GE)
              alwaysTrue = true;
          }
        }

        if (alwaysTrue) {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, SplatElementsAttr::get(type, rewriter.getBoolAttr(true)));
          return success();
        }
        if (alwaysFalse) {
          rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
              op, rewriter.getZeroAttr(type));
          return success();
        }
      }
      return failure();
    };

    if (rhsAttr && guaranteedNonNegativeResult(lhs, rewriter)) {
      if (succeeded(simplifyNonNegative(rhsAttr, direction, lhs))) {
        return success();
      }
    }
    if (lhsAttr && guaranteedNonNegativeResult(rhs, rewriter)) {
      if (succeeded(
              simplifyNonNegative(lhsAttr, invertDirection(direction), rhs))) {
        return success();
      }
    }

    return failure();
  }
};

struct CompareExt final
    : CheckedOpRewritePattern<stablehlo::CompareOp, CompareExt> {
  using CheckedOpRewritePattern<stablehlo::CompareOp,
                                CompareExt>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp op,
                                    PatternRewriter &rewriter) const {
    auto elemType =
        cast<RankedTensorType>(op.getLhs().getType()).getElementType();
    if (!elemType.isInteger())
      return failure();

    auto direction = op.getComparisonDirection();

    auto lhsConvert = op.getLhs().getDefiningOp<stablehlo::ConvertOp>();
    auto rhsConvert = op.getRhs().getDefiningOp<stablehlo::ConvertOp>();
    if (!lhsConvert && !rhsConvert)
      return failure();

    auto isConvertFromBool = [](stablehlo::ConvertOp cvtOp) -> bool {
      return cvtOp && cast<RankedTensorType>(cvtOp.getOperand().getType())
                          .getElementType()
                          .isInteger(1);
    };

    if (isConvertFromBool(lhsConvert) && isConvertFromBool(rhsConvert) &&
        direction == stablehlo::ComparisonDirection::EQ) {
      rewriter.replaceOpWithNewOp<stablehlo::AndOp>(op, lhsConvert.getOperand(),
                                                    rhsConvert.getOperand());
      return success();
    }

    if (isConvertFromBool(lhsConvert) &&
        direction == stablehlo::ComparisonDirection::EQ) {
      Attribute rhsAttr;
      if (matchPattern(op.getRhs(), m_Constant(&rhsAttr))) {
        if (matchPattern(rhsAttr, m_One())) {
          rewriter.replaceOp(op, lhsConvert.getOperand());
          return success();
        } else if (matchPattern(rhsAttr, m_Zero())) {
          rewriter.replaceOpWithNewOp<stablehlo::NotOp>(
              op, lhsConvert.getOperand());
          return success();
        }
      }
    }

    if (isConvertFromBool(rhsConvert) &&
        direction == stablehlo::ComparisonDirection::EQ) {
      Attribute lhsAttr;
      if (matchPattern(op.getLhs(), m_Constant(&lhsAttr))) {
        if (matchPattern(lhsAttr, m_One())) {
          rewriter.replaceOp(op, rhsConvert.getOperand());
          return success();
        } else if (matchPattern(lhsAttr, m_Zero())) {
          rewriter.replaceOpWithNewOp<stablehlo::NotOp>(
              op, rhsConvert.getOperand());
          return success();
        }
      }
    }

    return failure();
  }
};

// Extract a scalar integer constant from a splat tensor or a scalar
// broadcast_in_dim of a constant (dims=[]).
static bool extractSplatInt(Value v, int64_t &out) {
  DenseIntElementsAttr attr;
  if (matchPattern(v, m_Constant(&attr)) && attr.isSplat()) {
    out = attr.getSplatValue<IntegerAttr>().getValue().getSExtValue();
    return true;
  }
  if (auto bcast = v.getDefiningOp<stablehlo::BroadcastInDimOp>()) {
    if (bcast.getBroadcastDimensions().empty())
      return extractSplatInt(bcast.getOperand(), out);
  }
  return false;
}

// Matches: select(broadcast_in_dim(compare(iota_expr, K), [dim]), A, B)
//      or: select(compare(iota_expr, K), A, B)   (no broadcast)
// where iota_expr is either iota or add(iota, const_offset).
// Replaces with concat(slice(A|B, ...), ...) along the iota/broadcast
// dimension.
struct SelectCompIotaConstSimplify final
    : CheckedOpRewritePattern<stablehlo::SelectOp,
                              SelectCompIotaConstSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp selectOp,
                                    PatternRewriter &rewriter) const {
    Value trueTensor = selectOp.getOnTrue();
    Value falseTensor = selectOp.getOnFalse();

    // pred may come from a broadcast_in_dim wrapping a compare, or directly
    // from a compare (no broadcast).
    auto broadcast =
        selectOp.getPred().getDefiningOp<stablehlo::BroadcastInDimOp>();
    stablehlo::CompareOp compare;
    if (broadcast) {
      compare = broadcast.getOperand().getDefiningOp<stablehlo::CompareOp>();
    } else {
      compare = selectOp.getPred().getDefiningOp<stablehlo::CompareOp>();
    }
    if (!compare)
      return failure();

    Value cmpLHS = compare.getLhs();
    Value cmpRHS = compare.getRhs();
    stablehlo::ComparisonDirection direction = compare.getComparisonDirection();

    stablehlo::IotaOp iota;
    int64_t K = 0;

    // Try to match lhs as an iota expression, rhs as a constant.
    // Handles: iota vs const, and add(iota, offset) vs const.
    // Returns true and sets iota/K on success.
    auto tryMatchIotaVsConst = [&](Value lhs, Value rhs) -> bool {
      // Form A: direct iota vs splat constant
      if (auto iotaOp = lhs.getDefiningOp<stablehlo::IotaOp>()) {
        int64_t rhsConst = 0;
        if (extractSplatInt(rhs, rhsConst)) {
          iota = iotaOp;
          K = rhsConst;
          return true;
        }
      }
      // Form B: add(iota, offset) vs splat constant
      // iota + offset direction rhsConst  ⟺  iota direction rhsConst - offset
      if (auto add = lhs.getDefiningOp<stablehlo::AddOp>()) {
        for (int side = 0; side < 2; side++) {
          Value maybeIota = side == 0 ? add.getLhs() : add.getRhs();
          Value maybeConst = side == 0 ? add.getRhs() : add.getLhs();
          if (auto iotaOp = maybeIota.getDefiningOp<stablehlo::IotaOp>()) {
            int64_t addOffset = 0, rhsConst = 0;
            if (extractSplatInt(maybeConst, addOffset) &&
                extractSplatInt(rhs, rhsConst)) {
              iota = iotaOp;
              K = rhsConst - addOffset;
              return true;
            }
          }
        }
      }
      return false;
    };

    if (!tryMatchIotaVsConst(cmpLHS, cmpRHS)) {
      // Try commuted: swap lhs/rhs and flip ordering directions
      if (!tryMatchIotaVsConst(cmpRHS, cmpLHS))
        return failure();
      // Flip ordering direction because we swapped operands
      switch (direction) {
      case stablehlo::ComparisonDirection::LT:
        direction = stablehlo::ComparisonDirection::GT;
        break;
      case stablehlo::ComparisonDirection::LE:
        direction = stablehlo::ComparisonDirection::GE;
        break;
      case stablehlo::ComparisonDirection::GT:
        direction = stablehlo::ComparisonDirection::LT;
        break;
      case stablehlo::ComparisonDirection::GE:
        direction = stablehlo::ComparisonDirection::LE;
        break;
      default:
        break;
      }
    }

    // Find which output dimension the iota dimension maps to.
    // With broadcast: look up through getBroadcastDimensions().
    // Without broadcast: the iota dim is already the output dim.
    const int64_t iotaDim = iota.getIotaDimension();
    int64_t outputDim = iotaDim;
    if (broadcast) {
      auto bcastDims = broadcast.getBroadcastDimensions();
      outputDim = -1;
      for (auto [inDim, outDim] : llvm::enumerate(bcastDims)) {
        if ((int64_t)inDim == iotaDim) {
          outputDim = outDim;
          break;
        }
      }
      if (outputDim == -1)
        return failure();
    }

    auto outputShape = selectOp.getType().getShape();
    const int64_t endValue = outputShape[outputDim];

    if (K < 0 || K > endValue)
      return failure();

    struct slice_data {
      Value tensor;
      int64_t count;
    };
    SmallVector<slice_data, 3> slices;
    switch (direction) {
    case stablehlo::ComparisonDirection::LT:
      slices.push_back({trueTensor, K});
      slices.push_back({falseTensor, endValue - K});
      break;
    case stablehlo::ComparisonDirection::LE:
      slices.push_back({trueTensor, K + 1});
      slices.push_back({falseTensor, endValue - K - 1});
      break;
    case stablehlo::ComparisonDirection::GT:
      slices.push_back({falseTensor, K + 1});
      slices.push_back({trueTensor, endValue - K - 1});
      break;
    case stablehlo::ComparisonDirection::GE:
      slices.push_back({falseTensor, K});
      slices.push_back({trueTensor, endValue - K});
      break;
    case stablehlo::ComparisonDirection::EQ:
      slices.push_back({falseTensor, K});
      slices.push_back({trueTensor, 1});
      slices.push_back({falseTensor, endValue - K - 1});
      break;
    case stablehlo::ComparisonDirection::NE:
      slices.push_back({trueTensor, K});
      slices.push_back({falseTensor, 1});
      slices.push_back({trueTensor, endValue - K - 1});
      break;
    }

    auto validCount = std::count_if(slices.begin(), slices.end(),
                                    [](slice_data d) { return d.count > 0; });
    if (validCount == 1) {
      for (auto &e : slices)
        if (e.count > 0) {
          rewriter.replaceAllOpUsesWith(selectOp, e.tensor);
          return success();
        }
    }

    const auto loc = selectOp.getLoc();
    SmallVector<int64_t> starts(outputShape.size(), 0);
    SmallVector<int64_t> limits(outputShape);
    SmallVector<int64_t> strides(outputShape.size(), 1);
    SmallVector<Value> sliceValues;
    int64_t cursor = 0;
    for (auto &e : slices) {
      if (e.count > 0) {
        starts[outputDim] = cursor;
        limits[outputDim] = cursor + e.count;
        sliceValues.push_back(stablehlo::SliceOp::create(
            rewriter, loc, e.tensor, ArrayRef<int64_t>(starts),
            ArrayRef<int64_t>(limits), ArrayRef<int64_t>(strides)));
        cursor += e.count;
      }
    }

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
        selectOp, ValueRange{sliceValues}, outputDim);
    return success();
  }
};

struct SelectCompIotaConstToDUS final
    : CheckedOpRewritePattern<stablehlo::SelectOp, SelectCompIotaConstToDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp selectOp,
                                    PatternRewriter &rewriter) const {
    auto pred = selectOp.getPred().getDefiningOp<stablehlo::AndOp>();
    if (!pred)
      return failure();
    Value trueTensor = selectOp.getOnTrue();
    Value falseTensor = selectOp.getOnFalse();

    Value lhs[2];
    Value rhs[2];
    stablehlo::ComparisonDirection direction[2];
    for (int i = 0; i < 2; i++) {
      Value cmpLHS = nullptr;
      Value cmpRHS = nullptr;
      stablehlo::ComparisonDirection flag;
      auto compare = pred->getOperand(i);
      if (auto cmp = compare.getDefiningOp<stablehlo::CompareOp>()) {
        cmpLHS = cmp.getLhs();
        cmpRHS = cmp.getRhs();
        flag = cmp.getComparisonDirection();
      } else if (auto notop = compare.getDefiningOp<stablehlo::NotOp>()) {
        if (auto cmp = compare.getDefiningOp<stablehlo::CompareOp>()) {
          cmpLHS = cmp.getLhs();
          cmpRHS = cmp.getRhs();
          flag = negatedComparisonDirection(cmp.getComparisonDirection());
        }
      }

      if (!cmpLHS)
        return failure();
      lhs[i] = cmpLHS;
      rhs[i] = cmpRHS;
      direction[i] = flag;
    }

    if (lhs[0] != lhs[1])
      return failure();

    int dimension = -1;
    int start = 0;

    if (stablehlo::IotaOp iota = lhs[0].getDefiningOp<stablehlo::IotaOp>()) {
      dimension = iota.getIotaDimension();
      start = 0;
    } else if (auto sl = lhs[0].getDefiningOp<stablehlo::SliceOp>()) {
      if (stablehlo::IotaOp iota =
              sl.getOperand().getDefiningOp<stablehlo::IotaOp>()) {
        dimension = iota.getIotaDimension();
        start = sl.getStartIndices()[dimension];
      }
    }

    if (dimension == -1)
      return failure();

    int64_t constants[2];
    for (int i = 0; i < 2; i++) {
      DenseIntElementsAttr constant;
      if (!matchPattern(rhs[i], m_Constant(&constant)))
        return failure();
      if (!constant.isSplat())
        return failure();

      constants[i] =
          constant.getSplatValue<IntegerAttr>().getValue().getSExtValue();
    }

    for (int i = 0; i < 2; i++) {

      auto lb = constants[i] - start;
      auto ub = constants[1 - i] - start;

      auto lb_pred = direction[i];
      bool legalLB = false;
      if (lb_pred == stablehlo::ComparisonDirection::GE) {
        legalLB = true;
      } else if (lb_pred == stablehlo::ComparisonDirection::GT) {
        legalLB = true;
        lb++;
      } else if (lb_pred == stablehlo::ComparisonDirection::NE && lb == 0) {
        legalLB = true;
        lb++;
      }

      if (!legalLB)
        continue;

      auto ub_pred = direction[1 - i];
      if (ub_pred != stablehlo::ComparisonDirection::LT &&
          ub_pred != stablehlo::ComparisonDirection::LE)
        continue;

      if (ub_pred == stablehlo::ComparisonDirection::LE) {
        ub++;
      }

      if (ub > selectOp.getType().getShape()[dimension])
        ub = selectOp.getType().getShape()[dimension];
      if (lb >= ub)
        continue;
      if (lb < 0)
        lb = 0;
      if (ub < 0)
        ub = 0;

      auto ITy = RankedTensorType::get({}, rewriter.getI32Type());

      SmallVector<int64_t> startSlices(selectOp.getType().getShape().size(), 0);
      SmallVector<int64_t> limits =
          llvm::to_vector(selectOp.getType().getShape());
      SmallVector<int64_t> step(selectOp.getType().getShape().size(), 1);
      startSlices[dimension] = lb;
      limits[dimension] = ub;

      auto slicedTrueTensor = stablehlo::SliceOp::create(
          rewriter, selectOp.getLoc(), trueTensor, startSlices, limits, step);

      SmallVector<Value> starts(
          selectOp.getType().getShape().size(),
          stablehlo::ConstantOp::create(rewriter, selectOp.getLoc(), ITy,
                                        cast<ElementsAttr>(makeAttr(ITy, 0))));

      starts[dimension] =
          stablehlo::ConstantOp::create(rewriter, selectOp.getLoc(), ITy,
                                        cast<ElementsAttr>(makeAttr(ITy, lb)));

      rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
          selectOp, falseTensor, slicedTrueTensor, starts);
      return success();
    }

    return failure();
  }
};

struct SelectPadToDUS final
    : CheckedOpRewritePattern<stablehlo::SelectOp, SelectPadToDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp selectOp,
                                    PatternRewriter &rewriter) const {
    auto pad = selectOp.getPred().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();
    Value trueTensor = selectOp.getOnTrue();
    Value falseTensor = selectOp.getOnFalse();

    SplatElementsAttr operand, padded;
    if (!matchPattern(pad.getOperand(), m_Constant(&operand))) {
      return failure();
    }
    if (!matchPattern(pad.getPaddingValue(), m_Constant(&padded))) {
      return failure();
    }

    bool operandV = !operand.getSplatValue<IntegerAttr>().getValue().isZero();
    bool paddedV = !padded.getSplatValue<IntegerAttr>().getValue().isZero();
    if (operandV == paddedV)
      return failure();

    for (auto pv : pad.getInteriorPadding()) {
      if (pv != 0)
        return failure();
    }

    SmallVector<int64_t> startSlices = llvm::to_vector(pad.getEdgePaddingLow());
    SmallVector<int64_t> limits =
        llvm::to_vector(selectOp.getType().getShape());
    for (int i = 0; i < selectOp.getType().getShape().size(); i++) {
      limits[i] -= pad.getEdgePaddingHigh()[i];
    }
    SmallVector<int64_t> step(selectOp.getType().getShape().size(), 1);

    Value dusOperand = stablehlo::SliceOp::create(
        rewriter, selectOp.getLoc(), operandV ? trueTensor : falseTensor,
        startSlices, limits, step);

    auto ITy = RankedTensorType::get({}, rewriter.getI32Type());
    SmallVector<Value> starts;
    for (int i = 0; i < selectOp.getType().getShape().size(); i++) {
      starts.push_back(stablehlo::ConstantOp::create(
          rewriter, selectOp.getLoc(), ITy,
          cast<ElementsAttr>(makeAttr(ITy, pad.getEdgePaddingLow()[i]))));
    }

    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
        selectOp, operandV ? falseTensor : trueTensor, dusOperand, starts);
    return success();
  }
};

struct SelectSelectSameCond final
    : CheckedOpRewritePattern<stablehlo::SelectOp, SelectSelectSameCond> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    Value cond = op.getPred();

    // Case 1: false branch is another select with the same condition
    //   select(cond, c, select(cond, a, b)) -> select(cond, c, b)
    if (auto inner = op.getOnFalse().getDefiningOp<stablehlo::SelectOp>()) {
      if (inner.getPred() == cond) {
        rewriter.modifyOpInPlace(
            op, [&]() { op.getOnFalseMutable().assign(inner.getOnFalse()); });
        return success();
      }
    }

    // Case 2: true branch is another select with the same condition
    //   select(cond, select(cond, a, b), c) -> select(cond, a, c)
    if (auto inner = op.getOnTrue().getDefiningOp<stablehlo::SelectOp>()) {
      if (inner.getPred() == cond) {
        rewriter.modifyOpInPlace(
            op, [&]() { op.getOnTrueMutable().assign(inner.getOnTrue()); });
        return success();
      }
    }

    return failure();
  }
};

struct SelectSelectNegCond final
    : CheckedOpRewritePattern<stablehlo::SelectOp, SelectSelectNegCond> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  // Returns the operand if v == not(operand), else null Value.
  static Value getNegated(Value v) {
    if (auto notOp = v.getDefiningOp<stablehlo::NotOp>())
      return notOp.getOperand();
    return {};
  }

  // If pred1 == and(c1, c2) and pred2 == and(c1, not(c2)) (operands in any
  // order), returns c1. Otherwise returns null Value.
  static Value getAndNegComplement(Value pred1, Value pred2) {
    auto and1 = pred1.getDefiningOp<stablehlo::AndOp>();
    auto and2 = pred2.getDefiningOp<stablehlo::AndOp>();
    if (!and1 || !and2)
      return {};
    Value and1Ops[2] = {and1.getLhs(), and1.getRhs()};
    Value and2Ops[2] = {and2.getLhs(), and2.getRhs()};
    for (Value a1 : and1Ops)
      for (Value b1 : and1Ops) {
        if (a1 == b1)
          continue;
        for (Value a2 : and2Ops)
          for (Value b2 : and2Ops) {
            if (a2 == b2)
              continue;
            if (a1 == a2 && getNegated(b2) == b1)
              return a1;
          }
      }
    return {};
  }

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    Value cond = op.getPred();

    // Case 1: select(cond, c, select(not(cond), a, b)) -> select(cond, c, a)
    if (auto inner = op.getOnFalse().getDefiningOp<stablehlo::SelectOp>()) {
      if (getNegated(inner.getPred()) == cond) {
        rewriter.modifyOpInPlace(
            op, [&]() { op.getOnFalseMutable().assign(inner.getOnTrue()); });
        return success();
      }
      // Case 3: select(and(c1,c2), c, select(and(c1,not(c2)), a, b))
      //      -> select(and(c1,c2), c, select(c1, a, b))
      if (Value c1 = getAndNegComplement(cond, inner.getPred())) {
        auto newInner = stablehlo::SelectOp::create(
            rewriter, op.getLoc(), c1, inner.getOnTrue(), inner.getOnFalse());
        rewriter.modifyOpInPlace(
            op, [&]() { op.getOnFalseMutable().assign(newInner); });
        return success();
      }
    }

    // Case 2: select(cond, select(not(cond), a, b), c) -> select(cond, b, c)
    if (auto inner = op.getOnTrue().getDefiningOp<stablehlo::SelectOp>()) {
      if (getNegated(inner.getPred()) == cond) {
        rewriter.modifyOpInPlace(
            op, [&]() { op.getOnTrueMutable().assign(inner.getOnFalse()); });
        return success();
      }
      // Case 4: select(and(c1,c2), select(and(c1,not(c2)), a, b), c)
      //      -> select(and(c1,c2), b, c)
      // When outer is true: c1=T,c2=T => inner pred=and(T,F)=F => b
      if (getAndNegComplement(cond, inner.getPred())) {
        rewriter.modifyOpInPlace(
            op, [&]() { op.getOnTrueMutable().assign(inner.getOnFalse()); });
        return success();
      }
    }

    return failure();
  }
};

// A select whose operands are themselves i1 and whose true-or-false branch is
// an all-true/all-false constant is pure boolean logic. Rewriting it as such
// replaces a select with an and/or, which -- unlike the select -- is
// associative and commutative, so the downstream reassociation, CSE and
// pad-pushing patterns (reshuffle_ands_compares, and_simplify, cse_compare,
// and_pad_pad) can all get at it.
//
//   select(p, x, false) -> and(p, x)
//   select(p, true,  x) -> or(p, x)
//   select(p, false, x) -> and(not(p), x)
//   select(p, x,  true) -> or(not(p), x)
//
// The first two forms are strictly op-count non-increasing. The latter two
// introduce a not, which is normally absorbed straight back into whatever
// produced the predicate (not_compare folds not(compare) into the flipped
// compare, involution_not_simplify cancels a double negation).
struct SelectToLogical final
    : CheckedOpRewritePattern<stablehlo::SelectOp, SelectToLogical> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  // Matches an all-true or all-false constant.
  static bool matchBoolSplat(Value v, bool &out) {
    DenseIntElementsAttr c;
    if (!matchPattern(v, m_Constant(&c)) || !c.isSplat())
      return false;
    if (!c.getElementType().isInteger(1))
      return false;
    out = c.getSplatValue<bool>();
    return true;
  }

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    auto resTy = dyn_cast<RankedTensorType>(op.getType());
    if (!resTy || !resTy.getElementType().isInteger(1))
      return failure();

    // stablehlo::SelectOp permits a rank-0 predicate against ranked operands;
    // and/or have no such broadcasting, so the predicate has to already be
    // shaped like the result for the rewrite to be type-correct.
    Value pred = op.getPred();
    if (pred.getType() != op.getType())
      return failure();

    Value onTrue = op.getOnTrue(), onFalse = op.getOnFalse();
    bool cv;

    if (matchBoolSplat(onFalse, cv)) {
      if (!cv) { // select(p, x, false) -> and(p, x)
        rewriter.replaceOpWithNewOp<stablehlo::AndOp>(op, pred, onTrue);
        return success();
      }
      // select(p, x, true) -> or(not(p), x)
      auto notPred = stablehlo::NotOp::create(rewriter, op.getLoc(), pred);
      rewriter.replaceOpWithNewOp<stablehlo::OrOp>(op, notPred, onTrue);
      return success();
    }

    if (matchBoolSplat(onTrue, cv)) {
      if (cv) { // select(p, true, x) -> or(p, x)
        rewriter.replaceOpWithNewOp<stablehlo::OrOp>(op, pred, onFalse);
        return success();
      }
      // select(p, false, x) -> and(not(p), x)
      auto notPred = stablehlo::NotOp::create(rewriter, op.getLoc(), pred);
      rewriter.replaceOpWithNewOp<stablehlo::AndOp>(op, notPred, onFalse);
      return success();
    }

    return failure();
  }
};

struct AndPadPad final : CheckedOpRewritePattern<stablehlo::AndOp, AndPadPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AndOp andOp,
                                    PatternRewriter &rewriter) const {
    auto padLHS = andOp.getLhs().getDefiningOp<stablehlo::PadOp>();
    if (!padLHS)
      return failure();
    auto padRHS = andOp.getRhs().getDefiningOp<stablehlo::PadOp>();
    if (!padRHS)
      return failure();

    SplatElementsAttr operandLHS, paddedLHS;
    if (!matchPattern(padLHS.getOperand(), m_Constant(&operandLHS))) {
      return failure();
    }
    if (!matchPattern(padLHS.getPaddingValue(), m_Constant(&paddedLHS))) {
      return failure();
    }

    bool operandLHSV =
        !operandLHS.getSplatValue<IntegerAttr>().getValue().isZero();
    bool paddedLHSV =
        !paddedLHS.getSplatValue<IntegerAttr>().getValue().isZero();
    if (operandLHSV == paddedLHSV)
      return failure();

    for (auto pv : padLHS.getInteriorPadding()) {
      if (pv != 0)
        return failure();
    }

    SplatElementsAttr operandRHS, paddedRHS;
    if (!matchPattern(padRHS.getOperand(), m_Constant(&operandRHS))) {
      return failure();
    }
    if (!matchPattern(padRHS.getPaddingValue(), m_Constant(&paddedRHS))) {
      return failure();
    }

    bool operandRHSV =
        !operandRHS.getSplatValue<IntegerAttr>().getValue().isZero();
    bool paddedRHSV =
        !paddedRHS.getSplatValue<IntegerAttr>().getValue().isZero();
    if (operandRHSV == paddedRHSV)
      return failure();

    for (auto pv : padLHS.getInteriorPadding()) {
      if (pv != 0)
        return failure();
    }

    int64_t idx = -1;
    for (int i = 0; i < padLHS.getInteriorPadding().size(); i++) {
      if (padLHS.getEdgePaddingLow()[i] != 0 ||
          padRHS.getEdgePaddingLow()[i] != 0 ||
          padLHS.getEdgePaddingHigh()[i] != 0 ||
          padRHS.getEdgePaddingHigh()[i] != 0) {
        if (idx == -1) {
          idx = i;
        } else {
          return failure();
        }
      }
    }
    if (idx == -1)
      return failure();

    std::set<int64_t> boundaries;
    boundaries.insert(0);
    boundaries.insert(padLHS.getType().getShape()[idx]);
    boundaries.insert(padLHS.getEdgePaddingLow()[idx]);
    boundaries.insert(padRHS.getEdgePaddingLow()[idx]);

    boundaries.insert(padLHS.getType().getShape()[idx] -
                      padLHS.getEdgePaddingHigh()[idx]);
    boundaries.insert(padRHS.getType().getShape()[idx] -
                      padRHS.getEdgePaddingHigh()[idx]);
    SmallVector<int64_t> boundariesV = llvm::to_vector(boundaries);

    SmallVector<bool> toConcat;
    SmallVector<int64_t> start;
    for (int i = 0; i < boundariesV.size() - 1; i++) {
      auto lhsV = ((boundariesV[i] >= padLHS.getEdgePaddingLow()[idx]) &&
                   (boundariesV[i] < (padLHS.getType().getShape()[idx] -
                                      padLHS.getEdgePaddingHigh()[idx])))
                      ? operandLHSV
                      : paddedLHSV;
      auto rhsV = ((boundariesV[i] >= padRHS.getEdgePaddingLow()[idx]) &&
                   (boundariesV[i] < (padRHS.getType().getShape()[idx] -
                                      padRHS.getEdgePaddingHigh()[idx])))
                      ? operandRHSV
                      : paddedRHSV;
      auto newV = lhsV & rhsV;
      if (start.size() == 0 || newV != toConcat.back()) {
        toConcat.push_back(newV);
        start.push_back(boundariesV[i]);
      }
    }

    start.push_back(padLHS.getType().getShape()[idx]);

    if (toConcat.size() == 3) {
      auto shape = llvm::to_vector(andOp.getType().getShape());
      shape[idx] = start[2] - start[1];
      auto RT = RankedTensorType::get(shape, andOp.getType().getElementType());

      auto newInner = stablehlo::ConstantOp::create(
          rewriter, andOp.getLoc(), RT,
          cast<ElementsAttr>(makeAttr(RT, toConcat[1])));

      auto RT0D = RankedTensorType::get({}, andOp.getType().getElementType());
      auto newOuter = stablehlo::ConstantOp::create(
          rewriter, andOp.getLoc(), RT0D,
          cast<ElementsAttr>(makeAttr(RT0D, toConcat[0])));

      SmallVector<int64_t> low(padLHS.getInteriorPadding().size(), 0);
      low[idx] = start[1];

      SmallVector<int64_t> high(padLHS.getInteriorPadding().size(), 0);
      high[idx] = start[3] - start[2];

      rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
          andOp, newInner, newOuter, low, high, padLHS.getInteriorPadding());
      return success();
    }

    SmallVector<Value> toConcatV;
    for (int i = 0; i < start.size() - 1; i++) {
      auto shape = llvm::to_vector(andOp.getType().getShape());
      shape[idx] = start[i + 1] - start[i];
      auto RT = RankedTensorType::get(shape, andOp.getType().getElementType());
      toConcatV.push_back(stablehlo::ConstantOp::create(
          rewriter, andOp.getLoc(), RT,
          cast<ElementsAttr>(makeAttr(RT, toConcat[i]))));
    }

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(andOp, toConcatV,
                                                          idx);
    return success();
  }
};

struct SelectOpUsedWithinIf final
    : CheckedOpRewritePattern<stablehlo::SelectOp, SelectOpUsedWithinIf> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    Value pred = op.getPred();
    Value result = op.getResult();

    if (cast<TensorType>(pred.getType()).getRank() != 0)
      return failure();

    auto block = op->getBlock();

    bool anyModified = false;

    rewriter.replaceUsesWithIf(result, op.getOnTrue(), [&](auto &use) {
      Operation *user = use.getOwner();
      if (user->getBlock() == block)
        return false;

      Operation *p = user->getParentOp();
      while (p && p != op) {
        if (auto ifOp = dyn_cast<stablehlo::IfOp>(p)) {
          if (ifOp.getPred() == pred &&
              ifOp.getTrueBranch().isAncestor(user->getParentRegion())) {
            anyModified = true;
            return true;
          }
        }
        p = p->getParentOp();
      }
      return false;
    });

    rewriter.replaceUsesWithIf(result, op.getOnFalse(), [&](auto &use) {
      Operation *user = use.getOwner();
      if (user->getBlock() == block)
        return false;

      Operation *p = user->getParentOp();
      while (p && p != op) {
        if (auto ifOp = dyn_cast<stablehlo::IfOp>(p)) {
          if (ifOp.getPred() == pred &&
              ifOp.getFalseBranch().isAncestor(user->getParentRegion())) {
            anyModified = true;
            return true;
          }
        }
        p = p->getParentOp();
      }
      return false;
    });

    return success(anyModified);
  }
};

struct SelectOpCanon final
    : CheckedOpRewritePattern<stablehlo::SelectOp, SelectOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  size_t max_constant_expansion;
  SelectOpCanon(size_t max_constant_expansion, MLIRContext *context,
                PatternBenefit benefit = 1,
                ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    RankedTensorType type = op.getType();

    Value trueVal = op.getOnTrue();
    Value falseVal = op.getOnFalse();

    // Eliminate select with two identical outcomes.
    if (trueVal == falseVal) {
      rewriter.replaceOp(op, trueVal);
      return success();
    }

    // Simplify when the condition is a constant.
    Value pred = op.getPred();
    ElementsAttr cond;
    if (!matchPattern(pred, m_Constant(&cond)))
      return failure();

    // Handle splat predicate and select either `trueVal` or `falseVal`.
    if (cond.isSplat()) {
      rewriter.replaceOp(op, cond.getSplatValue<bool>() ? trueVal : falseVal);
      return success();
    }

    // Handle elementwise selection when both outcomes are also constants. This
    // will create a new, likely non-splat constant.
    if (cond.getNumElements() > max_constant_expansion)
      return failure();

    ElementsAttr trueAttr;
    if (!matchPattern(trueVal, m_Constant(&trueAttr)))
      return failure();

    ElementsAttr falseAttr;
    if (!matchPattern(falseVal, m_Constant(&falseAttr)))
      return failure();

    SmallVector<Attribute> newValues;
    newValues.reserve(cond.getNumElements());
    for (auto [condElem, trueElem, falseElem] : llvm::zip_equal(
             cond.getValues<bool>(), trueAttr.getValues<Attribute>(),
             falseAttr.getValues<Attribute>())) {
      newValues.push_back(condElem ? trueElem : falseElem);
    }

    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, DenseElementsAttr::get(type, newValues));
    return success();
  }
};

struct BroadcastInDimOpCanon final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                              BroadcastInDimOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    RankedTensorType type = op.getType();

    TypedValue<RankedTensorType> operand = op.getOperand();
    RankedTensorType operandTy = operand.getType();

    // Fold when broadcast is a noop.
    auto dims = op.getBroadcastDimensions();
    bool isDimsIota = isIotaRange(dims);
    if (type == operandTy && isDimsIota) {
      rewriter.replaceOp(op, operand);
      return success();
    }

    // Handle splat broadcasts.
    if (SplatElementsAttr cstAttr;
        matchPattern(operand, m_Constant(&cstAttr))) {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, SplatElementsAttr::get(op.getType(),
                                     cstAttr.getSplatValue<Attribute>()));
      return success();
    }

    if (operandTy.hasStaticShape() && type.hasStaticShape() &&
        type.getNumElements() == operandTy.getNumElements()) {
      // BroadcastInDim equivalent to reshape.
      if (isDimsIota) {
        rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, type, operand);
        return success();
      }
      // BroadcastInDim equivalent to transpose.
      if (type.getRank() == operandTy.getRank()) {
        rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(
            op, type, operand, getInversePermutation(dims));
        return success();
      }
    }

    // Eliminate redundant nested BroadcastInDim.
    if (auto definingOp =
            operand.getDefiningOp<stablehlo::BroadcastInDimOp>()) {
      auto newIndices = llvm::to_vector(
          llvm::map_range(definingOp.getBroadcastDimensions(),
                          [&dims](int64_t dim) { return dims[dim]; }));
      rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
          op, type, definingOp.getOperand(), newIndices);
      return success();
    }

    // Attempt to rewrite which dims are broadcast or not
    {
      bool changed = false;
      SmallVector<int64_t> dims = llvm::to_vector(op.getBroadcastDimensions());
      for (size_t i = 0; i < dims.size(); i++) {
        if (operandTy.getShape()[i] != 1)
          continue;
        if (type.getShape()[dims[i]] == 1)
          continue;
        bool found = false;
        size_t newidx = 0;
        for (size_t j = 0; j < type.getShape().size(); j++) {
          if (type.getShape()[j] != 1)
            continue;
          if (llvm::is_contained(dims, j))
            continue;
          found = true;
          newidx = j;
          break;
        }
        if (found) {
          changed = true;
          dims[i] = newidx;
        }
      }
      if (changed) {
        rewriter.modifyOpInPlace(op,
                                 [&]() { op.setBroadcastDimensions(dims); });
        return success();
      }
    }
    return failure();
  }
};

struct TransposeBroadcastInDimToBroadcastInDim final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                              TransposeBroadcastInDimToBroadcastInDim> {
  using CheckedOpRewritePattern<
      stablehlo::BroadcastInDimOp,
      TransposeBroadcastInDimToBroadcastInDim>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    auto transposeOp = op.getOperand().getDefiningOp<stablehlo::TransposeOp>();
    if (!transposeOp)
      return failure();

    auto broadcastDims = op.getBroadcastDimensions();
    auto permutation = transposeOp.getPermutation();

    // For each input dimension, find where it maps in final output
    SmallVector<int64_t> newBroadcastDims(
        transposeOp.getOperand().getType().getRank(), -1);
    for (auto [idx, oldDim] : llvm::enumerate(broadcastDims)) {
      newBroadcastDims[permutation[idx]] = oldDim;
    }

    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), transposeOp.getOperand(),
        rewriter.getDenseI64ArrayAttr(newBroadcastDims));

    return success();
  }
};

struct BroadcastInDimTransposeToBroadcastInDim final
    : CheckedOpRewritePattern<stablehlo::TransposeOp,
                              BroadcastInDimTransposeToBroadcastInDim> {
  using CheckedOpRewritePattern<
      stablehlo::TransposeOp,
      BroadcastInDimTransposeToBroadcastInDim>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto broadcastOp =
        op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!broadcastOp)
      return failure();

    bool singleUser = broadcastOp->getResult(0).hasOneUse();
    auto broadcastDims = broadcastOp.getBroadcastDimensions();
    auto permutation = op.getPermutation();

    // Compute the inverse permutation
    SmallVector<int64_t> inversePermutation(permutation.size());
    for (auto [idx, perm] : llvm::enumerate(permutation)) {
      inversePermutation[perm] = idx;
    }

    // Adjust the broadcast dimensions using the inverse permutation
    SmallVector<int64_t> newBroadcastDims;
    for (auto oldDim : broadcastDims) {
      newBroadcastDims.push_back(inversePermutation[oldDim]);
    }

    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), broadcastOp.getOperand(),
        rewriter.getDenseI64ArrayAttr(newBroadcastDims));
    if (singleUser)
      rewriter.eraseOp(broadcastOp);

    return success();
  }
};

struct ConcatenateOpCanon final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatenateOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  size_t max_constant_expansion;
  ConcatenateOpCanon(size_t max_constant_expansion, MLIRContext *context,
                     PatternBenefit benefit = 1,
                     ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        max_constant_expansion(max_constant_expansion) {}

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    RankedTensorType type = op.getType();
    if (!type.hasStaticShape())
      return failure();

    size_t numElems = type.getNumElements();
    if (numElems > max_constant_expansion)
      return failure();

    // Fold concatenate when all inputs are constants.
    OperandRange inputs = op.getInputs();
    SmallVector<DenseElementsAttr> constants(inputs.size());
    for (auto [input, constant] : llvm::zip_equal(inputs, constants)) {
      if (!matchPattern(input, m_Constant(&constant)))
        return failure();
    }

    uint64_t dim = op.getDimension();
    ArrayRef<int64_t> shape = type.getShape();
    int64_t topSize = std::accumulate(shape.begin(), shape.begin() + dim,
                                      int64_t{1}, std::multiplies<>{});

    SmallVector<Attribute> newElems;
    newElems.reserve(numElems);

    for (int64_t i = 0; i != topSize; ++i) {
      for (ElementsAttr attr : constants) {
        size_t bottomSize = attr.getNumElements() / topSize;
        auto begin = attr.value_begin<Attribute>() + (i * bottomSize);
        newElems.append(begin, begin + bottomSize);
      }
    }

    assert(newElems.size() == numElems);
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, DenseElementsAttr::get(op.getType(), newElems));
    return success();
  }
};

struct ConvertOpCanon final
    : CheckedOpRewritePattern<stablehlo::ConvertOp, ConvertOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    // Check if this convert is a noop.
    if (op.getOperand().getType() != op.getType())
      return failure();

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

struct DivideSqrtToMultiplyRsqrt final
    : CheckedOpRewritePattern<stablehlo::DivOp, DivideSqrtToMultiplyRsqrt> {
  using CheckedOpRewritePattern<
      stablehlo::DivOp, DivideSqrtToMultiplyRsqrt>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DivOp op,
                                    PatternRewriter &rewriter) const {
    auto rhsOp = op.getRhs().getDefiningOp<stablehlo::SqrtOp>();
    if ((!rhsOp) || !rhsOp->hasOneUse())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
        op, op.getLhs(),
        stablehlo::RsqrtOp::create(rewriter, op.getLoc(), rhsOp.getOperand()));
    return success();
  }
};

/// Does the same as PatternRewriter::replaceOpWithNewOp, but with a twist.
///
/// Sometimes, we want to replace an op with a new op and simultaneously refine
/// the result type from a dynamically-shaped type to a statically-shaped type.
/// (Search for usages of this function for examples).
//
/// Oftentimes, this works just fine because HLO is designed to accommodate
/// this kind of type refinements. But sometimes, this doesn't work - when
/// the op is used outside of the HLO dialect (e.g. in func.return). In these
/// cases, we insert a tensor.cast to smooth things out.
template <typename OpTy, typename... Args>
static OpTy refineOpWithNewOp(PatternRewriter &rewriter, Operation *op,
                              Args &&...args) {
  auto newOp =
      OpTy::create(rewriter, op->getLoc(), std::forward<Args>(args)...);

  llvm::SmallVector<Value> replacementResults;
  assert(op->getNumResults() == newOp->getNumResults() &&
         "replacement op doesn't match results of original op");
  for (auto [opResult, newOpResult] :
       llvm::zip(op->getResults(), newOp->getResults())) {
    Value replacementResult = newOpResult;
    if (llvm::any_of(opResult.getUsers(), [&](Operation *user) {
          return user->getDialect() != op->getDialect();
        }))
      replacementResult = mlir::tensor::CastOp::create(
          rewriter, op->getLoc(), opResult.getType(), newOpResult);
    replacementResults.push_back(replacementResult);
  }

  rewriter.replaceOp(op, replacementResults);
  return newOp;
}

/// If a DynamicBroadCastInDimOp is not actually dynamic, use an ordinary
/// BroadcastInDimOp.
struct DynamicBroadcastInDimOpNotActuallyDynamic final
    : CheckedOpRewritePattern<stablehlo::DynamicBroadcastInDimOp,
                              DynamicBroadcastInDimOpNotActuallyDynamic> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicBroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    RankedTensorType operandType = op.getOperand().getType();
    if (!operandType.hasStaticShape())
      return rewriter.notifyMatchFailure(op, "requires operand static shape");

    RankedTensorType type = op.getType();
    // output has static shape, replace with broadcast_in_dim
    if (type.hasStaticShape()) {
      rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
          op, type, op.getOperand(), op.getBroadcastDimensionsAttr());
      return success();
    }

    // output_dimensions are constant, set output shape with output_dimensions,
    // then replace with broadcast_in_dim
    if (llvm::SmallVector<int64_t> shape;
        succeeded(hlo::matchInts(op.getOutputDimensions(), shape))) {
      refineOpWithNewOp<stablehlo::BroadcastInDimOp>(
          rewriter, op, RankedTensorType::get(shape, type.getElementType()),
          op.getOperand(), op.getBroadcastDimensionsAttr());
      return success();
    }
    return rewriter.notifyMatchFailure(
        op, "requires output static shape or constant broadcast dimensions");
  }
};

struct ChainedDynamicBroadcastInDimCanonicalization final
    : CheckedOpRewritePattern<stablehlo::DynamicBroadcastInDimOp,
                              ChainedDynamicBroadcastInDimCanonicalization> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicBroadcastInDimOp bcast,
                                    PatternRewriter &rewriter) const {
    auto precedingBcast =
        bcast.getOperand().getDefiningOp<stablehlo::DynamicBroadcastInDimOp>();
    if (!precedingBcast)
      return failure();

    // Compose broadcast dimensions.
    SmallVector<int64_t> composition;
    for (int64_t precedingDim : precedingBcast.getBroadcastDimensions()) {
      composition.push_back(bcast.getBroadcastDimensions()[precedingDim]);
    }
    auto composedBcastDims = rewriter.getDenseI64ArrayAttr(composition);

    rewriter.replaceOpWithNewOp<stablehlo::DynamicBroadcastInDimOp>(
        bcast, bcast.getType(), precedingBcast.getOperand(),
        bcast.getOutputDimensions(), composedBcastDims);
    return success();
  }
};

// If all dimensions are known to be nonexpanding from the attribute, replace
// the dynamic broadcast with a cast.
struct DynamicBroadcastInDimAllDimsNonExpanding final
    : CheckedOpRewritePattern<stablehlo::DynamicBroadcastInDimOp,
                              DynamicBroadcastInDimAllDimsNonExpanding> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicBroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    RankedTensorType type = op.getType();

    if (!op.getKnownNonexpandingDimensions() ||
        static_cast<int64_t>(op.getKnownNonexpandingDimensions()->size()) !=
            type.getRank()) {
      return rewriter.notifyMatchFailure(
          op, "known_nonexpanding_dimensions don't cover all output dims");
    }

    auto cast = rewriter.createOrFold<tensor::CastOp>(op.getLoc(), type,
                                                      op.getOperand());
    rewriter.replaceOp(op, cast);
    return success();
  }
};

struct NoopReduceOpCanon final
    : CheckedOpRewritePattern<stablehlo::ReduceOp, NoopReduceOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    // No dimensions to reduce.
    if (op.getDimensions().empty()) {
      rewriter.replaceOp(op, op.getInputs());
      return success();
    }

    // If all returned values in the ReduceOp region exists outside the
    // region, replace the ReduceOp with those values.
    if (auto retOp = dyn_cast<stablehlo::ReturnOp>(
            op.getBody().front().getTerminator())) {
      Region *retRegion = retOp->getParentRegion();
      if (llvm::any_of(retOp.getResults(), [retRegion](Value result) {
            return result.getParentRegion() == retRegion;
          }))
        return failure();

      SmallVector<Value> vals;
      DenseI64ArrayAttr empty = rewriter.getDenseI64ArrayAttr({});
      for (auto [res, opres] : llvm::zip(retOp.getResults(), op.getResults()))
        vals.push_back(stablehlo::BroadcastInDimOp::create(
            rewriter, op.getLoc(), opres.getType(), res, empty));
      rewriter.replaceOp(op, vals);
      return success();
    }

    return failure();
  }
};

struct EmptyReduceOpCanon final
    : CheckedOpRewritePattern<stablehlo::ReduceOp, EmptyReduceOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    // We require all reduce shapes to be the same, up to the element types, so
    // we can just use the first operand and the first result as
    // representatives.
    auto elemTy = cast<RankedTensorType>(op.getInputs().getType().front());

    if (op.getDimensions().empty()) {
      rewriter.replaceAllOpUsesWith(op, op.getInputs());
      return success();
    }

    if (!llvm::is_contained(elemTy.getShape(), 0))
      return failure();

    Location loc = op.getLoc();
    DenseI64ArrayAttr empty = rewriter.getDenseI64ArrayAttr({});
    if (elemTy.hasStaticShape()) {
      SmallVector<Value> broadcasts(op.getNumResults());
      for (auto [bcast, init, outTy] : llvm::zip_equal(
               broadcasts, op.getInitValues(), op.getResultTypes())) {
        bcast = stablehlo::BroadcastInDimOp::create(rewriter, loc, outTy, init,
                                                    empty);
      }
      rewriter.replaceOp(op, broadcasts);
      return success();
    }

    SmallVector<Value> shapes;
    if (failed(op.reifyReturnTypeShapes(rewriter, op.getOperands(), shapes)))
      return failure();

    SmallVector<Value> broadcasts(op.getNumResults());
    for (auto [bcast, init, shape, outTy] : llvm::zip_equal(
             broadcasts, op.getInitValues(), shapes, op.getResultTypes())) {
      bcast = stablehlo::DynamicBroadcastInDimOp::create(rewriter, loc, outTy,
                                                         init, shape, empty);
    }
    rewriter.replaceOp(op, broadcasts);
    return success();
  }
};

struct DynamicReshapeOpCanon final
    : CheckedOpRewritePattern<stablehlo::DynamicReshapeOp,
                              DynamicReshapeOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // This is a noop when the output type is already a static shape.
    RankedTensorType type = op.getType();
    if (!type.hasStaticShape())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, type,
                                                      op.getOperand());
    return success();
  }
};

struct GetTupleElementOpCanon final
    : CheckedOpRewritePattern<stablehlo::GetTupleElementOp,
                              GetTupleElementOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::GetTupleElementOp op,
                                    PatternRewriter &rewriter) const {
    auto tuple = op.getOperand().getDefiningOp<stablehlo::TupleOp>();
    if (!tuple)
      return failure();

    Value result = tuple.getOperand(op.getIndex());
    rewriter.replaceOp(op, result);
    return success();
  }
};

struct RealOpCanon final
    : CheckedOpRewritePattern<stablehlo::RealOp, RealOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::RealOp op,
                                    PatternRewriter &rewriter) const {
    auto elTy = op.getOperand().getType().getElementType();
    if (!isa<ComplexType>(elTy)) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }

    if (auto complex = op.getOperand().getDefiningOp<stablehlo::ComplexOp>()) {
      rewriter.replaceOp(op, complex.getLhs());
      return success();
    }

    if (guaranteedPurelyImagResult(op.getOperand(), rewriter)) {
      rewriter.replaceOp(
          op, stablehlo::ConstantOp::create(rewriter, op->getLoc(),
                                            makeAttr(op.getType(), 0)));
      return success();
    }

    return failure();
  }
};

struct ImagOpCanon final
    : CheckedOpRewritePattern<stablehlo::ImagOp, ImagOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ImagOp op,
                                    PatternRewriter &rewriter) const {
    if (auto complex = op.getOperand().getDefiningOp<stablehlo::ComplexOp>()) {
      rewriter.replaceOp(op, complex.getRhs());
      return success();
    }

    auto elTy = op.getOperand().getType().getElementType();
    if (!isa<ComplexType>(elTy) ||
        guaranteedPurelyRealResult(op.getOperand(), rewriter)) {
      rewriter.replaceOp(
          op, stablehlo::ConstantOp::create(rewriter, op->getLoc(),
                                            makeAttr(op.getType(), 0)));
      return success();
    }

    return failure();
  }
};

// (conj (complex a, (neg b))) -> (complex a b)
struct ConjComplexNegate final
    : CheckedOpRewritePattern<chlo::ConjOp, ConjComplexNegate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::ConjOp op,
                                    PatternRewriter &rewriter) const {
    auto complex = op.getOperand().getDefiningOp<stablehlo::ComplexOp>();
    if (!complex)
      return failure();

    auto neg = complex.getRhs().getDefiningOp<stablehlo::NegOp>();
    if (!neg)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ComplexOp>(
        op, op.getType(), complex.getLhs(), neg.getOperand());
    return success();
  }
};

// (neg (imag (conj x))) -> (imag x)
struct NegateImagConj final
    : CheckedOpRewritePattern<stablehlo::NegOp, NegateImagConj> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::NegOp op,
                                    PatternRewriter &rewriter) const {
    auto imag = op.getOperand().getDefiningOp<stablehlo::ImagOp>();
    if (!imag)
      return failure();

    auto conj = imag.getOperand().getDefiningOp<chlo::ConjOp>();
    if (!conj)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ImagOp>(op, op.getType(),
                                                   conj.getOperand());
    return success();
  }
};

struct GetDimensionSizeOpCanon final
    : CheckedOpRewritePattern<stablehlo::GetDimensionSizeOp,
                              GetDimensionSizeOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::GetDimensionSizeOp op,
                                    PatternRewriter &rewriter) const {
    // Fold get_dimension_size when the queried dim is statically known.
    RankedTensorType operandTy = op.getOperand().getType();

    int64_t dimSize = operandTy.getDimSize(op.getDimension());
    if (dimSize < 0)
      return failure();

    auto elemTy = cast<IntegerType>(op.getType().getElementType());
    IntegerAttr elemVal = rewriter.getIntegerAttr(elemTy, dimSize);
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, DenseElementsAttr::get(op.getType(), elemVal));
    return success();
  }
};

struct NoopReverse final
    : CheckedOpRewritePattern<stablehlo::ReverseOp, NoopReverse> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReverseOp op,
                                    PatternRewriter &rewriter) const {
    SplatElementsAttr splat;
    if (matchPattern(op.getOperand(), m_Constant(&splat))) {
      rewriter.replaceAllUsesWith(op, op.getOperand());
      return success();
    }

    SmallVector<int64_t> newDimensions;
    auto dimensions = op.getDimensions();
    auto shape = op.getResult().getType().getShape();

    for (auto dim : dimensions) {
      auto size = shape[dim];
      if (size != 1)
        newDimensions.push_back(dim);
    }

    if (auto bcast =
            op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>()) {
      peelBroadcastedDimensions(bcast, newDimensions);
    }

    if (newDimensions.empty()) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }

    if (newDimensions.size() == dimensions.size())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ReverseOp>(op, op.getOperand(),
                                                      newDimensions);
    return success();
  }

private:
  void peelBroadcastedDimensions(stablehlo::BroadcastInDimOp op,
                                 SmallVectorImpl<int64_t> &dims) const {
    DenseMap<int64_t, int64_t> dimMap;
    for (auto [i, dim] : llvm::enumerate(op.getBroadcastDimensions())) {
      dimMap[dim] = i;
    }

    auto opShape = cast<RankedTensorType>(op.getOperand().getType()).getShape();

    auto newEnd = llvm::remove_if(dims, [&](int64_t dim) {
      auto it = dimMap.find(dim);
      if (it != dimMap.end()) {
        return opShape[it->second] == 1; // if 1 then trivially expanded
      }
      return true; // not in broadcast dims so it was expanded
    });
    dims.erase(newEnd, dims.end());
    return;
  }
};

// slice(reverse(x, dims=[...,d,...]), starts=[...,s,...], limits=[...,s+1,...])
// → slice(x, starts=[...,N-1-s,...], limits=[...,N-s,...])
// for each reversed dim d where the slice picks exactly one element.
// Reversed dims that still span multiple elements keep a residual ReverseOp.
struct SliceReverse final
    : CheckedOpRewritePattern<stablehlo::SliceOp, SliceReverse> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    auto reverse = op.getOperand().getDefiningOp<stablehlo::ReverseOp>();
    if (!reverse)
      return failure();

    auto operandType = cast<RankedTensorType>(reverse.getOperand().getType());

    SmallVector<int64_t> starts(op.getStartIndices());
    SmallVector<int64_t> limits(op.getLimitIndices());
    SmallVector<int64_t> strides(op.getStrides());

    SmallVector<int64_t> remainingDims;
    bool changed = false;
    for (int64_t dim : reverse.getDimensions()) {
      // Single element along this dim: output size == 1.
      if (limits[dim] - starts[dim] <= strides[dim]) {
        int64_t N = operandType.getDimSize(dim);
        int64_t newStart = N - 1 - starts[dim];
        starts[dim] = newStart;
        limits[dim] = newStart + 1;
        strides[dim] = 1;
        changed = true;
      } else {
        remainingDims.push_back(dim);
      }
    }

    if (!changed || !remainingDims.empty() &&
                        !llvm::hasSingleElement(reverse.getResult().getUses()))
      return failure();

    // If any reversed dims remain (not eliminated), we must re-apply the
    // reverse on the original operand *before* slicing so that the slice
    // indices for those dims still select the right elements.
    Value src = reverse.getOperand();
    if (!remainingDims.empty())
      src = stablehlo::ReverseOp::create(rewriter, op.getLoc(), src,
                                         remainingDims);
    rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(op, src, starts, limits,
                                                    strides);
    return success();
  }
};

/// Converts gather ops to slice ops in case we have a single set of constant
/// indices.
struct GatherOpCanon final
    : CheckedOpRewritePattern<stablehlo::GatherOp, GatherOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::GatherOp gather,
                                    PatternRewriter &rewriter) const {
    if (tryRewriteGatherWithConstantStartIndices(gather, rewriter)
            .succeeded()) {
      return success();
    }

    if (tryRewriteGatherWithIotaIndexing(gather, rewriter).succeeded()) {
      return success();
    }

    return failure();
  }

  // Recursively extracts components of a grid index expression (like additions
  // of iotas and splat constants) to represent them as a list of
  // IotaLikeTensors.
  static bool extractGridIotas(Value val,
                               SmallVectorImpl<IotaLikeTensor> &result) {
    if (auto iota = detectIotaLikeTensor(val)) {
      result.push_back(*iota);
      return true;
    }

    if (auto addOp = val.getDefiningOp<stablehlo::AddOp>()) {
      SmallVector<IotaLikeTensor> lhsIotas;
      if (!extractGridIotas(addOp.getLhs(), lhsIotas)) {
        return false;
      }
      SmallVector<IotaLikeTensor> rhsIotas;
      if (!extractGridIotas(addOp.getRhs(), rhsIotas)) {
        return false;
      }
      result.append(lhsIotas.begin(), lhsIotas.end());
      result.append(rhsIotas.begin(), rhsIotas.end());
      return true;
    }

    DenseElementsAttr constAttr;
    if (matchPattern(val, m_Constant(&constAttr))) {
      auto tensorType = cast<RankedTensorType>(constAttr.getType());
      auto elemType = tensorType.getElementType();
      if (!isa<IntegerType>(elemType)) {
        return false;
      }
      if (constAttr.isSplat()) {
        TypedAttr zeroScale = IntegerAttr::get(elemType, 0);
        IotaLikeTensor constIota;
        constIota.dimension = 0;
        constIota.scale = zeroScale;
        constIota.start = constAttr.getSplatValue<TypedAttr>();
        constIota.tensorType = tensorType;
        result.push_back(constIota);
        return true;
      }

      int64_t rank = tensorType.getRank();
      auto shape = tensorType.getShape();
      auto strides = computeStrides(shape);
      int64_t numElements = tensorType.getNumElements();

      auto values = constAttr.getValues<APInt>();

      // 1. Find offset (value at index 0,0,...,0)
      APInt offset = values[0];

      // 2. Find scales for each dimension
      SmallVector<APInt> scales;
      scales.reserve(rank);
      for (int64_t d = 0; d < rank; d++) {
        if (shape[d] > 1) {
          int64_t linearIdx = strides[d];
          scales.push_back(values[linearIdx] - offset);
        } else {
          scales.push_back(APInt(elemType.getIntOrFloatBitWidth(), 0));
        }
      }

      // 3. Verify that the linear relationship holds for all elements
      SmallVector<int64_t> multiIndex;
      for (int64_t idx = 0; idx < numElements; idx++) {
        linearToMultiIndex(idx, strides, multiIndex);
        APInt expected = offset;
        for (int64_t d = 0; d < rank; d++) {
          if (multiIndex[d] != 0 && !scales[d].isZero()) {
            expected +=
                scales[d] * APInt(scales[d].getBitWidth(), multiIndex[d]);
          }
        }
        if (values[idx] != expected) {
          return false;
        }
      }

      // 4. If verified, create IotaLikeTensors for offset and each non-zero
      // scale
      for (int64_t d = 0; d < rank; d++) {
        if (!scales[d].isZero()) {
          IotaLikeTensor iota;
          iota.dimension = d;
          iota.scale = IntegerAttr::get(elemType, scales[d]);
          iota.start = IntegerAttr::get(elemType, 0);
          iota.tensorType = tensorType;
          result.push_back(iota);
        }
      }
      IotaLikeTensor offsetIota;
      offsetIota.dimension = 0;
      offsetIota.scale = IntegerAttr::get(elemType, 0);
      offsetIota.start = IntegerAttr::get(elemType, offset);
      offsetIota.tensorType = tensorType;
      result.push_back(offsetIota);
      return true;
    }

    return false;
  }

  // Attempt to rewrite a gather operation with iota-like indexing into a slice.
  // This supports cases where indices are computed from one or more iota
  // operations, representing a multi-dimensional grid layout that indexes into
  // a 1D operand.
  LogicalResult
  tryRewriteGatherWithIotaIndexing(stablehlo::GatherOp op,
                                   PatternRewriter &rewriter) const {
    auto operand = op.getOperand();
    auto operandTy = cast<RankedTensorType>(operand.getType());
    // TODO: check if this optimization is possible for higher dimenional
    //       tensors?
    if (operandTy.getRank() != 1) {
      return failure();
    }

    for (auto size : op.getSliceSizes()) {
      if (size != 1) {
        return failure();
      }
    }

    auto indices = op.getStartIndices();

    // If the index tensor has only 1 element, it represents a single scalar
    // index. Rewrite it directly to a dynamic slice.
    if (indices.getType().getNumElements() == 1) {
      auto scalarIndex =
          stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), indices, {});
      auto dsOp = stablehlo::DynamicSliceOpCreate(rewriter, op.getLoc(),
                                                  operand, {scalarIndex}, {1});
      auto res =
          stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), dsOp,
                                     cast<ShapedType>(op.getType()).getShape());
      rewriter.replaceOp(op, res);
      return success();
    }

    // Unwrap any reshape to find the underlying grid index expression.
    Value grid = indices;
    if (auto reshapeOp = indices.getDefiningOp<stablehlo::ReshapeOp>()) {
      grid = reshapeOp.getOperand();
    }

    // Extract all component iotas and constant offsets from the grid
    // expression.
    SmallVector<IotaLikeTensor> iotas;
    if (!extractGridIotas(grid, iotas)) {
      return failure();
    }

    auto gridTy = cast<RankedTensorType>(grid.getType());
    int64_t gridRank = gridTy.getRank();
    if (gridRank == 0) {
      return failure();
    }

    auto dimNumbers = op.getDimensionNumbers();
    auto indicesTy = cast<RankedTensorType>(indices.getType());
    int64_t indexVectorDim = dimNumbers.getIndexVectorDim();
    if (indexVectorDim < indicesTy.getRank() &&
        indicesTy.getDimSize(indexVectorDim) != 1) {
      return failure();
    }
    if (dimNumbers.getStartIndexMap().size() != 1 ||
        dimNumbers.getStartIndexMap()[0] != 0) {
      return failure();
    }

    // Sum the scales and constant offsets for each grid dimension.
    SmallVector<int64_t> dimScales(gridRank, 0);
    int64_t totalStart = 0;

    for (auto &iota : iotas) {
      int64_t start = cast<IntegerAttr>(iota.start).getValue().getSExtValue();
      int64_t scale = cast<IntegerAttr>(iota.scale).getValue().getSExtValue();
      totalStart += start;
      if (scale != 0) {
        if (iota.dimension >= gridRank)
          return failure();
        dimScales[iota.dimension] += scale;
      }
    }

    // Map each grid dimension with a non-zero scale to its stride.
    struct DimMapping {
      int64_t gridDim;
      int64_t scale;
      int64_t absScale;
    };
    SmallVector<DimMapping> mappings;
    for (int64_t d = 0; d < gridRank; ++d) {
      if (dimScales[d] != 0) {
        mappings.push_back({d, dimScales[d], std::abs(dimScales[d])});
      }
    }

    if (mappings.empty())
      return failure();

    // Sort the mappings by descending absolute scale. We will reshape the 1D
    // operand tensor into a multi-dimensional tensor where each dimension has
    // strides matching these scales.
    llvm::sort(mappings, [](const DimMapping &a, const DimMapping &b) {
      return a.absScale > b.absScale;
    });

    int64_t N = operandTy.getDimSize(0);

    // Factorize the 1D operand size N into a multi-dimensional shape where each
    // stride corresponds to one of the grid dimension scales.
    int64_t currentDivisor = N;
    SmallVector<int64_t> operandShape;
    for (const auto &map : mappings) {
      if (map.absScale == 0 || currentDivisor % map.absScale != 0)
        return failure();
      operandShape.push_back(currentDivisor / map.absScale);
      currentDivisor = map.absScale;
    }
    if (currentDivisor > 1) {
      operandShape.push_back(currentDivisor);
    }

    int64_t opRank = operandShape.size();

    // Compute strides for the reshaped operand.
    SmallVector<int64_t> opStrides(opRank);
    int64_t st = 1;
    for (int64_t k = opRank - 1; k >= 0; --k) {
      opStrides[k] = st;
      st *= operandShape[k];
    }

    // Calculate the minimum total starting offset, adjusting for any negative
    // scales that effectively offset the slice starts in the grid.
    int64_t minTotalOffset = totalStart;
    for (int64_t d = 0; d < gridRank; ++d) {
      if (dimScales[d] < 0) {
        minTotalOffset += (gridTy.getDimSize(d) - 1) * dimScales[d];
      }
    }

    if (minTotalOffset < 0)
      return failure();

    SmallVector<int64_t> sliceSizes(opRank, 1);
    for (size_t m = 0; m < mappings.size(); ++m) {
      int64_t gridDim = mappings[m].gridDim;
      sliceSizes[m] = gridTy.getDimSize(gridDim);
    }

    // Special case: If there is only one active grid dimension mapping, we can
    // slice directly out of the 1D operand without reshaping/broadcasting.
    if (mappings.size() == 1) {
      int64_t stride = mappings[0].scale;
      int64_t count = sliceSizes[0];
      int64_t start = minTotalOffset;

      int64_t limit;
      bool needsReverse = false;
      if (stride > 0) {
        limit = start + count * stride;
      } else {
        needsReverse = true;
        int64_t start_original = start - (count - 1) * stride;
        // For negative stride, the iota accesses indices in descending order:
        //   start, start+stride, start+2*stride, ... (where stride < 0)
        // The smallest index is start + (count-1)*stride, the largest is start.
        // We slice in the forward direction with abs(stride), then reverse.
        limit = start_original + 1; // one past the largest index
        stride = -stride;
      }

      if (limit > operandTy.getDimSize(0)) { // gather clamps indices
        return failure();
      }

      Value result = stablehlo::SliceOpCreate(rewriter, op.getLoc(), operand,
                                              {start}, {limit}, {stride});
      if (needsReverse) {
        result = stablehlo::ReverseOp::create(
            rewriter, op.getLoc(), result, rewriter.getDenseI64ArrayAttr({0}));
      }

      if (result.getType() != op.getType()) {
        result = stablehlo::ReshapeOp::create(rewriter, op.getLoc(),
                                              op.getType(), result);
      }
      rewriter.replaceOp(op, result);
      return success();
    }

    // General case: We have multiple grid dimensions mapped.
    // Decompose the 1D starting offset into coordinates for each dimension of
    // the multi-dimensional reshaped operand.
    SmallVector<int64_t> sliceStarts(opRank, 0);
    int64_t rem = minTotalOffset;
    for (int64_t k = 0; k < opRank; ++k) {
      sliceStarts[k] = rem / opStrides[k];
      rem %= opStrides[k];
    }

    // Check that the sliced region is within bounds along all dimensions.
    for (int64_t k = 0; k < opRank; ++k) {
      if (sliceStarts[k] + sliceSizes[k] > operandShape[k]) {
        return failure();
      }
    }

    // Reshape the 1D operand to the factored multi-dimensional shape.
    Value reshapedOperand = stablehlo::ReshapeOpCreate(rewriter, op.getLoc(),
                                                       operand, operandShape);

    SmallVector<int64_t> sliceLimits(opRank);
    SmallVector<int64_t> sliceStrides(opRank, 1);
    for (int64_t k = 0; k < opRank; ++k) {
      sliceLimits[k] = sliceStarts[k] + sliceSizes[k];
    }

    // Slice the multi-dimensional sub-tensor.
    Value sliced =
        stablehlo::SliceOpCreate(rewriter, op.getLoc(), reshapedOperand,
                                 sliceStarts, sliceLimits, sliceStrides);

    // Reverse any slice dimensions that had a negative scale in the grid index.
    SmallVector<int64_t> reverseDims;
    for (size_t m = 0; m < mappings.size(); ++m) {
      if (mappings[m].scale < 0) {
        reverseDims.push_back(m);
      }
    }
    if (!reverseDims.empty()) {
      sliced = stablehlo::ReverseOp::create(
          rewriter, op.getLoc(), sliced,
          rewriter.getDenseI64ArrayAttr(reverseDims));
    }

    // Drop trailing singleton dimensions if the reshaped operand has higher
    // rank.
    SmallVector<int64_t> mappedShape;
    for (size_t m = 0; m < mappings.size(); ++m) {
      mappedShape.push_back(sliceSizes[m]);
    }
    if (opRank > (int64_t)mappings.size()) {
      sliced = stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), sliced,
                                          mappedShape);
    }

    // Broadcast the sliced result back to the dimensions of the full grid
    // shape.
    SmallVector<int64_t> broadcastDims;
    for (size_t m = 0; m < mappings.size(); ++m) {
      broadcastDims.push_back(mappings[m].gridDim);
    }

    auto gridResultTy =
        RankedTensorType::get(gridTy.getShape(), operandTy.getElementType());
    sliced = stablehlo::BroadcastInDimOp::create(
        rewriter, op.getLoc(), gridResultTy, sliced,
        rewriter.getDenseI64ArrayAttr(broadcastDims));

    // Finally, reshape to the expected gather output type.
    sliced = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), sliced,
        cast<RankedTensorType>(op.getType()).getShape());

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

  LogicalResult
  tryRewriteGatherWithConstantStartIndices(stablehlo::GatherOp gather,
                                           PatternRewriter &rewriter) const {
    DenseIntElementsAttr index;
    if (!matchPattern(gather.getStartIndices(), m_Constant(&index))) {
      return failure();
    }

    stablehlo::GatherDimensionNumbersAttr dnums = gather.getDimensionNumbers();
    if (dnums.getIndexVectorDim() != 0 || index.getType().getRank() > 1)
      return failure();

    // TODO: Remove when the verifier catches this case that is
    // invalid if all previous condition holds.
    if (index.getNumElements() !=
        static_cast<int64_t>(dnums.getStartIndexMap().size())) {
      return failure();
    }

    auto operandType = cast<RankedTensorType>(gather->getOperand(0).getType());
    if (!operandType.hasStaticShape())
      return failure();

    auto sliceEnd = llvm::to_vector(gather.getSliceSizes());
    SmallVector<int64_t> sliceStart(sliceEnd.size(), 0);
    for (auto [mapIndex, value] :
         llvm::zip_equal(dnums.getStartIndexMap(), index.getValues<APInt>())) {
      // Clamp the indices within bounds to faithfully mirror gather semantics.
      int64_t offset =
          std::clamp(value.getSExtValue(), static_cast<int64_t>(0),
                     operandType.getDimSize(mapIndex) - sliceEnd[mapIndex]);
      sliceStart[mapIndex] += offset;
      sliceEnd[mapIndex] += offset;
    }

    SmallVector<int64_t> sliceStride(sliceEnd.size(), 1);
    SmallVector<int64_t> sliceShape(sliceEnd.size());
    for (auto [shapeElem, startElem, endElem] :
         llvm::zip_equal(sliceShape, sliceStart, sliceEnd)) {
      shapeElem = endElem - startElem;
    }

    Type elementType = gather.getType().getElementType();
    auto sliceType = RankedTensorType::get(sliceShape, elementType);
    Value result = stablehlo::SliceOp::create(
        rewriter, gather.getLoc(), sliceType, gather.getOperand(),
        rewriter.getDenseI64ArrayAttr(sliceStart),
        rewriter.getDenseI64ArrayAttr(sliceEnd),
        rewriter.getDenseI64ArrayAttr(sliceStride));

    ArrayRef<int64_t> collapsedSliceDims = dnums.getCollapsedSliceDims();
    if (!collapsedSliceDims.empty()) {
      llvm::SmallVector<int64_t> reshapeShape;
      for (auto [idx, dim] : llvm::enumerate(sliceShape)) {
        if (!llvm::is_contained(collapsedSliceDims, idx))
          reshapeShape.push_back(dim);
      }
      result = stablehlo::ReshapeOpCreate(rewriter, gather.getLoc(), result,
                                          reshapeShape);
    }

    result.setType(gather.getType());
    rewriter.replaceOp(gather, result);
    return success();
  }
};

struct ReshapeOpCanon final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Fold noop reshape.
    if (op.getType() == op.getOperand().getType()) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }

    // Fold reshape of a constant.
    DenseElementsAttr cstAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&cstAttr)))
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, cstAttr.reshape(op.getType()));
    return success();
  }
};

struct MergeConsecutiveReshapes final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, MergeConsecutiveReshapes> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    // Fold noop reshape.
    auto operand = op.getOperand();
    if (op.getType() == operand.getType()) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }

    // Fold reshape(reshape(x)).
    auto reshapeOp = operand.getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshapeOp)
      return rewriter.notifyMatchFailure(
          op, "requires defining op of operand to be Reshape");

    op.setOperand(reshapeOp->getOperand(0));
    return success();
  }
};

struct TransposeIsReshape final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeIsReshape> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto input = op.getOperand();
    auto permutation = op.getPermutation();

    if (isIotaRange(permutation)) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }

    RankedTensorType inputTy = input.getType();
    if (!inputTy.hasStaticShape() || !op.getType().hasStaticShape()) {
      return rewriter.notifyMatchFailure(
          op, "requires input and output to be of a statically-shaped ranked "
              "tensor type");
    }

    if (!OpIsReshapeLike(op)) {
      return rewriter.notifyMatchFailure(op, "memory layout change");
    }

    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(), input);
    return success();
  }
};

struct IfRemoveUnused final
    : CheckedOpRewritePattern<stablehlo::IfOp, IfRemoveUnused> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IfOp op,
                                    PatternRewriter &rewriter) const {

    SmallVector<bool> resultUsed(op->getNumResults(), true);

    bool anyUnused = false;
    for (const auto &it : llvm::enumerate(op->getResults())) {
      bool unused = it.value().use_empty();
      resultUsed[it.index()] = !unused;
      anyUnused |= unused;
    }

    if (!anyUnused)
      return failure();

    SmallVector<Type> newResultTypes;

    Operation *trueTerm = op.getTrueBranch().front().getTerminator();
    Operation *falseTerm = op.getFalseBranch().front().getTerminator();

    unsigned removed = 0;
    for (const auto &it : llvm::enumerate(op->getResults())) {
      bool used = resultUsed[it.index()];
      if (used) {
        newResultTypes.push_back(it.value().getType());
        continue;
      }

      auto i = it.index() - removed;
      rewriter.modifyOpInPlace(trueTerm, [&] { trueTerm->eraseOperand(i); });
      rewriter.modifyOpInPlace(falseTerm, [&] { falseTerm->eraseOperand(i); });
      removed++;
    }

    auto newIf = stablehlo::IfOp::create(rewriter, op.getLoc(), newResultTypes,
                                         op.getPred());
    newIf.getTrueBranch().takeBody(op.getTrueBranch());
    newIf.getFalseBranch().takeBody(op.getFalseBranch());

    removed = 0;
    for (const auto &it : llvm::enumerate(resultUsed)) {
      bool used = it.value();
      if (!used) {
        removed++;
        continue;
      }
      auto res = op.getResult(it.index());
      auto newRes = newIf.getResult(it.index() - removed);
      rewriter.replaceAllUsesWith(res, newRes);
    }

    rewriter.eraseOp(op);

    return success();
  }
};

struct IfPredPropagation final
    : CheckedOpRewritePattern<stablehlo::IfOp, IfPredPropagation> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IfOp op,
                                    PatternRewriter &rewriter) const {
    Value pred = op.getPred();
    bool anyModified = false;

    auto makeBool = [&](bool value) {
      auto i1type = RankedTensorType::get({}, rewriter.getI1Type());
      return stablehlo::ConstantOp::create(
                 rewriter, pred.getLoc(), i1type,
                 SplatElementsAttr::get(i1type, rewriter.getBoolAttr(value)))
          .getResult();
    };

    Region *trueBranch = &op.getTrueBranch();
    Region *falseBranch = &op.getFalseBranch();

    for (OpOperand &use : pred.getUses()) {
      Operation *owner = use.getOwner();
      Region *useRegion = owner->getParentRegion();

      if (trueBranch->isAncestor(useRegion)) {
        anyModified = true;
        rewriter.modifyOpInPlace(owner, [&]() { use.assign(makeBool(true)); });
      } else if (falseBranch->isAncestor(useRegion)) {
        anyModified = true;
        rewriter.modifyOpInPlace(owner, [&]() { use.assign(makeBool(false)); });
      }
    }

    return success(anyModified);
  }
};

struct IfInline final : CheckedOpRewritePattern<stablehlo::IfOp, IfInline> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IfOp op,
                                    PatternRewriter &rewriter) const {

    Attribute predAttr;
    bool predConst = matchPattern(op.getPred(), m_Constant(&predAttr));
    auto iszero = predConst && matchPattern(predAttr, m_Zero());
    auto isone = predConst && matchPattern(predAttr, m_One());

    if (!iszero && !isone)
      return failure();

    auto &reg = isone ? op.getTrueBranch() : op.getFalseBranch();

    if (reg.empty()) {
      rewriter.eraseOp(op);
      return success();
    }
    assert(reg.hasOneBlock());  // stablehlo.if only allows 1 or 0 block in the
    auto *block = &reg.front(); // regions

    auto term = block->getTerminator();
    rewriter.inlineBlockBefore(block, op);
    rewriter.replaceOp(op, term->getOperands());
    rewriter.eraseOp(term);
    return success();
  }
};

// https:github.com/llvm/llvm-project/blob/74d8f3952c4acf6d57948983d7c5b0d0a7763c28/mlir/lib/Dialect/SCF/IR/SCF.cpp#L2313
struct IfToSelect final
    : public CheckedOpRewritePattern<stablehlo::IfOp, IfToSelect> {
  using CheckedOpRewritePattern<stablehlo::IfOp,
                                IfToSelect>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IfOp op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumResults() == 0 || op.getTrueBranch().empty() ||
        op.getFalseBranch().empty())
      return failure();

    auto pred = op.getPred();

    auto trueTerm = op.getTrueBranch().front().getTerminator();
    auto falseTerm = op.getFalseBranch().front().getTerminator();
    auto trueOperands = trueTerm->getOperands();
    auto falseOperands = falseTerm->getOperands();

    SmallVector<Type> nonHoistable;
    for (auto [trueVal, falseVal] : llvm::zip(trueOperands, falseOperands)) {
      if (&op.getTrueBranch() == trueVal.getParentRegion() ||
          &op.getFalseBranch() == falseVal.getParentRegion())
        nonHoistable.push_back(trueVal.getType());
    }

    // Early exit if there aren't any yielded values we can
    // hoist outside the if.
    if (nonHoistable.size() == op->getNumResults())
      return failure();

    auto replacement =
        stablehlo::IfOp::create(rewriter, op.getLoc(), nonHoistable, pred);
    replacement.getTrueBranch().takeBody(op.getTrueBranch());
    replacement.getFalseBranch().takeBody(op.getFalseBranch());

    SmallVector<Value> results(op->getNumResults());
    assert(trueOperands.size() == results.size());
    assert(falseOperands.size() == results.size());

    SmallVector<Value> trueReturns;
    SmallVector<Value> falseReturns;
    rewriter.setInsertionPoint(replacement);
    for (const auto &it :
         llvm::enumerate(llvm::zip(trueOperands, falseOperands))) {
      Value trueVal = std::get<0>(it.value());
      Value falseVal = std::get<1>(it.value());
      if (&replacement.getTrueBranch() == trueVal.getParentRegion() ||
          &replacement.getFalseBranch() == falseVal.getParentRegion()) {
        results[it.index()] = replacement.getResult(trueReturns.size());
        trueReturns.push_back(trueVal);
        falseReturns.push_back(falseVal);
      } else if (trueVal == falseVal)
        results[it.index()] = trueVal;
      else
        results[it.index()] = stablehlo::SelectOp::create(
            rewriter, op.getLoc(), pred, trueVal, falseVal);
    }

    rewriter.setInsertionPointToEnd(&replacement.getTrueBranch().front());
    rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(trueTerm, trueReturns);

    rewriter.setInsertionPointToEnd(&replacement.getFalseBranch().front());
    rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(falseTerm, falseReturns);

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

// https://github.com/llvm/llvm-project/blob/74d8f3952c4acf6d57948983d7c5b0d0a7763c28/mlir/lib/Dialect/SCF/IR/SCF.cpp#L2313
struct SpeculateIfPadToSelect final
    : public CheckedOpRewritePattern<stablehlo::IfOp, SpeculateIfPadToSelect> {
  using CheckedOpRewritePattern<
      stablehlo::IfOp, SpeculateIfPadToSelect>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IfOp op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumResults() == 0 || op.getTrueBranch().empty() ||
        op.getFalseBranch().empty())
      return failure();

    auto trueOperands =
        op.getTrueBranch().front().getTerminator()->getOperands();
    auto falseOperands =
        op.getFalseBranch().front().getTerminator()->getOperands();

    bool anyPad = false;
    for (auto [trueVal, falseVal] : llvm::zip(trueOperands, falseOperands)) {
      if (&op.getTrueBranch() != trueVal.getParentRegion() &&
          &op.getFalseBranch() == falseVal.getParentRegion()) {
        if (auto pad = falseVal.getDefiningOp<stablehlo::PadOp>()) {
          if (pad.getOperand().getParentRegion() != &op.getFalseBranch() &&
              pad.getPaddingValue().getParentRegion() != &op.getFalseBranch()) {
            rewriter.modifyOpInPlace(pad, [&]() { pad->moveBefore(op); });
            anyPad = true;
            continue;
          }
        }
      }

      if (&op.getFalseBranch() != falseVal.getParentRegion() &&
          &op.getTrueBranch() == trueVal.getParentRegion()) {
        if (auto pad = trueVal.getDefiningOp<stablehlo::PadOp>()) {
          if (pad.getOperand().getParentRegion() != &op.getTrueBranch() &&
              pad.getPaddingValue().getParentRegion() != &op.getTrueBranch()) {
            rewriter.modifyOpInPlace(pad, [&]() { pad->moveBefore(op); });
            anyPad = true;
            continue;
          }
        }
      }
    }

    if (anyPad)
      return success();
    else
      return failure();
  }
};

bool verifyInversePermutations(stablehlo::TransposeOp innerTrans,
                               stablehlo::TransposeOp outerTrans) {
  auto innerPerm = innerTrans.getPermutation();
  auto outerPerm = outerTrans.getPermutation();

  if (innerPerm.size() != outerPerm.size())
    return false;

  SmallVector<int64_t> composition(innerPerm.size());
  for (size_t i = 0; i < innerPerm.size(); ++i) {
    composition[i] = outerPerm[innerPerm[i]];
  }

  // Check if the composition is the identity permutation
  for (size_t i = 0; i < composition.size(); ++i) {
    if (composition[i] != static_cast<int64_t>(i))
      return false;
  }

  return true;
}

// Currently supports:
// 1. Identifies induction variable
// 2. Addition of constant step value
// 3. Less than comparision
struct WhileOpInductionReplacement
    : public CheckedOpRewritePattern<stablehlo::WhileOp,
                                     WhileOpInductionReplacement> {
  using CheckedOpRewritePattern<
      stablehlo::WhileOp, WhileOpInductionReplacement>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Only handle while loops with identifiable iteration patterns
    bool canonicalized = false;

    // Look for a loop counter variable (induction variable analog)
    Block &bodyBlock = whileOp.getBody().front();
    auto returnOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Find the counter variable and its limit before entering the loop
    // This is used to identify and optimize induction variables
    unsigned counterIdx = 0;
    Value limitValue = nullptr;
    bool hasCounter = findCounterAndLimit(whileOp, counterIdx, limitValue);

    // If we can't find a counter variable, we can't optimize induction
    // variables
    if (!hasCounter)
      return failure();

    if (!definedOutside(limitValue, whileOp))
      return failure();

    // Get the counter argument and its start value for later use
    Value counterArg = whileOp.getBody().getArgument(counterIdx);
    Value startValue = findCounterStartValue(whileOp, counterIdx);

    // Find the counter step value (how much the counter increments each
    // iteration) This is needed to correctly scale the induction variable
    // calculation
    Value counterStepValue = findCounterStepValue(whileOp, counterIdx);
    if (!counterStepValue)
      return failure();

    // Examine each iteration argument and result
    for (unsigned i = 0; i < whileOp.getOperands().size(); ++i) {
      // Get the input, the body argument, and the yielded value
      Value initValue = whileOp.getOperands()[i];
      BlockArgument iterArg = bodyBlock.getArgument(i);
      Value yieldedValue = returnOp.getOperand(i);
      Value result = whileOp.getResult(i);

      // Look for a simple step pattern: iter_arg + step, step + iter_arg, or
      // iter_arg - step. A counter that walks backwards is exactly as redundant
      // as one that walks forwards; reverse-mode AD emits the backwards form,
      // indexing rows from the end while the loop counter runs from the start.
      Operation *stepOp = yieldedValue.getDefiningOp();
      if (!isa_and_nonnull<stablehlo::AddOp, stablehlo::SubtractOp>(stepOp))
        continue;

      if (!cast<RankedTensorType>(stepOp->getResult(0).getType())
               .getElementType()
               .isInteger())
        continue;

      // Check which operand is the iteration argument and which is the step
      Value stepValue;
      bool negateStep = false;
      if (isa<stablehlo::SubtractOp>(stepOp)) {
        // Pattern: iter_arg - step. (step - iter_arg is not a counter.)
        if (stepOp->getOperand(0) != iterArg)
          continue;
        stepValue = stepOp->getOperand(1);
        negateStep = true;
      } else if (stepOp->getOperand(0) == iterArg) {
        // Pattern: iter_arg + step
        stepValue = stepOp->getOperand(1);
      } else if (stepOp->getOperand(1) == iterArg) {
        // Pattern: step + iter_arg
        stepValue = stepOp->getOperand(0);
      } else {
        // Neither operand is the iteration argument
        continue;
      }

      // Find if the step is a constant
      auto constOp = stepValue.getDefiningOp<stablehlo::ConstantOp>();
      if (!constOp)
        continue;

      // Fold the sign into the constant rather than emitting a negate, so the
      // rewritten index stays a plain affine expression of the counter.
      if (negateStep) {
        DenseIntElementsAttr stepAttr;
        if (!matchPattern(constOp.getResult(), m_Constant(&stepAttr)))
          continue;
        SmallVector<APInt> negated;
        for (auto v : stepAttr.getValues<APInt>())
          negated.push_back(-v);
        OpBuilder::InsertionGuard guard(rewriter);
        rewriter.setInsertionPoint(whileOp);
        constOp = stablehlo::ConstantOp::create(
            rewriter, constOp.getLoc(),
            DenseIntElementsAttr::get(cast<ShapedType>(constOp.getType()),
                                      negated));
        stepValue = constOp.getResult();
      }

      // Similarly replace uses of the result outside the loop
      // with a calculation based on the final counter value
      if (!result.use_empty() && limitValue) {
        rewriter.setInsertionPointAfter(whileOp);

        // Calculate total iterations: limit - start
        Value totalIters = stablehlo::SubtractOp::create(
            rewriter, whileOp.getLoc(), limitValue, startValue);

        // First multiply by the step value (using the same step value
        // identified earlier)
        Value scaledOffset = stablehlo::MulOp::create(
            rewriter, whileOp.getLoc(),
            getCorrectlySizedValue(totalIters, stepValue, rewriter), stepValue);

        // Then divide by the counter step value to get the correct scaling
        Value normalizedOffset = stablehlo::DivOp::create(
            rewriter, whileOp.getLoc(), scaledOffset,
            getCorrectlySizedValue(counterStepValue, scaledOffset, rewriter));

        Value finalValue = stablehlo::AddOp::create(
            rewriter, whileOp.getLoc(), initValue, normalizedOffset);

        rewriter.replaceAllUsesWith(result, finalValue);
        canonicalized = true;
      }

      // Skip the counter variable itself - we don't want to optimize it away
      if (i == counterIdx)
        continue;

      // Now we can replace uses of the iterArg inside the loop
      // with a direct calculation based on the counter:
      // replacement = init_value + ((counter - start_value) * step_value) /
      // counter_step_value
      if (!iterArg.use_empty()) {
        rewriter.setInsertionPointToStart(&bodyBlock);

        // Create the calculation for the current iteration
        Value iterOffset = stablehlo::SubtractOp::create(
            rewriter, whileOp.getLoc(), counterArg, startValue);

        // First multiply by the step value
        Value scaledOffset = stablehlo::MulOp::create(
            rewriter, whileOp.getLoc(),
            getCorrectlySizedValue(iterOffset, stepValue, rewriter), stepValue);

        // Then divide by the counter step value to get the correct scaling
        Value normalizedOffset = stablehlo::DivOp::create(
            rewriter, whileOp.getLoc(), scaledOffset,
            getCorrectlySizedValue(counterStepValue, scaledOffset, rewriter));

        Value replacement = stablehlo::AddOp::create(
            rewriter, whileOp.getLoc(), initValue, normalizedOffset);

        rewriter.modifyOpInPlace(
            whileOp, [&] { iterArg.replaceAllUsesWith(replacement); });
        canonicalized = true;
      }
    }

    return success(canonicalized);
  }

private:
  Value getCorrectlySizedValue(Value origValue, Value targetValue,
                               PatternRewriter &rewriter) const {
    return getCorrectlySizedValue(
        origValue, cast<RankedTensorType>(targetValue.getType()), rewriter);
  }

  Value getCorrectlySizedValue(Value origValue, RankedTensorType targetType,
                               PatternRewriter &rewriter) const {
    if (origValue.getType() == targetType)
      return origValue;

    if (cast<ShapedType>(origValue.getType()).getElementType() !=
        targetType.getElementType()) {
      origValue = stablehlo::ConvertOp::create(
          rewriter, origValue.getLoc(),
          RankedTensorType::get({}, targetType.getElementType()), origValue);
    }

    if (origValue.getType() != targetType) {
      origValue = stablehlo::BroadcastInDimOp::create(
          rewriter, origValue.getLoc(), targetType, origValue,
          rewriter.getDenseI64ArrayAttr({}));
    }

    return origValue;
  }

  // Helper function to identify the counter variable and its limit
  // Returns the index of the counter argument and the limit value
  bool findCounterAndLimit(stablehlo::WhileOp whileOp, unsigned &counterIdx,
                           Value &limitValue) const {
    // Look in the condition region for a comparison operation
    Block &condBlock = whileOp.getCond().front();
    Operation *terminator = condBlock.getTerminator();

    // Typical pattern: return %comparison
    if (auto returnOp = dyn_cast<stablehlo::ReturnOp>(terminator)) {
      if (returnOp.getNumOperands() != 1)
        return false;

      // Look for a comparison that controls the loop
      auto compareOp =
          returnOp.getOperand(0).getDefiningOp<stablehlo::CompareOp>();
      if (!compareOp)
        return false;

      if (compareOp.getComparisonDirection() !=
          stablehlo::ComparisonDirection::LT)
        return false;

      if (!compareOp.getLhs().getType().getElementType().isInteger())
        return false;

      // Check if one side is a block argument (our counter)
      if (auto blockArg = dyn_cast<BlockArgument>(compareOp.getLhs())) {
        if (blockArg.getOwner() == &condBlock) {
          counterIdx = blockArg.getArgNumber();
          limitValue = compareOp.getRhs();
          return true;
        }
      }

      if (auto blockArg = dyn_cast<BlockArgument>(compareOp.getRhs())) {
        if (blockArg.getOwner() == &condBlock) {
          counterIdx = blockArg.getArgNumber();
          limitValue = compareOp.getLhs();
          return true;
        }
      }
    }

    return false;
  }

  // Helper to find the initial value of the counter
  Value findCounterStartValue(stablehlo::WhileOp whileOp,
                              unsigned counterIdx) const {
    // The initial value is the corresponding operand to the while op
    return whileOp.getOperands()[counterIdx];
  }

  // Helper to find the counter step value (how much it increments each
  // iteration)
  Value findCounterStepValue(stablehlo::WhileOp whileOp,
                             unsigned counterIdx) const {
    // Get the block argument in the body region
    Block &bodyBlock = whileOp.getBody().front();
    BlockArgument counterArg = bodyBlock.getArgument(counterIdx);

    // Find the terminator to get the yielded value
    auto returnOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Get the yielded value for the counter
    Value yieldedCounter = returnOp.getOperand(counterIdx);

    // Look for addition pattern: counter + step or step + counter
    auto addOp =
        dyn_cast_or_null<stablehlo::AddOp>(yieldedCounter.getDefiningOp());
    if (!addOp)
      return nullptr;

    // Check both sides of the addition operation (since addition is
    // commutative)
    if (addOp.getLhs() == counterArg) {
      // Pattern: counter + step
      return addOp.getRhs();
    }

    if (addOp.getRhs() == counterArg) {
      // Pattern: step + counter
      return addOp.getLhs();
    }

    // Counter is not directly used in the addition
    return nullptr;
  }
};

struct TransposeWhile
    : public CheckedOpRewritePattern<stablehlo::WhileOp, TransposeWhile> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct TransposeCandidate {
      unsigned idx;
      stablehlo::TransposeOp outerTranspose;
    };

    llvm::SmallVector<TransposeCandidate, 4> outerTransposes;

    // For each yielded value, check if it comes from a transpose
    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {
      // Check if the corresponding while result is used by a transpose
      Value whileResult = whileOp.getResult(idx);

      // Check that the while result has exactly one use
      if (!whileResult.hasOneUse())
        continue;

      // Get the single user and verify it's a transpose
      Operation *singleUser = *whileResult.getUsers().begin();
      auto outerTransposeOp = dyn_cast<stablehlo::TransposeOp>(singleUser);
      if (!outerTransposeOp)
        continue;

      outerTransposes.emplace_back(TransposeCandidate{idx, outerTransposeOp});
    }

    // If no candidates found, no rewrite needed
    if (outerTransposes.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create input transposes for each candidate
    for (auto &candidate : outerTransposes) {
      unsigned idx = candidate.idx;
      stablehlo::TransposeOp outerTranspose = candidate.outerTranspose;

      // Create a new transpose before the while loop
      auto inputTranspose = stablehlo::TransposeOp::create(
          rewriter, whileOp.getLoc(),
          outerTranspose.getType(), // The type after transposition
          whileOperands[idx],       // Original input to while
          outerTranspose.getPermutation());

      // Update the while operand to use the new transposed value
      newOperands[idx] = inputTranspose;
    }

    // Update yield op to use the input of the inner transpose
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : outerTransposes) {
        unsigned idx = candidate.idx;
        newReturnValues[idx] = stablehlo::TransposeOp::create(
            rewriter, whileOp.getLoc(), newReturnValues[idx],
            candidate.outerTranspose.getPermutation());
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock.clear();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      SmallVector<Value> replacements;
      for (auto arg : newBodyBlock.getArguments())
        replacements.push_back(arg);
      for (auto &pair : outerTransposes) {
        replacements[pair.idx] = stablehlo::TransposeOp::create(
            rewriter, pair.outerTranspose.getLoc(), replacements[pair.idx],
            getInversePermutation(pair.outerTranspose.getPermutation()));
      }
      rewriter.mergeBlocks(&whileOp.getBody().front(), &newBodyBlock,
                           replacements);
    }

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      SmallVector<Value> replacements;
      for (auto arg : newCondBlock.getArguments())
        replacements.push_back(arg);
      for (auto &pair : outerTransposes) {
        replacements[pair.idx] = stablehlo::TransposeOp::create(
            rewriter, pair.outerTranspose.getLoc(), replacements[pair.idx],
            getInversePermutation(pair.outerTranspose.getPermutation()));
      }
      rewriter.mergeBlocks(&whileOp.getCond().front(), &newCondBlock,
                           replacements);
    }

    // Step 5. Replace outerTranspose with the newWhileOp results
    for (auto &candidate : outerTransposes) {
      unsigned idx = candidate.idx;
      stablehlo::TransposeOp outerTranspose = candidate.outerTranspose;

      // Update the while operand to use the new transposed value
      rewriter.replaceOp(outerTranspose, newWhileOp.getResult(idx));
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, newWhileOp.getResults());
    return success();
  }
};

bool isLegalConcatToOneDimDUS(stablehlo::ConcatenateOp outer,
                              stablehlo::SliceOp *lhsP,
                              stablehlo::SliceOp *rhsP = nullptr,
                              Value *operand = nullptr) {
  if (outer.getOperands().size() < 2)
    return false;
  SmallVector<stablehlo::ConcatenateOp> inners;

  stablehlo::SliceOp lhs = nullptr;
  if (auto lhsSlice =
          outer.getOperands()[0].getDefiningOp<stablehlo::SliceOp>()) {
    for (int i = 0; i < lhsSlice.getType().getShape().size(); i++) {
      if (lhsSlice.getStartIndices()[i] != 0) {
        return false;
      }
      if (i != outer.getDimension()) {
        if (cast<RankedTensorType>(lhsSlice.getOperand().getType())
                .getShape()[i] != outer.getType().getShape()[i]) {
          return false;
        }
      }
      if (lhsSlice.getStrides()[i] != 1) {
        return false;
      }
    }
    lhs = lhsSlice;
  }

  stablehlo::SliceOp rhs = nullptr;
  if (auto rhsSlice =
          outer.getOperands().back().getDefiningOp<stablehlo::SliceOp>()) {
    for (int i = 0; i < rhsSlice.getType().getShape().size(); i++) {
      if (rhsSlice.getLimitIndices()[i] != outer.getType().getShape()[i]) {
        return false;
      }
      if (i != outer.getDimension()) {
        if (cast<RankedTensorType>(rhsSlice.getOperand().getType())
                .getShape()[i] != outer.getType().getShape()[i]) {
          return false;
        }
      }
      if (rhsSlice.getStrides()[i] != 1) {
        return false;
      }
    }
    rhs = rhsSlice;
  }

  if (!lhs && !rhs) {
    return false;
  }
  if (!rhs && cast<RankedTensorType>(lhs.getOperand().getType())
                      .getShape()[outer.getDimension()] !=
                  outer.getType().getShape()[outer.getDimension()]) {
    return false;
  }
  if (!lhs && cast<RankedTensorType>(rhs.getOperand().getType())
                      .getShape()[outer.getDimension()] !=
                  outer.getType().getShape()[outer.getDimension()]) {
    return false;
  }

  if (lhs && rhs && outer.getOperands().size() == 2) {
    return false;
  }
  if (lhs && rhs) {
    if (lhs.getOperand() != rhs.getOperand()) {
      return false;
    }
    if (cast<RankedTensorType>(lhs.getOperand().getType())
            .getShape()[outer.getDimension()] !=
        outer.getType().getShape()[outer.getDimension()]) {
      return false;
    }
  }
  if (lhsP)
    *lhsP = lhs;
  if (rhsP)
    *rhsP = rhs;
  if (operand && lhs)
    *operand = lhs.getOperand();
  if (operand && rhs)
    *operand = rhs.getOperand();
  return true;
}

stablehlo::DynamicUpdateSliceOp
concatToOneDimDUS(PatternRewriter &rewriter, stablehlo::ConcatenateOp outer) {
  stablehlo::SliceOp lhs = nullptr;
  stablehlo::SliceOp rhs = nullptr;
  if (!isLegalConcatToOneDimDUS(outer, &lhs, &rhs))
    return nullptr;

  auto shard = sdy::getShardingPerValue(outer);

  SmallVector<Value> newOps;
  int start = lhs ? 1 : 0;
  int end = outer.getOperands().size() - (rhs ? 1 : 0);
  for (int i = start; i < end; i++) {
    newOps.push_back(outer.getOperands()[i]);
  }
  Value innerConcat = newOps[0];

  auto iTy = RankedTensorType::get({}, rewriter.getI64Type());
  Value operand = lhs ? lhs.getOperand() : rhs.getOperand();
  assert(operand.getType() == outer.getType());
  SmallVector<Value> starts(
      outer.getType().getShape().size(),
      stablehlo::ConstantOp::create(rewriter, outer.getLoc(), iTy,
                                    cast<ElementsAttr>(makeAttr(iTy, 0))));

  if (lhs) {
    starts[outer.getDimension()] = stablehlo::ConstantOp::create(
        rewriter, outer.getLoc(), iTy,
        cast<ElementsAttr>(
            makeAttr(iTy, lhs.getType().getShape()[outer.getDimension()])));
  }

  rewriter.setInsertionPointAfter(outer);
  if (newOps.size() != 1) {
    auto nConcat = stablehlo::ConcatenateOp::create(
        rewriter, outer.getLoc(), newOps, outer.getDimension());
    innerConcat = nConcat;
    if (shard) {
      sdy::setShardings(nConcat, shard);
    }
  }

  auto OT = outer.getType();
  assert(operand.getType() == OT);
  if (operand.getType() != OT) {
    llvm_unreachable("invalid replacement (0)\n");
  }
  auto dus = stablehlo::DynamicUpdateSliceOp::create(
      rewriter, outer.getLoc(), operand, innerConcat, starts);
  assert(dus.getType() == OT);
  if (dus.getType() != OT) {
    llvm_unreachable("invalid replacement\n");
  }
  if (shard) {
    sdy::setShardings(dus, shard);
  }
  rewriter.replaceOp(outer, dus);
  return dus;
}

struct ConcatToOneDimDUS final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatToOneDimDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp outer,
                                    PatternRewriter &rewriter) const {
    if (concatToOneDimDUS(rewriter, outer))
      return success();
    else
      return failure();
  }
};

struct WhileDUS : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct DUSCandidate {
      unsigned idx;
      stablehlo::DynamicUpdateSliceOp DUS;
      mlir::Value outerOperand;
      mlir::Value conditionalOperand;
    };

    llvm::SmallVector<DUSCandidate, 4> candidates;
    bool hasConditional = false;

    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      Value DUSOperand = nullptr;

      if (auto DUS = yieldOp.getOperand(idx)
                         .getDefiningOp<stablehlo::DynamicUpdateSliceOp>()) {
        DUSOperand = DUS.getOperand();

        bool legal = true;
        for (auto idx : DUS.getStartIndices()) {
          if (!definedOutside(idx, whileOp)) {
            legal = false;
          }
        }

        if (!legal)
          continue;
      } else if (auto concat = yieldOp.getOperand(idx)
                                   .getDefiningOp<stablehlo::ConcatenateOp>()) {
        Value operand;
        if (isLegalConcatToOneDimDUS(concat, nullptr, nullptr, &operand)) {
          DUSOperand = operand;
        }
      }

      // Check that the while result has exactly one use
      if (!DUSOperand)
        continue;

      mlir::Value conditionalOperand = nullptr;
      if (DUSOperand == whileOp.getBody().front().getArgument(idx)) {
      } else if (definedOutside(DUSOperand, whileOp)) {

        bool hasArgUse = !whileOp.getCond().getArgument(idx).use_empty() ||
                         !whileOp.getBody().getArgument(idx).use_empty();

        if (hasArgUse) {
          continue;
        }

        conditionalOperand = DUSOperand;
        hasConditional = true;
      } else {
        continue;
      }

      auto DUS = yieldOp.getOperand(idx)
                     .getDefiningOp<stablehlo::DynamicUpdateSliceOp>();

      if (!DUS) {
        mlir::OpBuilder::InsertionGuard guard(rewriter);
        auto concat =
            yieldOp.getOperand(idx).getDefiningOp<stablehlo::ConcatenateOp>();
        assert(concat);
        DUS = concatToOneDimDUS(rewriter, concat);
        assert(DUS);
        for (auto idx : DUS.getStartIndices()) {
          (void)idx;
          assert(definedOutside(idx, whileOp));
        }
      }

      candidates.emplace_back(DUSCandidate{idx, DUS, whileOp.getOperands()[idx],
                                           conditionalOperand});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create input transposes for each candidate
    for (auto &candidate : candidates) {
      // Create a new transpose before the while loop

      newOperands[candidate.idx] = stablehlo::DynamicSliceOp::create(
          rewriter, candidate.DUS.getLoc(), candidate.outerOperand,
          candidate.DUS.getStartIndices(),
          candidate.DUS.getUpdate().getType().getShape());
    }

    // Update yield op to use the input of the inner transpose
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : candidates) {
        newReturnValues[candidate.idx] = candidate.DUS.getUpdate();
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    SmallVector<Value> results;
    for (auto res : newWhileOp.getResults())
      results.push_back(res);

    {
      mlir::IRMapping mapper;
      Value useInner = nullptr;
      if (hasConditional) {

        for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
          mapper.map(whileOp.getCond().getArgument(i),
                     whileOp.getOperands()[i]);
        }
        for (auto &op : whileOp.getCond().front().getOperations()) {
          // Skip the terminator - we'll add it after all other operations
          if (isa<stablehlo::ReturnOp>(op))
            continue;

          // Clone the operation with the value mapping
          rewriter.clone(op, mapper);
        }
        useInner = whileOp.getCond().front().getTerminator()->getOperand(0);
        useInner = mapper.lookupOrDefault(useInner);
      }
      for (auto &candidate : candidates) {
        Value operand = candidate.outerOperand;
        if (candidate.conditionalOperand) {
          operand = stablehlo::SelectOp::create(
              rewriter, whileOp.getLoc(), useInner,
              candidate.conditionalOperand, operand);
        }

        results[candidate.idx] = stablehlo::DynamicUpdateSliceOp::create(
            rewriter, candidate.DUS.getLoc(), operand, results[candidate.idx],
            candidate.DUS.getStartIndices());
      }
    }

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    // Create an IR mapper to map values from old op to new op
    mlir::IRMapping mapper;

    // Clear the new body block but keep its arguments
    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock
        .clear(); // This clears operations but preserves block arguments

    // Clone operations from old body to new body
    Block &oldBodyBlock = whileOp.getBody().front();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      // Set up operand mapping for the body region
      for (unsigned i = 0; i < whileOp.getBody().getNumArguments(); ++i) {
        auto oldArg = whileOp.getBody().getArgument(i);
        Value newArg = newWhileOp.getBody().getArgument(i);
        for (auto &pair : candidates) {
          if (pair.idx == i) {
            newArg = stablehlo::DynamicUpdateSliceOp::create(
                rewriter, pair.DUS.getLoc(), pair.outerOperand, newArg,
                pair.DUS.getStartIndices());
            break;
          }
        }
        mapper.map(oldArg, newArg);
      }

      for (auto &op : oldBodyBlock.getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
    }

    // Create a new terminator for the body region using new values
    {
      OpBuilder::InsertionGuard guard(rewriter);
      SmallVector<Value> newReturnValues;

      // Map old return values to new values using the mapper
      for (auto oldRetVal : yieldOp.getOperands()) {
        Value newRetVal = mapper.lookupOrNull(oldRetVal);
        // If the value isn't in the mapper, maybe it was a block argument or
        // constant
        if (!newRetVal)
          newRetVal = oldRetVal; // Consider more robust handling if needed
        newReturnValues.push_back(newRetVal);
      }

      // Create the return op at the end of the body
      rewriter.setInsertionPointToEnd(&newBodyBlock);
      stablehlo::ReturnOp::create(rewriter, yieldOp.getLoc(), newReturnValues);
    }

    // Create condition region mapper
    mlir::IRMapping condMapper;

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    Block &oldCondBlock = whileOp.getCond().front();

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        auto oldArg = whileOp.getCond().getArgument(i);
        Value newArg = newWhileOp.getCond().getArgument(i);
        for (auto &pair : candidates) {
          if (pair.idx == i) {
            newArg = stablehlo::DynamicUpdateSliceOp::create(
                rewriter, pair.DUS.getLoc(), pair.outerOperand, newArg,
                pair.DUS.getStartIndices());
            break;
          }
        }
        condMapper.map(oldArg, newArg);
      }

      for (auto &op : oldCondBlock.getOperations()) {
        rewriter.clone(op, condMapper);
      }
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, results);

    return success();
  }
};

struct WhileUpdateWithoutCorners
    : public CheckedOpRewritePattern<stablehlo::WhileOp,
                                     WhileUpdateWithoutCorners> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct DUSCandidate {
      unsigned idx;
      enzymexla::UpdateWithoutCornersOp DUS;
      mlir::Value outerOperand;
      mlir::Value conditionalOperand;
    };

    llvm::SmallVector<DUSCandidate, 4> candidates;
    bool hasConditional = false;

    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      auto DUS = yieldOp.getOperand(idx)
                     .getDefiningOp<enzymexla::UpdateWithoutCornersOp>();
      if (!DUS)
        continue;
      Value DUSOperand = DUS.getOperand();

      // Check that the while result has exactly one use

      mlir::Value conditionalOperand = nullptr;
      if (DUSOperand == whileOp.getBody().front().getArgument(idx)) {
      } else if (definedOutside(DUSOperand, whileOp)) {

        bool hasArgUse = !whileOp.getCond().getArgument(idx).use_empty() ||
                         !whileOp.getBody().getArgument(idx).use_empty();

        if (hasArgUse) {
          continue;
        }

        conditionalOperand = DUSOperand;
        hasConditional = true;
      } else {
        continue;
      }

      candidates.emplace_back(DUSCandidate{idx, DUS, whileOp.getOperands()[idx],
                                           conditionalOperand});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create input transposes for each candidate
    for (auto &candidate : candidates) {
      // Create a new transpose before the while loop

      newOperands[candidate.idx] = candidate.outerOperand;
    }

    // Update yield op to use the input of the inner transpose
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : candidates) {
        newReturnValues[candidate.idx] = candidate.DUS.getUpdate();
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    SmallVector<Value> results;
    for (auto res : newWhileOp.getResults())
      results.push_back(res);

    {
      mlir::IRMapping mapper;
      Value useInner = nullptr;
      if (hasConditional) {

        for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
          mapper.map(whileOp.getCond().getArgument(i),
                     whileOp.getOperands()[i]);
        }
        for (auto &op : whileOp.getCond().front().getOperations()) {
          // Skip the terminator - we'll add it after all other operations
          if (isa<stablehlo::ReturnOp>(op))
            continue;

          // Clone the operation with the value mapping
          rewriter.clone(op, mapper);
        }
        useInner = whileOp.getCond().front().getTerminator()->getOperand(0);
        useInner = mapper.lookupOrDefault(useInner);
      }
      for (auto &candidate : candidates) {
        Value operand = candidate.outerOperand;
        if (candidate.conditionalOperand) {
          operand = stablehlo::SelectOp::create(
              rewriter, whileOp.getLoc(), useInner,
              candidate.conditionalOperand, operand);
        }

        results[candidate.idx] = enzymexla::UpdateWithoutCornersOp::create(
            rewriter, candidate.DUS.getLoc(), operand, results[candidate.idx],
            candidate.DUS.getDimensionX(), candidate.DUS.getX1(),
            candidate.DUS.getX2(), candidate.DUS.getDimensionY(),
            candidate.DUS.getY1(), candidate.DUS.getY2());
      }
    }

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    // Create an IR mapper to map values from old op to new op
    mlir::IRMapping mapper;

    // Clear the new body block but keep its arguments
    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock
        .clear(); // This clears operations but preserves block arguments

    // Clone operations from old body to new body
    Block &oldBodyBlock = whileOp.getBody().front();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      // Set up operand mapping for the body region
      for (unsigned i = 0; i < whileOp.getBody().getNumArguments(); ++i) {
        auto oldArg = whileOp.getBody().getArgument(i);
        Value newArg = newWhileOp.getBody().getArgument(i);
        for (auto &pair : candidates) {
          if (pair.idx == i) {
            newArg = enzymexla::UpdateWithoutCornersOp::create(
                rewriter, pair.DUS.getLoc(), pair.outerOperand, newArg,
                pair.DUS.getDimensionX(), pair.DUS.getX1(), pair.DUS.getX2(),
                pair.DUS.getDimensionY(), pair.DUS.getY1(), pair.DUS.getY2());
            break;
          }
        }
        mapper.map(oldArg, newArg);
      }

      for (auto &op : oldBodyBlock.getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
    }

    // Create a new terminator for the body region using new values
    {
      OpBuilder::InsertionGuard guard(rewriter);
      SmallVector<Value> newReturnValues;

      // Map old return values to new values using the mapper
      for (auto oldRetVal : yieldOp.getOperands()) {
        Value newRetVal = mapper.lookupOrNull(oldRetVal);
        // If the value isn't in the mapper, maybe it was a block argument or
        // constant
        if (!newRetVal)
          newRetVal = oldRetVal; // Consider more robust handling if needed
        newReturnValues.push_back(newRetVal);
      }

      // Create the return op at the end of the body
      rewriter.setInsertionPointToEnd(&newBodyBlock);
      stablehlo::ReturnOp::create(rewriter, yieldOp.getLoc(), newReturnValues);
    }

    // Create condition region mapper
    mlir::IRMapping condMapper;

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    Block &oldCondBlock = whileOp.getCond().front();

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        auto oldArg = whileOp.getCond().getArgument(i);
        Value newArg = newWhileOp.getCond().getArgument(i);
        for (auto &pair : candidates) {
          if (pair.idx == i) {
            newArg = enzymexla::UpdateWithoutCornersOp::create(
                rewriter, pair.DUS.getLoc(), pair.outerOperand, newArg,
                pair.DUS.getDimensionX(), pair.DUS.getX1(), pair.DUS.getX2(),
                pair.DUS.getDimensionY(), pair.DUS.getY1(), pair.DUS.getY2());
            break;
          }
        }
        condMapper.map(oldArg, newArg);
      }

      for (auto &op : oldCondBlock.getOperations()) {
        rewriter.clone(op, condMapper);
      }
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, results);

    return success();
  }
};

struct IVInfo {
  int index;    // Index of the induction variable in the while op arguments
  int64_t step; // Step size (how much IV increments each iteration)
  Value limit;
  Value start;
  bool isValid; // Whether we successfully identified the IV and step
  bool zeroStart;
  bool canonical;
};

IVInfo extractSimpleIVInfo(stablehlo::WhileOp whileOp) {
  IVInfo result = {-1, 0, nullptr, nullptr, false, false, false};

  // 1. Get the condition block and its return operation
  if (whileOp.getBody().empty() || whileOp.getCond().empty())
    return result;

  Block &condBlock = whileOp.getCond().front();
  if (condBlock.empty())
    return result;

  auto condReturn = dyn_cast<stablehlo::ReturnOp>(condBlock.getTerminator());
  if (!condReturn || condReturn.getNumOperands() != 1)
    return result;

  // 2. Check if the return operand comes from a compare operation
  auto compareOp =
      condReturn.getOperand(0).getDefiningOp<stablehlo::CompareOp>();
  if (!compareOp)
    return result;

  // 3. Identify which block argument is used in the comparison (the IV)
  Value lhs = compareOp.getLhs();
  Value rhs = compareOp.getRhs();

  // Check which side of the comparison is a block argument
  BlockArgument ivArg = nullptr;
  Value limit;
  for (BlockArgument arg : condBlock.getArguments()) {
    if (lhs == arg && (compareOp.getComparisonDirection() ==
                           stablehlo::ComparisonDirection::LT ||
                       compareOp.getComparisonDirection() ==
                           stablehlo::ComparisonDirection::NE)) {
      ivArg = arg;
      limit = rhs;
      break;
    }
    if (rhs == arg && (compareOp.getComparisonDirection() ==
                           stablehlo::ComparisonDirection::GT ||
                       compareOp.getComparisonDirection() ==
                           stablehlo::ComparisonDirection::NE)) {
      ivArg = arg;
      limit = rhs;
      break;
    }
  }

  if (!ivArg)
    return result;

  if (!definedOutside(limit, whileOp))
    return result;

  result.limit = limit;

  // Record the IV index
  result.index = ivArg.getArgNumber();

  // 4. Now analyze the body region to find the step
  Block &bodyBlock = whileOp.getBody().front();
  auto bodyReturn = dyn_cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
  if (!bodyReturn || bodyReturn.getNumOperands() <= result.index)
    return result;

  // Get the value being returned for the IV position
  Value updatedIV = bodyReturn.getOperand(result.index);

  // Look for an add operation that defines the updated IV
  auto addOp = updatedIV.getDefiningOp<stablehlo::AddOp>();
  if (!addOp)
    return result;

  // Check if one side is a constant (the step)
  for (int i = 0; i < 2; i++) {
    if (addOp->getOperand(i) !=
        whileOp.getBody().getArgument(ivArg.getArgNumber()))
      continue;

    DenseIntElementsAttr idxAttr;
    if (!matchPattern(addOp->getOperand(1 - i), m_Constant(&idxAttr)))
      continue;

    result.step = (*idxAttr.begin()).getSExtValue();
    result.isValid = true;
    break;
  }

  result.start = whileOp.getOperands()[result.index];

  DenseIntElementsAttr inpAttr;
  if (matchPattern(result.start, m_Constant(&inpAttr))) {
    result.zeroStart = (*inpAttr.begin()).getSExtValue() == 0;
  }

  result.canonical = result.zeroStart && result.step == 1;

  return result;
}

/// Detects if a value is a replication padding pattern.
stablehlo::ConcatenateOp detectReplicationPadding(Value yieldOperand) {

  // Check if the value is a concatenate operation
  auto concatOp = yieldOperand.getDefiningOp<stablehlo::ConcatenateOp>();
  if (!concatOp)
    return nullptr;

  // Check if we have exactly 3 operands (left border, main tensor, right
  // border)
  if (concatOp.getInputs().size() != 3)
    return nullptr;

  // Get the concat dimension
  int64_t concatDim = concatOp.getDimension();

  // We need to identify potential slice operations that extract borders
  // Check first and last operands
  auto firstSlice = concatOp.getOperand(0).getDefiningOp<stablehlo::SliceOp>();
  auto lastSlice = concatOp.getOperand(concatOp.getNumOperands() - 1)
                       .getDefiningOp<stablehlo::SliceOp>();

  // Check the middle tensor (should be the main tensor)
  Value middleTensor = concatOp.getOperand(1);

  // Both first and last operands must be slice operations
  if (!firstSlice || !lastSlice)
    return nullptr;

  // The slices must come from the same source tensor
  if (firstSlice.getOperand() != middleTensor ||
      lastSlice.getOperand() != middleTensor)
    return nullptr;

  // Verify the slices are taking border elements
  // Get the shapes
  auto middleTensorType = dyn_cast<RankedTensorType>(middleTensor.getType());
  if (!middleTensorType || !middleTensorType.hasStaticShape())
    return nullptr;

  // Check that first slice is taking the first element along concat dimension
  if (firstSlice.getStartIndices()[concatDim] != 0 ||
      firstSlice.getLimitIndices()[concatDim] != 1)
    return nullptr;

  // Check the rest of the firstSlice indices are the same size as the middle
  // tensor
  if (lastSlice.getStartIndices()[concatDim] !=
          middleTensorType.getShape()[concatDim] - 1 ||
      lastSlice.getLimitIndices()[concatDim] !=
          middleTensorType.getShape()[concatDim])
    return nullptr;

  for (int i = 0; i < firstSlice.getStartIndices().size(); i++) {
    if (i == concatDim) // Already checked
      continue;

    // check limit indices are the same size as the middle tensor
    if (firstSlice.getLimitIndices()[i] != middleTensorType.getShape()[i])
      return nullptr;
    // check start indices are 0
    if (firstSlice.getStartIndices()[i] != 0)
      return nullptr;

    // check limit indices are the same size as the middle tensor
    if (lastSlice.getLimitIndices()[i] != middleTensorType.getShape()[i])
      return nullptr;
    // check start indices are 0
    if (lastSlice.getStartIndices()[i] != 0)
      return nullptr;
  }

  return concatOp;
}

struct WhileRepeatedInductionReduction
    : public CheckedOpRewritePattern<stablehlo::WhileOp,
                                     WhileRepeatedInductionReduction> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  template <typename CRangeT, typename RangeT>
  static stablehlo::IfOp
  createConditional(PatternRewriter &rewriter, stablehlo::WhileOp whileOp,
                    IVInfo &ivInfo, CRangeT &candidates, Value iv,
                    const RangeT iterOperand, bool cloneCond) {
    SmallVector<Type> ifResultTypes;

    SmallVector<Value> oldReturns;
    for (auto &candidate : candidates) {
      ifResultTypes.push_back(whileOp.getOperands()[candidate.idx].getType());
      oldReturns.push_back(whileOp.getOperands()[candidate.idx]);
    }

    Value condition;

    if (cloneCond) {
      IRMapping mapper;

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        mapper.map(whileOp.getCond().getArgument(i), whileOp.getOperands()[i]);
      }
      for (auto &op : whileOp.getCond().front().getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
      condition = whileOp.getCond().front().getTerminator()->getOperand(0);
      condition = mapper.lookupOrDefault(condition);
    } else {
      condition = stablehlo::CompareOp::create(
          rewriter, whileOp.getLoc(), iv, ivInfo.start,
          stablehlo::ComparisonDirection::EQ);
    }

    auto ifOp = stablehlo::IfOp::create(rewriter, whileOp.getLoc(),
                                        ifResultTypes, condition);

    // Create the then and else regions for the if operation

    {
      Region &thenRegion =
          cloneCond ? ifOp.getFalseBranch() : ifOp.getTrueBranch();
      OpBuilder::InsertionGuard guard(rewriter);
      Block *thenBlock = rewriter.createBlock(&thenRegion);
      rewriter.setInsertionPointToStart(thenBlock);
      stablehlo::ReturnOp::create(rewriter, whileOp.getLoc(), oldReturns);
    }

    {
      Region &elseRegion =
          cloneCond ? ifOp.getTrueBranch() : ifOp.getFalseBranch();
      OpBuilder::InsertionGuard guard(rewriter);
      Block *elseBlock = rewriter.createBlock(&elseRegion);
      rewriter.setInsertionPointToStart(elseBlock);
      SmallVector<Value> newReturns;

      for (auto &candidate : candidates) {
        auto T =
            cast<RankedTensorType>(candidate.middleTensor.getType()).getShape();

        SmallVector<int64_t> interior(T.size(), 1);

        SmallVector<int64_t> firstLower(T.size(), 0);

        SmallVector<int64_t> firstUpper = llvm::to_vector(T);
        firstUpper[candidate.concat.getDimension()] = 1;

        SmallVector<int64_t> thirdLower(T.size(), 0);
        thirdLower[candidate.concat.getDimension()] =
            T[candidate.concat.getDimension()] - 1;

        SmallVector<int64_t> thirdUpper = llvm::to_vector(T);

        Value args[3] = {
            stablehlo::SliceOp::create(rewriter, whileOp.getLoc(),
                                       iterOperand[candidate.idx], firstLower,
                                       firstUpper, interior)
                .getResult(),
            iterOperand[candidate.idx],
            stablehlo::SliceOp::create(rewriter, whileOp.getLoc(),
                                       iterOperand[candidate.idx], thirdLower,
                                       thirdUpper, interior)
                .getResult(),
        };

        newReturns.push_back(stablehlo::ConcatenateOp::create(
            rewriter, candidate.concat.getLoc(), args,
            candidate.concat.getDimension()));
      }
      stablehlo::ReturnOp::create(rewriter, whileOp.getLoc(), newReturns);
    }

    return ifOp;
  }

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find the index of IV and the step to check for 1 iteration
    auto ivInfo = extractSimpleIVInfo(whileOp);
    if (!ivInfo.isValid)
      return failure();

    if (ivInfo.step == 0)
      return failure();

    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct Candidate {
      unsigned idx;
      stablehlo::ConcatenateOp concat;
      Value middleTensor;
      Value outerOperand;
      SmallVector<int64_t> innerStarts;
      SmallVector<int64_t> innerEnds;
      SmallVector<int64_t> innerStrides;
    };

    llvm::SmallVector<Candidate> candidates;
    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      auto concat = detectReplicationPadding(yieldOp.getOperand(idx));
      if (!concat)
        continue;

      auto T = concat.getType();

      SmallVector<int64_t> innerStarts(T.getShape().size(), 0);
      SmallVector<int64_t> innerEnds = llvm::to_vector(T.getShape());
      SmallVector<int64_t> innerStrides(T.getShape().size(), 1);
      innerStarts[concat.getDimension()] += 1;
      innerEnds[concat.getDimension()] -= 1;
      candidates.emplace_back(Candidate{idx, concat, concat.getOperand(1),
                                        whileOp.getOperands()[idx], innerStarts,
                                        innerEnds, innerStrides});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create input transposes for each candidate
    for (auto &candidate : candidates) {
      newOperands[candidate.idx] = stablehlo::SliceOp::create(
          rewriter, candidate.concat.getLoc(), candidate.outerOperand,
          candidate.innerStarts, candidate.innerEnds, candidate.innerStrides);
    }

    // Update yield op to use the input of the inner transpose
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : candidates) {
        newReturnValues[candidate.idx] = candidate.middleTensor;
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    SmallVector<Value> results;
    for (auto res : newWhileOp.getResults())
      results.push_back(res);

    auto ifOp =
        createConditional(rewriter, whileOp, ivInfo, candidates,
                          whileOp.getOperands()[ivInfo.index], results, true);
    for (auto res : ifOp->getResults()) {
      results[candidates[res.getResultNumber()].idx] = res;
    }

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    // Create an IR mapper to map values from old op to new op
    mlir::IRMapping mapper;

    // Clear the new body block but keep its arguments
    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock
        .clear(); // This clears operations but preserves block arguments

    // Clone operations from old body to new body
    Block &oldBodyBlock = whileOp.getBody().front();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      // Set up operand mapping for the body region
      for (unsigned i = 0; i < whileOp.getBody().getNumArguments(); ++i) {
        auto oldArg = whileOp.getBody().getArgument(i);
        Value newArg = newWhileOp.getBody().getArgument(i);
        mapper.map(oldArg, newArg);
      }

      for (auto res :
           createConditional(rewriter, whileOp, ivInfo, candidates,
                             newWhileOp.getBody().getArgument(ivInfo.index),
                             newWhileOp.getBody().getArguments(), false)
               ->getResults()) {
        auto idx = candidates[res.getResultNumber()].idx;
        mapper.map(whileOp.getBody().getArgument(idx), res);
      }

      for (auto &op : oldBodyBlock.getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
    }

    // Create a new terminator for the body region using new values
    {
      OpBuilder::InsertionGuard guard(rewriter);
      SmallVector<Value> newReturnValues;

      // Map old return values to new values using the mapper
      for (auto oldRetVal : yieldOp.getOperands()) {
        Value newRetVal = mapper.lookupOrNull(oldRetVal);
        // If the value isn't in the mapper, maybe it was a block argument or
        // constant
        if (!newRetVal)
          newRetVal = oldRetVal; // Consider more robust handling if needed
        newReturnValues.push_back(newRetVal);
      }

      // Create the return op at the end of the body
      rewriter.setInsertionPointToEnd(&newBodyBlock);
      stablehlo::ReturnOp::create(rewriter, yieldOp.getLoc(), newReturnValues);
    }

    // Create condition region mapper
    mlir::IRMapping condMapper;

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    Block &oldCondBlock = whileOp.getCond().front();

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        auto oldArg = whileOp.getCond().getArgument(i);
        Value newArg = newWhileOp.getCond().getArgument(i);
        condMapper.map(oldArg, newArg);
      }

      for (auto res :
           createConditional(rewriter, whileOp, ivInfo, candidates,
                             newWhileOp.getCond().getArgument(ivInfo.index),
                             newWhileOp.getCond().getArguments(), false)
               ->getResults()) {
        auto idx = candidates[res.getResultNumber()].idx;
        condMapper.map(whileOp.getBody().getArgument(idx), res);
      }

      for (auto &op : oldCondBlock.getOperations()) {
        rewriter.clone(op, condMapper);
      }
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, results);
    return success();
  }
};

struct WhilePadInductionReduction
    : public CheckedOpRewritePattern<stablehlo::WhileOp,
                                     WhilePadInductionReduction> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  template <typename CRangeT, typename RangeT>
  static stablehlo::IfOp
  createConditional(PatternRewriter &rewriter, stablehlo::WhileOp whileOp,
                    IVInfo &ivInfo, CRangeT &candidates, Value iv,
                    const RangeT iterOperand, bool cloneCond) {
    SmallVector<Type> ifResultTypes;

    SmallVector<Value> oldReturns;
    for (auto &candidate : candidates) {
      ifResultTypes.push_back(whileOp.getOperands()[candidate.idx].getType());
      oldReturns.push_back(whileOp.getOperands()[candidate.idx]);
    }

    Value condition;

    if (cloneCond) {
      IRMapping mapper;

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        mapper.map(whileOp.getCond().getArgument(i), whileOp.getOperands()[i]);
      }
      for (auto &op : whileOp.getCond().front().getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
      condition = whileOp.getCond().front().getTerminator()->getOperand(0);
      condition = mapper.lookupOrDefault(condition);
    } else {
      condition = stablehlo::CompareOp::create(
          rewriter, whileOp.getLoc(), iv, ivInfo.start,
          stablehlo::ComparisonDirection::EQ);
    }

    auto ifOp = stablehlo::IfOp::create(rewriter, whileOp.getLoc(),
                                        ifResultTypes, condition);

    // Create the then and else regions for the if operation

    {
      Region &thenRegion =
          cloneCond ? ifOp.getFalseBranch() : ifOp.getTrueBranch();
      OpBuilder::InsertionGuard guard(rewriter);
      Block *thenBlock = rewriter.createBlock(&thenRegion);
      rewriter.setInsertionPointToStart(thenBlock);
      stablehlo::ReturnOp::create(rewriter, whileOp.getLoc(), oldReturns);
    }

    {
      Region &elseRegion =
          cloneCond ? ifOp.getTrueBranch() : ifOp.getFalseBranch();
      OpBuilder::InsertionGuard guard(rewriter);
      Block *elseBlock = rewriter.createBlock(&elseRegion);
      rewriter.setInsertionPointToStart(elseBlock);
      SmallVector<Value> newReturns;

      for (auto &candidate : candidates) {
        // Create a pad op using these values
        auto padOp = stablehlo::PadOp::create(
            rewriter, candidate.pad.getLoc(),
            iterOperand[candidate.idx], // operand
            candidate.paddingValue,     // padding_value
            candidate.lowAttr,          // edge_padding_low
            candidate.highAttr,         // edge_padding_high
            candidate.interiorAttr);    // interior_padding
        newReturns.push_back(padOp.getResult());
      }
      stablehlo::ReturnOp::create(rewriter, whileOp.getLoc(), newReturns);
    }

    return ifOp;
  }

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find the index of IV and the step to check for 1 iteration
    auto ivInfo = extractSimpleIVInfo(whileOp);
    if (!ivInfo.isValid)
      return failure();

    if (ivInfo.step == 0)
      return failure();

    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct Candidate {
      unsigned idx;
      stablehlo::PadOp pad;
      Value operand;                  // The operand being padded
      Value paddingValue;             // The padding value
      DenseI64ArrayAttr lowAttr;      // Edge padding low attribute
      DenseI64ArrayAttr highAttr;     // Edge padding high attribute
      DenseI64ArrayAttr interiorAttr; // Interior padding attribute
    };

    llvm::SmallVector<Candidate> candidates;
    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      // check if operand is a pad operation
      auto pad = yieldOp.getOperand(idx).getDefiningOp<stablehlo::PadOp>();
      if (!pad)
        continue;

      if (!definedOutside(pad.getPaddingValue(), whileOp))
        continue;

      // Store original attributes directly
      candidates.emplace_back(
          Candidate{idx, pad, pad.getOperand(), pad.getPaddingValue(),
                    rewriter.getDenseI64ArrayAttr(pad.getEdgePaddingLow()),
                    rewriter.getDenseI64ArrayAttr(pad.getEdgePaddingHigh()),
                    rewriter.getDenseI64ArrayAttr(pad.getInteriorPadding())});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create sliced versions of each padded input
    for (auto &candidate : candidates) {

      // Create an undef op of the same type as the operand
      auto tensorType = cast<TensorType>(candidate.operand.getType());
      auto zeroAttr = DenseElementsAttr::get(
          tensorType, rewriter.getZeroAttr(tensorType.getElementType()));
      auto undef = stablehlo::ConstantOp::create(
          rewriter, candidate.operand.getLoc(), zeroAttr);
      newOperands[candidate.idx] = undef;
    }

    // Update yield op to use the unpadded values
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : candidates) {
        newReturnValues[candidate.idx] = candidate.operand;
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    SmallVector<Value> results;
    for (auto res : newWhileOp.getResults())
      results.push_back(res);

    auto ifOp =
        createConditional(rewriter, whileOp, ivInfo, candidates,
                          whileOp.getOperands()[ivInfo.index], results, true);
    for (auto res : ifOp->getResults()) {
      results[candidates[res.getResultNumber()].idx] = res;
    }

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock.clear();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      SmallVector<Value> replacements;
      for (auto arg : newBodyBlock.getArguments())
        replacements.push_back(arg);
      for (auto res :
           createConditional(rewriter, whileOp, ivInfo, candidates,
                             newWhileOp.getBody().getArgument(ivInfo.index),
                             newWhileOp.getBody().getArguments(), false)
               ->getResults()) {
        auto idx = candidates[res.getResultNumber()].idx;
        replacements[idx] = res;
      }
      rewriter.mergeBlocks(&whileOp.getBody().front(), &newBodyBlock,
                           replacements);
    }

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      SmallVector<Value> replacements;
      for (auto arg : newCondBlock.getArguments())
        replacements.push_back(arg);

      for (auto res :
           createConditional(rewriter, whileOp, ivInfo, candidates,
                             newWhileOp.getCond().getArgument(ivInfo.index),
                             newWhileOp.getCond().getArguments(), false)
               ->getResults()) {
        auto idx = candidates[res.getResultNumber()].idx;
        replacements[idx] = res;
      }

      rewriter.mergeBlocks(&whileOp.getCond().front(), &newCondBlock,
                           replacements);
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, results);
    return success();
  }
};

static bool mayReadMemoryWrittenTo(stablehlo::SliceOp slice,
                                   stablehlo::DynamicUpdateSliceOp DUS) {

  size_t i = 0;
  for (auto &&[sstart, slimit, dstartv] :
       llvm::zip(slice.getStartIndices(), slice.getLimitIndices(),
                 DUS.getStartIndices())) {
    DenseIntElementsAttr startattr;
    if (!matchPattern(dstartv, m_Constant(&startattr))) {
      i++;
      continue;
    }
    int64_t dstart = (*startattr.begin()).getSExtValue();
    int64_t dlimit = dstart + DUS.getUpdate().getType().getShape()[i];
    i++;
    if (slimit <= dstart)
      return false;
    if (dlimit <= sstart)
      return false;
  }
  return true;
}

static bool mustReadMemoryWrittenTo(stablehlo::SliceOp slice,
                                    stablehlo::DynamicUpdateSliceOp DUS) {

  size_t i = 0;
  for (auto &&[sstart, slimit, dstartv] :
       llvm::zip(slice.getStartIndices(), slice.getLimitIndices(),
                 DUS.getStartIndices())) {
    DenseIntElementsAttr startattr;
    if (!matchPattern(dstartv, m_Constant(&startattr))) {
      return false;
    }
    int64_t dstart = (*startattr.begin()).getSExtValue();
    int64_t dlimit = dstart + DUS.getUpdate().getType().getShape()[i];
    i++;
    if (sstart < dstart)
      return false;
    if (slimit > dlimit)
      return false;
  }
  return true;
}

// Returns true if the data returned by this slice could equivalently be written
// as a slice of this other operand
static bool sliceOperandIsEquivalentTo(stablehlo::SliceOp slice,
                                       Value operand) {
  while (true) {
    if (operand == slice.getOperand())
      return true;
    if (auto dus = operand.getDefiningOp<stablehlo::DynamicUpdateSliceOp>()) {
      if (!mayReadMemoryWrittenTo(slice, dus)) {
        operand = dus.getOperand();
        continue;
      }
    }
    return false;
  }
  return false;
}

struct DUSDUSToDUSExtend
    : public CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                                     DUSDUSToDUSExtend> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dus,
                                    PatternRewriter &rewriter) const {

    auto slice = dus.getUpdate().getDefiningOp<stablehlo::SliceOp>();

    if (!slice)
      return failure();
    auto dus2 =
        dus.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();

    if (!dus2)
      return failure();

    if (!sliceOperandIsEquivalentTo(slice, dus2.getOperand()))
      return failure();

    auto slice2 = dus2.getUpdate().getDefiningOp<stablehlo::SliceOp>();

    if (!slice2)
      return failure();

    if (!sliceOperandIsEquivalentTo(slice2, dus2.getOperand()))
      return failure();

    SmallVector<stablehlo::SliceOp> middleSlices;
    for (auto u : dus2.getResult().getUsers()) {
      if (u == dus)
        continue;
      if (auto midslice = dyn_cast<stablehlo::SliceOp>(u)) {
        if (mayReadMemoryWrittenTo(midslice, dus)) {
          return failure();
        }
        middleSlices.push_back(midslice);
        continue;
      }
      return failure();
    }

    RankedTensorType tys[2];
    stablehlo::DynamicUpdateSliceOp duses[2] = {dus, dus2};
    for (auto en : llvm::enumerate(duses)) {
      auto ty = dyn_cast<RankedTensorType>(en.value().getUpdate().getType());
      if (!ty)
        return failure();
      tys[en.index()] = ty;
    }

    if (dus.getOperand().getType() != dus2.getOperand().getType())
      return failure();

    ssize_t diffidx = -1;
    SmallVector<int64_t> starts;
    for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
      DenseIntElementsAttr startattr;
      if (!matchPattern(dus.getStartIndices()[i], m_Constant(&startattr)))
        return failure();
      int64_t ival = (*startattr.begin()).getSExtValue();
      starts.push_back(ival);

      if (dus.getStartIndices()[i] == dus2.getStartIndices()[i]) {
        continue;
      }
      if (diffidx != -1) {
        return failure();
      }
      diffidx = i;
    }

    if (diffidx == -1) {
      for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
        if (tys[0].getShape()[i] == tys[1].getShape()[i])
          continue;
        if (diffidx != -1) {
          return failure();
        }
        diffidx = i;
      }
    }

    if (diffidx == -1) {
      return failure();
    }

    // Sizes must be the same except for the differing index
    for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
      if (slice.getStrides()[i] != 1)
        return failure();
      if (slice2.getStrides()[i] != 1)
        return failure();

      if (i == diffidx)
        continue;
      if (tys[0].getShape()[i] != tys[1].getShape()[i])
        return failure();

      if (slice.getStartIndices()[i] != starts[i])
        return failure();
      if (slice2.getStartIndices()[i] != starts[i])
        return failure();
      if (slice.getType().getShape()[i] != tys[0].getShape()[i])
        return failure();
      if (slice2.getType().getShape()[i] != tys[0].getShape()[i])
        return failure();
    }

    int64_t idxs[2];

    for (auto en : llvm::enumerate(duses)) {
      auto val = en.value().getStartIndices()[diffidx];
      DenseIntElementsAttr startattr;
      if (!matchPattern(val, m_Constant(&startattr))) {
        return failure();
      }
      int64_t ival = (*startattr.begin()).getSExtValue();
      idxs[en.index()] = ival;
    }

    // We have dus(dus2(operand, update2, start=idxs[1]), update1,
    // start=idxs[0])
    //
    // only one index differs, which may differ in start, size, or both
    rewriter.setInsertionPoint(dus2);
    if (idxs[0] + tys[0].getShape()[diffidx] < idxs[1] &&
        slice.getStartIndices()[diffidx] ==
            idxs[0] + tys[0].getShape()[diffidx] &&
        slice2.getLimitIndices()[diffidx] == idxs[1]) {
      // Case 1:
      // [idxs[0]   ... update1 ...  idxs[0] + tys[0].size] .. data .. [idxs[1]
      // .... update2 ...   ]

      auto prestart = starts;
      prestart[diffidx] = idxs[0] + tys[0].getShape()[diffidx];
      auto prelimit = prestart;
      for (size_t i = 0; i < dus.getStartIndices().size(); i++)
        prelimit[i] += tys[0].getShape()[i];
      prelimit[diffidx] = idxs[1];
      auto preslice =
          stablehlo::SliceOp::create(rewriter, dus.getLoc(), dus2.getOperand(),
                                     prestart, prelimit, slice.getStrides());

      auto preextend = enzymexla::ExtendOp::create(
          rewriter, dus.getLoc(), preslice, tys[0].getShape()[diffidx],
          tys[1].getShape()[diffidx], diffidx);

      auto repdus =
          rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
              dus, dus2.getOperand(), preextend, dus.getStartIndices());
      for (auto midslice : middleSlices) {
        rewriter.modifyOpInPlace(
            midslice, [&]() { midslice.getOperandMutable().assign(repdus); });
      }
      rewriter.eraseOp(dus2);
      return success();
    } else if (idxs[1] + tys[1].getShape()[diffidx] < idxs[0] &&
               slice2.getStartIndices()[diffidx] ==
                   idxs[1] + tys[1].getShape()[diffidx] &&
               slice.getLimitIndices()[diffidx] == idxs[0]) {
      // Case 2:
      // [idxs[1]   ... update2 ...  idxs[1] + tys[1].size] .. data .. [idxs[0]
      // .... update1 ...   ]

      auto prestart = starts;
      prestart[diffidx] = idxs[1] + tys[1].getShape()[diffidx];
      auto prelimit = prestart;
      for (size_t i = 0; i < dus.getStartIndices().size(); i++)
        prelimit[i] += tys[0].getShape()[i];
      prelimit[diffidx] = idxs[0];
      auto preslice =
          stablehlo::SliceOp::create(rewriter, dus.getLoc(), dus2.getOperand(),
                                     prestart, prelimit, slice.getStrides());

      auto preextend = enzymexla::ExtendOp::create(
          rewriter, dus.getLoc(), preslice, tys[1].getShape()[diffidx],
          tys[0].getShape()[diffidx], diffidx);

      auto repdus =
          rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
              dus, dus2.getOperand(), preextend, dus2.getStartIndices());
      for (auto midslice : middleSlices) {
        rewriter.modifyOpInPlace(
            midslice, [&]() { midslice.getOperandMutable().assign(repdus); });
      }
      rewriter.eraseOp(dus2);
      return success();
    }

    return failure();
  }
};

struct DUSDUSToDUSPad
    : public CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                                     DUSDUSToDUSPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dus,
                                    PatternRewriter &rewriter) const {

    SplatElementsAttr cst;
    if (!matchPattern(dus.getUpdate(), m_Constant(&cst))) {
      return failure();
    }

    if (!cst)
      return failure();
    auto dus2 =
        dus.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();

    if (!dus2)
      return failure();

    SplatElementsAttr cst2;
    if (!matchPattern(dus2.getUpdate(), m_Constant(&cst2))) {
      return failure();
    }

    if (cst.getSplatValue<Attribute>() != cst2.getSplatValue<Attribute>()) {
      return failure();
    }

    SmallVector<stablehlo::SliceOp> middleSlices;
    for (auto u : dus2.getResult().getUsers()) {
      if (u == dus)
        continue;
      if (auto midslice = dyn_cast<stablehlo::SliceOp>(u)) {
        if (mayReadMemoryWrittenTo(midslice, dus)) {
          return failure();
        }
        middleSlices.push_back(midslice);
        continue;
      }
      return failure();
    }

    RankedTensorType tys[2];
    stablehlo::DynamicUpdateSliceOp duses[2] = {dus, dus2};
    for (auto en : llvm::enumerate(duses)) {
      auto ty = dyn_cast<RankedTensorType>(en.value().getUpdate().getType());
      if (!ty)
        return failure();
      tys[en.index()] = ty;
    }

    if (dus.getOperand().getType() != dus2.getOperand().getType())
      return failure();

    ssize_t diffidx = -1;
    SmallVector<int64_t> starts;
    for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
      DenseIntElementsAttr startattr;
      if (!matchPattern(dus.getStartIndices()[i], m_Constant(&startattr)))
        return failure();
      int64_t ival = (*startattr.begin()).getSExtValue();
      starts.push_back(ival);

      if (dus.getStartIndices()[i] == dus2.getStartIndices()[i]) {
        continue;
      }
      if (diffidx != -1) {
        return failure();
      }
      diffidx = i;
    }

    if (diffidx == -1) {
      for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
        if (tys[0].getShape()[i] == tys[1].getShape()[i])
          continue;
        if (diffidx != -1) {
          return failure();
        }
        diffidx = i;
      }
    }

    if (diffidx == -1) {
      return failure();
    }

    // Sizes must be the same except for the differing index
    for (size_t i = 0; i < dus.getStartIndices().size(); i++) {
      if (i == diffidx)
        continue;
      if (tys[0].getShape()[i] != tys[1].getShape()[i])
        return failure();
    }

    int64_t idxs[2];

    for (auto en : llvm::enumerate(duses)) {
      auto val = en.value().getStartIndices()[diffidx];
      DenseIntElementsAttr startattr;
      if (!matchPattern(val, m_Constant(&startattr))) {
        return failure();
      }
      int64_t ival = (*startattr.begin()).getSExtValue();
      idxs[en.index()] = ival;
    }

    // We have dus(dus2(operand, update2, start=idxs[1]), update1,
    // start=idxs[0])
    //
    // only one index differs, which may differ in start, size, or both
    rewriter.setInsertionPoint(dus2);
    if (idxs[0] + tys[0].getShape()[diffidx] < idxs[1]) {
      // Case 1:
      // [idxs[0]   ... update1 ...  idxs[0] + tys[0].size] .. data .. [idxs[1]
      // .... update2 ...   ]

      auto prestart = starts;
      prestart[diffidx] = idxs[0] + tys[0].getShape()[diffidx];
      auto prelimit = prestart;
      for (size_t i = 0; i < dus.getStartIndices().size(); i++)
        prelimit[i] += tys[0].getShape()[i];
      prelimit[diffidx] = idxs[1];
      SmallVector<int64_t> strides(tys[1].getShape().size(), 1);
      auto preslice =
          stablehlo::SliceOp::create(rewriter, dus.getLoc(), dus2.getOperand(),
                                     prestart, prelimit, strides);

      SmallVector<int64_t> padlow(tys[1].getShape().size(), 0);
      SmallVector<int64_t> padhigh(tys[1].getShape().size(), 0);
      SmallVector<int64_t> padinner(tys[1].getShape().size(), 0);

      padlow[diffidx] = tys[0].getShape()[diffidx];
      padhigh[diffidx] = tys[1].getShape()[diffidx];

      auto newcst = stablehlo::ConstantOp::create(
          rewriter, dus.getLoc(),
          cst.resizeSplat(RankedTensorType::get({}, cst.getElementType())));
      auto preextend = stablehlo::PadOp::create(
          rewriter, dus.getLoc(), preslice, newcst, padlow, padhigh, padinner);

      auto repdus =
          rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
              dus, dus2.getOperand(), preextend, dus.getStartIndices());
      for (auto midslice : middleSlices) {
        rewriter.modifyOpInPlace(
            midslice, [&]() { midslice.getOperandMutable().assign(repdus); });
      }
      rewriter.eraseOp(dus2);
      return success();
    } else if (idxs[1] + tys[1].getShape()[diffidx] < idxs[0]) {
      // Case 2:
      // [idxs[1]   ... update2 ...  idxs[1] + tys[1].size] .. data .. [idxs[0]
      // .... update1 ...   ]

      auto prestart = starts;
      prestart[diffidx] = idxs[1] + tys[1].getShape()[diffidx];
      auto prelimit = prestart;
      for (size_t i = 0; i < dus.getStartIndices().size(); i++)
        prelimit[i] += tys[0].getShape()[i];
      prelimit[diffidx] = idxs[0];
      SmallVector<int64_t> strides(tys[1].getShape().size(), 1);
      auto preslice =
          stablehlo::SliceOp::create(rewriter, dus.getLoc(), dus2.getOperand(),
                                     prestart, prelimit, strides);

      SmallVector<int64_t> padlow(tys[1].getShape().size(), 0);
      SmallVector<int64_t> padhigh(tys[1].getShape().size(), 0);
      SmallVector<int64_t> padinner(tys[1].getShape().size(), 0);

      padlow[diffidx] = tys[1].getShape()[diffidx];
      padhigh[diffidx] = tys[0].getShape()[diffidx];

      auto newcst = stablehlo::ConstantOp::create(
          rewriter, dus.getLoc(),
          cst.resizeSplat(RankedTensorType::get({}, cst.getElementType())));
      auto preextend = stablehlo::PadOp::create(
          rewriter, dus.getLoc(), preslice, newcst, padlow, padhigh, padinner);

      auto repdus =
          rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
              dus, dus2.getOperand(), preextend, dus2.getStartIndices());
      for (auto midslice : middleSlices) {
        rewriter.modifyOpInPlace(
            midslice, [&]() { midslice.getOperandMutable().assign(repdus); });
      }
      rewriter.eraseOp(dus2);
      return success();
    }

    return failure();
  }
};

// Information about the provenance of a tensor value. Contains a potentially
// disjoint relation mapping tensor subscripts to source tensor values and
// subscripts in it. The first dimension of the relation range is a position in
// the `sources` vector pointing to the source value, the remaining dimensions
// are subscripts. This currently assumes that the rank of the tensor for which
// the provenance is computed matches the rank of all source tensors, except for
// broadcasted constants. This means expand/collapse operations cannot be
// supported. This is fundamentally possible using Presburger relations, but
// requires handling a disjoint relation with potentially different number of
// range dimensions, which isn't directly supported by the Presburger library.
// (It is in theory possible to use the maximum of all range ranks and pin the
// unnecessary trailing dimensions to zero).
struct TensorValueProvenance {
  SmallVector<Value> sources;
  presburger::PresburgerRelation provenanceRelation;

  TensorValueProvenance(ArrayRef<Value> sources,
                        const presburger::PresburgerRelation &relation)
      : sources(sources.begin(), sources.end()), provenanceRelation(relation) {}
};

// Returns an integer set representing a hyperrectangle with the given start
// coordiantes and the given size.
static presburger::IntegerPolyhedron
getBoundedDomain(const presburger::PresburgerSpace &space,
                 ArrayRef<int64_t> starts, ArrayRef<int64_t> extents) {
  auto domain = presburger::IntegerPolyhedron::getUniverse(space);
  size_t rank = starts.size();
  for (size_t i = 0; i < rank; ++i) {
    domain.addBound(presburger::BoundType::LB, i, starts[i]);
    domain.addBound(presburger::BoundType::UB, i, starts[i] + extents[i] - 1);
  }
  return domain;
}

// Projects out local dimensions from the given relation. This is useful since
// the Presburger library, unlike isl, does not automatically remove original
// dimensions when a relation is applied to the domain or range of another
// relation.
static void
presburgerRelationProjectOutLocals(presburger::PresburgerRelation &relation) {
  auto updated = presburger::PresburgerRelation::getEmpty(
      relation.getSpace().getSpaceWithoutLocals());
  for (auto disjunct : relation.getAllDisjuncts()) {
    disjunct.projectOut(disjunct.getNumDimAndSymbolVars(),
                        disjunct.getNumLocalVars());
    updated.unionInPlace(disjunct);
  }
  relation = updated;
}

// Update the relation so that the indices of source values (leading range
// dimension) provided as `currentSources` match their positions in
// `allSources`, appending missing values when necessary.
static void alignSourceValues(presburger::PresburgerRelation &relation,
                              SmallVectorImpl<Value> &allSources,
                              ArrayRef<Value> currentSources) {
  if (allSources == currentSources)
    return;

  presburger::PresburgerSpace rangeSpace = relation.getSpace().getRangeSpace();
  size_t rank = rangeSpace.getNumSetDimVars() - 1;
  auto remappingSpace = presburger::PresburgerSpace::getRelationSpace(
      rangeSpace.getNumSetDimVars(), rangeSpace.getNumSetDimVars(),
      rangeSpace.getNumSymbolVars());
  presburger::PresburgerRelation remapping =
      presburger::PresburgerRelation::getEmpty(remappingSpace);
  for (auto [i, source] : llvm::enumerate(currentSources)) {
    auto it = llvm::find(allSources, source);
    size_t position;
    if (it == allSources.end()) {
      position = allSources.size();
      allSources.push_back(source);
    } else {
      position = std::distance(allSources.begin(), it);
    }

    auto remappingElement =
        presburger::IntegerRelation::getUniverse(remappingSpace);
    remappingElement.addBound(presburger::BoundType::EQ, 0, i);
    remappingElement.addBound(presburger::BoundType::EQ, rank + 1, position);
    for (size_t dim = 0; dim < rank; ++dim) {
      SmallVector<int64_t> equalityCoefficients(2 * rank + 3, 0);
      equalityCoefficients[dim + 1] = -1;
      equalityCoefficients[rank + dim + 2] = 1;
      remappingElement.addEquality(equalityCoefficients);
    }
    remapping.unionInPlace(remappingElement);
  }
  relation.applyRange(remapping);
  presburgerRelationProjectOutLocals(relation);
}

// Create an integer relation mapping dimensions in the given space to the
// co-indexed positions with the given offsets. That is, x_i is mapped to x_i +
// offset_i. If `positive` is not set, map to x_i - offset_i instead.
static presburger::IntegerRelation
getOffsetRelation(const presburger::PresburgerSpace &space,
                  ArrayRef<int64_t> offsets, bool positive = true) {
  assert(space.getNumSetDimVars() == offsets.size());
  auto offsetSpace = presburger::PresburgerSpace::getRelationSpace(
      space.getNumSetDimVars(), space.getNumSetDimVars(),
      space.getNumSymbolVars());
  auto offset = presburger::IntegerRelation::getUniverse(offsetSpace);
  size_t rank = offsets.size();
  for (size_t i = 0; i < rank; ++i) {
    SmallVector<int64_t> offsetCoefficients(2 * rank + 1, 0);
    offsetCoefficients[i] = -1;
    offsetCoefficients[rank + i] = 1;
    offsetCoefficients.back() = (positive ? 1 : -1) * offsets[i];
    offset.addEquality(offsetCoefficients);
  }
  return offset;
}

// Mark the value as the provenance source by indicating it as its own source
// for all valid subscripts.
static void
addSelfProvenance(Value value,
                  DenseMap<Value, TensorValueProvenance> &provenanceInfo) {
  assert(!provenanceInfo.contains(value));

  size_t rank = cast<RankedTensorType>(value.getType()).getRank();

  presburger::PresburgerSpace space =
      presburger::PresburgerSpace::getRelationSpace(
          /*numDomainVars=*/rank,
          /*numRangeVars=*/rank + 1,
          /*numSymbolVars=*/0);
  auto relation = presburger::IntegerRelation::getUniverse(space);
  for (size_t i = 0; i < rank; ++i) {
    SmallVector<int64_t> equalityCoefficients(2 * rank + 2, 0);
    equalityCoefficients[i] = -1;
    equalityCoefficients[rank + i + 1] = 1;
    relation.addEquality(equalityCoefficients);

    relation.addBound(presburger::BoundType::LB, i, 0);
    relation.addBound(presburger::BoundType::UB, i,
                      cast<RankedTensorType>(value.getType()).getShape()[i] -
                          1);
  }
  SmallVector<int64_t> sourceIndexCoefficients(2 * rank + 2, 0);
  sourceIndexCoefficients[rank] = -1;
  sourceIndexCoefficients.back() = 0;
  relation.addEquality(sourceIndexCoefficients);

  // Sanity check that all elements are present in the mapping.
  if (auto volume = relation.getDomainSet().computeVolume()) {
    assert(*volume == cast<RankedTensorType>(value.getType()).getNumElements());
  }

  provenanceInfo.try_emplace(value, ArrayRef<Value>(value),
                             presburger::PresburgerRelation(relation));
}

// Compute provenance for all given values and popupate the provenance info data
// structure with it.
void computeTensorValueProvenanceImpl(
    SmallVectorImpl<Value> &pendingValues,
    DenseMap<Value, TensorValueProvenance> &provenanceInfo2) {
  while (!pendingValues.empty()) {
    Value value = pendingValues.pop_back_val();
    if (provenanceInfo2.contains(value))
      continue;

    if (auto slice = value.getDefiningOp<stablehlo::SliceOp>()) {
      // If we don't know the provenance of the operand, figure it out before
      // coming back to this value.
      auto it = provenanceInfo2.find(slice.getOperand());
      if (it == provenanceInfo2.end()) {
        pendingValues.push_back(value);
        pendingValues.push_back(slice.getOperand());
        continue;
      }

      // Intersect domain with that defined by the slice, than shift by offsets
      // (start indices).
      presburger::PresburgerSpace domainSpace =
          it->second.provenanceRelation.getSpace().getDomainSpace();
      presburger::IntegerPolyhedron restriction = getBoundedDomain(
          domainSpace, slice.getStartIndices(), slice.getType().getShape());
      presburger::IntegerRelation offset =
          getOffsetRelation(domainSpace, slice.getStartIndices());

      auto [insertIt, _] = provenanceInfo2.try_emplace(
          value, it->second.sources, it->second.provenanceRelation);
      insertIt->second.provenanceRelation =
          insertIt->second.provenanceRelation.intersectDomain(
              presburger::PresburgerSet(restriction));
      insertIt->second.provenanceRelation.applyDomain(
          presburger::PresburgerRelation(offset));
      presburgerRelationProjectOutLocals(insertIt->second.provenanceRelation);
      insertIt->second.provenanceRelation =
          insertIt->second.provenanceRelation.simplify().coalesce();

      if (std::optional<DynamicAPInt> volume =
              insertIt->second.provenanceRelation.getDomainSet()
                  .computeVolume()) {
        assert(*volume ==
               cast<RankedTensorType>(value.getType()).getNumElements());
      }
      continue;
    }

    if (auto dus = value.getDefiningOp<stablehlo::DynamicUpdateSliceOp>()) {
      auto it = provenanceInfo2.find(dus.getOperand());
      if (it == provenanceInfo2.end()) {
        pendingValues.push_back(value);
        pendingValues.push_back(dus.getOperand());
        continue;
      }

      auto up = provenanceInfo2.find(dus.getUpdate());
      if (up == provenanceInfo2.end()) {
        pendingValues.push_back(value);
        pendingValues.push_back(dus.getUpdate());
        continue;
      }

      auto resultShape = cast<RankedTensorType>(dus.getType()).getShape();
      auto updateShape =
          cast<RankedTensorType>(dus.getUpdate().getType()).getShape();
      size_t rank = resultShape.size();

      // Try to extract constant start indices for each dimension
      SmallVector<std::optional<int64_t>> constantStarts(rank);
      for (size_t i = 0; i < rank; i++) {
        DenseIntElementsAttr startAttr;
        if (matchPattern(dus.getStartIndices()[i], m_Constant(&startAttr))) {
          constantStarts[i] = (*startAttr.begin()).getSExtValue();
        }
      }

      // Check if all start indices are constant (required for provenance
      // tracking)
      bool allConstant = true;
      for (const auto &start : constantStarts) {
        if (!start.has_value()) {
          allConstant = false;
          break;
        }
      }

      if (!allConstant) {
        // If start indices are dynamic, we can't precisely track provenance.
        // Assume constant propagation has run before this.
        addSelfProvenance(value, provenanceInfo2);
        continue;
      }

      auto knownConstantStarts = llvm::map_to_vector(
          constantStarts, [](std::optional<int64_t> v) { return v.value(); });

      // Update provenance to that of the DUS update operand for the slice.
      presburger::PresburgerRelation provenanceUnion(
          it->second.provenanceRelation);
      presburger::PresburgerSet provenanceRemainingDomain =
          provenanceUnion.getDomainSet();
      presburgerRelationProjectOutLocals(provenanceRemainingDomain);
      provenanceRemainingDomain =
          provenanceRemainingDomain.subtract(presburger::PresburgerRelation(
              getBoundedDomain(provenanceUnion.getSpace().getDomainSpace(),
                               knownConstantStarts, updateShape)));
      provenanceUnion =
          provenanceUnion.intersectDomain(provenanceRemainingDomain);

      presburger::PresburgerRelation update = up->second.provenanceRelation;
      update.applyDomain(presburger::PresburgerRelation(getOffsetRelation(
          update.getSpace().getDomainSpace(), knownConstantStarts, false)));
      presburgerRelationProjectOutLocals(update);

      SmallVector<Value> allSources = it->second.sources;
      alignSourceValues(update, allSources, up->second.sources);
      provenanceUnion.unionInPlace(update);
      auto [insertIt, _] = provenanceInfo2.try_emplace(
          value, allSources, provenanceUnion.simplify().coalesce());

      // Sanity check that all elements remain covered.
      if (std::optional<DynamicAPInt> volume =
              insertIt->second.provenanceRelation.getDomainSet()
                  .computeVolume()) {
        assert(*volume ==
               cast<RankedTensorType>(value.getType()).getNumElements());
      }
      continue;
    }

    if (auto pad = value.getDefiningOp<stablehlo::PadOp>()) {
      auto it = provenanceInfo2.find(pad.getOperand());
      if (it == provenanceInfo2.end()) {
        pendingValues.push_back(value);
        pendingValues.push_back(pad.getOperand());
        continue;
      }

      auto operandType = cast<RankedTensorType>(pad.getOperand().getType());
      auto resultType = cast<RankedTensorType>(pad.getResult().getType());

      if (!operandType.hasStaticShape() || !resultType.hasStaticShape()) {
        addSelfProvenance(value, provenanceInfo2);
        continue;
      }

      size_t rank = resultType.getRank();

      // 1. Create a relation for the entire output marking the padding constant
      // as source.
      SmallVector<Value> allSources = it->second.sources;
      auto sourcesIt = llvm::find(allSources, pad.getPaddingValue());
      size_t position;
      if (sourcesIt == allSources.end()) {
        position = allSources.size();
        allSources.push_back(pad.getPaddingValue());
      } else {
        position = std::distance(allSources.begin(), sourcesIt);
      }

      presburger::PresburgerSpace paddingSpace =
          presburger::PresburgerSpace::getRelationSpace(
              /*numDomainVars=*/rank,
              /*numRangeVars=*/rank + 1,
              /*numSymbolVars=*/0);
      auto paddingRelation =
          presburger::IntegerRelation::getUniverse(paddingSpace);

      // Add bounds for the entire output tensor.
      for (size_t i = 0; i < rank; ++i) {
        paddingRelation.addBound(presburger::BoundType::LB, i, 0);
        paddingRelation.addBound(presburger::BoundType::UB, i,
                                 resultType.getShape()[i] - 1);
      }

      // Set target offsets to zero as we are broadcasting.
      for (size_t i = 0; i < rank; ++i) {
        paddingRelation.addBound(presburger::BoundType::EQ, rank + 1 + i, 0);
      }

      // Mark source as the padding value.
      SmallVector<int64_t> paddingSourceCoefficients(2 * rank + 2, 0);
      paddingSourceCoefficients[rank] = -1;
      paddingSourceCoefficients.back() = position;
      paddingRelation.addEquality(paddingSourceCoefficients);

      // 2. Subtract the central region from it.
      presburger::PresburgerSpace centralSpace =
          paddingRelation.getSpace().getDomainSpace();
      presburger::IntegerPolyhedron centralRegion = getBoundedDomain(
          centralSpace, pad.getEdgePaddingLow(), operandType.getShape());

      auto [insertedIt, _] = provenanceInfo2.try_emplace(
          value, allSources, presburger::PresburgerRelation(paddingRelation));
      insertedIt->second.provenanceRelation =
          insertedIt->second.provenanceRelation.intersectDomain(
              insertedIt->second.provenanceRelation.getDomainSet().subtract(
                  presburger::PresburgerSet(centralRegion)));

      insertedIt->second.provenanceRelation =
          insertedIt->second.provenanceRelation.simplify();
      presburgerRelationProjectOutLocals(insertedIt->second.provenanceRelation);

      // 3. Add the provenance from the operand for the central region, shifted
      // by the low padding amounts. No need to adjust positions since we
      // accounted for it above, the original indexing is fine.
      presburger::PresburgerRelation operandProvenance =
          it->second.provenanceRelation;
      operandProvenance.applyDomain(presburger::PresburgerRelation(
          getOffsetRelation(operandProvenance.getSpace().getDomainSpace(),
                            pad.getEdgePaddingLow(),
                            /*positive=*/false)));
      presburgerRelationProjectOutLocals(operandProvenance);

      // Union the padding and operand provenances
      insertedIt->second.provenanceRelation.unionInPlace(operandProvenance);
      insertedIt->second.provenanceRelation =
          insertedIt->second.provenanceRelation.simplify().coalesce();

      if (std::optional<DynamicAPInt> volume =
              insertedIt->second.provenanceRelation.getDomainSet()
                  .computeVolume()) {
        assert(*volume ==
               cast<RankedTensorType>(value.getType()).getNumElements());
      }
      continue;
    }

    // Base case: no further provenance
    addSelfProvenance(value, provenanceInfo2);
  }
}

static Value
rematerializeByPiecewiseSelect(RewriterBase &rewriter, Location loc,
                               ArrayRef<int64_t> expectedShape,
                               const TensorValueProvenance &provenance) {
  unsigned rank = expectedShape.size();
  llvm::MapVector<int64_t, SmallVector<int64_t>> sourcePosToDisjunctIdx;
  SmallVector<Attribute> boxes;
  boxes.reserve(provenance.provenanceRelation.getNumDisjuncts());
  for (auto &&[i, disjunct] :
       llvm::enumerate(provenance.provenanceRelation.getAllDisjuncts())) {
    std::optional<DynamicAPInt> sourcePos =
        disjunct.getConstantBound(presburger::BoundType::EQ, rank);
    if (!sourcePos.has_value()) {
      return nullptr;
    }

    Value source = provenance.sources[static_cast<int64_t>(*sourcePos)];
    auto sourceType = cast<RankedTensorType>(source.getType());
    if (sourceType.getRank() != 0 && sourceType.getShape() != expectedShape)
      return nullptr;

    // If the source is not a broadcasted scalar, we need to check that it the
    // mapping is identity, otherwise we cannot simply update the value.
    // TODO: we can see whether we can then reconstruct it differently, e.g., if
    // it is a shifted equality, we can emit a slice. The general problem is
    // similar to "opening the polyhedral compiler black box" paper.
    if (sourceType.getRank() != 0) {
      auto identityRelation = presburger::IntegerRelation::getUniverse(
          disjunct.getSpace().getSpaceWithoutLocals());
      for (unsigned i = 0; i < rank; ++i) {
        SmallVector<int64_t> coefficients(2 * rank + 2, 0);
        coefficients[i] = -1;
        coefficients[rank + i + 1] = 1;
        identityRelation.addEquality(coefficients);
      }
      SmallVector<int64_t> coefficients(2 * rank + 2, 0);
      coefficients[rank] = -1;
      coefficients.back() = static_cast<int64_t>(*sourcePos);
      identityRelation.addEquality(coefficients);
      // The identity relation is not bounded unlike the disjunct, so we check
      // for subset rather than full equality. It is also cheaper to do so.
      if (!disjunct.isSubsetOf(identityRelation)) {
        return nullptr;
      }
    }

    sourcePosToDisjunctIdx[static_cast<int64_t>(*sourcePos)].push_back(i);

    SmallVector<int64_t> box;
    box.reserve(2 * rank);
    for (unsigned i = 0; i < rank; ++i) {
      auto lb = disjunct.getConstantBound(presburger::BoundType::LB, i);
      auto ub = disjunct.getConstantBound(presburger::BoundType::UB, i);
      if (!lb.has_value() || !ub.has_value()) {
        return nullptr;
      }
      box.push_back(static_cast<int64_t>(*lb));
      box.push_back(static_cast<int64_t>(*ub) + 1);
    }
    boxes.push_back(rewriter.getI64ArrayAttr(box));
  }

  auto materializeSource = [&](int64_t sourcePos) -> Value {
    Value source = provenance.sources[sourcePos];
    if (cast<RankedTensorType>(source.getType()).getRank() != 0)
      return source;
    return stablehlo::BroadcastOp::create(rewriter, loc, source, expectedShape);
  };

  Value current;
  for (auto &&[sourcePos, disjuncts] : sourcePosToDisjunctIdx) {
    Value source = materializeSource(sourcePos);
    if (!current) {
      current = source;
      continue;
    }
    SmallVector<Attribute> sourceBoxes;
    for (int64_t i = 0, e = boxes.size(); i < e; ++i) {
      if (llvm::is_contained(disjuncts, i)) {
        sourceBoxes.push_back(boxes[i]);
      }
    }
    current = enzymexla::PiecewiseSelectOp::create(
        rewriter, loc, current.getType(), source, current,
        rewriter.getArrayAttr(sourceBoxes));
  }
  return current;
}

// Materialize a value given its provenance by emitting concat operations to
// create a full value from individual chunks. It is under responsibility of the
// caller to ensure that all values are visible at the insertion point of the
// rewriter. Assumes that provenance disjuncts are adjacent hyperrectangles.
// NOTE: Presburger relation does not support coalescing beyond trivial
// overlaps. This function implements a simple coalescing algorithm for adjacent
// hyperrectangles.
static Value rematerializeByConcat(RewriterBase &rewriter, Location loc,
                                   const TensorValueProvenance &provenance) {
  auto materializeValueForProvenanceDisjunct =
      [](presburger::IntegerRelation disjunct, ArrayRef<Value> sources,
         Location loc, RewriterBase &rewriter) -> Value {
    std::optional<DynamicAPInt> sourcePos = disjunct.getConstantBound(
        presburger::BoundType::EQ, disjunct.getNumDomainVars());
    assert(sourcePos.has_value());
    Value source = sources[static_cast<int64_t>(*sourcePos)];

    unsigned rank = disjunct.getNumDomainVars();
    // If the source is a constant, create a broadcast.
    if (auto constantOp = source.getDefiningOp<stablehlo::ConstantOp>()) {
      assert(constantOp.getResult().getType().getNumElements() == 1);
      SmallVector<int64_t> shape;
      for (size_t i = 0; i < rank; ++i) {
        auto lb = disjunct.getConstantBound(presburger::BoundType::LB, i);
        auto ub = disjunct.getConstantBound(presburger::BoundType::UB, i);
        assert(lb.has_value() && ub.has_value());
        shape.push_back(static_cast<int64_t>((*ub - *lb) + 1));
      }
      return stablehlo::BroadcastOp::create(rewriter, loc, constantOp, shape);
    }

    // Otherwise create a slice.
    SmallVector<int64_t> starts, limits, shape;
    SmallVector<int64_t> strides(rank, 1);
    for (size_t i = 0; i < rank; ++i) {
      auto lb =
          disjunct.getConstantBound(presburger::BoundType::LB, rank + i + 1);
      auto ub =
          disjunct.getConstantBound(presburger::BoundType::UB, rank + i + 1);
      assert(lb.has_value() && ub.has_value());
      starts.push_back(static_cast<int64_t>(*lb));
      limits.push_back(static_cast<int64_t>((*ub) + 1));
      shape.push_back(static_cast<int64_t>((*ub - *lb) + 1));
    }
    return stablehlo::SliceOp::create(rewriter, loc, source, starts, limits,
                                      strides);
  };

  llvm::MapVector<Value, presburger::IntegerPolyhedron> rematerializedValues;
  for (const auto &disjunct : provenance.provenanceRelation.getAllDisjuncts()) {
    auto domain = disjunct.getDomainSet();
    domain.projectOut(domain.getNumDimAndSymbolVars(),
                      domain.getNumLocalVars());
    rematerializedValues.try_emplace(
        materializeValueForProvenanceDisjunct(disjunct, provenance.sources, loc,
                                              rewriter),
        domain);
  }

  while (rematerializedValues.size() != 1) {
    bool found = false;
    auto array = rematerializedValues.getArrayRef();
    for (size_t i = 0, e = array.size(); i < e; ++i) {
      for (size_t j = i + 1; j < e; ++j) {
        auto &&[left, leftDisjunct] = array[i];
        auto &&[right, rightDisjunct] = array[j];
        if (left == right)
          continue;

        // coalescing is not sufficiently powerful, it only handles subsets but
        // not adjacencies

        // Find the dimension along which they are adjacent.
        size_t concatDim = size_t(-1);
        std::optional<bool> concatLeft = std::nullopt;
        SmallVector<int64_t> lowerBounds, extents;

        for (unsigned j = 0, e = leftDisjunct.getNumRangeVars(); j < e; ++j) {
          auto leftLb =
              leftDisjunct.getConstantBound(presburger::BoundType::LB, j);
          auto leftUb =
              leftDisjunct.getConstantBound(presburger::BoundType::UB, j);
          auto rightLb =
              rightDisjunct.getConstantBound(presburger::BoundType::LB, j);
          auto rightUb =
              rightDisjunct.getConstantBound(presburger::BoundType::UB, j);
          assert(leftLb.has_value() && leftUb.has_value() &&
                 rightLb.has_value() && rightUb.has_value());

          // This dimension remains the same.
          if (*leftLb == *rightLb && *leftUb == *rightUb) {
            lowerBounds.push_back(static_cast<int64_t>(*leftLb));
            extents.push_back(static_cast<int64_t>((*leftUb - *leftLb) + 1));
            continue;
          }

          // If we already found a concat dimension, cannot have another one.
          if (concatLeft) {
            concatLeft = std::nullopt;
            break;
          }

          // Check if they are adjacent along this dimension and concat.
          if (*leftUb + 1 == *rightLb) {
            concatLeft = true;
            concatDim = j;
            lowerBounds.push_back(static_cast<int64_t>(*leftLb));
            extents.push_back(static_cast<int64_t>((*rightUb - *leftLb) + 1));
            continue;
          } else if (*rightUb + 1 == *leftLb) {
            concatLeft = false;
            concatDim = j;
            lowerBounds.push_back(static_cast<int64_t>(*rightLb));
            extents.push_back(static_cast<int64_t>((*leftUb - *rightLb) + 1));
            continue;
          }

          // If a dimension is neither equal nor concat, fail.
          concatLeft = std::nullopt;
          break;
        }
        if (!concatLeft)
          continue;

        // TODO: assert that the product of extents corresponds to the sum of
        // volumes of the domain to merge as a proxy of them being
        // hyper-rectangular.
        assert(concatDim != size_t(-1));

        Value joined = stablehlo::ConcatenateOp::create(
            rewriter, loc,
            {*concatLeft ? left : right, *concatLeft ? right : left},
            concatDim);

        rematerializedValues.erase(right);
        rematerializedValues.erase(left);
        rematerializedValues.try_emplace(
            joined,
            getBoundedDomain(leftDisjunct.getSpace(), lowerBounds, extents));

        found = true;
        break;
      }

      if (found)
        break;
    }
    assert(found);
  }
  return rematerializedValues.begin()->first;
}

struct DUSDUSSubsuming
    : public CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                                     DUSDUSSubsuming> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dus,
                                    PatternRewriter &rewriter) const {
    auto dus2 =
        dus.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();

    if (!dus2)
      return failure();

    SmallVector<Value> dusResults = {dus2.getResult(), dus.getResult()};
    DenseMap<Value, TensorValueProvenance> provenanceInfo;
    computeTensorValueProvenanceImpl(dusResults, provenanceInfo);

    DominanceInfo domInfo;
    llvm::MapVector<Operation *, Operation *> movedSlices;
    SmallVector<Operation *> originalUsers =
        llvm::to_vector(dus2.getResult().getUsers());
    for (Operation *user : originalUsers) {
      if (user == dus)
        continue;
      auto slice = dyn_cast<stablehlo::SliceOp>(user);
      if (!slice)
        continue;

      bool isSliceMovable = true;
      // Find if there is any user of the slice that would need to be moved.
      for (Operation *sliceUser : slice.getResult().getUsers()) {
        if (!domInfo.dominates(sliceUser, dus))
          continue;

        isSliceMovable = false;
        break;
      }
      if (!isSliceMovable)
        continue;

      IRMapping mapping;
      rewriter.setInsertionPointAfter(dus);
      mapping.map(dus2.getResult(), dus.getResult());
      Operation *clonedSlice = rewriter.clone(*slice, mapping);

      SmallVector<Value> slices = {slice.getResult(),
                                   clonedSlice->getResult(0)};
      computeTensorValueProvenanceImpl(slices, provenanceInfo);
      auto it = provenanceInfo.find(slice.getResult());
      assert(it != provenanceInfo.end());
      auto it2 = provenanceInfo.find(clonedSlice->getResult(0));
      assert(it2 != provenanceInfo.end());

      // Check that provenance of the sliced part would match its provenance
      // after move. if provenance matches, we can replace the old slice with
      // the new one, otherwise we just erase the new one
      presburger::PresburgerRelation originalProvenance =
          it->second.provenanceRelation;
      SmallVector<Value> originalSources = it->second.sources;
      alignSourceValues(originalProvenance, originalSources,
                        it2->second.sources);

      if (originalSources == it2->second.sources &&
          originalProvenance.isEqual(it2->second.provenanceRelation)) {
        movedSlices.insert({slice.getOperation(), clonedSlice});
      } else {
        rewriter.eraseOp(clonedSlice);

        // Don't forget to erase the provenance info for the op result we just
        // erased since that value address may be reused for something entirely
        // different!
        provenanceInfo.erase(clonedSlice->getResult(0));
      }
    }

    bool didSomething = false;
    for (auto [oldSliceOp, newSliceOp] : movedSlices) {
      rewriter.replaceOp(oldSliceOp, newSliceOp->getResult(0));
      didSomething = true;
    }

    auto dusProvenance = provenanceInfo.find(dus.getResult())->second;
    rewriter.setInsertionPoint(dus);
    Value replacement = rematerializeByPiecewiseSelect(
        rewriter, dus.getLoc(), dus.getResult().getType().getShape(),
        dusProvenance);
    if (replacement && replacement != dus) {
      rewriter.replaceOp(dus, replacement);
      didSomething = true;
      // TODO: this is a WORKING fallback to create a lot of concats
      // } else {
      //   Value rematerializedValue =
      //       rematerializeByConcat(rewriter, dus.getLoc(), dusProvenance);
      //   rewriter.replaceOp(dus, rematerializedValue);
      //   didSomething = true;
    }
    return success(didSomething);
  }
};

struct SliceBroadcast2
    : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceBroadcast2> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp sliceOp,
                                    PatternRewriter &rewriter) const {
    auto broadcastOp =
        sliceOp.getOperand().getDefiningOp<stablehlo::BroadcastOp>();
    if (!broadcastOp ||
        broadcastOp.getOperand().getType().getNumElements() != 1)
      return failure();

    SmallVector<int64_t> shape;
    for (auto [s, l, str] :
         llvm::zip_equal(sliceOp.getStartIndices(), sliceOp.getLimitIndices(),
                         sliceOp.getStrides())) {
      shape.push_back(l - s);
      if (str != 1)
        return failure();
    }
    auto newBroadcast = stablehlo::BroadcastOp::create(
        rewriter, sliceOp.getLoc(), broadcastOp.getOperand(), shape);
    rewriter.replaceOp(sliceOp, newBroadcast.getResult());
    return success();
  }
};

struct ConcatBroadcastToPad
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     ConcatBroadcastToPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concatOp,
                                    PatternRewriter &rewriter) const {
    unsigned axis = concatOp.getDimension();
    auto leftBroadcast =
        concatOp->getOperand(0).getDefiningOp<stablehlo::BroadcastOp>();
    auto rightBroadcast = concatOp->getOperand(concatOp->getNumOperands() - 1)
                              .getDefiningOp<stablehlo::BroadcastOp>();
    if (!leftBroadcast && !rightBroadcast)
      return failure();

    unsigned rank = concatOp.getResult().getType().getRank();
    SmallVector<int64_t> leftPad(rank, 0), rightPad(rank, 0),
        interiorPad(rank, 0);
    if (rightBroadcast && leftBroadcast &&
        rightBroadcast.getOperand() != leftBroadcast.getOperand()) {
      rightBroadcast = nullptr;
    }
    Value broadcastedValue;
    if (leftBroadcast) {
      leftPad[axis] = leftBroadcast.getResult().getType().getShape()[axis];
      broadcastedValue = leftBroadcast.getOperand();
    }
    if (rightBroadcast) {
      rightPad[axis] = rightBroadcast.getResult().getType().getShape()[axis];
      broadcastedValue = rightBroadcast.getOperand();
    }

    if (leftBroadcast && rightBroadcast && concatOp->getNumOperands() == 2) {
      // Just create a bigger broadcast.
      stablehlo::BroadcastOp::create(rewriter, concatOp.getLoc(),
                                     broadcastedValue,
                                     concatOp.getResult().getType().getShape());
      rewriter.replaceOp(concatOp, leftBroadcast.getResult());
      return success();
    }

    Value middle;
    unsigned start = leftBroadcast ? 1 : 0;
    unsigned end = rightBroadcast ? concatOp.getNumOperands() - 1
                                  : concatOp.getNumOperands();
    if (end - start == 1) {
      middle = concatOp->getOperand(start);
    } else {
      middle = stablehlo::ConcatenateOp::create(
                   rewriter, concatOp.getLoc(),
                   concatOp.getOperands().slice(start, end - start), axis)
                   .getResult();
    }

    auto pad = stablehlo::PadOp::create(rewriter, concatOp.getLoc(), middle,
                                        broadcastedValue, leftPad, rightPad,
                                        interiorPad);
    rewriter.replaceOp(concatOp, pad.getResult());
    return success();
  }
};

struct SinkDUS : public CheckedOpRewritePattern<stablehlo::WhileOp, SinkDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct Candidate {
      unsigned idx;
      stablehlo::DynamicUpdateSliceOp DUS;
      SmallVector<stablehlo::SliceOp> mustReaders;
      SmallVector<stablehlo::SliceOp> mayReaders;
    };

    llvm::SmallVector<Candidate, 4> candidates;
    bool anyReader = false;
    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      bool legal = true;
      // Skip DUS candidates which can be removed in a better form by WhileDUS
      auto DUS = yieldOp.getOperand(idx)
                     .getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
      if (!DUS)
        continue;
      if (!definedOutside(DUS.getUpdate(), whileOp))
        continue;

      for (auto s : DUS.getStartIndices()) {
        if (!definedOutside(s, whileOp)) {
          legal = false;
        }
      }
      if (!legal)
        continue;

      auto argOperand = whileOp.getBody().getArgument(idx);
      auto condOperand = whileOp.getCond().getArgument(idx);

      auto T = cast<RankedTensorType>(argOperand.getType());
      auto rank = T.getShape().size();

      if (rank == 0)
        continue;

      SmallVector<Value> todo = {argOperand, condOperand};
      bool selfYield = false;

      SmallVector<stablehlo::SliceOp> mayReaders;
      SmallVector<stablehlo::SliceOp> mustReaders;

      while (!todo.empty()) {
        auto cur = todo.pop_back_val();
        for (auto &u : cur.getUses()) {
          Operation *user = u.getOwner();

          if (auto use = dyn_cast<stablehlo::SliceOp>(user)) {
            if (!mayReadMemoryWrittenTo(use, DUS)) {
              continue;
            }
            bool mayReadOther = false;
            Value v = DUS.getOperand();
            while (v != argOperand && v != condOperand) {
              auto DUS2 = v.getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
              assert(DUS2);
              if (mayReadMemoryWrittenTo(use, DUS2)) {
                mayReadOther = true;
                break;
              }
              v = DUS2.getOperand();

              // Don't handle case where in cond region.
              if (v == condOperand) {
                mayReadOther = true;
              }
            }
            if (!mayReadOther) {
              if (mustReadMemoryWrittenTo(use, DUS)) {
                mustReaders.push_back(use);
                continue;
              } else {
                mayReaders.push_back(use);
                continue;
              }
            }
          }

          if (auto use = dyn_cast<stablehlo::DynamicUpdateSliceOp>(user)) {
            if (use.getUpdate() == cur) {
              legal = false;
              break;
            }
            todo.push_back(use);
            continue;
          }

          if (auto use = dyn_cast<stablehlo::ReturnOp>(user)) {
            if (use->getParentOp() != whileOp) {
              legal = false;
              break;
            }
            if (u.getOperandNumber() != idx) {
              legal = false;
              break;
            }
            selfYield = true;
            continue;
          }

          legal = false;
          break;
        }
        if (!legal)
          break;
      }
      if (!legal)
        continue;
      if (!selfYield)
        continue;

      if (mustReaders.size())
        anyReader = true;
      if (mayReaders.size())
        anyReader = true;
      candidates.emplace_back(Candidate{idx, DUS, mustReaders, mayReaders});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    if (anyReader) {
      // Find the index of IV and the step to check for 1 iteration
      auto ivInfo = extractSimpleIVInfo(whileOp);
      if (!ivInfo.isValid)
        return failure();

      if (ivInfo.step == 0)
        return failure();
    }

    Value inPost;

    rewriter.setInsertionPointAfter(whileOp);
    {
      mlir::IRMapping condMapper;
      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        auto oldArg = whileOp.getCond().getArgument(i);
        condMapper.map(oldArg, whileOp.getOperands()[i]);
      }

      for (auto &op : whileOp.getCond().front().without_terminator()) {
        rewriter.clone(op, condMapper);
      }

      inPost = condMapper.lookupOrDefault(
          whileOp.getCond().front().getTerminator()->getOperand(0));
    }

    for (auto &candidate : candidates) {
      auto newDUS = stablehlo::DynamicUpdateSliceOp::create(
          rewriter, candidate.DUS.getLoc(), whileOp.getResults()[candidate.idx],
          candidate.DUS.getUpdate(), candidate.DUS.getStartIndices());
      auto sel = stablehlo::SelectOp::create(
          rewriter, whileOp.getLoc(), inPost, newDUS,
          whileOp.getResults()[candidate.idx]);
      SmallPtrSet<Operation *, 2> except = {sel, newDUS};
      rewriter.replaceAllUsesExcept(whileOp.getResults()[candidate.idx], sel,
                                    except);
    }

    if (anyReader) {
      auto ivInfo = extractSimpleIVInfo(whileOp);
      rewriter.setInsertionPointToStart(&whileOp.getBody().front());
      auto firstIter = stablehlo::CompareOp::create(
          rewriter, whileOp.getLoc(),
          whileOp.getBody().getArgument(ivInfo.index), ivInfo.start,
          stablehlo::ComparisonDirection::EQ);

      for (auto &candidate : candidates) {
        for (auto slice : candidate.mustReaders) {
          auto newstarts = llvm::to_vector(slice.getStartIndices());
          auto newlimits = llvm::to_vector(slice.getLimitIndices());
          for (size_t i = 0; i < slice.getType().getShape().size(); i++) {
            DenseIntElementsAttr startattr;

            if (!matchPattern(candidate.DUS.getStartIndices()[i],
                              m_Constant(&startattr))) {
              llvm_unreachable("expected constant stride");
            }
            int64_t dstart = (*startattr.begin()).getSExtValue();
            newstarts[i] -= dstart;
            newlimits[i] -= dstart;
          }
          rewriter.setInsertionPoint(slice);
          auto newSlice = stablehlo::SliceOp::create(
              rewriter, slice.getLoc(), candidate.DUS.getUpdate(), newstarts,
              newlimits, slice.getStrides());
          auto cloneSlice = rewriter.clone(*slice)->getResult(0);
          auto sel = stablehlo::SelectOp::create(
              rewriter, whileOp.getLoc(), firstIter, cloneSlice, newSlice);
          rewriter.replaceOp(slice, sel);
        }
        for (auto slice : candidate.mayReaders) {
          rewriter.setInsertionPoint(slice);
          auto newDUS = stablehlo::DynamicUpdateSliceOp::create(
              rewriter, candidate.DUS.getLoc(), slice.getOperand(),
              candidate.DUS.getUpdate(), candidate.DUS.getStartIndices());
          auto newSlice = stablehlo::SliceOp::create(
              rewriter, slice.getLoc(), newDUS, slice.getStartIndices(),
              slice.getLimitIndices(), slice.getStrides());
          auto cloneSlice = rewriter.clone(*slice)->getResult(0);
          auto sel = stablehlo::SelectOp::create(
              rewriter, whileOp.getLoc(), firstIter, cloneSlice, newSlice);
          rewriter.replaceOp(slice, sel);
        }
      }
    }

    for (auto &candidate : candidates) {
      rewriter.replaceOp(candidate.DUS, candidate.DUS.getOperand());
    }

    return success();
  }
};

struct HoistSlice
    : public CheckedOpRewritePattern<stablehlo::SliceOp, HoistSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp slice,
                                    PatternRewriter &rewriter) const {
    auto arg = dyn_cast<BlockArgument>(slice.getOperand());
    if (!arg)
      return failure();

    auto whileOp = dyn_cast<stablehlo::WhileOp>(arg.getOwner()->getParentOp());

    if (!whileOp)
      return failure();
    if (arg.getOwner() != &whileOp.getBody().front())
      return failure();
    auto idx = arg.getArgNumber();

    // Find yield op in the body
    //
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    auto operand = yieldOp.getOperands()[idx];
    while (operand != arg) {
      if (auto DUS = operand.getDefiningOp<stablehlo::DynamicUpdateSliceOp>()) {
        if (!mayReadMemoryWrittenTo(slice, DUS)) {
          operand = DUS.getOperand();
          continue;
        }
      }

      return failure();
    }

    rewriter.modifyOpInPlace(slice, [&]() {
      slice.getOperandMutable().assign(whileOp.getOperands()[idx]);
    });

    return success();
  }
};

struct WhileInductionReduction
    : public CheckedOpRewritePattern<stablehlo::WhileOp,
                                     WhileInductionReduction> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct Candidate {
      unsigned idx;
      SmallVector<int64_t> lowerBounds;
      SmallVector<int64_t> lowerUpdateBounds;
      SmallVector<int64_t> upperBounds;
      SmallVector<int64_t> upperUpdateBounds;
      BlockArgument argOperand;
      BlockArgument condOperand;
      Value outerOperand;
      Value outerSlice;
    };

    llvm::SmallVector<Candidate, 4> candidates;
    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      bool legal = true;
      // Skip DUS candidates which can be removed in a better form by WhileDUS
      if (auto DUS = yieldOp.getOperand(idx)
                         .getDefiningOp<stablehlo::DynamicUpdateSliceOp>()) {

        bool legalDUS = true;
        if (DUS.getOperand() == whileOp.getBody().front().getArgument(idx)) {
        } else if (definedOutside(DUS.getOperand(), whileOp)) {

          bool hasArgUse = !whileOp.getCond().getArgument(idx).use_empty() ||
                           !whileOp.getBody().getArgument(idx).use_empty();

          if (hasArgUse) {
            legalDUS = false;
          }

        } else {
          legalDUS = false;
        }

        if (legalDUS)
          legal = false;
      }

      if (!legal)
        continue;

      auto argOperand = whileOp.getBody().getArgument(idx);
      auto condOperand = whileOp.getCond().getArgument(idx);

      auto T = cast<RankedTensorType>(argOperand.getType());
      auto rank = T.getShape().size();

      if (rank == 0)
        continue;

      SmallVector<int64_t> lowerBounds = llvm::to_vector(T.getShape());
      SmallVector<int64_t> upperBounds(T.getShape().size(), 0);
      SmallVector<int64_t> lowerUpdateBounds = llvm::to_vector(T.getShape());
      SmallVector<int64_t> upperUpdateBounds(T.getShape().size(), 0);

      // We want to find a smaller window to use as the inductive variable
      // This means that no slice can read from outside the window (otherwise we
      // need the data). We also want to ensure that no dynamic update slice
      // overwrites outside the window
      //. This is because we want to just do a single DUS of the original input
      // outside the window

      SmallVector<Value> todo = {argOperand, condOperand};
      bool selfYield = false;
      while (!todo.empty()) {
        auto cur = todo.pop_back_val();
        for (auto &u : cur.getUses()) {
          Operation *user = u.getOwner();
          if (auto use = dyn_cast<stablehlo::SliceOp>(user)) {
            for (int i = 0; i < rank; i++) {
              if (use.getStrides()[i] != 1) {
                legal = false;
                break;
              }
              lowerBounds[i] =
                  std::min(lowerBounds[i], use.getStartIndices()[i]);
              upperBounds[i] =
                  std::max(upperBounds[i], use.getLimitIndices()[i]);
            }
            continue;
          }

          if (auto use = dyn_cast<stablehlo::DynamicUpdateSliceOp>(user)) {
            if (use.getUpdate() == cur) {
              legal = false;
              break;
            }
            for (int i = 0; i < rank; i++) {

              DenseIntElementsAttr idxAttr;
              if (!matchPattern(use.getStartIndices()[i],
                                m_Constant(&idxAttr)) ||
                  idxAttr.getNumElements() != 1) {
                legal = false;
                break;
              }
              auto start = (*idxAttr.begin()).getSExtValue();
              auto end = start + use.getUpdate().getType().getShape()[i];

              lowerBounds[i] = std::min(lowerBounds[i], start);
              upperBounds[i] = std::max(upperBounds[i], end);
              lowerUpdateBounds[i] = std::min(lowerUpdateBounds[i], start);
              upperUpdateBounds[i] = std::max(upperUpdateBounds[i], end);
            }
            todo.push_back(use);
            continue;
          }

          if (auto use = dyn_cast<stablehlo::ReturnOp>(user)) {
            if (use->getParentOp() != whileOp) {
              legal = false;
              break;
            }
            if (u.getOperandNumber() != idx) {
              legal = false;
              break;
            }
            selfYield = true;
            continue;
          }

          legal = false;
          break;
        }
        if (!legal)
          break;
      }
      if (!legal)
        continue;
      if (!selfYield)
        continue;

      bool seenSlice = false;

      for (int i = 0; i < rank; i++) {
        if (lowerUpdateBounds[i] >= upperUpdateBounds[i]) {
          legal = false;
          continue;
        }
        if (lowerUpdateBounds[i] != 0) {
          seenSlice = true;
          break;
        }
        if (upperUpdateBounds[i] != T.getShape()[i]) {
          seenSlice = true;
          break;
        }
      }

      if (!legal || !seenSlice)
        continue;

      candidates.emplace_back(Candidate{
          idx, lowerBounds, lowerUpdateBounds, upperBounds, upperUpdateBounds,
          argOperand, condOperand, whileOp.getOperands()[idx], nullptr});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create input transposes for each candidate
    for (auto &candidate : candidates) {
      // Create a new transpose before the while loop
      SmallVector<int64_t> strides(candidate.lowerBounds.size(), 1);
      candidate.outerSlice = stablehlo::SliceOp::create(
          rewriter, candidate.argOperand.getLoc(), candidate.outerOperand,
          candidate.lowerBounds, candidate.upperBounds, strides);
      SmallVector<int64_t> lower = llvm::to_vector(candidate.lowerUpdateBounds);
      SmallVector<int64_t> upper = llvm::to_vector(candidate.upperUpdateBounds);
      for (int i = 0; i < lower.size(); i++) {
        lower[i] -= candidate.lowerBounds[i];
        upper[i] -= candidate.lowerBounds[i];
      }
      newOperands[candidate.idx] = stablehlo::SliceOp::create(
          rewriter, candidate.argOperand.getLoc(), candidate.outerSlice, lower,
          upper, strides);
    }

    // Update yield op to use the input of the inner transpose
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : candidates) {
        SmallVector<int64_t> strides(candidate.lowerBounds.size(), 1);
        newReturnValues[candidate.idx] = stablehlo::SliceOp::create(
            rewriter, candidate.argOperand.getLoc(),
            newReturnValues[candidate.idx], candidate.lowerUpdateBounds,
            candidate.upperUpdateBounds, strides);
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    SmallVector<Value> results;
    for (auto res : newWhileOp.getResults())
      results.push_back(res);

    {
      auto ctype = RankedTensorType::get({}, rewriter.getI64Type());
      for (auto &candidate : candidates) {
        SmallVector<Value> starts;
        for (auto idx : candidate.lowerUpdateBounds) {
          starts.push_back(stablehlo::ConstantOp::create(
              rewriter, candidate.argOperand.getLoc(), ctype,
              cast<ElementsAttr>(makeAttr(ctype, idx))));
        }

        results[candidate.idx] = stablehlo::DynamicUpdateSliceOp::create(
            rewriter, candidate.argOperand.getLoc(), candidate.outerOperand,
            results[candidate.idx], starts);
      }
    }

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    // Create an IR mapper to map values from old op to new op
    mlir::IRMapping mapper;

    // Clear the new body block but keep its arguments
    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock
        .clear(); // This clears operations but preserves block arguments

    // Clone operations from old body to new body
    Block &oldBodyBlock = whileOp.getBody().front();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      // Set up operand mapping for the body region
      for (unsigned i = 0; i < whileOp.getBody().getNumArguments(); ++i) {
        auto oldArg = whileOp.getBody().getArgument(i);
        Value newArg = newWhileOp.getBody().getArgument(i);
        for (auto &pair : candidates) {
          if (pair.idx == i) {

            auto itype = RankedTensorType::get({}, rewriter.getI64Type());

            SmallVector<Value> update_starts;
            for (int i = 0; i < pair.lowerBounds.size(); i++) {
              update_starts.push_back(stablehlo::ConstantOp::create(
                  rewriter, pair.argOperand.getLoc(), itype,
                  cast<ElementsAttr>(
                      makeAttr(itype, pair.lowerUpdateBounds[i] -
                                          pair.lowerBounds[i]))));
            }

            newArg = stablehlo::DynamicUpdateSliceOp::create(
                rewriter, pair.argOperand.getLoc(), pair.outerSlice, newArg,
                update_starts);

            auto ctype = RankedTensorType::get(
                {}, cast<RankedTensorType>(pair.argOperand.getType())
                        .getElementType());
            auto padVal = stablehlo::ConstantOp::create(
                rewriter, pair.argOperand.getLoc(), ctype,
                cast<ElementsAttr>(makeAttr(ctype, 0)));

            SmallVector<int64_t> slow = llvm::to_vector(pair.lowerBounds);
            SmallVector<int64_t> shigh = llvm::to_vector(
                cast<RankedTensorType>(pair.argOperand.getType()).getShape());
            for (int i = 0; i < shigh.size(); i++)
              shigh[i] -= pair.upperBounds[i];
            SmallVector<int64_t> sint(shigh.size(), 0);

            newArg =
                stablehlo::PadOp::create(rewriter, pair.argOperand.getLoc(),
                                         newArg, padVal, slow, shigh, sint);
            break;
          }
        }
        mapper.map(oldArg, newArg);
      }

      for (auto &op : oldBodyBlock.getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
    }

    // Create a new terminator for the body region using new values
    {
      OpBuilder::InsertionGuard guard(rewriter);
      SmallVector<Value> newReturnValues;

      // Map old return values to new values using the mapper
      for (auto oldRetVal : yieldOp.getOperands()) {
        Value newRetVal = mapper.lookupOrNull(oldRetVal);
        // If the value isn't in the mapper, maybe it was a block argument or
        // constant
        if (!newRetVal)
          newRetVal = oldRetVal; // Consider more robust handling if needed
        newReturnValues.push_back(newRetVal);
      }

      // Create the return op at the end of the body
      rewriter.setInsertionPointToEnd(&newBodyBlock);
      stablehlo::ReturnOp::create(rewriter, yieldOp.getLoc(), newReturnValues);
    }

    // Create condition region mapper
    mlir::IRMapping condMapper;

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    Block &oldCondBlock = whileOp.getCond().front();

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        auto oldArg = whileOp.getCond().getArgument(i);
        Value newArg = newWhileOp.getCond().getArgument(i);
        for (auto &pair : candidates) {
          if (pair.idx == i) {
            auto itype = RankedTensorType::get({}, rewriter.getI64Type());

            SmallVector<Value> update_starts;
            for (int i = 0; i < pair.lowerBounds.size(); i++) {
              update_starts.push_back(stablehlo::ConstantOp::create(
                  rewriter, pair.argOperand.getLoc(), itype,
                  cast<ElementsAttr>(
                      makeAttr(itype, pair.lowerUpdateBounds[i] -
                                          pair.lowerBounds[i]))));
            }

            newArg = stablehlo::DynamicUpdateSliceOp::create(
                rewriter, pair.argOperand.getLoc(), pair.outerSlice, newArg,
                update_starts);
            auto ctype = RankedTensorType::get(
                {}, cast<RankedTensorType>(pair.condOperand.getType())
                        .getElementType());
            auto padVal = stablehlo::ConstantOp::create(
                rewriter, pair.condOperand.getLoc(), ctype,
                cast<ElementsAttr>(makeAttr(ctype, 0)));

            SmallVector<int64_t> slow = llvm::to_vector(pair.lowerBounds);
            SmallVector<int64_t> shigh = llvm::to_vector(
                cast<RankedTensorType>(pair.condOperand.getType()).getShape());
            for (int i = 0; i < shigh.size(); i++)
              shigh[i] -= pair.upperBounds[i];
            SmallVector<int64_t> sint(shigh.size(), 0);

            newArg =
                stablehlo::PadOp::create(rewriter, pair.condOperand.getLoc(),
                                         newArg, padVal, slow, shigh, sint);
            break;
          }
        }
        condMapper.map(oldArg, newArg);
      }

      for (auto &op : oldCondBlock.getOperations()) {
        rewriter.clone(op, condMapper);
      }
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, results);
    return success();
  }
};

// TODO: this is not valid in general but presumes the inner structure is valid
// from the input
struct WhileConcat
    : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileConcat> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct Candidate {
      unsigned idx;
      stablehlo::ConcatenateOp concat;
      int lhsSize;
      int rhsSize;
      stablehlo::SliceOp ops[3];
    };

    llvm::SmallVector<Candidate, 4> candidates;

    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      auto concat =
          yieldOp.getOperand(idx).getDefiningOp<stablehlo::ConcatenateOp>();

      // Check that the while result has exactly one use
      if (!concat)
        continue;

      if (!concat->hasOneUse()) {
        continue;
      }

      if (concat.getOperands().size() != 3)
        continue;

      stablehlo::SliceOp ops[3] = {
          concat.getOperands()[0].getDefiningOp<stablehlo::SliceOp>(),
          concat.getOperands()[1].getDefiningOp<stablehlo::SliceOp>(),
          concat.getOperands()[2].getDefiningOp<stablehlo::SliceOp>(),
      };

      if (!ops[0])
        continue;
      if (!ops[1])
        continue;
      if (!ops[2])
        continue;

      if (ops[0].getOperand() != ops[1].getOperand())
        continue;
      if (ops[0].getOperand() != ops[2].getOperand())
        continue;

      bool legal = true;

      for (int opn = 0; opn < 3; opn++)
        for (int i = 0; i < concat.getType().getShape().size(); i++)
          if (i != concat.getDimension()) {
            if (ops[0].getStartIndices()[i] != ops[opn].getStartIndices()[i]) {
              legal = false;
              break;
            }
            if (ops[opn].getStrides()[i] != 1) {
              legal = false;
              break;
            }
            if (ops[0].getLimitIndices()[i] != ops[opn].getLimitIndices()[i]) {
              legal = false;
              break;
            }
          }

      if (!legal)
        continue;

      int lowerLim = ops[1].getStartIndices()[concat.getDimension()];
      int upperLim = ops[1].getLimitIndices()[concat.getDimension()];
      if (ops[2].getStartIndices()[concat.getDimension()] != lowerLim)
        continue;

      if (ops[0].getLimitIndices()[concat.getDimension()] != upperLim)
        continue;

      int lhsSize = ops[0].getLimitIndices()[concat.getDimension()] -
                    ops[0].getStartIndices()[concat.getDimension()];
      if (lhsSize !=
          ops[1].getOperand().getType().getShape()[concat.getDimension()] -
              upperLim)
        continue;
      int rhsSize = ops[2].getLimitIndices()[concat.getDimension()] -
                    ops[2].getStartIndices()[concat.getDimension()];
      if (rhsSize != lowerLim)
        continue;

      // TODO this is unsafe unless the input is verified to have this property
      // For now we will assume it

      candidates.emplace_back(
          Candidate{idx, concat, lhsSize, rhsSize, {ops[0], ops[1], ops[2]}});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create input transposes for each candidate
    for (auto &candidate : candidates) {
      // Create a new transpose before the while loop

      newOperands[candidate.idx] = stablehlo::SliceOp::create(
          rewriter, candidate.ops[1].getLoc(),
          whileOp.getOperands()[candidate.idx],
          candidate.ops[1].getStartIndices(),
          candidate.ops[1].getLimitIndices(), candidate.ops[1].getStrides());
    }

    // Update yield op to use the input of the inner transpose
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : candidates) {
        newReturnValues[candidate.idx] = candidate.ops[1];
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    SmallVector<Value> results;
    for (auto res : newWhileOp.getResults())
      results.push_back(res);

    {
      for (auto &candidate : candidates) {

        SmallVector<int64_t> lowerStarts(
            candidate.concat.getType().getShape().size(), 0);
        SmallVector<int64_t> upperStarts(
            candidate.concat.getType().getShape().size(), 0);
        SmallVector<int64_t> lowerEnds(
            cast<RankedTensorType>(results[candidate.idx].getType())
                .getShape()
                .begin(),
            cast<RankedTensorType>(results[candidate.idx].getType())
                .getShape()
                .end());
        SmallVector<int64_t> upperEnds(
            cast<RankedTensorType>(results[candidate.idx].getType())
                .getShape()
                .begin(),
            cast<RankedTensorType>(results[candidate.idx].getType())
                .getShape()
                .end());
        SmallVector<int64_t> strides(
            candidate.concat.getType().getShape().size(), 1);

        lowerEnds[candidate.concat.getDimension()] = candidate.lhsSize;

        upperStarts[candidate.concat.getDimension()] =
            upperEnds[candidate.concat.getDimension()] - candidate.rhsSize;

        Value ops[3] = {
            stablehlo::SliceOp::create(rewriter, candidate.concat.getLoc(),
                                       results[candidate.idx], lowerStarts,
                                       lowerEnds, strides),
            results[candidate.idx],
            stablehlo::SliceOp::create(rewriter, candidate.concat.getLoc(),
                                       results[candidate.idx], upperStarts,
                                       upperEnds, strides),

        };
        results[candidate.idx] = stablehlo::ConcatenateOp::create(
            rewriter, candidate.concat.getLoc(), ops,
            candidate.concat.getDimension());
      }
    }

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    // Create an IR mapper to map values from old op to new op
    mlir::IRMapping mapper;

    // Clear the new body block but keep its arguments
    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock
        .clear(); // This clears operations but preserves block arguments

    // Clone operations from old body to new body
    Block &oldBodyBlock = whileOp.getBody().front();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      // Set up operand mapping for the body region
      for (unsigned i = 0; i < whileOp.getBody().getNumArguments(); ++i) {
        auto oldArg = whileOp.getBody().getArgument(i);
        Value newArg = newWhileOp.getBody().getArgument(i);
        for (auto &candidate : candidates) {
          if (candidate.idx == i) {

            SmallVector<int64_t> lowerStarts(
                candidate.concat.getType().getShape().size(), 0);
            SmallVector<int64_t> upperStarts(
                candidate.concat.getType().getShape().size(), 0);
            SmallVector<int64_t> lowerEnds(
                cast<RankedTensorType>(newArg.getType()).getShape().begin(),
                cast<RankedTensorType>(newArg.getType()).getShape().end());
            SmallVector<int64_t> upperEnds(
                cast<RankedTensorType>(newArg.getType()).getShape().begin(),
                cast<RankedTensorType>(newArg.getType()).getShape().end());
            SmallVector<int64_t> strides(
                candidate.concat.getType().getShape().size(), 1);

            lowerEnds[candidate.concat.getDimension()] = candidate.lhsSize;

            upperStarts[candidate.concat.getDimension()] =
                upperEnds[candidate.concat.getDimension()] - candidate.rhsSize;

            Value ops[3] = {
                stablehlo::SliceOp::create(rewriter, candidate.concat.getLoc(),
                                           newArg, lowerStarts, lowerEnds,
                                           strides),
                newArg,
                stablehlo::SliceOp::create(rewriter, candidate.concat.getLoc(),
                                           newArg, upperStarts, upperEnds,
                                           strides),

            };
            newArg = stablehlo::ConcatenateOp::create(
                rewriter, candidate.concat.getLoc(), ops,
                candidate.concat.getDimension());
            break;
          }
        }
        mapper.map(oldArg, newArg);
      }

      for (auto &op : oldBodyBlock.getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
    }

    // Create a new terminator for the body region using new values
    {
      OpBuilder::InsertionGuard guard(rewriter);
      SmallVector<Value> newReturnValues;

      // Map old return values to new values using the mapper
      for (auto oldRetVal : yieldOp.getOperands()) {
        Value newRetVal = mapper.lookupOrNull(oldRetVal);
        // If the value isn't in the mapper, maybe it was a block argument or
        // constant
        if (!newRetVal)
          newRetVal = oldRetVal; // Consider more robust handling if needed
        newReturnValues.push_back(newRetVal);
      }

      // Create the return op at the end of the body
      rewriter.setInsertionPointToEnd(&newBodyBlock);
      stablehlo::ReturnOp::create(rewriter, yieldOp.getLoc(), newReturnValues);
    }

    // Create condition region mapper
    mlir::IRMapping condMapper;

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    Block &oldCondBlock = whileOp.getCond().front();

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        auto oldArg = whileOp.getCond().getArgument(i);
        Value newArg = newWhileOp.getCond().getArgument(i);
        for (auto &candidate : candidates) {
          if (candidate.idx == i) {
            SmallVector<int64_t> lowerStarts(
                candidate.concat.getType().getShape().size(), 0);
            SmallVector<int64_t> upperStarts(
                candidate.concat.getType().getShape().size(), 0);
            SmallVector<int64_t> lowerEnds(
                cast<RankedTensorType>(newArg.getType()).getShape().begin(),
                cast<RankedTensorType>(newArg.getType()).getShape().end());
            SmallVector<int64_t> upperEnds(
                cast<RankedTensorType>(newArg.getType()).getShape().begin(),
                cast<RankedTensorType>(newArg.getType()).getShape().end());
            SmallVector<int64_t> strides(
                candidate.concat.getType().getShape().size(), 1);

            lowerEnds[candidate.concat.getDimension()] = candidate.lhsSize;

            upperStarts[candidate.concat.getDimension()] =
                upperEnds[candidate.concat.getDimension()] - candidate.rhsSize;

            Value ops[3] = {
                stablehlo::SliceOp::create(rewriter, candidate.concat.getLoc(),
                                           newArg, lowerStarts, lowerEnds,
                                           strides),
                newArg,
                stablehlo::SliceOp::create(rewriter, candidate.concat.getLoc(),
                                           newArg, upperStarts, upperEnds,
                                           strides),

            };
            newArg = stablehlo::ConcatenateOp::create(
                rewriter, candidate.concat.getLoc(), ops,
                candidate.concat.getDimension());
            break;
          }
        }
        condMapper.map(oldArg, newArg);
      }

      for (auto &op : oldCondBlock.getOperations()) {
        rewriter.clone(op, condMapper);
      }
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, results);

    return success();
  }
};

// TODO: this is not valid in general but presumes the inner structure is valid
// from the input
template <typename T>
struct WhileWrap
    : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileWrap<T>> {
  using CheckedOpRewritePattern<stablehlo::WhileOp,
                                WhileWrap<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find yield op in the body
    auto &bodyBlock = whileOp.getBody().front();
    auto yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());

    // Step 1: Track which results need to be transformed
    struct Candidate {
      unsigned idx;
      T concat;
    };

    llvm::SmallVector<Candidate, 4> candidates;

    for (unsigned idx = 0; idx < yieldOp.getNumOperands(); ++idx) {

      auto concat = yieldOp.getOperand(idx).getDefiningOp<T>();

      // Check that the while result has exactly one use
      if (!concat)
        continue;

      candidates.emplace_back(Candidate{idx, concat});
    }

    // If no candidates found, no rewrite needed
    if (candidates.empty())
      return failure();

    // Step 2 : Make transformations in the original while op
    // Get the operands of the while op to use later
    auto whileOperands = llvm::to_vector(whileOp.getOperands());

    // New operands
    SmallVector<Value> newOperands(whileOp.getOperands().begin(),
                                   whileOp.getOperands().end());

    // Create input transposes for each candidate
    for (auto &candidate : candidates) {
      // Create a new transpose before the while loop
      SmallVector<int64_t> starts(candidate.concat.getType().getShape().size(),
                                  0);
      SmallVector<int64_t> steps(candidate.concat.getType().getShape().size(),
                                 1);
      SmallVector<int64_t> limits = llvm::to_vector(
          cast<RankedTensorType>(whileOp.getOperands()[candidate.idx].getType())
              .getShape());
      starts[candidate.concat.getDimension()] += candidate.concat.getLhs();
      limits[candidate.concat.getDimension()] -= candidate.concat.getRhs();

      newOperands[candidate.idx] = stablehlo::SliceOp::create(
          rewriter, candidate.concat.getLoc(),
          whileOp.getOperands()[candidate.idx], starts, limits, steps);
    }

    // Update yield op to use the input of the inner transpose
    {
      // Save the current insertion point
      mlir::OpBuilder::InsertionGuard guard(rewriter);

      // New return values
      SmallVector<Value> newReturnValues(yieldOp.getOperands().begin(),
                                         yieldOp.getOperands().end());

      rewriter.setInsertionPoint(yieldOp);
      for (auto &candidate : candidates) {
        newReturnValues[candidate.idx] = candidate.concat.getOperand();
      }
      rewriter.replaceOpWithNewOp<stablehlo::ReturnOp>(yieldOp,
                                                       newReturnValues);
      // Update the yieldOp
      yieldOp = cast<stablehlo::ReturnOp>(bodyBlock.getTerminator());
    }

    // Step 3 : Create a new while op with the new operands and move the body of
    // original whileOp
    SmallVector<Type> newResultTypes;
    newResultTypes.reserve(newOperands.size());

    for (auto operand : newOperands) {
      newResultTypes.push_back(operand.getType());
    }
    auto newWhileOp =
        stablehlo::WhileOp::create(rewriter, whileOp.getLoc(), newResultTypes,
                                   newOperands, whileOp->getAttrs());

    SmallVector<Value> results;
    for (auto res : newWhileOp.getResults())
      results.push_back(res);

    {
      for (auto &candidate : candidates) {
        results[candidate.idx] = T::create(
            rewriter, candidate.concat.getLoc(), results[candidate.idx],
            candidate.concat.getLhs(), candidate.concat.getRhs(),
            candidate.concat.getDimension());
      }
    }

    // Create blocks in both regions first
    {
      // Create a block in the condition region
      Block *condBlock = rewriter.createBlock(&newWhileOp.getCond());

      // Add arguments to the condition block matching operand types
      for (auto type : newResultTypes) {
        condBlock->addArgument(type, whileOp.getLoc());
      }

      // Create a block in the body region
      Block *bodyBlock = rewriter.createBlock(&newWhileOp.getBody());

      // Add arguments to the body block matching operand types
      for (auto type : newResultTypes) {
        bodyBlock->addArgument(type, whileOp.getLoc());
      }
    }

    // Create an IR mapper to map values from old op to new op
    mlir::IRMapping mapper;

    // Clear the new body block but keep its arguments
    Block &newBodyBlock = newWhileOp.getBody().front();
    newBodyBlock
        .clear(); // This clears operations but preserves block arguments

    // Clone operations from old body to new body
    Block &oldBodyBlock = whileOp.getBody().front();
    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newBodyBlock);

      // Set up operand mapping for the body region
      for (unsigned i = 0; i < whileOp.getBody().getNumArguments(); ++i) {
        auto oldArg = whileOp.getBody().getArgument(i);
        Value newArg = newWhileOp.getBody().getArgument(i);
        for (auto &candidate : candidates) {
          if (candidate.idx == i) {
            newArg =
                T::create(rewriter, candidate.concat.getLoc(), newArg,
                          candidate.concat.getLhs(), candidate.concat.getRhs(),
                          candidate.concat.getDimension());
            break;
          }
        }
        mapper.map(oldArg, newArg);
      }

      for (auto &op : oldBodyBlock.getOperations()) {
        // Skip the terminator - we'll add it after all other operations
        if (isa<stablehlo::ReturnOp>(op))
          continue;

        // Clone the operation with the value mapping
        rewriter.clone(op, mapper);
      }
    }

    // Create a new terminator for the body region using new values
    {
      OpBuilder::InsertionGuard guard(rewriter);
      SmallVector<Value> newReturnValues;

      // Map old return values to new values using the mapper
      for (auto oldRetVal : yieldOp.getOperands()) {
        Value newRetVal = mapper.lookupOrNull(oldRetVal);
        // If the value isn't in the mapper, maybe it was a block argument or
        // constant
        if (!newRetVal)
          newRetVal = oldRetVal; // Consider more robust handling if needed
        newReturnValues.push_back(newRetVal);
      }

      // Create the return op at the end of the body
      rewriter.setInsertionPointToEnd(&newBodyBlock);
      stablehlo::ReturnOp::create(rewriter, yieldOp.getLoc(), newReturnValues);
    }

    // Create condition region mapper
    mlir::IRMapping condMapper;

    // Clear and clone condition region
    Block &newCondBlock = newWhileOp.getCond().front();
    newCondBlock.clear();
    Block &oldCondBlock = whileOp.getCond().front();

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(&newCondBlock);

      for (unsigned i = 0; i < whileOp.getCond().getNumArguments(); ++i) {
        auto oldArg = whileOp.getCond().getArgument(i);
        Value newArg = newWhileOp.getCond().getArgument(i);
        for (auto &candidate : candidates) {
          if (candidate.idx == i) {
            newArg =
                T::create(rewriter, candidate.concat.getLoc(), newArg,
                          candidate.concat.getLhs(), candidate.concat.getRhs(),
                          candidate.concat.getDimension());
            break;
          }
        }
        condMapper.map(oldArg, newArg);
      }

      for (auto &op : oldCondBlock.getOperations()) {
        rewriter.clone(op, condMapper);
      }
    }

    // Finally, replace all uses of the old while op with the new one
    rewriter.replaceOp(whileOp, results);

    return success();
  }
};

// A while-carried accumulator of the form
//
//   init      = <zero splat>
//   yield[k]  = add(bodyArg[k], dynamic_update_slice(<zero>, row, i))
//
// only ever adds into rows that are still zero, provided the scatter index i is
// distinct on every iteration and never clamped. The add is then a no-op
// everywhere except the written row, where it adds to zero, so the whole thing
// collapses to a read-modify-write of that row:
//
//   yield[k] = dynamic_update_slice(bodyArg[k], row, i)
//
// That turns a full-tensor materialise-and-add -- roughly 4 x S0 x M of traffic
// per iteration -- into a single row write, which XLA can perform in place
// since the accumulator is a loop carry with no other reader.
//
// This needs only injectivity, not full coverage of [0, S0): rows that are
// never written keep the zero they started with under either form.
struct WhileScatterAccumulatorNoAdd final
    : CheckedOpRewritePattern<stablehlo::AddOp, WhileScatterAccumulatorNoAdd> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AddOp addOp,
                                    PatternRewriter &rewriter) const {
    auto whileOp = dyn_cast<stablehlo::WhileOp>(addOp->getParentOp());
    if (!whileOp || addOp->getBlock() != &whileOp.getBody().front())
      return failure();

    // The accumulator must be yielded straight back, and read by nothing else.
    if (!addOp->hasOneUse())
      return failure();
    auto yieldOp = dyn_cast<stablehlo::ReturnOp>(*addOp->getUsers().begin());
    if (!yieldOp)
      return failure();

    BlockArgument bodyArg;
    Value dusSide;
    for (auto [a, b] : {std::make_pair(addOp.getLhs(), addOp.getRhs()),
                        std::make_pair(addOp.getRhs(), addOp.getLhs())}) {
      if (auto arg = dyn_cast<BlockArgument>(a)) {
        if (arg.getOwner() == &whileOp.getBody().front()) {
          bodyArg = arg;
          dusSide = b;
          break;
        }
      }
    }
    if (!bodyArg || !bodyArg.hasOneUse())
      return failure();

    unsigned idx = bodyArg.getArgNumber();
    if (yieldOp.getOperand(idx) != addOp.getResult())
      return failure();
    if (!enzyme::isZero(whileOp.getOperands()[idx]))
      return failure();

    auto dus = dusSide.getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
    if (!dus || !dus->hasOneUse() || !enzyme::isZero(dus.getOperand()))
      return failure();

    auto accTy = cast<RankedTensorType>(dus.getType());
    auto updateTy = cast<RankedTensorType>(dus.getUpdate().getType());
    if (updateTy.getRank() != accTy.getRank())
      return failure();

    // Exactly one dimension may be scattered; the rest must be written whole,
    // at offset zero, so that "the written row" is well defined.
    int64_t scatterDim = -1;
    auto starts = dus.getStartIndices();
    for (int64_t d = 0; d < accTy.getRank(); ++d) {
      if (updateTy.getDimSize(d) == accTy.getDimSize(d)) {
        if (!enzyme::isZero(starts[d]))
          return failure();
        continue;
      }
      if (updateTy.getDimSize(d) != 1 || scatterDim != -1)
        return failure();
      scatterDim = d;
    }
    if (scatterDim == -1)
      return failure();

    enzyme::WhileLoopInfo info(whileOp);
    if (failed(info.computeInfo()) || !info.isValid() || !info.isConstant())
      return failure();
    info.propagateAffineIndexInfo();
    info.propagateBounds();

    Value index = starts[scatterDim];

    // Distinct on every iteration: the index must actually vary with the
    // induction variable.
    auto affine = info.getAffineIndexInfo();
    auto affineIt = affine.find(index);
    if (affineIt == affine.end() || affineIt->second.scale.isZero())
      return failure();

    // Never clamped: dynamic_update_slice would otherwise fold out-of-range
    // iterations onto the boundary row and genuinely accumulate there.
    auto bounds = info.getBounds(index);
    if (!bounds || bounds->min.isNegative() ||
        bounds->max.getSExtValue() > accTy.getDimSize(scatterDim) - 1)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
        addOp, bodyArg, dus.getUpdate(), starts);
    return success();
  }
};

// Replace while op iteration variables which are not updated with their
// upcoming value
struct WhileSimplify
    : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  bool hoist_all;
  WhileSimplify(bool hoist_all, MLIRContext *context,
                PatternBenefit benefit = 1,
                ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        hoist_all(hoist_all) {}

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp op,
                                    PatternRewriter &rewriter) const {
    SmallVector<unsigned> operands;

    Block *cond = &op.getCond().front(), *body = &op.getBody().front();
    Operation *bodyTerm = body->getTerminator();

    int deleted = 0;

    for (auto &opOperand : op->getOpOperands()) {
      Value inputValue = opOperand.get();

      auto i = opOperand.getOperandNumber() - deleted;
      Value bodyArg = body->getArgument(i);
      Value condArg = cond->getArgument(i);

      bool canHoist = inputValue.getDefiningOp<stablehlo::ConstantOp>();
      if (hoist_all) {
        canHoist = true;
      } else if (auto BA = dyn_cast<BlockArgument>(inputValue)) {
        canHoist |= isa<FunctionOpInterface>(BA.getOwner()->getParentOp());
      }

      Value bodyRes = bodyTerm->getOperand(i);

      if (canHoist && bodyArg == bodyRes) {
        // This variable is not updated during iterations
        rewriter.replaceAllUsesWith(bodyArg, inputValue);
        rewriter.replaceAllUsesWith(condArg, inputValue);
        rewriter.modifyOpInPlace(bodyTerm,
                                 [&] { bodyTerm->setOperands(i, 1, {}); });
        rewriter.replaceAllUsesWith(op.getResult(opOperand.getOperandNumber()),
                                    inputValue);

        body->eraseArgument(i);
        cond->eraseArgument(i);

        deleted++;
      } else {
        operands.push_back(opOperand.getOperandNumber());
      }
    }

    if (operands.size() == op->getNumOperands())
      return failure();

    SmallVector<Value> newOperands;
    newOperands.reserve(operands.size());

    for (auto opOperand : operands) {
      newOperands.push_back(op->getOperand(opOperand));
    }

    auto newWhile = stablehlo::WhileOp::create(rewriter, op.getLoc(),
                                               newOperands, op->getAttrs());
    newWhile.getCond().takeBody(op.getCond());
    newWhile.getBody().takeBody(op.getBody());

    // Replace uses for remaining results.
    for (const auto &it : llvm::enumerate(operands)) {
      Value oldRes = op->getResult(it.value());
      Value newRes = newWhile->getResult(it.index());

      rewriter.replaceAllUsesWith(oldRes, newRes);
    }

    rewriter.eraseOp(op);

    return success();
  }
};

// Replace while op iteration variables which are not updated with their
// upcoming value
struct WhileLICM
    : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileLICM> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  bool hoist_all;
  WhileLICM(bool hoist_all, MLIRContext *context, PatternBenefit benefit = 1,
            ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        hoist_all(hoist_all) {}

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp op,
                                    PatternRewriter &rewriter) const {
    SmallVector<unsigned> operands;

    Block *cond = &op.getCond().front(), *body = &op.getBody().front();
    Operation *bodyTerm = body->getTerminator();

    int deleted = 0;

    // Find the index of IV and the step to check for 1 iteration
    auto ivInfo = extractSimpleIVInfo(op);

    for (auto &opOperand : op->getOpOperands()) {
      Value inputValue = opOperand.get();

      auto i = opOperand.getOperandNumber() - deleted;
      Value bodyArg = body->getArgument(i);
      Value condArg = cond->getArgument(i);

      bool canHoist = inputValue.getDefiningOp<stablehlo::ConstantOp>();
      if (auto BA = dyn_cast<BlockArgument>(inputValue)) {
        canHoist |= isa<FunctionOpInterface>(BA.getOwner()->getParentOp());
      } else if (hoist_all) {
        canHoist = true;
      }

      Value bodyRes = bodyTerm->getOperand(i);

      if (canHoist && definedOutside(bodyRes, op) && ivInfo.isValid &&
          ivInfo.step != 0) {

        Value resultReplacement;
        {
          OpBuilder::InsertionGuard guard(rewriter);
          IRMapping mapper;

          for (unsigned i = 0; i < op.getCond().getNumArguments(); ++i) {
            mapper.map(op.getCond().getArgument(i), op.getOperands()[i]);
          }
          for (auto &op : op.getCond().front().getOperations()) {
            // Skip the terminator - we'll add it after all other operations
            if (isa<stablehlo::ReturnOp>(op))
              continue;

            // Clone the operation with the value mapping
            rewriter.clone(op, mapper);
          }
          Value useInner = op.getCond().front().getTerminator()->getOperand(0);
          useInner = mapper.lookupOrDefault(useInner);

          resultReplacement = stablehlo::SelectOp::create(
              rewriter, op.getLoc(), useInner, bodyRes, inputValue);
        }

        // This variable is not updated during iterations
        {
          OpBuilder::InsertionGuard guard(rewriter);
          rewriter.setInsertionPointToStart(&op.getBody().front());
          auto replacement = stablehlo::SelectOp::create(
              rewriter, op.getLoc(),
              stablehlo::CompareOp::create(
                  rewriter, op.getLoc(), op.getBody().getArgument(ivInfo.index),
                  ivInfo.start, stablehlo::ComparisonDirection::EQ),
              inputValue, bodyRes);
          rewriter.replaceAllUsesWith(bodyArg, replacement);
        }

        {
          OpBuilder::InsertionGuard guard(rewriter);
          rewriter.setInsertionPointToStart(&op.getCond().front());
          auto replacement = stablehlo::SelectOp::create(
              rewriter, op.getLoc(),
              stablehlo::CompareOp::create(
                  rewriter, op.getLoc(), op.getCond().getArgument(ivInfo.index),
                  ivInfo.start, stablehlo::ComparisonDirection::EQ),
              inputValue, bodyRes);
          rewriter.replaceAllUsesWith(condArg, replacement);
        }

        rewriter.modifyOpInPlace(bodyTerm,
                                 [&] { bodyTerm->setOperands(i, 1, {}); });

        rewriter.replaceAllUsesWith(op.getResult(opOperand.getOperandNumber()),
                                    resultReplacement);

        body->eraseArgument(i);
        cond->eraseArgument(i);

        deleted++;
      } else {
        operands.push_back(opOperand.getOperandNumber());
      }
    }

    if (operands.size() == op->getNumOperands())
      return failure();

    SmallVector<Value> newOperands;
    newOperands.reserve(operands.size());

    for (auto opOperand : operands) {
      newOperands.push_back(op->getOperand(opOperand));
    }

    auto newWhile = stablehlo::WhileOp::create(rewriter, op.getLoc(),
                                               newOperands, op->getAttrs());
    newWhile.getCond().takeBody(op.getCond());
    newWhile.getBody().takeBody(op.getBody());

    // Replace uses for remaining results.
    for (const auto &it : llvm::enumerate(operands)) {
      Value oldRes = op->getResult(it.value());
      Value newRes = newWhile->getResult(it.index());

      rewriter.replaceAllUsesWith(oldRes, newRes);
    }

    rewriter.eraseOp(op);

    return success();
  }
};

// For a stablehlo.while whose body contains a stablehlo.if with a
// loop-invariant predicate, hoist the conditional out of the loop by creating
// two specialised while loops — one per branch — wrapped in an outer
// stablehlo.if.  The pattern fires once per qualifying if; the greedy driver
// iterates until no more loop-invariant ifs remain.
struct LoopUnswitch
    : public CheckedOpRewritePattern<stablehlo::WhileOp, LoopUnswitch> {
  int64_t maxOpsOutside;
  LoopUnswitch(int64_t maxOpsOutside, MLIRContext *context,
               PatternBenefit benefit = 1,
               ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        maxOpsOutside(maxOpsOutside) {}

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    // Find the first stablehlo.if in the body whose predicate is defined
    // outside the while loop.
    stablehlo::IfOp targetIf;
    for (Operation &op : whileOp.getBody().front().without_terminator()) {
      if (auto ifOp = dyn_cast<stablehlo::IfOp>(&op)) {
        if (definedOutside(ifOp.getPred(), whileOp)) {
          targetIf = ifOp;
          break;
        }
      }
    }
    if (!targetIf)
      return failure();

    // Count ops in the body outside the target if (excluding the terminator).
    // Unswitching duplicates all of these; skip if too many.
    int64_t opsOutside = 0;
    for (Operation &op : whileOp.getBody().front().without_terminator())
      if (&op != targetIf.getOperation())
        ++opsOutside;
    if (opsOutside > maxOpsOutside)
      return failure();

    Value pred = targetIf.getPred();
    Location loc = whileOp.getLoc();

    // Clone the while and inline one branch of the cloned if in place of the
    // if itself.  Because pred is defined outside the while, the clone does not
    // remap it, so we can locate the cloned if by predicate identity.
    auto buildVersionedWhile = [&](bool useTrue) -> stablehlo::WhileOp {
      IRMapping mapper;
      auto newWhile =
          cast<stablehlo::WhileOp>(rewriter.clone(*whileOp, mapper));

      stablehlo::IfOp clonedIf;
      for (Operation &op : newWhile.getBody().front().without_terminator()) {
        if (auto ifOp = dyn_cast<stablehlo::IfOp>(&op);
            ifOp && ifOp.getPred() == pred) {
          clonedIf = ifOp;
          break;
        }
      }
      assert(clonedIf && "cloned while must contain the target if");

      Region &inlineBranch =
          useTrue ? clonedIf.getTrueBranch() : clonedIf.getFalseBranch();
      Operation *branchTerm = inlineBranch.front().getTerminator();
      SmallVector<Value> branchReturns(branchTerm->getOperands());

      rewriter.setInsertionPoint(clonedIf);
      for (Operation &innerOp : llvm::make_early_inc_range(
               inlineBranch.front().without_terminator()))
        rewriter.modifyOpInPlace(&innerOp,
                                 [&] { innerOp.moveBefore(clonedIf); });

      rewriter.replaceOp(clonedIf, branchReturns);
      return newWhile;
    };

    rewriter.setInsertionPoint(whileOp);
    auto outerIf =
        stablehlo::IfOp::create(rewriter, loc, whileOp->getResultTypes(), pred);

    // IfOp::create leaves both regions empty; we must add blocks before use.
    {
      OpBuilder::InsertionGuard g(rewriter);
      Block *trueBlock = rewriter.createBlock(&outerIf.getTrueBranch());
      auto trueWhile = buildVersionedWhile(/*useTrue=*/true);
      rewriter.setInsertionPointToEnd(trueBlock);
      stablehlo::ReturnOp::create(rewriter, loc, trueWhile.getResults());
    }
    {
      OpBuilder::InsertionGuard g(rewriter);
      Block *falseBlock = rewriter.createBlock(&outerIf.getFalseBranch());
      auto falseWhile = buildVersionedWhile(/*useTrue=*/false);
      rewriter.setInsertionPointToEnd(falseBlock);
      stablehlo::ReturnOp::create(rewriter, loc, falseWhile.getResults());
    }

    rewriter.replaceOp(whileOp, outerIf.getResults());
    return success();
  }
};

// Replace a while op consisting of DUS chains for each iter arg where each DUS
// doesn't depend on values evolving in the loop. This can be later extended to
// support other similarly idempotent ops. The loop is expected to have a
// counter-style condition that will be used as the if condition.
struct WhileIdempotentDUS
    : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileIdempotentDUS> {
  using CheckedOpRewritePattern<stablehlo::WhileOp,
                                WhileIdempotentDUS>::CheckedOpRewritePattern;

  // Returns true if the value is the block argument of the while body at the
  // given position.
  bool isSameBlockArgument(Value value, stablehlo::WhileOp whileOp,
                           unsigned position) const {
    auto blockArg = dyn_cast<BlockArgument>(value);
    if (!blockArg)
      return false;
    auto foundWhileOp =
        dyn_cast<stablehlo::WhileOp>(blockArg.getOwner()->getParentOp());
    if (whileOp != foundWhileOp ||
        blockArg.getOwner() != &whileOp.getBody().front())
      return false;
    return blockArg.getArgNumber() == position;
  }

  // Returns true if the value is transitively defined by DUS ops that do not
  // depend on other iter args than the one at the given position.
  bool isIdempotent(Value value, stablehlo::WhileOp whileOp, unsigned position,
                    SmallPtrSetImpl<Operation *> &updates) const {
    if (areValuesDefinedAbove(ValueRange{value}, whileOp.getBody()))
      return true;

    if (auto dus = value.getDefiningOp<stablehlo::DynamicUpdateSliceOp>()) {
      for (auto ind : dus.getStartIndices()) {
        if (!isIdempotent(ind, whileOp, position, updates))
          return false;
      }
      if (!isIdempotent(dus.getUpdate(), whileOp, position, updates))
        return false;
      if (!isIdempotent(dus.getOperand(), whileOp, position, updates) &&
          !isSameBlockArgument(dus.getOperand(), whileOp, position))
        return false;

      updates.insert(dus);
      return true;
    }

    if (auto dus = value.getDefiningOp<enzymexla::UpdateWithoutCornersOp>()) {
      if (!isIdempotent(dus.getUpdate(), whileOp, position, updates))
        return false;

      if (!isIdempotent(dus.getOperand(), whileOp, position, updates) &&
          !isSameBlockArgument(dus.getOperand(), whileOp, position))
        return false;

      updates.insert(dus);
      return true;
    }

    return false;
  }

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp op,
                                    PatternRewriter &rewriter) const {

    // Find the index of IV and the step to check for 1 iteration
    auto ivInfo = extractSimpleIVInfo(op);
    if (!ivInfo.isValid)
      return failure();

    if (ivInfo.step == 0)
      return failure();

    // check yielded value for being the iter arg potentially going through a
    // chain of DUS that do not depend on other iter args
    Block &body = op.getBody().front();
    auto returnOp = cast<stablehlo::ReturnOp>(body.getTerminator());
    SmallPtrSet<Operation *, 1> toMove;
    SmallVector<size_t> inds;
    for (OpOperand &operand : returnOp->getOpOperands()) {

      // Don't bother with arguments with no uses.
      if (body.getArguments()[operand.getOperandNumber()].use_empty() &&
          op.getCond()
              .front()
              .getArguments()[operand.getOperandNumber()]
              .use_empty() &&
          op.getResults()[operand.getOperandNumber()].use_empty())
        continue;

      if (areValuesDefinedAbove(ValueRange{operand.get()}, op.getBody()))
        continue;

      // if all operands are idempotent, we can replace the loop with a
      // conditional
      SmallPtrSet<Operation *, 1> localToMove;
      if (!isIdempotent(operand.get(), op, operand.getOperandNumber(),
                        localToMove))
        continue;

      SmallVector<Value, 2> vals = {
          body.getArguments()[operand.getOperandNumber()],
          op.getCond().front().getArguments()[operand.getOperandNumber()]};
      for (auto op : localToMove) {
        for (auto res : op->getResults())
          vals.push_back(res);
      }
      bool legal = true;
      for (auto val : vals) {
        for (auto &use : val.getUses()) {
          if (localToMove.contains(use.getOwner())) {
            continue;
          }
          if (use.getOwner() == returnOp) {
            if (use.getOperandNumber() == operand.getOperandNumber()) {
              continue;
            }
          }
          legal = false;
          break;
        }
      }
      if (!legal)
        continue;
      for (auto op : localToMove)
        toMove.insert(op);
      inds.push_back(operand.getOperandNumber());
    }

    if (inds.empty())
      return failure();

    IRMapping mapping;
    mapping.map(op.getBody().getArguments(), op.getOperands());
    for (auto &op : body) {
      if (toMove.contains(&op))
        rewriter.clone(op, mapping);
    }

    auto ogops = llvm::to_vector(returnOp.getOperands());
    rewriter.modifyOpInPlace(returnOp, [&]() {
      auto ops = ogops;
      for (auto ind : inds) {
        ops[ind] = mapping.lookupOrDefault(ops[ind]);
      }
      returnOp.getResultsMutable().assign(ops);
    });

    {
      rewriter.setInsertionPointToStart(&op.getCond().front());
      auto firstIter = stablehlo::CompareOp::create(
          rewriter, op.getLoc(), op.getCond().getArgument(ivInfo.index),
          ivInfo.start, stablehlo::ComparisonDirection::EQ);
      for (auto ind : inds) {
        if (!op.getCond().getArguments()[ind].use_empty()) {
          auto sel = stablehlo::SelectOp::create(
              rewriter, op.getLoc(), firstIter, op.getOperands()[ind],
              mapping.lookupOrDefault(ogops[ind]));
          rewriter.replaceAllUsesWith(op.getCond().getArguments()[ind], sel);
        }
      }
    }

    {
      rewriter.setInsertionPointToStart(&op.getBody().front());
      auto firstIter = stablehlo::CompareOp::create(
          rewriter, op.getLoc(), op.getBody().getArgument(ivInfo.index),
          ivInfo.start, stablehlo::ComparisonDirection::EQ);

      for (auto ind : inds) {
        if (!op.getBody().getArguments()[ind].use_empty()) {
          auto sel = stablehlo::SelectOp::create(
              rewriter, op.getLoc(), firstIter, op.getOperands()[ind],
              mapping.lookupOrDefault(ogops[ind]));
          rewriter.replaceAllUsesWith(op.getBody().getArguments()[ind], sel);
        }
      }
    }

    rewriter.setInsertionPointAfter(op);

    auto zeroIters = stablehlo::CompareOp::create(
        rewriter, op.getLoc(), op.getResults()[ivInfo.index],
        op.getOperands()[ivInfo.index], stablehlo::ComparisonDirection::EQ);

    for (auto ind : inds) {
      if (!op.getResults()[ind].use_empty()) {
        auto sel = stablehlo::SelectOp::create(
            rewriter, op.getLoc(), zeroIters, op.getOperands()[ind],
            mapping.lookupOrDefault(ogops[ind]));
        rewriter.replaceAllUsesWith(op.getResults()[ind], sel);
      }
    }
    return success();
  }
};

struct DynamicGatherOpIsNotDynamic
    : public CheckedOpRewritePattern<stablehlo::DynamicGatherOp,
                                     DynamicGatherOpIsNotDynamic> {
  using CheckedOpRewritePattern<
      stablehlo::DynamicGatherOp,
      DynamicGatherOpIsNotDynamic>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicGatherOp op,
                                    PatternRewriter &rewriter) const {
    // Check if slice sizes are constant.
    DenseIntElementsAttr sliceSizesAttr;
    if (!matchPattern(op.getSliceSizes(), m_Constant(&sliceSizesAttr))) {
      return failure();
    }

    // dynamic_gather allows non-int64 slice sizes, but we need to convert them
    // to int64 for the gather.
    if (!sliceSizesAttr.getType().getElementType().isInteger(64)) {
      SmallVector<APInt> sliceSizes;
      for (auto size : sliceSizesAttr.getValues<APInt>()) {
        sliceSizes.push_back(size);
      }
      auto newSliceSizesAttr = DenseElementsAttr::get(
          RankedTensorType::get(sliceSizesAttr.getType().getShape(),
                                rewriter.getIntegerType(64)),
          sliceSizes);
      sliceSizesAttr = cast<DenseIntElementsAttr>(newSliceSizesAttr);
    }

    SmallVector<int64_t> sliceSizes;
    for (auto size : sliceSizesAttr.getValues<int64_t>()) {
      sliceSizes.push_back(size);
    }
    auto sliceSizesArrayAttr =
        DenseI64ArrayAttr::get(op.getContext(), sliceSizes);

    rewriter.replaceOpWithNewOp<stablehlo::GatherOp>(
        op, op.getType(), op.getOperand(), op.getStartIndices(),
        stablehlo::GatherDimensionNumbersAttr::get(
            op.getContext(), op.getDimensionNumbers().getOffsetDims(),
            op.getDimensionNumbers().getCollapsedSliceDims(),
            /*operandBatchingDims=*/{},
            /*startIndicesBatchingDims=*/{},
            op.getDimensionNumbers().getStartIndexMap(),
            op.getDimensionNumbers().getIndexVectorDim()),
        sliceSizesArrayAttr);

    return success();
  }
};

/// Check if a `t` is a tensor with zero extents.
static std::optional<RankedTensorType> isZeroExtent(Type t) {
  auto type = dyn_cast<RankedTensorType>(t);
  if (type && type.hasStaticShape() && type.getNumElements() == 0)
    return type;
  return std::nullopt;
}

// Replace instances of zero extent tensors with empty tensors of the same
// type.
struct ZeroExtentTensorCanon final : RewritePattern {
  ZeroExtentTensorCanon(MLIRContext *context,
                        PatternBenefit benefit = PatternBenefit(1))
      : RewritePattern(MatchAnyOpTypeTag(), benefit, context) {}

  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const {
    auto disabledByAttr =
        failIfFuncOpInterfaceHasAttr(op, kDisablePatternAttrName, rewriter);
    if (disabledByAttr.failed())
      return disabledByAttr;

    auto loc = op->getLoc();

    if (!isa_and_present<stablehlo::StablehloDialect>(op->getDialect()))
      return rewriter.notifyMatchFailure(op, "not stablehlo");

    // If the result is a zero-extent tensor, replace the whole op with an empty
    // tensor.
    bool didUpdate = false;
    for (auto result : op->getResults()) {
      auto resultType = isZeroExtent(result.getType());
      if (!resultType || result.use_empty())
        continue;
      rewriter.replaceAllUsesWith(
          result, tensor::EmptyOp::create(rewriter, loc, resultType->getShape(),
                                          resultType->getElementType()));
      didUpdate = true;
    }

    // If one of the operands is a zero-extent tensor, replace the operand with
    // an empty tensor.
    for (OpOperand &operand : op->getOpOperands()) {
      auto operandType = isZeroExtent(operand.get().getType());
      if (!operandType || operand.get().getDefiningOp<tensor::EmptyOp>())
        continue;
      Operation *owner = operand.getOwner();
      int operandNum = operand.getOperandNumber();
      auto emptyTensorOp =
          tensor::EmptyOp::create(rewriter, loc, operandType->getShape(),
                                  operandType->getElementType());
      rewriter.modifyOpInPlace(
          owner, [&]() { owner->setOperand(operandNum, emptyTensorOp); });
      didUpdate = true;
    }
    return success(didUpdate);
  }
};

struct ReorderElementwiseAndShapeOp final
    : CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                   ReorderElementwiseAndShapeOp> {
  using CheckedOpTraitRewritePattern::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (op->getOperands().size() != 1)
      return rewriter.notifyMatchFailure(op, "expected to be unary");

    auto definingOp = op->getOperand(0).getDefiningOp();
    if (!definingOp)
      return rewriter.notifyMatchFailure(
          op, "expected to have an op before elementise op");

    if (!isa<stablehlo::ReshapeOp>(definingOp) &&
        !isa<stablehlo::TransposeOp>(definingOp) &&
        !isa<stablehlo::BroadcastOp>(definingOp))
      return rewriter.notifyMatchFailure(
          op, "defining operation of unexpected type");

    // Only reorder if the defining op has no other uses.
    if (!llvm::hasSingleElement(definingOp->getResult(0).getUses()))
      return rewriter.notifyMatchFailure(op, "operation has more than one use");

    Value input = definingOp->getOperand(0);
    Value result = op->getResult(0);
    auto intermediateType = cast<ShapedType>(input.getType())
                                .clone(getElementTypeOrSelf(result.getType()));

    // avoid reordering if we have high-priority fusions
    auto inputDefOp = input.getDefiningOp();
    if (inputDefOp) {
      if (auto reshapeDefOp = dyn_cast<stablehlo::ReshapeOp>(definingOp)) {
        if (isFusible(inputDefOp, reshapeDefOp)) {
          return failure();
        }
      } else if (auto bcastDefOp =
                     dyn_cast<stablehlo::BroadcastInDimOp>(definingOp)) {
        if (isFusible(inputDefOp, bcastDefOp)) {
          return failure();
        }
      }
    }

    // Reorder the operation and rewire the inputs/outputs.
    op->moveBefore(definingOp);
    definingOp->getResult(0).setType(result.getType());
    rewriter.replaceAllUsesWith(result, definingOp->getResult(0));
    result.setType(intermediateType);
    op->setOperands(input);
    definingOp->setOperands(result);
    return success();
  }
};

struct NoNanMulSimplify final
    : public NoNanCheckedOpRewritePattern<stablehlo::MulOp, NoNanMulSimplify> {
  using NoNanCheckedOpRewritePattern<
      stablehlo::MulOp, NoNanMulSimplify>::NoNanCheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    Attribute lhsAttr, rhsAttr;
    bool lhsConst = matchPattern(op.getLhs(), m_Constant(&lhsAttr));
    bool rhsConst = matchPattern(op.getRhs(), m_Constant(&rhsAttr));

    auto isZero = [](Attribute attr) {
      return matchPattern(attr, m_AnyZeroFloat()) ||
             matchPattern(attr, m_Zero()) ||
             matchPattern(attr, m_AnyZeroComplex());
    };

    // 0 * x -> 0
    if (lhsConst && isZero(lhsAttr)) {
      if (canApplyNoNanPattern(allowOnFloatingPointMath,
                               op.getResult().getType(),
                               op.getOperand(0).getType(), op, rewriter)) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            op, cast<ElementsAttr>(makeAttr(op.getType(), 0)));
        return success();
      }
    }
    // x * 0 -> 0
    if (rhsConst && isZero(rhsAttr)) {
      if (canApplyNoNanPattern(allowOnFloatingPointMath,
                               op.getResult().getType(),
                               op.getOperand(0).getType(), op, rewriter)) {
        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
            op, cast<ElementsAttr>(makeAttr(op.getType(), 0)));
        return success();
      }
    }

    return failure();
  }
};

// c = a + b; d = c - b => d = a
// c = a + b; d = b - c => d = -a
struct NoNanAddSubSimplify final
    : public NoNanCheckedOpRewritePattern<stablehlo::SubtractOp,
                                          NoNanAddSubSimplify> {
  using NoNanCheckedOpRewritePattern<
      stablehlo::SubtractOp, NoNanAddSubSimplify>::NoNanCheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SubtractOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();
    auto subOutTy = op.getResult().getType();

    // Check if LHS is defined by an AddOp
    if (auto lhsAddOp = lhs.getDefiningOp<stablehlo::AddOp>()) {
      auto addOutTy = lhsAddOp.getResult().getType();

      // Case: c = a + b; d = c - b -> d = a
      if (lhsAddOp.getRhs() == rhs) {
        if (canApplyNoNanPattern(allowOnFloatingPointMath, addOutTy, subOutTy,
                                 op, rewriter)) {
          rewriter.replaceOp(op, lhsAddOp.getLhs());
          return success();
        }
      }

      // Case: c = a + b; d = c - a -> d = b
      if (lhsAddOp.getLhs() == rhs) {
        if (canApplyNoNanPattern(allowOnFloatingPointMath, addOutTy, subOutTy,
                                 op, rewriter)) {
          rewriter.replaceOp(op, lhsAddOp.getRhs());
          return success();
        }
      }
    }

    // Check if RHS is defined by an AddOp
    if (auto rhsAddOp = rhs.getDefiningOp<stablehlo::AddOp>()) {
      auto addOutTy = rhsAddOp.getResult().getType();

      // Case: c = a + b; d = b - c -> d = -a
      if (rhsAddOp.getLhs() == lhs) {
        if (canApplyNoNanPattern(allowOnFloatingPointMath, addOutTy, subOutTy,
                                 op, rewriter)) {
          rewriter.replaceOpWithNewOp<stablehlo::NegOp>(op, rhsAddOp.getRhs());
          return success();
        }
      }

      // Case: c = a + b; d = a - c -> d = -b
      if (rhsAddOp.getRhs() == lhs) {
        if (canApplyNoNanPattern(allowOnFloatingPointMath, addOutTy, subOutTy,
                                 op, rewriter)) {
          rewriter.replaceOpWithNewOp<stablehlo::NegOp>(op, rhsAddOp.getLhs());
          return success();
        }
      }
    }

    // No simplification pattern matched
    return failure();
  }
};

// a > b ? a : b or a >= b ? a : b ---> maximum(a, b)
// a < b ? a : b or a <= b ? a : b ---> minimum(a, b)
struct CompareSelectSimplify
    : public CheckedOpRewritePattern<stablehlo::SelectOp,
                                     CompareSelectSimplify> {
  using CheckedOpRewritePattern<stablehlo::SelectOp,
                                CompareSelectSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    auto compOp = op.getPred().getDefiningOp<stablehlo::CompareOp>();
    if (!compOp)
      return failure();

    auto selectlhs = op.getOnTrue();
    auto selectrhs = op.getOnFalse();

    auto complhs = compOp.getLhs();
    auto comprhs = compOp.getRhs();

    if ((compOp.getComparisonDirection() ==
         stablehlo::ComparisonDirection::GT) ||
        (compOp.getComparisonDirection() ==
         stablehlo::ComparisonDirection::GE)) {
      // select(a > b || a >= b, a, b)
      if (complhs == selectlhs && comprhs == selectrhs) {
        rewriter.replaceOpWithNewOp<stablehlo::MaxOp>(op, selectlhs, selectrhs);
        return success();
      }
      // select(a > b || a >= b, b, a)
      if (complhs == selectrhs && comprhs == selectlhs) {
        rewriter.replaceOpWithNewOp<stablehlo::MinOp>(op, selectlhs, selectrhs);
        return success();
      }
    }

    if ((compOp.getComparisonDirection() ==
         stablehlo::ComparisonDirection::LT) ||
        (compOp.getComparisonDirection() ==
         stablehlo::ComparisonDirection::LE)) {
      // select(a < b || a <= b, a, b)
      if (complhs == selectlhs && comprhs == selectrhs) {
        rewriter.replaceOpWithNewOp<stablehlo::MinOp>(op, selectlhs, selectrhs);
        return success();
      }
      // select(a < b || a <= b, b, a)
      if (complhs == selectrhs && comprhs == selectlhs) {
        rewriter.replaceOpWithNewOp<stablehlo::MaxOp>(op, selectlhs, selectrhs);
        return success();
      }
    }

    return failure();
  }
};

// select(!op, lhs, rhs) --> select(op, rhs, lhs)
struct NotSelectSimplify
    : public CheckedOpRewritePattern<stablehlo::SelectOp, NotSelectSimplify> {
  using CheckedOpRewritePattern<stablehlo::SelectOp,
                                NotSelectSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    auto notOp = op.getPred().getDefiningOp<stablehlo::NotOp>();
    if (!notOp)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(
        op, notOp.getOperand(), op.getOnFalse(), op.getOnTrue());
    return success();
  }
};

struct NotCompare
    : public CheckedOpRewritePattern<stablehlo::NotOp, NotCompare> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::NotOp op,
                                    PatternRewriter &rewriter) const {
    auto cmp = op.getOperand().getDefiningOp<stablehlo::CompareOp>();
    if (!cmp)
      return failure();

    if (!llvm::hasSingleElement(cmp->getUsers()))
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::CompareOp>(
        op, cmp.getLhs(), cmp.getRhs(),
        negatedComparisonDirection(cmp.getComparisonDirection()),
        cmp.getCompareTypeAttr());

    rewriter.eraseOp(cmp);

    return success();
  }
};

struct CommonCompareExpressionRewrite
    : public CheckedOpRewritePattern<stablehlo::CompareOp,
                                     CommonCompareExpressionRewrite> {
  using CheckedOpRewritePattern<
      stablehlo::CompareOp,
      CommonCompareExpressionRewrite>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto negDir = negatedComparisonDirection(op.getComparisonDirection());

    // Check for equivalent users for the value with the fewest other users to
    // check. For ease here we simplify just checking if not constant (since
    // getNumUsers is O(n)).
    auto userCheck = lhs;
    if (matchPattern(userCheck, m_Constant()))
      userCheck = rhs;

    for (auto user : userCheck.getUsers()) {
      auto userCompareOp = dyn_cast<stablehlo::CompareOp>(user);
      if (!userCompareOp || userCompareOp.getComparisonDirection() != negDir)
        continue;

      if (user->getBlock() != op->getBlock())
        continue;

      if (userCompareOp.getLhs() == lhs && userCompareOp.getRhs() == rhs) {
        if (user->isBeforeInBlock(op)) {
          auto negatedCondition = stablehlo::NotOp::create(
              rewriter, op.getLoc(), userCompareOp.getResult());
          rewriter.replaceOp(op, negatedCondition);
          return success();
        } else {
          auto negatedCondition = stablehlo::NotOp::create(
              rewriter, userCompareOp.getLoc(), op.getResult());
          rewriter.replaceOp(user, negatedCondition);
          return success();
        }
      }
    }

    return failure();
  }
};

stablehlo::ComparisonDirection
reorderComparisionDirection(stablehlo::ComparisonDirection direction) {
  switch (direction) {
  case stablehlo::ComparisonDirection::EQ:
    return stablehlo::ComparisonDirection::EQ;
  case stablehlo::ComparisonDirection::NE:
    return stablehlo::ComparisonDirection::NE;
  case stablehlo::ComparisonDirection::GE:
    return stablehlo::ComparisonDirection::LE;
  case stablehlo::ComparisonDirection::GT:
    return stablehlo::ComparisonDirection::LT;
  case stablehlo::ComparisonDirection::LE:
    return stablehlo::ComparisonDirection::GE;
  case stablehlo::ComparisonDirection::LT:
    return stablehlo::ComparisonDirection::GT;
  default:
    llvm_unreachable("could not reorder comparison direction");
  }
}

struct CompareCleanup
    : public CheckedOpRewritePattern<stablehlo::CompareOp, CompareCleanup> {
  using CheckedOpRewritePattern<stablehlo::CompareOp,
                                CompareCleanup>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CompareOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();

    if (!cast<RankedTensorType>(lhs.getType())
             .getElementType()
             .isSignlessInteger(64)) {
      return failure();
    }

    DenseIntElementsAttr rhs;
    if (!matchPattern(op.getRhs(), m_Constant(&rhs))) {
      return failure();
    }

    if (!rhs.isSplat())
      return failure();

    auto rhsv = rhs.getSplatValue<IntegerAttr>().getValue().getSExtValue();

    if (auto add = lhs.getDefiningOp<stablehlo::AddOp>()) {
      DenseIntElementsAttr c;
      if (matchPattern(add.getRhs(), m_Constant(&c)) && c.isSplat()) {
        auto cv = c.getSplatValue<IntegerAttr>().getValue().getSExtValue();

        auto off = stablehlo::ConstantOp::create(
            rewriter, add.getLoc(), add.getType(),
            cast<ElementsAttr>(makeAttr(add.getType(), rhsv - cv)));

        // x + cv ?= rhsv -> x ?= rhs - cv
        rewriter.modifyOpInPlace(op, [&]() {
          op.getLhsMutable().assign(add.getLhs());
          op.getRhsMutable().assign(off);
        });

        return success();
      }
    }

    if (auto mul = lhs.getDefiningOp<stablehlo::MulOp>()) {
      DenseIntElementsAttr c;
      if (matchPattern(mul.getRhs(), m_Constant(&c)) && c.isSplat()) {
        auto cv = c.getSplatValue<IntegerAttr>().getValue().getSExtValue();
        if (cv == -1) {
          auto off = stablehlo::ConstantOp::create(
              rewriter, mul.getLoc(), mul.getType(),
              cast<ElementsAttr>(makeAttr(mul.getType(), -rhsv)));

          // x * -1 ?= rhsv -> x =? -rhs
          rewriter.modifyOpInPlace(op, [&]() {
            op.getLhsMutable().assign(mul.getLhs());
            op.getRhsMutable().assign(off);
            op.setComparisonDirection(
                reorderComparisionDirection(op.getComparisonDirection()));
          });

          return success();
        }
      }
    }

    return failure();
  }
};

struct ScatterUpdateComputationConstProp
    : public CheckedOpRewritePattern<stablehlo::ScatterOp,
                                     ScatterUpdateComputationConstProp> {
  using CheckedOpRewritePattern<
      stablehlo::ScatterOp,
      ScatterUpdateComputationConstProp>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp op,
                                    PatternRewriter &rewriter) const {
    if (!op.getUniqueIndices())
      return failure();

    auto &region = op.getUpdateComputation();
    auto &block = region.front();

    // Check all inputs are constant and splat and their values are the same.
    auto [constInput, inputSplatAttr] =
        isConstantSplatValueRange(op.getInputs());

    // Check all updates are constant and splat and their values are the same.
    auto [constUpdate, updateSplatAttr] =
        isConstantSplatValueRange(op.getUpdates());

    if (constInput || constUpdate) {
      bool inputTransformed = false;
      bool updateTransformed = false;
      auto blockArgInput = block.getArgument(0);
      auto blockArgUpdate = block.getArgument(1);

      if (constInput && !blockArgInput.getUses().empty()) {
        inputTransformed = true;
        auto denseAttr = DenseElementsAttr::get(
            cast<ShapedType>(blockArgInput.getType()), inputSplatAttr);
        auto constInputOp =
            stablehlo::ConstantOp::create(rewriter, op.getLoc(), denseAttr);
        blockArgInput.replaceAllUsesWith(constInputOp);
      }

      if (constUpdate && !blockArgUpdate.getUses().empty()) {
        updateTransformed = true;
        auto denseAttr = DenseElementsAttr::get(
            cast<ShapedType>(blockArgUpdate.getType()), updateSplatAttr);
        auto constUpdateOp =
            stablehlo::ConstantOp::create(rewriter, op.getLoc(), denseAttr);
        blockArgUpdate.replaceAllUsesWith(constUpdateOp);
      }

      if (!inputTransformed && !updateTransformed)
        return failure();

      auto newOp = stablehlo::ScatterOp::create(
          rewriter, op.getLoc(), op.getResultTypes(), op.getInputs(),
          op.getScatterIndices(), op.getUpdates(),
          op.getScatterDimensionNumbers(), op.getIndicesAreSorted(),
          op.getUniqueIndices());
      newOp.getUpdateComputation().takeBody(region);
      rewriter.replaceOp(op, newOp);

      return success();
    }

    return failure();
  }

private:
  std::tuple<bool, Attribute>
  isConstantSplatValueRange(ValueRange range) const {
    Attribute splatAttr = nullptr;
    bool isConstant = true;
    for (auto val : range) {
      DenseElementsAttr attr;
      if (matchPattern(val, m_Constant(&attr))) {
        if (attr.isSplat()) {
          if (!splatAttr) {
            splatAttr = attr.getSplatValue<Attribute>();
            continue;
          } else if (splatAttr != attr.getSplatValue<Attribute>()) {
            isConstant = false;
            break;
          }
        } else {
          isConstant = false;
          break;
        }
      } else {
        isConstant = false;
        break;
      }
    }
    return std::make_tuple(isConstant, splatAttr);
  };
};

struct ScatterIndicesAreUnique
    : public CheckedOpRewritePattern<stablehlo::ScatterOp,
                                     ScatterIndicesAreUnique> {
  using CheckedOpRewritePattern<
      stablehlo::ScatterOp, ScatterIndicesAreUnique>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getUniqueIndices())
      return failure(); // already unique, no need to do anything

    auto scatterIndices = op.getScatterIndices();
    auto dimNumbers = op.getScatterDimensionNumbers();
    Attribute scatterIndicesAttr;
    bool uniqueIndices = false;

    if (scatterIndices.getType().getNumElements() == 1) {
      uniqueIndices = true;
    } else if (auto iotaOp =
                   scatterIndices.getDefiningOp<stablehlo::IotaOp>()) {
      // IotaOp produces [0, 1, ..., N-1] along its iota_dimension and
      // repeats the same values along all other dimensions. We can only
      // mark indices as unique when every scatter point gets a distinct
      // index tuple.
      auto rank = scatterIndices.getType().getRank();
      if (rank == 1) {
        // 1-D iota is always unique: either it is a single index tuple
        // (indexVectorDim==0, one scatter point) or N scalar indices
        // [0, 1, ..., N-1] (indexVectorDim==1).
        uniqueIndices = true;
      } else {
        int64_t iotaDim = iotaOp.getIotaDimension();
        int64_t indexVectorDim = dimNumbers.getIndexVectorDim();
        if (iotaDim != indexVectorDim) {
          // Iota varies along a scatter (non-index-vector) dimension.
          // Each position along that axis gets a distinct iota value,
          // but positions differing only in *other* scatter dimensions
          // see the same value. Unique iff every other scatter dimension
          // (i.e. every dim that is neither iotaDim nor indexVectorDim)
          // has size 1.
          // Note: when indexVectorDim == rank (scalar indices), it never
          // equals any valid dim index, so the loop naturally checks all
          // dims except iotaDim.
          auto shape = scatterIndices.getType().getShape();
          bool allOtherDimsUnit = true;
          for (int64_t i = 0; i < rank; ++i) {
            if (i != iotaDim && i != indexVectorDim && shape[i] != 1) {
              allOtherDimsUnit = false;
              break;
            }
          }
          if (allOtherDimsUnit) {
            uniqueIndices = true;
          }
        }
        // If iotaDim == indexVectorDim, all scatter points read the same
        // tuple [0, 1, ..., shape[v]-1] → not unique (the single-
        // scatter-point case is already handled by getNumElements()==1).
      }
    } else if (matchPattern(scatterIndices, m_Constant(&scatterIndicesAttr))) {
      auto denseAttr = dyn_cast<DenseIntElementsAttr>(scatterIndicesAttr);

      auto shape = cast<ShapedType>(scatterIndices.getType()).getShape();
      if (shape.empty())
        return failure();

      int64_t indexVectorDim = dimNumbers.getIndexVectorDim();

      int64_t tupleSize = shape[indexVectorDim];

      SmallVector<int64_t> strides(shape.size());
      strides[shape.size() - 1] = 1;
      for (int64_t i = shape.size() - 2; i >= 0; --i) {
        strides[i] = strides[i + 1] * shape[i + 1];
      }

      SmallVector<int64_t> nonIndexVectorShape;
      for (int64_t i = 0; i < shape.size(); ++i) {
        if (i != indexVectorDim) {
          nonIndexVectorShape.push_back(shape[i]);
        }
      }

      // Iterate over the scatter indices tensor to extract tuples
      SmallVector<SmallVector<int64_t>> indexTuples;
      auto values = denseAttr.getValues<APInt>();

      std::function<void(SmallVector<int64_t>, int64_t)> extractTuples =
          [&](SmallVector<int64_t> currentIndices, int64_t dim) {
            if (dim == nonIndexVectorShape.size()) {
              SmallVector<int64_t> indexTuple;
              for (int64_t component = 0; component < tupleSize; ++component) {
                // Build full multi-dimensional index
                SmallVector<int64_t> fullIndex;
                int64_t nonIndexDim = 0;
                for (int64_t d = 0; d < shape.size(); ++d) {
                  if (d == indexVectorDim) {
                    fullIndex.push_back(component);
                  } else {
                    fullIndex.push_back(currentIndices[nonIndexDim++]);
                  }
                }

                // Convert to linear index
                int64_t linearIdx = 0;
                for (int64_t d = 0; d < shape.size(); ++d) {
                  linearIdx += fullIndex[d] * strides[d];
                }

                auto it = values.begin();
                std::advance(it, linearIdx);
                indexTuple.push_back((*it).getSExtValue());
              }
              indexTuples.push_back(indexTuple);
              return;
            }

            for (int64_t i = 0; i < nonIndexVectorShape[dim]; ++i) {
              SmallVector<int64_t> newIndices = currentIndices;
              newIndices.push_back(i);
              extractTuples(newIndices, dim + 1);
            }
          };

      extractTuples({}, 0);

      uniqueIndices = areIndexTuplesUnique(indexTuples);
    }

    if (!uniqueIndices) {
      return failure();
    }

    auto newOp = stablehlo::ScatterOp::create(
        rewriter, op.getLoc(), op.getResultTypes(), op.getInputs(),
        scatterIndices, op.getUpdates(), dimNumbers,
        op.getIndicesAreSortedAttr(), rewriter.getBoolAttr(uniqueIndices));
    newOp.getUpdateComputation().takeBody(op.getUpdateComputation());
    rewriter.replaceOp(op, newOp);
    return success();
  }

private:
  bool areIndexTuplesUnique(
      const SmallVector<SmallVector<int64_t>> &indexTuples) const {
    std::set<SmallVector<int64_t>> uniqueSet;
    for (const auto &tuple : indexTuples) {
      if (!uniqueSet.insert(tuple).second) {
        return false; // Duplicate found
      }
    }
    return true;
  }
};

// (add (mul a x) (mul a y)) -> (mul a (add x y))
template <typename Op>
struct AssociativeCommonMulOpReordering final
    : public CheckedOpRewritePattern<Op, AssociativeCommonMulOpReordering<Op>> {
  using CheckedOpRewritePattern<
      Op, AssociativeCommonMulOpReordering<Op>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(Op op, PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsMul = lhs.template getDefiningOp<stablehlo::MulOp>();
    auto rhsMul = rhs.template getDefiningOp<stablehlo::MulOp>();

    if (!lhsMul || !rhsMul)
      return failure();

    Value common = nullptr, lhsVal, rhsVal;

    if (lhsMul.getLhs() == rhsMul.getRhs()) {
      common = lhsMul.getLhs();
      rhsVal = rhsMul.getLhs();
      lhsVal = lhsMul.getRhs();
    }

    if (lhsMul.getLhs() == rhsMul.getLhs()) {
      common = lhsMul.getLhs();
      rhsVal = rhsMul.getRhs();
      lhsVal = lhsMul.getRhs();
    }

    if (lhsMul.getRhs() == rhsMul.getRhs()) {
      common = lhsMul.getRhs();
      rhsVal = rhsMul.getLhs();
      lhsVal = lhsMul.getLhs();
    }

    if (lhsMul.getRhs() == rhsMul.getLhs()) {
      common = lhsMul.getRhs();
      rhsVal = rhsMul.getRhs();
      lhsVal = lhsMul.getLhs();
    }

    if (!common)
      return failure();

    auto newMul = Op::create(rewriter, op.getLoc(), lhsVal, rhsVal);
    rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, common,
                                                  newMul.getResult());

    return success();
  }
};

// Helper to check if two SplatElementsAttrs are negations of each other.
static bool areSplatNegationsOf(SplatElementsAttr a, SplatElementsAttr b) {

  auto elementType = a.getElementType();
  if (elementType != b.getElementType())
    return false;

  if (isa<FloatType>(elementType)) {
    APFloat aVal = a.getSplatValue<APFloat>();
    APFloat bVal = b.getSplatValue<APFloat>();
    APFloat negA = aVal;
    negA.changeSign();
    return negA.bitwiseIsEqual(bVal);
  }

  if (isa<IntegerType>(elementType)) {
    APInt aVal = a.getSplatValue<APInt>();
    APInt bVal = b.getSplatValue<APInt>();
    return (-aVal) == bVal;
  }

  return false;
}

// Optimizes the following patterns:
// OuterOp=AddOp:
//   (a * cst) + (b * -cst) -> (a - b) * cst
// OuterOp=SubtractOp:
//   (a * cst) - (b * -cst) -> (a + b) * cst
//
// All commutative combinations of the multiply operands are considered:
//   (cst * a), (a * cst), (-cst * b), (b * -cst)
template <typename OuterOp>
struct NegatedConstantMulFactoring final
    : public CheckedOpRewritePattern<OuterOp,
                                     NegatedConstantMulFactoring<OuterOp>> {
  using CheckedOpRewritePattern<
      OuterOp, NegatedConstantMulFactoring<OuterOp>>::CheckedOpRewritePattern;

  // AddOp -> SubtractOp, SubtractOp -> AddOp
  using ResultOp = std::conditional_t<std::is_same_v<OuterOp, stablehlo::AddOp>,
                                      stablehlo::SubtractOp, stablehlo::AddOp>;

  LogicalResult matchAndRewriteImpl(OuterOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsMul = lhs.template getDefiningOp<stablehlo::MulOp>();
    auto rhsMul = rhs.template getDefiningOp<stablehlo::MulOp>();

    if (!lhsMul || !rhsMul) {
      return failure();
    }

    // For each mul op, try to extract (nonConst, constAttr) where the constant
    // side is a splat DenseElementsAttr. We try both operand orders.
    struct MulInfo {
      Value nonConst;
      Value constOperand;
      SplatElementsAttr constAttr;
    };

    auto extractMulInfo = [](stablehlo::MulOp mul) -> std::optional<MulInfo> {
      SplatElementsAttr attr;
      // Try rhs as constant
      if (matchPattern(mul.getRhs(), m_Constant(&attr))) {
        return MulInfo{mul.getLhs(), mul.getRhs(), attr};
      }
      // Try lhs as constant
      if (matchPattern(mul.getLhs(), m_Constant(&attr))) {
        return MulInfo{mul.getRhs(), mul.getLhs(), attr};
      }
      return std::nullopt;
    };

    auto lhsInfo = extractMulInfo(lhsMul);
    auto rhsInfo = extractMulInfo(rhsMul);

    if (!lhsInfo || !rhsInfo) {
      return failure();
    }

    // Check if the two constants are negations of each other.
    // areSplatNegationsOf checks if lhsConst == -rhsConst.
    // This is symmetric, so we only need one check.
    if (!areSplatNegationsOf(lhsInfo->constAttr, rhsInfo->constAttr)) {
      return failure();
    }

    // For OuterOp=AddOp (ResultOp=SubtractOp):
    //   a * lhsConst + b * rhsConst  (where rhsConst = -lhsConst)
    //   = a * lhsConst - b * lhsConst = (a - b) * lhsConst
    // For OuterOp=SubtractOp (ResultOp=AddOp):
    //   a * lhsConst - b * rhsConst  (where rhsConst = -lhsConst)
    //   = a * lhsConst + b * lhsConst = (a + b) * lhsConst
    auto newBinop = ResultOp::create(rewriter, op.getLoc(), lhsInfo->nonConst,
                                     rhsInfo->nonConst);
    rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, newBinop.getResult(),
                                                  lhsInfo->constOperand);
    return success();
  }
};

// This lets us reorder the following
// Case 1: (op x (op (op y x) y)) -> (op (op x y) (op x y))
// Case 2: (op x (op (op x y) y)) -> (op (op x y) (op x y))
// Case 3: (op x (op y (op x y))) -> (op (op x y) (op x y))
// Case 4: (op x (op y (op y x))) -> (op (op x y) (op x y))
// Case 5: (op (op x cst1) cst2) -> (op x (op cst1 cst2))
// Case 6: (op (op cst1 x) cst2) -> (op x (op cst1 cst2)) only if op is
// commutative
// Case 7: (op cst2 (op x cst1)) -> (op (op cst1 cst2) x) only if op is
// commutative
// Case 8: (op cst2 (op cst1 x)) -> (op (op cst1 cst2) x)
struct AssociativeBinaryOpReordering
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          AssociativeBinaryOpReordering> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise,
      AssociativeBinaryOpReordering>::CheckedOpTraitRewritePattern;

  // TODO: generalize to the case where lhs is an Op
  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() != 2 || !stablehlo::isAssociativeOp(op))
      return failure();

    auto lhs = op->getOperand(0);
    auto rhs = op->getOperand(1);
    auto lhsOp = lhs.getDefiningOp();
    auto rhsOp = rhs.getDefiningOp();

    if (rhsOp && rhsOp->getName() == op->getName()) {
      auto rhslhs = rhsOp->getOperand(0);
      auto rhsrhs = rhsOp->getOperand(1);

      auto rhslhsOp = rhslhs.getDefiningOp();
      if (rhslhsOp && rhslhsOp->getName() == op->getName()) {
        auto rhslhslhs = rhslhsOp->getOperand(0);
        auto rhslhsrhs = rhslhsOp->getOperand(1);

        // Case 1 || Case 2
        if ((lhs == rhslhsrhs && rhslhslhs == rhsrhs) ||
            (lhs == rhslhslhs && rhslhsrhs == rhsrhs)) {
          rewriter.modifyOpInPlace(op, [&] {
            op->setOperand(0, rhslhsOp->getResult(0));
            op->setOperand(1, rhslhsOp->getResult(0));
          });
          return success();
        }
      }

      auto rhsrhsOp = rhsrhs.getDefiningOp();
      if (rhsrhsOp && rhsrhsOp->getName() == op->getName()) {
        auto rhsrhslhs = rhsrhsOp->getOperand(0);
        auto rhsrhsrhs = rhsrhsOp->getOperand(1);

        // Case 3 || Case 4
        if ((lhs == rhsrhslhs && rhslhs == rhsrhsrhs) ||
            (lhs == rhsrhsrhs && rhslhs == rhsrhslhs)) {
          rewriter.modifyOpInPlace(op, [&] {
            op->setOperand(0, rhsrhsOp->getResult(0));
            op->setOperand(1, rhsrhsOp->getResult(0));
          });
          return success();
        }
      }
    }

    bool isCommutative = op->hasTrait<mlir::hlo::OpTrait::IsCommutative>() ||
                         op->hasTrait<mlir::OpTrait::IsCommutative>();

    auto isConstant = [](Value v) -> bool {
      return matchPattern(v, m_Constant());
    };

    // Case 5: (op (op x cst1) cst2) -> (op x (op cst1 cst2))
    if (lhsOp && lhsOp->getName() == op->getName() && isConstant(rhs)) {
      auto innerRhs = lhsOp->getOperand(1);
      auto innerLhs = lhsOp->getOperand(0);
      if (isConstant(innerRhs) && !isConstant(innerLhs)) {
        // Create (op cst1 cst2) as the new rhs
        auto cstCombined = rewriter.clone(*op);
        cstCombined->setOperand(0, innerRhs);
        cstCombined->setOperand(1, rhs);
        rewriter.modifyOpInPlace(op, [&] {
          op->setOperand(0, innerLhs);
          op->setOperand(1, cstCombined->getResult(0));
        });
        return success();
      }
    }

    // Case 6: (op (op cst1 x) cst2) -> (op x (op cst1 cst2))
    // only if op is commutative
    if (isCommutative && lhsOp && lhsOp->getName() == op->getName() &&
        isConstant(rhs)) {
      auto innerLhs = lhsOp->getOperand(0);
      auto innerRhs = lhsOp->getOperand(1);
      if (isConstant(innerLhs) && !isConstant(innerRhs)) {
        // Create (op cst1 cst2) as the new rhs
        auto cstCombined = rewriter.clone(*op);
        cstCombined->setOperand(0, innerLhs);
        cstCombined->setOperand(1, rhs);
        rewriter.modifyOpInPlace(op, [&] {
          op->setOperand(0, innerRhs);
          op->setOperand(1, cstCombined->getResult(0));
        });
        return success();
      }
    }

    // Case 7: (op cst2 (op x cst1)) -> (op (op cst2 cst1) x)
    // only if op is commutative
    if (isCommutative && rhsOp && rhsOp->getName() == op->getName() &&
        isConstant(lhs)) {
      auto innerRhs = rhsOp->getOperand(1);
      auto innerLhs = rhsOp->getOperand(0);
      if (isConstant(innerRhs) && !isConstant(innerLhs)) {
        // Create (op cst2 cst1) as the new lhs
        auto cstCombined = rewriter.clone(*op);
        cstCombined->setOperand(0, lhs);
        cstCombined->setOperand(1, innerRhs);
        rewriter.modifyOpInPlace(op, [&] {
          op->setOperand(0, cstCombined->getResult(0));
          op->setOperand(1, innerLhs);
        });
        return success();
      }
    }

    // Case 8: (op cst2 (op cst1 x)) -> (op (op cst2 cst1) x)
    if (rhsOp && rhsOp->getName() == op->getName() && isConstant(lhs)) {
      auto innerLhs = rhsOp->getOperand(0);
      auto innerRhs = rhsOp->getOperand(1);
      if (isConstant(innerLhs) && !isConstant(innerRhs)) {
        // Create (op cst2 cst1) as the new lhs
        auto cstCombined = rewriter.clone(*op);
        cstCombined->setOperand(0, lhs);
        cstCombined->setOperand(1, innerLhs);
        rewriter.modifyOpInPlace(op, [&] {
          op->setOperand(0, cstCombined->getResult(0));
          op->setOperand(1, innerRhs);
        });
        return success();
      }
    }

    return failure();
  }
};

struct ReduceTransposeSimplify
    : public CheckedOpRewritePattern<stablehlo::ReduceOp,
                                     ReduceTransposeSimplify> {
  using CheckedOpRewritePattern<
      stablehlo::ReduceOp, ReduceTransposeSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1) // TODO: support for multiple inputs
      return failure();

    auto input = op.getInputs()[0];
    auto transposeOp = input.getDefiningOp<stablehlo::TransposeOp>();
    if (!transposeOp)
      return failure();

    auto transposeInput = transposeOp.getOperand();
    auto transposeInputType = cast<ShapedType>(transposeInput.getType());
    auto transposePermutation = transposeOp.getPermutation();
    auto reduceDimensions = op.getDimensions();

    SmallVector<int64_t> newReduceDimensions;
    for (auto dim : reduceDimensions) {
      int64_t newDim = transposePermutation[dim];
      newReduceDimensions.push_back(newDim);
    }

    // Calculate result shape after reduction
    SmallVector<int64_t> resultShape;
    for (int64_t i = 0; i < transposeInputType.getRank(); ++i) {
      if (!llvm::is_contained(newReduceDimensions, i)) {
        resultShape.push_back(transposeInputType.getDimSize(i));
      }
    }
    auto elementType = cast<ShapedType>(op.getType(0)).getElementType();
    auto newResultType = RankedTensorType::get(resultShape, elementType);

    // Create a new reduce operation with the adjusted dimensions
    auto newReduceOp = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), TypeRange(newResultType),
        ValueRange(transposeInput), op.getInitValues(), newReduceDimensions);
    newReduceOp.getRegion().takeBody(op.getRegion());

    // Map non-reduced dimensions
    SmallVector<bool> isReduced(transposeInputType.getRank(), false);
    for (auto dim : newReduceDimensions) {
      isReduced[dim] = true;
    }

    // Count the number of reduced dimensions before the transpose dim
    SmallVector<int64_t> reducedDimsBeforeTranspose(
        transposePermutation.size());
    reducedDimsBeforeTranspose[0] = isReduced[0];
    for (int64_t i = 1; i < transposePermutation.size(); ++i) {
      reducedDimsBeforeTranspose[i] =
          reducedDimsBeforeTranspose[i - 1] + isReduced[i];
    }

    // Create final permutation
    // original permutation - reduced dimensions before transposed dim
    SmallVector<int64_t> finalPermutation(transposePermutation.size() -
                                          newReduceDimensions.size());
    int64_t j = 0;
    for (int64_t i = 0; i < transposePermutation.size(); ++i) {
      if (!isReduced[transposePermutation[i]]) {
        finalPermutation[j] =
            transposePermutation[i] -
            reducedDimsBeforeTranspose[transposePermutation[i]];
        ++j;
      }
    }

    rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(
        op, newReduceOp.getResult(0), finalPermutation);
    return success();
  }
};

// (mul (sign x) (abs x)) -> x
// (mul (abs x) (sign x)) -> x
struct SignAbsSimplify
    : public CheckedOpRewritePattern<stablehlo::MulOp, SignAbsSimplify> {
  using CheckedOpRewritePattern<stablehlo::MulOp,
                                SignAbsSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getOperand(0);
    auto rhs = op.getOperand(1);

    auto lhsSignOp = lhs.getDefiningOp<stablehlo::SignOp>();
    if (lhsSignOp) {
      auto rhsAbsOp = rhs.getDefiningOp<stablehlo::AbsOp>();
      if (!rhsAbsOp)
        return failure();

      if (lhsSignOp.getOperand() != rhsAbsOp.getOperand())
        return failure();

      rewriter.replaceOp(op, lhsSignOp.getOperand());
      return success();
    }

    auto rhsSignOp = rhs.getDefiningOp<stablehlo::SignOp>();
    if (rhsSignOp) {
      auto lhsAbsOp = lhs.getDefiningOp<stablehlo::AbsOp>();
      if (!lhsAbsOp)
        return failure();

      if (rhsSignOp.getOperand() != lhsAbsOp.getOperand())
        return failure();

      rewriter.replaceOp(op, rhsSignOp.getOperand());
      return success();
    }

    return failure();
  }
};

struct AbsPositiveSimplify
    : public CheckedOpRewritePattern<stablehlo::AbsOp, AbsPositiveSimplify> {
  using CheckedOpRewritePattern<stablehlo::AbsOp,
                                AbsPositiveSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AbsOp op,
                                    PatternRewriter &rewriter) const {

    auto operand = op.getOperand();
    if (isa<ComplexType>(operand.getType().getElementType()))
      return failure();

    if (guaranteedNonNegativeResult(operand.getDefiningOp(), rewriter)) {
      rewriter.replaceOp(op, op.getOperand());
      return success();
    }
    return failure();
  }
};

struct TransposeReshapeToBroadcast final
    : CheckedOpRewritePattern<stablehlo::TransposeOp,
                              TransposeReshapeToBroadcast> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto reshapeOp = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshapeOp)
      return failure();

    RankedTensorType reshapeOpInputType = reshapeOp.getOperand().getType();
    RankedTensorType reshapeOpOutputType = reshapeOp.getResult().getType();

    SmallVector<int64_t> insertionDims =
        findReshapeInsertionDims(reshapeOpInputType, reshapeOpOutputType);
    if (insertionDims.empty()) {
      return failure();
    }
    llvm::sort(insertionDims);

    auto permutation = op.getPermutation();
    if (permutation.size() != reshapeOpOutputType.getRank()) {
      return failure();
    }

    SmallVector<int64_t> broadcastDimensions;
    for (int64_t i = 0; i < reshapeOpOutputType.getRank(); ++i) {
      if (llvm::is_contained(insertionDims, i)) {
        continue;
      }

      auto it = llvm::find(permutation, i);
      if (it == permutation.end()) {
        return failure(); // The index was not found in the permutation
      }
      int64_t outputIdx = std::distance(permutation.begin(), it);
      broadcastDimensions.push_back(outputIdx);
    }

    // Create a single broadcast_in_dim operation to replace the reshape +
    // transpose sequence.
    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getResult().getType(), reshapeOp.getOperand(),
        rewriter.getDenseI64ArrayAttr(broadcastDimensions));

    return success();
  }
};

struct ReshapeTransposeToBroadcast final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp,
                              ReshapeTransposeToBroadcast> {
  using CheckedOpRewritePattern<
      stablehlo::ReshapeOp,
      ReshapeTransposeToBroadcast>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto transposeOp = op.getOperand().getDefiningOp<stablehlo::TransposeOp>();
    if (!transposeOp)
      return failure();

    RankedTensorType reshapeOpInputType = op.getOperand().getType();
    RankedTensorType reshapeOpOutputType = op.getResult().getType();

    SmallVector<int64_t> insertionDims =
        findReshapeInsertionDims(reshapeOpInputType, reshapeOpOutputType);
    if (insertionDims.empty()) {
      return failure();
    }

    auto permutation = transposeOp.getPermutation();

    SmallVector<int64_t> broadcastDimensions(reshapeOpInputType.getRank());
    int64_t origTensorDim = 0;
    for (int64_t i = 0; i < reshapeOpOutputType.getRank(); ++i) {
      if (llvm::is_contained(insertionDims, i)) {
        continue;
      }

      broadcastDimensions[permutation[origTensorDim]] = i;
      origTensorDim++;
    }

    // Create a single broadcast_in_dim operation to replace the reshape +
    // transpose sequence.
    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getResult().getType(), transposeOp.getOperand(),
        rewriter.getDenseI64ArrayAttr(broadcastDimensions));
    return success();
  }
};

struct BroadcastInDimIsReshape final
    : CheckedOpRewritePattern<stablehlo::BroadcastInDimOp,
                              BroadcastInDimIsReshape> {
  using CheckedOpRewritePattern<
      stablehlo::BroadcastInDimOp,
      BroadcastInDimIsReshape>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::BroadcastInDimOp op,
                                    PatternRewriter &rewriter) const {
    if (stablehlo::OpIsReshapeLike(op)) {
      rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                        op.getOperand());
      return success();
    }
    return failure();
  }
};

struct ReshapeToBroadcast final
    : CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeToBroadcast> {
  using CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                ReshapeToBroadcast>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp reshape,
                                    PatternRewriter &rewriter) const {

    auto inShape = reshape.getOperand().getType().getShape();
    auto outShape = reshape.getResult().getType().getShape();
    if (inShape.size() + 1 != outShape.size())
      return failure();

    if (outShape[0] != 1) {
      return failure();
    }
    for (auto &&[lhs, rhs] : llvm::zip_equal(outShape.slice(1), inShape)) {
      if (lhs != rhs) {
        return failure();
      }
    }
    SmallVector<int64_t> vals(inShape.size(), 0);
    for (int i = 0; i < inShape.size(); i++) {
      vals[i] = i + 1;
    }
    auto shard = sdy::getShardingPerValue(reshape);
    auto rep = rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        reshape, reshape.getType(), reshape.getOperand(), vals);
    if (shard) {
      sdy::setShardings(rep, shard);
    }
    return success();
  }
};

struct PadConcatToConcatPad
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     PadConcatToConcatPad> {
  using CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                PadConcatToConcatPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concatOp,
                                    PatternRewriter &rewriter) const {

    if (concatOp.getNumOperands() <= 1) {
      return failure();
    }

    // Check if all operands are pad ops with the same padding value
    SmallVector<stablehlo::PadOp> padOps;
    Value padValue;

    for (Value operand : concatOp.getOperands()) {
      auto padOp = operand.getDefiningOp<stablehlo::PadOp>();
      if (!padOp)
        return failure();

      if (!padOp->hasOneUse())
        return failure();

      if (padOps.empty()) {
        padValue = padOp.getPaddingValue();
      } else if (padValue != padOp.getPaddingValue()) {
        return failure(); // Different padding values not supported
      }

      padOps.push_back(padOp);
    }

    int64_t concatDim = concatOp.getDimension();
    int64_t rank = padOps[0].getEdgePaddingLow().size();

    if (rank <= 0)
      return failure();

    // Compute smallest common padding for all tensors
    SmallVector<int64_t> commonLowPadding(rank,
                                          std::numeric_limits<int64_t>::max());
    SmallVector<int64_t> commonHighPadding(rank,
                                           std::numeric_limits<int64_t>::max());
    SmallVector<int64_t> interiorPadding(rank, 0);

    // Find minimum padding across all inputs (conservative common padding)
    for (auto padOp : padOps) {
      for (int64_t dim = 0; dim < rank; ++dim) {
        commonLowPadding[dim] =
            std::min(commonLowPadding[dim], padOp.getEdgePaddingLow()[dim]);
        commonHighPadding[dim] =
            std::min(commonHighPadding[dim], padOp.getEdgePaddingHigh()[dim]);
      }
    }

    bool commonPad = false;

    for (int64_t dim = 0; dim < rank; ++dim) {
      if (commonLowPadding[dim] != 0 || commonHighPadding[dim] != 0) {
        commonPad = true;
        break;
      }
    }

    if (!commonPad) {
      return failure();
    }

    // Collect original operands with adjusted padding
    SmallVector<Value> adjOperands;

    for (auto padOp : padOps) {

      SmallVector<int64_t> diffLowPadding(rank);
      SmallVector<int64_t> diffHighPadding(rank);

      for (int64_t dim = 0; dim < rank; ++dim) {
        diffLowPadding[dim] =
            padOp.getEdgePaddingLow()[dim] - commonLowPadding[dim];
        diffHighPadding[dim] =
            padOp.getEdgePaddingHigh()[dim] - commonHighPadding[dim];
      }

      bool needsExtraPad = false;
      for (int64_t dim = 0; dim < rank; ++dim) {
        if (diffLowPadding[dim] > 0 || diffHighPadding[dim] > 0) {
          needsExtraPad = true;
          break;
        }
      }

      if (needsExtraPad) {

        auto adjustedOp = stablehlo::PadOp::create(
            rewriter, padOp.getLoc(),
            padOp.getOperand(), // we pad the input operand
            padOp.getPaddingValue(), diffLowPadding, diffHighPadding,
            padOp.getInteriorPaddingAttr());

        adjOperands.push_back(adjustedOp);
      } else {
        // No extra padding needed, use original tensor
        adjOperands.push_back(padOp.getOperand());
      }
    }

    auto newConcatOp = stablehlo::ConcatenateOp::create(
        rewriter, concatOp.getLoc(), adjOperands, concatDim);

    // Apply the common padding to get the final result
    auto result = stablehlo::PadOp::create(
        rewriter, concatOp.getLoc(), newConcatOp, padValue, commonLowPadding,
        commonHighPadding, interiorPadding);

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

struct SelectPad final
    : public CheckedOpRewritePattern<stablehlo::SelectOp, SelectPad> {
  using CheckedOpRewritePattern<stablehlo::SelectOp,
                                SelectPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    // Get the two pad operations from the on_true and on_false branches.
    auto padTrue = op.getOnTrue().getDefiningOp<stablehlo::PadOp>();
    if (!padTrue || anyPadSizesNegative(padTrue))
      return failure();

    auto padFalse = op.getOnFalse().getDefiningOp<stablehlo::PadOp>();
    if (!padFalse || anyPadSizesNegative(padFalse))
      return failure();

    // Ensure the padding configurations are identical.
    if (padTrue.getEdgePaddingLow() != padFalse.getEdgePaddingLow() ||
        padTrue.getEdgePaddingHigh() != padFalse.getEdgePaddingHigh() ||
        padTrue.getInteriorPadding() != padFalse.getInteriorPadding()) {
      return rewriter.notifyMatchFailure(op, "padding configurations mismatch");
    }

    // Only proceed if the pad ops have a single user.
    if (!padTrue->hasOneUse() || !padFalse->hasOneUse()) {
      return rewriter.notifyMatchFailure(op, "pad op has multiple users");
    }

    auto operandTrue = padTrue.getOperand();
    auto operandFalse = padFalse.getOperand();

    auto pred = op.getPred();
    auto predType = cast<RankedTensorType>(pred.getType());

    Value newOperandSelect;
    Value newPadValSelect;

    // Case 1: Predicate is a scalar (0-rank tensor).
    if (predType.getRank() == 0) {
      newOperandSelect = stablehlo::SelectOp::create(
          rewriter, op.getLoc(), pred, operandTrue, operandFalse);

      auto padValTrue = padTrue.getPaddingValue();
      auto padValFalse = padFalse.getPaddingValue();

      newPadValSelect = stablehlo::SelectOp::create(rewriter, op.getLoc(), pred,
                                                    padValTrue, padValFalse);

      // Case 2: Predicate is a tensor, but padding values are identical.
    } else {
      auto padValTrue = padTrue.getPaddingValue();
      auto padValFalse = padFalse.getPaddingValue();
      if (padValTrue != padValFalse) {
        return rewriter.notifyMatchFailure(
            op, "selecting different padding values with a tensor predicate is "
                "not supported");
      }
      newPadValSelect = padValTrue; // Use the common padding value.

      // To create the new select for operands, the predicate needs to be
      // "un-padded" to match the shape of the operands. This is only
      // straightforward if interior padding is zero.
      if (llvm::any_of(padTrue.getInteriorPadding(),
                       [](int64_t p) { return p != 0; })) {
        return rewriter.notifyMatchFailure(
            op, "interior padding is not supported for tensor predicates");
      }

      auto operandType = cast<RankedTensorType>(operandTrue.getType());
      auto starts = padTrue.getEdgePaddingLow();
      auto limits = llvm::to_vector(starts);
      for (size_t i = 0; i < limits.size(); ++i) {
        limits[i] += operandType.getShape()[i];
      }
      auto strides = SmallVector<int64_t>(operandType.getRank(), 1);

      auto slicedPred = stablehlo::SliceOp::create(rewriter, op.getLoc(), pred,
                                                   starts, limits, strides);

      newOperandSelect = stablehlo::SelectOp::create(
          rewriter, op.getLoc(), slicedPred, operandTrue, operandFalse);
    }

    // Create the new pad operation.
    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, op.getType(), newOperandSelect, newPadValSelect,
        padTrue.getEdgePaddingLowAttr(), padTrue.getEdgePaddingHighAttr(),
        padTrue.getInteriorPaddingAttr());

    return success();
  }
};

struct SliceSelect
    : public CheckedOpRewritePattern<stablehlo::SliceOp, SliceSelect> {
  using CheckedOpRewritePattern<stablehlo::SliceOp,
                                SliceSelect>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp sliceOp,
                                    PatternRewriter &rewriter) const {

    auto selOp = sliceOp.getOperand().getDefiningOp<stablehlo::SelectOp>();

    if (!selOp)
      return failure();

    if (!selOp->hasOneUse())
      return failure();

    bool scalar_pred = false;
    Value pred = selOp.getPred();
    Value on_true = selOp.getOnTrue();
    Value on_false = selOp.getOnFalse();

    if (dyn_cast<RankedTensorType>(pred.getType()).getRank() == 0) {
      scalar_pred = true;
    }

    Value slicedPred;
    if (!scalar_pred) {
      // slice predicate
      slicedPred = stablehlo::SliceOp::create(
          rewriter, sliceOp.getLoc(), pred, sliceOp.getStartIndices(),
          sliceOp.getLimitIndices(), sliceOp.getStrides());
    } else {
      slicedPred = pred;
    }
    Value slicedOnTrue = stablehlo::SliceOp::create(
        rewriter, sliceOp.getLoc(), on_true, sliceOp.getStartIndices(),
        sliceOp.getLimitIndices(), sliceOp.getStrides());

    Value slicedOnFalse = stablehlo::SliceOp::create(
        rewriter, sliceOp.getLoc(), on_false, sliceOp.getStartIndices(),
        sliceOp.getLimitIndices(), sliceOp.getStrides());

    auto newSelectOp = stablehlo::SelectOp::create(
        rewriter, sliceOp.getLoc(), slicedPred, slicedOnTrue, slicedOnFalse);

    rewriter.replaceOp(sliceOp, newSelectOp.getResult());

    return success();
    ;
  }
};

struct ConstPadConcatToConcat
    : public CheckedOpRewritePattern<stablehlo::PadOp, ConstPadConcatToConcat> {
  using CheckedOpRewritePattern<
      stablehlo::PadOp, ConstPadConcatToConcat>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::PadOp padOp,
                                    PatternRewriter &rewriter) const {
    auto concatOp =
        padOp.getOperand().getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concatOp)
      return failure();

    if (!concatOp->hasOneUse())
      return failure();

    DenseElementsAttr padConst;
    if (!matchPattern(padOp.getPaddingValue(), m_Constant(&padConst)))
      return failure();

    // Ensure interior padding is all zero.
    for (int64_t ip : padOp.getInteriorPadding())
      if (ip != 0)
        return failure();

    // Count how many dimensions have nonzero pad.
    int nonZeroPads = 0;
    int padDim = -1;
    for (int i = 0, e = padOp.getEdgePaddingLow().size(); i < e; ++i) {
      int64_t low = padOp.getEdgePaddingLow()[i];
      int64_t high = padOp.getEdgePaddingHigh()[i];
      if (low != 0 || high != 0) {
        ++nonZeroPads;
        padDim = i;
      }
    }
    // Only handle the case with a single padded dimension.
    if (nonZeroPads != 1)
      return failure();

    // Ensure the padded dimension (padDim) matches the concatenate op's
    // dimension.
    if (padDim != concatOp.getDimension())
      return failure();

    // For simplicity, require padding on only one side.
    bool padAtLow = (padOp.getEdgePaddingLow()[padDim] > 0);
    bool padAtHigh = (padOp.getEdgePaddingHigh()[padDim] > 0);
    if (padAtLow && padAtHigh)
      return failure();

    // Create a constant tensor to replace the pad.
    // Its shape is the same as the output except in padDim where its size
    // equals the pad amount.
    auto outShape = padOp.getType().getShape();
    int64_t padAmount = padAtLow ? padOp.getEdgePaddingLow()[padDim]
                                 : padOp.getEdgePaddingHigh()[padDim];
    SmallVector<int64_t> constShape(outShape.begin(), outShape.end());
    constShape[padDim] = padAmount;
    auto elemType = padOp.getType().getElementType();
    auto constTensorType = RankedTensorType::get(constShape, elemType);
    // Use the pad constant (splat) to create the constant tensor.
    auto newSplattedConstOp =
        stablehlo::ConstantOp::create(rewriter, padOp.getLoc(), constTensorType,
                                      padConst.resizeSplat(constTensorType));

    // Now, instead of padding the concatenation result, we insert the pad
    // slice via a new concatenate.
    auto origOperands = concatOp.getOperands();
    SmallVector<Value, 4> newOperands(origOperands.begin(), origOperands.end());
    if (padAtLow) {
      // Insert the pad slice at the beginning.
      newOperands.insert(newOperands.begin(), newSplattedConstOp);
    } else {
      // Insert the pad slice at the end.
      newOperands.push_back(newSplattedConstOp);
    }

    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
        padOp, padOp.getResult().getType(), newOperands,
        concatOp.getDimension());
    return success();
  }
};

template <typename T> struct Term {
  APFloat constantFactor;
  Value valFactor;
  T term;
  Term(APFloat constantFactor, Value valFactor, T term)
      : constantFactor(constantFactor), valFactor(valFactor), term(term) {}
};

template <typename ST, typename Child>
struct SumToReductionBase
    : public CheckedOpRewritePattern<ST, SumToReductionBase<ST, Child>> {
  using CheckedOpRewritePattern<
      ST, SumToReductionBase<ST, Child>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(ST op, PatternRewriter &rewriter) const {
    if (!op.getType().getElementType().isFloat())
      return failure();

    SmallVector<Term<stablehlo::SliceOp>> done0;
    SmallVector<Term<Value>> todo;
    APFloat zero(
        cast<FloatType>(op.getType().getElementType()).getFloatSemantics(),
        "0");
    APFloat one(
        cast<FloatType>(op.getType().getElementType()).getFloatSemantics(),
        "1");
    todo.emplace_back(one, (Value) nullptr, (Value)op.getResult());
    SmallVector<Value> intermediates;
    SmallPtrSet<Operation *, 1> seen;
    while (!todo.empty()) {
      auto cur = todo.pop_back_val();
      if (auto add = cur.term.getDefiningOp<stablehlo::AddOp>()) {
        todo.emplace_back(cur.constantFactor, cur.valFactor, add.getLhs());
        todo.emplace_back(cur.constantFactor, cur.valFactor, add.getRhs());
        intermediates.push_back(add);
        seen.insert(add);
        continue;
      }
      if (auto sub = cur.term.getDefiningOp<stablehlo::SubtractOp>()) {
        todo.emplace_back(cur.constantFactor, cur.valFactor, sub.getLhs());
        todo.emplace_back(-cur.constantFactor, cur.valFactor, sub.getRhs());
        intermediates.push_back(sub);
        seen.insert(sub);
        continue;
      }
      if (auto neg = cur.term.getDefiningOp<stablehlo::NegOp>()) {
        todo.emplace_back(-cur.constantFactor, cur.valFactor, neg.getOperand());
        intermediates.push_back(neg);
        seen.insert(neg);
        continue;
      }
      if (auto mul = cur.term.getDefiningOp<stablehlo::MulOp>()) {
        bool legal = false;
        for (int i = 0; i < 2; i++) {
          SplatElementsAttr other;
          if (!matchPattern(mul->getOperand(1 - i), m_Constant(&other)))
            continue;
          todo.emplace_back(cur.constantFactor * other.getSplatValue<APFloat>(),
                            cur.valFactor, mul->getOperand(i));
          legal = true;
          break;
        }
        if (legal) {
          intermediates.push_back(mul);
          seen.insert(mul);
          continue;
        }
      }
      if (auto div = cur.term.getDefiningOp<stablehlo::DivOp>()) {
        SplatElementsAttr other;
        if (matchPattern(div.getRhs(), m_Constant(&other))) {
          todo.emplace_back(cur.constantFactor / other.getSplatValue<APFloat>(),
                            cur.valFactor, div.getLhs());
          intermediates.push_back(div);
          seen.insert(div);
          continue;
        }
      }
      if (auto sl = cur.term.getDefiningOp<stablehlo::SliceOp>()) {
        done0.emplace_back(cur.constantFactor, cur.valFactor, sl);
        seen.insert(sl);
        continue;
      }
      if (auto dfop = cur.term.getDefiningOp())
        return rewriter.notifyMatchFailure(
            dfop, "operand is not a linear slice term");
      else
        return rewriter.notifyMatchFailure(
            op, "had operand that is not a linear slice term");
    }

    if (done0.size() < 2)
      return failure();

    auto T = op.getType();

    for (auto v : intermediates) {
      if (v == op)
        continue;
      for (auto u : v.getUsers()) {
        if (!seen.contains(u)) {
          return rewriter.notifyMatchFailure(
              u, "Had use outside of intermediates");
        }
      }
    }

    llvm::MapVector<Value, SmallVector<Term<stablehlo::SliceOp>>> doneMapping;

    for (auto result : done0) {
      doneMapping[result.term.getOperand()].push_back(result);
    }

    bool hasMerge = false;

    SmallVector<Term<Value>> finalToAdd;

    for (auto &&[_, done] : doneMapping) {
      if (done.size() == 1) {
        for (auto val : done)
          finalToAdd.emplace_back(val.constantFactor, val.valFactor, val.term);
        continue;
      }

      ssize_t offsetDim = -1;
      bool legal = true;
      for (int i = 0; i < done.size(); i++) {
        assert(done[i].term.getOperand() == done[0].term.getOperand());
        if (done[i].valFactor) {
          legal = false;
          break;
        }
        ssize_t mismatch = -1;
        for (int j = 0; j < T.getShape().size(); j++) {
          if (done[i].term.getStrides()[j] != 1) {
            LLVM_DEBUG(llvm::dbgs()
                       << "Non-one stride for " << done[i].term << "\n");
            legal = false;
            break;
          }
          if (done[i].term.getStartIndices()[j] !=
                  done[0].term.getStartIndices()[j] ||
              done[i].term.getLimitIndices()[j] !=
                  done[0].term.getLimitIndices()[j]) {
            if (mismatch != -1) {
              if (mismatch != j) {
                LLVM_DEBUG(llvm::dbgs()
                           << "Multi-dimensional mismatch of slice between "
                           << done[i].term << " and " << done[0].term << "\n");
                legal = false;
                break;
              }
            } else {
              mismatch = j;
            }
          }
        }
        if (mismatch == -1) {
          continue;
        }
        if (offsetDim == -1) {
          offsetDim = mismatch;
          continue;
        }
        if (offsetDim == mismatch) {
          continue;
        }

        LLVM_DEBUG(llvm::dbgs()
                   << "Multi-dimensional(2) mismatch of slice between "
                   << done[i].term << " and " << done[0].term << "\n");
        legal = false;
        break;
      }

      if (!legal) {
        for (auto val : done)
          finalToAdd.emplace_back(val.constantFactor, val.valFactor, val.term);
        continue;
      }

      if (legal && offsetDim == -1) {
        std::optional<APFloat> tally;
        for (auto term : done) {
          if (tally)
            tally = *tally + term.constantFactor;
          else
            tally = term.constantFactor;
        }
        finalToAdd.emplace_back(*tally, nullptr, done[0].term);
        hasMerge = true;
        continue;
      }

      std::map<int, APFloat> terms;
      for (int i = 0; i < done.size(); i++) {
        int offset = (int)done[i].term.getStartIndices()[offsetDim] -
                     (int)done[0].term.getStartIndices()[offsetDim];
        auto found = terms.find(offset);
        if (found != terms.end())
          found->second = found->second + done[i].constantFactor;
        else
          terms.emplace(offset, done[i].constantFactor);
      }
      assert(terms.size() > 1);

      int startidx = terms.begin()->first;
      auto lastItr = terms.end();
      lastItr--;
      int lastidx = lastItr->first;
      assert(lastidx != startidx);

      SmallVector<APFloat> pad(lastidx + 1 - startidx, zero);
      // Check contiguous
      for (auto &term : terms) {
        pad[term.first - startidx] = pad[term.first - startidx] + term.second;
      }

      if (!((Child *)this)->applies(offsetDim, T, terms, startidx, lastidx)) {
        for (auto val : done)
          finalToAdd.emplace_back(val.constantFactor, val.valFactor, val.term);
        continue;
      }

      SmallVector<int64_t> newStart =
          llvm::to_vector(done[0].term.getStartIndices());
      SmallVector<int64_t> newLimit =
          llvm::to_vector(done[0].term.getLimitIndices());
      SmallVector<int64_t> newStride =
          llvm::to_vector(done[0].term.getStrides());
      newStart[offsetDim] += startidx;
      newLimit[offsetDim] += lastidx;

      Value input = stablehlo::SliceOp::create(rewriter, done[0].term.getLoc(),
                                               done[0].term.getOperand(),
                                               newStart, newLimit, newStride);

      auto fty = RankedTensorType::get({lastidx + 1 - startidx},
                                       rewriter.getF64Type());
      Value filter =
          stablehlo::ConstantOp::create(rewriter, op.getLoc(), fty,
                                        DenseFPElementsAttr::get(fty, pad))
              .getResult();
      filter = stablehlo::ConvertOp::create(
                   rewriter, op.getLoc(),
                   RankedTensorType::get({lastidx + 1 - startidx},
                                         op.getType().getElementType()),
                   filter)
                   .getResult();

      auto conv = ((Child *)this)
                      ->makeReduction(rewriter, input, offsetDim, T, terms,
                                      startidx, lastidx, filter);
      finalToAdd.emplace_back(one, nullptr, conv);
      hasMerge = true;
    }

    if (!hasMerge)
      return failure();

    assert(finalToAdd.size());

    Value result = nullptr;
    for (auto term : finalToAdd) {
      Value intermediate = term.term;
      if (!term.constantFactor.isExactlyValue(1.0)) {
        intermediate = stablehlo::MulOp::create(
            rewriter, op.getLoc(), intermediate,
            stablehlo::ConstantOp::create(
                rewriter, op.getLoc(), intermediate.getType(),
                cast<ElementsAttr>(
                    makeAttr(intermediate.getType(), term.constantFactor))));
      }
      if (result == nullptr)
        result = intermediate;
      else
        result = stablehlo::AddOp::create(rewriter, op.getLoc(), result,
                                          intermediate);
    }
    assert(result);
    rewriter.replaceOp(op, ValueRange{result});
    return success();
  }

  bool reduceWindowApplies(size_t offsetDim, RankedTensorType T,
                           const std::map<int, APFloat> &pad, int startidx,
                           int lastidx) {
    auto start = pad.find(startidx)->second;
    for (int i = startidx; i <= lastidx; i++) {
      auto found = pad.find(i);
      if (found == pad.end())
        return false;
      if (found->second != start)
        return false;
    }
    return true;
  }
};

template <typename ST>
struct SumToConv : public SumToReductionBase<ST, SumToConv<ST>> {
  bool collapseDims;
  SumToConv(bool collapseDims, MLIRContext *context, PatternBenefit benefit = 1,
            ArrayRef<StringRef> generatedNames = {})
      : SumToReductionBase<ST, SumToConv<ST>>(context, benefit, generatedNames),
        collapseDims(collapseDims) {}

  using SumToReductionBase<ST, SumToConv<ST>>::matchAndRewrite;

  bool applies(size_t offsetDim, RankedTensorType T,
               const std::map<int, APFloat> &pad, int startidx, int lastidx) {
    if (SumToReductionBase<ST, SumToConv<ST>>::reduceWindowApplies(
            offsetDim, T, pad, startidx, lastidx)) {
      return false;
    }
    if (collapseDims)
      return true;
    else {
      if (T.getShape().size() > 2) {
        return false;
      } else {
        return true;
      }
    }
  }

  Value makeReduction(PatternRewriter &rewriter, Value input, size_t offsetDim,
                      RankedTensorType T, const std::map<int, APFloat> &pad,
                      int startidx, int lastidx, Value filter) {

    size_t newOffsetDim = offsetDim;
    RankedTensorType pre_reshape = T;
    size_t reshapeOffsetDim = 0;
    SmallVector<int64_t> permutation;

    if (T.getShape().size() > 2) {
      if (newOffsetDim != 0 && newOffsetDim != T.getShape().size() - 1) {
        for (int i = 0; i < T.getShape().size(); i++)
          permutation.push_back(i);
        permutation[newOffsetDim] = 0;
        permutation[0] = newOffsetDim;
        input = stablehlo::TransposeOp::create(rewriter, input.getLoc(), input,
                                               permutation);
        newOffsetDim = 0;
      }
      if (newOffsetDim == 0) {
        auto RT = cast<RankedTensorType>(input.getType());
        pre_reshape = RT;
        reshapeOffsetDim = newOffsetDim;
        int64_t newDims[3] = {RT.getShape()[newOffsetDim], 1, 1};
        for (int i = 1; i < RT.getShape().size(); i++) {
          newDims[1] *= RT.getShape()[i];
        }
        input = stablehlo::ReshapeOpCreate(rewriter, input.getLoc(), input,
                                           newDims);
      } else {
        assert(newOffsetDim == T.getShape().size() - 1);

        auto RT = cast<RankedTensorType>(input.getType());
        pre_reshape = RT;
        reshapeOffsetDim = newOffsetDim;
        int64_t newDims[3] = {1, 1, RT.getShape()[newOffsetDim]};
        for (int i = 0; i < RT.getShape().size() - 1; i++) {
          newDims[0] *= RT.getShape()[i];
        }
        input = stablehlo::ReshapeOpCreate(rewriter, input.getLoc(), input,
                                           newDims);
        newOffsetDim = 2;
      }
    } else if (T.getShape().size() < 3) {
      SmallVector<int64_t> newDims =
          llvm::to_vector(cast<RankedTensorType>(input.getType()).getShape());
      reshapeOffsetDim = newOffsetDim;
      while (newDims.size() < 3) {
        newDims.insert(newDims.begin(), 1);
        newOffsetDim++;
      }
      pre_reshape = cast<RankedTensorType>(input.getType());
      input =
          stablehlo::ReshapeOpCreate(rewriter, input.getLoc(), input, newDims);
    }
    SmallVector<int64_t> nonOffsetDims;
    auto CT = cast<RankedTensorType>(input.getType()).getShape();
    for (int i = 0; i < CT.size(); i++) {
      if (i != newOffsetDim)
        nonOffsetDims.push_back(i);
    }
    if (CT[nonOffsetDims[1]] != 1 && CT[nonOffsetDims[0]] == 1) {
      std::swap(nonOffsetDims[1], nonOffsetDims[0]);
    }

    filter = stablehlo::ReshapeOpCreate(rewriter, input.getLoc(), filter,
                                        {lastidx + 1 - startidx, 1, 1});

    // Create convolution dimension numbers
    auto convDims = stablehlo::ConvDimensionNumbersAttr::get(
        rewriter.getContext(),
        /*input_batch_dimension=*/nonOffsetDims[0],
        /*input_feature_dimension=*/nonOffsetDims[1],
        /*input_spatial_dimensions=*/{(int64_t)newOffsetDim},
        /*kernel_input_feature_dimension=*/1,
        /*kernel_output_feature_dimension=*/2,
        /*kernel_spatial_dimensions=*/{0},
        /*output_batch_dimension=*/nonOffsetDims[0],
        /*output_feature_dimension=*/nonOffsetDims[1],
        /*output_spatial_dimensions=*/{(int64_t)newOffsetDim});

    // Create the convolution operation
    SmallVector<int64_t> outShape =
        llvm::to_vector(cast<RankedTensorType>(input.getType()).getShape());
    outShape[newOffsetDim] -= (lastidx - startidx);
    auto convOutType = RankedTensorType::get(outShape, T.getElementType());
    Value conv = stablehlo::ConvolutionOp::create(
        rewriter, input.getLoc(), convOutType, input, filter,
        /*window_strides=*/nullptr,
        /*padding=*/nullptr,
        /*lhs_dilation=*/nullptr,
        /*rhs_dilation=*/nullptr,
        /*window_reversal=*/nullptr,
        /*conv_dimension_numbers=*/convDims,
        /*feature_group_count=*/rewriter.getI64IntegerAttr(1),
        /*batch_group_count=*/rewriter.getI64IntegerAttr(1),
        /*precision_config=*/nullptr);

    if (conv.getType() != pre_reshape) {
      SmallVector<int64_t> post_shape = llvm::to_vector(pre_reshape.getShape());
      post_shape[reshapeOffsetDim] -= (lastidx - startidx);
      conv = stablehlo::ReshapeOpCreate(rewriter, input.getLoc(), conv,
                                        post_shape);
    }
    if (permutation.size())
      conv = stablehlo::TransposeOp::create(rewriter, input.getLoc(), conv,
                                            permutation);
    return conv;
  }
};

template <typename ST>
struct SumToReduceWindow
    : public SumToReductionBase<ST, SumToReduceWindow<ST>> {
  SumToReduceWindow(MLIRContext *context, PatternBenefit benefit = 1,
                    ArrayRef<StringRef> generatedNames = {})
      : SumToReductionBase<ST, SumToReduceWindow<ST>>(context, benefit,
                                                      generatedNames) {}

  using SumToReductionBase<ST, SumToReduceWindow<ST>>::matchAndRewrite;

  bool applies(size_t offsetDim, RankedTensorType T,
               const std::map<int, APFloat> &pad, int startidx, int lastidx) {
    return SumToReductionBase<ST, SumToReduceWindow<ST>>::reduceWindowApplies(
        offsetDim, T, pad, startidx, lastidx);
  }

  Value makeReduction(PatternRewriter &rewriter, Value input, size_t offsetDim,
                      RankedTensorType T, const std::map<int, APFloat> &pad,
                      int startidx, int lastidx, Value filter) {
    auto factor = pad.begin()->second;

    SmallVector<int64_t> outShape =
        llvm::to_vector(cast<RankedTensorType>(input.getType()).getShape());
    outShape[offsetDim] -= (lastidx - startidx);

    auto unrankedTensorType = RankedTensorType::get({}, T.getElementType());
    Value init_values[1] = {stablehlo::ConstantOp::create(
        rewriter, input.getLoc(), rewriter.getZeroAttr(unrankedTensorType))};

    SmallVector<int64_t> win_dim(outShape.size(), 1);
    win_dim[offsetDim] = lastidx - startidx + 1;

    SmallVector<int64_t> win_strides(outShape.size(), 1);
    SmallVector<int64_t> win_dialations(outShape.size(), 1);
    SmallVector<int64_t> base_dialations(outShape.size(), 1);
    SmallVector<int64_t> padding_dialations(2 * outShape.size(), 0);

    int64_t padding_shape[2] = {(int64_t)outShape.size(), 2};

    Value operands[1] = {input};
    Type restys[1] = {RankedTensorType::get(outShape, T.getElementType())};
    auto redwin = stablehlo::ReduceWindowOp::create(
        rewriter, input.getLoc(), restys, operands, init_values,
        rewriter.getDenseI64ArrayAttr(win_dim),
        rewriter.getDenseI64ArrayAttr(win_strides),
        rewriter.getDenseI64ArrayAttr(base_dialations),
        rewriter.getDenseI64ArrayAttr(win_dialations),
        DenseIntElementsAttr::get(
            RankedTensorType::get(padding_shape, rewriter.getIntegerType(64)),
            padding_dialations));

    Type tys[2] = {unrankedTensorType, unrankedTensorType};
    Location locs[2] = {input.getLoc(), input.getLoc()};

    Value result = redwin->getResult(0);
    if (!factor.isExactlyValue(1.0)) {
      result = stablehlo::MulOp::create(
          rewriter, input.getLoc(), result,
          stablehlo::ConstantOp::create(
              rewriter, input.getLoc(), result.getType(),
              cast<ElementsAttr>(makeAttr(result.getType(), factor))));
    }

    {
      OpBuilder::InsertionGuard guard(rewriter);
      auto block = rewriter.createBlock(&redwin.getBody(), {}, tys, locs);
      rewriter.setInsertionPointToStart(block);
      auto addOp = stablehlo::AddOp::create(rewriter, input.getLoc(),
                                            block->getArgument(0),
                                            block->getArgument(1));
      stablehlo::ReturnOp::create(rewriter, input.getLoc(), addOp.getResult());
    }

    return result;
  }
};

struct TransposeSelect
    : public CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeSelect> {
  using CheckedOpRewritePattern<stablehlo::TransposeOp,
                                TransposeSelect>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp transposeOp,
                                    PatternRewriter &rewriter) const {
    auto selectOp =
        transposeOp.getOperand().getDefiningOp<stablehlo::SelectOp>();
    if (!selectOp)
      return failure();

    if (!selectOp->hasOneUse())
      return failure();

    Value pred = selectOp.getPred();
    bool scalar_pred =
        dyn_cast<RankedTensorType>(pred.getType()).getRank() == 0;

    Value newPred;
    if (!scalar_pred) {
      newPred = stablehlo::TransposeOp::create(
          rewriter, transposeOp.getLoc(), pred, transposeOp.getPermutation());
    } else {
      newPred = pred;
    }

    SmallVector<int64_t> permutation;
    for (int i = 0; i < transposeOp.getPermutation().size(); i++) {
      permutation.push_back(transposeOp.getPermutation()[i]);
    }

    auto onTrueTransposed = stablehlo::TransposeOp::create(
        rewriter, transposeOp.getLoc(), selectOp.getOnTrue(), permutation);
    auto onFalseTransposed = stablehlo::TransposeOp::create(
        rewriter, transposeOp.getLoc(), selectOp.getOnFalse(), permutation);

    rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(
        transposeOp, newPred, onTrueTransposed, onFalseTransposed);
    return success();
  }
};

struct ReshapeSelect
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp, ReshapeSelect> {
  using CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                ReshapeSelect>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp reshapeOp,
                                    PatternRewriter &rewriter) const {
    auto selectOp = reshapeOp.getOperand().getDefiningOp<stablehlo::SelectOp>();
    if (!selectOp)
      return failure();

    if (!selectOp->hasOneUse())
      return failure();

    Value pred = selectOp.getPred();
    bool scalar_pred =
        dyn_cast<RankedTensorType>(pred.getType()).getRank() == 0;

    Value newPred;
    if (!scalar_pred) {
      newPred = stablehlo::ReshapeOpCreate(rewriter, reshapeOp.getLoc(), pred,
                                           reshapeOp.getType().getShape());
    } else {
      newPred = pred;
    }

    auto onTrueReshaped = stablehlo::ReshapeOpCreate(
        rewriter, reshapeOp.getLoc(), selectOp.getOnTrue(),
        reshapeOp.getType().getShape());
    auto onFalseReshaped = stablehlo::ReshapeOpCreate(
        rewriter, reshapeOp.getLoc(), selectOp.getOnFalse(),
        reshapeOp.getType().getShape());

    rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(
        reshapeOp, newPred, onTrueReshaped, onFalseReshaped);
    return success();
  }
};

template <typename T>
struct GroupComms : public CheckedOpRewritePattern<T, GroupComms<T>> {
  using CheckedOpRewritePattern<T, GroupComms<T>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(T end, PatternRewriter &rewriter) const {
    if (end->template getParentOfType<enzymexla::CommRegionOp>())
      return failure();
    if (end->template getParentOfType<sdy::ManualComputationOp>())
      return failure();
    SetVector<Operation *> done;
    done.insert(end);
    SmallVector<Operation *> todo = {end};
    while (todo.size()) {
      auto cur = todo.pop_back_val();
      if (cur != end && done.contains(cur))
        continue;
      if (!isa<stablehlo::SliceOp, stablehlo::ConcatenateOp, stablehlo::PadOp,
               stablehlo::TransposeOp, stablehlo::ReshapeOp,
               stablehlo::DynamicUpdateSliceOp, enzymexla::RotateOp,
               enzymexla::ExtendOp, enzymexla::WrapOp>(cur))
        continue;
      bool allWithin = true;
      if (cur != end)
        for (auto res : cur->getResults()) {
          for (auto u : res.getUsers()) {
            if (!done.contains(u)) {
              allWithin = false;
              break;
            }
          }
        }
      if (!allWithin)
        continue;
      done.insert(cur);
      for (auto op : cur->getOperands()) {
        if (auto v = op.getDefiningOp()) {
          todo.push_back(v);
        }
      }
    }

    done = mlir::topologicalSort(done);

    auto newOp = enzymexla::CommRegionOp::create(rewriter, end.getLoc(),
                                                 end.getResult().getType());
    rewriter.createBlock(&newOp.getBody(), newOp.getBody().begin());
    IRMapping map;
    for (auto op : done) {
      rewriter.clone(*op, map);
    }
    stablehlo::ReturnOp::create(rewriter, end.getLoc(),
                                map.lookup(end.getResult()));
    for (auto op : llvm::reverse(done)) {
      if (op == end) {
        rewriter.replaceOp(op, newOp);
      } else {
        rewriter.eraseOp(op);
      }
    }
    return success();
  }
};

struct LowerCommRegion
    : public CheckedOpRewritePattern<enzymexla::CommRegionOp, LowerCommRegion> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::CommRegionOp end,
                                    PatternRewriter &rewriter) const {
    IRMapping map;
    for (auto &op : end.getBody().front()) {
      if (isa<stablehlo::ReturnOp>(op)) {
        SmallVector<Value> operands;
        for (auto v : op.getOperands())
          operands.push_back(map.lookup(v));
        rewriter.replaceOp(end, operands);
        return success();
      } else {
        rewriter.clone(op, map);
      }
    }
    return failure();
  }
};

bool isRotateLike(int dimension, Value lhs, Value rhs, stablehlo::SliceOp *sl0P,
                  stablehlo::SliceOp *sl1P) {
  auto sl0 = lhs.getDefiningOp<stablehlo::SliceOp>();
  if (!sl0)
    return false;
  auto sl1 = rhs.getDefiningOp<stablehlo::SliceOp>();
  if (!sl1)
    return false;
  if (sl0.getOperand() != sl1.getOperand())
    return false;

  if (sl0P)
    *sl0P = sl0;
  if (sl1P)
    *sl1P = sl1;

  if (sl0.getOperand() != sl1.getOperand())
    return false;

  // sl0[A:end], sl1[start:A]
  for (int j = 0; j < sl0.getType().getShape().size(); j++) {
    if (j == dimension) {
      if (sl0.getStrides()[j] != 1 || sl1.getStrides()[j] != 1) {
        return false;
      }
      if (sl0.getStartIndices()[j] != sl1.getLimitIndices()[j]) {
        return false;
      }
    } else {
      if (sl0.getLimitIndices()[j] != sl1.getLimitIndices()[j]) {
        return false;
      }
      if (sl0.getStartIndices()[j] != sl1.getStartIndices()[j]) {
        return false;
      }
      if (sl0.getStrides()[j] != sl1.getStrides()[j]) {
        return false;
      }
    }
  }
  return true;
}

bool isExtendRotateLike(int dimension, Value lhs, Value rhs,
                        stablehlo::ConcatenateOp *concat = nullptr,
                        PatternRewriter *rewriter = nullptr) {

  auto shapeL = cast<RankedTensorType>(lhs.getType()).getShape();
  auto shapeR = cast<RankedTensorType>(rhs.getType()).getShape();

  // The size [and data] of non relevant dimension must be the same.
  for (size_t i = 0; i < shapeL.size(); i++) {
    if (shapeL[i] != shapeR[i] && i != dimension) {
      return false;
    }
  }

  if (shapeL[dimension] == 0 || shapeR[dimension] == 0)
    return false;

  auto sl0 = lhs.getDefiningOp<stablehlo::SliceOp>();
  auto sl1 = rhs.getDefiningOp<stablehlo::SliceOp>();

  auto r0 = lhs.getDefiningOp<enzymexla::RotateOp>();
  auto r1 = rhs.getDefiningOp<enzymexla::RotateOp>();

  if (sl0 && sl1 && sl0.getOperand() == sl1.getOperand()) {

    // [0...3056], [0]
    // %1258 = stablehlo.slice %1257 [8:12, 7:1527, 0:3056] :
    // (tensor<20x1536x3056xf32>) -> tensor<4x1520x3056xf32> loc(#loc4507) %1262
    // = stablehlo.slice %1257 [8:12, 7:1527, 0:1] : (tensor<20x1536x3056xf32>)
    // -> tensor<4x1520x1xf32> loc(#loc3748) %2878 = stablehlo.concatenate
    // %1258, %1262, dim = 2 : (tensor<4x1520x3056xf32>, tensor<4x1520x1xf32>)
    // -> tensor<4x1520x3057xf32> loc(#loc2107)
    //
    // [1, 2, 3, ... 3056], [0, 1]
    // %1187 = stablehlo.slice %1179 [0:4, 0:1519, 1:3056] :
    // (tensor<4x1519x3056xf32>) -> tensor<4x1519x3055xf32> loc(#loc2852) %1183
    // = stablehlo.slice %1179 [0:4, 0:1519, 0:2] : (tensor<4x1519x3056xf32>) ->
    // tensor<4x1519x2xf32> loc(#loc2852) %2936 = stablehlo.concatenate %1187,
    // %1183, dim = 2 : (tensor<4x1519x3055xf32>, tensor<4x1519x2xf32>) ->
    // tensor<4x1519x3057xf32> loc(#loc3581)
    //
    // [2, 3, ... 3056], [0, 1, 2]
    // %1184 = stablehlo.slice %1179 [0:4, 0:1519, 2:3056] :
    // (tensor<4x1519x3056xf32>) -> tensor<4x1519x3054xf32> loc(#loc2852) %1182
    // = stablehlo.slice %1179 [0:4, 0:1519, 0:3] : (tensor<4x1519x3056xf32>) ->
    // tensor<4x1519x3xf32> loc(#loc2852) %2984 = stablehlo.concatenate %1184,
    // %1182, dim = 2 : (tensor<4x1519x3054xf32>, tensor<4x1519x3xf32>) ->
    // tensor<4x1519x3057xf32> loc(#loc3581)

    // sl0[A:end], sl1[start:A+ext]
    for (int j = 0; j < shapeL.size(); j++) {
      if (j == dimension) {
        if (sl0.getStrides()[j] != 1 || sl1.getStrides()[j] != 1) {
          return false;
        }
        if (sl0.getStartIndices()[j] >= sl1.getLimitIndices()[j]) {
          return false;
        }
      } else {
        if (sl0.getLimitIndices()[j] != sl1.getLimitIndices()[j]) {
          return false;
        }
        if (sl0.getStartIndices()[j] != sl1.getStartIndices()[j]) {
          return false;
        }
        if (sl0.getStrides()[j] != sl1.getStrides()[j]) {
          return false;
        }
      }
    }

    auto extAmt =
        sl1.getLimitIndices()[dimension] - sl0.getStartIndices()[dimension];
    assert(extAmt > 0);
    if (extAmt != 1)
      return false;

    if (rewriter) {
      SmallVector<int64_t> starts = llvm::to_vector(sl1.getStartIndices());
      SmallVector<int64_t> limits = llvm::to_vector(sl0.getLimitIndices());
      SmallVector<int64_t> stride = llvm::to_vector(sl0.getStrides());
      bool needs_slice = false;
      for (size_t i = 0; i < starts.size(); i++) {
        if (starts[i] != 0 || stride[i] != 1 ||
            limits[i] != sl0.getOperand().getType().getShape()[i]) {
          needs_slice = true;
        }
      }
      Value base = sl0.getOperand();
      auto shard = sdy::getShardingPerValue(*concat);
      if (needs_slice) {
        auto slice2 = stablehlo::SliceOp::create(*rewriter, concat->getLoc(),
                                                 base, starts, limits, stride);
        if (shard) {
          sdy::setShardings(slice2, shard);
        }
        base = slice2;
      }

      auto A = sl0.getStartIndices()[dimension] - starts[dimension];

      if (A == 0) {
        auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
            *concat, base, /*lhs*/ 0, /*rhs*/ 1, dimension);
        if (shard) {
          sdy::setShardings(wrap, shard);
        }
        return true;
      }

      if (A ==
          cast<RankedTensorType>(base.getType()).getShape()[dimension] - 1) {
        auto shard = sdy::getShardingPerValue(*concat);
        auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
            *concat, base, /*lhs*/ 1, /*rhs*/ 0, dimension);
        if (shard) {
          sdy::setShardings(wrap, shard);
        }
        return true;
      }

      auto rot = enzymexla::RotateOp::create(*rewriter, concat->getLoc(), base,
                                             A, dimension);
      if (shard) {
        sdy::setShardings(rot, shard);
      }
      base = rot;

      auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
          *concat, base, /*lhs*/ 0, /*rhs*/ 1, dimension);
      if (shard) {
        sdy::setShardings(wrap, shard);
      }
    }
    return true;
  }

  if (sl0 && sl0.getOperand() == rhs) {

    //  %2668 = stablehlo.slice %1171 [0:4, 0:1520, 3055:3056] :
    //  (tensor<4x1520x3056xf32>) -> tensor<4x1520x1xf32> loc(#loc2557)
    //  ...
    //  %2669 = stablehlo.concatenate %2668, %1171, dim = 2 :
    //  (tensor<4x1520x1xf32>, tensor<4x1520x3056xf32>) ->
    //  tensor<4x1520x3057xf32> loc(#loc2557)

    // sl0[end-ext, end], X[0:end]

    for (int j = 0; j < shapeL.size(); j++) {
      if (j == dimension) {
        if (sl0.getStrides()[j] != 1) {
          return false;
        }
        if (sl0.getLimitIndices()[j] != shapeR[j]) {
          return false;
        }
      } else {
        if (sl0.getLimitIndices()[j] != shapeR[j]) {
          return false;
        }
        if (sl0.getStartIndices()[j] != 0) {
          return false;
        }
        if (sl0.getStrides()[j] != 1) {
          return false;
        }
      }
    }

    auto extAmt = shapeL[dimension];
    assert(extAmt > 0);

    if (rewriter) {
      Value base = sl0.getOperand();
      auto shard = sdy::getShardingPerValue(*concat);
      auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
          *concat, base, /*lhs*/ extAmt, /*rhs*/ 0, dimension);
      if (shard) {
        sdy::setShardings(wrap, shard);
      }
    }

    return true;
  }

  if (sl1 && sl1.getOperand() == lhs) {
    // ...
    // %1186 = stablehlo.slice %1179 [0:4, 0:1519, 0:1] :
    // (tensor<4x1519x3056xf32>) -> tensor<4x1519x1xf32> loc(#loc2852) %2933 =
    // stablehlo.concatenate %1179, %1186, dim = 2 : (tensor<4x1519x3056xf32>,
    // tensor<4x1519x1xf32>) -> tensor<4x1519x3057xf32> loc(#loc3581)

    // X[0:end], sl1[0, ext]

    for (int j = 0; j < shapeL.size(); j++) {
      if (j == dimension) {
        if (sl1.getStrides()[j] != 1) {
          return false;
        }
        if (sl1.getStartIndices()[j] != 0) {
          return false;
        }
      } else {
        if (sl1.getLimitIndices()[j] != shapeL[j]) {
          return false;
        }
        if (sl1.getStartIndices()[j] != 0) {
          return false;
        }
        if (sl1.getStrides()[j] != 1) {
          return false;
        }
      }
    }

    auto extAmt = shapeR[dimension];
    assert(extAmt > 0);

    if (rewriter) {
      Value base = lhs;

      auto shard = sdy::getShardingPerValue(*concat);
      auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
          *concat, base, /*lhs*/ 0, /*rhs*/ extAmt, dimension);
      if (shard) {
        sdy::setShardings(wrap, shard);
      }
    }
    return true;
  }

  if (sl0 && r1 && sl0.getOperand() == r1.getOperand() &&
      r1.getDimension() == dimension) {

    //       [3054], [3055, end, ... 0, 3054]
    //       %2665 = stablehlo.slice %1171 [0:4, 0:1520, 3054:3055] :
    //       (tensor<4x1520x3056xf32>) -> tensor<4x1520x1xf32> loc(#loc2557)
    //       %2290 = "enzymexla.rotate"(%1171) <{amount = 3055 : si32, dimension
    //       = 2 : si32}> : (tensor<4x1520x3056xf32>) -> tensor<4x1520x3056xf32>
    //       loc(#loc3734) %2666 = stablehlo.concatenate %2665, %2290, dim = 2 :
    //       (tensor<4x1520x1xf32>, tensor<4x1520x3056xf32>) ->
    //       tensor<4x1520x3057xf32> loc(#loc2557)

    // sl0[amt-ext, amt], rotateleft(x, amt)
    for (int j = 0; j < shapeL.size(); j++) {
      if (j == dimension) {
        if (sl0.getStrides()[j] != 1) {
          return false;
        }
        if (sl0.getLimitIndices()[j] != r1.getAmount()) {
          return false;
        }
      } else {
        if (sl0.getLimitIndices()[j] != shapeR[j]) {
          return false;
        }
        if (sl0.getStartIndices()[j] != 0) {
          return false;
        }
        if (sl0.getStrides()[j] != 1) {
          return false;
        }
      }
    }

    auto extAmt = shapeL[dimension];
    assert(extAmt > 0);
    if (extAmt != 1)
      return false;

    if (rewriter) {
      auto shard = sdy::getShardingPerValue(*concat);
      auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
          *concat, r1, /*lhs*/ extAmt, /*rhs*/ 0, dimension);
      if (shard) {
        sdy::setShardings(wrap, shard);
      }
    }
    return true;
  }

  if (sl0 && r1 && r1.getDimension() == dimension) {
    sl1 = r1.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (sl1 && sl1.getOperand() == sl0.getOperand()) {

      // %22 = stablehlo.slice %15 [8:12, 8:6136, 12269:12270] :
      // (tensor<20x6144x12272xf32>) -> tensor<4x6128x1xf32> loc(#loc3742) %23 =
      // stablehlo.slice %15 [8:12, 8:6136, 0:12272] :
      // (tensor<20x6144x12272xf32>) -> tensor<4x6128x12272xf32> loc(#loc3742)
      // [12270, 12271, 0, 1, ... 12269]
      // %648 = "enzymexla.rotate"(%23) <{amount = 12270 : si32, dimension = 2 :
      // si32}> : (tensor<4x6128x12272xf32>) -> tensor<4x6128x12272xf32>
      // loc(#loc3742) [12269], [12270, 12271, 0, 1, ... 12269] %1847 =
      // stablehlo.concatenate %22, %648, dim = 2 : (tensor<4x6128x1xf32>,
      // tensor<4x6128x12272xf32>) -> tensor<4x6128x12273xf32> loc(#loc3603)

      // sl0[start+A-ext:start+A], rotate(x[start:end], leftAmt=A)
      for (int j = 0; j < shapeL.size(); j++) {
        if (j == dimension) {
          if (sl0.getStrides()[j] != 1 || sl1.getStrides()[j] != 1) {
            return false;
          }
          if (sl0.getLimitIndices()[j] !=
              r1.getAmount() + sl1.getStartIndices()[j]) {
            return false;
          }
          if (sl0.getStartIndices()[j] <= sl1.getStartIndices()[j]) {
            return false;
          }
          if (sl0.getLimitIndices()[j] >= sl1.getLimitIndices()[j]) {
            return false;
          }
        } else {
          if (sl0.getLimitIndices()[j] != sl1.getLimitIndices()[j]) {
            return false;
          }
          if (sl0.getStartIndices()[j] != sl1.getStartIndices()[j]) {
            return false;
          }
          if (sl0.getStrides()[j] != sl1.getStrides()[j]) {
            return false;
          }
        }
      }

      auto extAmt = sl0.getType().getShape()[dimension];
      assert(extAmt > 0);

      if (rewriter) {
        Value base = r1;

        auto shard = sdy::getShardingPerValue(*concat);
        auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
            *concat, base, /*lhs*/ extAmt, /*rhs*/ 0, dimension);
        if (shard) {
          sdy::setShardings(wrap, shard);
        }
      }
      return true;
    }
  }

  if (sl1 && r0 && r0.getDimension() == dimension) {
    sl0 = r0.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (sl0 && sl0.getOperand() == sl1.getOperand()) {

      //  %685 = stablehlo.slice %arg20 [8:12, 8:6136, 8:12280] :
      //  (tensor<20x6144x12288xf32>) -> tensor<4x6128x12272xf32> [10, 11, ...,
      //  12280, 8, 9] %1235 = "enzymexla.rotate"(%685) <{amount = 2 : si32,
      //  dimension = 2 : si32}> : (tensor<4x6128x12272xf32>) ->
      //  tensor<4x6128x12272xf32> [10] %1635 = stablehlo.slice %arg20 [8:12,
      //  8:6136, 10:11] : (tensor<20x6144x12288xf32>) -> tensor<4x6128x1xf32>
      //  [10, 11, ..., 12280, 8, 9], [10]
      //  %1636 = stablehlo.concatenate %1235, %1635, dim = 2 :
      //  (tensor<4x6128x12272xf32>, tensor<4x6128x1xf32>) ->
      //  tensor<4x6128x12273xf32>

      // rotate(x[start:end], leftAmt=A), sl1[start+A:start+A+ext]
      for (int j = 0; j < shapeL.size(); j++) {
        if (j == dimension) {
          if (sl0.getStrides()[j] != 1 || sl1.getStrides()[j] != 1) {
            return false;
          }
          if (sl1.getStartIndices()[j] !=
              r0.getAmount() + sl0.getStartIndices()[j]) {
            return false;
          }
          if (sl1.getStartIndices()[j] <= sl0.getStartIndices()[j]) {
            return false;
          }
          if (sl1.getLimitIndices()[j] >= sl0.getLimitIndices()[j]) {
            return false;
          }
        } else {
          if (sl0.getLimitIndices()[j] != sl1.getLimitIndices()[j]) {
            return false;
          }
          if (sl0.getStartIndices()[j] != sl1.getStartIndices()[j]) {
            return false;
          }
          if (sl0.getStrides()[j] != sl1.getStrides()[j]) {
            return false;
          }
        }
      }

      auto extAmt = sl1.getType().getShape()[dimension];
      assert(extAmt > 0);

      if (rewriter) {
        Value base = r0;
        auto shard = sdy::getShardingPerValue(*concat);
        auto wrap = rewriter->replaceOpWithNewOp<enzymexla::WrapOp>(
            *concat, base, /*lhs*/ 0, /*rhs*/ extAmt, dimension);
        if (shard) {
          sdy::setShardings(wrap, shard);
        }
      }
      return true;
    }
  }

  if (sl1 && r0 && sl1.getOperand() == r0.getOperand() &&
      r0.getDimension() == dimension) {
    // TODO
  }
  return false;
}

struct RecognizeRotate
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     RecognizeRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concat,
                                    PatternRewriter &rewriter) const {
    if (concat.getOperands().size() < 2)
      return failure();
    if (concat.getOperands().size() == 2 &&
        isExtendRotateLike(concat.getDimension(), concat.getOperands()[0],
                           concat.getOperands()[1], &concat, &rewriter))
      return success();
    for (int i = 1; i < concat.getOperands().size(); i++) {
      stablehlo::SliceOp sl0, sl1;
      if (!isRotateLike(concat.getDimension(), concat.getOperands()[i - 1],
                        concat.getOperands()[i], &sl0, &sl1)) {
        continue;
      }
      auto starts = llvm::to_vector(sl1.getStartIndices());
      auto limits = llvm::to_vector(sl0.getLimitIndices());
      Value outerSlice = sl0.getOperand();

      bool needsSlice = false;
      for (int j = 0; j < sl0.getType().getShape().size(); j++) {
        if (starts[j] != 0) {
          needsSlice = true;
          break;
        }
        if (limits[j] !=
            cast<RankedTensorType>(sl0.getOperand().getType()).getShape()[j]) {
          needsSlice = true;
          break;
        }
        if (sl0.getStrides()[j] != 1) {
          needsSlice = true;
          break;
        }
      }
      if (needsSlice) {
        outerSlice =
            stablehlo::SliceOp::create(rewriter, sl0.getLoc(), sl0.getOperand(),
                                       starts, limits, sl0.getStrides());
        if (auto shard = sdy::getShardingPerValue(sl0)) {
          sdy::setShardings(outerSlice.getDefiningOp(), shard);
        }
      }
      auto rotate = enzymexla::RotateOp::create(
          rewriter, sl1.getLoc(), outerSlice,
          sl1.getType().getShape()[concat.getDimension()],
          concat.getDimension());
      if (auto shard = sdy::getShardingPerValue(concat)) {
        sdy::setShardings(rotate, shard);
      }
      SmallVector<Value> toConcat;
      for (int j = 0; j < i - 1; j++)
        toConcat.push_back(concat.getOperands()[j]);
      toConcat.push_back(rotate);
      for (int j = i + 1; j < concat.getOperands().size(); j++)
        toConcat.push_back(concat.getOperands()[j]);
      if (toConcat.size() == 1) {
        rewriter.replaceOp(concat, toConcat[0]);
      } else {
        auto shard = sdy::getShardingPerValue(concat);
        auto newConcat = rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
            concat, toConcat, concat.getDimension());
        if (shard) {
          sdy::setShardings(newConcat, shard);
        }
      }
      return success();
    }

    return failure();
  }
};

struct RecognizeUpdateWithoutCorners
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     RecognizeUpdateWithoutCorners> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concat,
                                    PatternRewriter &rewriter) const {
    if (concat.getOperands().size() != 3)
      return failure();

    auto concat0 =
        concat.getOperands()[0].getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concat0)
      return failure();
    if (concat0.getOperands().size() != 3)
      return failure();
    if (concat0.getDimension() == concat.getDimension())
      return failure();

    auto concat2 =
        concat.getOperands()[2].getDefiningOp<stablehlo::ConcatenateOp>();
    if (!concat2)
      return failure();
    if (concat2.getOperands().size() != 3)
      return failure();
    if (concat0.getDimension() != concat2.getDimension())
      return failure();

    auto slice00 = concat0.getOperands()[0].getDefiningOp<stablehlo::SliceOp>();
    if (!slice00)
      return failure();
    auto slice02 = concat0.getOperands()[2].getDefiningOp<stablehlo::SliceOp>();
    if (!slice02)
      return failure();
    auto slice20 = concat2.getOperands()[0].getDefiningOp<stablehlo::SliceOp>();
    if (!slice20)
      return failure();
    auto slice22 = concat2.getOperands()[2].getDefiningOp<stablehlo::SliceOp>();
    if (!slice22)
      return failure();

    auto dimX = concat.getDimension();
    auto dimY = concat0.getDimension();

    if (slice00.getOperand() != slice02.getOperand())
      return failure();

    if (slice00.getOperand() != slice02.getOperand())
      return failure();

    if (slice00.getOperand() != slice20.getOperand())
      return failure();

    if (slice00.getOperand() != slice22.getOperand())
      return failure();
    if (slice00.getOperand() != slice22.getOperand())
      return failure();

    ssize_t x1 = -1, x2 = -1, y1 = -1, y2 = -1;

    ///
    // | dimY [ aka dim1]
    // v
    //
    // -> dimX [aka dim0]
    //
    //   [ slice00 ]        [ slice20 ]
    //                data
    //   [ slice02 ]        [ slice22 ]
    //

    for (size_t i = 0; i < concat.getType().getShape().size(); i++) {
      for (auto slice : {slice00, slice02, slice20, slice22}) {
        size_t expectedStart = 0;
        size_t expectedLimit = concat.getType().getShape()[i];

        if (i == dimX) {
          if (slice == slice00 || slice == slice02) {
            if (x1 == -1) {
              x1 = slice.getType().getShape()[i];
            }
            expectedLimit = x1;
          }

          if (slice == slice20 || slice == slice22) {
            if (x2 == -1) {
              x2 = slice.getStartIndices()[i];
            }
            expectedStart = x2;
          }
        } else if (i == dimY) {
          if (slice == slice00 || slice == slice20) {
            if (y1 == -1) {
              y1 = slice.getType().getShape()[i];
            }
            expectedLimit = y1;
          }

          if (slice == slice02 || slice == slice22) {
            if (y2 == -1) {
              y2 = slice.getStartIndices()[i];
            }
            expectedStart = y2;
          }
        }

        if (slice.getStartIndices()[i] != expectedStart) {
          return failure();
        }
        if (slice.getLimitIndices()[i] != expectedLimit) {
          return failure();
        }
        if (slice.getStrides()[i] != 1) {
          return failure();
        }
      }
    }

    // We've now established it is indeed UpdateWithoutCorners-like.
    // Now we must figure out a good same-sized update to use.
    if (auto extend =
            concat.getOperands()[1].getDefiningOp<enzymexla::ExtendOp>()) {
      ///
      // | dimY [ aka dim1]
      // v
      //
      // -> dimX [aka dim0]
      //
      //   [ slice00 ]    lhs        [ slice20 ]
      //     data[0]    extend(data)   data[end]
      //   [ slice02 ]    rhs        [ slice22 ]
      //

      auto eslice0 =
          concat0.getOperands()[1].getDefiningOp<stablehlo::SliceOp>();
      auto eslice2 =
          concat2.getOperands()[1].getDefiningOp<stablehlo::SliceOp>();
      auto ETy = cast<RankedTensorType>(extend.getOperand().getType());

      if (eslice0 && eslice2 && eslice0.getOperand() == extend.getOperand() &&
          eslice2.getOperand() == extend.getOperand() &&
          extend.getDimension() == dimY && extend.getLhs() == y1 &&
          extend.getRhs() == extend.getType().getShape()[dimY] - y2) {

        for (size_t i = 0; i < concat.getType().getShape().size(); i++) {
          for (auto eslice : {eslice0, eslice2}) {
            size_t expectedStart = 0;
            size_t expectedLimit = ETy.getShape()[i];

            if (i == dimX) {
              if (eslice == eslice0) {
                expectedLimit = x1;
              }

              if (eslice == eslice2) {
                expectedStart =
                    ETy.getShape()[i] - (concat.getType().getShape()[i] - x2);
              }
            }

            if (eslice.getStartIndices()[i] != expectedStart) {
              return failure();
            }
            if (eslice.getLimitIndices()[i] != expectedLimit) {
              return failure();
            }
            if (eslice.getStrides()[i] != 1) {
              return failure();
            }
          }
        }

        auto shard = sdy::getShardingPerValue(concat);

        auto extend2 = enzymexla::ExtendOp::create(
            rewriter, concat0.getLoc(), extend, x1,
            concat.getType().getShape()[dimX] - x2, dimX);
        enzymexla::UpdateWithoutCornersOp newUpdate;

        if (dimX >= dimY) {
          newUpdate =
              rewriter.replaceOpWithNewOp<enzymexla::UpdateWithoutCornersOp>(
                  concat, slice00.getOperand(), extend2, dimY, y1, y2, dimX, x1,
                  x2);
        } else {
          newUpdate =
              rewriter.replaceOpWithNewOp<enzymexla::UpdateWithoutCornersOp>(
                  concat, slice00.getOperand(), extend2, dimX, x1, x2, dimY, y1,
                  y2);
        }

        if (shard) {
          sdy::setShardings(extend2, shard);
          sdy::setShardings(newUpdate, shard);
        }
        return success();
      }
    }
    return failure();
  }
};

bool isWrapLike(int dim, Value lhs, Value mid, Value rhs,
                stablehlo::SliceOp *sl0P = nullptr,
                stablehlo::SliceOp *sl1P = nullptr) {
  auto sl0 = lhs.getDefiningOp<stablehlo::SliceOp>();
  if (!sl0)
    return false;
  auto midT = cast<RankedTensorType>(mid.getType());
  auto sl1 = rhs.getDefiningOp<stablehlo::SliceOp>();
  if (!sl1)
    return false;
  if (sl0.getOperand() != sl1.getOperand())
    return false;

  if (lhs == mid || mid == rhs) {
    return false;
  }

  if (sl0P)
    *sl0P = sl0;
  if (sl1P)
    *sl1P = sl1;

  SmallVector<int64_t> midStarts(midT.getShape().size(), 0);
  SmallVector<int64_t> midLimits = llvm::to_vector(midT.getShape());
  SmallVector<int64_t> midStrides(midT.getShape().size(), 1);
  if (sl0.getOperand() != mid) {
    auto slM = mid.getDefiningOp<stablehlo::SliceOp>();
    if (!slM)
      return false;
    if (slM.getOperand() != sl0.getOperand())
      return false;
    midStarts = llvm::to_vector(slM.getStartIndices());
    midLimits = llvm::to_vector(slM.getLimitIndices());
    midStrides = llvm::to_vector(slM.getStrides());
  }

  // sl0[B-lhs:B], mid[A:B] sl1[A:A+rhs]
  for (int j = 0; j < sl0.getType().getShape().size(); j++) {
    if (j == dim) {
      if (sl0.getStrides()[j] != 1 || sl1.getStrides()[j] != 1 ||
          midStrides[j] != 1) {
        return false;
      }
      if (sl0.getLimitIndices()[j] != midLimits[j]) {
        return false;
      }
      if (midStarts[j] != sl1.getStartIndices()[j]) {
        return false;
      }
    } else {
      if (sl0.getLimitIndices()[j] != sl1.getLimitIndices()[j]) {
        return false;
      }
      if (sl0.getLimitIndices()[j] != midLimits[j]) {
        return false;
      }
      if (sl0.getStartIndices()[j] != sl1.getStartIndices()[j]) {
        return false;
      }
      if (sl0.getStartIndices()[j] != midStarts[j]) {
        return false;
      }
      if (sl0.getStrides()[j] != sl1.getStrides()[j]) {
        return false;
      }
      if (sl0.getStrides()[j] != midStrides[j]) {
        return false;
      }
    }
  }
  return true;
}

struct RecognizeWrap
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp, RecognizeWrap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concat,
                                    PatternRewriter &rewriter) const {

    int concatDim = concat.getDimension();
    SmallVector<Value> operands = llvm::to_vector(concat.getOperands());
    auto shard = sdy::getShardingPerValue(concat);

    if (operands.size() < 2)
      return failure();
    for (int i = 2; i < operands.size(); i++) {
      stablehlo::SliceOp sl0;
      auto mid = operands[i - 1];
      stablehlo::SliceOp sl1;
      if (isWrapLike(concatDim, operands[i - 2], mid, operands[i], &sl0,
                     &sl1)) {
        auto wrap = enzymexla::WrapOp::create(
            rewriter, sl0.getLoc(), mid, sl0.getType().getShape()[concatDim],
            sl1.getType().getShape()[concatDim], concatDim);
        if (auto shard = sdy::getShardingPerValue(sl0)) {
          sdy::setShardings(wrap, shard);
        }
        SmallVector<Value> toConcat;
        for (int j = 0; j < i - 2; j++)
          toConcat.push_back(operands[j]);
        toConcat.push_back(wrap);
        for (int j = i + 1; j < operands.size(); j++)
          toConcat.push_back(operands[j]);
        if (toConcat.size() == 1) {
          rewriter.replaceOp(concat, toConcat[0]);
        } else {
          auto newConcat =
              rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
                  concat, toConcat, concatDim);
          if (shard) {
            sdy::setShardings(newConcat, shard);
          }
        }
        return success();
      }
      auto rs0 = operands[i - 2].getDefiningOp<stablehlo::ReshapeOp>();
      auto rsmid = operands[i - 1].getDefiningOp<stablehlo::ReshapeOp>();
      auto rs1 = operands[i].getDefiningOp<stablehlo::ReshapeOp>();
      if (rs0 && rsmid && rs1 && isOuterReducingReshape(rs0) &&
          isOuterReducingReshape(rsmid) && isOuterReducingReshape(rs1)) {
        if (isWrapLike(concatDim + 1, rs0.getOperand(), rsmid.getOperand(),
                       rs1.getOperand(), &sl0, &sl1)) {
          auto wrap = enzymexla::WrapOp::create(
              rewriter, sl0.getLoc(), rsmid.getOperand(),
              sl0.getType().getShape()[concatDim + 1],
              sl1.getType().getShape()[concatDim + 1], concatDim + 1);
          if (auto shard = sdy::getShardingPerValue(sl0)) {
            sdy::setShardings(wrap, shard);
          }
          SmallVector<Value> toConcat;
          for (int j = 0; j < i - 2; j++)
            toConcat.push_back(operands[j]);
          SmallVector<int64_t> newShape =
              llvm::to_vector(wrap.getType().getShape());
          assert(newShape[0] == 1);
          newShape.erase(newShape.begin());
          auto reshape = stablehlo::ReshapeOpCreate(rewriter, concat.getLoc(),
                                                    wrap, newShape);
          if (auto shard = sdy::getShardingPerValue(rs0)) {
            sdy::setShardings(reshape.getDefiningOp(), shard);
          }
          toConcat.push_back(reshape);
          for (int j = i + 1; j < operands.size(); j++)
            toConcat.push_back(operands[j]);
          if (toConcat.size() == 1) {
            rewriter.replaceOp(concat, toConcat[0]);
          } else {
            auto shard = sdy::getShardingPerValue(rs0);
            rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
                concat, toConcat, concatDim);
            if (shard) {
              sdy::setShardings(reshape.getDefiningOp(), shard);
            }
          }
          return success();
        }
      }
    }
    return failure();
  }
};

struct ExtendSplat
    : public CheckedOpRewritePattern<enzymexla::ExtendOp, ExtendSplat> {
  using CheckedOpRewritePattern<enzymexla::ExtendOp,
                                ExtendSplat>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::ExtendOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr cstAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&cstAttr)))
      return failure();

    if (!cstAttr.isSplat())
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
        op, SplatElementsAttr::get(op.getType(),
                                   cstAttr.getSplatValue<Attribute>()));

    return success();
  }
};

struct ExtendPad
    : public CheckedOpRewritePattern<enzymexla::ExtendOp, ExtendPad> {
  using CheckedOpRewritePattern<enzymexla::ExtendOp,
                                ExtendPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::ExtendOp extend,
                                    PatternRewriter &rewriter) const {

    auto pad = extend.getOperand().getDefiningOp<stablehlo::PadOp>();
    if (!pad)
      return failure();

    if (pad.getEdgePaddingLow()[extend.getDimension()] != 0)
      return failure();

    if (pad.getEdgePaddingHigh()[extend.getDimension()] != 0)
      return failure();

    if (pad.getInteriorPadding()[extend.getDimension()] != 0)
      return failure();

    auto newExtend = enzymexla::ExtendOp::create(
        rewriter, extend.getLoc(), pad.getOperand(), extend.getLhs(),
        extend.getRhs(), extend.getDimension());

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        extend, newExtend, pad.getPaddingValue(), pad.getEdgePaddingLow(),
        pad.getEdgePaddingHigh(), pad.getInteriorPadding());

    return success();
  }
};

struct ExtendToBroadcast
    : public CheckedOpRewritePattern<enzymexla::ExtendOp, ExtendToBroadcast> {
  using CheckedOpRewritePattern<enzymexla::ExtendOp,
                                ExtendToBroadcast>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::ExtendOp extend,
                                    PatternRewriter &rewriter) const {
    // Check if the input is a ranked tensor
    auto inputType = dyn_cast<RankedTensorType>(extend.getOperand().getType());
    if (!inputType)
      return failure();

    // Get the dimension being extended
    int64_t dim = extend.getDimension();

    // Check if the dimension is a singleton (size = 1)
    if (inputType.getDimSize(dim) != 1)
      return failure();

    // Build the broadcast_dimensions attribute
    // This maps each dimension of the input to the corresponding dimension in
    // the output
    SmallVector<int64_t> broadcastDims;
    for (int64_t i = 0; i < inputType.getRank(); i++) {
      broadcastDims.push_back(i);
    }

    // Replace the extend with a broadcast_in_dim
    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        extend, extend.getType(), extend.getOperand(),
        rewriter.getDenseI64ArrayAttr(broadcastDims));

    return success();
  }
};

template <typename EnzymeOp>
LogicalResult commUnaryOpElementwise(bool onlySingleUser, EnzymeOp op,
                                     PatternRewriter &rewriter) {
  if (onlySingleUser && !llvm::hasSingleElement(op->getUsers()))
    return failure();

  bool anyModified = false;
  for (auto elem : llvm::make_early_inc_range(op->getUsers())) {
    if (!hasTraitElementwise(elem) || elem->getNumResults() != 1 ||
        elem->getNumOperands() != 1)
      continue;

    auto newOp = rewriter.create(
        elem->getLoc(), elem->getName().getIdentifier(),
        ValueRange(op.getOperand()),
        TypeRange(cast<RankedTensorType>(elem->getResult(0).getType())
                      .clone(cast<RankedTensorType>(op.getOperand().getType())
                                 .getShape())),
        elem->getAttrs(), {}, {});
    rewriter.replaceOpWithNewOp<EnzymeOp>(
        elem, newOp->getResult(0), op.getLhs(), op.getRhs(), op.getDimension());
    anyModified = true;
  }

  return success(anyModified);
}

struct ExtendUnaryElementwise
    : public CheckedOpRewritePattern<enzymexla::ExtendOp,
                                     ExtendUnaryElementwise> {
  bool onlySingleUser;

  ExtendUnaryElementwise(bool onlySingleUser, MLIRContext *context,
                         PatternBenefit benefit = 1,
                         ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        onlySingleUser(onlySingleUser) {}

  LogicalResult matchAndRewriteImpl(enzymexla::ExtendOp op,
                                    PatternRewriter &rewriter) const {
    return commUnaryOpElementwise(onlySingleUser, op, rewriter);
  }
};

struct WrapUnaryElementwise
    : public CheckedOpRewritePattern<enzymexla::WrapOp, WrapUnaryElementwise> {
  bool onlySingleUser;

  WrapUnaryElementwise(bool onlySingleUser, MLIRContext *context,
                       PatternBenefit benefit = 1,
                       ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        onlySingleUser(onlySingleUser) {}

  LogicalResult matchAndRewriteImpl(enzymexla::WrapOp op,
                                    PatternRewriter &rewriter) const {
    return commUnaryOpElementwise(onlySingleUser, op, rewriter);
  }
};

template <typename Op, typename EnzymeOp>
LogicalResult commBinOpElementWise(Op op, PatternRewriter &rewriter) {
  auto lhs = op.getLhs();
  auto rhs = op.getRhs();

  auto lhsExtend = lhs.template getDefiningOp<EnzymeOp>();
  auto rhsExtend = rhs.template getDefiningOp<EnzymeOp>();

  if (!lhsExtend || !rhsExtend)
    return failure();

  if (lhsExtend.getLhs() != rhsExtend.getLhs() ||
      lhsExtend.getRhs() != rhsExtend.getRhs() ||
      lhsExtend.getDimension() != rhsExtend.getDimension())
    return failure();

  if (!llvm::hasSingleElement(lhs.getUsers()) ||
      !llvm::hasSingleElement(rhs.getUsers()))
    return failure();

  auto elementWise = Op::create(rewriter, op.getLoc(), lhsExtend.getOperand(),
                                rhsExtend.getOperand());
  rewriter.replaceOpWithNewOp<EnzymeOp>(op, elementWise, lhsExtend.getLhs(),
                                        lhsExtend.getRhs(),
                                        lhsExtend.getDimension());

  return success();
};

template <typename Op>
struct ExtendElementwise
    : public CheckedOpRewritePattern<Op, ExtendElementwise<Op>> {
  using CheckedOpRewritePattern<Op,
                                ExtendElementwise<Op>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(Op op, PatternRewriter &rewriter) const {
    return commBinOpElementWise<Op, enzymexla::ExtendOp>(op, rewriter);
  }
};

template <typename Op>
struct WrapElementwise
    : public CheckedOpRewritePattern<Op, WrapElementwise<Op>> {
  using CheckedOpRewritePattern<Op,
                                WrapElementwise<Op>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(Op op, PatternRewriter &rewriter) const {
    return commBinOpElementWise<Op, enzymexla::WrapOp>(op, rewriter);
  }
};

LogicalResult isExtendLike(int dim, Value _lhs, Value _mid, Value _rhs,
                           Location loc, RewriterBase &rewriter,
                           StaticSlice *lhsSS = nullptr,
                           StaticSlice *midSS = nullptr,
                           StaticSlice *rhsSS = nullptr) {
  std::optional<StaticSlice> lhs, mid, rhs;
  if (_lhs)
    lhs = StaticSlice::get(_lhs);
  if (_mid)
    mid = StaticSlice::get(_mid);
  if (_rhs)
    rhs = StaticSlice::get(_rhs);

  if (!mid)
    return rewriter.notifyMatchFailure(loc, "lhs or mid or rhs not slice");

  if (!lhs && !rhs) {
    return rewriter.notifyMatchFailure(loc, "lhs or rhs must be slice");
  }

  if (_lhs && !StaticSlice::isPrefixInDim(*lhs, *mid, dim))
    return rewriter.notifyMatchFailure(loc, "lhs not a prefix of mid");
  if (_lhs && lhs->getOutputShape(dim) != 1) {
    return failure();
  }
  if (_rhs && rhs->getOutputShape(dim) != 1) {
    return failure();
  }
  if (_rhs && !StaticSlice::isSuffixInDim(*rhs, *mid, dim))
    return rewriter.notifyMatchFailure(loc, "rhs is not suffix of mid");
  if (_rhs && !rhs->isStrideOneAtDim(dim))
    return rewriter.notifyMatchFailure(loc, "RHS not stride one");
  if (_lhs && !lhs->isStrideOneAtDim(dim))
    return rewriter.notifyMatchFailure(loc, "LHS not stride one");
  if (!mid->isStrideOneAtDim(dim))
    return rewriter.notifyMatchFailure(loc, "Mid not stride one");
  if (_rhs && rhs->getInput() != mid->getInput())
    return rewriter.notifyMatchFailure(loc,
                                       "mid and rhs not on the same input");
  if (_lhs && lhs->getInput() != mid->getInput())
    return rewriter.notifyMatchFailure(loc,
                                       "mid and lhs not on the same input");

  if (_lhs && lhsSS)
    *lhsSS = *lhs;
  if (_rhs && rhsSS)
    *rhsSS = *rhs;
  if (midSS)
    *midSS = *mid;

  return success();
}

struct RecognizeExtend
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     RecognizeExtend> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concat,
                                    PatternRewriter &rewriter) const {
    unsigned dim = concat.getDimension();
    if (concat.getNumOperands() == 2) {
      StaticSlice lhs;
      StaticSlice mid;

      if (succeeded(isExtendLike(dim, concat.getOperand(0),
                                 concat.getOperand(1), nullptr, concat.getLoc(),
                                 rewriter, &lhs, &mid, nullptr))) {
        auto extend = enzymexla::ExtendOp::create(
            rewriter, concat.getLoc(), mid.getOutput(), lhs.getOutputShape(dim),
            0, dim);
        if (auto shard = sdy::getShardingPerValue(concat)) {
          sdy::setShardings(extend, shard);
        }
        rewriter.replaceOp(concat, extend);
        return success();
      }
    }
    if (concat.getNumOperands() == 2) {
      StaticSlice rhs;
      StaticSlice mid;

      if (succeeded(isExtendLike(dim, nullptr, concat.getOperand(0),
                                 concat.getOperand(1), concat.getLoc(),
                                 rewriter, nullptr, &mid, &rhs))) {
        auto extend = enzymexla::ExtendOp::create(rewriter, concat.getLoc(),
                                                  mid.getOutput(), 0,
                                                  rhs.getOutputShape(dim), dim);
        if (auto shard = sdy::getShardingPerValue(concat)) {
          sdy::setShardings(extend, shard);
        }
        rewriter.replaceOp(concat, extend);
        return success();
      }
    }
    for (unsigned i = 2; i < concat.getNumOperands(); i++) {
      auto finish = [&](Value extend) {
        SmallVector<Value> toConcat;
        for (unsigned j = 0; j < i - 2; j++)
          toConcat.push_back(concat.getOperand(j));
        toConcat.push_back(extend);
        for (unsigned j = i + 1; j < concat.getNumOperands(); j++)
          toConcat.push_back(concat.getOperand(j));

        if (toConcat.size() == 1)
          rewriter.replaceOp(concat, extend);
        else {
          auto shard = sdy::getShardingPerValue(concat);
          auto newConcat =
              rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
                  concat, toConcat, dim);
          if (shard) {
            sdy::setShardings(newConcat, shard);
          }
        }
      };
      Value lhsv = concat.getOperand(i - 2);
      Value midv = concat.getOperand(i - 1);
      Value rhsv = concat.getOperand(i - 0);

      StaticSlice lhs;
      StaticSlice mid;
      StaticSlice rhs;

      if (succeeded(isExtendLike(dim, lhsv, midv, rhsv, concat.getLoc(),
                                 rewriter, &lhs, &mid, &rhs))) {
        auto extend = enzymexla::ExtendOp::create(
            rewriter, concat.getLoc(), mid.getOutput(), lhs.getOutputShape(dim),
            rhs.getOutputShape(dim), dim);
        if (auto shard = sdy::getShardingPerValue(concat)) {
          sdy::setShardings(extend, shard);
        }
        finish(extend);
        return success();
      }

      std::optional<unsigned> removedDim = std::nullopt;
      auto peelReshape = [&](Value v) -> Value {
        if (auto reshape = v.getDefiningOp<stablehlo::ReshapeOp>()) {
          auto inShape = reshape.getOperand().getType().getShape();
          auto outShape = reshape.getResult().getType().getShape();
          if (inShape.size() != outShape.size() + 1)
            return nullptr;

          if (inShape[0] == 1) {
            bool legal = true;
            for (auto &&[lhs, rhs] :
                 llvm::zip_equal(outShape, inShape.slice(1))) {
              if (lhs != rhs) {
                legal = false;
                break;
              }
            }
            if (legal) {
              if (removedDim) {
                if (*removedDim != 0) {
                  return nullptr;
                }
              } else {
                removedDim = 0;
              }
              return reshape.getOperand();
            }
          }

          for (unsigned inI = 0, outI = 0; inI < inShape.size();) {
            if (outI == outShape.size())
              return nullptr;
            if (inShape[inI] == outShape[outI]) {
              inI++;
              outI++;
            } else if (inShape[inI] == 1) {
              if (removedDim == inI) {
                inI++;
              } else if (!removedDim) {
                removedDim = inI;
                inI++;
              } else {
                return nullptr;
              }
            } else {
              return nullptr;
            }
          }
          return reshape.getOperand();
        }
        return nullptr;
      };
      lhsv = peelReshape(lhsv);
      midv = peelReshape(midv);
      rhsv = peelReshape(rhsv);
      if (!lhsv || !midv || !rhsv || !removedDim) {
        continue;
      }
      unsigned reshapedDim = dim;
      if (reshapedDim >= *removedDim)
        reshapedDim++;

      if (succeeded(isExtendLike(reshapedDim, lhsv, midv, rhsv, concat.getLoc(),
                                 rewriter, &lhs, &mid, &rhs))) {
        auto midO = mid.getOutput();
        auto extend = enzymexla::ExtendOp::create(
            rewriter, concat.getLoc(), midO, lhs.getOutputShape(reshapedDim),
            rhs.getOutputShape(reshapedDim), reshapedDim);
        if (auto midOp = midO.getDefiningOp())
          if (auto shard = sdy::getShardingPerValue(midOp)) {
            sdy::setShardings(extend, shard);
          }
        auto shape = llvm::to_vector(extend.getResult().getType().getShape());
        assert(shape[*removedDim] == 1);
        shape.erase(std::next(shape.begin(), *removedDim),
                    std::next(shape.begin(), *removedDim + 1));
        auto reshape = stablehlo::ReshapeOpCreate(rewriter, concat.getLoc(),
                                                  extend, shape);
        if (auto shard = sdy::getShardingPerValue(concat)) {
          sdy::setShardings(reshape.getDefiningOp(), shard);
        }
        finish(reshape);
        return success();
      }
    }
    return rewriter.notifyMatchFailure(concat, "Could not find extend pattern");
  }
};

bool canFuseIntoSingleSlice(int dimension, ArrayRef<Value> vals) {
  for (int i = 1; i < vals.size(); i++) {
    auto sl0 = vals[i - 1].getDefiningOp<stablehlo::SliceOp>();
    auto sl1 = vals[i].getDefiningOp<stablehlo::SliceOp>();
    if (!sl0 || !sl1 || !canMergeSlicesAlongAxis(dimension, sl0, sl1)) {
      return false;
    }
  }
  return true;
}

bool canFuseIntoSingleSlice(stablehlo::ConcatenateOp concatOp) {
  return canFuseIntoSingleSlice(concatOp.getDimension(),
                                llvm::to_vector(concatOp.getOperands()));
}

bool canFuseIntoSingleSlice(enzymexla::ExtendOp op) {
  return op.getLhs() == 0 && op.getRhs() == 0;
}

bool isAxisFusible(int dimension, ArrayRef<Value> vals) {
  assert(vals.size());

  for (int i = 0; i < vals.size(); i++) {
    if (auto concat = vals[i].getDefiningOp<stablehlo::ConcatenateOp>()) {
      if (concat.getDimension() == dimension)
        return true;
    }
  }

  for (int i = 1; i < vals.size(); i++) {
    auto sl0 = vals[i - 1].getDefiningOp<stablehlo::SliceOp>();
    auto sl1 = vals[i].getDefiningOp<stablehlo::SliceOp>();
    if (sl0 && sl1 && canMergeSlicesAlongAxis(dimension, sl0, sl1)) {
      return true;
    }
  }

  for (int i = 1; i < vals.size(); i++) {
    auto pad0 = vals[i - 1].getDefiningOp<stablehlo::PadOp>();
    auto pad1 = vals[i].getDefiningOp<stablehlo::PadOp>();
    if (pad0 && pad1 && canMergePadsAlongAxis(dimension, pad0, pad1)) {
      return true;
    }
  }

  for (int i = 1; i < vals.size(); i++) {
    auto sl0 = vals[i - 1].getDefiningOp<enzymexla::WrapOp>();
    auto sl1 = vals[i].getDefiningOp<enzymexla::WrapOp>();
    if (sl0 && sl1 && canMergeWrapsAlongAxis(dimension, sl0, sl1)) {
      return true;
    }
  }

  for (int i = 1; i < vals.size(); i++) {
    if (isRotateLike(dimension, vals[i - 1], vals[i])) {
      return true;
    }
    if (isExtendRotateLike(dimension, vals[i - 1], vals[i])) {
      return true;
    }
  }

  for (int i = 2; i < vals.size(); i++) {
    if (isWrapLike(dimension, vals[i - 2], vals[i - 1], vals[i])) {
      return true;
    }
  }

  auto rewriter = IRRewriter(vals[0].getContext());
  for (int i = 2; i < vals.size(); i++) {
    if (succeeded(isExtendLike(dimension, vals[i - 2], vals[i - 1], vals[i],
                               vals[i].getLoc(), rewriter))) {
      return true;
    }
  }

  return false;
}

struct SliceExtend final
    : CheckedOpRewritePattern<enzymexla::ExtendOp, SliceExtend> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  struct CandidateInfo {
    enzymexla::ExtendOp extendOp;
    stablehlo::SliceOp sliceOp; // Null if it's a direct extend user
  };

  LogicalResult matchAndRewriteImpl(enzymexla::ExtendOp triggerExtendOp,
                                    PatternRewriter &rewriter) const {

    Value triggerOperand = triggerExtendOp.getOperand();
    auto triggerSliceOp = triggerOperand.getDefiningOp<stablehlo::SliceOp>();

    if (!triggerSliceOp) {
      return failure();
    }

    Value baseOperand = triggerSliceOp.getOperand();

    // --- Get Target Extension Parameters ---
    int64_t targetExtendDim = triggerExtendOp.getDimension();
    int64_t targetLhs = triggerExtendOp.getLhs();
    int64_t targetRhs = triggerExtendOp.getRhs();
    Location loc = triggerExtendOp.getLoc();

    // --- Check Validity of Trigger Slice ---
    auto baseOperandType = dyn_cast<RankedTensorType>(baseOperand.getType());
    if (!baseOperandType || !baseOperandType.hasStaticShape()) {
      return rewriter.notifyMatchFailure(triggerExtendOp,
                                         "Base operand requires static shape");
    }
    if (triggerSliceOp.getStartIndices()[targetExtendDim] != 0 ||
        triggerSliceOp.getLimitIndices()[targetExtendDim] !=
            baseOperandType.getShape()[targetExtendDim] ||
        triggerSliceOp.getStrides()[targetExtendDim] != 1) {
      return rewriter.notifyMatchFailure(
          triggerExtendOp,
          "Trigger SliceOp modifies the dimension being extended");
    }

    llvm::SmallVector<CandidateInfo> candidates;
    candidates.push_back({triggerExtendOp, triggerSliceOp});

    for (auto const &userOp : baseOperand.getUsers()) {
      // Skip the slice that defines the trigger operand
      if (userOp == triggerSliceOp.getOperation())
        continue;

      // Case 1: Direct Extend
      if (auto directExtend = dyn_cast<enzymexla::ExtendOp>(userOp)) {
        if (directExtend.getDimension() == targetExtendDim &&
            directExtend.getLhs() == targetLhs &&
            directExtend.getRhs() == targetRhs) {
          candidates.push_back({directExtend, nullptr});
        }
      }
      // Case 2: Extend of Slice
      else if (auto sliceUser = dyn_cast<stablehlo::SliceOp>(userOp)) {
        if (!sliceUser->hasOneUse())
          continue;
        auto extendOfSlice =
            dyn_cast<enzymexla::ExtendOp>(*sliceUser->user_begin());
        if (!extendOfSlice)
          continue;

        if (extendOfSlice.getDimension() == targetExtendDim &&
            extendOfSlice.getLhs() == targetLhs &&
            extendOfSlice.getRhs() == targetRhs) {
          // Check validity: sliceUser must not modify targetExtendDim
          if (sliceUser.getStartIndices()[targetExtendDim] == 0 &&
              sliceUser.getLimitIndices()[targetExtendDim] ==
                  baseOperandType.getShape()[targetExtendDim] &&
              sliceUser.getStrides()[targetExtendDim] == 1) {
            candidates.push_back({extendOfSlice, sliceUser});
          }
        }
      }
    }

    if (candidates.size() <= 1) {
      return rewriter.notifyMatchFailure(
          triggerExtendOp,
          "Rewrite condition not met (only found the trigger candidate)");
    }

    SmallVector<int64_t> newBaseExtendShape =
        llvm::to_vector(baseOperandType.getShape());
    newBaseExtendShape[targetExtendDim] += (targetLhs + targetRhs);
    auto newBaseExtendType = RankedTensorType::get(
        newBaseExtendShape, baseOperandType.getElementType());

    if (auto subOp = baseOperand.getDefiningOp())
      rewriter.setInsertionPointAfter(subOp);
    else
      rewriter.setInsertionPointToStart(
          cast<BlockArgument>(baseOperand).getOwner());

    auto newBaseExtendOp = enzymexla::ExtendOp::create(
        rewriter, loc, newBaseExtendType, baseOperand, targetLhs, targetRhs,
        targetExtendDim);
    RankedTensorType newBaseExtendResultType = newBaseExtendOp.getType();

    for (const auto &candidate : candidates) {
      enzymexla::ExtendOp oldExtendOp = candidate.extendOp;
      stablehlo::SliceOp oldSliceOp = candidate.sliceOp; // Might be null

      if (!oldSliceOp) {
        // Direct Extend - Replace directly
        if (oldExtendOp.getResult().getType() == newBaseExtendOp.getType()) {
          rewriter.replaceOp(oldExtendOp, newBaseExtendOp);
        } else {
          auto castOp = tensor::CastOp::create(
              rewriter, loc, oldExtendOp.getResult().getType(),
              newBaseExtendOp);
          rewriter.replaceOp(oldExtendOp, castOp);
        }
      } else {
        SmallVector<int64_t> newSliceStarts =
            llvm::to_vector(oldSliceOp.getStartIndices());
        SmallVector<int64_t> newSliceLimits =
            llvm::to_vector(oldSliceOp.getLimitIndices());
        SmallVector<int64_t> newSliceStrides =
            llvm::to_vector(oldSliceOp.getStrides());

        newSliceStarts[targetExtendDim] = 0;
        newSliceLimits[targetExtendDim] =
            newBaseExtendResultType.getDimSize(targetExtendDim);
        newSliceStrides[targetExtendDim] = 1;
        assert(newBaseExtendOp);
        rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
            oldExtendOp, newBaseExtendOp, newSliceStarts, newSliceLimits,
            newSliceStrides);
      }
    }

    return success();
  }
};

struct SliceWrap final : CheckedOpRewritePattern<enzymexla::WrapOp, SliceWrap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  struct CandidateInfo {
    enzymexla::WrapOp wrapOp;
    stablehlo::SliceOp sliceOp; // Null if it's a direct wrap user
  };

  LogicalResult matchAndRewriteImpl(enzymexla::WrapOp triggerWrapOp,
                                    PatternRewriter &rewriter) const {

    Value triggerOperand = triggerWrapOp.getOperand();
    auto triggerSliceOp = triggerOperand.getDefiningOp<stablehlo::SliceOp>();

    if (!triggerSliceOp) {
      return failure();
    }

    Value baseOperand = triggerSliceOp.getOperand();

    // --- Get Target Wrap Parameters ---
    int64_t targetWrapDim = triggerWrapOp.getDimension();
    int64_t targetLhs = triggerWrapOp.getLhs();
    int64_t targetRhs = triggerWrapOp.getRhs();
    Location loc = triggerWrapOp.getLoc();

    // --- Check Validity of Trigger Slice ---
    auto baseOperandType = dyn_cast<RankedTensorType>(baseOperand.getType());
    if (!baseOperandType || !baseOperandType.hasStaticShape()) {
      return rewriter.notifyMatchFailure(triggerWrapOp,
                                         "Base operand requires static shape");
    }
    if (triggerSliceOp.getStartIndices()[targetWrapDim] != 0 ||
        triggerSliceOp.getLimitIndices()[targetWrapDim] !=
            baseOperandType.getShape()[targetWrapDim] ||
        triggerSliceOp.getStrides()[targetWrapDim] != 1) {
      return rewriter.notifyMatchFailure(
          triggerWrapOp,
          "Trigger SliceOp modifies the dimension being wrapped");
    }

    llvm::SmallVector<CandidateInfo> candidates;
    candidates.push_back({triggerWrapOp, triggerSliceOp});

    for (auto const &userOp : baseOperand.getUsers()) {
      // Skip the slice that defines the trigger operand
      if (userOp == triggerSliceOp.getOperation())
        continue;

      // Case 1: Direct Wrap
      if (auto directWrap = dyn_cast<enzymexla::WrapOp>(userOp)) {
        if (directWrap.getDimension() == targetWrapDim &&
            directWrap.getLhs() == targetLhs &&
            directWrap.getRhs() == targetRhs) {
          candidates.push_back({directWrap, nullptr});
        }
      }
      // Case 2: Wrap of Slice
      else if (auto sliceUser = dyn_cast<stablehlo::SliceOp>(userOp)) {
        if (!sliceUser->hasOneUse())
          continue;
        auto wrapOfSlice =
            dyn_cast<enzymexla::WrapOp>(*sliceUser->user_begin());
        if (!wrapOfSlice)
          continue;

        if (wrapOfSlice.getDimension() == targetWrapDim &&
            wrapOfSlice.getLhs() == targetLhs &&
            wrapOfSlice.getRhs() == targetRhs) {
          // Check validity: sliceUser must not modify targetWrapDim
          if (sliceUser.getStartIndices()[targetWrapDim] == 0 &&
              sliceUser.getLimitIndices()[targetWrapDim] ==
                  baseOperandType.getShape()[targetWrapDim] &&
              sliceUser.getStrides()[targetWrapDim] == 1) {
            candidates.push_back({wrapOfSlice, sliceUser});
          }
        }
      }
    }

    if (candidates.size() <= 1) {
      return rewriter.notifyMatchFailure(
          triggerWrapOp,
          "Rewrite condition not met (only found the trigger candidate)");
    }

    SmallVector<int64_t> newBaseWrapShape =
        llvm::to_vector(baseOperandType.getShape());
    newBaseWrapShape[targetWrapDim] += (targetLhs + targetRhs);
    auto newBaseWrapType = RankedTensorType::get(
        newBaseWrapShape, baseOperandType.getElementType());

    if (auto subOp = baseOperand.getDefiningOp())
      rewriter.setInsertionPointAfter(subOp);
    else
      rewriter.setInsertionPointToStart(
          cast<BlockArgument>(baseOperand).getOwner());

    auto newBaseWrapOp =
        enzymexla::WrapOp::create(rewriter, loc, newBaseWrapType, baseOperand,
                                  targetLhs, targetRhs, targetWrapDim);
    Value newBaseWrapResult = newBaseWrapOp.getResult();
    RankedTensorType newBaseWrapResultType =
        cast<RankedTensorType>(newBaseWrapResult.getType());

    for (const auto &candidate : candidates) {
      enzymexla::WrapOp oldWrapOp = candidate.wrapOp;
      stablehlo::SliceOp oldSliceOp = candidate.sliceOp; // Might be null

      if (!oldSliceOp) {
        // Direct Wrap - Replace directly
        if (oldWrapOp.getResult().getType() == newBaseWrapResult.getType()) {
          rewriter.replaceOp(oldWrapOp, newBaseWrapResult);
        } else {
          auto castOp = tensor::CastOp::create(rewriter, loc,
                                               oldWrapOp.getResult().getType(),
                                               newBaseWrapResult);
          rewriter.replaceOp(oldWrapOp, castOp);
        }
      } else {
        SmallVector<int64_t> newSliceStarts =
            llvm::to_vector(oldSliceOp.getStartIndices());
        SmallVector<int64_t> newSliceLimits =
            llvm::to_vector(oldSliceOp.getLimitIndices());
        SmallVector<int64_t> newSliceStrides =
            llvm::to_vector(oldSliceOp.getStrides());

        newSliceStarts[targetWrapDim] = 0;
        newSliceLimits[targetWrapDim] =
            newBaseWrapResultType.getDimSize(targetWrapDim);
        newSliceStrides[targetWrapDim] = 1;

        auto newSlice = stablehlo::SliceOp::create(
            rewriter, oldWrapOp.getLoc(),
            oldWrapOp.getResult()
                .getType(), // Use original wrap op's result type
            newBaseWrapResult, newSliceStarts, newSliceLimits, newSliceStrides);
        rewriter.replaceAllOpUsesWith(oldWrapOp, newSlice.getResult());
      }
    }

    return success();
  }
};

// transpose(wrap) -> wrap(transpose)
struct TransposeWrap final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeWrap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    // Check if the operand is a wrap operation
    auto wrapOp = op.getOperand().getDefiningOp<enzymexla::WrapOp>();
    if (!wrapOp)
      return failure();

    // Get wrap operation parameters
    int64_t wrapDim = wrapOp.getDimension();
    int64_t lhs = wrapOp.getLhs();
    int64_t rhs = wrapOp.getRhs();

    // Get permutation array
    SmallVector<int64_t> permutation;
    for (auto val : op.getPermutation()) {
      permutation.push_back(static_cast<int64_t>(val));
    }

    // The new wrap dimension will be the permuted dimension
    int64_t newWrapDim = -1;
    for (size_t i = 0; i < permutation.size(); ++i) {
      if (permutation[i] == wrapDim) {
        newWrapDim = i;
        break;
      }
    }

    if (newWrapDim == -1)
      return failure();

    // First transpose the wrap's operand
    auto newTranspose = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), wrapOp.getOperand(), op.getPermutation());

    // Then create a new wrap operation on the transposed data
    auto newWrapType = op.getType();
    auto newWrapOp = enzymexla::WrapOp::create(
        rewriter, op.getLoc(), newWrapType, newTranspose.getResult(), lhs, rhs,
        newWrapDim);

    // Replace the original op with the new wrap operation
    rewriter.replaceOp(op, newWrapOp);
    return success();
  }
};

// transpose(extend) -> extend(transpose)
struct TransposeExtend final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeExtend> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    // Check if the operand is an extend operation
    auto extendOp = op.getOperand().getDefiningOp<enzymexla::ExtendOp>();
    if (!extendOp)
      return failure();

    // Get extend operation parameters
    int64_t extendDim = extendOp.getDimension();
    int64_t lhs = extendOp.getLhs();
    int64_t rhs = extendOp.getRhs();

    // Get permutation array
    SmallVector<int64_t> permutation;
    for (auto val : op.getPermutation()) {
      permutation.push_back(static_cast<int64_t>(val));
    }

    // The new extend dimension will be the permuted dimension
    int64_t newExtendDim = -1;
    for (size_t i = 0; i < permutation.size(); ++i) {
      if (permutation[i] == extendDim) {
        newExtendDim = i;
        break;
      }
    }

    if (newExtendDim == -1)
      return failure();

    // First transpose the extend's operand
    auto newTranspose = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), extendOp.getOperand(), op.getPermutation());

    // Then create a new extend operation on the transposed data
    auto newExtendOp = enzymexla::ExtendOp::create(rewriter, op.getLoc(),
                                                   newTranspose.getResult(),
                                                   lhs, rhs, newExtendDim);

    // Replace the original op with the new extend operation
    rewriter.replaceOp(op, newExtendOp);
    return success();
  }
};

// transpose(rotate) -> rotate(transpose)
struct TransposeRotate final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto type = dyn_cast<RankedTensorType>(op.getType());
    if (!type)
      return failure();

    // Check if the operand is a rotate operation
    auto rotateOp = op.getOperand().getDefiningOp<enzymexla::RotateOp>();
    if (!rotateOp)
      return failure();

    // Get rotate operation parameters
    int64_t rotateDim = rotateOp.getDimension();
    int64_t amount = rotateOp.getAmount();

    // Get permutation array
    SmallVector<int64_t> permutation;
    for (auto val : op.getPermutation()) {
      permutation.push_back(static_cast<int64_t>(val));
    }

    // The new rotate dimension will be the permuted dimension
    int64_t newRotateDim = -1;
    for (size_t i = 0; i < permutation.size(); ++i) {
      if (permutation[i] == rotateDim) {
        newRotateDim = i;
        break;
      }
    }

    if (newRotateDim == -1)
      return failure();

    // First transpose the rotate's operand
    auto newTranspose = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), rotateOp.getOperand(), op.getPermutation());

    // Then create a new rotate operation on the transposed data
    auto newRotateOp = enzymexla::RotateOp::create(
        rewriter, op.getLoc(), newTranspose.getResult(), amount, newRotateDim);

    // Replace the original op with the new rotate operation
    rewriter.replaceOp(op, newRotateOp);
    return success();
  }
};

struct ConcatConcatAxisSwap final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatConcatAxisSwap> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp outer,
                                    PatternRewriter &rewriter) const {
    if (outer.getOperands().size() < 2) {
      return failure();
    }

    SmallVector<Operation *> inners;
    int64_t innerConcatDim = -1, nInnerOperands = -1;
    SmallVector<int64_t> innerConcatDimSizes;

    // Get a full square of concat of compatible dimensions for axis swap
    for (auto v : outer.getOperands()) {
      auto defOp = v.getDefiningOp();
      if (!defOp) {
        return failure();
      }

      Operation *innerOp;
      bool supported =
          TypeSwitch<Operation *, bool>(defOp)
              .Case<stablehlo::ConcatenateOp>([&](auto op) {
                auto concatDim = op.getDimension();
                if (outer.getDimension() == concatDim ||
                    op.getOperands().size() < 2) {
                  return false;
                }
                if (!inners.empty()) {
                  if (nInnerOperands != op->getNumOperands() ||
                      innerConcatDim != concatDim) {
                    return false;
                  }
                  for (size_t i = 0; i < nInnerOperands; i++) {
                    if (cast<RankedTensorType>(op->getOperand(i).getType())
                            .getDimSize(concatDim) != innerConcatDimSizes[i]) {
                      return false;
                    }
                  }
                } else {
                  innerConcatDim = concatDim;
                  nInnerOperands = op->getNumOperands();
                  for (size_t i = 0; i < nInnerOperands; i++) {
                    innerConcatDimSizes.push_back(
                        cast<RankedTensorType>(op->getOperand(i).getType())
                            .getDimSize(concatDim));
                  }
                }
                innerOp = op;
                return true;
              })
              .Case<enzymexla::ExtendOp>([&](auto op) {
                auto concatDim = op.getDimension();
                if (outer.getDimension() == concatDim ||
                    (op.getLhs() == 0 && op.getRhs() == 0)) {
                  return false;
                }
                auto noperands = 1 + (op.getLhs() != 0) + (op.getRhs() != 0);
                if (!inners.empty()) {
                  if (nInnerOperands != noperands ||
                      innerConcatDim != concatDim) {
                    return false;
                  }
                  size_t idx = 0;
                  if (op.getLhs() != 0) {
                    if (op.getLhs() != innerConcatDimSizes[idx]) {
                      return false;
                    }
                    idx++;
                  }
                  if (cast<RankedTensorType>(op.getOperand().getType())
                          .getDimSize(concatDim) != innerConcatDimSizes[idx]) {
                    return false;
                  }
                  idx++;
                  if (op.getRhs() != 0) {
                    if (op.getRhs() != innerConcatDimSizes[idx]) {
                      return false;
                    }
                  }
                } else {
                  innerConcatDim = concatDim;
                  nInnerOperands = noperands;
                  if (op.getLhs() != 0) {
                    innerConcatDimSizes.push_back(op.getLhs());
                  }
                  innerConcatDimSizes.push_back(
                      cast<RankedTensorType>(op.getOperand().getType())
                          .getDimSize(concatDim));
                  if (op.getRhs() != 0) {
                    innerConcatDimSizes.push_back(op.getRhs());
                  }
                }
                innerOp = op;
                return true;
              })
              .Default([](auto op) { return false; });

      if (!supported) {
        return failure();
      }

      inners.push_back(innerOp);
    }

    if (llvm::any_of(inners, ([=](auto op) {
                       return TypeSwitch<Operation *, bool>(op)
                           .Case<stablehlo::ConcatenateOp, enzymexla::ExtendOp>(
                               [&](auto op) {
                                 return canFuseIntoSingleSlice(op);
                               })
                           .Default([](auto op) { return false; });
                     }))) {
      return failure();
    }

    // high-priority for fusion that creates a larger contiguous slice
    // We can only check this for ConcatenateOps directly. ExtendOps don't have
    // their virtual operands materialized until lowering, so we compute virtual
    // slice bounds for them.
    bool highPriorityFusion = false;
    bool allConcats = llvm::all_of(inners, [](Operation *op) {
      return isa<stablehlo::ConcatenateOp>(op);
    });

    if (allConcats) {
      // All inners are ConcatenateOps, we can directly check operands
      for (int i = 0; i < nInnerOperands && !highPriorityFusion; i++) {
        SmallVector<Value> newOperands;
        for (int j = 0; j < outer->getNumOperands(); j++) {
          newOperands.push_back(inners[j]->getOperand(i));
        }
        if (canFuseIntoSingleSlice(outer.getDimension(), newOperands)) {
          highPriorityFusion = true;
          break;
        }
      }
    } else {
      // Some inners are ExtendOps. For each virtual operand index, compute
      // virtual slice bounds and check if consecutive ones can be merged.
      // Virtual operand layout for ExtendOp: [lhs_pad, operand, rhs_pad]
      // where lhs_pad and rhs_pad are zero-padding regions (no defining op).
      for (int i = 0; i < nInnerOperands && !highPriorityFusion; i++) {
        // For each pair of consecutive values in outer, check if the virtual
        // operands at index i can be merged as slices along
        // outer.getDimension()
        bool canMergeAll = true;
        for (int j = 1; j < outer->getNumOperands() && canMergeAll; j++) {
          // Get the virtual operand info for inners[j-1] and inners[j] at index
          // i
          auto getSliceInfo =
              [&](Operation *op, int idx) -> std::optional<stablehlo::SliceOp> {
            if (auto concat = dyn_cast<stablehlo::ConcatenateOp>(op)) {
              auto slice =
                  concat.getOperands()[idx].getDefiningOp<stablehlo::SliceOp>();
              if (!slice) {
                return std::nullopt;
              }
              return slice;
            } else if (auto extend = dyn_cast<enzymexla::ExtendOp>(op)) {
              // ExtendOp virtual operands:
              // idx 0 (if lhs != 0): padding region of size lhs
              // idx (lhs != 0 ? 1 : 0): the actual operand
              // idx (last, if rhs != 0): padding region of size rhs
              // Padding regions have no source value so cannot be merged as
              // slices.
              int operandIdx = extend.getLhs() != 0 ? 1 : 0;
              if (idx != operandIdx) {
                // This is a padding region, not a real slice
                return std::nullopt;
              }
              // The actual operand - get slice info from it if it's a slice
              auto slice =
                  extend.getOperand().getDefiningOp<stablehlo::SliceOp>();
              if (!slice) {
                return std::nullopt;
              }
              return slice;
            }
            return std::nullopt;
          };

          auto info0 = getSliceInfo(inners[j - 1], i);
          auto info1 = getSliceInfo(inners[j], i);

          if (!info0.has_value() || !info1.has_value()) {
            canMergeAll = false;
            continue;
          }

          if (!canMergeSlicesAlongAxis(outer.getDimension(), info0.value(),
                                       info1.value())) {
            canMergeAll = false;
          }
        }

        if (canMergeAll) {
          highPriorityFusion = true;
          break;
        }
      }
    }

    // Check that we don't have a current axis fuse opportunity, and wait for
    // those fusions
    if (!highPriorityFusion) {
      for (auto inner : inners) {
        auto innerConcat = dyn_cast<stablehlo::ConcatenateOp>(inner);
        if (!innerConcat) {
          return failure();
        }
        if (isAxisFusible(innerConcat.getDimension(),
                          llvm::to_vector(innerConcat.getOperands()))) {
          return failure();
        }
      }
    }

    bool anyFusible = highPriorityFusion;
    for (int i = 0; i < nInnerOperands && !anyFusible; i++) {
      SmallVector<Value> newOperands;
      for (int j = 0; j < outer.getOperands().size(); j++) {
        auto innerConcat = dyn_cast<stablehlo::ConcatenateOp>(inners[j]);
        if (!innerConcat) {
          return failure();
        }
        newOperands.push_back(innerConcat.getOperands()[i]);
      }
      if (isAxisFusible(outer.getDimension(), newOperands)) {
        anyFusible = true;
        break;
      }
    }

    if (!anyFusible) {
      return failure();
    }

    SmallVector<Value> newOuters;

    // lower all the extends (if any)
    for (size_t i = 0; i < inners.size(); i++) {
      auto extendOp = dyn_cast<enzymexla::ExtendOp>(inners[i]);
      if (extendOp) {
        inners[i] = lowerExtend(extendOp, rewriter, true);
      }
    }

    for (int i = 0; i < nInnerOperands; i++) {
      SmallVector<Value> newOperands;
      for (int j = 0; j < outer.getNumOperands(); j++) {
        newOperands.push_back(inners[j]->getOperand(i));
      }
      newOuters.push_back(stablehlo::ConcatenateOp::create(
          rewriter, outer.getLoc(), newOperands, outer.getDimension()));
    }
    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(outer, newOuters,
                                                          innerConcatDim);
    return success();
  }
};

struct SliceRotate final
    : CheckedOpRewritePattern<enzymexla::RotateOp, SliceRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  struct CandidateInfo {
    enzymexla::RotateOp rotateOp;
    stablehlo::SliceOp sliceOp; // Null if it's a direct extend user
  };

  LogicalResult matchAndRewriteImpl(enzymexla::RotateOp triggerRotateOp,
                                    PatternRewriter &rewriter) const {

    Value triggerOperand = triggerRotateOp.getOperand();
    auto triggerSliceOp = triggerOperand.getDefiningOp<stablehlo::SliceOp>();

    if (!triggerSliceOp) {
      return failure();
    }

    Value baseOperand = triggerSliceOp.getOperand();

    // --- Get Target Rotate Parameters ---
    int targetAmount = triggerRotateOp.getAmount();
    int targetRotateDim = triggerRotateOp.getDimension();
    Location loc = triggerRotateOp.getLoc();

    // --- Check Validity of Trigger Slice ---
    auto baseOperandType = dyn_cast<RankedTensorType>(baseOperand.getType());
    if (!baseOperandType || !baseOperandType.hasStaticShape()) {
      return rewriter.notifyMatchFailure(triggerRotateOp,
                                         "Base operand requires static shape");
    }
    if (triggerSliceOp.getStartIndices()[targetRotateDim] != 0 ||
        triggerSliceOp.getLimitIndices()[targetRotateDim] !=
            baseOperandType.getShape()[targetRotateDim] ||
        triggerSliceOp.getStrides()[targetRotateDim] != 1) {
      return rewriter.notifyMatchFailure(
          triggerRotateOp,
          "Trigger SliceOp modifies the dimension being extended");
    }

    llvm::SmallVector<CandidateInfo> candidates;
    candidates.push_back({triggerRotateOp, triggerSliceOp});

    for (auto const &userOp : baseOperand.getUsers()) {
      // Skip the slice that defines the trigger operand
      if (userOp == triggerSliceOp.getOperation())
        continue;

      // Case 1: Direct Rotate
      if (auto directRotate = dyn_cast<enzymexla::RotateOp>(userOp)) {
        if (directRotate.getDimension() == targetRotateDim &&
            directRotate.getAmount() == targetAmount) {
          candidates.push_back({directRotate, nullptr});
        }
      }
      // Case 2: Rotate of Slice
      else if (auto sliceUser = dyn_cast<stablehlo::SliceOp>(userOp)) {
        if (!sliceUser->hasOneUse())
          continue;
        auto rotateOfSlice =
            dyn_cast<enzymexla::RotateOp>(*sliceUser->user_begin());
        if (!rotateOfSlice)
          continue;

        if (rotateOfSlice.getDimension() == targetRotateDim &&
            rotateOfSlice.getAmount() == targetAmount) {
          // Check validity: sliceUser must not modify targetRotateDim
          if (sliceUser.getStartIndices()[targetRotateDim] == 0 &&
              sliceUser.getLimitIndices()[targetRotateDim] ==
                  baseOperandType.getShape()[targetRotateDim] &&
              sliceUser.getStrides()[targetRotateDim] == 1) {
            candidates.push_back({rotateOfSlice, sliceUser});
          }
        }
      }
    }

    if (candidates.size() <= 1) {
      return rewriter.notifyMatchFailure(
          triggerRotateOp,
          "Rewrite condition not met (only found the trigger candidate)");
    }

    SmallVector<int64_t> newBaseExtendShape =
        llvm::to_vector(baseOperandType.getShape());

    auto newBaseRotateType = baseOperandType;

    if (auto subOp = baseOperand.getDefiningOp())
      rewriter.setInsertionPointAfter(subOp);
    else
      rewriter.setInsertionPointToStart(
          cast<BlockArgument>(baseOperand).getOwner());

    auto newBaseRotateOp =
        enzymexla::RotateOp::create(rewriter, loc, newBaseRotateType,
                                    baseOperand, targetAmount, targetRotateDim);
    Value newBaseRotateResult = newBaseRotateOp.getResult();
    RankedTensorType newBaseRotateResultType = newBaseRotateType;

    for (const auto &candidate : candidates) {
      enzymexla::RotateOp oldRotateOp = candidate.rotateOp;
      stablehlo::SliceOp oldSliceOp = candidate.sliceOp; // Might be null

      if (!oldSliceOp) {
        // Direct Rotate - Replace directly
        rewriter.replaceOp(oldRotateOp, newBaseRotateResult);
      } else {
        SmallVector<int64_t> newSliceStarts =
            llvm::to_vector(oldSliceOp.getStartIndices());
        SmallVector<int64_t> newSliceLimits =
            llvm::to_vector(oldSliceOp.getLimitIndices());
        SmallVector<int64_t> newSliceStrides =
            llvm::to_vector(oldSliceOp.getStrides());

        newSliceStarts[targetRotateDim] = 0;
        newSliceLimits[targetRotateDim] =
            newBaseRotateResultType.getDimSize(targetRotateDim);
        newSliceStrides[targetRotateDim] = 1;

        auto newSlice = stablehlo::SliceOp::create(
            rewriter, oldRotateOp.getLoc(),
            oldRotateOp.getResult()
                .getType(), // Use original extend op's result type
            newBaseRotateResult, newSliceStarts, newSliceLimits,
            newSliceStrides);
        rewriter.replaceAllOpUsesWith(oldRotateOp, newSlice.getResult());
      }
    }

    return success();
  }
};

struct SquareAbsSimplify
    : public CheckedOpRewritePattern<stablehlo::MulOp, SquareAbsSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    if (lhs != rhs)
      return failure();

    auto absOp = lhs.getDefiningOp<stablehlo::AbsOp>();
    if (!absOp)
      return failure();

    auto operand = absOp.getOperand();
    auto operandType = dyn_cast<RankedTensorType>(operand.getType());
    if (!operandType)
      return failure();

    if (isa<ComplexType>(operandType.getElementType())) {
      // abs(z)^2 = real(z)^2 + imag(z)^2 -- only applied if abs(z) is used in
      // this operation
      if (!isOnlyUsedInOperation(absOp, op)) {
        return failure();
      }
      auto real = stablehlo::RealOp::create(rewriter, op.getLoc(), operand);
      auto imag = stablehlo::ImagOp::create(rewriter, op.getLoc(), operand);
      auto realSq = stablehlo::MulOp::create(rewriter, op.getLoc(), real, real);
      auto imagSq = stablehlo::MulOp::create(rewriter, op.getLoc(), imag, imag);
      rewriter.replaceOpWithNewOp<stablehlo::AddOp>(op, realSq, imagSq);
      return success();
    } else {
      // abs(x)^2 = x * x
      rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, operand, operand);
      return success();
    }
  }
};

struct ConcatBroadcastSlice
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     ConcatBroadcastSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto dim = op.getDimension();
    SmallVector<Value> newOperands;

    SmallVector<Value> oldOperands(op.getOperands());
    auto res =
        concatBroadcastSliceSimplify(rewriter, oldOperands, dim, newOperands);
    if (!res.succeeded())
      return res;

    rewriter.modifyOpInPlace(op, [&]() { op->setOperands(newOperands); });
    return success();
  }
};

struct ConcatReshapeSlice
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     ConcatReshapeSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto dim = op.getDimension();
    SmallVector<Value> newOperands;

    auto oldOperands = llvm::to_vector(op.getOperands());
    auto res =
        concatReshapeSliceSimplify(rewriter, oldOperands, dim, newOperands);
    if (!res.succeeded())
      return res;

    rewriter.modifyOpInPlace(op, [&]() { op->setOperands(newOperands); });
    return success();
  }
};

bool reshapeOfEquivalentReduces(stablehlo::ReshapeOp reshapeOp,
                                SmallVector<stablehlo::ReduceOp> &allOperands,
                                SmallVector<Value> &reduceOpOperands) {
  auto rank =
      cast<RankedTensorType>(reshapeOp.getOperand().getType()).getRank();
  if (rank != 0)
    return false;

  auto reduceOp = reshapeOp.getOperand().getDefiningOp<stablehlo::ReduceOp>();
  if (!reduceOp)
    return false;

  if (!isOnlyUsedInOperation(reduceOp, reshapeOp))
    return false;

  if (reduceOp.getInputs().size() != 1)
    return false;

  if (allOperands.size() >= 1 &&
      !OperationEquivalence::isEquivalentTo(
          reduceOp, allOperands[0],
          OperationEquivalence::ignoreValueEquivalence, nullptr,
          OperationEquivalence::IgnoreLocations, nullptr))
    return false;

  reduceOpOperands.push_back(reduceOp.getInputs()[0]);
  allOperands.push_back(reduceOp);
  return true;
}

struct ConcatElementwise final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatElementwise> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concatOp,
                                    PatternRewriter &rewriter) const {
    if (concatOp.getNumOperands() <= 1)
      return failure();

    SmallVector<Operation *> concatOpOperands;

    for (auto [i, v] : llvm::enumerate(concatOp.getOperands())) {
      auto vdefOp = v.getDefiningOp();
      if (!vdefOp)
        return failure();

      if (isa<stablehlo::ConvertOp>(vdefOp)) // Conflicts with ConvertConcat
        return failure();

      if (hasTraitElementwise(vdefOp)) {
        if (concatOpOperands.size() != 0) {
          if (!OperationEquivalence::isEquivalentTo(
                  concatOpOperands[0], vdefOp,
                  OperationEquivalence::ignoreValueEquivalence, nullptr,
                  OperationEquivalence::IgnoreLocations, nullptr))
            return failure();
        }

        if (!isOnlyUsedInOperation(vdefOp, concatOp))
          return failure();

        concatOpOperands.push_back(vdefOp);
      } else {
        return failure();
      }
    }

    SmallVector<Value> elementwiseOperands;

    for (int i = 0; i < concatOpOperands[0]->getNumOperands(); i++) {
      SmallVector<Value> newConcatOperands;
      for (auto v : concatOpOperands) {
        newConcatOperands.push_back(v->getOperand(i));
      }
      auto newConcatOp = stablehlo::ConcatenateOp::create(
          rewriter, concatOp.getLoc(), newConcatOperands,
          concatOp.getDimension());
      elementwiseOperands.push_back(newConcatOp.getResult());
    }

    auto newElementwiseOp = rewriter.create(
        concatOp.getLoc(), concatOpOperands[0]->getName().getIdentifier(),
        ValueRange(elementwiseOperands), TypeRange(concatOp.getType()),
        concatOpOperands[0]->getAttrs(), {}, {});

    rewriter.replaceOp(concatOp, newElementwiseOp);
    return success();
  }
};

struct ConcatReshapeReduce final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatReshapeReduce> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concatOp,
                                    PatternRewriter &rewriter) const {
    if (concatOp.getNumOperands() <= 1)
      return failure();

    SmallVector<stablehlo::ReduceOp> allOperands;
    SmallVector<Value> reduceOpOperands;
    for (auto v : concatOp.getOperands()) {
      if (auto reshapeOp = v.getDefiningOp<stablehlo::ReshapeOp>()) {
        if (!isOnlyUsedInOperation(reshapeOp, concatOp))
          return failure();
        if (!reshapeOfEquivalentReduces(reshapeOp, allOperands,
                                        reduceOpOperands))
          return failure();
      } else {
        return rewriter.notifyMatchFailure(concatOp,
                                           "Operand is not a reshape.");
      }
    }

    auto reduceDims = llvm::to_vector(allOperands[0].getDimensions());

    auto concatDim = concatOp.getDimension();
    for (int64_t i = 0; i < reduceDims.size(); i++) {
      if (reduceDims[i] >= concatDim)
        reduceDims[i]++;
    }

    SmallVector<Value> newConcatOperands;
    auto elemTy =
        cast<RankedTensorType>(reduceOpOperands[0].getType()).getElementType();

    SmallVector<int64_t> oldShape = llvm::to_vector(
        cast<RankedTensorType>(reduceOpOperands[0].getType()).getShape());

    SmallVector<int64_t> preConcatShape(oldShape.size() + 1, 0);
    for (int64_t i = 0; i < preConcatShape.size(); i++) {
      if (i == concatDim) {
        preConcatShape[i] = 1;
      } else if (i < concatDim) {
        preConcatShape[i] = oldShape[i];
      } else {
        preConcatShape[i] = oldShape[i - 1];
      }
    }

    for (int64_t i = 0; i < reduceOpOperands.size(); i++) {
      newConcatOperands.push_back(stablehlo::ReshapeOpCreate(
          rewriter, concatOp.getLoc(), reduceOpOperands[i], preConcatShape));
    }

    auto newConcatOp = stablehlo::ConcatenateOp::create(
        rewriter, concatOp.getLoc(), newConcatOperands, concatDim);

    auto reduceOp = stablehlo::ReduceOp::create(
        rewriter, concatOp.getLoc(),
        TypeRange(RankedTensorType::get(
            {static_cast<int64_t>(reduceOpOperands.size())}, elemTy)),
        ValueRange(newConcatOp), ValueRange(allOperands[0].getInitValues()[0]),
        rewriter.getDenseI64ArrayAttr(reduceDims));

    // Clone the reduction body
    rewriter.inlineRegionBefore(allOperands[0].getBody(), reduceOp.getBody(),
                                reduceOp.getBody().end());

    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(
        concatOp, concatOp.getType(), reduceOp.getResult(0));
    return success();
  }
};

// reverse(transpose x) -> transpose(reverse x)
struct ReverseTranspose final
    : CheckedOpRewritePattern<stablehlo::ReverseOp, ReverseTranspose> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReverseOp op,
                                    PatternRewriter &rewriter) const {
    auto transposeOp = op.getOperand().getDefiningOp<stablehlo::TransposeOp>();
    if (!transposeOp)
      return failure();

    if (!transposeOp->getResult(0).hasOneUse())
      return failure();

    auto permutation = transposeOp.getPermutation();
    SmallVector<int64_t> newReverseDims(op.getDimensions().size());
    for (auto [i, dim] : llvm::enumerate(op.getDimensions()))
      newReverseDims[i] = permutation[dim];

    auto newReverseOp = stablehlo::ReverseOp::create(
        rewriter, op.getLoc(), transposeOp.getOperand(), newReverseDims);
    rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(op, newReverseOp,
                                                        permutation);
    return success();
  }
};

// transpose(reverse x) -> reverse(transpose x)
struct TransposeReverse final
    : CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeReverse> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto reverseOp = op.getOperand().getDefiningOp<stablehlo::ReverseOp>();
    if (!reverseOp)
      return failure();

    if (!reverseOp->getResult(0).hasOneUse())
      return failure();

    auto invPerm = getInversePermutation(op.getPermutation());
    SmallVector<int64_t> newReverseDims(reverseOp.getDimensions().size());
    for (auto [i, dim] : llvm::enumerate(reverseOp.getDimensions()))
      newReverseDims[i] = invPerm[dim];

    auto newTranspose = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), reverseOp.getOperand(), op.getPermutation());
    rewriter.replaceOpWithNewOp<stablehlo::ReverseOp>(op, newTranspose,
                                                      newReverseDims);
    return success();
  }
};

struct ElementwiseReshapeLike
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          ElementwiseReshapeLike> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise,
      ElementwiseReshapeLike>::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    SmallVector<Value> parentOperands;
    Operation *operandOp = nullptr;
    for (auto operand : op->getOperands()) {
      auto defOp = operand.getDefiningOp();
      if (!defOp)
        return failure();

      if (!isa<stablehlo::ReshapeOp, stablehlo::BroadcastInDimOp>(defOp))
        return failure();

      auto defOpOperation = defOp->getOperand(0).getDefiningOp();
      if (defOpOperation) {
        bool canFuse =
            TypeSwitch<Operation *, bool>(defOp)
                .Case<stablehlo::BroadcastInDimOp, stablehlo::ReshapeOp>(
                    [&](auto op) { return isFusible(defOpOperation, op); })
                .Default([](auto op) { return false; });
        if (canFuse) { // other passes get priority here
          return failure();
        }
      }

      if (!isOnlyUsedInOperation(defOp, op))
        return rewriter.notifyMatchFailure(
            op, "operand is used in more than one op");

      if (!operandOp) {
        operandOp = defOp;
      } else {
        if (!OperationEquivalence::isEquivalentTo(
                operandOp, defOp, OperationEquivalence::ignoreValueEquivalence,
                nullptr, OperationEquivalence::IgnoreLocations, nullptr)) {
          return rewriter.notifyMatchFailure(
              op, "operand operations are not equivalent");
        }
      }

      parentOperands.push_back(defOp->getOperand(0));
    }

    if (parentOperands.size() == 0)
      return failure();

    auto elemOp = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(), ValueRange(parentOperands),
        TypeRange{RankedTensorType::get(
            cast<RankedTensorType>(parentOperands[0].getType()).getShape(),
            cast<RankedTensorType>(op->getResult(0).getType())
                .getElementType())},
        op->getAttrs(), {}, {});
    auto reshapeLikeOp = rewriter.create(
        op->getLoc(), operandOp->getName().getIdentifier(),
        ValueRange{elemOp->getResult(0)}, TypeRange{op->getResult(0).getType()},
        operandOp->getAttrs(), {}, {});
    rewriter.replaceOp(op, reshapeLikeOp);
    return success();
  }
};

struct ConcatTranspose final
    : CheckedOpRewritePattern<stablehlo::ConcatenateOp, ConcatTranspose> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getNumOperands() < 2)
      return failure();

    SmallVector<Value> parentOperands;
    stablehlo::TransposeOp transposeOp;
    for (auto operand : op.getOperands()) {
      auto defOp = operand.getDefiningOp<stablehlo::TransposeOp>();
      if (!defOp)
        return failure();

      if (!isOnlyUsedInOperation(defOp, op))
        return rewriter.notifyMatchFailure(
            op, "operand is used in more than one op");

      if (!transposeOp) {
        transposeOp = defOp;
      } else {
        if (!OperationEquivalence::isEquivalentTo(
                transposeOp, defOp,
                OperationEquivalence::ignoreValueEquivalence, nullptr,
                OperationEquivalence::IgnoreLocations, nullptr)) {
          return rewriter.notifyMatchFailure(
              op, "operand operations are not equivalent");
        }
      }

      parentOperands.push_back(defOp->getOperand(0));
    }

    auto permutation = transposeOp.getPermutation();
    auto newConcatDim = permutation[op.getDimension()];

    auto newConcat = stablehlo::ConcatenateOp::create(
        rewriter, op.getLoc(), parentOperands, newConcatDim);
    rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(op, newConcat,
                                                        permutation);
    return success();
  }
};

struct ReduceReduce final
    : CheckedOpRewritePattern<stablehlo::ReduceOp, ReduceReduce> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1)
      return rewriter.notifyMatchFailure(
          op, "reduce op has more than one input. not yet supported");

    auto redOp = op.getInputs()[0].getDefiningOp<stablehlo::ReduceOp>();
    if (!redOp)
      return failure();

    if (redOp.getInputs().size() != 1)
      return rewriter.notifyMatchFailure(
          op, "reduce op has more than one input. not yet supported");

    if (!llvm::hasSingleElement(op.getInputs()[0].getUses()))
      return failure();

    if (!OperationEquivalence::isEquivalentTo(
            redOp.getInitValues()[0].getDefiningOp(),
            op.getInitValues()[0].getDefiningOp(),
            OperationEquivalence::IgnoreLocations))
      return rewriter.notifyMatchFailure(
          op, "reduce op init values are not equivalent");

    if (!OperationEquivalence::isRegionEquivalentTo(
            &redOp.getBody(), &op.getBody(),
            OperationEquivalence::IgnoreLocations))
      return failure();

    SmallVector<int64_t> innerDimensions =
        llvm::to_vector(redOp.getDimensions());
    SmallVector<int64_t> outerDimensions = llvm::to_vector(op.getDimensions());
    SmallVector<int64_t> mergedDimensions = innerDimensions;

    int64_t inputRank =
        cast<RankedTensorType>(redOp.getInputs()[0].getType()).getRank();

    llvm::SmallBitVector reduced(inputRank, false);
    SmallVector<int64_t> dimensionMap;

    for (auto dim : innerDimensions)
      reduced.set(dim);
    for (int64_t i = 0; i < inputRank; ++i) {
      if (!reduced[i])
        dimensionMap.push_back(i);
    }
    for (auto dim : outerDimensions) {
      if (dim >= dimensionMap.size())
        return rewriter.notifyMatchFailure(op, "out of bounds dimension map");

      mergedDimensions.push_back(dimensionMap[dim]);
    }

    auto newReduceOp = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), TypeRange(op.getType(0)),
        ValueRange(redOp.getInputs()), ValueRange(redOp.getInitValues()),
        mergedDimensions);
    rewriter.inlineRegionBefore(redOp.getBody(), newReduceOp.getBody(),
                                newReduceOp.getBody().end());
    rewriter.replaceOp(op, newReduceOp.getResult(0));
    rewriter.eraseOp(redOp);
    return success();
  }
};

struct ConjReal final : public CheckedOpRewritePattern<chlo::ConjOp, ConjReal> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  bool supportsDynamicShapes() { return true; }

  LogicalResult matchAndRewriteImpl(chlo::ConjOp op,
                                    PatternRewriter &rewriter) const {
    auto input = op.getOperand();

    if (guaranteedPurelyRealResult(input, rewriter)) {
      rewriter.replaceOp(op, input);
      return success();
    }

    return failure();
  }
};

struct ConcatReshapeElementwise final
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     ConcatReshapeElementwise> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp concatOp,
                                    PatternRewriter &rewriter) const {
    if (concatOp.getNumOperands() <= 1)
      return failure();

    auto concatDim = concatOp.getDimension();

    SmallVector<Operation *> concatOpOperands;

    for (auto [i, v] : llvm::enumerate(concatOp.getOperands())) {
      auto reshapeOp = v.getDefiningOp<stablehlo::ReshapeOp>();
      if (!reshapeOp)
        return rewriter.notifyMatchFailure(concatOp, "not a reshape op");

      RankedTensorType reshapeOpInputType = reshapeOp.getOperand().getType();
      RankedTensorType reshapeOpOutputType = reshapeOp.getResult().getType();

      SmallVector<int64_t> insertionDims =
          findReshapeInsertionDims(reshapeOpInputType, reshapeOpOutputType);

      if (insertionDims.size() != 1)
        return rewriter.notifyMatchFailure(
            concatOp, "reshape op has more than one insertion dim");

      if (insertionDims[0] != concatDim)
        return rewriter.notifyMatchFailure(
            concatOp, "concat dim is not same as insertion dim");

      auto vdefOp = reshapeOp.getOperand().getDefiningOp();
      if (!vdefOp)
        return failure();

      if (hasTraitElementwise(vdefOp)) {
        if (concatOpOperands.size() != 0) {
          if (!OperationEquivalence::isEquivalentTo(
                  concatOpOperands[0], vdefOp,
                  OperationEquivalence::ignoreValueEquivalence, nullptr,
                  OperationEquivalence::IgnoreLocations, nullptr))
            return rewriter.notifyMatchFailure(
                concatOp, "elementwise op is not equivalent to first");
        }

        if (!isOnlyUsedInOperation(vdefOp, reshapeOp))
          return rewriter.notifyMatchFailure(
              concatOp, "elementwise op is not only used in reshape op");

        concatOpOperands.push_back(vdefOp);
      } else {
        return rewriter.notifyMatchFailure(concatOp, "not a valid elementwise");
      }
    }

    SmallVector<Value> elementwiseOperands;

    for (int i = 0; i < concatOpOperands[0]->getNumOperands(); i++) {
      SmallVector<Value> newConcatOperands;
      for (auto v : concatOpOperands) {
        auto inputOp = v->getOperand(i);

        auto inputType = cast<RankedTensorType>(inputOp.getType());
        auto inputShape = inputType.getShape();
        SmallVector<int64_t> outputShape;
        for (int j = 0; j < concatDim; j++)
          outputShape.push_back(inputShape[j]);
        outputShape.push_back(1);
        for (int j = concatDim + 1; j < (inputShape.size() + 1); j++)
          outputShape.push_back(inputShape[j - 1]);

        auto newReshapeOp = stablehlo::ReshapeOpCreate(
            rewriter, concatOp.getLoc(), inputOp, outputShape);
        newConcatOperands.push_back(newReshapeOp);
      }
      auto newConcatOp = stablehlo::ConcatenateOp::create(
          rewriter, concatOp.getLoc(), newConcatOperands, concatDim);
      elementwiseOperands.push_back(newConcatOp.getResult());
    }

    auto newElementwiseOp = rewriter.create(
        concatOp.getLoc(), concatOpOperands[0]->getName().getIdentifier(),
        ValueRange(elementwiseOperands), TypeRange(concatOp.getType()),
        concatOpOperands[0]->getAttrs(), {}, {});

    rewriter.replaceOp(concatOp, newElementwiseOp);
    return success();
  }
};

struct TransposeBatchNormTraining final
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeBatchNormTraining> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto batchnormTrainingOp =
        op.getOperand().getDefiningOp<stablehlo::BatchNormTrainingOp>();
    if (!batchnormTrainingOp)
      return failure();

    if (op.getOperand() != batchnormTrainingOp.getOutput())
      return failure();

    auto permutation = op.getPermutation();
    auto newFeatureIndex = permutation[batchnormTrainingOp.getFeatureIndex()];

    rewriter.setInsertionPoint(batchnormTrainingOp);
    auto newTransposeOp = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), batchnormTrainingOp.getOperand(), permutation);
    auto newBatchNormOp = stablehlo::BatchNormTrainingOp::create(
        rewriter, op.getLoc(), newTransposeOp.getResult(),
        batchnormTrainingOp.getScale(), batchnormTrainingOp.getOffset(),
        batchnormTrainingOp.getEpsilon(), newFeatureIndex);

    rewriter.replaceAllUsesWith(op->getResult(0), newBatchNormOp->getResult(0));
    rewriter.replaceAllUsesWith(batchnormTrainingOp->getResult(1),
                                newBatchNormOp->getResult(1));
    rewriter.replaceAllUsesWith(batchnormTrainingOp->getResult(2),
                                newBatchNormOp->getResult(2));

    return success();
  }
};

struct TransposeBatchNormInference final
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeBatchNormInference> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto batchnormInferenceOp =
        op.getOperand().getDefiningOp<stablehlo::BatchNormInferenceOp>();
    if (!batchnormInferenceOp)
      return failure();

    if (op.getOperand() != batchnormInferenceOp.getResult())
      return failure();

    auto permutation = op.getPermutation();
    auto newFeatureIndex = permutation[batchnormInferenceOp.getFeatureIndex()];

    auto newTransposeOp = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), batchnormInferenceOp.getOperand(), permutation);
    rewriter.replaceOpWithNewOp<stablehlo::BatchNormInferenceOp>(
        op, newTransposeOp.getResult(), batchnormInferenceOp.getScale(),
        batchnormInferenceOp.getOffset(), batchnormInferenceOp.getMean(),
        batchnormInferenceOp.getVariance(), batchnormInferenceOp.getEpsilon(),
        newFeatureIndex);
    return success();
  }
};

struct TransposeBatchNormGrad final
    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
                                     TransposeBatchNormGrad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto batchnormGradOp =
        op.getOperand().getDefiningOp<stablehlo::BatchNormGradOp>();
    if (!batchnormGradOp)
      return failure();

    if (op.getOperand() != batchnormGradOp.getGradOperand())
      return failure();

    auto permutation = op.getPermutation();
    auto newFeatureIndex = permutation[batchnormGradOp.getFeatureIndex()];

    rewriter.setInsertionPoint(batchnormGradOp);
    auto transposeOperand = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), batchnormGradOp.getOperand(), permutation);
    auto transposeGradOutput = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), batchnormGradOp.getGradOutput(), permutation);

    auto newBatchNormOp = stablehlo::BatchNormGradOp::create(
        rewriter, op.getLoc(), transposeOperand, batchnormGradOp.getScale(),
        batchnormGradOp.getMean(), batchnormGradOp.getVariance(),
        transposeGradOutput, batchnormGradOp.getEpsilon(), newFeatureIndex);

    rewriter.replaceAllUsesWith(op->getResult(0), newBatchNormOp->getResult(0));
    rewriter.replaceAllUsesWith(batchnormGradOp->getResult(1),
                                newBatchNormOp->getResult(1));
    rewriter.replaceAllUsesWith(batchnormGradOp->getResult(2),
                                newBatchNormOp->getResult(2));
    return success();
  }
};

struct SelectBroadcastInDim final
    : public CheckedOpRewritePattern<stablehlo::SelectOp,
                                     SelectBroadcastInDim> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    auto bcastOp = op.getPred().getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcastOp)
      return failure();

    auto bcastOpOperand = bcastOp.getOperand();
    auto bcastOpOperandType = cast<RankedTensorType>(bcastOpOperand.getType());
    auto bcastOpOperandShape = bcastOpOperandType.getShape();

    if (bcastOpOperandShape.size() == 0) {
      // 0-dim tensor
      rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(
          op, bcastOpOperand, op.getOnTrue(), op.getOnFalse());
      return success();
    } else if (bcastOpOperandShape.size() == 1) {
      if (bcastOpOperandShape[0] != 1)
        return failure();

      // 1-dim tensor of size (1,)
      auto reshapedPred = stablehlo::ReshapeOpCreate(
          rewriter, op.getLoc(), bcastOpOperand, ArrayRef<int64_t>{});
      rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(
          op, reshapedPred, op.getOnTrue(), op.getOnFalse());
      return success();
    }

    return failure();
  }
};

struct TransposeIf final
    : public CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeIf> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto ifOp = op.getOperand().getDefiningOp<stablehlo::IfOp>();
    if (!ifOp)
      return failure();
    if (!llvm::hasSingleElement(op.getOperand().getUses()))
      return failure();

    auto opRes = cast<OpResult>(op.getOperand());
    ssize_t opIdx = opRes.getResultNumber();

    SmallVector<Type> ifResultTypes = llvm::to_vector(ifOp.getResultTypes());
    ifResultTypes[opIdx] = op.getType();

    rewriter.setInsertionPoint(ifOp);
    auto newIfOp = stablehlo::IfOp::create(rewriter, op.getLoc(), ifResultTypes,
                                           ifOp.getPred());

    Operation *trueTerm = ifOp.getTrueBranch().front().getTerminator();
    Operation *falseTerm = ifOp.getFalseBranch().front().getTerminator();

    rewriter.setInsertionPoint(trueTerm);
    auto newTrue = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), trueTerm->getOperands()[opIdx],
        op.getPermutation());
    rewriter.modifyOpInPlace(trueTerm,
                             [&] { trueTerm->setOperand(opIdx, newTrue); });

    rewriter.setInsertionPoint(falseTerm);
    auto newFalse = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), falseTerm->getOperands()[opIdx],
        op.getPermutation());
    rewriter.modifyOpInPlace(falseTerm,
                             [&] { falseTerm->setOperand(opIdx, newFalse); });

    newIfOp.getTrueBranch().takeBody(ifOp.getTrueBranch());
    newIfOp.getFalseBranch().takeBody(ifOp.getFalseBranch());

    for (int i = 0; i < ifOp.getNumResults(); i++) {
      if (i == opIdx) {
        rewriter.replaceAllUsesWith(op.getResult(), newIfOp.getResult(i));
      } else {
        rewriter.replaceAllUsesWith(ifOp.getResult(i), newIfOp.getResult(i));
      }
    }

    rewriter.eraseOp(op);
    rewriter.eraseOp(ifOp);
    return success();
  }
};

struct IfOpLiftCommonOps final
    : public CheckedOpRewritePattern<stablehlo::IfOp, IfOpLiftCommonOps> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::IfOp op,
                                    PatternRewriter &rewriter) const {
    mlir::Region &trueRegion = op.getTrueBranch();
    mlir::Region &falseRegion = op.getFalseBranch();

    SmallVector<std::pair<Operation *, Operation *>> opsToLift;
    llvm::SmallPtrSet<Operation *, 8> falseOpsMatched;

    auto &trueBlock = trueRegion.front();
    auto &falseBlock = falseRegion.front();

    for (auto trueIt = trueBlock.begin();
         !trueIt->hasTrait<mlir::OpTrait::IsTerminator>(); ++trueIt) {
      Operation *trueOp = &*trueIt;

      for (auto falseIt = falseBlock.begin();
           !falseIt->hasTrait<mlir::OpTrait::IsTerminator>(); ++falseIt) {
        Operation *falseOp = &*falseIt;

        if (falseOpsMatched.contains(falseOp))
          continue;

        if (OperationEquivalence::isEquivalentTo(
                trueOp, falseOp, OperationEquivalence::IgnoreLocations)) {
          opsToLift.emplace_back(trueOp, falseOp);
          falseOpsMatched.insert(falseOp);
          break;
        }
      }
    }

    if (opsToLift.empty())
      return rewriter.notifyMatchFailure(op, "no common ops found");

    for (auto [trueOp, falseOp] : opsToLift) {
      Operation *trueOpCopy = trueOp;
      rewriter.modifyOpInPlace(trueOp, [&]() { trueOpCopy->moveBefore(op); });
      rewriter.replaceOp(falseOp, trueOp);
    }

    return success();
  }
};

// used for ops that dont define the Involution trait
template <typename OpTy>
struct InvolutionSimplify
    : public CheckedOpRewritePattern<OpTy, InvolutionSimplify<OpTy>> {
  using CheckedOpRewritePattern<
      OpTy, InvolutionSimplify<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    auto operandOp = op.getOperand().template getDefiningOp<OpTy>();
    if (!operandOp)
      return failure();

    rewriter.replaceOp(op, operandOp.getOperand());
    return success();
  }
};

struct RealConjSimplify final
    : public CheckedOpRewritePattern<stablehlo::RealOp, RealConjSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::RealOp op,
                                    PatternRewriter &rewriter) const {
    auto operandOp = op.getOperand().getDefiningOp<chlo::ConjOp>();
    if (!operandOp)
      return failure();

    rewriter.replaceOpWithNewOp<stablehlo::RealOp>(op, operandOp.getOperand());
    return success();
  }
};

struct RealConvertSimplify final
    : public CheckedOpRewritePattern<stablehlo::RealOp, RealConvertSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::RealOp op,
                                    PatternRewriter &rewriter) const {
    auto operandOp = op.getOperand().getDefiningOp<stablehlo::ConvertOp>();
    if (!operandOp || isa<ComplexType>(cast<RankedTensorType>(
                                           operandOp.getOperand().getType())
                                           .getElementType())) {
      return failure();
    }

    rewriter.replaceOp(op, operandOp.getOperand());
    return success();
  }
};

struct ConjComplexSimplify final
    : public CheckedOpRewritePattern<chlo::ConjOp, ConjComplexSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(chlo::ConjOp op,
                                    PatternRewriter &rewriter) const {
    auto operandOp = op.getOperand().getDefiningOp<stablehlo::ComplexOp>();
    if (!operandOp)
      return failure();

    auto rhs = operandOp.getRhs();
    if (!matchPattern(rhs, m_Constant())) {
      return failure();
    }

    auto negateRhs = stablehlo::NegOp::create(rewriter, op.getLoc(), rhs);
    rewriter.replaceOpWithNewOp<stablehlo::ComplexOp>(op, operandOp.getLhs(),
                                                      negateRhs);
    return success();
  }
};

// elementwise_op(all operands have zero imag) -> do the op in real domain and
// convert to complex
struct ElementwiseComplexSimplify final
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          ElementwiseComplexSimplify> {
  using CheckedOpTraitRewritePattern::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (isa<stablehlo::ConvertOp, stablehlo::ComplexOp>(op)) {
      return failure();
    }

    if (op->getNumResults() != 1) {
      return failure();
    }

    auto origRes = op->getResult(0);
    auto resElemType =
        cast<RankedTensorType>(origRes.getType()).getElementType();
    auto complexEltype = dyn_cast<ComplexType>(resElemType);
    if (!complexEltype) {
      return failure();
    }
    auto realEltype = complexEltype.getElementType();

    auto isComplexEltype = [](auto val) {
      auto elemType = cast<RankedTensorType>(val.getType()).getElementType();
      return isa<ComplexType>(elemType);
    };

    if (!llvm::all_of(op->getOperands(), [&](auto val) {
          return guaranteedPurelyRealResult(val, rewriter) &&
                 isValueOnlyUsedInOperation(val, op);
        })) {
      return failure();
    }

    // extract the real operands
    auto newOperands =
        llvm::map_to_vector(op->getOperands(), [&](auto operand) -> Value {
          if (!isComplexEltype(operand)) {
            return operand;
          }
          return stablehlo::RealOp::create(rewriter, operand.getLoc(), operand)
              .getResult();
        });

    auto nex = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(), ValueRange(newOperands),
        TypeRange(RankedTensorType::get(
            cast<ShapedType>(op->getResult(0).getType()).getShape(),
            realEltype)),
        op->getAttrs(), {}, {});

    rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(
        op, op->getResultTypes()[0], nex->getResult(0));
    return success();
  }
};

// workaround for https://github.com/openxla/xla/issues/27446
struct SplitConvolutionIntoReverseConvolution final
    : public CheckedOpRewritePattern<stablehlo::ConvolutionOp,
                                     SplitConvolutionIntoReverseConvolution> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvolutionOp op,
                                    PatternRewriter &rewriter) const {
    auto windowReversal = op.getWindowReversal();
    if (!windowReversal)
      return rewriter.notifyMatchFailure(op, "no window reversal");

    if (!llvm::any_of(windowReversal.value(), [](bool b) { return b; }))
      return rewriter.notifyMatchFailure(op, "all window reversals are false");

    auto convDims = op.getDimensionNumbers();

    SmallVector<int64_t> reversalDims;
    for (auto [dim, reversal] : llvm::zip(convDims.getKernelSpatialDimensions(),
                                          windowReversal.value())) {
      if (reversal) {
        reversalDims.push_back(dim);
      }
    }

    auto reverseOp = stablehlo::ReverseOp::create(rewriter, op.getLoc(),
                                                  op.getRhs(), reversalDims);

    rewriter.replaceOpWithNewOp<stablehlo::ConvolutionOp>(
        op, op.getType(), op.getLhs(), reverseOp.getResult(),
        op.getWindowStridesAttr(), op.getPaddingAttr(), op.getLhsDilationAttr(),
        op.getRhsDilationAttr(), nullptr, op.getDimensionNumbersAttr(),
        op.getFeatureGroupCountAttr(), op.getBatchGroupCountAttr(),
        op.getPrecisionConfigAttr());
    return success();
  }
};

// Check whether a user of a scatter result can be simplified by one of the
// ScatterBinaryOpSimplify patterns given the scatter's isAllZeros/isAllOnes.
static bool canScatterUserBeSimplified(Operation *user, Value scatterResult,
                                       bool isAllZeros, bool isAllOnes) {
  if (isAllZeros) {
    if (isa<stablehlo::AddOp, stablehlo::SubtractOp, stablehlo::MulOp>(user))
      return true;
  }
  if (isAllOnes) {
    if (isa<stablehlo::MulOp>(user))
      return true;
    if (isa<stablehlo::DivOp>(user)) {
      // DivOp only simplifies when the scatter is on the rhs (x / scatter(1))
      bool scatterIsLhs = (user->getOperand(0) == scatterResult);
      return !scatterIsLhs;
    }
  }
  return false;
}

template <typename OpTy, typename Child>
struct ScatterBinaryOpSimplifyBase
    : public CheckedOpRewritePattern<OpTy, Child> {
  using CheckedOpRewritePattern<OpTy, Child>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    stablehlo::ScatterOp scatterOp;
    mlir::Value otherValue;

    auto lhsScatterOp = lhs.template getDefiningOp<stablehlo::ScatterOp>();
    auto rhsScatterOp = rhs.template getDefiningOp<stablehlo::ScatterOp>();
    bool lhsIsScatter;
    if (!lhsScatterOp && !rhsScatterOp) {
      return failure();
    } else {
      if (lhsScatterOp) {
        lhsIsScatter = true;
        scatterOp = lhsScatterOp;
        otherValue = rhs;
      } else {
        lhsIsScatter = false;
        scatterOp = rhsScatterOp;
        otherValue = lhs;
      }
    }

    if (!scatterOp.getUniqueIndices()) {
      return failure();
    }

    bool isAllZeros = false, isAllOnes = false;

    SplatElementsAttr constSetIndexValue = nullptr;
    auto status = detectConstantSetindexScatterOp(
        scatterOp, /*allowedMultipleUses*/ true,
        [&isAllZeros, &isAllOnes](mlir::Value input) {
          Attribute attr;
          if (matchPattern(input, m_Constant(&attr))) {
            isAllZeros = matchPattern(attr, m_AnyZeroFloat()) ||
                         matchPattern(attr, m_Zero());
            if (!isAllZeros)
              isAllOnes = matchPattern(attr, m_OneFloat()) ||
                          matchPattern(attr, m_One());
          }
          return isAllZeros || isAllOnes;
        },
        constSetIndexValue);
    if (!status.ok()) {
      return rewriter.notifyMatchFailure(op, status.message());
    }

    // If the scatter has multiple uses, verify that all other uses can be
    // simplified by a later iteration of the same pass.
    if (!scatterOp.getResult(0).hasOneUse()) {
      for (auto *user : scatterOp.getResult(0).getUsers()) {
        if (user == op.getOperation())
          continue;
        if (!canScatterUserBeSimplified(user, scatterOp.getResult(0),
                                        isAllZeros, isAllOnes)) {
          return rewriter.notifyMatchFailure(
              op, "scatter has uses that cannot be simplified by a later pass");
        }
      }
    }

    SmallVector<int64_t> sliceSizes = computeGatherSliceSizes(scatterOp);

    return ((Child *)this)
        ->rewriteScatterElementwise(op, rewriter, scatterOp, otherValue,
                                    isAllZeros, isAllOnes, lhsIsScatter,
                                    constSetIndexValue, sliceSizes);
  }

  // if !fuseIntoScatter
  //     %g = gather(%other_value, %scatter_indices)
  //     %tmp = elementwise_op(%g, %updates) if !lhsIsScatter else
  //            elementwise_op(%updates, %g)
  //     %new_scatter_op = scatter(%input, %scatter_indices, %tmp)
  // else
  //     %new_scatter_op = scatter(%input, %scatter_indices, %updates) {
  //         elementwise_op(%arg1, %arg2) // based on lhsIsScatter
  //     }
  template <auto CreateOpFn>
  void gatherElementwiseSetIndex(OpTy op, PatternRewriter &rewriter,
                                 stablehlo::ScatterOp scatterOp,
                                 Value otherValue, Value scatterInput,
                                 bool lhsIsScatter,
                                 SplatElementsAttr constSetIndexValue,
                                 SmallVectorImpl<int64_t> &sliceSizes,
                                 bool fuseIntoScatter) const {
    Value updateVal;
    if (constSetIndexValue) {
      updateVal = stablehlo::ConstantOp::create(
          rewriter, scatterOp.getLoc(),
          constSetIndexValue.resizeSplat(
              cast<ShapedType>(scatterOp.getUpdates()[0].getType())));
    } else {
      updateVal = scatterOp.getUpdates()[0];
    }

    if (fuseIntoScatter) {
      assert(scatterInput == otherValue);

      auto newScatterOp = stablehlo::ScatterOp::create(
          rewriter, op.getLoc(), scatterOp.getResultTypes(),
          ValueRange(scatterInput), scatterOp.getScatterIndices(),
          ValueRange(updateVal), scatterOp.getScatterDimensionNumbersAttr(),
          scatterOp.getIndicesAreSortedAttr(),
          scatterOp.getUniqueIndicesAttr());

      auto &updateRegion = newScatterOp.getUpdateComputation();
      auto *block = rewriter.createBlock(&updateRegion);
      auto elemType = cast<RankedTensorType>(scatterOp.getResultTypes()[0])
                          .getElementType();
      auto argType = RankedTensorType::get({}, elemType);
      block->addArgument(argType, op.getLoc());
      block->addArgument(argType, op.getLoc());
      rewriter.setInsertionPointToStart(block);

      stablehlo::ReturnOp::create(
          rewriter, op.getLoc(),
          CreateOpFn(rewriter, op.getLoc(), block->getArgument(lhsIsScatter),
                     block->getArgument(!lhsIsScatter), std::nullopt));

      rewriter.replaceOp(op, newScatterOp);

      return;
    }

    auto gatheredValues =
        stablehlo::GatherOp::create(
            rewriter, scatterOp.getLoc(), otherValue,
            scatterOp.getScatterIndices(),
            getGatherDims(rewriter.getContext(),
                          scatterOp.getScatterDimensionNumbers()),
            rewriter.getDenseI64ArrayAttr(sliceSizes),
            scatterOp.getIndicesAreSortedAttr())
            .getResult();

    Value newUpdates = CreateOpFn(
        rewriter, scatterOp.getLoc(), lhsIsScatter ? updateVal : gatheredValues,
        lhsIsScatter ? gatheredValues : updateVal, std::nullopt);

    auto newScatterOp = stablehlo::ScatterOp::create(
        rewriter, op.getLoc(), scatterOp.getResultTypes(),
        ValueRange(scatterInput), scatterOp.getScatterIndices(),
        ValueRange(newUpdates), scatterOp.getScatterDimensionNumbersAttr(),
        scatterOp.getIndicesAreSortedAttr(), scatterOp.getUniqueIndicesAttr());

    auto &updateRegion = newScatterOp.getUpdateComputation();
    auto *block = rewriter.createBlock(&updateRegion);
    auto elemType =
        cast<RankedTensorType>(scatterOp.getResultTypes()[0]).getElementType();
    auto argType = RankedTensorType::get({}, elemType);
    block->addArgument(argType, op.getLoc());
    block->addArgument(argType, op.getLoc());
    rewriter.setInsertionPointToStart(block);
    stablehlo::ReturnOp::create(rewriter, op.getLoc(), block->getArgument(1));

    rewriter.replaceOp(op, newScatterOp);
  }
};

struct ScatterMultiplySimplify final
    : public ScatterBinaryOpSimplifyBase<stablehlo::MulOp,
                                         ScatterMultiplySimplify> {
  using ScatterBinaryOpSimplifyBase<
      stablehlo::MulOp, ScatterMultiplySimplify>::ScatterBinaryOpSimplifyBase;

  LogicalResult
  rewriteScatterElementwise(stablehlo::MulOp op, PatternRewriter &rewriter,
                            stablehlo::ScatterOp scatterOp, Value otherValue,
                            bool isAllZeros, bool isAllOnes, bool lhsIsScatter,
                            SplatElementsAttr constSetIndexValue,
                            SmallVectorImpl<int64_t> &sliceSizes) const {
    if (isAllZeros) { // non scattered values before zeros
      gatherElementwiseSetIndex<stablehlo::MulOpCreate>(
          op, rewriter, scatterOp, otherValue, scatterOp.getInputs()[0],
          lhsIsScatter, constSetIndexValue, sliceSizes, false);
      return success();
    }

    if (isAllOnes) { // non-scattered values stay as is
      auto newScatterOp = stablehlo::ScatterOp::create(
          rewriter, op.getLoc(), scatterOp.getResultTypes(),
          ValueRange(otherValue), scatterOp.getScatterIndices(),
          scatterOp.getUpdates(), scatterOp.getScatterDimensionNumbersAttr(),
          scatterOp.getIndicesAreSortedAttr(),
          scatterOp.getUniqueIndicesAttr());

      auto &updateRegion = newScatterOp.getUpdateComputation();
      auto *block = rewriter.createBlock(&updateRegion);
      auto elemType = cast<RankedTensorType>(scatterOp.getResultTypes()[0])
                          .getElementType();
      auto argType = RankedTensorType::get({}, elemType);
      block->addArgument(argType, op.getLoc());
      block->addArgument(argType, op.getLoc());
      rewriter.setInsertionPointToStart(block);

      Value mulRhs;
      if (constSetIndexValue) {
        mulRhs = stablehlo::ConstantOp::create(
            rewriter, op.getLoc(),
            constSetIndexValue.resizeSplat(
                RankedTensorType::get({}, elemType)));
      } else {
        mulRhs = block->getArgument(1);
      }
      auto mulOp = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                            block->getArgument(0), mulRhs);
      stablehlo::ReturnOp::create(rewriter, op.getLoc(), mulOp.getResult());

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

    return failure();
  }
};

struct ScatterDivSimplify final
    : public ScatterBinaryOpSimplifyBase<stablehlo::DivOp, ScatterDivSimplify> {
  using ScatterBinaryOpSimplifyBase<
      stablehlo::DivOp, ScatterDivSimplify>::ScatterBinaryOpSimplifyBase;

  LogicalResult
  rewriteScatterElementwise(stablehlo::DivOp op, PatternRewriter &rewriter,
                            stablehlo::ScatterOp scatterOp, Value otherValue,
                            bool isAllZeros, bool isAllOnes, bool lhsIsScatter,
                            SplatElementsAttr constSetIndexValue,
                            SmallVectorImpl<int64_t> &sliceSizes) const {
    if (isAllOnes && !lhsIsScatter) {
      // x / 1 -> setindex into x
      gatherElementwiseSetIndex<stablehlo::DivOpCreate>(
          op, rewriter, scatterOp, otherValue, otherValue, lhsIsScatter,
          constSetIndexValue, sliceSizes, true);
      return success();
    }

    return failure();
  }
};

struct ScatterAddSimplify final
    : public ScatterBinaryOpSimplifyBase<stablehlo::AddOp, ScatterAddSimplify> {
  using ScatterBinaryOpSimplifyBase<
      stablehlo::AddOp, ScatterAddSimplify>::ScatterBinaryOpSimplifyBase;

  LogicalResult
  rewriteScatterElementwise(stablehlo::AddOp op, PatternRewriter &rewriter,
                            stablehlo::ScatterOp scatterOp, Value otherValue,
                            bool isAllZeros, bool isAllOnes, bool lhsIsScatter,
                            SplatElementsAttr constSetIndexValue,
                            SmallVectorImpl<int64_t> &sliceSizes) const {
    if (isAllZeros) {
      gatherElementwiseSetIndex<stablehlo::AddOpCreate>(
          op, rewriter, scatterOp, otherValue, otherValue, lhsIsScatter,
          constSetIndexValue, sliceSizes, true);
      return success();
    }

    return failure();
  }
};

struct ScatterSubSimplify final
    : public ScatterBinaryOpSimplifyBase<stablehlo::SubtractOp,
                                         ScatterSubSimplify> {
  using ScatterBinaryOpSimplifyBase<
      stablehlo::SubtractOp, ScatterSubSimplify>::ScatterBinaryOpSimplifyBase;

  LogicalResult
  rewriteScatterElementwise(stablehlo::SubtractOp op, PatternRewriter &rewriter,
                            stablehlo::ScatterOp scatterOp, Value otherValue,
                            bool isAllZeros, bool isAllOnes, bool lhsIsScatter,
                            SplatElementsAttr constSetIndexValue,
                            SmallVectorImpl<int64_t> &sliceSizes) const {
    if (isAllZeros) {
      if (lhsIsScatter) {
        otherValue =
            stablehlo::NegOp::create(rewriter, op.getLoc(), otherValue);
        gatherElementwiseSetIndex<stablehlo::AddOpCreate>(
            op, rewriter, scatterOp, otherValue, otherValue, lhsIsScatter,
            constSetIndexValue, sliceSizes, true);
      } else {
        gatherElementwiseSetIndex<stablehlo::SubtractOpCreate>(
            op, rewriter, scatterOp, otherValue, otherValue, lhsIsScatter,
            constSetIndexValue, sliceSizes, true);
      }
      return success();
    }

    return failure();
  }
};

// Fuse scatter(scatter(input, indices, updates1) { body1 }, indices, updates2)
// { body2 } into a single scatter when both use the same indices and scatter
// dimension numbers.
struct ScatterOfScatterSimplify final
    : public CheckedOpRewritePattern<stablehlo::ScatterOp,
                                     ScatterOfScatterSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp outerScatter,
                                    PatternRewriter &rewriter) const {
    // The outer scatter must have a single input
    if (outerScatter.getInputs().size() != 1)
      return failure();

    // The input to the outer scatter must be the result of another scatter
    auto innerScatter =
        outerScatter.getInputs()[0].getDefiningOp<stablehlo::ScatterOp>();
    if (!innerScatter)
      return failure();

    // Inner scatter must have a single use (the outer scatter)
    if (!innerScatter.getResult(0).hasOneUse())
      return failure();

    if (innerScatter.getInputs().size() != 1)
      return failure();

    // Both must have unique_indices
    if (!outerScatter.getUniqueIndices() || !innerScatter.getUniqueIndices())
      return failure();

    // Same scatter indices
    if (outerScatter.getScatterIndices() != innerScatter.getScatterIndices())
      return failure();

    // Same scatter dimension numbers
    if (outerScatter.getScatterDimensionNumbers() !=
        innerScatter.getScatterDimensionNumbers())
      return failure();

    using ScatterOpKind = stablehlo::ScatterOpKind;
    auto innerInput = innerScatter.getInputs()[0];
    auto loc = outerScatter.getLoc();
    auto resultType = outerScatter.getResultTypes()[0];
    auto elemType = cast<RankedTensorType>(resultType).getElementType();
    auto scalarType = RankedTensorType::get({}, elemType);

    auto innerCheck = stablehlo::CheckCommonScatterOp(innerScatter);
    auto outerCheck = stablehlo::CheckCommonScatterOp(outerScatter);

    // When the outer body completely ignores the current value (Setindex or
    // ConstantSetindex), the inner scatter is irrelevant — just replace the
    // input with the inner's input.
    if (outerCheck.kind == ScatterOpKind::Setindex ||
        outerCheck.kind == ScatterOpKind::ConstantSetindex) {
      auto newScatter = stablehlo::ScatterOp::create(
          rewriter, loc, outerScatter.getResultTypes(), ValueRange(innerInput),
          outerScatter.getScatterIndices(), outerScatter.getUpdates(),
          outerScatter.getScatterDimensionNumbersAttr(),
          outerScatter.getIndicesAreSortedAttr(),
          outerScatter.getUniqueIndicesAttr());

      // Clone the outer update computation
      rewriter.cloneRegionBefore(outerScatter.getUpdateComputation(),
                                 newScatter.getUpdateComputation(),
                                 newScatter.getUpdateComputation().end());

      rewriter.replaceOp(outerScatter, newScatter);
      return success();
    }

    // When the inner scatter is ConstantSetindex(c1), the value at scattered
    // positions after the inner is always c1. So the outer body sees c1 as
    // arg0 (the current value). Clone the outer body and replace arg0 with
    // the constant — later passes will constant-fold as needed.
    if (innerCheck.kind == ScatterOpKind::ConstantSetindex) {
      SplatElementsAttr c1 = innerCheck.constant;

      auto newScatter = stablehlo::ScatterOp::create(
          rewriter, loc, outerScatter.getResultTypes(), ValueRange(innerInput),
          outerScatter.getScatterIndices(), outerScatter.getUpdates(),
          outerScatter.getScatterDimensionNumbersAttr(),
          outerScatter.getIndicesAreSortedAttr(),
          outerScatter.getUniqueIndicesAttr());

      rewriter.cloneRegionBefore(outerScatter.getUpdateComputation(),
                                 newScatter.getUpdateComputation(),
                                 newScatter.getUpdateComputation().end());

      // Replace uses of the first block argument (the current value at the
      // scatter position, which is always c1) with the constant.
      auto &block = newScatter.getUpdateComputation().front();
      rewriter.setInsertionPointToStart(&block);
      auto c1Const = stablehlo::ConstantOp::create(rewriter, loc,
                                                   c1.resizeSplat(scalarType));
      block.getArgument(0).replaceAllUsesWith(c1Const);

      rewriter.replaceOp(outerScatter, newScatter);
      return success();
    }

    return failure();
  }
};

struct GatherConstProp final
    : public CheckedOpRewritePattern<stablehlo::GatherOp, GatherConstProp> {
  using CheckedOpRewritePattern<stablehlo::GatherOp,
                                GatherConstProp>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::GatherOp op,
                                    PatternRewriter &rewriter) const {
    DenseElementsAttr operandAttr;
    if (!matchPattern(op.getOperand(), m_Constant(&operandAttr)))
      return rewriter.notifyMatchFailure(op,
                                         "GatherOp with non-constant input");

    if (operandAttr.isSplat()) {
      // In this case the indices don't matter and we can construct a new
      // splatted result
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, op.getType(), operandAttr.resizeSplat(op.getType()));
      return success();
    }

    DenseElementsAttr startIndicesAttr;
    if (!matchPattern(op.getStartIndices(), m_Constant(&startIndicesAttr))) {
      return rewriter.notifyMatchFailure(
          op, "GatherOp with non-constant start indices and unsplatted input");
    }

    stablehlo::Tensor operandTensor = stablehlo::constantOp(operandAttr);
    stablehlo::Tensor startIndicesTensor =
        stablehlo::constantOp(startIndicesAttr);
    auto gatherDims = op.getDimensionNumbers();

    auto sliceSizes = op.getSliceSizes();
    auto elementType = rewriter.getIntegerType(64);
    auto attrType =
        RankedTensorType::get({(int64_t)sliceSizes.size()}, elementType);
    auto sliceSizesAttr = DenseElementsAttr::get(attrType, sliceSizes);

    auto result = stablehlo::gatherOp(
        operandTensor, startIndicesTensor,
        stablehlo::Axes(gatherDims.getOffsetDims()),
        stablehlo::Axes(gatherDims.getCollapsedSliceDims()),
        stablehlo::Axes(gatherDims.getOperandBatchingDims()),
        stablehlo::Axes(gatherDims.getStartIndicesBatchingDims()),
        stablehlo::Axes(gatherDims.getStartIndexMap()),
        stablehlo::Axis(gatherDims.getIndexVectorDim()),
        stablehlo::makeSizes(stablehlo::constantOp(sliceSizesAttr)),
        op.getIndicesAreSorted(), op.getType());
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(),
                                                       fromTensor(result));
    return success();
  }
};

struct UnaryElementwiseScatterSimplify final
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          UnaryElementwiseScatterSimplify> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise,
      UnaryElementwiseScatterSimplify>::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() != 1)
      return rewriter.notifyMatchFailure(op, "not a unary elementwise op");

    auto input = op->getOperand(0);
    auto scatterOp = input.getDefiningOp<stablehlo::ScatterOp>();
    if (!scatterOp)
      return rewriter.notifyMatchFailure(op, "not a scatter op");

    auto status = detectConstantSetindexScatterOp(
        scatterOp, false,
        [](mlir::Value input) { return matchPattern(input, m_Constant()); });
    if (!status.ok()) {
      return rewriter.notifyMatchFailure(op, status.message());
    }

    auto scatterInput = scatterOp.getInputs()[0];

    auto elemType =
        cast<RankedTensorType>(op->getResult(0).getType()).getElementType();

    // should get constant propagated
    auto scatterInputElem = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(), ValueRange(scatterInput),
        TypeRange{RankedTensorType::get(
            cast<RankedTensorType>(scatterInput.getType()).getShape(),
            elemType)},
        op->getAttrs(), {}, {});

    auto scatterUpdatesElem = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(),
        ValueRange(scatterOp.getUpdates()),
        TypeRange{RankedTensorType::get(
            cast<RankedTensorType>(scatterOp.getUpdates()[0].getType())
                .getShape(),
            elemType)},
        op->getAttrs(), {}, {});

    auto resultType = RankedTensorType::get(
        cast<RankedTensorType>(scatterOp.getResultTypes()[0]).getShape(),
        elemType);

    auto newScatterOp = stablehlo::ScatterOp::create(
        rewriter, op->getLoc(), TypeRange(resultType),
        ValueRange(scatterInputElem->getResult(0)),
        scatterOp.getScatterIndices(),
        ValueRange(scatterUpdatesElem->getResult(0)),
        scatterOp.getScatterDimensionNumbersAttr(),
        scatterOp.getIndicesAreSortedAttr(), scatterOp.getUniqueIndicesAttr());

    auto &updateRegion = newScatterOp.getUpdateComputation();
    auto *block = rewriter.createBlock(&updateRegion);
    auto argType = RankedTensorType::get({}, elemType);
    block->addArgument(argType, op->getLoc());
    block->addArgument(argType, op->getLoc());
    rewriter.setInsertionPointToStart(block);
    stablehlo::ReturnOp::create(rewriter, op->getLoc(), block->getArgument(1));

    rewriter.replaceOp(op, newScatterOp->getResult(0));
    return success();
  }
};

struct GatherElementwise
    : public CheckedOpRewritePattern<stablehlo::GatherOp, GatherElementwise> {
  using CheckedOpRewritePattern<stablehlo::GatherOp,
                                GatherElementwise>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::GatherOp op,
                                    PatternRewriter &rewriter) const {
    auto gatherInput = op.getOperand();
    auto defOp = gatherInput.getDefiningOp();
    if (!defOp || !hasTraitElementwise(defOp))
      return rewriter.notifyMatchFailure(op,
                                         "GatherOp with non-elementwise input");

    if (!isOnlyUsedInOperation(defOp, op))
      return failure();

    if (isa<stablehlo::RealOp, stablehlo::ImagOp, stablehlo::ComplexOp>(defOp))
      return failure();

    int64_t outElemCount =
        cast<RankedTensorType>(op.getType()).getNumElements();
    int64_t inElemCount =
        cast<RankedTensorType>(gatherInput.getType()).getNumElements();

    if (outElemCount >= inElemCount)
      return rewriter.notifyMatchFailure(
          op, "GatherOp has more output elements than input elements");

    SmallVector<Value> newElementwiseInputs;
    for (auto input : defOp->getOperands()) {
      auto neeGatherOp = stablehlo::GatherOp::create(
          rewriter, op.getLoc(), input, op.getStartIndices(),
          op.getDimensionNumbersAttr(), op.getSliceSizesAttr(),
          op.getIndicesAreSortedAttr());
      newElementwiseInputs.push_back(neeGatherOp->getResult(0));
    }

    auto newElemOp = rewriter.create(
        op.getLoc(), defOp->getName().getIdentifier(),
        ValueRange(newElementwiseInputs), TypeRange{op.getResult().getType()},
        defOp->getAttrs(), {}, {});
    rewriter.replaceOp(op, newElemOp->getResult(0));
    return success();
  }
};

struct ElementwiseGather
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          ElementwiseGather> {
  using CheckedOpTraitRewritePattern::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumResults() != 1) {
      return failure();
    }
    if (isa<stablehlo::ComplexOp>(op)) {
      return failure();
    }

    stablehlo::GatherOp gatherOp;
    SmallVector<Value> gatherOperands;

    for (auto [i, operand] : llvm::enumerate(op->getOperands())) {
      auto defOp = operand.getDefiningOp<stablehlo::GatherOp>();
      if (!defOp || !isOnlyUsedInOperation(defOp, op)) {
        return failure();
      }

      auto gatherOutType = cast<RankedTensorType>(defOp.getType());
      auto gatherInType = cast<RankedTensorType>(defOp.getOperand().getType());
      if (!gatherOutType.hasStaticShape() || !gatherInType.hasStaticShape()) {
        return failure();
      }

      if (i == 0) {
        gatherOp = defOp;

        int64_t gatherOutElemCount = gatherOutType.getNumElements();
        int64_t gatherInElemCount = gatherInType.getNumElements();
        if (gatherOutElemCount <= gatherInElemCount) {
          return failure();
        }
      } else {
        // operands can be different, start_indices must be the same
        if (!OperationEquivalence::isEquivalentTo(
                gatherOp, defOp, OperationEquivalence::ignoreValueEquivalence,
                nullptr, OperationEquivalence::IgnoreLocations, nullptr) ||
            gatherOp.getStartIndices() != defOp.getStartIndices()) {
          return failure();
        }
      }

      gatherOperands.push_back(defOp.getOperand());
    }

    if (gatherOperands.empty()) {
      return failure();
    }

    auto newElemOp = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(), ValueRange(gatherOperands),
        TypeRange{RankedTensorType::get(
            cast<RankedTensorType>(gatherOperands[0].getType()).getShape(),
            cast<TensorType>(op->getResult(0).getType()).getElementType())},
        op->getAttrs(), {}, {});
    auto newGatherOp = stablehlo::GatherOp::create(
        rewriter, op->getLoc(), newElemOp->getResult(0),
        gatherOp.getStartIndices(), gatherOp.getDimensionNumbersAttr(),
        gatherOp.getSliceSizesAttr(), gatherOp.getIndicesAreSortedAttr());
    rewriter.replaceOp(op, newGatherOp);
    return success();
  }
};

// (a1 * A) × (b1 * B) → (a1 * b1) * (A × B)
struct FactorScalarsInDotGeneral final
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                     FactorScalarsInDotGeneral> {
  using CheckedOpRewritePattern<
      stablehlo::DotGeneralOp,
      FactorScalarsInDotGeneral>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    // first check if atleast one of lhs or rhs has a
    // mutliplication with a scalar
    if (!canRewriteOperation(op, lhs, rhs))
      return failure();

    Value lhsScalar, lhsZ;
    Value rhsScalar, rhsZ;

    stablehlo::extractMultiplicationFactor(lhs, lhsScalar, lhsZ, op, rewriter);
    stablehlo::extractMultiplicationFactor(rhs, rhsScalar, rhsZ, op, rewriter);

    auto newDot = stablehlo::DotGeneralOp::create(
        rewriter, op.getLoc(), op.getType(), lhsZ, rhsZ,
        op.getDotDimensionNumbers(), op.getPrecisionConfigAttr(),
        op.getAlgorithmAttr());

    Value combinedScalar;
    if (lhsScalar && rhsScalar) {
      combinedScalar =
          stablehlo::MulOp::create(rewriter, op.getLoc(), lhsScalar, rhsScalar);
    } else {
      combinedScalar = lhsScalar ? lhsScalar : rhsScalar;
    }

    auto mulResult =
        stablehlo::MulOpCreate(rewriter, op.getLoc(), combinedScalar, newDot);
    rewriter.replaceOp(op, mulResult);
    return success();
  }

private:
  bool canRewriteOperation(stablehlo::DotGeneralOp op, Value lhs,
                           Value rhs) const {
    auto lhsMulOp = lhs.getDefiningOp<stablehlo::MulOp>();
    auto rhsMulOp = rhs.getDefiningOp<stablehlo::MulOp>();
    if (!lhsMulOp && !rhsMulOp)
      return false; // nothing to do

    if ((lhsMulOp && !isOnlyUsedInOperation(lhsMulOp, op)) ||
        (rhsMulOp && !isOnlyUsedInOperation(rhsMulOp, op)))
      return false; // better to not do anything

    bool lhsHasScalar = false;
    if (lhsMulOp) {
      lhsHasScalar =
          isScalarValue(lhsMulOp.getLhs()) || isScalarValue(lhsMulOp.getRhs());
    }

    bool rhsHasScalar = false;
    if (rhsMulOp) {
      rhsHasScalar =
          isScalarValue(rhsMulOp.getLhs()) || isScalarValue(rhsMulOp.getRhs());
    }

    return lhsHasScalar || rhsHasScalar;
  }
};

struct ChainedMultiplyToPower final
    : public CheckedOpRewritePattern<stablehlo::MulOp, ChainedMultiplyToPower> {
  using CheckedOpRewritePattern<
      stablehlo::MulOp, ChainedMultiplyToPower>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsDefOp = lhs.getDefiningOp<stablehlo::MulOp>();
    auto rhsDefOp = rhs.getDefiningOp<stablehlo::MulOp>();

    if (lhsDefOp && isOnlyUsedInOperation(lhsDefOp, op)) {
      auto lhsLhs = lhsDefOp.getLhs();
      auto lhsRhs = lhsDefOp.getRhs();
      if (lhsLhs == lhsRhs) {
        if (rhsDefOp && isOnlyUsedInOperation(rhsDefOp, op)) {
          auto rhsLhs = rhsDefOp.getLhs();
          auto rhsRhs = rhsDefOp.getRhs();
          if (rhsLhs == rhsRhs &&
              rhsLhs == lhsRhs) { // (mul (mul x x) (mul x x))
            auto powType = RankedTensorType::get(
                cast<RankedTensorType>(lhs.getType()).getShape(),
                lhs.getType().getElementType());
            auto powVal = stablehlo::ConstantOp::create(
                rewriter, op.getLoc(), powType,
                cast<ElementsAttr>(makeAttr(powType, 4)));
            rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, rhsLhs, powVal);
            return success();
          }
        }

        if (lhsLhs == rhs) { // (mul (mul x x) x)
          auto powType = RankedTensorType::get(
              cast<RankedTensorType>(rhs.getType()).getShape(),
              rhs.getType().getElementType());
          auto powVal = stablehlo::ConstantOp::create(
              rewriter, op.getLoc(), powType,
              cast<ElementsAttr>(makeAttr(powType, 3)));
          rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, rhs, powVal);
          return success();
        }
      }
    }

    if (rhsDefOp && isOnlyUsedInOperation(rhsDefOp, op)) { // mul x (mul x x)
      auto rhsLhs = rhsDefOp.getLhs();
      auto rhsRhs = rhsDefOp.getRhs();
      if (rhsLhs == rhsRhs) {
        if (rhsLhs == lhs) {
          auto powType = RankedTensorType::get(
              cast<RankedTensorType>(lhs.getType()).getShape(),
              lhs.getType().getElementType());
          auto powVal = stablehlo::ConstantOp::create(
              rewriter, op.getLoc(), powType,
              cast<ElementsAttr>(makeAttr(powType, 3)));
          rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, lhs, powVal);
          return success();
        }
      }
    }

    return failure();
  }
};

struct PowerMultiplyToPower final
    : public CheckedOpRewritePattern<stablehlo::MulOp, PowerMultiplyToPower> {
  using CheckedOpRewritePattern<stablehlo::MulOp,
                                PowerMultiplyToPower>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsDefOp = lhs.getDefiningOp<stablehlo::PowOp>();
    auto rhsDefOp = rhs.getDefiningOp<stablehlo::PowOp>();

    auto rType = cast<RankedTensorType>(op.getType());

    if (lhsDefOp && isOnlyUsedInOperation(lhsDefOp, op)) {
      auto lhsLhs = lhsDefOp.getLhs();

      if (rhsDefOp && isOnlyUsedInOperation(rhsDefOp, op)) {
        auto rhsLhs = rhsDefOp.getLhs();

        if (lhsLhs ==
            rhsLhs) { // (mul (pow a b) (pow a c)) => (pow a (add b c))
          auto newPowVal = stablehlo::AddOp::create(
              rewriter, op.getLoc(), lhsDefOp.getRhs(), rhsDefOp.getRhs());
          rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, lhsLhs, newPowVal);
          return success();
        }
      }

      auto rhsMulDefOp = rhs.getDefiningOp<stablehlo::MulOp>();
      if (rhsMulDefOp && isOnlyUsedInOperation(rhsMulDefOp, op)) {
        auto rhsMulLhs = rhsMulDefOp.getLhs();
        auto rhsMulRhs = rhsMulDefOp.getRhs();

        if (rhsMulLhs == rhsMulRhs &&
            rhsMulLhs ==
                lhsLhs) { // (mul (pow a b) (mul a a)) => (pow a (add b 2))
          auto newPowVal = stablehlo::AddOp::create(
              rewriter, op.getLoc(), lhsDefOp.getRhs(),
              stablehlo::ConstantOp::create(
                  rewriter, op.getLoc(), rType,
                  cast<ElementsAttr>(makeAttr(rType, 2))));
          rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, lhsLhs, newPowVal);
          return success();
        }
      }

      if (lhsLhs == rhs) { // (mul (pow x y) x) => (pow x (add y 1))
        auto newPowVal = stablehlo::AddOp::create(
            rewriter, op.getLoc(), lhsDefOp.getRhs(),
            stablehlo::ConstantOp::create(
                rewriter, op.getLoc(), rType,
                cast<ElementsAttr>(makeAttr(rType, 1))));
        rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, lhsLhs, newPowVal);
        return success();
      }
    }

    if (rhsDefOp && isOnlyUsedInOperation(rhsDefOp, op)) {
      auto rhsLhs = rhsDefOp.getLhs();

      if (rhsLhs == lhs) { // (mul x (pow x y)) => (pow x (add y 1))
        auto newPowVal = stablehlo::AddOp::create(
            rewriter, op.getLoc(), rhsDefOp.getRhs(),
            stablehlo::ConstantOp::create(
                rewriter, op.getLoc(), rType,
                cast<ElementsAttr>(makeAttr(rType, 1))));
        rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, rhsLhs, newPowVal);
        return success();
      }

      auto lhsMulDefOp = lhs.getDefiningOp<stablehlo::MulOp>();
      if (lhsMulDefOp && isOnlyUsedInOperation(lhsMulDefOp, op)) {
        auto lhsMulLhs = lhsMulDefOp.getLhs();
        auto lhsMulRhs = lhsMulDefOp.getRhs();

        if (lhsMulLhs == lhsMulRhs &&
            lhsMulLhs ==
                rhsLhs) { // (mul (mul a a) (pow a b)) => (pow a (add b 2))
          auto newPowVal = stablehlo::AddOp::create(
              rewriter, op.getLoc(), rhsDefOp.getRhs(),
              stablehlo::ConstantOp::create(
                  rewriter, op.getLoc(), rType,
                  cast<ElementsAttr>(makeAttr(rType, 2))));
          rewriter.replaceOpWithNewOp<stablehlo::PowOp>(op, rhsLhs, newPowVal);
          return success();
        }
      }
    }

    return failure();
  }
};

struct CommonAssociativeCommutativeOpReorder final
    : public CheckedOpTraitRewritePattern<
          OpTrait::Elementwise, CommonAssociativeCommutativeOpReorder> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise,
      CommonAssociativeCommutativeOpReorder>::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (!(op->hasTrait<OpTrait::IsCommutative>() ||
          op->hasTrait<hlo::OpTrait::IsCommutative>()) ||
        !stablehlo::isAssociativeOp(op) || op->getNumOperands() != 2)
      return failure();

    auto createNewOp = [&](Value lhs, Value rhs) {
      auto newOp = Operation::create(
          op->getLoc(), op->getName(), {op->getResult(0).getType()}, {lhs, rhs},
          op->getAttrs(), mlir::PropertyRef(), op->getSuccessors(), 0);
      rewriter.insert(newOp);
      return newOp;
    };

    auto lhs = op->getOperand(0);
    auto rhs = op->getOperand(1);

    if (lhs == rhs)
      return failure(); // already in a good form

    DenseFPElementsAttr constAttr;
    if (matchPattern(lhs, m_Constant(&constAttr)) &&
        matchPattern(rhs, m_Constant(&constAttr)))
      return rewriter.notifyMatchFailure(op, "const prop has higher priority");

    auto lhsDefOp = lhs.getDefiningOp();
    if (lhsDefOp && lhsDefOp->getName() != op->getName())
      return failure();

    auto rhsDefOp = rhs.getDefiningOp();
    if (rhsDefOp && rhsDefOp->getName() != op->getName())
      return failure();

    if (lhsDefOp && isOnlyUsedInOperation(lhsDefOp, op)) {
      auto lhsLhs = lhsDefOp->getOperand(0);
      auto lhsRhs = lhsDefOp->getOperand(1);

      if (lhsLhs == lhsRhs)
        return failure(); // already in a good form

      if (matchPattern(lhsRhs, m_Constant(&constAttr)) &&
          matchPattern(lhsLhs, m_Constant(&constAttr)))
        return rewriter.notifyMatchFailure(op,
                                           "const prop has higher priority");

      if (rhsDefOp && isOnlyUsedInOperation(rhsDefOp, op)) {
        auto rhsLhs = rhsDefOp->getOperand(0);
        auto rhsRhs = rhsDefOp->getOperand(1);

        if (rhsLhs == rhsRhs)
          return failure(); // already in a good form

        if (matchPattern(rhsRhs, m_Constant(&constAttr)) &&
            matchPattern(rhsLhs, m_Constant(&constAttr)))
          return rewriter.notifyMatchFailure(op,
                                             "const prop has higher priority");

        Value commonOperand, otherOperand1, otherOperand2;
        bool foundCommonOperand = false;

        if (lhsLhs ==
            rhsLhs) { // (op (op x y) (op x z)) => (op (op x x) (op y z))
          if (lhsRhs == rhsRhs)
            return failure(); // we can CSE this case
          commonOperand = lhsLhs;
          otherOperand1 = lhsRhs;
          otherOperand2 = rhsRhs;
          foundCommonOperand = true;
        }

        if (!foundCommonOperand &&
            lhsLhs == rhsRhs) { // (op (op x y) (op z x)) => (op (op x x)
                                // (op y z))
          if (lhsRhs == rhsLhs)
            return failure(); // we can CSE this case
          commonOperand = lhsLhs;
          otherOperand1 = rhsLhs;
          otherOperand2 = lhsRhs;
          foundCommonOperand = true;
        }

        if (!foundCommonOperand &&
            lhsRhs == rhsLhs) { // (op (op y x) (op x z)) => (op (op x x)
                                // (op y z))
          if (lhsLhs == rhsRhs)
            return failure(); // we can CSE this case
          commonOperand = lhsRhs;
          otherOperand1 = lhsLhs;
          otherOperand2 = rhsRhs;
          foundCommonOperand = true;
        }

        if (!foundCommonOperand &&
            lhsRhs == rhsRhs) { // (op (op y x) (op z x)) => (op (op x x)
                                // (op y z))
          if (lhsLhs == rhsLhs)
            return failure(); // we can CSE this case
          commonOperand = lhsRhs;
          otherOperand1 = rhsLhs;
          otherOperand2 = lhsLhs;
          foundCommonOperand = true;
        }

        if (foundCommonOperand) {
          auto newLhs = createNewOp(commonOperand, commonOperand);
          auto newRhs = createNewOp(otherOperand1, otherOperand2);
          auto newReplacementOp =
              createNewOp(newLhs->getResult(0), newRhs->getResult(0));
          rewriter.replaceOp(op, newReplacementOp);
          return success();
        }
      }

      if (lhsLhs == rhs) { // (op (op x y) x) => (op (op x x) y)
        auto newLhs = createNewOp(lhsLhs, lhsLhs);
        auto newReplacementOp = createNewOp(newLhs->getResult(0), lhsRhs);
        rewriter.replaceOp(op, newReplacementOp);
        return success();
      }

      if (lhsRhs == rhs) { // (op (op y x) x) => (op y (op x x))
        auto newRhs = createNewOp(lhsRhs, lhsRhs);
        auto newReplacementOp = createNewOp(lhsLhs, newRhs->getResult(0));
        rewriter.replaceOp(op, newReplacementOp);
        return success();
      }
    }

    if (rhsDefOp && isOnlyUsedInOperation(rhsDefOp, op)) {
      auto rhsLhs = rhsDefOp->getOperand(0);
      auto rhsRhs = rhsDefOp->getOperand(1);

      if (rhsLhs == rhsRhs)
        return failure(); // already in a good form

      if (matchPattern(rhsRhs, m_Constant(&constAttr)) &&
          matchPattern(rhsLhs, m_Constant(&constAttr)))
        return rewriter.notifyMatchFailure(op,
                                           "const prop has higher priority");

      if (rhsLhs == lhs) { // (op x (op x y)) => (op (op x x) y)
        auto newLhs = createNewOp(rhsLhs, rhsLhs);
        auto newReplacementOp = createNewOp(newLhs->getResult(0), rhsRhs);
        rewriter.replaceOp(op, newReplacementOp);
        return success();
      }

      if (rhsRhs == lhs) { // (op x (op y x)) => (op y (op x x))
        auto newRhs = createNewOp(rhsRhs, rhsRhs);
        auto newReplacementOp = createNewOp(rhsLhs, newRhs->getResult(0));
        rewriter.replaceOp(op, newReplacementOp);
        return success();
      }
    }

    return failure();
  }
};

struct LogSimplify final
    : public CheckedOpRewritePattern<stablehlo::LogOp, LogSimplify> {
  using CheckedOpRewritePattern<stablehlo::LogOp,
                                LogSimplify>::CheckedOpRewritePattern;

  // True iff `value` is a finite float constant whose every element is
  // non-negative (and, unless allowZero, strictly positive). Excludes NaN,
  // negatives, and infinities, which would let the log-splits below narrow the
  // domain. See #2570.
  static bool isPositiveFiniteConstant(Value value, bool allowZero) {
    DenseElementsAttr attr;
    if (!matchPattern(value, m_Constant(&attr)) ||
        !isa<FloatType>(attr.getElementType()))
      return false;

    return llvm::all_of(attr.getValues<APFloat>(),
                        [allowZero](const APFloat &element) {
                          return element.isFinite() && !element.isNegative() &&
                                 (allowZero || !element.isZero());
                        });
  }

  LogicalResult matchAndRewriteImpl(stablehlo::LogOp op,
                                    PatternRewriter &rewriter) const {
    { // log(exp(x)) -> x
      auto defOp = op.getOperand().getDefiningOp<stablehlo::ExpOp>();
      if (defOp) {
        rewriter.replaceAllUsesWith(op.getResult(), defOp.getOperand());
        return success();
      }
    }

    {
      auto defOp = op.getOperand().getDefiningOp<stablehlo::MulOp>();
      if (defOp) {
        auto lhs = defOp.getLhs();
        auto rhs = defOp.getRhs();
        // log(mul(a, a)) -> 2 * log(a). Only when a >= 0: for a < 0, 2*log(a)
        // is NaN while log(a*a) is finite. Gating on the analysis avoids an abs
        // (possibly slower than the mul). See #2570.
        if (lhs == rhs && guaranteedNonNegativeResult(lhs, rewriter)) {
          auto two = stablehlo::ConstantOp::create(
              rewriter, op.getLoc(), lhs.getType(),
              cast<ElementsAttr>(makeAttr(lhs.getType(), 2)));
          auto logLhs = stablehlo::LogOp::create(rewriter, op.getLoc(), lhs);
          rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, two, logLhs);
          return success();
        }

        if (anyOperandIsConstant(defOp) &&
            !allOperandsAreConstant(defOp)) { // log(mul(a, b)) -> log(a) +
                                              // log(b) if a or b is constant
          // Split only for a strictly positive constant; negative/zero narrows
          // the domain, and the split isn't bit-exact under FP rounding. See
          // #2570.
          Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs;
          if (isPositiveFiniteConstant(cst, /*allowZero=*/false)) {
            auto logLhs = stablehlo::LogOp::create(rewriter, op.getLoc(), lhs);
            auto logRhs = stablehlo::LogOp::create(rewriter, op.getLoc(), rhs);
            rewriter.replaceOpWithNewOp<stablehlo::AddOp>(op, logLhs, logRhs);
            return success();
          }
        }
      }
    }

    {
      auto defOp = op.getOperand().getDefiningOp<stablehlo::AddOp>();
      if (defOp) {
        auto lhs = defOp.getLhs();
        auto rhs = defOp.getRhs();
        if (lhs == rhs) { // log(add(a, a)) -> log(2) + log(a)
          auto logTwo = stablehlo::LogOp::create(
              rewriter, op.getLoc(),
              stablehlo::ConstantOp::create(
                  rewriter, op.getLoc(), lhs.getType(),
                  cast<ElementsAttr>(makeAttr(lhs.getType(), 2))));
          auto logLhs = stablehlo::LogOp::create(rewriter, op.getLoc(), lhs);
          rewriter.replaceOpWithNewOp<stablehlo::AddOp>(op, logTwo, logLhs);
          return success();
        }

        Attribute rhsAttr, lhsAttr;
        matchPattern(rhs, m_Constant(&rhsAttr));
        matchPattern(lhs, m_Constant(&lhsAttr));
        if (rhsAttr &&
            (matchPattern(rhsAttr, m_One()) ||
             matchPattern(rhsAttr, m_OneFloat()))) { // log(x + 1) -> log1p(x)
          rewriter.replaceOpWithNewOp<stablehlo::Log1pOp>(op, lhs);
          return success();
        }

        if (lhsAttr &&
            (matchPattern(lhsAttr, m_One()) ||
             matchPattern(lhsAttr, m_OneFloat()))) { // log(1 + x) -> log1p(x)
          rewriter.replaceOpWithNewOp<stablehlo::Log1pOp>(op, rhs);
          return success();
        }
      }
    }

    {
      auto defOp = op.getOperand().getDefiningOp<stablehlo::DivOp>();
      if (defOp) {
        auto lhs = defOp.getLhs();
        auto rhs = defOp.getRhs();

        if (anyOperandIsConstant(defOp) &&
            !allOperandsAreConstant(defOp)) { // log(div(a, b)) -> log(a) -
                                              // log(b) if a or b is constant
          // Constant numerator must be strictly positive: a zero numerator
          // turns a finite log(-0) into NaN after splitting. A zero denominator
          // is safe, so it only needs to be non-negative. See #2570.
          bool constantIsLhs = matchPattern(lhs, m_Constant());
          Value cst = constantIsLhs ? lhs : rhs;
          if (isPositiveFiniteConstant(cst, /*allowZero=*/!constantIsLhs)) {
            auto logLhs = stablehlo::LogOp::create(rewriter, op.getLoc(), lhs);
            auto logRhs = stablehlo::LogOp::create(rewriter, op.getLoc(), rhs);
            rewriter.replaceOpWithNewOp<stablehlo::SubtractOp>(op, logLhs,
                                                               logRhs);
            return success();
          }
        }
      }
    }

    {
      auto defOp = op.getOperand().getDefiningOp<stablehlo::SqrtOp>();
      if (defOp &&
          isOnlyUsedInOperation(defOp, op)) { // log(sqrt(x)) -> log(x) / 2
        auto logOperand =
            stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getOperand());
        auto two = stablehlo::ConstantOp::create(
            rewriter, op.getLoc(), defOp.getType(),
            cast<ElementsAttr>(makeAttr(defOp.getType(), 2)));
        rewriter.replaceOpWithNewOp<stablehlo::DivOp>(op, logOperand, two);
        return success();
      }
    }

    {
      auto defOp = op.getOperand().getDefiningOp<stablehlo::CbrtOp>();
      if (defOp &&
          isOnlyUsedInOperation(defOp, op)) { // log(cbrt(x)) -> log(x) / 3
        auto logOperand =
            stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getOperand());
        auto three = stablehlo::ConstantOp::create(
            rewriter, op.getLoc(), defOp.getType(),
            cast<ElementsAttr>(makeAttr(defOp.getType(), 3)));
        rewriter.replaceOpWithNewOp<stablehlo::DivOp>(op, logOperand, three);
        return success();
      }
    }

    {
      auto defOp = op.getOperand().getDefiningOp<stablehlo::RsqrtOp>();
      if (defOp &&
          isOnlyUsedInOperation(defOp, op)) { // log(rsqrt(x)) -> -log(x) / 2
        auto logOperand =
            stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getOperand());
        auto negTwo = stablehlo::ConstantOp::create(
            rewriter, op.getLoc(), defOp.getType(),
            cast<ElementsAttr>(makeAttr(defOp.getType(), -2)));
        rewriter.replaceOpWithNewOp<stablehlo::DivOp>(op, logOperand, negTwo);
        return success();
      }
    }

    return failure();
  }
};

struct NegMulConstSimplify final
    : public CheckedOpRewritePattern<stablehlo::NegOp, NegMulConstSimplify> {
  using CheckedOpRewritePattern<stablehlo::NegOp,
                                NegMulConstSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::NegOp op,
                                    PatternRewriter &rewriter) const {
    auto mulOp = op.getOperand().getDefiningOp<stablehlo::MulOp>();
    if (!mulOp)
      return failure();

    auto lhs = mulOp.getLhs();
    auto rhs = mulOp.getRhs();

    DenseElementsAttr lhsAttr;
    bool lhsIsConst = matchPattern(lhs, m_Constant(&lhsAttr));

    DenseElementsAttr rhsAttr;
    bool rhsIsConst = matchPattern(rhs, m_Constant(&rhsAttr));

    if (lhsIsConst && rhsIsConst)
      return failure(); // const prop will evaluate this

    if (lhsIsConst) {
      rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
          op, stablehlo::NegOp::create(rewriter, op.getLoc(), lhs), rhs);
      return success();
    }

    if (rhsIsConst) {
      rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
          op, lhs, stablehlo::NegOp::create(rewriter, op.getLoc(), rhs));
      return success();
    }

    return failure();
  }
};

struct NegDivConstSimplify final
    : public CheckedOpRewritePattern<stablehlo::NegOp, NegDivConstSimplify> {
  using CheckedOpRewritePattern<stablehlo::NegOp,
                                NegDivConstSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::NegOp op,
                                    PatternRewriter &rewriter) const {
    auto divOp = op.getOperand().getDefiningOp<stablehlo::DivOp>();
    if (!divOp)
      return failure();

    auto lhs = divOp.getLhs();
    auto rhs = divOp.getRhs();

    DenseElementsAttr lhsAttr;
    bool lhsIsConst = matchPattern(lhs, m_Constant(&lhsAttr));

    DenseElementsAttr rhsAttr;
    bool rhsIsConst = matchPattern(rhs, m_Constant(&rhsAttr));

    if (lhsIsConst && rhsIsConst)
      return failure(); // const prop will evaluate this

    if (lhsIsConst) {
      rewriter.replaceOpWithNewOp<stablehlo::DivOp>(
          op, stablehlo::NegOp::create(rewriter, op.getLoc(), lhs), rhs);
      return success();
    }

    if (rhsIsConst) {
      rewriter.replaceOpWithNewOp<stablehlo::DivOp>(
          op, lhs, stablehlo::NegOp::create(rewriter, op.getLoc(), rhs));
      return success();
    }

    return failure();
  }
};

struct ReshapeDeletionsBroadcastInDimSimplify final
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                     ReshapeDeletionsBroadcastInDimSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto bcastInDimOp =
        op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcastInDimOp)
      return failure();

    if (!isOnlyUsedInOperation(bcastInDimOp, op))
      return failure();

    auto deletionDims =
        findReshapeInsertionDims(op.getType(), op.getOperand().getType());
    if (deletionDims.empty())
      return failure();

    SmallVector<int64_t> newBcastDims;
    for (auto dim : bcastInDimOp.getBroadcastDimensions()) {
      int64_t nDeleted = 0;
      for (auto delDim : deletionDims) {
        if (delDim == dim)
          return failure();
        if (delDim < dim)
          nDeleted++;
      }
      int64_t newDim = dim - nDeleted;
      if (newDim < 0)
        return failure(); // this should not happen
      newBcastDims.push_back(newDim);
    }

    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), bcastInDimOp.getOperand(), newBcastDims);
    return success();
  }
};

struct ReshapeInsertionsBroadcastInDimSimplify final
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                     ReshapeInsertionsBroadcastInDimSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto bcastInDimOp =
        op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcastInDimOp)
      return failure();

    if (!isOnlyUsedInOperation(bcastInDimOp, op))
      return failure();

    auto insertionDims =
        findReshapeInsertionDims(op.getOperand().getType(), op.getType());
    if (insertionDims.empty())
      return failure();

    SmallVector<int64_t> newPositions;
    for (int i = 0; i < cast<ShapedType>(op.getType()).getRank(); i++) {
      if (!llvm::is_contained(insertionDims, i))
        newPositions.push_back(i);
    }

    SmallVector<int64_t> newBcastDims;
    for (auto dim : bcastInDimOp.getBroadcastDimensions()) {
      if (dim >= newPositions.size())
        return failure(); // this should not happen
      newBcastDims.push_back(newPositions[dim]);
    }

    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), bcastInDimOp.getOperand(), newBcastDims);
    return success();
  }
};

struct TransposeFFT final
    : public CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeFFT> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto fftOp = op.getOperand().getDefiningOp<stablehlo::FftOp>();
    if (!fftOp)
      return failure();

    auto fftLength = llvm::to_vector(fftOp.getFftLength());

    auto [preTransposePerm, postTransposePerm, newFftLength, needsPostTranspose,
          invalidMove] =
        splitPermutation(op.getPermutation(),
                         cast<ShapedType>(op.getOperand().getType()).getRank(),
                         fftLength.size(), fftOp.getFftType(), fftLength);

    if (invalidMove)
      return rewriter.notifyMatchFailure(fftOp,
                                         "Can't move transpose above FFT");

    auto preTransposeOp = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), fftOp.getOperand(),
        rewriter.getDenseI64ArrayAttr(preTransposePerm));
    auto newFftOp = stablehlo::FftOp::create(rewriter, fftOp.getLoc(),
                                             preTransposeOp.getResult(),
                                             fftOp.getFftType(), newFftLength);

    if (needsPostTranspose) {
      rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(
          op, newFftOp.getResult(),
          rewriter.getDenseI64ArrayAttr(postTransposePerm));
    } else {
      rewriter.replaceOp(op, newFftOp.getResult());
    }

    return success();
  }

private:
  bool isIotaVector(SmallVector<int64_t> v) const {
    for (int64_t i = 0; i < v.size(); ++i) {
      if (v[i] != i)
        return false;
    }
    return true;
  }

  std::tuple<SmallVector<int64_t>, SmallVector<int64_t>, SmallVector<int64_t>,
             bool, bool>
  splitPermutation(ArrayRef<int64_t> permutation, int64_t rank, int64_t nDimFft,
                   stablehlo::FftType fftType,
                   SmallVector<int64_t> fftLength) const {
    SmallVector<int64_t> preTransposePerm(rank), postTransposePerm(rank);
    bool needsPostTranspose = false, invalidMove = false;
    SmallVector<int64_t> newFftLength(fftLength.size());

    std::iota(preTransposePerm.begin(), preTransposePerm.end(), 0);
    std::iota(postTransposePerm.begin(), postTransposePerm.end(), 0);

    int64_t nonFftDims = rank - nDimFft;

    for (int64_t i = 0; i < rank; ++i) {
      int64_t target = permutation[i];
      bool sourceIsNonFft = (i < nonFftDims);
      bool targetIsNonFft = (target < nonFftDims);

      if (sourceIsNonFft != targetIsNonFft) {
        // Cross-boundary permutation - this is not a valid move
        invalidMove = true;
        break;
      }
    }

    if (invalidMove)
      return {preTransposePerm, postTransposePerm, newFftLength,
              needsPostTranspose, invalidMove};

    SmallVector<int64_t> nonFftTargets, nonFftSources;
    for (int64_t i = 0; i < nonFftDims; i++) {
      int64_t target = permutation[i];
      if (target < nonFftDims) {
        nonFftTargets.push_back(target);
        nonFftSources.push_back(i);
      }
    }

    for (size_t i = 0; i < nonFftSources.size(); ++i) {
      preTransposePerm[nonFftTargets[i]] = nonFftSources[i];
    }

    if (fftType == stablehlo::FftType::FFT ||
        fftType == stablehlo::FftType::IFFT) { // can permute fft dims without
                                               // restriction
      for (int64_t i = 0; i < nDimFft; ++i) {
        int64_t fftDimIdx = nonFftDims + i;
        if (permutation[fftDimIdx] != fftDimIdx) {
          preTransposePerm[fftDimIdx] = permutation[fftDimIdx];
          newFftLength[i] = fftLength[permutation[fftDimIdx] - nonFftDims];
        } else {
          newFftLength[i] = fftLength[i];
        }
      }
    } else { // for IRFFT & RFFT we can't move the last dim
      for (int64_t i = 0; i < nDimFft - 1; ++i) {
        int64_t fftDimIdx = nonFftDims + i;
        if (permutation[fftDimIdx] != fftDimIdx) {
          preTransposePerm[fftDimIdx] = permutation[fftDimIdx];
          newFftLength[i] = fftLength[permutation[fftDimIdx] - nonFftDims];
        } else {
          newFftLength[i] = fftLength[i];
        }
      }

      int64_t lastDimIdx = nonFftDims + nDimFft - 1;
      newFftLength[nDimFft - 1] = fftLength[nDimFft - 1];
      if (permutation[lastDimIdx] != lastDimIdx) {
        needsPostTranspose = true;
        postTransposePerm[lastDimIdx] = permutation[lastDimIdx];
      }
    }

    // If we determined iota then we can't move the transpose above the FFT
    invalidMove = invalidMove || isIotaVector(preTransposePerm);
    // Adding a couple of sanity checks below to ensure that the permutation
    // is valid
    if (!invalidMove) {
      SmallVector<int64_t> preTransposePerm2(preTransposePerm.size());
      for (int i = 0; i < preTransposePerm.size(); ++i) {
        preTransposePerm2[i] = preTransposePerm[i];
      }
      std::sort(preTransposePerm2.begin(), preTransposePerm2.end());
      invalidMove = !isIotaVector(preTransposePerm2);
    }
    if (!invalidMove) {
      SmallVector<int64_t> postTransposePerm2(postTransposePerm.size());
      for (int i = 0; i < postTransposePerm.size(); ++i) {
        postTransposePerm2[i] = postTransposePerm[i];
      }
      std::sort(postTransposePerm2.begin(), postTransposePerm2.end());
      invalidMove = !isIotaVector(postTransposePerm2);
    }

    return {preTransposePerm, postTransposePerm, newFftLength,
            needsPostTranspose, invalidMove};
  }
};

// This applies if reshape expands certain dimensions and transpose doesn't move
// those dimensions around. See lit_tests/transpose_reshape.mlir for examples of
// where this is applicable
struct TransposeReshape final
    : public CheckedOpRewritePattern<stablehlo::TransposeOp, TransposeReshape> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
                                    PatternRewriter &rewriter) const {
    auto reshapeOp = op.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshapeOp)
      return failure();

    auto inputShape = llvm::to_vector(
        cast<RankedTensorType>(reshapeOp.getOperand().getType()).getShape());
    auto outputShape =
        llvm::to_vector(cast<RankedTensorType>(reshapeOp.getType()).getShape());

    auto dimOrdering = dimOrderingToPreserve(inputShape, outputShape);

    if (!dimOrdering.valid)
      return rewriter.notifyMatchFailure(op, "Invalid expansion for reshape");

    auto permutation = op.getPermutation();
    SmallVector<int64_t> newPermutation;

    if (dimOrdering.dimOrdering.size() +
            dimOrdering.singletonCollapsedDims.size() !=
        inputShape.size())
      return rewriter.notifyMatchFailure(op, "Invalid dimOrdering");

    for (auto [i, dimList] : llvm::enumerate(dimOrdering.dimOrdering)) {
      if (dimList.size() == 0)
        return rewriter.notifyMatchFailure(op,
                                           "Empty dimList. Should not happen!");

      auto firstDim = dimList[0];
      auto idx = std::find(permutation.begin(), permutation.end(), firstDim);
      if (idx == permutation.end())
        return rewriter.notifyMatchFailure(op, "Invalid permutation");

      int64_t index = std::distance(permutation.begin(), idx);

      for (int i = 1; i < dimList.size(); ++i) {
        if (i + index >= permutation.size() ||
            permutation[i + index] != dimList[i])
          return rewriter.notifyMatchFailure(op, "unsupported permutation");
      }
    }

    int64_t outputIdx = 0;
    for (int i = 0; i < dimOrdering.dimOrdering.size(); ++i) {
      auto transposeIdx = permutation[outputIdx];

      auto idx = std::find(dimOrdering.outputIdxStarts.begin(),
                           dimOrdering.outputIdxStarts.end(), transposeIdx);
      auto index = std::distance(dimOrdering.outputIdxStarts.begin(), idx);

      auto permVal = transposeIdx - dimOrdering.extrasList[index];
      for (int j = 0; j < dimOrdering.singletonCollapsedDims.size(); ++j) {
        if (dimOrdering.singletonCollapsedDims[j] <= permVal) {
          permVal++;
        }
      }
      newPermutation.push_back(permVal);

      for (int j = 0; j < dimOrdering.dimOrdering.size(); ++j) {
        if (dimOrdering.dimOrdering[j][0] == transposeIdx) {
          outputIdx += dimOrdering.dimOrdering[j].size();
          break;
        }
      }
    }

    for (int i = 0; i < dimOrdering.singletonCollapsedDims.size(); ++i) {
      auto dim = dimOrdering.singletonCollapsedDims[i];
      newPermutation.insert(newPermutation.begin() + dim, dim);
    }

    if (!isValidPermutation(newPermutation))
      return rewriter.notifyMatchFailure(op, "Invalid permutation");

    auto newTranspose = stablehlo::TransposeOp::create(
        rewriter, op.getLoc(), reshapeOp.getOperand(),
        rewriter.getDenseI64ArrayAttr(newPermutation));
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      newTranspose.getResult());
    return success();
  }

private:
  bool isValidPermutation(const SmallVector<int64_t> permutation) const {
    const size_t n = permutation.size();
    SmallVector<bool> seen(n, false);

    for (int64_t v : permutation) {
      if (v < 0 || v >= static_cast<int64_t>(n))
        return false;
      if (seen[v])
        return false;
      seen[v] = true;
    }
    return true;
  }

  struct DimOrdering {
    SmallVector<int64_t> outputIdxStarts;
    SmallVector<int64_t> extrasList;
    bool valid;
    SmallVector<SmallVector<int64_t>> dimOrdering;
    SmallVector<int64_t> singletonCollapsedDims;
  };

  DimOrdering dimOrderingToPreserve(SmallVector<int64_t> inputShape,
                                    SmallVector<int64_t> outputShape) const {
    SmallVector<SmallVector<int64_t>> dimOrdering;
    SmallVector<int64_t> outputIdxStarts, extrasList, singletonCollapsedDims;
    bool valid = true;

    int64_t outputIdx = 0, i = 0, extras = 0;
    for (; i < inputShape.size() && valid && outputIdx < outputShape.size();
         ++i) {
      int64_t outputSize = 1;
      SmallVector<int64_t> dims;

      if (inputShape[i] == 1 && outputShape[outputIdx] != 1) {
        singletonCollapsedDims.push_back(i);
        continue;
      }

      outputIdxStarts.push_back(outputIdx);
      extrasList.push_back(extras);

      do {
        outputSize *= outputShape[outputIdx];
        dims.push_back(outputIdx);
        outputIdx++;
      } while (outputIdx < outputShape.size() && outputSize < inputShape[i]);

      if (outputSize != inputShape[i]) {
        valid = false;
        break;
      }

      dimOrdering.push_back(dims);
      extras += dims.size() - 1;
    }

    if (outputIdx != outputShape.size())
      valid = false;

    if (valid && i != inputShape.size()) {
      for (; i < inputShape.size(); ++i) {
        if (inputShape[i] != 1) {
          valid = false;
          break;
        }
        singletonCollapsedDims.push_back(i);
      }
    }

    return DimOrdering{outputIdxStarts, extrasList, valid, dimOrdering,
                       singletonCollapsedDims};
  }
};

struct DotGeneralReshape final
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                     DotGeneralReshape> {
  using CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                DotGeneralReshape>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    auto lhsReshapeOp = op.getLhs().getDefiningOp<stablehlo::ReshapeOp>();
    auto rhsReshapeOp = op.getRhs().getDefiningOp<stablehlo::ReshapeOp>();
    if (!lhsReshapeOp && !rhsReshapeOp) {
      return failure();
    }

    auto dotGeneralDims = op.getDotDimensionNumbers();

    SmallVector<int64_t> adjustedLhsBatchingDims, adjustedLhsContractingDims,
        adjustedRhsBatchingDims, adjustedRhsContractingDims;
    SmallVector<int64_t> lhsInsertionDims, rhsInsertionDims;
    Value newLhs, newRhs;
    bool changed = false;

    newLhs = op.getLhs();
    if (lhsReshapeOp &&
        canFuseOp(lhsReshapeOp, dotGeneralDims.getLhsBatchingDimensions(),
                  dotGeneralDims.getLhsContractingDimensions(),
                  lhsInsertionDims)) {
      changed = true;
      newLhs = lhsReshapeOp.getOperand();
      adjustedLhsBatchingDims = adjustDimsForInsertion(
          dotGeneralDims.getLhsBatchingDimensions(), lhsInsertionDims);
      adjustedLhsContractingDims = adjustDimsForInsertion(
          dotGeneralDims.getLhsContractingDimensions(), lhsInsertionDims);
    } else {
      adjustedLhsBatchingDims =
          llvm::to_vector(dotGeneralDims.getLhsBatchingDimensions());
      adjustedLhsContractingDims =
          llvm::to_vector(dotGeneralDims.getLhsContractingDimensions());
    }

    newRhs = op.getRhs();
    if (rhsReshapeOp &&
        canFuseOp(rhsReshapeOp, dotGeneralDims.getRhsBatchingDimensions(),
                  dotGeneralDims.getRhsContractingDimensions(),
                  rhsInsertionDims)) {
      changed = true;
      newRhs = rhsReshapeOp.getOperand();
      adjustedRhsBatchingDims = adjustDimsForInsertion(
          dotGeneralDims.getRhsBatchingDimensions(), rhsInsertionDims);
      adjustedRhsContractingDims = adjustDimsForInsertion(
          dotGeneralDims.getRhsContractingDimensions(), rhsInsertionDims);
    } else {
      adjustedRhsBatchingDims =
          llvm::to_vector(dotGeneralDims.getRhsBatchingDimensions());
      adjustedRhsContractingDims =
          llvm::to_vector(dotGeneralDims.getRhsContractingDimensions());
    }

    if (!changed) {
      return failure();
    }

    auto lhsShape = cast<RankedTensorType>(newLhs.getType()).getShape();
    auto rhsShape = cast<RankedTensorType>(newRhs.getType()).getShape();
    SmallVector<int64_t> newShape;
    for (int64_t idx : adjustedLhsBatchingDims) {
      newShape.push_back(lhsShape[idx]);
    }
    for (auto [idx, dim] : llvm::enumerate(lhsShape)) {
      if (llvm::is_contained(adjustedLhsContractingDims, idx) ||
          llvm::is_contained(adjustedLhsBatchingDims, idx)) {
        continue;
      }
      newShape.push_back(dim);
    }
    for (auto [idx, dim] : llvm::enumerate(rhsShape)) {
      if (llvm::is_contained(adjustedRhsContractingDims, idx) ||
          llvm::is_contained(adjustedRhsBatchingDims, idx)) {
        continue;
      }
      newShape.push_back(dim);
    }

    auto newDotGeneral = stablehlo::DotGeneralOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(
            newShape,
            cast<RankedTensorType>(newLhs.getType()).getElementType()),
        newLhs, newRhs,
        stablehlo::DotDimensionNumbersAttr::get(
            rewriter.getContext(), adjustedLhsBatchingDims,
            adjustedRhsBatchingDims, adjustedLhsContractingDims,
            adjustedRhsContractingDims),
        op.getPrecisionConfigAttr(), op.getAlgorithmAttr());
    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
                                                      newDotGeneral);
    return success();
  }

private:
  bool canFuseOp(stablehlo::ReshapeOp op, ArrayRef<int64_t> batchingDims,
                 ArrayRef<int64_t> contractingDims,
                 SmallVectorImpl<int64_t> &dims) const {
    SmallVector<int64_t> insertionDims;
    if (!op) {
      return false;
    }

    insertionDims = findReshapeInsertionDims(
        cast<RankedTensorType>(op.getOperand().getType()),
        cast<RankedTensorType>(op.getType()));

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

    for (auto dim : insertionDims) {
      if (llvm::is_contained(batchingDims, dim) ||
          llvm::is_contained(contractingDims, dim)) {
        return false;
      }
    }

    dims.append(insertionDims.begin(), insertionDims.end());
    return true;
  }

  SmallVector<int64_t>
  adjustDimsForInsertion(ArrayRef<int64_t> dims,
                         ArrayRef<int64_t> insertionDims) const {
    SmallVector<int64_t> adjustedDims;
    adjustedDims.reserve(dims.size());

    for (int64_t dim : dims) {
      int64_t adjustedDim = dim;
      for (int64_t insertionDim : insertionDims) {
        if (insertionDim <= adjustedDim) {
          adjustedDim--;
        }
      }
      adjustedDims.push_back(adjustedDim);
    }

    return adjustedDims;
  }
};

struct DiagonalTensorDotGeneralRewrite final
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                     DiagonalTensorDotGeneralRewrite> {
  using CheckedOpRewritePattern<
      stablehlo::DotGeneralOp,
      DiagonalTensorDotGeneralRewrite>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    auto rhsType = op.getRhs().getType();
    auto lhsType = op.getLhs().getType();
    auto dotDimNumbers = op.getDotDimensionNumbers();

    if (dotDimNumbers.getLhsBatchingDimensions().size() != 0 ||
        dotDimNumbers.getRhsBatchingDimensions().size() != 0 ||
        dotDimNumbers.getLhsContractingDimensions().size() != 1 ||
        dotDimNumbers.getRhsContractingDimensions().size() != 1) {
      return rewriter.notifyMatchFailure(
          op, "batching is not supported for diagonal tensor dot general");
    }

    int64_t rhsRank = cast<RankedTensorType>(rhsType).getRank();
    int64_t lhsRank = cast<RankedTensorType>(lhsType).getRank();

    // Support 2D x 2D, 2D x 1D, and 1D x 2D cases
    if ((lhsRank != 2 && lhsRank != 1) || (rhsRank != 2 && rhsRank != 1)) {
      return rewriter.notifyMatchFailure(op,
                                         "only 1D and 2D tensors are supported "
                                         "for diagonal tensor dot general");
    }

    if (lhsRank == 1 && rhsRank == 1) {
      return rewriter.notifyMatchFailure(
          op, "1D x 1D dot product not supported in this pattern");
    }

    int64_t lhsContractingDim = dotDimNumbers.getLhsContractingDimensions()[0];
    int64_t rhsContractingDim = dotDimNumbers.getRhsContractingDimensions()[0];

    {
      // LHS Diagonal Matrix Multiply: Scale rows by diagonal elements
      mlir::Value diagonalTensor;
      if (auto scatterOp = op.getLhs().getDefiningOp<stablehlo::ScatterOp>()) {
        auto status = detectDiagonalTensor(scatterOp, &diagonalTensor);
        if (status.ok()) {
          // TODO: support batched diagonal tensor
          if (cast<RankedTensorType>(diagonalTensor.getType()).getRank() == 1) {
            // LHS is diagonal (2D), RHS can be 2D or 1D
            if (rhsRank == 2) {
              int64_t rhsNonContractingDim = 1 - rhsContractingDim;
              auto scaling = stablehlo::BroadcastInDimOp::create(
                  rewriter, op.getLoc(), rhsType, diagonalTensor,
                  ArrayRef<int64_t>({rhsContractingDim}));
              auto result = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                                     scaling, op.getRhs());
              if (rhsNonContractingDim == 0) {
                rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(
                    op, op.getType(), result, ArrayRef<int64_t>({1, 0}));
              } else {
                rewriter.replaceOp(op, result);
              }
            } else {
              rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, op.getRhs(),
                                                            diagonalTensor);
            }
            return success();
          }
        }
      }
    }

    {
      // RHS Diagonal Matrix Multiply: Scale columns by diagonal elements
      mlir::Value diagonalTensor;
      if (auto scatterOp = op.getRhs().getDefiningOp<stablehlo::ScatterOp>()) {
        auto status = detectDiagonalTensor(scatterOp, &diagonalTensor);
        if (status.ok()) {
          // TODO: support batched diagonal tensor
          if (cast<RankedTensorType>(diagonalTensor.getType()).getRank() == 1) {
            // RHS is diagonal (2D), LHS can be 2D or 1D
            if (lhsRank == 2) {
              int64_t lhsNonContractingDim = 1 - lhsContractingDim;
              auto scaling = stablehlo::BroadcastInDimOp::create(
                  rewriter, op.getLoc(), lhsType, diagonalTensor,
                  ArrayRef<int64_t>({lhsContractingDim}));
              auto result = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                                     scaling, op.getLhs());
              if (lhsNonContractingDim == 1) {
                rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(
                    op, op.getType(), result, ArrayRef<int64_t>({1, 0}));
              } else {
                rewriter.replaceOp(op, result);
              }
            } else {
              rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, op.getLhs(),
                                                            diagonalTensor);
            }
            return success();
          }
        }
      }
    }

    return failure();
  }
};

// pred ? trues : falses --> pred
// pred ? falses : trues --> !pred
struct SelectSimplify
    : public CheckedOpRewritePattern<stablehlo::SelectOp, SelectSimplify> {
  using CheckedOpRewritePattern<stablehlo::SelectOp,
                                SelectSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SelectOp op,
                                    PatternRewriter &rewriter) const {
    if (!cast<ShapedType>(op.getType()).getElementType().isInteger(1))
      return failure();

    auto shardingAttr = sdy::getShardingPerValue(op);

    auto predType = cast<RankedTensorType>(op.getPred().getType());
    bool needsBroadcast = predType.getRank() == 0;

    Attribute onTrueAttr, onFalseAttr;
    matchPattern(op.getOnTrue(), m_Constant(&onTrueAttr));
    matchPattern(op.getOnFalse(), m_Constant(&onFalseAttr));

    if (onTrueAttr && matchPattern(onTrueAttr, m_One()) && onFalseAttr &&
        matchPattern(onFalseAttr, m_Zero())) {
      if (needsBroadcast) {
        auto bcast = rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
            op, op.getType(), op.getPred(), ArrayRef<int64_t>({}));
        if (shardingAttr)
          sdy::setShardings(bcast, shardingAttr);
      } else {
        rewriter.replaceAllUsesWith(op, op.getPred());
      }
      return success();
    }

    if (onTrueAttr && matchPattern(onTrueAttr, m_Zero()) && onFalseAttr &&
        matchPattern(onFalseAttr, m_One())) {
      auto notOp =
          stablehlo::NotOp::create(rewriter, op.getLoc(), op.getPred());
      if (shardingAttr)
        sdy::setShardings(notOp, shardingAttr);

      if (needsBroadcast) {
        auto bcast = rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
            op, op.getType(), notOp, ArrayRef<int64_t>({}));
        if (shardingAttr)
          sdy::setShardings(bcast, shardingAttr);
      } else {
        rewriter.replaceAllUsesWith(op, notOp);
      }
      return success();
    }

    return failure();
  }
};

struct ElementwisePad
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          ElementwisePad> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise, ElementwisePad>::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    if (op->getNumOperands() != 1)
      return failure();

    auto elemType =
        cast<RankedTensorType>(op->getResultTypes()[0]).getElementType();

    auto operand = op->getOperand(0);
    auto padOp = operand.getDefiningOp<stablehlo::PadOp>();
    if (!padOp)
      return failure();
    if (!llvm::hasSingleElement(padOp->getUsers()))
      return failure();

    auto padOperand = padOp.getOperand();
    auto padValue = padOp.getPaddingValue();

    auto elemOperand = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(), ValueRange(padOperand),
        TypeRange{RankedTensorType::get(
            cast<RankedTensorType>(padOperand.getType()).getShape(), elemType)},
        op->getAttrs(), {}, {});
    auto elemPadValue = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(), ValueRange(padValue),
        TypeRange{RankedTensorType::get({}, elemType)}, op->getAttrs(), {}, {});

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, elemOperand->getResult(0), elemPadValue->getResult(0),
        padOp.getEdgePaddingLow(), padOp.getEdgePaddingHigh(),
        padOp.getInteriorPadding());
    return success();
  }
};

// See https://github.com/EnzymeAD/Enzyme-JAX/issues/1204 for details
struct ConcatenateSubtractToSubtractPad
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     ConcatenateSubtractToSubtractPad> {
  using CheckedOpRewritePattern<
      stablehlo::ConcatenateOp,
      ConcatenateSubtractToSubtractPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto operands = op.getOperands();
    if (operands.size() != 3)
      return failure();

    auto left = operands[0];
    auto mid = operands[1];
    auto right = operands[2];
    auto concatDim = op.getDimension();

    auto midSubtractOp = mid.getDefiningOp<stablehlo::SubtractOp>();
    if (!midSubtractOp)
      return failure();

    auto rightNegateOp = right.getDefiningOp<stablehlo::NegOp>();
    if (!rightNegateOp)
      return failure();

    auto outputShape = cast<ShapedType>(op.getType()).getShape();
    auto opRank = outputShape.size();

    auto zeroType = RankedTensorType::get(
        {}, cast<ShapedType>(op.getType()).getElementType());
    auto zero = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), zeroType,
        cast<ElementsAttr>(makeAttr(zeroType, 0)));

    auto newLhsPrePadding = stablehlo::ConcatenateOp::create(
        rewriter, op.getLoc(), ValueRange({left, midSubtractOp.getLhs()}),
        concatDim);

    SmallVector<int64_t> lhsPadLow(opRank, 0);
    SmallVector<int64_t> lhsPadHigh(opRank, 0);
    lhsPadHigh[concatDim] =
        outputShape[concatDim] -
        cast<ShapedType>(newLhsPrePadding.getType()).getShape()[concatDim];
    SmallVector<int64_t> lhsPadInner(opRank, 0);
    auto fullNewLhs =
        stablehlo::PadOp::create(rewriter, op.getLoc(), newLhsPrePadding, zero,
                                 lhsPadLow, lhsPadHigh, lhsPadInner);

    auto newRhsPrePadding = stablehlo::ConcatenateOp::create(
        rewriter, op.getLoc(),
        ValueRange({midSubtractOp.getRhs(), rightNegateOp.getOperand()}),
        concatDim);

    SmallVector<int64_t> rhsPadLow(opRank, 0);
    rhsPadLow[concatDim] =
        outputShape[concatDim] -
        cast<ShapedType>(newRhsPrePadding.getType()).getShape()[concatDim];
    SmallVector<int64_t> rhsPadHigh(opRank, 0);
    SmallVector<int64_t> rhsPadInner(opRank, 0);
    auto fullNewRhs =
        stablehlo::PadOp::create(rewriter, op.getLoc(), newRhsPrePadding, zero,
                                 rhsPadLow, rhsPadHigh, rhsPadInner);

    rewriter.replaceOpWithNewOp<stablehlo::SubtractOp>(op, fullNewLhs,
                                                       fullNewRhs);
    return success();
  }
};

struct ConcatenateAddToAddPad
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     ConcatenateAddToAddPad> {
  using CheckedOpRewritePattern<
      stablehlo::ConcatenateOp,
      ConcatenateAddToAddPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto operands = op.getOperands();
    if (operands.size() != 2)
      return failure();

    auto left = operands[0];
    auto mid = operands[1];
    auto concatDim = op.getDimension();

    auto midSubtractOp = mid.getDefiningOp<stablehlo::AddOp>();
    if (!midSubtractOp)
      return failure();

    if (midSubtractOp.getType().getShape()[concatDim] <=
        cast<RankedTensorType>(left.getType()).getShape()[concatDim])
      return failure();

    auto leftSlice = left.getDefiningOp<stablehlo::SliceOp>();

    if (!leftSlice)
      return failure();

    auto outputShape = cast<ShapedType>(op.getType()).getShape();
    auto opRank = outputShape.size();

    for (int j = 0; j < 2; j++) {
      auto rightSlice =
          midSubtractOp->getOperand(j).getDefiningOp<stablehlo::SliceOp>();
      if (!rightSlice)
        continue;
      if (rightSlice.getOperand() != leftSlice.getOperand())
        continue;

      bool legal = true;
      for (size_t i = 0; i < opRank; i++) {
        if (i == concatDim) {
          if (leftSlice.getLimitIndices()[i] !=
              rightSlice.getStartIndices()[i]) {
            legal = false;
            break;
          }
          if (leftSlice.getStrides()[i] != 1) {
            legal = false;
            break;
          }
          if (rightSlice.getStrides()[i] != 1) {
            legal = false;
            break;
          }
        } else {
          if (leftSlice.getStartIndices()[i] !=
              rightSlice.getStartIndices()[i]) {
            legal = false;
            break;
          }
          if (leftSlice.getLimitIndices()[i] !=
              rightSlice.getLimitIndices()[i]) {
            legal = false;
            break;
          }
          if (leftSlice.getStrides()[i] != rightSlice.getStrides()[i]) {
            legal = false;
            break;
          }
        }
      }
      if (!legal)
        continue;

      SmallVector<int64_t> starts =
          llvm::to_vector(leftSlice.getStartIndices());
      SmallVector<int64_t> limits =
          llvm::to_vector(rightSlice.getLimitIndices());
      SmallVector<int64_t> strides = llvm::to_vector(rightSlice.getStrides());

      auto preSlice = stablehlo::SliceOpCreate(rewriter, op.getLoc(),
                                               leftSlice.getOperand(), starts,
                                               limits, strides);

      auto zeroType = RankedTensorType::get(
          {}, cast<ShapedType>(op.getType()).getElementType());
      auto zero = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), zeroType,
          cast<ElementsAttr>(makeAttr(zeroType, 0)));
      SmallVector<int64_t> lhsPadLow(opRank, 0);
      SmallVector<int64_t> lhsPadHigh(opRank, 0);
      SmallVector<int64_t> lhsPadInner(opRank, 0);
      lhsPadLow[concatDim] = leftSlice.getType().getShape()[concatDim];
      auto newRight = stablehlo::PadOp::create(
          rewriter, op.getLoc(), midSubtractOp->getOperand(1 - j), zero,
          lhsPadLow, lhsPadHigh, lhsPadInner);
      rewriter.replaceOpWithNewOp<stablehlo::AddOp>(op, preSlice, newRight);
      return success();
    }
    return failure();
  }
};

struct ConcatenateBroadcastInDim
    : public CheckedOpRewritePattern<stablehlo::ConcatenateOp,
                                     ConcatenateBroadcastInDim> {
  using CheckedOpRewritePattern<
      stablehlo::ConcatenateOp,
      ConcatenateBroadcastInDim>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConcatenateOp op,
                                    PatternRewriter &rewriter) const {
    auto operands = op.getOperands();
    if (operands.size() == 1)
      return failure();

    int input_dim = 0;
    SmallVector<Value> operandOperands;
    for (auto operand : operands) {
      if (auto broadcastInDimOp =
              operand.getDefiningOp<stablehlo::BroadcastInDimOp>()) {
        auto bcastInDimDimensions = broadcastInDimOp.getBroadcastDimensions();
        if (bcastInDimDimensions.size() != 1)
          return failure();
        if (bcastInDimDimensions[input_dim] != op.getDimension())
          return failure();
        if (broadcastInDimOp.getOperand().getType().getShape()[input_dim] !=
            broadcastInDimOp.getType().getShape()[op.getDimension()])
          return failure();
        operandOperands.push_back(broadcastInDimOp.getOperand());
        continue;
      }
      return failure();
    }

    auto newConcatOp = stablehlo::ConcatenateOp::create(
        rewriter, op.getLoc(), ValueRange(operandOperands), input_dim);
    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        op, op.getType(), newConcatOp,
        ArrayRef<int64_t>({static_cast<int64_t>(op.getDimension())}));
    return success();
  }
};

template <typename OpTy, typename Child>
struct ElementwiseConcatLike
    : public CheckedOpTraitRewritePattern<OpTrait::Elementwise,
                                          ElementwiseConcatLike<OpTy, Child>> {
  using CheckedOpTraitRewritePattern<
      OpTrait::Elementwise,
      ElementwiseConcatLike<OpTy, Child>>::CheckedOpTraitRewritePattern;

  LogicalResult matchAndRewriteImpl(Operation *op,
                                    PatternRewriter &rewriter) const {
    auto operands = op->getOperands();
    if (operands.size() <= 1)
      return failure();

    SmallVector<Value> operandOperands;
    SmallVector<Operation *> concatLikeOps;
    for (auto operand : operands) {
      if (auto concatLikeOp = operand.template getDefiningOp<OpTy>()) {
        concatLikeOps.push_back(concatLikeOp);
        operandOperands.push_back(concatLikeOp.getOperand());

        if (!llvm::hasSingleElement(operand.getUsers())) {
          return failure();
        } else if (concatLikeOps.size() > 1) {
          if (!OperationEquivalence::isEquivalentTo(
                  concatLikeOps[0], concatLikeOp,
                  OperationEquivalence::ignoreValueEquivalence, nullptr,
                  OperationEquivalence::IgnoreLocations, nullptr)) {
            return failure();
          }
        }
      } else {
        return failure();
      }
    }

    // Do the elementwise first
    auto newElem = rewriter.create(
        op->getLoc(), op->getName().getIdentifier(),
        ValueRange(operandOperands),
        TypeRange{RankedTensorType::get(
            cast<RankedTensorType>(operandOperands[0].getType()).getShape(),
            cast<RankedTensorType>(op->getResultTypes()[0]).getElementType())},
        op->getAttrs(), {}, {});

    // Then the concatlike
    ((Child *)this)
        ->createSimilarConcatLike(rewriter, op, newElem->getResult(0),
                                  concatLikeOps[0]);
    return success();
  }
};

struct ElementwiseRotate
    : public ElementwiseConcatLike<enzymexla::RotateOp, ElementwiseRotate> {
  using ElementwiseConcatLike<enzymexla::RotateOp,
                              ElementwiseRotate>::ElementwiseConcatLike;

  void createSimilarConcatLike(PatternRewriter &rewriter, Operation *op,
                               Value elem, Operation *baseRotateOp) const {
    auto rotateOp = dyn_cast<enzymexla::RotateOp>(baseRotateOp);
    assert(rotateOp);
    rewriter.replaceOpWithNewOp<enzymexla::RotateOp>(
        op, elem, rotateOp.getAmount(), rotateOp.getDimension());
    return;
  }
};

struct ElementwiseWrap
    : public ElementwiseConcatLike<enzymexla::WrapOp, ElementwiseWrap> {
  using ElementwiseConcatLike<enzymexla::WrapOp,
                              ElementwiseWrap>::ElementwiseConcatLike;

  void createSimilarConcatLike(PatternRewriter &rewriter, Operation *op,
                               Value elem, Operation *baseWrapOp) const {
    auto wrapOp = dyn_cast<enzymexla::WrapOp>(baseWrapOp);
    assert(wrapOp);
    rewriter.replaceOpWithNewOp<enzymexla::WrapOp>(
        op, elem, wrapOp.getLhs(), wrapOp.getRhs(), wrapOp.getDimension());
    return;
  }
};

struct ElementwiseExtend
    : public ElementwiseConcatLike<enzymexla::ExtendOp, ElementwiseExtend> {
  using ElementwiseConcatLike<enzymexla::ExtendOp,
                              ElementwiseExtend>::ElementwiseConcatLike;

  void createSimilarConcatLike(PatternRewriter &rewriter, Operation *op,
                               Value elem, Operation *baseExtendOp) const {
    auto extendOp = dyn_cast<enzymexla::ExtendOp>(baseExtendOp);
    assert(extendOp);
    rewriter.replaceOpWithNewOp<enzymexla::ExtendOp>(
        op, elem, extendOp.getLhs(), extendOp.getRhs(),
        extendOp.getDimension());
    return;
  }
};

struct SubtractMultiplyConstToAddMulConst
    : public CheckedOpRewritePattern<stablehlo::SubtractOp,
                                     SubtractMultiplyConstToAddMulConst> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SubtractOp op,
                                    PatternRewriter &rewriter) const {
    auto rhsMulOp = op.getRhs().getDefiningOp<stablehlo::MulOp>();
    if (!rhsMulOp)
      return failure();

    DenseElementsAttr rhsMulLhsConstant, rhsMulRhsConstant;
    bool isRhsMulLhsConstant =
        matchPattern(rhsMulOp.getLhs(), m_Constant(&rhsMulLhsConstant));
    bool isRhsMulRhsConstant =
        matchPattern(rhsMulOp.getRhs(), m_Constant(&rhsMulRhsConstant));

    if (isRhsMulLhsConstant && isRhsMulRhsConstant) {
      return failure(); // constant prop will handle this
    }

    if (isRhsMulLhsConstant) {
      rewriter.replaceOpWithNewOp<stablehlo::AddOp>(
          op, op.getLhs(),
          stablehlo::MulOp::create(rewriter, op.getLoc(),
                                   stablehlo::NegOp::create(rewriter,
                                                            op.getLoc(),
                                                            rhsMulOp.getLhs()),
                                   rhsMulOp.getRhs()));
      return success();
    } else if (isRhsMulRhsConstant) {
      rewriter.replaceOpWithNewOp<stablehlo::AddOp>(
          op, op.getLhs(),
          stablehlo::MulOp::create(
              rewriter, op.getLoc(), rhsMulOp.getLhs(),
              stablehlo::NegOp::create(rewriter, op.getLoc(),
                                       rhsMulOp.getRhs())));
      return success();
    }

    return failure();
  }
};

template <typename OpTy>
struct SelfElementwiseToConvolutionLike
    : public CheckedOpRewritePattern<OpTy,
                                     SelfElementwiseToConvolutionLike<OpTy>> {
  using Base =
      CheckedOpRewritePattern<OpTy, SelfElementwiseToConvolutionLike<OpTy>>;
  using Base::Base;

  SelfElementwiseToConvolutionLike<OpTy>(bool allowEmitConvolution,
                                         MLIRContext *ctx,
                                         PatternBenefit benefit = 1)
      : Base(ctx, benefit), allowEmitConvolution(allowEmitConvolution) {}

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    if (op->getNumOperands() != 2)
      return rewriter.notifyMatchFailure(op, "expected two operands");

    auto outType = cast<RankedTensorType>(op.getType());
    auto outShape = outType.getShape();
    auto outRank = outType.getRank();

    // Extract info about both operands
    auto lhsInfo = extractOperandInfo(op.getLhs());
    auto rhsInfo = extractOperandInfo(op.getRhs());

    // For reduce_window, we need the following:
    //   - op is commutative and associative
    //   - *.constantScalar & *.constantAttr must be empty
    bool canEmitReduceWindow = false;
    bool isCommutative =
        op->template hasTrait<mlir::hlo::OpTrait::IsCommutative>() ||
        op->template hasTrait<mlir::OpTrait::IsCommutative>();
    // mlir has no trait for isAssociative??
    bool isAssociative = false;
    if constexpr (std::is_base_of_v<stablehlo::AddOp, OpTy> ||
                  std::is_base_of_v<stablehlo::MulOp, OpTy>) {
      isAssociative = true;
    }
    if (isAssociative && isCommutative) {
      if (!lhsInfo.constantAttr.has_value() &&
          !rhsInfo.constantAttr.has_value() &&
          !lhsInfo.constantScalar.has_value() &&
          !rhsInfo.constantScalar.has_value()) {
        canEmitReduceWindow = true;
      }
    }

    if (!canEmitReduceWindow) {
      if constexpr (std::is_base_of_v<stablehlo::MulOp, OpTy>) {
        return rewriter.notifyMatchFailure(
            op, "MulOp is only supported if we can emit a reduce window.");
      }

      if (!allowEmitConvolution) {
        return rewriter.notifyMatchFailure(
            op, "Emitting Convolution is not Allowed.");
      }
    }

    Operation *newBaseOp = nullptr;
    int64_t windowDilation = -1, dimension = -1;
    bool flipped = false;

    {
      // Try to see if the parent op is a rotate
      bool matched = false;
      flipped = false;
      windowDilation = -1;
      dimension = -1;
      Value operand;

      // TODO: generalize the pattern. We can deal with both branches being a
      // rotate
      if (auto lhsRotateOp =
              lhsInfo.base.template getDefiningOp<enzymexla::RotateOp>()) {
        if (lhsRotateOp.getOperand() == rhsInfo.base) {
          dimension = lhsRotateOp.getDimension();
          windowDilation = lhsRotateOp.getAmount();
          operand = rhsInfo.base;
          matched = true;
          flipped = true;
        }
      }

      if (!matched) {
        if (auto rhsRotateOp =
                rhsInfo.base.template getDefiningOp<enzymexla::RotateOp>()) {
          if (rhsRotateOp.getOperand() == lhsInfo.base) {
            dimension = rhsRotateOp.getDimension();
            windowDilation = rhsRotateOp.getAmount();
            operand = lhsInfo.base;
            matched = true;
          }
        }
      }

      if (matched)
        newBaseOp = enzymexla::WrapOp::create(rewriter, op.getLoc(), operand, 0,
                                              windowDilation, dimension);
    }

    if (!newBaseOp) {
      // match pads to convolutions. see pad_to_convlike.mlir
      auto lhsPadOp = lhsInfo.base.template getDefiningOp<stablehlo::PadOp>();
      auto rhsPadOp = rhsInfo.base.template getDefiningOp<stablehlo::PadOp>();
      bool matched = false;
      flipped = false;
      windowDilation = -1;
      dimension = -1;

      if (lhsPadOp && rhsPadOp) {
        if (lhsPadOp.getOperand() == rhsPadOp.getOperand() &&
            lhsPadOp.getPaddingValue() == rhsPadOp.getPaddingValue()) {
          auto lhsLowPadding = lhsPadOp.getEdgePaddingLow();
          auto rhsLowPadding = rhsPadOp.getEdgePaddingLow();
          auto lhsHighPadding = lhsPadOp.getEdgePaddingHigh();
          auto rhsHighPadding = rhsPadOp.getEdgePaddingHigh();
          auto lhsInteriorPadding = lhsPadOp.getInteriorPadding();
          auto rhsInteriorPadding = rhsPadOp.getInteriorPadding();

          int64_t counter = 0;
          for (auto [lL, rL, lH, rH, lI, rI] : llvm::zip(
                   lhsLowPadding, rhsLowPadding, lhsHighPadding, rhsHighPadding,
                   lhsInteriorPadding, rhsInteriorPadding)) {
            // all interior paddings must be zero
            if (lI != 0 || rI != 0) {
              matched = false;
              break;
            }

            // if lL is non zero, corresponding rL must be non zero. All others
            // must be zero
            if (lL != 0) {
              if (lL != rH || lH != 0 || rL != 0 || dimension != -1) {
                matched = false;
                break;
              }
              dimension = counter;
              windowDilation = lL;
              matched = true;
            }

            // if rL is non zero, corresponding lL must be non zero. All others
            // must be zero
            if (rL != 0) {
              if (lH != rL || lL != 0 || rH != 0 || dimension != -1) {
                matched = false;
                break;
              }
              dimension = counter;
              windowDilation = rL;
              matched = true;
              flipped = true;
            }

            counter++;
          }
        }
      }

      if (matched) {
        SmallVector<int64_t> paddingLow(outRank, 0);
        SmallVector<int64_t> paddingHigh(outRank, 0);
        SmallVector<int64_t> paddingInterior(outRank, 0);
        paddingLow[dimension] = windowDilation;
        paddingHigh[dimension] = windowDilation;

        newBaseOp = stablehlo::PadOp::create(
            rewriter, op.getLoc(), lhsPadOp.getOperand(),
            lhsPadOp.getPaddingValue(), paddingLow, paddingHigh,
            paddingInterior);
      }
    }

    if (!newBaseOp) {
      // matching slice to convolution. see slice_to_convlike.mlir for example
      auto lhsSliceOp =
          lhsInfo.base.template getDefiningOp<stablehlo::SliceOp>();
      auto rhsSliceOp =
          rhsInfo.base.template getDefiningOp<stablehlo::SliceOp>();
      SmallVector<int64_t> newSliceStarts(outRank, -1),
          newSliceLimits(outRank, -1), newSliceStrides(outRank, 1);
      bool matched = false;
      flipped = false;
      windowDilation = -1;
      dimension = -1;

      if (lhsSliceOp && rhsSliceOp &&
          lhsSliceOp.getOperand() == rhsSliceOp.getOperand()) {
        auto lhsSliceStarts = lhsSliceOp.getStartIndices();
        auto lhsSliceLimits = lhsSliceOp.getLimitIndices();

        auto rhsSliceStarts = rhsSliceOp.getStartIndices();
        auto rhsSliceLimits = rhsSliceOp.getLimitIndices();

        // for now we don't support strided slices
        if (llvm::all_of(lhsSliceOp.getStrides(),
                         [](int64_t i) { return i == 1; }) &&
            llvm::all_of(rhsSliceOp.getStrides(),
                         [](int64_t i) { return i == 1; })) {
          for (int i = 0; i < outRank; ++i) {
            if (lhsSliceStarts[i] == rhsSliceStarts[i] &&
                lhsSliceLimits[i] == rhsSliceLimits[i]) {
              newSliceStarts[i] = lhsSliceStarts[i];
              newSliceLimits[i] = lhsSliceLimits[i];
            } else if (matched) {
              // If we have already matched, we can't match anymore.
              matched = false;
              break;
            } else {
              if (lhsSliceStarts[i] > rhsSliceStarts[i]) {
                windowDilation = lhsSliceStarts[i] - rhsSliceStarts[i];
                flipped = true;
              } else {
                windowDilation = rhsSliceStarts[i] - lhsSliceStarts[i];
              }
              dimension = i;
              newSliceStarts[i] =
                  std::min(lhsSliceStarts[i], rhsSliceStarts[i]);
              newSliceLimits[i] =
                  std::max(lhsSliceLimits[i], rhsSliceLimits[i]);
              matched = true;
            }
          }
        }
      }

      if (matched) {
        newBaseOp = stablehlo::SliceOp::create(
            rewriter, op.getLoc(), lhsSliceOp.getOperand(), newSliceStarts,
            newSliceLimits, newSliceStrides);
      }
    }

    if (!newBaseOp && canEmitReduceWindow) {
      // Match add(reduce_window(slice(X)), slice(X)) or the symmetric case.
      // This allows iteratively folding chains of add(slice, slice, ..., slice)
      // into a single reduce_window across multiple greedy rewriter iterations.
      //
      // Given: reduce_window(slice[s:e], window_dim=N) + slice[s+N : e+1]
      // Produce: reduce_window(slice[s : e+1], window_dim=N+1)
      //
      // Similarly for the prepend case:
      // Given: reduce_window(slice[s:e], window_dim=N) + slice[s-1 : e-N]
      // Produce: reduce_window(slice[s-1 : e], window_dim=N+1)

      auto tryMatchReduceWindowSlice = [&](Value rwVal,
                                           Value sliceVal) -> bool {
        auto rwOp = rwVal.template getDefiningOp<stablehlo::ReduceWindowOp>();
        if (!rwOp || rwOp.getInputs().size() != 1)
          return false;

        // Check that the reduce_window body matches OpTy
        auto commonRW = CheckCommonReduceWindowOp(rwOp);
        ReduceOpKind expectedKind = ReduceOpKind::Unknown;
        if constexpr (std::is_base_of_v<stablehlo::AddOp, OpTy>)
          expectedKind = ReduceOpKind::Add;
        else if constexpr (std::is_base_of_v<stablehlo::MulOp, OpTy>)
          expectedKind = ReduceOpKind::Mul;
        if (commonRW.kind != expectedKind)
          return false;

        // The reduce_window input must be a slice
        auto rwInputSlice =
            rwOp.getInputs()[0].template getDefiningOp<stablehlo::SliceOp>();
        if (!rwInputSlice)
          return false;

        // The other operand must also be a slice from the same tensor
        auto otherSlice = sliceVal.template getDefiningOp<stablehlo::SliceOp>();
        if (!otherSlice)
          return false;

        if (rwInputSlice.getOperand() != otherSlice.getOperand())
          return false;

        // Both slices must have stride 1
        if (!llvm::all_of(rwInputSlice.getStrides(),
                          [](int64_t i) { return i == 1; }) ||
            !llvm::all_of(otherSlice.getStrides(),
                          [](int64_t i) { return i == 1; }))
          return false;

        // reduce_window must have no strides, no base_dilations, no padding,
        // and all window_dilations == 1
        auto rwWindowStrides = rwOp.getWindowStrides();
        auto rwBaseDilations = rwOp.getBaseDilations();
        auto rwWindowDilations = rwOp.getWindowDilations();
        auto rwPadding = rwOp.getPadding();

        if (rwWindowStrides && !llvm::all_of(rwWindowStrides.value(),
                                             [](int64_t i) { return i == 1; }))
          return false;
        if (rwBaseDilations && !llvm::all_of(rwBaseDilations.value(),
                                             [](int64_t i) { return i == 1; }))
          return false;
        if (rwWindowDilations &&
            !llvm::all_of(rwWindowDilations.value(),
                          [](int64_t i) { return i == 1; }))
          return false;
        if (rwPadding && !llvm::all_of(rwPadding.value(), [](auto pad_val) {
              return pad_val == 0;
            }))
          return false;

        auto rwWindowDims = rwOp.getWindowDimensions();

        auto rwSliceStarts = rwInputSlice.getStartIndices();
        auto rwSliceLimits = rwInputSlice.getLimitIndices();
        auto otherStarts = otherSlice.getStartIndices();
        auto otherLimits = otherSlice.getLimitIndices();

        // Find the single dimension where the window_dim > 1
        int64_t rwDim = -1;
        int64_t rwWindowSize = -1;
        for (int64_t i = 0; i < outRank; ++i) {
          if (rwWindowDims[i] != 1) {
            if (rwDim != -1)
              return false; // multiple active window dims not supported
            rwDim = i;
            rwWindowSize = rwWindowDims[i];
          }
        }
        if (rwDim == -1)
          return false;

        // All non-active dims must have identical slice bounds
        for (int64_t i = 0; i < outRank; ++i) {
          if (i == rwDim)
            continue;
          if (rwSliceStarts[i] != otherStarts[i] ||
              rwSliceLimits[i] != otherLimits[i])
            return false;
        }

        // Check if the other slice is adjacent on the right:
        //   otherStarts[rwDim] == rwSliceStarts[rwDim] + rwWindowSize
        //   otherLimits[rwDim] == rwSliceLimits[rwDim] + 1
        bool appendRight =
            (otherStarts[rwDim] == rwSliceStarts[rwDim] + rwWindowSize) &&
            (otherLimits[rwDim] == rwSliceLimits[rwDim] + 1);

        // Check if the other slice is adjacent on the left:
        //   otherStarts[rwDim] == rwSliceStarts[rwDim] - 1
        //   otherLimits[rwDim] == rwSliceLimits[rwDim] - rwWindowSize
        bool prependLeft =
            (otherStarts[rwDim] == rwSliceStarts[rwDim] - 1) &&
            (otherLimits[rwDim] == rwSliceLimits[rwDim] - rwWindowSize);

        if (!appendRight && !prependLeft)
          return false;

        // Build the new extended slice
        SmallVector<int64_t> newStarts(rwSliceStarts.begin(),
                                       rwSliceStarts.end());
        SmallVector<int64_t> newLimits(rwSliceLimits.begin(),
                                       rwSliceLimits.end());
        SmallVector<int64_t> newStrides(rwInputSlice.getStrides().begin(),
                                        rwInputSlice.getStrides().end());
        if (appendRight) {
          newLimits[rwDim] = rwSliceLimits[rwDim] + 1;
        } else {
          newStarts[rwDim] = rwSliceStarts[rwDim] - 1;
        }

        auto newSlice = stablehlo::SliceOp::create(
            rewriter, op.getLoc(), rwInputSlice.getOperand(), newStarts,
            newLimits, newStrides);

        // Build extended window dimensions
        SmallVector<int64_t> newWindowDims(rwWindowDims.begin(),
                                           rwWindowDims.end());
        newWindowDims[rwDim] = rwWindowSize + 1;

        // Reuse the init value from the original reduce_window
        auto newReduceWindowOp =
            rewriter.replaceOpWithNewOp<stablehlo::ReduceWindowOp>(
                op, TypeRange(outType), ValueRange(newSlice.getResult()),
                rwOp.getInitValues(),
                rewriter.getDenseI64ArrayAttr(newWindowDims),
                rwOp.getWindowStridesAttr(), rwOp.getBaseDilationsAttr(),
                rwOp.getWindowDilationsAttr(), rwOp.getPaddingAttr());

        newReduceWindowOp.getBody().takeBody(rwOp.getBody());
        return true;
      };

      // Try both orderings: lhs=reduce_window, rhs=slice and vice versa
      if (tryMatchReduceWindowSlice(lhsInfo.base, rhsInfo.base) ||
          tryMatchReduceWindowSlice(rhsInfo.base, lhsInfo.base)) {
        return success();
      }
    }

    if (!newBaseOp || windowDilation < 0 || dimension < 0)
      return rewriter.notifyMatchFailure(
          op, "None of the parent op conditions match.");

    auto loc = op.getLoc();
    SmallVector<int64_t> windowDilations(outRank, 1);
    windowDilations[dimension] = windowDilation;

    if (canEmitReduceWindow) {
      SmallVector<int64_t> windowDims(outRank, 1);
      windowDims[dimension] = 2;

      int64_t initScalar = 0;
      if constexpr (std::is_same_v<OpTy, stablehlo::MulOp>) {
        initScalar = 1;
      }

      auto scalarType = RankedTensorType::get({}, outType.getElementType());
      auto initVal = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), scalarType,
          cast<ElementsAttr>(makeAttr(scalarType, initScalar)));

      auto reduceWindowOp =
          rewriter.replaceOpWithNewOp<stablehlo::ReduceWindowOp>(
              op, TypeRange(outType), ValueRange(newBaseOp->getResult(0)),
              ValueRange(initVal), rewriter.getDenseI64ArrayAttr(windowDims),
              DenseI64ArrayAttr(), DenseI64ArrayAttr(),
              llvm::all_of(windowDilations, [](int64_t i) { return i == 1; })
                  ? DenseI64ArrayAttr()
                  : rewriter.getDenseI64ArrayAttr(windowDilations),
              DenseIntElementsAttr());

      {
        auto *block = rewriter.createBlock(&reduceWindowOp.getBody());
        auto argType = RankedTensorType::get({}, outType.getElementType());
        block->addArgument(argType, loc);
        block->addArgument(argType, loc);
        rewriter.setInsertionPointToStart(block);

        stablehlo::ReturnOp::create(
            rewriter, loc,
            ValueRange(OpTy::create(rewriter, loc, block->getArgument(0),
                                    block->getArgument(1))));
      }

      return success();
    }

    // We need to reshape the base op to have 2 additional dimensions in the
    // beginning
    auto lhsBaseType =
        cast<RankedTensorType>(newBaseOp->getResult(0).getType());
    auto lhsBaseShape = lhsBaseType.getShape();

    SmallVector<int64_t> spatialDims(outRank, 1);
    for (int64_t i = 0; i < outRank; ++i)
      spatialDims[i] = i + 2;

    SmallVector<int64_t> newShape(2, 1);
    SmallVector<int64_t> newOutShape(2, 1);
    for (int64_t i = 0; i < outRank; ++i) {
      newShape.push_back(lhsBaseShape[i]);
      newOutShape.push_back(outShape[i]);
    }
    auto convOutType =
        RankedTensorType::get(newOutShape, outType.getElementType());

    auto convLhs = stablehlo::ReshapeOpCreate(
        rewriter, loc, newBaseOp->getResult(0), newShape);

    // construct the conv rhs
    SmallVector<int64_t> filterShape(outRank + 2, 1);

    // filter value
    Value firstElement, secondElement, lhsElement, rhsElement;

    if (lhsInfo.constantAttr.has_value()) {
      lhsElement = stablehlo::ConstantOp::create(
          rewriter, loc, RankedTensorType::get({}, outType.getElementType()),
          lhsInfo.constantAttr.value());
    } else if (lhsInfo.constantScalar.has_value()) {
      lhsElement = lhsInfo.constantScalar.value();
    } else {
      lhsElement = stablehlo::ConstantOp::create(
          rewriter, loc, RankedTensorType::get({}, outType.getElementType()),
          cast<ElementsAttr>(makeAttr(
              RankedTensorType::get({}, outType.getElementType()), 1)));
    }

    if (rhsInfo.constantAttr.has_value()) {
      rhsElement = stablehlo::ConstantOp::create(
          rewriter, loc, RankedTensorType::get({}, outType.getElementType()),
          rhsInfo.constantAttr.value());
    } else if (rhsInfo.constantScalar.has_value()) {
      rhsElement = rhsInfo.constantScalar.value();
    } else {
      rhsElement = stablehlo::ConstantOp::create(
          rewriter, loc, RankedTensorType::get({}, outType.getElementType()),
          cast<ElementsAttr>(makeAttr(
              RankedTensorType::get({}, outType.getElementType()), 1)));
    }

    if constexpr (std::is_base_of_v<stablehlo::SubtractOp, OpTy>) {
      rhsElement = stablehlo::NegOp::create(rewriter, loc, rhsElement);
    }

    if (flipped) {
      firstElement = rhsElement;
      secondElement = lhsElement;
    } else {
      firstElement = lhsElement;
      secondElement = rhsElement;
    }

    firstElement =
        stablehlo::ReshapeOpCreate(rewriter, loc, firstElement, filterShape);
    secondElement =
        stablehlo::ReshapeOpCreate(rewriter, loc, secondElement, filterShape);

    auto filter = stablehlo::ConcatenateOp::create(
        rewriter, loc, ValueRange{firstElement, secondElement},
        rewriter.getI64IntegerAttr(dimension + 2));

    auto convDims = stablehlo::ConvDimensionNumbersAttr::get(
        rewriter.getContext(),
        /*input_batch_dimension=*/0,
        /*input_feature_dimension=*/1,
        /*input_spatial_dimensions=*/spatialDims,
        /*kernel_input_feature_dimension=*/0,
        /*kernel_output_feature_dimension=*/1,
        /*kernel_spatial_dimensions=*/spatialDims,
        /*output_batch_dimension=*/0,
        /*output_feature_dimension=*/1,
        /*output_spatial_dimensions=*/spatialDims);

    // we do the additional checks here for pretty printing
    auto conv = stablehlo::ConvolutionOp::create(
        rewriter, loc, convOutType, convLhs, filter,
        /*window_strides=*/nullptr,
        /*padding=*/nullptr,
        /*lhs_dilation=*/nullptr,
        /*rhs_dilation=*/
        llvm::all_of(windowDilations, [](int64_t i) { return i == 1; })
            ? nullptr
            : rewriter.getDenseI64ArrayAttr(windowDilations),
        /*window_reversal=*/nullptr,
        /*conv_dimension_numbers=*/convDims,
        /*feature_group_count=*/rewriter.getI64IntegerAttr(1),
        /*batch_group_count=*/rewriter.getI64IntegerAttr(1),
        /*precision_config=*/nullptr);

    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(), conv);
    return success();
  }

private:
  struct OperandInfo {
    Value base;
    std::optional<ElementsAttr> constantAttr;
    std::optional<Value> constantScalar;
  };

  std::tuple<std::optional<ElementsAttr>, std::optional<Value>>
  checkConstant(Value val) const {
    if (auto bcastInDimOp = val.getDefiningOp<stablehlo::BroadcastInDimOp>()) {
      auto operand = bcastInDimOp.getOperand();
      auto type = cast<RankedTensorType>(operand.getType());
      if (type) {
        if (type.getRank() == 0) {
          return std::make_tuple(std::nullopt, operand);
        }
      }
    }

    if (auto constOp = val.getDefiningOp<stablehlo::ConstantOp>()) {
      auto attr = constOp.getValue();
      if (attr.isSplat()) {
        return std::make_tuple(
            cast<ElementsAttr>(
                cast<SplatElementsAttr>(attr).resizeSplat(RankedTensorType::get(
                    {},
                    cast<RankedTensorType>(val.getType()).getElementType()))),
            std::nullopt);
      }
    }
    return std::make_tuple(std::nullopt, std::nullopt);
  }

  OperandInfo extractOperandInfo(Value val) const {
    if (auto mulOp = val.getDefiningOp<stablehlo::MulOp>()) {
      {
        auto [constantAttr, constantScalar] = checkConstant(mulOp.getLhs());
        if (constantAttr.has_value() || constantScalar.has_value()) {
          return {mulOp.getRhs(), constantAttr, constantScalar};
        }
      }

      {
        auto [constantAttr, constantScalar] = checkConstant(mulOp.getRhs());
        if (constantAttr.has_value() || constantScalar.has_value()) {
          return {mulOp.getLhs(), constantAttr, constantScalar};
        }
      }
    }
    return {val, std::nullopt};
  }

  bool allowEmitConvolution;
};

// Applies distributive property to op(dot_general(a, b), dot_general(a, c))
// operations and converts them to dot_general(a, op(b, c)) or
// dot_general(op(b, c), a)
template <typename OpTy>
LogicalResult
dotGeneralDistributiveSimplify(OpTy op, stablehlo::DotGeneralOp lhsDotGeneralOp,
                               stablehlo::DotGeneralOp rhsDotGeneralOp,
                               PatternRewriter &rewriter) {
  // TODO: generalize this to handle cases where the dot general dim numbers are
  // different but the ops can still be distributed
  if (!OperationEquivalence::isEquivalentTo(
          lhsDotGeneralOp, rhsDotGeneralOp,
          OperationEquivalence::ignoreValueEquivalence, nullptr,
          OperationEquivalence::IgnoreLocations, nullptr))
    return failure();

  auto lhsDotGeneralLhs = lhsDotGeneralOp.getLhs();
  auto rhsDotGeneralLhs = rhsDotGeneralOp.getLhs();

  if (lhsDotGeneralLhs == rhsDotGeneralLhs) {
    auto newDotRhs =
        OpTy::create(rewriter, op.getLoc(), lhsDotGeneralOp.getRhs(),
                     rhsDotGeneralOp.getRhs());
    rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
        op, TypeRange{op.getType()},
        ValueRange{lhsDotGeneralLhs, newDotRhs.getResult()},
        lhsDotGeneralOp->getAttrs());
    return success();
  }

  auto lhsDotGeneralRhs = lhsDotGeneralOp.getRhs();
  auto rhsDotGeneralRhs = rhsDotGeneralOp.getRhs();

  if (lhsDotGeneralRhs == rhsDotGeneralRhs) {
    auto newDotLhs =
        OpTy::create(rewriter, op.getLoc(), lhsDotGeneralOp.getLhs(),
                     rhsDotGeneralOp.getLhs());
    rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
        op, TypeRange{op.getType()},
        ValueRange{newDotLhs.getResult(), rhsDotGeneralRhs},
        lhsDotGeneralOp->getAttrs());
    return success();
  }

  return failure();
}

template <typename OpTy>
LogicalResult dotGeneralDistributiveSimplify(OpTy op,
                                             PatternRewriter &rewriter) {
  if (op.getNumOperands() != 2)
    return failure();

  auto lhs = op.getLhs(), rhs = op.getRhs();
  auto lhsDotGeneralOp = lhs.template getDefiningOp<stablehlo::DotGeneralOp>();
  auto rhsDotGeneralOp = rhs.template getDefiningOp<stablehlo::DotGeneralOp>();

  if (lhsDotGeneralOp && !isOnlyUsedInOperation(lhsDotGeneralOp, op))
    return failure();
  if (rhsDotGeneralOp && !isOnlyUsedInOperation(rhsDotGeneralOp, op))
    return failure();

  if (lhsDotGeneralOp == rhsDotGeneralOp) {
    return failure(); // only a single op
  }

  if (lhsDotGeneralOp && rhsDotGeneralOp)
    return dotGeneralDistributiveSimplify<OpTy>(op, lhsDotGeneralOp,
                                                rhsDotGeneralOp, rewriter);

  // FIXME: This is currently incorrect. We can't replace with a simple
  //        elementwise, instead we need to use scatter_op to perform
  //        this operation (on the diagonal indices of Y)
  // Value newLhs, newRhs;
  // stablehlo::DotGeneralOp dotGeneralOp;

  // auto constructOtherOp = [&](Value other, Value scalar, bool scalarIsLhs) {
  //   auto bcastedScalar = stablehlo::BroadcastInDimOp::create(
  //       rewriter, op.getLoc(), other.getType(), scalar,
  //       rewriter.getDenseI64ArrayAttr({}));
  //   auto lhs = scalarIsLhs ? bcastedScalar : other;
  //   auto rhs = scalarIsLhs ? other : bcastedScalar;
  //   return OpTy::create(rewriter, op.getLoc(), lhs, rhs);
  // };

  // if (lhsDotGeneralOp) {
  //   Value scalar, other;
  //   dotGeneralOp = lhsDotGeneralOp;
  //   if (extractMultiplicationFactor(rhs, other, op, rewriter)) {
  //     if (other == lhsDotGeneralOp.getLhs()) {
  //       extractMultiplicationFactor(rhs, scalar, other, op, rewriter);

  //       // op(X × Y, a * X) -> X × op(Y, a)
  //       newLhs = lhsDotGeneralOp.getLhs();
  //       newRhs = constructOtherOp(lhsDotGeneralOp.getRhs(), scalar, false);
  //     } else if (other == lhsDotGeneralOp.getRhs()) {
  //       extractMultiplicationFactor(rhs, scalar, other, op, rewriter);

  //       // op(Y × X, a * X) -> op(Y, a) × X
  //       newLhs = constructOtherOp(lhsDotGeneralOp.getLhs(), scalar, false);
  //       newRhs = lhsDotGeneralOp.getRhs();
  //     }
  //   }
  // }

  // if (rhsDotGeneralOp) {
  //   Value scalar, other;
  //   dotGeneralOp = rhsDotGeneralOp;
  //   if (extractMultiplicationFactor(lhs, other, op, rewriter)) {
  //     if (other == rhsDotGeneralOp.getLhs()) {
  //       extractMultiplicationFactor(lhs, scalar, other, op, rewriter);

  //       // op(a * X, X × Y) -> X × op(a, Y)
  //       newLhs = rhsDotGeneralOp.getLhs();
  //       newRhs = constructOtherOp(rhsDotGeneralOp.getRhs(), scalar, true);
  //     } else if (other == rhsDotGeneralOp.getRhs()) {
  //       extractMultiplicationFactor(lhs, scalar, other, op, rewriter);

  //       // op(a * X, Y × X) -> op(a, Y) × X
  //       newLhs = constructOtherOp(rhsDotGeneralOp.getLhs(), scalar, true);
  //       newRhs = rhsDotGeneralOp.getRhs();
  //     }
  //   }
  // }

  // if (newLhs && newRhs) {
  //   rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
  //       op, TypeRange{op.getType()}, ValueRange{newLhs, newRhs},
  //       dotGeneralOp->getAttrs());
  //   return success();
  // }

  return failure();
}

template <typename OpTy>
struct DotGeneralDistributiveSimplify
    : public CheckedOpRewritePattern<OpTy,
                                     DotGeneralDistributiveSimplify<OpTy>> {
  using CheckedOpRewritePattern<
      OpTy, DotGeneralDistributiveSimplify<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    return dotGeneralDistributiveSimplify<OpTy>(op, rewriter);
  }
};

struct TrivialReduceWindowToReduceOp
    : public CheckedOpRewritePattern<stablehlo::ReduceWindowOp,
                                     TrivialReduceWindowToReduceOp> {
  using CheckedOpRewritePattern<
      stablehlo::ReduceWindowOp,
      TrivialReduceWindowToReduceOp>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceWindowOp op,
                                    PatternRewriter &rewriter) const {
    auto operandType = cast<RankedTensorType>(op.getOperand(0).getType());
    if (!operandType)
      return failure();

    ArrayRef<int64_t> operandShape = operandType.getShape();
    int64_t rank = operandShape.size();

    // window dims should either be 1 or have size same as that dim size
    auto windowDims = llvm::to_vector(op.getWindowDimensions());
    if (windowDims.size() != rank)
      return failure();

    SmallVector<int64_t> reductionDims;
    SetVector<int64_t> reductionDimsSet;
    for (int64_t i = 0; i < rank; i++) {
      int64_t windowDim = windowDims[i];
      if (windowDim == 1) {
        continue;
      } else if (windowDim == operandShape[i]) {
        reductionDims.push_back(i);
        reductionDimsSet.insert(i);
      } else {
        return failure();
      }
    }

    if (reductionDims.empty())
      return failure();

    auto windowStrides = op.getWindowStrides();
    if (windowStrides.has_value()) {
      for (auto stride : windowStrides.value()) {
        if (stride != 1)
          return failure();
      }
    }

    auto baseDilations = op.getBaseDilations();
    if (baseDilations.has_value()) {
      for (auto dilation : baseDilations.value()) {
        if (dilation != 1)
          return failure();
      }
    }

    auto windowDilations = op.getWindowDilations();
    if (windowDilations.has_value()) {
      for (auto dilation : windowDilations.value()) {
        if (dilation != 1)
          return failure();
      }
    }

    auto padding = op.getPadding();
    if (padding.has_value()) {
      for (auto pad : padding.value()) {
        if (pad != 0)
          return failure();
      }
    }

    SmallVector<Type> resultTypes;
    for (int64_t i = 0; i < op.getNumResults(); i++) {
      auto origType = cast<RankedTensorType>(op.getType(i));
      SmallVector<int64_t> newShape;
      auto oldShape = origType.getShape();
      for (int64_t j = 0; j < origType.getRank(); j++) {
        if (!reductionDimsSet.contains(j))
          newShape.push_back(oldShape[j]);
      }
      resultTypes.push_back(
          RankedTensorType::get(newShape, origType.getElementType()));
    }

    auto newReduceOp = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), TypeRange{resultTypes}, op.getInputs(),
        op.getInitValues(), rewriter.getDenseI64ArrayAttr(reductionDims));
    newReduceOp.getRegion().takeBody(op.getRegion());

    for (int64_t i = 0; i < op.getNumResults(); i++) {
      auto newReshape = stablehlo::ReshapeOpCreate(
          rewriter, op.getLoc(), newReduceOp.getResults()[i],
          cast<RankedTensorType>(op.getType(i)).getShape());
      rewriter.replaceAllUsesWith(op.getResult(i), newReshape);
    }
    return success();
  }
};

// Returns true if `val` is the identity element for the given `kind` of reduce
// operation, i.e. reduce_body(val, x) == x for all x.
static bool isIdentityValueForReduceOp(Value val, ReduceOpKind kind) {
  Attribute attr;
  if (!matchPattern(val, m_Constant(&attr)))
    return false;
  switch (kind) {
  case ReduceOpKind::Add:
  case ReduceOpKind::Or:
  case ReduceOpKind::Xor:
    // Both integer zero and float +0.0 / -0.0 are identity for add/or/xor.
    return matchPattern(attr, m_AnyZeroFloat()) || matchPattern(attr, m_Zero());
  case ReduceOpKind::Mul:
    return matchPattern(attr, m_One()) || matchPattern(attr, m_OneFloat());
  case ReduceOpKind::And:
    return matchPattern(attr, m_AllOnes());
  case ReduceOpKind::Max:
    return matchPattern(attr, m_NegInfFloat());
  case ReduceOpKind::Min:
    return matchPattern(attr, m_PosInfFloat());
  default:
    return false;
  }
}

// Eliminate reduce_window ops that are no-ops: all window dimensions are 1,
// there is no padding, all strides and dilations are 1, and the init value is
// the identity element for the reduction body.  In that case every output
// element equals the corresponding input element and the op can be replaced
// by the input.
struct NoopReduceWindowOpCanon final
    : CheckedOpRewritePattern<stablehlo::ReduceWindowOp,
                              NoopReduceWindowOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceWindowOp op,
                                    PatternRewriter &rewriter) const {
    // Only handle single-input / single-init ops for now.
    if (op.getInputs().size() != 1 || op.getInitValues().size() != 1)
      return failure();

    auto inputType = dyn_cast<RankedTensorType>(op.getInputs()[0].getType());
    if (!inputType || !inputType.hasStaticShape())
      return failure();

    auto resultType = dyn_cast<RankedTensorType>(op.getType(0));
    if (!resultType || !resultType.hasStaticShape())
      return failure();

    // The result must have the same shape as the input.
    if (inputType != resultType)
      return failure();

    // All window dimensions must be 1.
    auto windowDims = op.getWindowDimensions();
    if (!llvm::all_of(windowDims, [](int64_t d) { return d == 1; }))
      return failure();

    // Padding must be absent or all-zero.
    auto padding = op.getPadding();
    if (padding && !llvm::all_of(padding->getValues<int64_t>(),
                                 [](int64_t p) { return p == 0; }))
      return failure();

    // Window strides must be absent or all-one.
    auto strides = op.getWindowStrides();
    if (strides && !llvm::all_of(*strides, [](int64_t s) { return s == 1; }))
      return failure();

    // Base dilations must be absent or all-one.
    auto baseDilations = op.getBaseDilations();
    if (baseDilations &&
        !llvm::all_of(*baseDilations, [](int64_t d) { return d == 1; }))
      return failure();

    // Window dilations must be absent or all-one (irrelevant for window=1 but
    // check anyway to keep the semantics clear).
    auto windowDilations = op.getWindowDilations();
    if (windowDilations &&
        !llvm::all_of(*windowDilations, [](int64_t d) { return d == 1; }))
      return failure();

    // Determine the reduction kind from the body.
    auto commonRW = CheckCommonReduceWindowOp(op);
    if (commonRW.kind == ReduceOpKind::Unknown)
      return failure();

    // The init value must be the identity element for the reduction so that
    // reduce_body(init, x) == x.
    if (!isIdentityValueForReduceOp(op.getInitValues()[0], commonRW.kind))
      return failure();

    // All checks passed – replace the reduce_window with its input.
    rewriter.replaceOp(op, op.getInputs()[0]);
    return success();
  }
};

template <typename BinaryOpType, typename Child>
struct ReduceSliceFusionBase
    : public CheckedOpRewritePattern<
          BinaryOpType, ReduceSliceFusionBase<BinaryOpType, Child>> {
  using CheckedOpRewritePattern<
      BinaryOpType,
      ReduceSliceFusionBase<BinaryOpType, Child>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(BinaryOpType binaryOp,
                                    PatternRewriter &rewriter) const {
    SmallVector<SliceInfo, 4> slices;
    SmallVector<Value> extraValues;

    if (!collectSlicesInChain(binaryOp, slices, extraValues)) {
      return failure();
    }

    int64_t sliceDim = extractSliceDim(slices);
    if (sliceDim == -1) {
      return failure();
    }

    return createFusedReduce(rewriter, binaryOp, slices, extraValues, sliceDim);
  };

private:
  struct SliceInfo {
    stablehlo::SliceOp sliceOp;
    llvm::SmallSetVector<int64_t, 4> sliceDims;
    Value initValue;
    Value source; // needed for partial fusion
  };

  int64_t extractSliceDim(SmallVector<SliceInfo, 4> &slices) const {
    int64_t sliceDim = -1;
    auto sliceDims = findIntersection(slices);

    for (auto dim : sliceDims) {
      // Sort slices by start index
      llvm::sort(slices, [&](SliceInfo &a, SliceInfo &b) {
        return a.sliceOp.getStartIndices()[dim] <
               b.sliceOp.getStartIndices()[dim];
      });

      // overlapping slices are optimized with a different fusion
      bool overlapping = false;
      for (int i = 0; i < slices.size() - 1; i++) {
        auto info1 = slices[i];
        auto info2 = slices[i + 1];
        if (info2.sliceOp.getStartIndices()[dim] <
            info1.sliceOp.getLimitIndices()[dim]) {
          overlapping = true;
          break;
        }
      }

      if (overlapping) {
        continue;
      }

      sliceDim = dim;
      break;
    }

    return sliceDim;
  }

  llvm::SmallSetVector<int64_t, 4>
  findIntersection(SmallVector<SliceInfo, 4> &slices) const {
    if (slices.empty()) {
      return {};
    }

    auto intersection = slices.front().sliceDims;
    for (size_t i = 1; i < slices.size(); i++) {
      llvm::set_intersect(intersection, slices[i].sliceDims);
    }
    return intersection;
  }

  std::tuple<bool, SliceInfo>
  matchReduceSlice(stablehlo::ReduceOp reduceOp) const {
    if (reduceOp.getDimensions().size() == 0)
      return {false, SliceInfo()};

    if (reduceOp.getInputs().size() != 1 ||
        reduceOp.getDimensions().size() != 1 ||
        !((Child *)this)->isCompatibleReduction(reduceOp)) {
      return {false, SliceInfo()};
    }

    auto slice =
        reduceOp.getInputs()[0].template getDefiningOp<stablehlo::SliceOp>();
    if (!slice) {
      return {false, SliceInfo()};
    }

    SmallVector<int64_t> reductionDims = {reduceOp.getDimensions()[0]};
    SliceInfo info = extractSliceInfo(slice, reductionDims);
    info.initValue = reduceOp.getInitValues()[0];
    info.source = reduceOp.getResults()[0];
    return {true, info};
  }

  std::tuple<bool, SliceInfo>
  matchReshapeReduceSlice(stablehlo::ReshapeOp reshapeOp) const {
    auto reduce =
        reshapeOp.getOperand().template getDefiningOp<stablehlo::ReduceOp>();
    if (!reduce || reduce.getDimensions().empty()) {
      return {false, SliceInfo()};
    }

    auto inputType = cast<RankedTensorType>(reshapeOp.getOperand().getType());
    auto outputType = cast<RankedTensorType>(reshapeOp.getType());
    auto reduceDim = reduce.getDimensions()[0];
    if (!areValidInsertionDims(inputType, outputType, {reduceDim})) {
      return {false, SliceInfo()};
    }

    auto [match, info] = matchReduceSlice(reduce);
    info.source = reshapeOp.getResult();
    return {match, info};
  }

  std::tuple<bool, SliceInfo>
  matchReshapeSlice(stablehlo::ReshapeOp reshapeOp) const {
    auto slice =
        reshapeOp.getOperand().template getDefiningOp<stablehlo::SliceOp>();
    if (!slice) {
      return {false, SliceInfo()};
    }

    SliceInfo info = extractSliceInfo(slice);
    if (info.sliceOp) {
      info.source = reshapeOp.getResult();
      return {true, info};
    }
    return {false, info};
  }

  bool matchingSourceOperand(SmallVectorImpl<SliceInfo> &slices,
                             stablehlo::SliceOp slice) const {
    if (!slice) {
      return false;
    }
    if (slices.size() == 0) {
      return true;
    }
    return slices[0].sliceOp.getOperand() == slice.getOperand();
  }

  // Collect all slices in the binary operation chain
  bool collectSlicesInChain(BinaryOpType startOp,
                            SmallVectorImpl<SliceInfo> &slices,
                            SmallVectorImpl<Value> &extraValues) const {
    // Use a worklist to traverse the binary operation chain
    SmallVector<Value> worklist;
    DenseSet<Value> visited;

    worklist.push_back(startOp.getResult());

    while (!worklist.empty()) {
      Value current = worklist.pop_back_val();

      if (visited.contains(current)) {
        if constexpr (!std::is_same_v<BinaryOpType, stablehlo::MinOp> &&
                      !std::is_same_v<BinaryOpType, stablehlo::MaxOp>)
          extraValues.push_back(current);
        continue;
      }
      visited.insert(current);

      if (auto binaryOp = current.template getDefiningOp<BinaryOpType>()) {
        worklist.push_back(binaryOp.getLhs());
        worklist.push_back(binaryOp.getRhs());
      } else if (auto reshape =
                     current.template getDefiningOp<stablehlo::ReshapeOp>()) {
        // slice -> reduce -> single insert dim along reduction dim -> ...
        auto [isMatchRRS, infoRRS] = matchReshapeReduceSlice(reshape);
        if (isMatchRRS && matchingSourceOperand(slices, infoRRS.sliceOp)) {
          slices.push_back(infoRRS);
        } else {
          // slice -> drop dim -> ...
          auto [isMatchRS, infoRS] = matchReshapeSlice(reshape);
          if (isMatchRS && matchingSourceOperand(slices, infoRS.sliceOp)) {
            slices.push_back(infoRS);
          } else {
            extraValues.push_back(current);
            continue;
          }
        }
      } else if (auto slice =
                     current.template getDefiningOp<stablehlo::SliceOp>()) {
        if (matchingSourceOperand(slices, slice)) {
          SliceInfo info = extractSliceInfo(slice);
          if (info.sliceOp) {
            slices.push_back(info);
          } else {
            extraValues.push_back(current);
            continue;
          }
        } else {
          extraValues.push_back(current);
          continue;
        }
      } else if (auto reduce =
                     current.template getDefiningOp<stablehlo::ReduceOp>()) {
        auto [isMatchRR, infoRR] = matchReduceSlice(reduce);
        if (isMatchRR && matchingSourceOperand(slices, infoRR.sliceOp)) {
          slices.push_back(infoRR);
        } else {
          extraValues.push_back(current);
          continue;
        }
      } else {
        extraValues.push_back(current);
        continue;
      }
    }

    return slices.size() > 1;
  }

  // Extract slice information
  SliceInfo extractSliceInfo(stablehlo::SliceOp slice) const {
    SliceInfo info{slice};
    auto outputType = cast<RankedTensorType>(slice.getType());
    for (size_t i = 0; i < outputType.getRank(); ++i) {
      if (outputType.getDimSize(i) == 1) {
        info.sliceDims.insert(i);
      }
    }
    info.source = slice.getResult();
    return info;
  }

  SliceInfo extractSliceInfo(stablehlo::SliceOp slice,
                             SmallVectorImpl<int64_t> &potentialDims) const {
    SliceInfo info{slice};
    for (auto dim : potentialDims)
      info.sliceDims.insert(dim);
    return info;
  }

  // Create the fused reduce operation
  LogicalResult createFusedReduce(PatternRewriter &rewriter,
                                  BinaryOpType binaryOp,
                                  SmallVectorImpl<SliceInfo> &slices,
                                  SmallVectorImpl<Value> &extraValues,
                                  int64_t sliceDim) const {
    Value initValue;
    for (auto sliceInfo : slices) {
      if (sliceInfo.initValue) {
        if (!initValue) {
          initValue = sliceInfo.initValue;
        } else {
          if (sliceInfo.initValue != initValue) {
            return rewriter.notifyMatchFailure(binaryOp,
                                               "init values are not the same");
          }
        }
      }
    }

    Value sourceOperand = slices[0].sliceOp.getOperand();

    {
      SmallVector<Value> newConcatInputs;
      for (auto sliceInfo : slices) {
        newConcatInputs.push_back(sliceInfo.sliceOp);
      }
      SmallVector<Value> finalConcatInputs;
      auto result = concatSliceSimplify(rewriter, newConcatInputs, sliceDim,
                                        finalConcatInputs);
      if (!result.succeeded()) {
        return result;
      }
      if (finalConcatInputs.size() == 1) {
        sourceOperand = finalConcatInputs[0];
      } else {
        sourceOperand = stablehlo::ConcatenateOp::create(
            rewriter, binaryOp.getLoc(), finalConcatInputs, sliceDim);
      }
    }

    auto elemType = cast<ShapedType>(binaryOp.getType()).getElementType();
    if (!initValue) {
      initValue = stablehlo::getIdentityValueForOp<BinaryOpType>(
          rewriter, binaryOp.getLoc(), elemType);
      if (!initValue) {
        return rewriter.notifyMatchFailure(
            binaryOp, "could not find identity value for element type");
      }
    }

    SmallVector<int64_t> newShape = llvm::to_vector(
        cast<RankedTensorType>(sourceOperand.getType()).getShape());
    newShape.erase(newShape.begin() + sliceDim);

    auto newReduce = stablehlo::ReduceOp::create(
        rewriter, binaryOp.getLoc(),
        TypeRange{RankedTensorType::get(newShape, elemType)},
        ValueRange{sourceOperand}, ValueRange{initValue},
        ArrayRef<int64_t>{sliceDim});

    auto scalarType = RankedTensorType::get({}, elemType);
    Block *block = rewriter.createBlock(&newReduce.getBody());
    block->addArgument(scalarType, binaryOp.getLoc());
    block->addArgument(scalarType, binaryOp.getLoc());

    {
      IRRewriter::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(block);
      auto elemOp =
          BinaryOpType::create(rewriter, binaryOp.getLoc(),
                               block->getArgument(0), block->getArgument(1));
      stablehlo::ReturnOp::create(rewriter, binaryOp.getLoc(),
                                  elemOp.getResult());
    }

    rewriter.setInsertionPointAfter(newReduce);
    auto binaryResultType =
        cast<RankedTensorType>(binaryOp.getResult().getType());
    Value result = stablehlo::ReshapeOpCreate(rewriter, binaryOp.getLoc(),
                                              newReduce.getResult(0),
                                              binaryResultType.getShape());

    for (auto &value : extraValues) {
      result = BinaryOpType::create(rewriter, binaryOp.getLoc(), result, value);
    }
    rewriter.replaceAllUsesWith(binaryOp.getResult(), result);
    return success();
  }
};

// Macro to generate ReduceSliceFusion structs for common binary ops.
#define REDUCESLICEFUSIONSPECIALIZE(OP)                                        \
  struct OP##ReduceSliceFusion                                                 \
      : public ReduceSliceFusionBase<stablehlo::OP##Op,                        \
                                     OP##ReduceSliceFusion> {                  \
    using ReduceSliceFusionBase<stablehlo::OP##Op,                             \
                                OP##ReduceSliceFusion>::ReduceSliceFusionBase; \
                                                                               \
    bool isCompatibleReduction(stablehlo::ReduceOp reduceOp) const {           \
      return mlir::stablehlo::CheckCommonReduceOp(reduceOp).kind ==            \
             mlir::stablehlo::ReduceOpKind::OP;                                \
    }                                                                          \
  };

REDUCESLICEFUSIONSPECIALIZE(Add)
REDUCESLICEFUSIONSPECIALIZE(Mul)
REDUCESLICEFUSIONSPECIALIZE(Min)
REDUCESLICEFUSIONSPECIALIZE(Max)
REDUCESLICEFUSIONSPECIALIZE(And)
REDUCESLICEFUSIONSPECIALIZE(Or)
REDUCESLICEFUSIONSPECIALIZE(Xor)

#undef REDUCESLICEFUSIONSPECIALIZE

struct CaseToIf : public CheckedOpRewritePattern<stablehlo::CaseOp, CaseToIf> {
  using CheckedOpRewritePattern<stablehlo::CaseOp,
                                CaseToIf>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::CaseOp caseOp,
                                    PatternRewriter &rewriter) const {
    if (caseOp.getBranches().size() != 2)
      return rewriter.notifyMatchFailure(caseOp,
                                         "case op does not have 2 branches");

    auto index = caseOp.getIndex();
    auto convertOp = index.getDefiningOp<stablehlo::ConvertOp>();
    if (!convertOp)
      return rewriter.notifyMatchFailure(caseOp, "index is not a convert op");
    auto boolValue = convertOp.getOperand();
    auto boolType = cast<RankedTensorType>(boolValue.getType());
    if (!boolType || !boolType.getElementType().isInteger(1) ||
        boolType.getRank() != 0)
      return rewriter.notifyMatchFailure(caseOp,
                                         "index is not a boolean tensor");

    auto ifOp = stablehlo::IfOp::create(rewriter, caseOp.getLoc(),
                                        caseOp.getResultTypes(), boolValue);

    // true -> 1 and false -> 0
    ifOp.getTrueBranch().takeBody(caseOp.getBranches()[1]);
    ifOp.getFalseBranch().takeBody(caseOp.getBranches()[0]);
    rewriter.replaceOp(caseOp, ifOp.getResults());

    return success();
  }
};

struct DUSToDynamicPad
    : public CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                                     DUSToDynamicPad> {
  using CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp,
                                DUSToDynamicPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto operand = op.getOperand();
    auto update = op.getUpdate();
    auto indices = op.getStartIndices();

    for (auto [i, index] : llvm::enumerate(indices)) {
      if (!matchPattern(index, m_Constant())) {
        return rewriter.notifyMatchFailure(
            op, "not all indices are constant. currently we don't support this "
                "case");
      }
    }

    Value scalarOperand = getScalarPadValue(rewriter, operand);
    if (!scalarOperand)
      return rewriter.notifyMatchFailure(op, "operand is not a scalar pad");

    auto updateShape = cast<RankedTensorType>(update.getType()).getShape();
    auto operandShape = cast<RankedTensorType>(operand.getType()).getShape();

    SmallVector<Value> edgePaddingLowValues, edgePaddingHighValues;
    for (auto [i, index] : llvm::enumerate(indices)) {
      auto cType = RankedTensorType::get(
          {}, cast<RankedTensorType>(index.getType()).getElementType());
      auto lowerBound = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), cType, cast<ElementsAttr>(makeAttr(cType, 0)));
      auto upperBound = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), cType,
          cast<ElementsAttr>(
              makeAttr(cType, operandShape[i] - updateShape[i])));
      auto clampedIndex = stablehlo::ClampOp::create(
          rewriter, op.getLoc(), lowerBound, index, upperBound);

      auto reshapedIndex = stablehlo::ReshapeOpCreate(
          rewriter, op.getLoc(), clampedIndex, ArrayRef<int64_t>{1});
      edgePaddingLowValues.push_back(reshapedIndex);

      auto iType = RankedTensorType::get(
          {1}, cast<RankedTensorType>(index.getType()).getElementType());
      auto tmp = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), iType,
          cast<ElementsAttr>(
              makeAttr(iType, operandShape[i] - updateShape[i])));
      auto paddingHigh = stablehlo::SubtractOp::create(rewriter, op.getLoc(),
                                                       tmp, reshapedIndex);
      edgePaddingHighValues.push_back(paddingHigh);
    }

    auto edgePaddingLow = stablehlo::ConcatenateOp::create(
        rewriter, op.getLoc(), edgePaddingLowValues, 0);
    auto edgePaddingHigh = stablehlo::ConcatenateOp::create(
        rewriter, op.getLoc(), edgePaddingHighValues, 0);
    auto interiorPadding = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), edgePaddingLow.getType(),
        cast<ElementsAttr>(makeAttr(edgePaddingLow.getType(), 0)));

    rewriter.replaceOpWithNewOp<stablehlo::DynamicPadOp>(
        op, op.getType(), update, scalarOperand, edgePaddingLow,
        edgePaddingHigh, interiorPadding);
    return success();
  }

private:
  Value getScalarPadValue(PatternRewriter &rewriter, Value operand) const {
    Value scalarOperand = getScalarPadValueViaBcastInDim(rewriter, operand);
    if (scalarOperand)
      return scalarOperand;

    scalarOperand = getScalarPadValueViaSplattedConstant(rewriter, operand);
    if (scalarOperand)
      return scalarOperand;

    return nullptr;
  }

  Value getScalarPadValueViaBcastInDim(PatternRewriter &rewriter,
                                       Value operand) const {
    auto bcastInDimOp = operand.getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcastInDimOp)
      return nullptr;

    auto bcastOperand = bcastInDimOp.getOperand();
    auto bcastOperandType = cast<RankedTensorType>(bcastOperand.getType());
    if (bcastOperandType.getRank() != 0)
      return nullptr;

    return bcastOperand;
  }

  Value getScalarPadValueViaSplattedConstant(PatternRewriter &rewriter,
                                             Value operand) const {
    SplatElementsAttr splatAttr;
    if (!matchPattern(operand, m_Constant(&splatAttr)))
      return nullptr;

    return stablehlo::ConstantOp::create(
        rewriter, operand.getLoc(),
        splatAttr.resizeSplat(
            RankedTensorType::get({}, splatAttr.getType().getElementType())));
  }
};

struct DynamicPadToPad
    : public CheckedOpRewritePattern<stablehlo::DynamicPadOp, DynamicPadToPad> {
  using CheckedOpRewritePattern<stablehlo::DynamicPadOp,
                                DynamicPadToPad>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicPadOp op,
                                    PatternRewriter &rewriter) const {
    auto operand = op.getOperand();
    auto paddingValue = op.getPaddingValue();
    auto edgePaddingLow = op.getEdgePaddingLow();
    auto edgePaddingHigh = op.getEdgePaddingHigh();
    auto interiorPadding = op.getInteriorPadding();

    DenseIntElementsAttr edgePaddingLowAttr, edgePaddingHighAttr,
        interiorPaddingAttr;
    if (!matchPattern(edgePaddingLow, m_Constant(&edgePaddingLowAttr)) ||
        !matchPattern(edgePaddingHigh, m_Constant(&edgePaddingHighAttr)) ||
        !matchPattern(interiorPadding, m_Constant(&interiorPaddingAttr)))
      return rewriter.notifyMatchFailure(op, "edge padding is not a constant");

    rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
        op, op.getType(), operand, paddingValue,
        convertToDenseI64ArrayAttr(edgePaddingLowAttr),
        convertToDenseI64ArrayAttr(edgePaddingHighAttr),
        convertToDenseI64ArrayAttr(interiorPaddingAttr));
    return success();
  }

private:
  DenseI64ArrayAttr
  convertToDenseI64ArrayAttr(DenseIntElementsAttr attr) const {
    auto values = attr.getValues<APInt>();
    llvm::SmallVector<int64_t> denseValues;
    for (auto value : values)
      denseValues.push_back(value.getSExtValue());
    return DenseI64ArrayAttr::get(attr.getContext(), denseValues);
  }
};

struct RemoveNoOpsFromWhileLoop
    : public CheckedOpRewritePattern<stablehlo::WhileOp,
                                     RemoveNoOpsFromWhileLoop> {
  using CheckedOpRewritePattern<
      stablehlo::WhileOp, RemoveNoOpsFromWhileLoop>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp whileOp,
                                    PatternRewriter &rewriter) const {
    auto info = WhileLoopInfo(whileOp);
    auto computeInfoSuccess = info.computeInfo();
    if (computeInfoSuccess.failed()) {
      return computeInfoSuccess;
    }

    if (!info.isValid() || !info.isConstant() ||
        info.getConstantNumIters() <= 0) {
      return failure();
    }

    // Propagate bounds using WhileLoopInfo
    info.propagateBounds();

    auto &boundsMap = info.getBoundsMap();
    unsigned bitWidth = info.getBoundsBitWidth();

    bool anyOpRewritten = false;
    // Annotate the IR with bounds
    for (auto &[value, bounds] : boundsMap) {
      auto defOp = value.getDefiningOp();
      if (!defOp)
        continue;

      // Build bounds attribute as ArrayAttr for each result
      SmallVector<Attribute> boundsAttrs;
      for (unsigned i = 0; i < defOp->getNumResults(); ++i) {
        auto result = defOp->getResult(i);
        if (boundsMap.count(result)) {
          auto &resultBounds = boundsMap[result];
          boundsAttrs.push_back(ArrayAttr::get(
              value.getContext(),
              {IntegerAttr::get(rewriter.getIntegerType(bitWidth),
                                resultBounds.min),
               IntegerAttr::get(rewriter.getIntegerType(bitWidth),
                                resultBounds.max)}));
        } else {
          // Use empty array for results without bounds
          boundsAttrs.push_back(ArrayAttr::get(value.getContext(), {}));
        }
      }
      auto arattr = ArrayAttr::get(value.getContext(), boundsAttrs);
      // Only annotate ops we haven't already annotated
      if (!defOp->hasAttr("enzymexla.bounds") ||
          defOp->getAttr("enzymexla.bounds") != arattr) {
        rewriter.startOpModification(defOp);
        defOp->setAttr("enzymexla.bounds", arattr);
        anyOpRewritten = true;
        rewriter.finalizeOpModification(defOp);
      }
    }

    // Rewrite ops based on computed bounds
    llvm::SetVector<Operation *> toProcess;
    for (auto &[value, bounds] : boundsMap) {
      for (auto user : value.getUsers()) {
        toProcess.insert(user);
      }
      auto defOp = value.getDefiningOp();
      if (!defOp) {
        continue;
      }
      toProcess.insert(defOp);
    }

    for (auto op : toProcess) {
      bool rewritten =
          llvm::TypeSwitch<Operation *, bool>(op)
              .Case<stablehlo::RemOp, stablehlo::CompareOp, stablehlo::AbsOp,
                    stablehlo::ClampOp>([&](auto op) {
                auto allBounds = getBoundsOfAllOperands(op, info);
                if (!allBounds.has_value())
                  return false;
                return rewriteOperation(rewriter, op, allBounds.value());
              })
              .Default([&](Operation *op) { return false; });
      anyOpRewritten |= rewritten;
    }
    return anyOpRewritten ? success() : failure();
  }

private:
  using Bounds = WhileLoopInfo::Bounds;

  APInt min(APInt a, APInt b) const { return a.slt(b) ? a : b; }
  APInt max(APInt a, APInt b) const { return a.sgt(b) ? a : b; }
  APInt abs(APInt a) const { return a.sgt(0) ? a : -a; }

  std::optional<llvm::DenseMap<Value, Bounds>>
  getBoundsOfAllOperands(Operation *op, WhileLoopInfo &info) const {
    llvm::DenseMap<Value, Bounds> newBoundsMap;
    for (auto operand : op->getOperands()) {
      auto bounds = info.getBounds(operand);
      if (!bounds.has_value()) {
        return std::nullopt;
      }
      newBoundsMap[operand] = bounds.value();
    }
    return newBoundsMap;
  }

  bool rewriteOperation(PatternRewriter &rewriter, stablehlo::RemOp remOp,
                        const DenseMap<Value, Bounds> &boundsMap) const {
    auto lhs_bounds = boundsMap.lookup(remOp.getLhs());
    auto rhs_bounds = boundsMap.lookup(remOp.getRhs());

    auto rhs_abs_min = abs(rhs_bounds.min);
    auto rhs_abs_max = abs(rhs_bounds.max);
    auto minboth = min(rhs_abs_min, rhs_abs_max);
    if (lhs_bounds.max.slt(minboth) && lhs_bounds.min.sgt(-minboth)) {
      auto lhs = remOp.getLhs();
      rewriter.replaceOp(remOp, lhs);
      return true;
    }

    return false;
  };

  bool rewriteOperation(PatternRewriter &rewriter,
                        stablehlo::CompareOp compareOp,
                        const DenseMap<Value, Bounds> &boundsMap) const {
    auto lhs_bounds = boundsMap.lookup(compareOp.getLhs());
    auto rhs_bounds = boundsMap.lookup(compareOp.getRhs());

    auto replaceWithTrue = [&]() {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          compareOp, compareOp.getType(),
          cast<ElementsAttr>(makeAttr(compareOp.getType(), true)));
      return true;
    };
    auto replaceWithFalse = [&]() {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          compareOp, compareOp.getType(),
          cast<ElementsAttr>(makeAttr(compareOp.getType(), false)));
      return true;
    };

    switch (compareOp.getComparisonDirection()) {
    case stablehlo::ComparisonDirection::EQ: // lhs == rhs
      // all values are equal
      if (lhs_bounds.min == lhs_bounds.max &&
          rhs_bounds.min == rhs_bounds.max &&
          lhs_bounds.min == rhs_bounds.min) {
        return replaceWithTrue();
      }
      // outside of range: lhs_max < rhs_min or rhs_max < lhs_min
      if (lhs_bounds.max.slt(rhs_bounds.min) ||
          rhs_bounds.max.slt(lhs_bounds.min)) {
        return replaceWithFalse();
      }
      break;
    case stablehlo::ComparisonDirection::NE: // lhs != rhs
      // outside of range: lhs_max < rhs_min or rhs_max < lhs_min
      if (lhs_bounds.max.slt(rhs_bounds.min) ||
          rhs_bounds.max.slt(lhs_bounds.min)) {
        return replaceWithFalse();
      }
      break;
    case stablehlo::ComparisonDirection::GE:    // lhs >= rhs
      if (lhs_bounds.max.slt(rhs_bounds.min)) { // lhs_max < rhs_min
        return replaceWithFalse();
      }
      if (lhs_bounds.min.sge(rhs_bounds.max)) { // lhs_min >= rhs_max
        return replaceWithTrue();
      }
      break;
    case stablehlo::ComparisonDirection::GT:    // lhs > rhs
      if (lhs_bounds.max.sle(rhs_bounds.min)) { // lhs_max <= rhs_min
        return replaceWithFalse();
      }
      if (lhs_bounds.min.sgt(rhs_bounds.max)) { // lhs_min > rhs_max
        return replaceWithTrue();
      }
      break;
    case stablehlo::ComparisonDirection::LE:    // lhs <= rhs
      if (rhs_bounds.max.slt(lhs_bounds.min)) { // rhs_max < lhs_min
        return replaceWithFalse();
      }
      if (rhs_bounds.min.sge(lhs_bounds.max)) { // rhs_min >= lhs_max
        return replaceWithTrue();
      }
      break;
    case stablehlo::ComparisonDirection::LT:    // lhs < rhs
      if (rhs_bounds.max.sle(lhs_bounds.min)) { // rhs_max <= lhs_min
        return replaceWithFalse();
      }
      if (rhs_bounds.min.sgt(lhs_bounds.max)) { // rhs_min > lhs_max
        return replaceWithTrue();
      }
      break;
    }

    return false;
  };

  bool rewriteOperation(PatternRewriter &rewriter, stablehlo::AbsOp absOp,
                        const DenseMap<Value, Bounds> &boundsMap) const {
    auto bounds = boundsMap.lookup(absOp.getOperand());

    APInt zero(bounds.min.getBitWidth(), 0, true);
    if (bounds.min.sge(zero)) {
      rewriter.replaceOp(absOp, absOp.getOperand());
      return true;
    }
    return false;
  };

  bool rewriteOperation(PatternRewriter &rewriter, stablehlo::ClampOp clampOp,
                        const DenseMap<Value, Bounds> &boundsMap) const {
    auto min_bounds = boundsMap.lookup(clampOp.getMin());
    auto max_bounds = boundsMap.lookup(clampOp.getMax());
    auto operand_bounds = boundsMap.lookup(clampOp.getOperand());

    if (min_bounds.max.sle(operand_bounds.min) &&
        max_bounds.min.sge(operand_bounds.max)) {
      rewriter.replaceOp(clampOp, clampOp.getOperand());
      return true;
    }
    return false;
  };
};

struct SplitVariadicScatterOp
    : public CheckedOpRewritePattern<stablehlo::ScatterOp,
                                     SplitVariadicScatterOp> {
  using CheckedOpRewritePattern<
      stablehlo::ScatterOp, SplitVariadicScatterOp>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp scatterOp,
                                    PatternRewriter &rewriter) const {
    size_t N = scatterOp.getInputs().size();
    if (N <= 1)
      return rewriter.notifyMatchFailure(scatterOp, "not variadic scatter");

    Region &region = scatterOp.getUpdateComputation();
    if (!isUpdateComputationValid(region, N))
      return rewriter.notifyMatchFailure(
          scatterOp, "invalid update computation for splitting");

    for (size_t i = 0; i < N; ++i) {
      rewriter.setInsertionPoint(scatterOp);
      auto newScatterOp = stablehlo::ScatterOp::create(
          rewriter, scatterOp.getLoc(), scatterOp.getInputs()[i].getType(),
          ValueRange{scatterOp.getInputs()[i]}, scatterOp.getScatterIndices(),
          ValueRange{scatterOp.getUpdates()[i]},
          scatterOp.getScatterDimensionNumbersAttr(),
          scatterOp.getIndicesAreSortedAttr(),
          scatterOp.getUniqueIndicesAttr());
      addUpdateComputationForIndex(newScatterOp, rewriter, scatterOp.getLoc(),
                                   region, i, N);
      rewriter.replaceAllUsesWith(scatterOp.getResult(i),
                                  newScatterOp.getResult(0));
    }
    return success();
  }

private:
  bool isUpdateComputationValid(Region &region, size_t numOperands) const {
    if (region.empty())
      return false;

    Block &block = region.front();

    if (block.getNumArguments() != numOperands * 2)
      return false;

    // Check if the computation mixes different operand indices
    // A valid computation should only operate on corresponding pairs:
    // (results[i], updates[i]) -> result[i]
    Operation *terminator = block.getTerminator();
    if (!terminator || terminator->getNumOperands() != numOperands)
      return false;

    // Trace data flow for each return value to ensure it only depends
    // on the corresponding input pair
    for (size_t i = 0; i < numOperands; ++i) {
      Value returnVal = terminator->getOperand(i);

      // Expected dependencies: args[i] and args[N+i]
      llvm::SmallSet<size_t, 4> allowedArgIndices = {i, numOperands + i};

      if (!checkValueDependencies(returnVal, block, allowedArgIndices)) {
        return false;
      }
    }

    return true;
  }

  // Check that a value only depends on allowed block arguments
  bool checkValueDependencies(
      Value val, Block &block,
      const llvm::SmallSet<size_t, 4> &allowedArgIndices) const {
    SmallVector<Value> worklist;
    llvm::SmallSet<Value, 16> visited;

    worklist.push_back(val);

    while (!worklist.empty()) {
      Value current = worklist.pop_back_val();
      if (!visited.insert(current).second)
        continue;

      if (auto blockArg = dyn_cast<BlockArgument>(current)) {
        if (blockArg.getOwner() != &block)
          continue;

        if (!allowedArgIndices.contains(blockArg.getArgNumber()))
          return false;

        continue;
      }

      Operation *defOp = current.getDefiningOp();
      if (!defOp)
        return false;

      if (defOp->getBlock() != &block)
        continue;

      for (Value operand : defOp->getOperands()) {
        worklist.push_back(operand);
      }

      // For operations with regions, add the region's results/yields to
      // worklist
      for (Region &region : defOp->getRegions()) {
        for (Block &nestedBlock : region.getBlocks()) {
          Operation *terminator = nestedBlock.getTerminator();
          if (terminator) {
            for (Value operand : terminator->getOperands()) {
              worklist.push_back(operand);
            }
          }
        }
      }
    }

    return true;
  }

  void addUpdateComputationForIndex(stablehlo::ScatterOp newScatterOp,
                                    PatternRewriter &rewriter, Location loc,
                                    Region &origRegion, size_t index,
                                    size_t numOperands) const {
    Block &origBlock = origRegion.front();

    // Create a new region for the single-operand update computation
    Region &newRegion = newScatterOp.getUpdateComputation();
    Block *newBlock = rewriter.createBlock(&newRegion);

    // Add 2 block arguments: one from result, one from update
    BlockArgument resultArg =
        newBlock->addArgument(origBlock.getArgument(index).getType(), loc);
    BlockArgument updateArg = newBlock->addArgument(
        origBlock.getArgument(numOperands + index).getType(), loc);

    // Map original arguments to new ones
    IRMapping mapper;
    mapper.map(origBlock.getArgument(index), resultArg);
    mapper.map(origBlock.getArgument(numOperands + index), updateArg);

    // Clone operations from original computation that affect this index
    Operation *terminator = origBlock.getTerminator();

    if (auto returnOp = dyn_cast<stablehlo::ReturnOp>(terminator)) {
      // Extract the i-th return value's defining operations
      Value returnVal = returnOp.getOperand(index);
      cloneComputationSlice(rewriter, returnVal, mapper, newBlock);

      // Create return with single value
      rewriter.setInsertionPointToEnd(newBlock);
      Value mappedResult = mapper.lookupOrDefault(returnVal);
      stablehlo::ReturnOp::create(rewriter, loc, mappedResult);
    }
  }

  void cloneComputationSlice(PatternRewriter &rewriter, Value val,
                             IRMapping &mapper, Block *targetBlock) const {
    if (mapper.contains(val))
      return;

    Operation *defOp = val.getDefiningOp();
    if (!defOp)
      return;

    for (Value operand : defOp->getOperands()) {
      cloneComputationSlice(rewriter, operand, mapper, targetBlock);
    }

    rewriter.setInsertionPointToEnd(targetBlock);
    Operation *clonedOp = rewriter.clone(*defOp, mapper);

    for (auto [orig, cloned] :
         llvm::zip(defOp->getResults(), clonedOp->getResults())) {
      mapper.map(orig, cloned);
    }
  }
};

struct DynamicSliceSimplify
    : public CheckedOpRewritePattern<stablehlo::DynamicSliceOp,
                                     DynamicSliceSimplify> {
  using CheckedOpRewritePattern<stablehlo::DynamicSliceOp,
                                DynamicSliceSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    // If slice is a splatted constant, then it doesn't matter which indices we
    // are slicing.
    SplatElementsAttr splat;
    if (matchPattern(op.getOperand(), m_Constant(&splat))) {
      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
          op, op.getType(), splat.resizeSplat(op.getType()));
      return success();
    }

    // If it is from a broadcasted scalar, we can replace it with a broadcast
    if (auto bcast =
            op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>()) {
      auto bcastInput = bcast.getOperand();
      auto bcastInputTy = cast<RankedTensorType>(bcastInput.getType());
      if (bcastInputTy.getRank() == 0) {
        rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
            op, op.getType(), bcastInput, rewriter.getDenseI64ArrayAttr({}));
        return success();
      }
    }

    return failure();
  }
};

// TODO: generalize to higher ranked tensors
// TODO: if we determine that all accesses are on some offset diagonal,
// we can still replace it will a multiply combined with pad/slice
// If we prove that only the diagonal elements of a dot_general are accessed,
// we replace the dot_general with a cheaper multiply op. Note that
// this implies `diag(new_op(A, B)) = diag(A x B)` however
// `new_op(A, B) != A x B`
struct DotGeneralOnlyDiagonalAccess
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                     DotGeneralOnlyDiagonalAccess> {
  using CheckedOpRewritePattern<
      stablehlo::DotGeneralOp,
      DotGeneralOnlyDiagonalAccess>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    auto resTy = cast<RankedTensorType>(op.getType());
    if (resTy.getRank() != 2)
      return failure();

    auto M = resTy.getDimSize(0);
    auto N = resTy.getDimSize(1);
    auto diagLen = std::min(M, N);

    auto lhs = op.getLhs();
    auto rhs = op.getRhs();
    auto dotDimNumbers = op.getDotDimensionNumbers();
    auto lhsContractingDims = dotDimNumbers.getLhsContractingDimensions();
    auto rhsContractingDims = dotDimNumbers.getRhsContractingDimensions();
    auto lhsBatchingDims = dotDimNumbers.getLhsBatchingDimensions();
    auto rhsBatchingDims = dotDimNumbers.getRhsBatchingDimensions();

    if (lhsContractingDims.size() != 1 || rhsContractingDims.size() != 1 ||
        lhsBatchingDims.size() != 0 || rhsBatchingDims.size() != 0)
      return failure();

    llvm::SetVector<Operation *> opsToReplace;
    llvm::SmallPtrSet<Operation *, 4> seenOps;
    for (auto user : op->getUsers()) {
      if (seenOps.count(user))
        continue;
      if (!enzyme::allAccessesAreOnMainDiagonal(user, opsToReplace))
        return failure();
      seenOps.insert(user);
    }

    if (opsToReplace.empty())
      return failure();

    // rewrite the dot_general to a multiply.
    // we insert transpose ops here, but those will get removed later
    auto lhsContractDim = lhsContractingDims[0];
    auto rhsContractDim = rhsContractingDims[0];
    // result[i, i] = sum_k (lhs[i, k] * rhs[k, i])
    //              = reduce_sum(lhs[i, :] * rhs[:, i])
    if (lhsContractDim == 0) {
      // move to dim = 1
      lhs = stablehlo::TransposeOp::create(
          rewriter, op.getLoc(), lhs, rewriter.getDenseI64ArrayAttr({1, 0}));
    }
    lhs = stablehlo::SliceOp::create(
        rewriter, op.getLoc(), lhs, rewriter.getDenseI64ArrayAttr({0, 0}),
        rewriter.getDenseI64ArrayAttr(
            {diagLen, cast<ShapedType>(lhs.getType()).getDimSize(1)}),
        rewriter.getDenseI64ArrayAttr({1, 1})); // [DiagSize, C]

    if (rhsContractDim == 0) {
      // move to dim = 1
      rhs = stablehlo::TransposeOp::create(
          rewriter, op.getLoc(), rhs, rewriter.getDenseI64ArrayAttr({1, 0}));
    }
    rhs = stablehlo::SliceOp::create(
        rewriter, op.getLoc(), rhs, rewriter.getDenseI64ArrayAttr({0, 0}),
        rewriter.getDenseI64ArrayAttr(
            {diagLen, cast<ShapedType>(rhs.getType()).getDimSize(1)}),
        rewriter.getDenseI64ArrayAttr({1, 1})); // [DiagSize, C]

    auto newMul = stablehlo::MulOp::create(rewriter, op.getLoc(), lhs,
                                           rhs); // [DiagSize, C]

    auto elemTy = cast<RankedTensorType>(newMul.getType()).getElementType();
    auto tenElemTy = RankedTensorType::get({}, elemTy);
    auto reduceOp = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), ValueRange(newMul.getResult()),
        ValueRange(stablehlo::ConstantOp::create(
                       rewriter, op.getLoc(), tenElemTy,
                       cast<ElementsAttr>(makeAttr(tenElemTy, 0)))
                       .getResult()),
        {1});

    {
      Region &region = reduceOp.getBody();
      Block *block = rewriter.createBlock(&region);
      block->addArgument(tenElemTy, op.getLoc());
      block->addArgument(tenElemTy, op.getLoc());

      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(block);
      auto addOp = stablehlo::AddOp::create(
          rewriter, op.getLoc(), block->getArgument(0), block->getArgument(1));
      stablehlo::ReturnOp::create(rewriter, op.getLoc(), addOp.getResult());
    }

    for (auto &opToReplace : opsToReplace) {
      if (auto sliceOp = dyn_cast<stablehlo::SliceOp>(opToReplace)) {
        replaceSliceOp(rewriter, sliceOp, reduceOp, M, N, diagLen);
      } else {
        assert(false && "Unknown op to replace. open an issue on github");
      }
    }

    return success();
  }

private:
  void replaceSliceOp(PatternRewriter &rewriter, stablehlo::SliceOp sliceOp,
                      stablehlo::ReduceOp reduceOp, int64_t M, int64_t N,
                      int64_t diagLen) const {
    int64_t start = sliceOp.getStartIndices()[0];
    int64_t limit = sliceOp.getLimitIndices()[0];
    int64_t stride = sliceOp.getStrides()[0];
    int64_t diagStride = N + 1;

    int64_t newStart = start / diagStride;
    int64_t newLimit = (limit - 1) / diagStride + 1;
    int64_t newStride = stride / diagStride;

    rewriter.setInsertionPoint(sliceOp);
    rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
        sliceOp, reduceOp.getResult(0),
        rewriter.getDenseI64ArrayAttr({newStart}),
        rewriter.getDenseI64ArrayAttr({newLimit}),
        rewriter.getDenseI64ArrayAttr({newStride}));
  }
};

template <typename OpTy>
struct BinaryNegatedOperandsSimplify
    : public CheckedOpRewritePattern<OpTy,
                                     BinaryNegatedOperandsSimplify<OpTy>> {
  using CheckedOpRewritePattern<
      OpTy, BinaryNegatedOperandsSimplify<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsInfo = getOperandInfo(lhs, op);
    if (lhsInfo.kind == NegKind::NOT)
      return failure();

    auto rhsInfo = getOperandInfo(rhs, op);
    if (rhsInfo.kind == NegKind::NOT)
      return failure();

    // prevent an infinite loop here op(sub, sub) -> op(sub, sub) -> ...
    if (lhsInfo.kind == NegKind::SUB && rhsInfo.kind == NegKind::SUB)
      return failure();

    Value negatedLhs = negateOperand(lhsInfo, rewriter);
    Value negatedRhs = negateOperand(rhsInfo, rewriter);
    rewriter.replaceOpWithNewOp<OpTy>(op, negatedLhs, negatedRhs);
    return success();
  }

private:
  enum class NegKind { NEG, SUB, NOT };

  struct OperandInfo {
    NegKind kind;
    Value lhs;
    Value rhs;
  };

  Value negateOperand(OperandInfo info, PatternRewriter &rewriter) const {
    switch (info.kind) {
    case NegKind::NEG:
      return info.lhs;
    case NegKind::SUB:
      return stablehlo::SubtractOp::create(rewriter, info.lhs.getLoc(),
                                           info.rhs, info.lhs)
          .getResult();
    default:
      llvm_unreachable("Unhandled negation");
    }
  }

  OperandInfo getOperandInfo(Value val, OpTy op) const {
    auto defOp = val.getDefiningOp();
    if (!defOp || !isOnlyUsedInOperation(defOp, op))
      return OperandInfo{NegKind::NOT, nullptr, nullptr};

    if (auto negOp = dyn_cast<stablehlo::NegOp>(defOp)) {
      return OperandInfo{NegKind::NEG, negOp.getOperand(), nullptr};
    }

    if (auto subOp = dyn_cast<stablehlo::SubtractOp>(defOp)) {
      return OperandInfo{NegKind::SUB, subOp.getLhs(), subOp.getRhs()};
    }

    return OperandInfo{NegKind::NOT, nullptr, nullptr};
  }
};

static void ReduceToDotGeneral(Location loc, OpBuilder &builder,
                               RankedTensorType resTy, RankedTensorType inTy,
                               ArrayRef<int64_t> dims, Value lhs, Value rhs,
                               Value &result) {
  SmallVector<int64_t> batchDims;
  for (int i = 0; i < inTy.getRank(); i++) {
    if (!llvm::is_contained(dims, i)) {
      batchDims.push_back(i);
    }
  }

  auto dotDims = stablehlo::DotDimensionNumbersAttr::get(
      builder.getContext(), batchDims, batchDims, dims, dims);
  result = stablehlo::DotGeneralOp::create(builder, loc, resTy, lhs, rhs,
                                           dotDims, nullptr, nullptr);
}

struct ReduceMulToDotGeneral
    : public CheckedOpRewritePattern<stablehlo::ReduceOp,
                                     ReduceMulToDotGeneral> {
  using CheckedOpRewritePattern<stablehlo::ReduceOp,
                                ReduceMulToDotGeneral>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1 || op.getInitValues().size() != 1) {
      return rewriter.notifyMatchFailure(
          op, "only single-operand single-init reduce is supported");
    }

    auto checkCommonReduce = mlir::stablehlo::CheckCommonReduceOp(op);
    if (checkCommonReduce.kind != ReduceOpKind::Add ||
        !enzyme::isZero(op.getInitValues()[0])) {
      return rewriter.notifyMatchFailure(op, "reduction is not add");
    }

    auto input = op.getInputs()[0];
    auto mul = input.getDefiningOp<stablehlo::MulOp>();
    if (!mul) {
      return rewriter.notifyMatchFailure(op, "input source is not a mul op");
    }
    if (!llvm::hasSingleElement(mul->getUsers())) {
      return rewriter.notifyMatchFailure(op, "mul op has multiple users");
    }

    Value result;
    ReduceToDotGeneral(op->getLoc(), rewriter,
                       cast<RankedTensorType>(op.getResult(0).getType()),
                       cast<RankedTensorType>(input.getType()),
                       op.getDimensions(), mul.getLhs(), mul.getRhs(), result);
    rewriter.replaceOp(op, result);
    return success();
  }
};

struct SplitReduceAddMulToAddDotGeneral final
    : public CheckedOpRewritePattern<stablehlo::ReduceOp,
                                     SplitReduceAddMulToAddDotGeneral> {
  using CheckedOpRewritePattern<
      stablehlo::ReduceOp,
      SplitReduceAddMulToAddDotGeneral>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1 || op.getInitValues().size() != 1) {
      return rewriter.notifyMatchFailure(
          op, "only single-operand single-init reduce is supported");
    }

    auto checkCommonReduce = mlir::stablehlo::CheckCommonReduceOp(op);
    auto initVal = op.getInitValues()[0];
    if (checkCommonReduce.kind != ReduceOpKind::Add ||
        !enzyme::isZero(initVal)) {
      return rewriter.notifyMatchFailure(op, "reduction is not add");
    }

    auto input = op.getInputs()[0];
    auto addOp = input.getDefiningOp<stablehlo::AddOp>();
    if (!addOp || !llvm::hasSingleElement(addOp->getUsers())) {
      return failure();
    }

    auto reduceOutTy = cast<RankedTensorType>(op.getResult(0).getType());
    auto reduceInTy = cast<RankedTensorType>(input.getType());

    // reduce(add, add(lhs, rhs))
    // atleast one of lhs or rhs must satisfy
    //    1. is a mul
    //    2. either of the mul operands is a bcast_in_dim (to enable
    //       DotGeneralBroadcastInDim simplifications)
    Value lhs = addOp.getLhs(), rhs = addOp.getRhs();
    bool lhsRewritten = false, rhsRewritten = false, lhsSameAsRhs = rhs == lhs;

    Value lhsLhs, lhsRhs, rhsLhs, rhsRhs;
    if (operandSatisfiesSplitCriteria(rewriter, lhs, lhsLhs, lhsRhs, addOp)) {
      ReduceToDotGeneral(op->getLoc(), rewriter, reduceOutTy, reduceInTy,
                         op.getDimensions(), lhsLhs, lhsRhs, lhs);
      lhsRewritten = true;
    }

    if (lhsSameAsRhs) {
      rhsRewritten = lhsRewritten;
      rhs = lhs;
    } else {
      if (operandSatisfiesSplitCriteria(rewriter, rhs, rhsLhs, rhsRhs, addOp)) {
        ReduceToDotGeneral(op->getLoc(), rewriter, reduceOutTy, reduceInTy,
                           op.getDimensions(), rhsLhs, rhsRhs, rhs);
        rhsRewritten = true;
      }

      if (rhsRewritten && !lhsRewritten) { // insert a reduceOp
        rewriter.setInsertionPointAfterValue(lhs);
        auto newLhsReduce = stablehlo::ReduceOp::create(
            rewriter, op->getLoc(), ValueRange(lhs), ValueRange(initVal),
            op.getDimensions());
        IRMapping mapper;
        op.getRegion().cloneInto(&newLhsReduce.getRegion(), mapper);
        lhs = newLhsReduce->getResult(0);
      }

      if (lhsRewritten && !rhsRewritten) { // insert a reduceOp
        rewriter.setInsertionPointAfterValue(rhs);
        auto newRhsReduce = stablehlo::ReduceOp::create(
            rewriter, op->getLoc(), ValueRange(rhs), ValueRange(initVal),
            op.getDimensions());
        IRMapping mapper;
        op.getRegion().cloneInto(&newRhsReduce.getRegion(), mapper);
        rhs = newRhsReduce->getResult(0);
      }
    }

    if (lhsRewritten || rhsRewritten) {
      rewriter.setInsertionPoint(op);
      auto newAddOp =
          stablehlo::AddOp::create(rewriter, op->getLoc(), lhs, rhs);
      rewriter.replaceOp(op, newAddOp);
      rewriter.eraseOp(addOp);
      return success();
    }
    return failure();
  }

private:
  bool operandSatisfiesSplitCriteria(PatternRewriter &rewriter, Value operand,
                                     Value &lhs, Value &rhs,
                                     Operation *op) const {
    auto mulOp = operand.getDefiningOp<stablehlo::MulOp>();
    if (!mulOp || !isOnlyUsedInOperation(mulOp, op)) {
      return false;
    }

    lhs = mulOp.getLhs();
    rhs = mulOp.getRhs();

    rewriter.setInsertionPoint(mulOp);
    auto lhsIsBcast = lhs.getDefiningOp<stablehlo::BroadcastInDimOp>();
    auto rhsIsBcast = rhs.getDefiningOp<stablehlo::BroadcastInDimOp>();
    return lhsIsBcast || rhsIsBcast;
  }
};

// returns input A and transpose param to syrk if can convert dot_general op to
// syrk
std::pair<Value, enzymexla::LapackTranspose>
extractSyrkInput(stablehlo::DotGeneralOp op, PatternRewriter &rewriter) {

  Value syrkInput;
  enzymexla::LapackTranspose lapackTranspose;

  auto failure =
      std::pair<Value, enzymexla::LapackTranspose>(syrkInput, lapackTranspose);

  auto dotDims = op.getDotDimensionNumbers();
  auto lhs = op.getLhs();
  auto rhs = op.getRhs();

  auto lhsType = cast<RankedTensorType>(lhs.getType());
  auto rhsType = cast<RankedTensorType>(rhs.getType());
  auto outType = cast<RankedTensorType>(op.getResult().getType());
  if (lhsType.getRank() != 2 || rhsType.getRank() != 2 ||
      outType.getRank() != 2) {
    return failure;
  }

  if (dotDims.getLhsBatchingDimensions().size() != 0 ||
      dotDims.getRhsBatchingDimensions().size() != 0) {
    return failure;
  }

  if (dotDims.getLhsContractingDimensions().size() != 1 ||
      dotDims.getRhsContractingDimensions().size() != 1) {
    return failure;
  }

  // check that transpose dimensions are [1,0]
  auto isTrueTranspose = [](stablehlo::TransposeOp tOp) -> bool {
    auto perm = tOp.getPermutation();
    return perm.size() == 2 && perm[0] == 1 && perm[1] == 0;
  };

  auto lhsContractingDim = dotDims.getLhsContractingDimensions()[0];
  auto rhsContractingDim = dotDims.getRhsContractingDimensions()[0];

  if (lhs == rhs) {
    if (lhsContractingDim == rhsContractingDim) {
      syrkInput = lhs;
      if (lhsContractingDim == 1) {
        lapackTranspose = enzymexla::LapackTranspose::none;
      } else {
        lapackTranspose = enzymexla::LapackTranspose::transpose;
      }
    } else {
      // if lhsContractingDim != rhsContractingDim, then we can fuse iff
      // lhs/rhs is a symmetric matrix
      if (canApplySymmetricPattern(lhs, rewriter)) {
        syrkInput = lhs;
        if (lhsContractingDim == 1) {
          lapackTranspose = enzymexla::LapackTranspose::none;
        } else {
          lapackTranspose = enzymexla::LapackTranspose::transpose;
        }
      }
    }
  }

  if (auto lhsT = lhs.getDefiningOp<stablehlo::TransposeOp>()) {
    if (isTrueTranspose(lhsT) && lhsT.getOperand() == rhs &&
        lhsContractingDim == 1 - rhsContractingDim) {
      syrkInput = rhs;
      if (rhsContractingDim == 1) {
        lapackTranspose = enzymexla::LapackTranspose::none;
      } else {
        lapackTranspose = enzymexla::LapackTranspose::transpose;
      }
    }
  }

  if (auto rhsT = rhs.getDefiningOp<stablehlo::TransposeOp>()) {
    if (isTrueTranspose(rhsT) && rhsT.getOperand() == lhs &&
        rhsContractingDim == 1 - lhsContractingDim) {
      syrkInput = lhs;
      if (lhsContractingDim == 0) {
        lapackTranspose = enzymexla::LapackTranspose::transpose;
      } else {
        lapackTranspose = enzymexla::LapackTranspose::none;
      }
    }
  }

  return std::pair<Value, enzymexla::LapackTranspose>(syrkInput,
                                                      lapackTranspose);
}

// currently limited to non-batched dot_general
struct DotGeneralToSyrk
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                     DotGeneralToSyrk> {
  using CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                DotGeneralToSyrk>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {

    auto [syrkInput, lapackTranspose] = extractSyrkInput(op, rewriter);
    if (!syrkInput)
      return failure();

    auto elemType =
        cast<RankedTensorType>(syrkInput.getType()).getElementType();
    auto alphaType = RankedTensorType::get({}, elemType);

    auto cMatrix = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), op.getType(),
        cast<ElementsAttr>(makeAttr(op.getType(), 0)));
    auto alpha = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), alphaType,
        cast<ElementsAttr>(makeAttr(alphaType, 1)));
    auto beta = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), alphaType,
        cast<ElementsAttr>(makeAttr(alphaType, 0)));
    auto syrkOp = enzymexla::SyrkOp::create(
        rewriter, op.getLoc(), op.getResult().getType(), syrkInput, cMatrix,
        alpha, beta,
        enzymexla::LapackUploAttr::get(op.getContext(),
                                       enzymexla::LapackUplo::F),
        enzymexla::LapackUploAttr::get(op.getContext(),
                                       enzymexla::LapackUplo::F),
        enzymexla::LapackTransposeAttr::get(op.getContext(), lapackTranspose));
    rewriter.replaceOp(op, syrkOp.getResult());
    return success();
  }
};

// currently limited to non-batched dot_general
struct DotGeneralToSymm
    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                     DotGeneralToSymm> {
  using CheckedOpRewritePattern<stablehlo::DotGeneralOp,
                                DotGeneralToSymm>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    // if can rewrite as an syrk operation instead, don't rewrite as symm
    auto [syrkInput, _] = extractSyrkInput(op, rewriter);
    if (syrkInput) {
      return failure();
    }

    auto dotDims = op.getDotDimensionNumbers();
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsType = cast<RankedTensorType>(lhs.getType());
    auto rhsType = cast<RankedTensorType>(rhs.getType());

    if (lhsType.getRank() != 2 || rhsType.getRank() != 2) {
      return failure();
    }

    if (dotDims.getLhsBatchingDimensions().size() != 0 ||
        dotDims.getRhsBatchingDimensions().size() != 0) {
      return failure();
    }

    if (dotDims.getLhsContractingDimensions().size() != 1 ||
        dotDims.getRhsContractingDimensions().size() != 1) {
      return failure();
    }

    Value symmMatrix, genMatrix;
    enzymexla::LapackSide side;

    auto lhsContractingDim = dotDims.getLhsContractingDimensions()[0];
    auto rhsContractingDim = dotDims.getRhsContractingDimensions()[0];

    // can only replace with symm if either lhs or rhs is symmetric
    if (canApplySymmetricPattern(lhs, rewriter)) {
      side = enzymexla::LapackSide::left;
      symmMatrix = lhs;
      if (rhsContractingDim == 1) {
        auto perm = rewriter.getDenseI64ArrayAttr({1, 0});
        rhs = stablehlo::TransposeOp::create(rewriter, op.getLoc(), rhs, perm);
      }
      genMatrix = rhs;
    } else if (canApplySymmetricPattern(rhs, rewriter)) {
      side = enzymexla::LapackSide::right;
      symmMatrix = rhs;
      if (lhsContractingDim == 0) {
        auto perm = rewriter.getDenseI64ArrayAttr({1, 0});
        lhs = stablehlo::TransposeOp::create(rewriter, op.getLoc(), lhs, perm);
      }
      genMatrix = lhs;
    } else {
      return failure();
    }

    auto elemType = cast<RankedTensorType>(lhs.getType()).getElementType();
    auto alphaType = RankedTensorType::get({}, elemType);

    auto cMatrix = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), op.getType(),
        cast<ElementsAttr>(makeAttr(op.getType(), 0)));
    auto alpha = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), alphaType,
        cast<ElementsAttr>(makeAttr(alphaType, 1)));
    auto beta = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), alphaType,
        cast<ElementsAttr>(makeAttr(alphaType, 0)));
    auto symmOp = enzymexla::SymmOp::create(
        rewriter, op.getLoc(), op.getResult().getType(),
        symmMatrix, // A (symmetric)
        genMatrix,  // B (general)
        cMatrix, alpha, beta,
        enzymexla::LapackSideAttr::get(op.getContext(), side),
        enzymexla::LapackUploAttr::get(op.getContext(),
                                       enzymexla::LapackUplo::F));
    rewriter.replaceOp(op, symmOp.getResult());
    return success();
  }
};

template <typename ST, typename Child>
struct FuseMulBase
    : public CheckedOpRewritePattern<stablehlo::MulOp, FuseMulBase<ST, Child>> {
  using CheckedOpRewritePattern<
      stablehlo::MulOp, FuseMulBase<ST, Child>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    ST targetOp;
    Value other;

    if (auto lhsTarget = lhs.getDefiningOp<ST>()) {
      targetOp = lhsTarget;
      other = rhs;
    } else if (auto rhsTarget = rhs.getDefiningOp<ST>()) {
      targetOp = rhsTarget;
      other = lhs;
    } else {
      return failure();
    }

    if (!isOnlyUsedInOperation(targetOp, op)) {
      return failure();
    }

    Value scalarVal =
        stablehlo::getScalarValue(other.getDefiningOp(), rewriter);
    if (!scalarVal) {
      return failure();
    }

    return ((Child *)this)->fuseMul(rewriter, targetOp, op, scalarVal);
  }
};

struct FuseMulIntoSymm
    : public FuseMulBase<enzymexla::SymmOp, FuseMulIntoSymm> {
  using FuseMulBase<enzymexla::SymmOp, FuseMulIntoSymm>::FuseMulBase;

  LogicalResult fuseMul(PatternRewriter &rewriter, enzymexla::SymmOp symmOp,
                        stablehlo::MulOp op, Value scalarVal) {
    auto newBeta = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                            symmOp.getBeta(), scalarVal);
    auto newAlpha = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                             symmOp.getAlpha(), scalarVal);

    rewriter.replaceOpWithNewOp<enzymexla::SymmOp>(
        op, symmOp.getType(), symmOp.getA(), symmOp.getB(), symmOp.getC(),
        newAlpha, newBeta, symmOp.getSideAttr(), symmOp.getUploAttr());
    return success();
  }
};

struct FuseMulIntoSyrk
    : public FuseMulBase<enzymexla::SyrkOp, FuseMulIntoSyrk> {
  using FuseMulBase<enzymexla::SyrkOp, FuseMulIntoSyrk>::FuseMulBase;

  LogicalResult fuseMul(PatternRewriter &rewriter, enzymexla::SyrkOp syrkOp,
                        stablehlo::MulOp op, Value scalarVal) {
    auto newBeta = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                            syrkOp.getBeta(), scalarVal);
    auto newAlpha = stablehlo::MulOp::create(rewriter, op.getLoc(),
                                             syrkOp.getAlpha(), scalarVal);

    rewriter.replaceOpWithNewOp<enzymexla::SyrkOp>(
        op, syrkOp.getType(), syrkOp.getA(), syrkOp.getC(), newAlpha, newBeta,
        syrkOp.getUploAttr(), syrkOp.getOutputUploAttr(),
        syrkOp.getTransposeAttr());
    return success();
  }
};

template <typename ST, typename Child>
struct FuseAddBase
    : public CheckedOpRewritePattern<stablehlo::AddOp, FuseAddBase<ST, Child>> {
  using CheckedOpRewritePattern<
      stablehlo::AddOp, FuseAddBase<ST, Child>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::AddOp op,
                                    PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    ST targetOp;
    Value other;

    if (auto lhsTarget = lhs.getDefiningOp<ST>()) {
      targetOp = lhsTarget;
      other = rhs;
    } else if (auto rhsTarget = rhs.getDefiningOp<ST>()) {
      targetOp = rhsTarget;
      other = lhs;
    } else {
      return failure();
    }

    if (!isOnlyUsedInOperation(targetOp, op)) {
      return failure();
    }

    return ((Child *)this)->fuseAdd(rewriter, targetOp, op, other);
  }

  std::pair<mlir::Value, mlir::Value> computeNewCBeta(PatternRewriter &rewriter,
                                                      Value oldBeta, Type type,
                                                      stablehlo::AddOp op,
                                                      Value c, Value other) {
    auto scaledC = stablehlo::MulOpCreate(rewriter, op.getLoc(), c, oldBeta);

    auto newC = stablehlo::AddOpCreate(rewriter, op.getLoc(), scaledC, other);

    auto newBeta = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(), oldBeta.getType(),
        cast<ElementsAttr>(makeAttr(oldBeta.getType(), 1)));
    return {newC, newBeta};
  }
};

struct FuseAddIntoSymm
    : public FuseAddBase<enzymexla::SymmOp, FuseAddIntoSymm> {
  using FuseAddBase<enzymexla::SymmOp, FuseAddIntoSymm>::FuseAddBase;

  LogicalResult fuseAdd(PatternRewriter &rewriter, enzymexla::SymmOp symmOp,
                        stablehlo::AddOp op, Value other) {
    auto [newC, newBeta] = computeNewCBeta(
        rewriter, symmOp.getBeta(), symmOp.getType(), op, symmOp.getC(), other);

    rewriter.replaceOpWithNewOp<enzymexla::SymmOp>(
        op, symmOp.getType(), symmOp.getA(), symmOp.getB(), newC,
        symmOp.getAlpha(), newBeta, symmOp.getSideAttr(), symmOp.getUploAttr());
    return success();
  }
};

struct FuseAddIntoSyrk
    : public FuseAddBase<enzymexla::SyrkOp, FuseAddIntoSyrk> {
  using FuseAddBase<enzymexla::SyrkOp, FuseAddIntoSyrk>::FuseAddBase;

  LogicalResult fuseAdd(PatternRewriter &rewriter, enzymexla::SyrkOp syrkOp,
                        stablehlo::AddOp op, Value other) {

    // we can fuse this addition iff the other operand is a symmetric matrix
    if (!canApplySymmetricPattern(other, rewriter)) {
      return failure();
    }

    auto [newC, newBeta] = computeNewCBeta(
        rewriter, syrkOp.getBeta(), syrkOp.getType(), op, syrkOp.getC(), other);

    rewriter.replaceOpWithNewOp<enzymexla::SyrkOp>(
        op, syrkOp.getType(), syrkOp.getA(), newC, syrkOp.getAlpha(), newBeta,
        syrkOp.getUploAttr(), syrkOp.getOutputUploAttr(),
        syrkOp.getTransposeAttr());
    return success();
  }
};

struct TransposeSyrkToSyrk
    : public CheckedOpRewritePattern<enzymexla::SyrkOp, TransposeSyrkToSyrk> {
  using CheckedOpRewritePattern<enzymexla::SyrkOp,
                                TransposeSyrkToSyrk>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::SyrkOp op,
                                    PatternRewriter &rewriter) const {
    auto input = op.getA();
    if (cast<RankedTensorType>(input.getType()).getRank() != 2) {
      return failure(); // support only rank 2 matrices for now
    }

    auto transposeOp = input.getDefiningOp<stablehlo::TransposeOp>();
    if (!transposeOp) {
      return failure();
    }

    auto perm = transposeOp.getPermutation();
    if (perm.size() != 2 || perm[0] != 1 || perm[1] != 0) {
      return failure();
    }

    rewriter.replaceOpWithNewOp<enzymexla::SyrkOp>(
        op, op.getResult().getType(), transposeOp.getOperand(), op.getC(),
        op.getAlpha(), op.getBeta(), op.getUploAttr(), op.getOutputUploAttr(),
        enzymexla::LapackTransposeAttr::get(
            op.getContext(),
            enzyme::transposeLapackTranspose(op.getTranspose(), false)));
    return success();
  }
};

struct SyrkSimplifyOutputUplo final
    : public CheckedOpRewritePattern<enzymexla::SyrkOp,
                                     SyrkSimplifyOutputUplo> {
  using CheckedOpRewritePattern<
      enzymexla::SyrkOp, SyrkSimplifyOutputUplo>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::SyrkOp op,
                                    PatternRewriter &rewriter) const {
    // we can also try to align the uplos for other cases but that would be
    // less common
    if (op.getOutputUplo() != enzymexla::LapackUplo::F) {
      return rewriter.notifyMatchFailure(op, "output_uplo is not F");
    }

    // Track all users of syrk. If these end at syrk ops (via supported
    // elementwise operations) we can potentially set a common uplo.
    SmallVector<enzymexla::SyrkOp> childSyrkOps;
    SmallVector<enzymexla::SymmOp> childSymmOps;

    // Worklist approach: track all children that need to be processed
    SmallVector<Value> worklist;
    llvm::SmallPtrSet<Value, 8> visited;

    // Start with the result of the current syrk op
    worklist.push_back(op.getResult());
    visited.insert(op.getResult());

    while (!worklist.empty()) {
      Value current = worklist.pop_back_val();

      // Check all users of the current value
      for (Operation *user : current.getUsers()) {
        // If user is a syrk op, add it to our list of child syrk ops
        if (auto childSyrk = dyn_cast<enzymexla::SyrkOp>(user)) {
          if (childSyrk.getC() != current) {
            return failure();
          }
          childSyrkOps.push_back(childSyrk);
          continue;
        }

        // If user is a symm op, add it to our list of child symm ops
        if (auto childSymm = dyn_cast<enzymexla::SymmOp>(user)) {
          if (childSymm.getA() != current) {
            return failure();
          }
          childSymmOps.push_back(childSymm);
          continue;
        }

        // Check if the operation is a supported elementwise operation
        // that preserves symmetry (elementwise ops preserve the symmetric
        // structure)
        if (!stablehlo::hasTraitElementwise(user)) {
          // Found a non-elementwise user that is not a syrk op
          // This is not supported, so we fail
          return rewriter.notifyMatchFailure(
              op, "found non-elementwise, non-syrk user");
        }

        // For elementwise operations, add their results to the worklist
        for (Value result : user->getResults()) {
          if (!visited.contains(result)) {
            visited.insert(result);
            worklist.push_back(result);
          }
        }
      }
    }

    // If no child syrk ops were found, nothing to optimize
    if (childSyrkOps.empty() && childSymmOps.empty()) {
      return rewriter.notifyMatchFailure(op, "no child syrk ops found");
    }

    // Check the uplos for all child syrk ops
    // Collect all the uplos we see
    bool hasUpper = false, hasLower = false;
    int countOutputUpper = 0, countOutputLower = 0;

    for (enzymexla::SyrkOp childSyrk : childSyrkOps) {
      switch (childSyrk.getUplo()) {
      case enzymexla::LapackUplo::U:
        hasUpper = true;
        break;
      case enzymexla::LapackUplo::L:
        hasLower = true;
        break;
      case enzymexla::LapackUplo::F:
        break;
      }

      switch (childSyrk.getOutputUplo()) {
      case enzymexla::LapackUplo::U:
        countOutputUpper++;
        break;
      case enzymexla::LapackUplo::L:
        countOutputLower++;
        break;
      case enzymexla::LapackUplo::F:
        break;
      }
    }

    for (enzymexla::SymmOp childSymm : childSymmOps) {
      switch (childSymm.getUplo()) {
      case enzymexla::LapackUplo::U:
        hasUpper = true;
        break;
      case enzymexla::LapackUplo::L:
        hasLower = true;
        break;
      case enzymexla::LapackUplo::F:
        break;
      }
    }

    // Check for conflict: if we have both U and L, we cannot satisfy both
    if (hasUpper && hasLower) {
      return rewriter.notifyMatchFailure(
          op, "conflicting uplos among child syrk ops (both U and L)");
    }

    // Determine the common uplo
    enzymexla::LapackUplo newOutputUplo;
    if (hasUpper) {
      // At least one child requires U, set output to U
      newOutputUplo = enzymexla::LapackUplo::U;
    } else if (hasLower) {
      // At least one child requires L, set output to L
      newOutputUplo = enzymexla::LapackUplo::L;
    } else {
      // All children have uplo F
      // Check the output_uplo of each child to find the uplo that minimizes
      // copying. According to LowerEnzymeXLABlas::resolveUplo:
      // - When uplo is F, needsCopy = (output_uplo == F)
      // - If output_uplo is U or L, no copy is needed when input matches

      // Choose the uplo that matches the most children's output_uplo
      // This minimizes the number of copies needed
      if (countOutputUpper > countOutputLower) {
        newOutputUplo = enzymexla::LapackUplo::U;
      } else if (countOutputLower > countOutputUpper) {
        newOutputUplo = enzymexla::LapackUplo::L;
      } else {
        // Tied or all have output_uplo F
        // Default to U as per standardizeUplo (see LowerEnzymeXLABlas)
        newOutputUplo = enzymexla::LapackUplo::U;
      }
    }

    // If the new output uplo is the same as what we already have, no change
    if (newOutputUplo == op.getOutputUplo()) {
      return rewriter.notifyMatchFailure(op, "output_uplo already optimal");
    }

    auto newUploAttr =
        enzymexla::LapackUploAttr::get(rewriter.getContext(), newOutputUplo);

    // Create a new syrk op with the updated output_uplo
    rewriter.modifyOpInPlace(op, [&]() { op.setOutputUploAttr(newUploAttr); });

    // Also update the child syrk ops that have uplo F to use the new uplo
    // This ensures they don't need to copy
    for (enzymexla::SyrkOp &childSyrk : childSyrkOps) {
      rewriter.modifyOpInPlace(childSyrk,
                               [&]() { childSyrk.setUploAttr(newUploAttr); });
    }
    for (enzymexla::SymmOp &childSymm : childSymmOps) {
      rewriter.modifyOpInPlace(childSymm,
                               [&]() { childSymm.setUploAttr(newUploAttr); });
    }
    return success();
  }
};

// if a dimension was created by broadcasting for both lhs and rhs, then we move
// the broadcasting after the dot_general
struct DotGeneralBroadcastInDim
    : public CheckedOpRewritePattern<
          stablehlo::DotGeneralOp,
          DotGeneralBroadcastInDim>::CheckedOpRewritePattern {
  using CheckedOpRewritePattern<
      stablehlo::DotGeneralOp,
      DotGeneralBroadcastInDim>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp dotOp,
                                    PatternRewriter &rewriter) const {
    auto lhsBcast = dotOp.getLhs().getDefiningOp<stablehlo::BroadcastInDimOp>();
    auto rhsBcast = dotOp.getRhs().getDefiningOp<stablehlo::BroadcastInDimOp>();

    if (!lhsBcast || !rhsBcast)
      return failure();

    auto lhsDims = lhsBcast.getBroadcastDimensions();
    auto rhsDims = rhsBcast.getBroadcastDimensions();

    auto lhsInputType = cast<RankedTensorType>(lhsBcast.getOperand().getType());
    auto rhsInputType = cast<RankedTensorType>(rhsBcast.getOperand().getType());
    auto lhsBcastType = cast<RankedTensorType>(lhsBcast.getResult().getType());
    auto rhsBcastType = cast<RankedTensorType>(rhsBcast.getResult().getType());

    auto dotDims = dotOp.getDotDimensionNumbers();
    auto lhsBatchDims(dotDims.getLhsBatchingDimensions());
    auto rhsBatchDims(dotDims.getRhsBatchingDimensions());
    auto lhsContractDims(dotDims.getLhsContractingDimensions());
    auto rhsContractDims(dotDims.getRhsContractingDimensions());

    SmallVector<int64_t> lhsCreatedBatchDims, rhsCreatedBatchDims,
        lhsRealBatchDims, rhsRealBatchDims;
    llvm::SmallDenseSet<int64_t> dimsToRemove;
    for (size_t i = 0; i < lhsBatchDims.size(); ++i) {
      int64_t lhsDim = lhsBatchDims[i];
      int64_t rhsDim = rhsBatchDims[i];

      // Check if this batch dimension was created by broadcast
      bool lhsIsNew =
          std::find(lhsDims.begin(), lhsDims.end(), lhsDim) == lhsDims.end();
      bool rhsIsNew =
          std::find(rhsDims.begin(), rhsDims.end(), rhsDim) == rhsDims.end();

      if (lhsIsNew && rhsIsNew) {
        dimsToRemove.insert(i);
        lhsCreatedBatchDims.push_back(lhsDim);
        rhsCreatedBatchDims.push_back(rhsDim);
      } else {
        lhsRealBatchDims.push_back(lhsDim);
        rhsRealBatchDims.push_back(rhsDim);
      }
    }

    if (lhsCreatedBatchDims.empty())
      return failure();

    SmallVector<int64_t> newDotShape;
    auto dotResultType = cast<RankedTensorType>(dotOp.getResult().getType());
    for (size_t i = 0; i < dotResultType.getRank(); ++i) {
      if (llvm::is_contained(dimsToRemove, i))
        continue;
      newDotShape.push_back(dotResultType.getDimSize(i));
    }

    SmallVector<int64_t> lhsNewBcastDims, rhsNewBcastDims;
    DenseMap<int64_t, int64_t> lhsOldToNew, rhsOldToNew;
    auto newLhsBcastType =
        getNewBroadcastType(lhsBcastType, lhsInputType, lhsDims,
                            lhsCreatedBatchDims, lhsNewBcastDims, lhsOldToNew);
    auto newRhsBcastType =
        getNewBroadcastType(rhsBcastType, rhsInputType, rhsDims,
                            rhsCreatedBatchDims, rhsNewBcastDims, rhsOldToNew);

    auto newLhsBcast = stablehlo::BroadcastInDimOp::create(
        rewriter, dotOp.getLoc(), newLhsBcastType, lhsBcast.getOperand(),
        rewriter.getDenseI64ArrayAttr(lhsNewBcastDims));
    auto newRhsBcast = stablehlo::BroadcastInDimOp::create(
        rewriter, dotOp.getLoc(), newRhsBcastType, rhsBcast.getOperand(),
        rewriter.getDenseI64ArrayAttr(rhsNewBcastDims));

    SmallVector<int64_t> newLhsContractDims, newRhsContractDims,
        newLhsBatchDims, newRhsBatchDims;
    for (auto dim : lhsContractDims) {
      newLhsContractDims.push_back(lhsOldToNew[dim]);
    }
    for (auto dim : rhsContractDims) {
      newRhsContractDims.push_back(rhsOldToNew[dim]);
    }
    for (auto dim : lhsRealBatchDims) {
      newLhsBatchDims.push_back(lhsOldToNew[dim]);
    }
    for (auto dim : rhsRealBatchDims) {
      newRhsBatchDims.push_back(rhsOldToNew[dim]);
    }

    auto newDotDims = DotDimensionNumbersAttr::get(
        rewriter.getContext(), newLhsBatchDims, newRhsBatchDims,
        newLhsContractDims, newRhsContractDims);
    auto newDotOp = stablehlo::DotGeneralOp::create(
        rewriter, dotOp.getLoc(),
        RankedTensorType::get(newDotShape, dotResultType.getElementType()),
        newLhsBcast, newRhsBcast, newDotDims, dotOp.getPrecisionConfigAttr(),
        dotOp.getAlgorithmAttr());

    SmallVector<int64_t> newMapping;
    for (size_t i = 0; i < cast<RankedTensorType>(dotOp.getType()).getRank();
         ++i) {
      if (llvm::is_contained(dimsToRemove, i))
        continue;
      newMapping.push_back(i);
    }
    rewriter.replaceOpWithNewOp<stablehlo::BroadcastInDimOp>(
        dotOp, dotOp.getType(), newDotOp.getResult(),
        rewriter.getDenseI64ArrayAttr(newMapping));

    return success();
  }

private:
  RankedTensorType
  getNewBroadcastType(RankedTensorType oldOutputType,
                      RankedTensorType inputType, ArrayRef<int64_t> oldMapping,
                      ArrayRef<int64_t> removeDims,
                      SmallVectorImpl<int64_t> &newMapping,
                      DenseMap<int64_t, int64_t> &oldToNew) const {
    SmallVector<int64_t> newOutShape;
    for (size_t i = 0; i < oldOutputType.getRank(); ++i) {
      if (llvm::is_contained(removeDims, i))
        continue;
      oldToNew[i] = newOutShape.size();
      newOutShape.push_back(oldOutputType.getDimSize(i));
    }
    for (auto dim : oldMapping) {
      auto numDel =
          llvm::count_if(removeDims, [&](int64_t d) { return d < dim; });
      newMapping.push_back(dim - numDel);
    }
    return RankedTensorType::get(newOutShape, oldOutputType.getElementType());
  }
};

// Shared helper function to compute the transpose permutation needed after
// changing dot_general operand dimensions.
//
// This computes the permutation that, when applied to the new dot_general
// result, restores the original output dimension ordering.
//
// Parameters:
// - origLhsType, origRhsType: original operand types (before transformation)
// - origDotDims: original dot dimension numbers
// - newLhsBatchDims, newRhsBatchDims: new batch dimensions
// - newLhsContractDims, newRhsContractDims: new contracting dimensions
// - lhsOldToNew, rhsOldToNew: maps from original operand dims to new dims
// - newLhsType, newRhsType: new operand types (after transformation)
//
// Returns the permutation vector to apply via transpose.
static SmallVector<int64_t> computeDotGeneralOutputPermutation(
    RankedTensorType origLhsType, RankedTensorType origRhsType,
    stablehlo::DotDimensionNumbersAttr origDotDims,
    ArrayRef<int64_t> newLhsBatchDims, ArrayRef<int64_t> newRhsBatchDims,
    ArrayRef<int64_t> newLhsContractDims, ArrayRef<int64_t> newRhsContractDims,
    const DenseMap<int64_t, int64_t> &lhsOldToNew,
    const DenseMap<int64_t, int64_t> &rhsOldToNew, RankedTensorType newLhsType,
    RankedTensorType newRhsType) {

  auto origLhsBatchDims = origDotDims.getLhsBatchingDimensions();
  auto origRhsBatchDims = origDotDims.getRhsBatchingDimensions();
  auto origLhsContractDims = origDotDims.getLhsContractingDimensions();
  auto origRhsContractDims = origDotDims.getRhsContractingDimensions();

  int64_t numBatch = origLhsBatchDims.size();
  int64_t origResultRank =
      numBatch +
      (origLhsType.getRank() - numBatch - origLhsContractDims.size()) +
      (origRhsType.getRank() - numBatch - origRhsContractDims.size());

  // Build mapping from original lhs/rhs dim -> original output position
  DenseMap<int64_t, int64_t> origLhsDimToOutputPos;
  DenseMap<int64_t, int64_t> origRhsDimToOutputPos;

  // Batch dims come first in output, in order of lhs batch dims
  for (int64_t i = 0; i < numBatch; ++i) {
    origLhsDimToOutputPos[origLhsBatchDims[i]] = i;
    origRhsDimToOutputPos[origRhsBatchDims[i]] = i;
  }

  // LHS remaining dims (non-batch, non-contract) come next
  int64_t outputPos = numBatch;
  for (int64_t i = 0; i < origLhsType.getRank(); ++i) {
    if (!llvm::is_contained(origLhsBatchDims, i) &&
        !llvm::is_contained(origLhsContractDims, i)) {
      origLhsDimToOutputPos[i] = outputPos++;
    }
  }

  // RHS remaining dims come after LHS remaining
  for (int64_t i = 0; i < origRhsType.getRank(); ++i) {
    if (!llvm::is_contained(origRhsBatchDims, i) &&
        !llvm::is_contained(origRhsContractDims, i)) {
      origRhsDimToOutputPos[i] = outputPos++;
    }
  }

  // Build reverse maps: new dim -> old dim
  DenseMap<int64_t, int64_t> lhsNewToOld, rhsNewToOld;
  for (auto &[orig, newDim] : lhsOldToNew) {
    lhsNewToOld[newDim] = orig;
  }
  for (auto &[orig, newDim] : rhsOldToNew) {
    rhsNewToOld[newDim] = orig;
  }

  // Build the new output position -> original output position mapping
  SmallVector<int64_t> newToOrigOutput(origResultRank);
  int64_t newOutputPos = 0;

  // New batch dims (in order of new lhs batch dims)
  for (int64_t newLhsBatchDim : newLhsBatchDims) {
    int64_t origLhsDim = lhsNewToOld[newLhsBatchDim];
    // Check if this was originally a batch dim or something else
    if (origLhsDimToOutputPos.count(origLhsDim)) {
      newToOrigOutput[newOutputPos++] = origLhsDimToOutputPos[origLhsDim];
    }
  }

  // New LHS remaining dims
  for (int64_t i = 0; i < newLhsType.getRank(); ++i) {
    if (!llvm::is_contained(newLhsBatchDims, i) &&
        !llvm::is_contained(newLhsContractDims, i)) {
      int64_t origLhsDim = lhsNewToOld[i];
      newToOrigOutput[newOutputPos++] = origLhsDimToOutputPos[origLhsDim];
    }
  }

  // New RHS remaining dims
  for (int64_t i = 0; i < newRhsType.getRank(); ++i) {
    if (!llvm::is_contained(newRhsBatchDims, i) &&
        !llvm::is_contained(newRhsContractDims, i)) {
      int64_t origRhsDim = rhsNewToOld[i];
      newToOrigOutput[newOutputPos++] = origRhsDimToOutputPos[origRhsDim];
    }
  }

  // Compute inverse permutation: perm[orig_pos] = new_pos
  SmallVector<int64_t> perm(origResultRank);
  for (int64_t i = 0; i < origResultRank; ++i) {
    perm[newToOrigOutput[i]] = i;
  }

  return perm;
}

struct DotGeneralBroadcastInDimSortDims
    : public CheckedOpRewritePattern<
          stablehlo::DotGeneralOp,
          DotGeneralBroadcastInDimSortDims>::CheckedOpRewritePattern {
  using CheckedOpRewritePattern<
      stablehlo::DotGeneralOp,
      DotGeneralBroadcastInDimSortDims>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp dotOp,
                                    PatternRewriter &rewriter) const {
    auto lhsBcast = dotOp.getLhs().getDefiningOp<stablehlo::BroadcastInDimOp>();
    auto rhsBcast = dotOp.getRhs().getDefiningOp<stablehlo::BroadcastInDimOp>();

    if (!lhsBcast && !rhsBcast)
      return failure();

    // if any of the operands has an implicit transpose, we absorb it into the
    // dot general op
    Value lhs = dotOp.getLhs(), rhs = dotOp.getRhs();
    bool changed = false;

    auto dotDims = dotOp.getDotDimensionNumbers();
    SmallVector<int64_t> newLhsBatchDims(dotDims.getLhsBatchingDimensions());
    SmallVector<int64_t> newLhsContractDims(
        dotDims.getLhsContractingDimensions());
    SmallVector<int64_t> newRhsBatchDims(dotDims.getRhsBatchingDimensions());
    SmallVector<int64_t> newRhsContractDims(
        dotDims.getRhsContractingDimensions());

    // Track the old-to-new dimension mappings for computing output transpose
    DenseMap<int64_t, int64_t> lhsOldToNew, rhsOldToNew;
    auto lhsType = cast<RankedTensorType>(dotOp.getLhs().getType());
    auto rhsType = cast<RankedTensorType>(dotOp.getRhs().getType());

    // Initialize to identity mapping
    for (int64_t i = 0; i < lhsType.getRank(); ++i)
      lhsOldToNew[i] = i;
    for (int64_t i = 0; i < rhsType.getRank(); ++i)
      rhsOldToNew[i] = i;

    if (lhsBcast && !llvm::is_sorted(lhsBcast.getBroadcastDimensions())) {
      changed = true;
      lhs = getNewBroadcastWithSortedDims(rewriter, lhsBcast, newLhsBatchDims,
                                          newLhsContractDims, lhsOldToNew);
    }

    if (rhsBcast && !llvm::is_sorted(rhsBcast.getBroadcastDimensions())) {
      changed = true;
      rhs = getNewBroadcastWithSortedDims(rewriter, rhsBcast, newRhsBatchDims,
                                          newRhsContractDims, rhsOldToNew);
    }

    if (changed) {
      // Compute the new result type based on the new operand types
      auto newLhsType = cast<RankedTensorType>(lhs.getType());
      auto newRhsType = cast<RankedTensorType>(rhs.getType());
      auto origResultType = cast<RankedTensorType>(dotOp.getType());

      auto newDotDims = stablehlo::DotDimensionNumbersAttr::get(
          rewriter.getContext(), newLhsBatchDims, newRhsBatchDims,
          newLhsContractDims, newRhsContractDims);

      auto newResultType = stablehlo::GetDotGeneralResultType(
          lhs, rhs, origResultType.getElementType(), newDotDims);

      auto newDotOp = stablehlo::DotGeneralOp::create(
          rewriter, dotOp.getLoc(), newResultType, lhs, rhs, newDotDims,
          dotOp.getPrecisionConfigAttr(), dotOp.getAlgorithmAttr());

      // Compute transpose permutation using shared helper
      auto perm = computeDotGeneralOutputPermutation(
          lhsType, rhsType, dotDims, newLhsBatchDims, newRhsBatchDims,
          newLhsContractDims, newRhsContractDims, lhsOldToNew, rhsOldToNew,
          newLhsType, newRhsType);

      Value transposed = TransposeOpCreate(rewriter, dotOp.getLoc(),
                                           newDotOp.getResult(), perm);
      rewriter.replaceOp(dotOp, transposed);
      return success();
    }
    return failure();
  }

private:
  Value
  getNewBroadcastWithSortedDims(PatternRewriter &rewriter,
                                stablehlo::BroadcastInDimOp oldBroadcast,
                                SmallVectorImpl<int64_t> &newBatchDims,
                                SmallVectorImpl<int64_t> &newContractDims,
                                DenseMap<int64_t, int64_t> &oldToNewOut) const {
    auto inputType =
        cast<RankedTensorType>(oldBroadcast.getOperand().getType());
    auto outputType = cast<RankedTensorType>(oldBroadcast.getType());
    SmallVector<int64_t> oldToNew(outputType.getRank());
    int64_t newDimsStart = inputType.getRank();
    SmallVector<int64_t> newOutputShape(outputType.getRank());
    for (size_t i = 0; i < outputType.getRank(); ++i) {
      if (!llvm::is_contained(oldBroadcast.getBroadcastDimensions(), i)) {
        oldToNew[i] = newDimsStart++;
        newOutputShape[oldToNew[i]] = outputType.getDimSize(i);
        continue;
      }
    }
    for (auto [i, dim] :
         llvm::enumerate(oldBroadcast.getBroadcastDimensions())) {
      oldToNew[dim] = i;
      newOutputShape[i] = outputType.getDimSize(dim);
    }

    for (auto &dim : newBatchDims)
      dim = oldToNew[dim];
    for (auto &dim : newContractDims)
      dim = oldToNew[dim];

    // Populate the output oldToNew mapping for the caller
    for (size_t i = 0; i < oldToNew.size(); ++i) {
      oldToNewOut[i] = oldToNew[i];
    }

    SmallVector<int64_t> newMapping(inputType.getRank());
    std::iota(newMapping.begin(), newMapping.end(), 0);
    return stablehlo::BroadcastInDimOp::create(
        rewriter, oldBroadcast.getLoc(),
        RankedTensorType::get(newOutputShape, inputType.getElementType()),
        oldBroadcast.getOperand(), rewriter.getDenseI64ArrayAttr(newMapping));
  }
};

// If a dot general batches over a dimension that was newly created for exactly
// one operand, we can remove batching over that dimension, we do need to make
// sure to insert a transpose to get the correct dimension ordering.
//
// For example:
//
// %x : tensor<64x32x1x64xf32>
// %y = bcast %arg3, dims = [0, 1] : (tensor<64x32xf32>) ->
//                                    tensor<64x32x64x1xf32>
// dot_general %x, %y, batching_dims = [0, 3] x [2, 0],
//                     contracting_dims = [1, 2] x [1, 3]
//
// since dimension 2 for %y was created by a broadcast while %x's dimension 0
// existed from before, we can remove that dimension of %y and the corresponding
// batching_dims.
//
// %x : tensor<64x32x1x64xf32>
// %y = bcast %arg3, dims = [0, 1] : (tensor<64x32xf32>) -> tensor<64x32x1xf32>
// %z = dot_general %x, %y, batching_dims = [3] x [0], contracting_dims = [1, 2]
//      x [1, 3]
// %res = transpose %z to bring the dimensions to match original
//        ordering
struct DotGeneralRemoveBatchDimensions
    : public CheckedOpRewritePattern<
          stablehlo::DotGeneralOp,
          DotGeneralRemoveBatchDimensions>::CheckedOpRewritePattern {
  using CheckedOpRewritePattern<
      stablehlo::DotGeneralOp,
      DotGeneralRemoveBatchDimensions>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp dotOp,
                                    PatternRewriter &rewriter) const {
    auto lhsBcast = dotOp.getLhs().getDefiningOp<stablehlo::BroadcastInDimOp>();
    auto rhsBcast = dotOp.getRhs().getDefiningOp<stablehlo::BroadcastInDimOp>();

    // At least one operand must have a broadcast
    // while we could do the same for splatted tensors, we apply
    // DotGeneralSimplify to those
    if (!lhsBcast && !rhsBcast) {
      return failure();
    }

    ArrayRef<int64_t> lhsDims =
        lhsBcast ? lhsBcast.getBroadcastDimensions() : ArrayRef<int64_t>();
    ArrayRef<int64_t> rhsDims =
        rhsBcast ? rhsBcast.getBroadcastDimensions() : ArrayRef<int64_t>();

    auto lhsType = cast<RankedTensorType>(dotOp.getLhs().getType());
    auto rhsType = cast<RankedTensorType>(dotOp.getRhs().getType());

    auto dotDims = dotOp.getDotDimensionNumbers();
    auto lhsBatchDims = dotDims.getLhsBatchingDimensions();
    auto rhsBatchDims = dotDims.getRhsBatchingDimensions();
    auto lhsContractDims = dotDims.getLhsContractingDimensions();
    auto rhsContractDims = dotDims.getRhsContractingDimensions();

    // Find batch dimensions where exactly one side has a created (broadcast)
    // dimension.
    // lhsCreatedBatchIndices: indices in the batch dim arrays where lhs has
    // created dim
    // rhsCreatedBatchIndices: indices in the batch dim arrays where rhs has
    // created dim
    SmallVector<int64_t> lhsOnlyCreatedBatchIndices, rhsOnlyCreatedBatchIndices,
        lhsCreatedBatchDims, rhsCreatedBatchDims;

    for (int64_t i = 0; i < lhsBatchDims.size(); ++i) {
      int64_t lhsDim = lhsBatchDims[i];
      int64_t rhsDim = rhsBatchDims[i];

      // Check if this batch dimension was created by broadcast (not in the
      // original input)
      bool lhsIsNew = lhsBcast && !llvm::is_contained(lhsDims, lhsDim);
      bool rhsIsNew = rhsBcast && !llvm::is_contained(rhsDims, rhsDim);

      // We only handle the case where exactly one side has a created dimension
      if (lhsIsNew && !rhsIsNew) {
        // LHS has created dimension, RHS has real dimension
        lhsOnlyCreatedBatchIndices.push_back(i);
        lhsCreatedBatchDims.push_back(lhsDim);
      } else if (!lhsIsNew && rhsIsNew) {
        // RHS has created dimension, LHS has real dimension
        rhsOnlyCreatedBatchIndices.push_back(i);
        rhsCreatedBatchDims.push_back(rhsDim);
      }
    }

    if (lhsCreatedBatchDims.empty() && rhsCreatedBatchDims.empty()) {
      return failure();
    }

    // Build new broadcast operation for the side with created dimensions
    Value newLhs = dotOp.getLhs();
    Value newRhs = dotOp.getRhs();
    DenseMap<int64_t, int64_t> lhsOldToNew, rhsOldToNew, lhsNewToOld,
        rhsNewToOld;

    if (lhsBcast && !lhsCreatedBatchDims.empty()) {
      newLhs = createReducedBroadcast(rewriter, dotOp.getLoc(), lhsBcast,
                                      lhsType, lhsCreatedBatchDims, lhsDims,
                                      lhsOldToNew, lhsNewToOld);
    } else {
      for (int64_t i = 0; i < lhsType.getRank(); ++i) {
        lhsOldToNew[i] = i;
        lhsNewToOld[i] = i;
      }
    }

    if (rhsBcast && !rhsCreatedBatchDims.empty()) {
      newRhs = createReducedBroadcast(rewriter, dotOp.getLoc(), rhsBcast,
                                      rhsType, rhsCreatedBatchDims, rhsDims,
                                      rhsOldToNew, rhsNewToOld);
    } else {
      for (int64_t i = 0; i < rhsType.getRank(); ++i) {
        rhsOldToNew[i] = i;
        rhsNewToOld[i] = i;
      }
    }

    // Build new dimension numbers, removing the batch dimensions that had
    // created dims
    SmallVector<int64_t> newLhsBatchDims, newRhsBatchDims;
    for (int64_t i = 0; i < lhsBatchDims.size(); ++i) {
      if (llvm::is_contained(lhsOnlyCreatedBatchIndices, i) ||
          llvm::is_contained(rhsOnlyCreatedBatchIndices, i)) {
        continue;
      }
      newLhsBatchDims.push_back(lhsOldToNew[lhsBatchDims[i]]);
      newRhsBatchDims.push_back(rhsOldToNew[rhsBatchDims[i]]);
    }

    SmallVector<int64_t> newLhsContractDims, newRhsContractDims;
    for (auto dim : lhsContractDims) {
      newLhsContractDims.push_back(lhsOldToNew[dim]);
    }
    for (auto dim : rhsContractDims) {
      newRhsContractDims.push_back(rhsOldToNew[dim]);
    }

    auto newDotDims = stablehlo::DotDimensionNumbersAttr::get(
        rewriter.getContext(), newLhsBatchDims, newRhsBatchDims,
        newLhsContractDims, newRhsContractDims);
    auto oldDotResultType = cast<RankedTensorType>(dotOp.getType());

    auto newLhsType = cast<RankedTensorType>(newLhs.getType());
    auto newRhsType = cast<RankedTensorType>(newRhs.getType());

    auto newDotOp = stablehlo::DotGeneralOp::create(
        rewriter, dotOp.getLoc(),
        stablehlo::GetDotGeneralResultType(
            newLhs, newRhs, oldDotResultType.getElementType(), newDotDims),
        newLhs, newRhs, newDotDims, dotOp.getPrecisionConfigAttr(),
        dotOp.getAlgorithmAttr());

    // Compute transpose permutation using shared helper
    auto perm = computeDotGeneralOutputPermutation(
        lhsType, rhsType, dotDims, newLhsBatchDims, newRhsBatchDims,
        newLhsContractDims, newRhsContractDims, lhsOldToNew, rhsOldToNew,
        newLhsType, newRhsType);

    Value transposed =
        TransposeOpCreate(rewriter, dotOp.getLoc(), newDotOp.getResult(), perm);
    rewriter.replaceOp(dotOp, transposed);
    return success();
  }

private:
  // Helper function to create a new broadcast with certain dimensions removed.
  // Updates the oldToNew map to reflect the dimension remapping.
  Value createReducedBroadcast(PatternRewriter &rewriter, Location loc,
                               stablehlo::BroadcastInDimOp bcast,
                               RankedTensorType origType,
                               ArrayRef<int64_t> dimsToRemove,
                               ArrayRef<int64_t> origBcastDims,
                               DenseMap<int64_t, int64_t> &oldToNew,
                               DenseMap<int64_t, int64_t> &newToOld) const {
    SmallVector<int64_t> newOutputShape;
    SmallVector<int64_t> newBcastDims;
    oldToNew.clear();

    for (int64_t i = 0; i < origType.getRank(); ++i) {
      if (llvm::is_contained(dimsToRemove, i)) {
        continue;
      }
      int64_t dim = newOutputShape.size();
      oldToNew[i] = dim;
      newToOld[dim] = i;
      newOutputShape.push_back(origType.getDimSize(i));
    }

    for (auto dim : origBcastDims) {
      if (llvm::is_contained(dimsToRemove, dim)) {
        continue;
      }
      newBcastDims.push_back(oldToNew[dim]);
    }

    auto newType =
        RankedTensorType::get(newOutputShape, origType.getElementType());
    return stablehlo::BroadcastInDimOp::create(
        rewriter, loc, newType, bcast.getOperand(),
        rewriter.getDenseI64ArrayAttr(newBcastDims));
  }
};

// If there is some overlap between the update and the ds, then we can
// replace the ds with a slice op (with a potentially padding)
LogicalResult DUSDSSimplifyWithSomeUpdateOverlapHelper(
    stablehlo::DynamicSliceOp dsOp, stablehlo::DynamicUpdateSliceOp dusOp,
    PatternRewriter &rewriter, std::optional<WhileLoopInfo> loopInfo) {
  auto dsStartIndices = dsOp.getStartIndices();
  auto dusStartIndices = dusOp.getStartIndices();

  SmallVector<int64_t> sliceStarts(dusOp.getStartIndices().size());
  SmallVector<Value> dynamicSliceStarts(dusOp.getStartIndices().size());
  bool canReplace = true, needsDynamicSlice = false;

  llvm::MapVector<Value, WhileLoopInfo::AffineIndexInfo> affineIndexInfoMap;
  if (loopInfo.has_value()) {
    affineIndexInfoMap = loopInfo.value().getAffineIndexInfo();
  }

  auto dusOperandTy = cast<RankedTensorType>(dusOp.getOperand().getType());
  auto dusOperandShape = llvm::to_vector(dusOperandTy.getShape());
  auto updateTy = cast<RankedTensorType>(dusOp.getUpdate().getType());
  auto updateShape = llvm::to_vector(updateTy.getShape());
  auto dsSliceSizes = llvm::to_vector(dsOp.getSliceSizes());

  for (size_t i = 0; i < dusOp.getStartIndices().size(); ++i) {
    auto dsStart = dsStartIndices[i];
    auto dusStart = dusStartIndices[i];

    if (dsStart == dusStart) {
      sliceStarts[i] = 0;
      continue;
    } else {
      APInt dsStartAP, dusStartAP;
      bool dusStartIsConstant =
          matchPattern(dusStart, m_ConstantInt(&dusStartAP));
      if (matchPattern(dsStart, m_ConstantInt(&dsStartAP)) &&
          dusStartIsConstant) {
        int64_t dsStartInt = dsStartAP.getSExtValue();
        int64_t dusStartInt = dusStartAP.getSExtValue();

        if (dusStartInt + updateShape[i] <= dsStartInt) { // no overlap
          canReplace = false;
          break;
        }

        sliceStarts[i] = dsStartInt - dusStartInt;
        continue;
      }

      if (loopInfo.has_value()) {
        if (affineIndexInfoMap.contains(dsStart) &&
            affineIndexInfoMap.contains(dusStart)) {
          auto dsIndexInfo = affineIndexInfoMap[dsStart];
          auto dusIndexInfo = affineIndexInfoMap[dusStart];
          if (dsIndexInfo.scale.isOne() && dusIndexInfo.scale.isOne()) {
            sliceStarts[i] = dsIndexInfo.offset.getSExtValue() -
                             dusIndexInfo.offset.getSExtValue();
            continue;
          }
        }

        if (dusStartIsConstant) {
          int64_t dusStartInt = dusStartAP.getSExtValue();

          // if the DUS is doing a full update along this dimension
          // then it doesn't matter what the exact dynamic slice index is
          if (dusStartInt == 0 && updateShape[i] == dusOperandShape[i]) {
            needsDynamicSlice = true;
            dynamicSliceStarts[i] = dsStart;
            continue;
          }

          LLVM_DEBUG(
              auto dsAccessBounds = loopInfo.value().getBounds(dsStart);
              if (dsAccessBounds.has_value()) {
                needsDynamicSlice = true;
                llvm::dbgs() << "  ds bound: [" << dsAccessBounds.value().min
                             << ", " << dsAccessBounds.value().max << ")\n";
                dsOp->emitError("TODO: support this case where we check that "
                                "all of the values are in the range for DUS");
                return failure();
              });
        }
      }
    }

    canReplace = false;
    break;
  }

  if (!canReplace) {
    return failure();
  }

  if (needsDynamicSlice) {
    // TODO: maybe we can support the padding case here as well? ignore for now

    for (size_t i = 0; i < dusOp.getStartIndices().size(); ++i) {
      if (!dynamicSliceStarts[i]) {
        if (sliceStarts[i] < 0) {
          return failure();
        }
        int64_t limit = dsSliceSizes[i] + sliceStarts[i];
        if (limit > updateShape[i]) {
          return failure();
        }
      }
    }

    // now we know we can emit a DS op
    auto dsIndexTy =
        cast<RankedTensorType>(dsOp.getStartIndices()[0].getType());

    for (size_t i = 0; i < dusOp.getStartIndices().size(); ++i) {
      if (!dynamicSliceStarts[i]) {
        dynamicSliceStarts[i] = stablehlo::ConstantOp::create(
            rewriter, dsOp->getLoc(),
            cast<ElementsAttr>(makeAttr(dsIndexTy, sliceStarts[i])));
      }
    }

    rewriter.setInsertionPoint(dsOp);
    auto newDS = stablehlo::DynamicSliceOpCreate(
        rewriter, dsOp.getLoc(), dusOp.getUpdate(), dynamicSliceStarts,
        dsOp.getSliceSizes());
    rewriter.replaceOp(dsOp, newDS);
    return success();
  }

  bool allOffsetsZero =
      llvm::all_of(sliceStarts, [](int64_t offset) { return offset == 0; });

  // simple case
  if (allOffsetsZero && updateShape == dsSliceSizes) {
    rewriter.replaceOp(dsOp, dusOp.getUpdate());
    return success();
  }

  SmallVector<int64_t> sliceLimits(dusOp.getStartIndices().size());
  SmallVector<int64_t> sliceStrides(dusOp.getStartIndices().size(), 1);
  SmallVector<int64_t> padLow(dusOp.getStartIndices().size(), 0);
  SmallVector<int64_t> padHigh(dusOp.getStartIndices().size(), 0);
  SmallVector<int64_t> padInner(dusOp.getStartIndices().size(), 0);
  bool needsPad = false;

  for (size_t i = 0; i < dusOp.getStartIndices().size(); ++i) {
    if (sliceStarts[i] < 0) {
      needsPad = true;
      padLow[i] = -sliceStarts[i];
      sliceStarts[i] = 0;
    }
    int64_t limit = dsSliceSizes[i] + sliceStarts[i] - padLow[i];
    if (limit > updateShape[i]) {
      needsPad = true;
      padHigh[i] = limit - updateShape[i];
      limit = updateShape[i];
    }
    sliceLimits[i] = limit;
  }

  if (needsPad && !stablehlo::isScalarValue(dusOp.getOperand())) {
    return failure();
  }

  rewriter.setInsertionPoint(dsOp);
  Value result =
      stablehlo::SliceOp::create(rewriter, dsOp.getLoc(), dusOp.getUpdate(),
                                 sliceStarts, sliceLimits, sliceStrides);
  if (needsPad) {
    auto padValue = stablehlo::getScalarValue(dusOp.getOperand(), rewriter);
    assert(padValue);
    result = stablehlo::PadOp::create(rewriter, dsOp.getLoc(), result, padValue,
                                      rewriter.getDenseI64ArrayAttr(padLow),
                                      rewriter.getDenseI64ArrayAttr(padHigh),
                                      rewriter.getDenseI64ArrayAttr(padInner));
  }
  rewriter.replaceOp(dsOp, result);
  return success();
}

// if we are slicing a region that has no overlap with the update, we can
// simply replace with a dynamic slice of the original value
LogicalResult DUSDSSimplifyWithNoUpdateOverlapHelper(
    stablehlo::DynamicSliceOp dsOp, stablehlo::DynamicUpdateSliceOp dusOp,
    PatternRewriter &rewriter, std::optional<WhileLoopInfo> loopInfo) {
  auto dsStartIndices = dsOp.getStartIndices();
  auto dusStartIndices = dusOp.getStartIndices();

  auto sliceSizes = dsOp.getSliceSizes();
  auto updateTy = cast<RankedTensorType>(dusOp.getUpdate().getType());
  auto updateShape = llvm::to_vector(updateTy.getShape());

  llvm::MapVector<Value, WhileLoopInfo::AffineIndexInfo> affineIndexInfoMap;
  if (loopInfo.has_value()) {
    affineIndexInfoMap = loopInfo.value().getAffineIndexInfo();
  }

  // TODO: support the no overlap cases with bounds checking as well

  bool overlapWithUpdate = true;
  for (size_t i = 0; i < dusStartIndices.size(); ++i) {
    auto dsStart = dsStartIndices[i];
    auto dusStart = dusStartIndices[i];

    if (dsStart == dusStart) {
      continue;
    }

    APInt dsStartAP, dusStartAP;
    if (matchPattern(dsStart, m_ConstantInt(&dsStartAP)) &&
        matchPattern(dusStart, m_ConstantInt(&dusStartAP))) {
      if ((dsStartAP.getSExtValue() + // ds ends before dus starts
               sliceSizes[i] <=
           dusStartAP.getSExtValue()) ||
          (dusStartAP.getSExtValue() + // dus ends before ds starts
               updateTy.getDimSize(i) <=
           dsStartAP.getSExtValue())) {
        overlapWithUpdate = false;
        break;
      }
    } else {
      if (loopInfo.has_value()) {
        if (affineIndexInfoMap.contains(dsStart) &&
            affineIndexInfoMap.contains(dusStart)) {
          auto dsIndexInfo = affineIndexInfoMap[dsStart];
          auto dusIndexInfo = affineIndexInfoMap[dusStart];
          if (dsIndexInfo.scale.isOne() && dusIndexInfo.scale.isOne()) {
            if (dsIndexInfo.offset.getSExtValue() + // ds ends before dus starts
                        sliceSizes[i] <=
                    dusIndexInfo.offset.getSExtValue() ||
                (dusIndexInfo.offset
                         .getSExtValue() + // dus ends before ds starts
                     updateTy.getDimSize(i) <=
                 dsIndexInfo.offset.getSExtValue())) {
              overlapWithUpdate = false;
              break;
            }
          }
        }
      }
    }
  }

  if (overlapWithUpdate) {
    return failure();
  }

  rewriter.modifyOpInPlace(dsOp,
                           [&]() { dsOp->setOperand(0, dusOp.getOperand()); });
  return success();
}

// meta function to consider all possible cases
LogicalResult
DUSDSSimplifyHelper(stablehlo::DynamicSliceOp dsOp,
                    stablehlo::DynamicUpdateSliceOp dusOp,
                    PatternRewriter &rewriter,
                    std::optional<WhileLoopInfo> loopInfo = std::nullopt) {
  auto rewriteWithSomeUpdateOverlap =
      DUSDSSimplifyWithSomeUpdateOverlapHelper(dsOp, dusOp, rewriter, loopInfo);
  if (rewriteWithSomeUpdateOverlap.succeeded()) {
    return rewriteWithSomeUpdateOverlap;
  }
  return DUSDSSimplifyWithNoUpdateOverlapHelper(dsOp, dusOp, rewriter,
                                                loopInfo);
}

struct DUSDynamicSliceSimplify final
    : public CheckedOpRewritePattern<stablehlo::DynamicSliceOp,
                                     DUSDynamicSliceSimplify> {
  using CheckedOpRewritePattern<
      stablehlo::DynamicSliceOp,
      DUSDynamicSliceSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicSliceOp op,
                                    PatternRewriter &rewriter) const {
    auto dusOp =
        op.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
    if (!dusOp) {
      return failure();
    }

    return DUSDSSimplifyHelper(op, dusOp, rewriter);
  }
};

struct WhileDUSDSSimplify final
    : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileDUSDSSimplify> {
  using CheckedOpRewritePattern<stablehlo::WhileOp,
                                WhileDUSDSSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp op,
                                    PatternRewriter &rewriter) const {
    WhileLoopInfo info(op);
    if (info.computeInfo().failed() || !info.isConstantStep() ||
        info.getConstantStep() == 0) {
      return failure();
    }
    info.propagateAffineIndexInfo(); // we don't care about computeInfo here
    info.propagateBounds();

    bool anyRewritten = false;

    op.getBody().walk([&](stablehlo::DynamicSliceOp dsOp) {
      auto dusOp =
          dsOp.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
      if (!dusOp) {
        return WalkResult::advance();
      }

      rewriter.setInsertionPoint(dsOp);
      if (DUSDSSimplifyHelper(dsOp, dusOp, rewriter, info).succeeded()) {
        anyRewritten = true;
      }

      return WalkResult::advance();
    });

    return anyRewritten ? success() : failure();
  }
};

// Helper function for (While)DUSDUSSimplify
// Analyzes two DynamicUpdateSlice operations to determine their overlap
// relationship
LogicalResult
DUSDUSSimplifyHelper(stablehlo::DynamicUpdateSliceOp outerDUS,
                     stablehlo::DynamicUpdateSliceOp innerDUS,
                     PatternRewriter &rewriter,
                     std::optional<WhileLoopInfo> loopInfo = std::nullopt) {
  auto outerStartIndices = outerDUS.getStartIndices();
  auto innerStartIndices = innerDUS.getStartIndices();

  auto outerUpdateTy = cast<RankedTensorType>(outerDUS.getUpdate().getType());
  auto innerUpdateTy = cast<RankedTensorType>(innerDUS.getUpdate().getType());
  auto innerOperandTy = cast<RankedTensorType>(innerDUS.getOperand().getType());
  auto outerUpdateShape = llvm::to_vector(outerUpdateTy.getShape());
  auto innerUpdateShape = llvm::to_vector(innerUpdateTy.getShape());
  auto innerOperandShape = llvm::to_vector(innerOperandTy.getShape());

  llvm::MapVector<Value, WhileLoopInfo::AffineIndexInfo> affineIndexInfoMap;
  if (loopInfo.has_value()) {
    affineIndexInfoMap = loopInfo.value().getAffineIndexInfo();
  }

  // Track start indices and whether we could determine them
  SmallVector<int64_t> outerStarts(outerStartIndices.size());
  SmallVector<int64_t> innerStarts(outerStartIndices.size());
  // For each dimension, track if:
  // - we know the exact indices (hasKnownIndices[i] = true)
  // - the inner DUS does a full update along this dimension
  //   (isFullInnerUpdate[i] = true)
  // - the outer DUS does a full update along this dimension
  //   (isFullOuterUpdate[i] = true)
  SmallVector<bool> hasKnownIndices(outerStartIndices.size(), false);
  SmallVector<bool> isFullInnerUpdate(outerStartIndices.size(), false);
  SmallVector<bool> isFullOuterUpdate(outerStartIndices.size(), false);

  for (size_t i = 0; i < outerStartIndices.size(); ++i) {
    auto outerStart = outerStartIndices[i];
    auto innerStart = innerStartIndices[i];

    // Check if inner DUS does a full update along this dimension
    // (inner start = 0 and inner update size = operand size)
    APInt innerStartAP;
    bool innerIsConstant =
        matchPattern(innerStart, m_ConstantInt(&innerStartAP));
    bool innerStartIsZero = innerIsConstant && innerStartAP.isZero();

    if (innerStartIsZero && innerUpdateShape[i] == innerOperandShape[i]) {
      isFullInnerUpdate[i] = true;
      innerStarts[i] = 0;
    }

    // Try to get exact indices
    APInt outerStartAP;
    bool outerIsConstant =
        matchPattern(outerStart, m_ConstantInt(&outerStartAP));
    bool outerStartIsZero = outerIsConstant && outerStartAP.isZero();

    if (outerStartIsZero && outerUpdateShape[i] == innerOperandShape[i]) {
      isFullOuterUpdate[i] = true;
      outerStarts[i] = 0;
    }

    if (outerIsConstant && innerIsConstant) {
      outerStarts[i] = outerStartAP.getSExtValue();
      innerStarts[i] = innerStartAP.getSExtValue();
      hasKnownIndices[i] = true;
    } else if (loopInfo.has_value() &&
               affineIndexInfoMap.contains(outerStart) &&
               affineIndexInfoMap.contains(innerStart)) {
      auto outerIndexInfo = affineIndexInfoMap[outerStart];
      auto innerIndexInfo = affineIndexInfoMap[innerStart];
      if (outerIndexInfo.scale.isOne() && innerIndexInfo.scale.isOne()) {
        outerStarts[i] = outerIndexInfo.offset.getSExtValue();
        innerStarts[i] = innerIndexInfo.offset.getSExtValue();
        hasKnownIndices[i] = true;
      }
    }
  }

  // Check for complete overlap: outer DUS completely covers inner DUS
  // i.e., outer starts at or before inner and outer ends at or after inner
  // This can be determined if:
  // - We know exact indices and outer covers inner, OR
  // - Outer update is a full slice update, OR
  // - The start indices are the same Value and outer update >= inner update
  bool outerCompletelyCoversInner = true;
  for (size_t i = 0; i < outerStartIndices.size(); ++i) {
    auto outerStart = outerStartIndices[i];
    auto innerStart = innerStartIndices[i];

    if (outerStart == innerStart) {
      // Same start index - check if outer update size >= inner update size
      if (outerUpdateShape[i] < innerUpdateShape[i]) {
        outerCompletelyCoversInner = false;
        break;
      }
    } else if (isFullOuterUpdate[i]) {
      // inner update indices don't matter in this case.
      continue;
    } else if (hasKnownIndices[i]) {
      // Different start indices but we know their values
      int64_t outerEnd = outerStarts[i] + outerUpdateShape[i];
      int64_t innerEnd = innerStarts[i] + innerUpdateShape[i];

      if (outerStarts[i] > innerStarts[i] || outerEnd < innerEnd) {
        outerCompletelyCoversInner = false;
        break;
      }
    } else {
      // Can't determine relationship
      outerCompletelyCoversInner = false;
      break;
    }
  }

  if (outerCompletelyCoversInner) {
    // Case 1: Outer DUS completely overlaps inner DUS - skip the inner DUS
    rewriter.modifyOpInPlace(outerDUS, [&]() {
      outerDUS.getOperandMutable().set(innerDUS.getOperand());
    });
    return success();
  }

  // Check for partial overlap: outer update is entirely within inner update
  // This can be determined if:
  // - We know the exact indices and outer is within inner, OR
  // - The inner DUS does a full update along that dimension (then any outer is
  //   within)
  bool outerWithinInner = true;
  for (size_t i = 0; i < outerStartIndices.size(); ++i) {
    if (isFullInnerUpdate[i]) {
      // Inner does full update, so outer is guaranteed to be within
      continue;
    }

    if (!hasKnownIndices[i]) {
      outerWithinInner = false;
      break;
    }

    int64_t outerEnd = outerStarts[i] + outerUpdateShape[i];
    int64_t innerEnd = innerStarts[i] + innerUpdateShape[i];

    if (outerStarts[i] < innerStarts[i] || outerEnd > innerEnd) {
      outerWithinInner = false;
      break;
    }
  }

  if (outerWithinInner) {
    // Case 2: Outer update is entirely within inner update
    // Create a DUS of outer update into inner update, then DUS into inner's
    // operand

    auto itype = innerStartIndices[0].getType();

    // Calculate the offset of outer update within inner update
    SmallVector<Value> offsetIndices;
    for (size_t i = 0; i < outerStartIndices.size(); ++i) {
      if (isFullInnerUpdate[i] && !hasKnownIndices[i]) {
        // Inner starts at 0, so offset is just the outer start index
        offsetIndices.push_back(outerStartIndices[i]);
      } else {
        int64_t offset = outerStarts[i] - innerStarts[i];
        auto offsetConst = stablehlo::ConstantOp::create(
            rewriter, outerDUS.getLoc(), itype,
            cast<ElementsAttr>(makeAttr(itype, offset)));
        offsetIndices.push_back(offsetConst);
      }
    }

    // Create a DUS of outer update into inner update
    auto combinedUpdate = stablehlo::DynamicUpdateSliceOp::create(
        rewriter, innerDUS.getLoc(), innerDUS.getUpdate(), outerDUS.getUpdate(),
        offsetIndices);

    // Replace outer DUS with a DUS using the combined update
    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
        outerDUS, innerDUS.getOperand(), combinedUpdate, innerStartIndices);
    return success();
  }

  return failure();
}

struct DUSDUS final
    : CheckedOpRewritePattern<stablehlo::DynamicUpdateSliceOp, DUSDUS> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DynamicUpdateSliceOp dus,
                                    PatternRewriter &rewriter) const {
    auto dusInner =
        dus.getOperand().getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
    if (!dusInner || !llvm::hasSingleElement(dusInner->getUsers())) {
      return failure();
    }

    return DUSDUSSimplifyHelper(dus, dusInner, rewriter, std::nullopt);
  }
};

// Pattern to simplify DUS(DUS) patterns inside a while body
// Both DUS operations must have the WhileOp as their owner (be in the while
// body)
struct WhileDUSDUSSimplify final
    : public CheckedOpRewritePattern<stablehlo::WhileOp, WhileDUSDUSSimplify> {
  using CheckedOpRewritePattern<stablehlo::WhileOp,
                                WhileDUSDUSSimplify>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::WhileOp op,
                                    PatternRewriter &rewriter) const {
    WhileLoopInfo info(op);
    if (info.computeInfo().failed() || !info.isConstantStep() ||
        info.getConstantStep() == 0) {
      return failure();
    }
    info.propagateAffineIndexInfo();
    info.propagateBounds();

    bool anyRewritten = false;

    op.getBody().walk([&](stablehlo::DynamicUpdateSliceOp outerDUS) {
      auto innerDUS = outerDUS.getOperand()
                          .getDefiningOp<stablehlo::DynamicUpdateSliceOp>();
      if (!innerDUS) {
        return WalkResult::advance();
      }

      // Check that both DUS are in the while body
      if (innerDUS->getParentOp() != op) {
        return WalkResult::advance();
      }

      rewriter.setInsertionPoint(outerDUS);
      if (DUSDUSSimplifyHelper(outerDUS, innerDUS, rewriter, info)
              .succeeded()) {
        anyRewritten = true;
      }

      return WalkResult::advance();
    });

    return anyRewritten ? success() : failure();
  }
};

struct ReshapeSliceReshape final
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                     ReshapeSliceReshape> {
  using CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                ReshapeSliceReshape>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp bottomReshapeOp,
                                    PatternRewriter &rewriter) const {
    auto sliceOp =
        bottomReshapeOp.getOperand().getDefiningOp<stablehlo::SliceOp>();
    if (!sliceOp || !llvm::hasSingleElement(sliceOp->getUsers())) {
      return failure();
    }

    auto topReshapeOp =
        sliceOp.getOperand().getDefiningOp<stablehlo::ReshapeOp>();
    if (!topReshapeOp) {
      return failure();
    }

    // compute which dims were inserted by the reshape op
    auto topReshapeInTy =
        cast<RankedTensorType>(topReshapeOp.getOperand().getType());
    auto topReshapeOutTy = cast<RankedTensorType>(topReshapeOp.getType());
    auto topInsertionDims =
        findReshapeInsertionDims(topReshapeInTy, topReshapeOutTy);
    if (topInsertionDims.empty()) {
      return failure();
    }

    auto bottomReshapeInTy =
        cast<RankedTensorType>(bottomReshapeOp.getOperand().getType());
    auto bottomReshapeOutTy = cast<RankedTensorType>(bottomReshapeOp.getType());
    auto bottomDeletionDims =
        findReshapeInsertionDims(bottomReshapeOutTy, bottomReshapeInTy);
    if (bottomDeletionDims.empty()) {
      return failure();
    }

    llvm::SetVector<int64_t> intersection;
    for (auto i : topInsertionDims) {
      if (llvm::is_contained(bottomDeletionDims, i)) {
        intersection.insert(i);
      }
    }

    if (intersection.empty()) {
      return failure();
    }

    SmallVector<int64_t> topReshapeShape, bottomReshapeShape, sliceStarts,
        sliceLimits, sliceStrides;
    for (size_t i = 0; i < topReshapeOutTy.getRank(); i++) {
      if (llvm::is_contained(intersection, i)) {
        continue;
      }
      topReshapeShape.push_back(topReshapeOutTy.getDimSize(i));
      sliceStarts.push_back(sliceOp.getStartIndices()[i]);
      sliceLimits.push_back(sliceOp.getLimitIndices()[i]);
      sliceStrides.push_back(sliceOp.getStrides()[i]);
    }

    rewriter.setInsertionPointAfter(topReshapeOp);
    auto newTopReshape =
        stablehlo::ReshapeOpCreate(rewriter, bottomReshapeOp.getLoc(),
                                   topReshapeOp.getOperand(), topReshapeShape);

    // Create the new slice op on the new top reshape
    rewriter.setInsertionPointAfter(sliceOp);
    auto newSlice =
        stablehlo::SliceOp::create(rewriter, sliceOp.getLoc(), newTopReshape,
                                   rewriter.getDenseI64ArrayAttr(sliceStarts),
                                   rewriter.getDenseI64ArrayAttr(sliceLimits),
                                   rewriter.getDenseI64ArrayAttr(sliceStrides));

    // Create the new bottom reshape and replace users
    rewriter.setInsertionPointAfter(bottomReshapeOp);
    auto newBottomReshape =
        stablehlo::ReshapeOpCreate(rewriter, bottomReshapeOp.getLoc(), newSlice,
                                   bottomReshapeOutTy.getShape());

    rewriter.replaceAllUsesWith(bottomReshapeOp.getResult(), newBottomReshape);
    return success();
  }
};

// If a reduce is followed by a reshape that deletes certain dimensions,
// we can absorb those deleted dimensions into the reduce op
struct ReduceDeleteDims final
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp, ReduceDeleteDims> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp reshapeOp,
                                    PatternRewriter &rewriter) const {
    auto reduceOp = reshapeOp.getOperand().getDefiningOp<stablehlo::ReduceOp>();
    if (!reduceOp)
      return failure();

    if (reduceOp.getInputs().size() != 1 ||
        reduceOp.getInitValues().size() != 1)
      return failure();

    if (!llvm::hasSingleElement(reduceOp->getUsers()))
      return failure();

    auto reduceOutTy = cast<RankedTensorType>(reduceOp.getResult(0).getType());
    auto reshapeOutTy = cast<RankedTensorType>(reshapeOp.getType());

    // Find which dimensions are being deleted by the reshape
    // findReshapeInsertionDims(output, input) returns dims in input not in
    // output
    auto deletedDims = findReshapeInsertionDims(reshapeOutTy, reduceOutTy);
    if (deletedDims.empty())
      return failure();

    // Map reduce output dims to input dims
    // The reduce output has dimensions: all input dims except the reduced ones
    auto reduceInputTy =
        cast<RankedTensorType>(reduceOp.getInputs()[0].getType());
    auto reduceDims = reduceOp.getDimensions();

    // Build mapping from reduce output dim -> reduce input dim
    SmallVector<int64_t> reduceOutToInMap;
    for (int64_t inIdx = 0; inIdx < reduceInputTy.getRank(); ++inIdx) {
      if (llvm::is_contained(reduceDims, inIdx))
        continue;
      reduceOutToInMap.push_back(inIdx);
    }

    // Map the deleted dimensions from reduce output back to reduce input
    SmallVector<int64_t> newReduceDims(reduceDims.begin(), reduceDims.end());
    for (auto dim : deletedDims) {
      if (dim >= reduceOutToInMap.size())
        return failure();
      newReduceDims.push_back(reduceOutToInMap[dim]);
    }

    // Sort the dimensions since reduce expects them in order
    llvm::sort(newReduceDims);

    // Create the new reduce op with the additional dimensions
    auto newReduceOp = stablehlo::ReduceOp::create(
        rewriter, reduceOp.getLoc(), TypeRange(reshapeOp.getType()),
        ValueRange(reduceOp.getInputs()), ValueRange(reduceOp.getInitValues()),
        newReduceDims);
    newReduceOp.getRegion().takeBody(reduceOp.getRegion());

    rewriter.replaceOp(reshapeOp, newReduceOp.getResults());
    return success();
  };
};

// If reshape deletes certain dimensions followed by a reduce,
// we can absorb those into the dimensions of the reduce op
struct DeleteDimsReduce final
    : public CheckedOpRewritePattern<stablehlo::ReduceOp, DeleteDimsReduce> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
                                    PatternRewriter &rewriter) const {
    if (op.getInputs().size() != 1 || op.getInitValues().size() != 1)
      return failure();

    auto reshapeOp = op.getInputs()[0].getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshapeOp)
      return failure();

    if (!llvm::hasSingleElement(reshapeOp->getUsers()))
      return failure();

    auto reshapeInTy = cast<RankedTensorType>(reshapeOp.getOperand().getType());
    auto reshapeOutTy = cast<RankedTensorType>(reshapeOp.getType());

    // Find which dimensions were deleted by the reshape (size 1 dims)
    // When we call findReshapeInsertionDims(output, input), we get the
    // indices in the input that were "inserted" to create the output,
    // which corresponds to dims that were deleted when going from input to
    // output
    auto deletedDims = findReshapeInsertionDims(reshapeOutTy, reshapeInTy);
    if (deletedDims.empty())
      return failure();

    // Build a mapping from reshapeOut dimensions to reshapeIn dimensions
    // For each dimension in reshapeOut, find the corresponding dimension
    // in reshapeIn
    SmallVector<int64_t> outToInMap;
    size_t inIdx = 0;
    for (int64_t outIdx = 0; outIdx < reshapeOutTy.getRank(); ++outIdx) {
      // Skip any deleted dimensions
      while (inIdx < reshapeInTy.getRank() &&
             llvm::is_contained(deletedDims, inIdx)) {
        ++inIdx;
      }
      if (inIdx >= reshapeInTy.getRank())
        return failure();
      outToInMap.push_back(inIdx);
      ++inIdx;
    }

    // Map the reduce dimensions from the reshaped tensor back to the
    // original tensor
    SmallVector<int64_t> newReduceDims;
    for (auto reduceDim : op.getDimensions()) {
      if (reduceDim >= outToInMap.size())
        return failure();
      newReduceDims.push_back(outToInMap[reduceDim]);
    }

    // Add all the deleted dimensions to the reduce dimensions (they're
    // size 1, so reducing them is a no-op mathematically)
    for (auto dim : deletedDims) {
      newReduceDims.push_back(dim);
    }

    // Sort the dimensions since reduce expects them in order
    llvm::sort(newReduceDims);

    // Create the new reduce op on the original input
    auto newReduceOp = stablehlo::ReduceOp::create(
        rewriter, op.getLoc(), TypeRange(op.getResultTypes()),
        ValueRange(reshapeOp.getOperand()), ValueRange(op.getInitValues()),
        newReduceDims);
    newReduceOp.getRegion().takeBody(op.getRegion());

    rewriter.replaceOp(op, newReduceOp.getResults());
    return success();
  }
};

struct DotGeneralInsertDimContractionSimplification final
    : public CheckedOpRewritePattern<
          stablehlo::DotGeneralOp,
          DotGeneralInsertDimContractionSimplification> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
                                    PatternRewriter &rewriter) const {
    auto dotDims = op.getDotDimensionNumbers();
    auto lhs = op.getLhs();
    auto lhsType = cast<RankedTensorType>(lhs.getType());
    auto rhs = op.getRhs();
    auto rhsType = cast<RankedTensorType>(rhs.getType());

    // Find singleton contracting dimensions (size 1 on both sides)
    SmallVector<int64_t> newLhsContractingDims, newRhsContractingDims;
    SmallVector<int64_t> lhsSingletonContractDims, rhsSingletonContractDims;

    for (auto [lhsDim, rhsDim] :
         llvm::zip(dotDims.getLhsContractingDimensions(),
                   dotDims.getRhsContractingDimensions())) {
      if (lhsType.getDimSize(lhsDim) == 1 && rhsType.getDimSize(rhsDim) == 1) {
        lhsSingletonContractDims.push_back(lhsDim);
        rhsSingletonContractDims.push_back(rhsDim);
        continue;
      }
      newLhsContractingDims.push_back(lhsDim);
      newRhsContractingDims.push_back(rhsDim);
    }

    if (lhsSingletonContractDims.empty()) {
      return failure();
    }

    // Check that both operands can have their singleton dims removed
    if (!canRemoveSingletonDim(lhs, lhsSingletonContractDims) ||
        !canRemoveSingletonDim(rhs, rhsSingletonContractDims)) {
      return failure();
    }

    // Build new LHS shape (remove singleton contracting dims)
    SmallVector<int64_t> newLhsShape;
    DenseMap<int64_t, int64_t> lhsOldToNew;
    for (int64_t i = 0; i < lhsType.getRank(); ++i) {
      if (llvm::is_contained(lhsSingletonContractDims, i))
        continue;
      lhsOldToNew[i] = newLhsShape.size();
      newLhsShape.push_back(lhsType.getDimSize(i));
    }

    // Build new RHS shape (remove singleton contracting dims)
    SmallVector<int64_t> newRhsShape;
    DenseMap<int64_t, int64_t> rhsOldToNew;
    for (int64_t i = 0; i < rhsType.getRank(); ++i) {
      if (llvm::is_contained(rhsSingletonContractDims, i))
        continue;
      rhsOldToNew[i] = newRhsShape.size();
      newRhsShape.push_back(rhsType.getDimSize(i));
    }

    // Map old batch/contracting dims to new positions
    SmallVector<int64_t> mappedLhsBatchDims, mappedRhsBatchDims;
    for (auto dim : dotDims.getLhsBatchingDimensions()) {
      if (lhsOldToNew.count(dim))
        mappedLhsBatchDims.push_back(lhsOldToNew[dim]);
    }
    for (auto dim : dotDims.getRhsBatchingDimensions()) {
      if (rhsOldToNew.count(dim))
        mappedRhsBatchDims.push_back(rhsOldToNew[dim]);
    }

    SmallVector<int64_t> mappedLhsContractDims, mappedRhsContractDims;
    for (auto dim : newLhsContractingDims) {
      if (lhsOldToNew.count(dim))
        mappedLhsContractDims.push_back(lhsOldToNew[dim]);
    }
    for (auto dim : newRhsContractingDims) {
      if (rhsOldToNew.count(dim))
        mappedRhsContractDims.push_back(rhsOldToNew[dim]);
    }

    // Create reshaped operands. This will be cleaned up later
    Value newLhs =
        stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), lhs, newLhsShape);
    Value newRhs =
        stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), rhs, newRhsShape);

    // Create new dot dimension numbers
    auto newDotDims = stablehlo::DotDimensionNumbersAttr::get(
        op.getContext(), mappedLhsBatchDims, mappedRhsBatchDims,
        mappedLhsContractDims, mappedRhsContractDims);

    // Create the new dot_general
    auto newDotOp = stablehlo::DotGeneralOp::create(
        rewriter, op.getLoc(), op.getType(), newLhs, newRhs, newDotDims,
        op.getPrecisionConfigAttr(), op.getAlgorithmAttr());

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

private:
  // Check if an operand can have a singleton dimension removed at the given
  // position. This is safe if the operand is:
  // 1. A splatted constant (all elements same value)
  // 2. A broadcast_in_dim that creates the singleton dimension (dim not in
  //    broadcast_dimensions)
  // 3. A reshape operation (can be adjusted)
  bool canRemoveSingletonDim(Value operand,
                             ArrayRef<int64_t> singletonDims) const {
    auto defOp = operand.getDefiningOp();
    if (!defOp) {
      return false;
    }

    // Case 1: Splatted constant
    SplatElementsAttr splat;
    if (matchPattern(defOp, m_Constant(&splat))) {
      return true;
    }

    // Case 2: broadcast_in_dim that creates the singleton dimension
    if (auto bcastOp = dyn_cast<stablehlo::BroadcastInDimOp>(defOp)) {
      auto broadcastDims = bcastOp.getBroadcastDimensions();
      // If dimToRemove is not in broadcastDims, it was created by the broadcast
      return llvm::all_of(singletonDims, [&](auto dim) {
        return !llvm::is_contained(broadcastDims, dim);
      });
    }

    // Case 3: Reshape operation
    if (dyn_cast<stablehlo::ReshapeOp>(defOp)) {
      return true;
    }

    return false;
  }
};

// This is different from ReshapeBroadcast. This tries to push the reshape up
// when deleting dims that were present in the original broadcast input.
// reshape(broadcast(x)) -> broadcast(reshape(x))
struct DeleteDimsBroadcast final
    : public CheckedOpRewritePattern<stablehlo::ReshapeOp,
                                     DeleteDimsBroadcast> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReshapeOp op,
                                    PatternRewriter &rewriter) const {
    auto operand = op.getOperand();
    auto bcastOp = operand.getDefiningOp<stablehlo::BroadcastInDimOp>();
    if (!bcastOp || !llvm::hasSingleElement(bcastOp->getUsers())) {
      return failure();
    }

    auto bcastInputTy = cast<RankedTensorType>(bcastOp.getOperand().getType());
    auto bcastOutputTy = cast<RankedTensorType>(bcastOp.getType());
    auto reshapeOutputTy = cast<RankedTensorType>(op.getType());

    // Find which dimensions are being deleted by the reshape
    auto deleteDims = findReshapeInsertionDims(reshapeOutputTy, bcastOutputTy);
    if (deleteDims.empty()) {
      return failure();
    }

    auto bcastDims = bcastOp.getBroadcastDimensions();

    // Check if all deleted dims are in the broadcast_dimensions
    // (i.e., they came from the original input)
    llvm::SetVector<int64_t> inputDimsToDelete;
    for (auto dim : deleteDims) {
      auto it = llvm::find(bcastDims, dim);
      if (it == bcastDims.end()) {
        // This dimension was introduced by the broadcast, not from input
        return failure();
      }
      // Map back to the input dimension
      inputDimsToDelete.insert(std::distance(bcastDims.begin(), it));
    }

    // Create new input shape by removing the deleted dimensions
    SmallVector<int64_t> newInputShape;
    for (int64_t i = 0; i < bcastInputTy.getRank(); ++i) {
      if (!llvm::is_contained(inputDimsToDelete, i)) {
        newInputShape.push_back(bcastInputTy.getDimSize(i));
      }
    }

    // Create the reshape on the input
    auto reshapeInput = stablehlo::ReshapeOpCreate(
        rewriter, op.getLoc(), bcastOp.getOperand(), newInputShape);

    // Compute new broadcast dimensions by:
    // 1. Removing deleted dims from bcastDims
    // 2. Adjusting remaining dims for the deleted dims in the output
    SmallVector<int64_t> newBroadcastDims;
    for (int64_t i = 0; i < bcastDims.size(); ++i) {
      if (llvm::is_contained(inputDimsToDelete, i)) {
        continue; // Skip deleted input dims
      }
      int64_t bcastDim = bcastDims[i];
      // Adjust for deleted output dims
      int64_t offset = 0;
      for (int64_t deletedDim : deleteDims) {
        if (deletedDim < bcastDim) {
          offset++;
        }
      }
      newBroadcastDims.push_back(bcastDim - offset);
    }

    // Create a new broadcast with the reshaped input
    auto newBcast = stablehlo::BroadcastInDimOp::create(
        rewriter, op.getLoc(), reshapeOutputTy, reshapeInput, newBroadcastDims);
    rewriter.replaceOp(op, newBcast.getResult());
    return success();
  }
};

struct FuseReshapeCollapseOrExpandDimsIntoReduce final
    : CheckedOpRewritePattern<stablehlo::ReduceOp,
                              FuseReshapeCollapseOrExpandDimsIntoReduce> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp reduceOp,
                                    PatternRewriter &rewriter) {
    if (reduceOp.getInputs().size() != 1) {
      return failure();
    }

    auto reshapeOp =
        reduceOp.getInputs()[0].getDefiningOp<stablehlo::ReshapeOp>();
    if (!reshapeOp || !llvm::hasSingleElement(reshapeOp->getUsers())) {
      return failure();
    }

    auto reshapeInTy = cast<RankedTensorType>(reshapeOp.getOperand().getType());
    auto reshapeOutTy = cast<RankedTensorType>(reshapeOp.getType());

    DenseMap<int64_t, SmallVector<int64_t, 2>> mapping;
    bool collapsing = false, valid = false;
    SmallVector<int64_t> newReduceDims;

    if (enzyme::getCollapsingMapping(reshapeInTy.getShape(),
                                     reshapeOutTy.getShape(), mapping)) {
      // collapsing dimensions
      collapsing = true;
      int64_t originalReduceInDims = reshapeOutTy.getRank();
      for (auto dim : reduceOp.getDimensions()) {
        auto mappedDims = mapping[dim];
        originalReduceInDims += mappedDims.size() - 1;
        newReduceDims.append(mappedDims.begin(), mappedDims.end());
      }
      valid = originalReduceInDims == reshapeInTy.getRank();
    }

    mapping.clear();
    if (!collapsing &&
        enzyme::getCollapsingMapping(reshapeOutTy.getShape(),
                                     reshapeInTy.getShape(), mapping)) {
      // expanding dimensions
      DenseMap<int64_t, int64_t> invMap;
      DenseMap<int64_t, SmallVector<int64_t, 2>> expandStartToFull;
      for (auto [key, val] : mapping) {
        invMap[val[0]] = key;
        expandStartToFull[val[0]] = val;
      }

      int64_t originalReduceInDims = reshapeOutTy.getRank();
      for (int64_t redDimIdx = 0; redDimIdx < reduceOp.getDimensions().size();
           redDimIdx++) {
        auto dim = reduceOp.getDimensions()[redDimIdx];
        if (!expandStartToFull.contains(dim)) {
          return failure();
        }
        auto expandFull = expandStartToFull[dim];
        if (expandFull.size() == 1) {
          newReduceDims.push_back(invMap[dim]);
        } else {
          auto NExpand = expandFull.size();
          originalReduceInDims -= NExpand - 1;
          newReduceDims.push_back(invMap[dim]);
          if (redDimIdx + NExpand > reduceOp.getDimensions().size()) {
            return failure();
          }
          for (int64_t j = 1; j < NExpand; j++) {
            if (reduceOp.getDimensions()[redDimIdx + j] != expandFull[j]) {
              return failure();
            }
          }
          redDimIdx += NExpand - 1;
        }
      }
      valid = originalReduceInDims == reshapeInTy.getRank();
    }

    if (valid) {
      llvm::sort(newReduceDims);
      auto newReduceOp = stablehlo::ReduceOp::create(
          rewriter, reduceOp.getLoc(), reduceOp.getResultTypes(),
          reshapeOp.getOperand(), reduceOp.getInitValues(), newReduceDims);
      newReduceOp.getRegion().takeBody(reduceOp.getRegion());
      rewriter.replaceOp(reduceOp, newReduceOp.getResults());
      return success();
    }

    return failure();
  }
};

struct GatherOfScatterSimplify final
    : CheckedOpRewritePattern<stablehlo::GatherOp, GatherOfScatterSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::GatherOp gatherOp,
                                    PatternRewriter &rewriter) {
    auto input = gatherOp.getOperand();
    auto scatterOp = input.getDefiningOp<stablehlo::ScatterOp>();

    if (!scatterOp ||
        scatterOp.getScatterIndices() != gatherOp.getStartIndices() ||
        computeGatherSliceSizes(scatterOp) != gatherOp.getSliceSizes() ||
        getGatherDims(scatterOp->getContext(),
                      scatterOp.getScatterDimensionNumbersAttr()) !=
            gatherOp.getDimensionNumbersAttr()) {
      return failure();
    }

    auto opResult = cast<OpResult>(input);
    auto opNum = opResult.getResultNumber();

    SplatElementsAttr constSetIndexValue;
    if (!detectConstantSetindexScatterOp(
             scatterOp, true, [](auto input) { return true; },
             constSetIndexValue)
             .ok()) {
      return failure();
    }

    if (constSetIndexValue) {
      auto constResult = stablehlo::ConstantOp::create(
          rewriter, gatherOp.getLoc(),
          constSetIndexValue.resizeSplat(cast<ShapedType>(gatherOp.getType())));
      rewriter.replaceOp(gatherOp, constResult);
      return success();
    }

    if (!scatterOp.getUniqueIndices()) {
      return failure();
    }

    auto newResult = scatterOp.getUpdates()[opNum];
    rewriter.replaceOp(gatherOp, newResult);
    return success();
  }
};

// jax doesn't control size of constants, we try to recover patterns that can
// produce these constants. We do the following:
//    1. iota detection
struct RecognizeFromConstant final
    : CheckedOpRewritePattern<stablehlo::ConstantOp, RecognizeFromConstant> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  size_t min_fold_size;
  RecognizeFromConstant(size_t min_fold_size, MLIRContext *context,
                        PatternBenefit benefit = 1,
                        ArrayRef<StringRef> generatedNames = {})
      : CheckedOpRewritePattern(context, benefit, generatedNames),
        min_fold_size(min_fold_size) {}

  LogicalResult matchAndRewriteImpl(stablehlo::ConstantOp op,
                                    PatternRewriter &rewriter) {
    DenseElementsAttr val = cast<DenseElementsAttr>(op.getValue());
    if (val.isSplat() || op.getType().getNumElements() < min_fold_size) {
      return failure();
    }

    auto maybeIotaDetect = detectIotaLikeTensor(val);
    if (maybeIotaDetect.has_value()) {
      auto iotaDetect = maybeIotaDetect.value();

      // Helper to create a splat constant from a TypedAttr
      auto createSplatConstant = [&](TypedAttr attr) -> Value {
        auto splatAttr = SplatElementsAttr::get(iotaDetect.tensorType, attr);
        return stablehlo::ConstantOp::create(rewriter, op.getLoc(), splatAttr);
      };

      Value iotaOp =
          stablehlo::IotaOp::create(rewriter, op.getLoc(),
                                    iotaDetect.tensorType, iotaDetect.dimension)
              .getResult();
      if (!isOneAttr(iotaDetect.scale)) {
        iotaOp = stablehlo::MulOpCreate(rewriter, op.getLoc(), iotaOp,
                                        createSplatConstant(iotaDetect.scale));
      }
      if (!isZeroAttr(iotaDetect.start)) {
        iotaOp = stablehlo::AddOpCreate(rewriter, op.getLoc(), iotaOp,
                                        createSplatConstant(iotaDetect.start));
      }

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

    // Check if the constant can be represented as a pad of a smaller constant
    auto maybePaddedDetect = enzyme::detectPaddedTensor(val);
    if (maybePaddedDetect.has_value()) {
      auto paddedDetect = maybePaddedDetect.value();

      // Create the inner constant
      auto innerConstant = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), paddedDetect.innerTensorAttr);

      // Create the padding value as a scalar constant
      auto paddingValueType =
          RankedTensorType::get({}, paddedDetect.resultType.getElementType());
      auto paddingValueAttr =
          DenseElementsAttr::get(paddingValueType, paddedDetect.paddingValue);
      auto paddingValue = stablehlo::ConstantOp::create(rewriter, op.getLoc(),
                                                        paddingValueAttr);

      // Interior padding is always zeros
      SmallVector<int64_t> interiorPadding(paddedDetect.lowPadding.size(), 0);

      // Create the pad operation
      auto padOp = stablehlo::PadOp::create(
          rewriter, op.getLoc(), paddedDetect.resultType, innerConstant,
          paddingValue, rewriter.getDenseI64ArrayAttr(paddedDetect.lowPadding),
          rewriter.getDenseI64ArrayAttr(paddedDetect.highPadding),
          rewriter.getDenseI64ArrayAttr(interiorPadding));

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

    return failure();
  }
};

struct ReduceWindowWrapSimplify final
    : CheckedOpRewritePattern<stablehlo::ReduceWindowOp,
                              ReduceWindowWrapSimplify> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ReduceWindowOp op,
                                    PatternRewriter &rewriter) {
    if (op.getInputs().size() != 1) {
      return failure();
    }

    auto input = op.getInputs()[0];
    auto wrapOp = input.getDefiningOp<enzymexla::WrapOp>();
    if (!wrapOp || !llvm::hasSingleElement(wrapOp->getUsers())) {
      return failure();
    }

    auto wrapInput = wrapOp.getOperand();
    auto wrapInputTy = cast<RankedTensorType>(wrapInput.getType());

    auto inputTy = cast<RankedTensorType>(input.getType());
    auto rank = inputTy.getRank();

    auto windowDims = op.getWindowDimensions();
    auto windowDilations = op.getWindowDilations();
    auto windowStrides = op.getWindowStrides();
    auto padding = op.getPadding();
    auto baseDilations = op.getBaseDilations();

    auto wrapDim = wrapOp.getDimension();

    if (padding && !llvm::all_of(padding.value(),
                                 [](auto pad_val) { return pad_val == 0; })) {
      return failure();
    }

    int64_t wrapSize;
    bool origSideLhs;
    if (wrapOp.getLhs() == 0) {
      wrapSize = wrapOp.getRhs();
      origSideLhs = false;
    } else if (wrapOp.getRhs() == 0) {
      wrapSize = wrapOp.getLhs();
      origSideLhs = true;
    } else {
      return rewriter.notifyMatchFailure(wrapOp, "needs rhs/lhs to be 0");
    }

    auto dimSize = wrapInputTy.getDimSize(wrapDim);

    if (wrapSize <= dimSize / 2 || wrapSize > dimSize) {
      return rewriter.notifyMatchFailure(wrapOp,
                                         "no savings by switching side");
    }

    // being overly conservative for now, we can expand as needed
    // All values must be 1s except:
    //   1. windowDims along wrapDim == 2
    //   2. windowDilations along wrapDim == wrapSize
    for (size_t i = 0; i < rank; i++) {
      if (i == wrapDim) {
        if (windowDims[i] != 2 ||
            (windowStrides && windowStrides.value()[i] != 1)) {
          return failure();
        }
        if (!windowDilations || windowDilations.value()[i] != wrapSize) {
          return failure();
        }
        continue;
      }

      if (windowDims[i] != 1 ||
          (windowDilations && windowDilations.value()[i] != 1) ||
          (windowStrides && windowStrides.value()[i] != 1) ||
          (baseDilations && baseDilations.value()[i] != 1)) {
        return failure();
      }
    }

    auto commonReduceWindowOp = CheckCommonReduceWindowOp(op);

    // We can optimize only for commutative ops
    if (!commonReduceWindowOp.isCommutativeOp()) {
      return failure();
    }

    if (wrapSize == dimSize) {
      // this is effectively op(a, a)
      auto result = commonReduceWindowOp.createEquivalentOperation(
          rewriter, op.getLoc(), wrapInput, wrapInput);
      rewriter.replaceOp(op, result);
      return success();
    }

    int64_t newLhs = 0, newRhs = 0, newDilation = dimSize - wrapSize;
    if (origSideLhs) {
      newRhs = newDilation;
    } else {
      newLhs = newDilation;
    }

    auto newWrapOp = enzymexla::WrapOp::create(rewriter, op.getLoc(), wrapInput,
                                               newLhs, newRhs, wrapDim);

    SmallVector<int64_t> newWindowDilations(rank, 1);
    newWindowDilations[wrapDim] = newDilation;

    auto newReduceWindowOp = stablehlo::ReduceWindowOp::create(
        rewriter, op.getLoc(), ValueRange(newWrapOp), op.getInitValues(),
        op.getWindowDimensionsAttr(), op.getWindowStridesAttr(),
        op.getBaseDilationsAttr(),
        rewriter.getDenseI64ArrayAttr(newWindowDilations), op.getPaddingAttr());
    newReduceWindowOp.getRegion().takeBody(op.getRegion());
    rewriter.replaceOp(op, newReduceWindowOp);
    return success();
  }
};

// Pattern to recognize multiple SliceOps on the same input that differ only
// by offset along one dimension and combine them into a single MultiSliceOp
struct RecognizeMultiSlice
    : public CheckedOpRewritePattern<stablehlo::SliceOp, RecognizeMultiSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::SliceOp op,
                                    PatternRewriter &rewriter) const {
    Value input = op.getOperand();
    auto startIndices = SmallVector<int64_t>(op.getStartIndices());
    auto limitIndices = SmallVector<int64_t>(op.getLimitIndices());
    auto strides = SmallVector<int64_t>(op.getStrides());

    // Find all SliceOps on the same input
    SmallVector<stablehlo::SliceOp> slices;

    for (Operation *user : input.getUsers()) {
      if (auto slice = dyn_cast<stablehlo::SliceOp>(user)) {
        if (!slice.use_empty()) {
          slices.push_back(slice);
        }
      }
    }

    // Need at least 2 slices to consider combining
    if (slices.size() < 2)
      return failure();

    int64_t rank = startIndices.size();

    // Try each dimension as the potential varying dimension
    for (int64_t dim = 0; dim < rank; dim++) {
      // Collect slices that match this op except for offset along `dim`
      SmallVector<std::pair<stablehlo::SliceOp, int64_t>> matchingSlices;
      DenseSet<int64_t> offsetSet;

      for (auto slice : slices) {
        auto sliceStart = SmallVector<int64_t>(slice.getStartIndices());
        auto sliceLimit = SmallVector<int64_t>(slice.getLimitIndices());
        auto sliceStrides = SmallVector<int64_t>(slice.getStrides());

        // Check if strides match
        if (sliceStrides != strides)
          continue;

        // Check if all dimensions except `dim` match exactly
        bool matches = true;
        for (int64_t d = 0; d < rank; d++) {
          if (d == dim)
            continue;
          if (sliceStart[d] != startIndices[d] ||
              sliceLimit[d] != limitIndices[d]) {
            matches = false;
            break;
          }
        }
        if (!matches)
          continue;

        // Check if the slice shape along `dim` matches
        int64_t thisSliceSize = limitIndices[dim] - startIndices[dim];
        int64_t otherSliceSize = sliceLimit[dim] - sliceStart[dim];
        if (thisSliceSize != otherSliceSize)
          continue;

        // Calculate offset relative to this slice (this slice has offset 0)
        int64_t offset = sliceStart[dim] - startIndices[dim];
        matchingSlices.push_back({slice, offset});
        offsetSet.insert(offset);
      }

      // Need at least 2 matching slices along this dimension
      if (matchingSlices.size() < 2)
        continue;

      // Sort offsets to find contiguous groups
      SmallVector<int64_t> sortedOffsets(offsetSet.begin(), offsetSet.end());
      llvm::sort(sortedOffsets);

      // Find all contiguous groups (no gaps in offsets)
      SmallVector<std::pair<int64_t, int64_t>>
          contiguousGroups; // (start, end) inclusive
      int64_t groupStart = sortedOffsets[0];
      int64_t groupEnd = sortedOffsets[0];

      for (size_t i = 1; i < sortedOffsets.size(); i++) {
        if (sortedOffsets[i] == groupEnd + 1) {
          // Extend current group
          groupEnd = sortedOffsets[i];
        } else {
          // Save current group and start new one
          contiguousGroups.push_back({groupStart, groupEnd});
          groupStart = sortedOffsets[i];
          groupEnd = sortedOffsets[i];
        }
      }
      contiguousGroups.push_back({groupStart, groupEnd});

      // Collect all offsets from groups that contain identity (0) or neighbor
      // it
      SmallVector<int64_t> qualifyingOffsets;

      for (auto &[start, end] : contiguousGroups) {
        bool containsIdentity = (start <= 0 && end >= 0);
        bool neighborsIdentity = (end == -1 || start == 1);

        if (containsIdentity || neighborsIdentity) {
          for (int64_t o = start; o <= end; o++) {
            if (offsetSet.contains(o)) {
              qualifyingOffsets.push_back(o);
            }
          }
        }
      }

      // No qualifying groups found
      if (qualifyingOffsets.size() < 2)
        continue;

      // Determine the multiSlice range from qualifying offsets, extended to
      // include 0
      int64_t rangeStart =
          *std::min_element(qualifyingOffsets.begin(), qualifyingOffsets.end());
      int64_t rangeEnd =
          *std::max_element(qualifyingOffsets.begin(), qualifyingOffsets.end());

      // Extend to include identity (0) if not already included
      if (rangeStart > 0)
        rangeStart = 0;
      if (rangeEnd < 0)
        rangeEnd = 0;

      // Check if this op (offset 0) is part of the selected range
      if (0 < rangeStart || 0 > rangeEnd)
        continue;

      // Find the minimum offset among actual slices in the selected range
      // to use as the canonical trigger (prevents processing same group
      // multiple times)
      int64_t minOffsetInRange = INT64_MAX;
      int slicesToCombine = 0;
      for (auto &[slice, offset] : matchingSlices) {
        if (offset >= rangeStart && offset <= rangeEnd) {
          minOffsetInRange = std::min(minOffsetInRange, offset);
          slicesToCombine++;
        }
      }

      // Only proceed if this op has the minimum offset in the range
      if (0 != minOffsetInRange)
        return failure();

      // Need at least 2 slices in the selected range to combine
      if (slicesToCombine < 2)
        continue;

      // Check that all slices in the range have the same sharding
      std::optional<sdy::TensorShardingPerValueAttr> commonSharding;
      bool shardingMismatch = false;
      for (auto &[slice, offset] : matchingSlices) {
        if (offset >= rangeStart && offset <= rangeEnd) {
          auto shardPerValue = sdy::getShardingPerValue(slice);
          if (!commonSharding.has_value()) {
            commonSharding = shardPerValue;
          } else if (commonSharding.value() != shardPerValue) {
            shardingMismatch = true;
            break;
          }
        }
      }
      if (shardingMismatch)
        return failure();

      // Calculate the shift amount
      // leftShift covers negative offsets (slices shifted left/earlier)
      int32_t leftShift = rangeStart < 0 ? (int32_t)(-rangeStart) : 0;
      int32_t rightExtent = rangeEnd > 0 ? (int32_t)rangeEnd : 0;
      int32_t amount = leftShift + rightExtent;
      int32_t totalResults = amount + 1;

      // Adjust start and limit indices by shifting left along the dimension
      SmallVector<int64_t> adjustedStartIndices = startIndices;
      SmallVector<int64_t> adjustedLimitIndices = limitIndices;
      adjustedStartIndices[dim] -= leftShift;
      adjustedLimitIndices[dim] -= leftShift;

      // Create the MultiSliceOp
      auto resultType = op.getResult().getType();
      SmallVector<Type> resultTypes(totalResults, resultType);

      rewriter.setInsertionPointAfterValue(input);
      auto newOp = enzymexla::MultiSliceOp::create(
          rewriter, op.getLoc(), resultTypes, input, adjustedStartIndices,
          adjustedLimitIndices, strides, (int32_t)dim, amount);

      // Propagate sharding if present (all slices have the same sharding)
      if (commonSharding.has_value() && commonSharding.value()) {
        auto shardings = commonSharding.value().getShardings();
        if (!shardings.empty()) {
          sdy::TensorShardingAttr singleShard = shardings[0];
          SmallVector<sdy::TensorShardingAttr> newShardings(totalResults,
                                                            singleShard);
          sdy::setShardings(newOp, sdy::TensorShardingPerValueAttr::get(
                                       op.getContext(), newShardings));
        }
      }

      // Replace slices that fall within the selected range
      for (auto &[slice, offset] : matchingSlices) {
        if (offset >= rangeStart && offset <= rangeEnd) {
          // Result index for offset 'o' is: leftShift + o
          // (since we shifted start_indices left by leftShift)
          int32_t resultIdx = leftShift + (int32_t)offset;
          rewriter.replaceOp(slice, newOp.getResult(resultIdx));
        }
        // Slices outside the range are left unchanged
      }

      return success();
    }

    return failure();
  }
};

// Pattern to reduce MultiSliceOp when some results are unused
struct ReduceUnusedMultiSlice final
    : CheckedOpRewritePattern<enzymexla::MultiSliceOp, ReduceUnusedMultiSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::MultiSliceOp op,
                                    PatternRewriter &rewriter) const {
    int32_t amount = op.getAmount();
    int32_t totalResults = amount + 1;

    // Check which results are actually used
    SmallVector<bool> used(totalResults, false);
    int usedCount = 0;
    for (int i = 0; i < totalResults; i++) {
      if (!op.getResult(i).use_empty()) {
        used[i] = true;
        usedCount++;
      }
    }

    // If all results are used, nothing to optimize
    if (usedCount == totalResults)
      return failure();

    // If no results are used, this should be handled by dead code elimination
    if (usedCount == 0) {
      rewriter.eraseOp(op);
      return success();
    }

    // Find the range of used results
    int firstUsed = -1, lastUsed = -1;
    for (int i = 0; i < totalResults; i++) {
      if (used[i]) {
        if (firstUsed == -1)
          firstUsed = i;
        lastUsed = i;
      }
    }

    // Calculate new amount
    int newAmount = lastUsed - firstUsed;

    // If only one result is used, replace with a single SliceOp
    if (usedCount == 1) {
      auto startIndices = SmallVector<int64_t>(op.getStartIndices());
      auto limitIndices = SmallVector<int64_t>(op.getLimitIndices());
      auto strides = SmallVector<int64_t>(op.getStrides());
      int32_t dim = op.getDimension();

      // Adjust start and limit indices for the used result's offset
      if (dim >= 0 && dim < (int64_t)startIndices.size()) {
        startIndices[dim] += firstUsed;
        limitIndices[dim] += firstUsed;
      }

      auto sliceOp = stablehlo::SliceOp::create(
          rewriter, op.getLoc(), op.getOperand(),
          rewriter.getDenseI64ArrayAttr(startIndices),
          rewriter.getDenseI64ArrayAttr(limitIndices),
          rewriter.getDenseI64ArrayAttr(strides));

      // Propagate sharding if present
      if (auto shard = sdy::getShardingPerValue(op)) {
        sdy::setShardings(sliceOp, shard);
      }

      rewriter.replaceAllUsesWith(op.getResult(firstUsed), sliceOp.getResult());
      rewriter.eraseOp(op);
      return success();
    }

    // Otherwise, create a smaller MultiSliceOp if we can reduce the range
    if (newAmount != amount) {
      auto startIndices = SmallVector<int64_t>(op.getStartIndices());
      auto limitIndices = SmallVector<int64_t>(op.getLimitIndices());
      int32_t dim = op.getDimension();

      // Shift indices to account for trimmed results at the beginning
      if (dim >= 0 && dim < (int64_t)startIndices.size()) {
        startIndices[dim] += firstUsed;
        limitIndices[dim] += firstUsed;
      }

      // Determine result types
      auto resultType = cast<RankedTensorType>(op.getResultTypes().front());
      SmallVector<Type> resultTypes;
      for (int i = 0; i < newAmount + 1; i++) {
        resultTypes.push_back(resultType);
      }

      auto newOp = enzymexla::MultiSliceOp::create(
          rewriter, op.getLoc(), resultTypes, op.getOperand(), startIndices,
          limitIndices, op.getStrides(), op.getDimension(), newAmount);

      // Propagate sharding if present
      if (auto shard = sdy::getShardingPerValue(op)) {
        sdy::setShardings(newOp, shard);
      }

      // Map old results to new results
      int newIdx = 0;
      for (int oldIdx = firstUsed; oldIdx <= lastUsed; oldIdx++) {
        if (used[oldIdx]) {
          op.getResult(oldIdx).replaceAllUsesWith(newOp.getResult(newIdx));
        }
        newIdx++;
      }

      rewriter.eraseOp(op);
      return success();
    }

    return failure();
  }
};

// Pattern to lower MultiSliceOp into individual SliceOps
struct LowerMultiSlice final
    : CheckedOpRewritePattern<enzymexla::MultiSliceOp, LowerMultiSlice> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::MultiSliceOp op,
                                    PatternRewriter &rewriter) const {
    int32_t amount = op.getAmount();
    int32_t totalResults = amount + 1;
    int32_t dim = op.getDimension();

    auto baseStartIndices = SmallVector<int64_t>(op.getStartIndices());
    auto baseLimitIndices = SmallVector<int64_t>(op.getLimitIndices());
    auto strides = SmallVector<int64_t>(op.getStrides());

    // Get sharding info if present
    auto shard = sdy::getShardingPerValue(op);

    SmallVector<Value> replacements(totalResults);

    for (int i = 0; i < totalResults; i++) {
      // Copy and adjust indices for this slice
      auto startIndices = baseStartIndices;
      auto limitIndices = baseLimitIndices;

      if (dim >= 0 && dim < (int64_t)startIndices.size()) {
        startIndices[dim] += i;
        limitIndices[dim] += i;
      }

      auto sliceOp = stablehlo::SliceOp::create(
          rewriter, op.getLoc(), op.getOperand(),
          rewriter.getDenseI64ArrayAttr(startIndices),
          rewriter.getDenseI64ArrayAttr(limitIndices),
          rewriter.getDenseI64ArrayAttr(strides));

      // Propagate sharding if present
      if (shard) {
        sdy::TensorShardingAttr shards[1] = {shard.getShardings()[i]};
        auto shard2 =
            sdy::TensorShardingPerValueAttr::get(sliceOp.getContext(), shards);
        sdy::setShardings(sliceOp, shard2);
      }

      replacements[i] = sliceOp.getResult();
    }

    rewriter.replaceOp(op, replacements);

    return success();
  }
};

struct LowerMultiPad final
    : CheckedOpRewritePattern<enzymexla::MultiPadOp, LowerMultiPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::MultiPadOp op,
                                    PatternRewriter &rewriter) const {
    int32_t amount = op.getAmount();
    int32_t totalResults = amount + 1;
    int32_t dim = op.getDimension();

    auto input = op.getOperand();
    auto paddingValue = op.getPaddingValue();

    auto inputType = cast<RankedTensorType>(input.getType());
    int64_t rank = inputType.getRank();

    auto shard = sdy::getShardingPerValue(op);

    SmallVector<Value> replacements(totalResults);

    for (int idx = 0; idx < totalResults; idx++) {
      SmallVector<int64_t> low(rank, 0);
      SmallVector<int64_t> high(rank, 0);
      SmallVector<int64_t> interior(rank, 0);

      high[dim] = idx;
      low[dim] = amount - idx;

      auto padOp =
          stablehlo::PadOp::create(rewriter, op.getLoc(), input, paddingValue,
                                   rewriter.getDenseI64ArrayAttr(low),
                                   rewriter.getDenseI64ArrayAttr(high),
                                   rewriter.getDenseI64ArrayAttr(interior));

      if (shard) {
        sdy::TensorShardingAttr shards[1] = {shard.getShardings()[idx]};
        auto shard2 =
            sdy::TensorShardingPerValueAttr::get(padOp.getContext(), shards);
        sdy::setShardings(padOp, shard2);
      }

      replacements[idx] = padOp.getResult();
    }

    rewriter.replaceOp(op, replacements);

    return success();
  }
};

struct RecognizeMultiRotate
    : public CheckedOpRewritePattern<enzymexla::RotateOp,
                                     RecognizeMultiRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  // Convert rotate amounts to a (smaller) amount around using negative avlues
  // to indicate right rotation.
  static int32_t getBidirectionalAmount(enzymexla::RotateOp op) {
    RankedTensorType type = op.getOperand().getType();
    int32_t amount = op.getAmount();
    int64_t dimSize = type.getDimSize(op.getDimension());
    if (ShapedType::isDynamic(dimSize) || dimSize > INT32_MAX ||
        dimSize < INT32_MIN) {
      // Can't determine size, return original amount
      return amount;
    }

    if (amount > dimSize / 2) {
      amount = amount - (int32_t)dimSize; // rotate right instead of left
    } else if (amount < -dimSize / 2) {
      amount = amount + (int32_t)dimSize; // rotate left instead of right
    }
    return amount;
  }

  LogicalResult matchAndRewriteImpl(enzymexla::RotateOp op,
                                    PatternRewriter &rewriter) const {
    Value input = op.getOperand();
    int64_t dimension = op.getDimension();
    int32_t thisAmount = getBidirectionalAmount(op);

    // Collect all RotateOps operating on the same input with the same dimension
    SmallVector<enzymexla::RotateOp> rotates;
    DenseSet<int32_t> amountSet;

    for (Operation *user : input.getUsers()) {
      if (auto rotate = dyn_cast<enzymexla::RotateOp>(user)) {
        if (rotate.getDimension() == dimension && !rotate.use_empty()) {
          rotates.push_back(rotate);
          amountSet.insert(getBidirectionalAmount(rotate));
        }
      }
    }

    // Need at least 2 rotations to consider combining
    if (rotates.size() < 2)
      return failure();

    // Sort amounts to find contiguous groups
    SmallVector<int32_t> sortedAmounts(amountSet.begin(), amountSet.end());
    llvm::sort(sortedAmounts);

    // Find all contiguous groups (no gaps in amounts)
    SmallVector<std::pair<int32_t, int32_t>>
        contiguousGroups; // (start, end) inclusive
    int32_t groupStart = sortedAmounts[0];
    int32_t groupEnd = sortedAmounts[0];

    for (size_t i = 1; i < sortedAmounts.size(); i++) {
      if (sortedAmounts[i] == groupEnd + 1) {
        // Extend current group
        groupEnd = sortedAmounts[i];
      } else {
        // Save current group and start new one
        contiguousGroups.push_back({groupStart, groupEnd});
        groupStart = sortedAmounts[i];
        groupEnd = sortedAmounts[i];
      }
    }
    contiguousGroups.push_back({groupStart, groupEnd});

    // Collect all amounts from groups that contain identity or neighbor it
    // Groups that neighbor identity will be connected through 0
    SmallVector<int32_t> qualifyingAmounts;

    for (auto &[start, end] : contiguousGroups) {
      bool containsIdentity = (start <= 0 && end >= 0);
      bool neighborsIdentity = (end == -1 || start == 1);

      if (containsIdentity || neighborsIdentity) {
        // Add all actual rotate amounts from this group
        for (int32_t a = start; a <= end; a++) {
          if (amountSet.contains(a)) {
            qualifyingAmounts.push_back(a);
          }
        }
      }
    }

    // No qualifying groups found
    if (qualifyingAmounts.size() < 2)
      return failure();

    // Determine the multirotate range from qualifying amounts, extended to
    // include 0
    int32_t rangeStart =
        *std::min_element(qualifyingAmounts.begin(), qualifyingAmounts.end());
    int32_t rangeEnd =
        *std::max_element(qualifyingAmounts.begin(), qualifyingAmounts.end());

    // Extend to include identity (0) if not already included
    if (rangeStart > 0)
      rangeStart = 0;
    if (rangeEnd < 0)
      rangeEnd = 0;

    // Check if this op is part of the selected range
    if (thisAmount < rangeStart || thisAmount > rangeEnd)
      return failure();

    // Find the minimum amount among actual rotates in the selected range
    // to use as the canonical trigger (prevents processing same group multiple
    // times)
    int32_t minAmountInRange = INT32_MAX;
    int rotatesToCombine = 0;
    for (auto rotate : rotates) {
      int32_t amt = getBidirectionalAmount(rotate);
      if (amt >= rangeStart && amt <= rangeEnd) {
        minAmountInRange = std::min(minAmountInRange, amt);
        rotatesToCombine++;
      }
    }

    // Only proceed if this op has the minimum amount in the range
    if (thisAmount != minAmountInRange)
      return failure();

    // Need at least 2 rotations in the selected range to combine
    if (rotatesToCombine < 2)
      return failure();

    // Check that all rotates in the range have the same sharding
    std::optional<sdy::TensorShardingPerValueAttr> commonSharding;
    bool shardingMismatch = false;
    for (auto rotate : rotates) {
      int32_t amt = getBidirectionalAmount(rotate);
      if (amt >= rangeStart && amt <= rangeEnd) {
        auto shardPerValue = sdy::getShardingPerValue(rotate);
        if (!commonSharding.has_value()) {
          commonSharding = shardPerValue;
        } else if (commonSharding.value() != shardPerValue) {
          shardingMismatch = true;
          break;
        }
      }
    }
    if (shardingMismatch)
      return failure();

    // Calculate left and right amounts for MultiRotateOp
    // leftAmount covers positive rotations (rotate left)
    // rightAmount covers negative rotations (rotate right)
    int32_t leftAmount = rangeEnd > 0 ? rangeEnd : 0;
    int32_t rightAmount = rangeStart < 0 ? -rangeStart : 0;
    int32_t totalResults = leftAmount + rightAmount + 1;

    // Create the MultiRotateOp
    rewriter.setInsertionPointAfterValue(input);
    auto newOp = enzymexla::MultiRotateOp::create(
        rewriter, op.getLoc(), SmallVector<Type>(totalResults, input.getType()),
        input, op.getDimension(), leftAmount, rightAmount);

    // Propagate sharding if present (all rotates have the same sharding)
    if (commonSharding.has_value() && commonSharding.value()) {
      auto shardings = commonSharding.value().getShardings();
      if (!shardings.empty()) {
        sdy::TensorShardingAttr singleShard = shardings[0];
        SmallVector<sdy::TensorShardingAttr> newShardings(totalResults,
                                                          singleShard);
        sdy::setShardings(newOp, sdy::TensorShardingPerValueAttr::get(
                                     op.getContext(), newShardings));
      }
    }

    // Replace rotations that fall within the selected range
    for (auto rotate : rotates) {
      int32_t amt = getBidirectionalAmount(rotate);
      if (amt >= rangeStart && amt <= rangeEnd) {
        // Result index for amount 'a' is: leftAmount - a
        int32_t resultIdx = leftAmount - amt;
        rewriter.replaceOp(rotate, newOp.getResult(resultIdx));
      }
      // Rotations outside the range are left unchanged
    }

    return success();
  }
};

// Pattern to lower MultiRotateOp into individual RotateOps
struct LowerMultiRotate final
    : CheckedOpRewritePattern<enzymexla::MultiRotateOp, LowerMultiRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::MultiRotateOp op,
                                    PatternRewriter &rewriter) const {
    return lowerMultiRotateToRotates(op, rewriter);
  }
};

// Pattern to reduce MultiRotateOp when some results are unused
struct ReduceUnusedMultiRotate final
    : CheckedOpRewritePattern<enzymexla::MultiRotateOp,
                              ReduceUnusedMultiRotate> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::MultiRotateOp op,
                                    PatternRewriter &rewriter) const {
    int32_t leftAmount = op.getLeftAmount();
    int32_t rightAmount = op.getRightAmount();
    int32_t totalResults = leftAmount + rightAmount + 1;

    // Check which results are actually used
    SmallVector<bool> used(totalResults, false);
    int usedCount = 0;
    for (int i = 0; i < totalResults; i++) {
      if (!op.getResult(i).use_empty()) {
        used[i] = true;
        usedCount++;
      }
    }

    // If all results are used, nothing to optimize
    if (usedCount == totalResults)
      return failure();

    // If no results are used, this should be handled by dead code elimination
    if (usedCount == 0) {
      rewriter.eraseOp(op);
      return success();
    }

    // Find the range of used results
    int firstUsed = -1, lastUsed = -1;
    for (int i = 0; i < totalResults; i++) {
      if (used[i]) {
        if (firstUsed == -1)
          firstUsed = i;
        lastUsed = i;
      }
    }

    // Calculate new left and right amounts
    int centerIdx = leftAmount;
    int newLeftAmount = centerIdx - firstUsed;
    int newRightAmount = lastUsed - centerIdx;

    // If only one result is used, replace with a single RotateOp
    if (usedCount == 1) {
      int usedIdx = firstUsed;
      int amount = centerIdx -
                   usedIdx; // Positive = rotate left, negative = rotate right
      auto operandType = op.getOperand().getType();
      if (amount < 0) {
        amount += operandType.getShape()[op.getDimension()];
      }

      auto rotateOp = enzymexla::RotateOp::create(
          rewriter, op.getLoc(), op.getOperand().getType(), op.getOperand(),
          amount, op.getDimension());
      // Propagate sharding if present
      if (auto shard = sdy::getShardingPerValue(op)) {
        sdy::setShardings(rotateOp, shard);
      }

      rewriter.replaceAllUsesWith(op.getResult(usedIdx), rotateOp.getResult());
      rewriter.eraseOp(op);
      return success();
    }

    // Otherwise, create a smaller MultiRotateOp
    if (newLeftAmount != leftAmount || newRightAmount != rightAmount) {
      auto newOp = enzymexla::MultiRotateOp::create(
          rewriter, op.getLoc(),
          SmallVector<Type>(newLeftAmount + newRightAmount + 1,
                            op.getOperand().getType()),
          op.getOperand(), op.getDimension(), newLeftAmount, newRightAmount);

      // Propagate sharding if present
      if (auto shard = sdy::getShardingPerValue(op)) {
        sdy::setShardings(newOp, shard);
      }

      // Map old results to new results
      SmallVector<Value> replacements(totalResults);
      int newIdx = 0;
      for (int oldIdx = firstUsed; oldIdx <= lastUsed; oldIdx++) {
        replacements[oldIdx] = newOp.getResult(newIdx++);
      }

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

    return failure();
  }
};

// Replace other uses of multirotate with the unrotated result of the
// multirotate when possible.
struct UseMultiRotateNeutralResult
    : public CheckedOpRewritePattern<enzymexla::MultiRotateOp,
                                     UseMultiRotateNeutralResult> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(enzymexla::MultiRotateOp op,
                                    PatternRewriter &rewriter) const {
    Value neutral = op.getResult(op.getLeftAmount());
    DominanceInfo domInfo;
    bool anyReplaced = false;
    rewriter.replaceUsesWithIf(op.getOperand(), neutral, [&](OpOperand &use) {
      Operation *user = use.getOwner();
      if (!domInfo.properlyDominates(op, user))
        return false;

      anyReplaced = true;
      return true;
    });
    return success(anyReplaced);
  }
};

struct RecognizeMultiPad final
    : CheckedOpRewritePattern<stablehlo::PadOp, RecognizeMultiPad> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::PadOp op,
                                    PatternRewriter &rewriter) const {
    Value input = op.getOperand();
    Value paddingValue = op.getPaddingValue();

    auto opLow = op.getEdgePaddingLow();
    auto opHigh = op.getEdgePaddingHigh();
    auto opInterior = op.getInteriorPadding();

    int64_t rank = cast<RankedTensorType>(input.getType()).getRank();

    // Find the single dimension that is padded.
    int32_t dimension = -1;
    for (int64_t i = 0; i < rank; ++i) {
      if (opLow[i] != 0 || opHigh[i] != 0 || opInterior[i] != 0) {
        if (dimension != -1) {
          return failure(); // Pads multiple dimensions, not supported by
                            // MultiPadOp
        }
        dimension = i;
      }
    }

    if (dimension == -1) {
      return failure(); // No padding at all?
    }

    if (opInterior[dimension] != 0) {
      return failure(); // Non-zero interior padding not supported
    }

    int64_t rootTotalPadding = opLow[dimension] + opHigh[dimension];

    // Collect all PadOps operating on the same input with the same padding
    // value and that only pad the same dimension and have 0 interior padding.
    SmallVector<stablehlo::PadOp> pads;

    for (Operation *user : input.getUsers()) {
      if (auto pad = dyn_cast<stablehlo::PadOp>(user)) {
        if (pad.getPaddingValue() != paddingValue)
          continue;

        auto low = pad.getEdgePaddingLow();
        auto high = pad.getEdgePaddingHigh();
        auto interior = pad.getInteriorPadding();

        int64_t totalPadding = low[dimension] + high[dimension];
        if (totalPadding != rootTotalPadding)
          continue;

        bool match = true;
        for (int64_t i = 0; i < rank; ++i) {
          if (i == dimension) {
            if (interior[i] != 0) {
              match = false;
              break;
            }
          } else {
            if (low[i] != 0 || high[i] != 0 || interior[i] != 0) {
              match = false;
              break;
            }
          }
        }

        if (match) {
          pads.push_back(pad);
        }
      }
    }

    if (pads.size() < 2)
      return failure(); // Need at least 2 pads to combine

    std::optional<sdy::TensorShardingPerValueAttr> commonSharding;
    bool shardingMismatch = false;
    for (auto pad : pads) {
      auto shardPerValue = sdy::getShardingPerValue(pad);
      if (!commonSharding.has_value()) {
        commonSharding = shardPerValue;
      } else if (commonSharding.value() != shardPerValue) {
        shardingMismatch = true;
        break;
      }
    }
    if (shardingMismatch)
      return failure();

    rewriter.setInsertionPointAfterValue(input);

    int64_t amount = rootTotalPadding;
    int64_t numResults = amount + 1;
    SmallVector<Type> resultTypes(numResults);

    auto inputType = cast<RankedTensorType>(input.getType());
    auto shape = llvm::to_vector(inputType.getShape());
    shape[dimension] += amount;
    Type resultType = RankedTensorType::get(shape, inputType.getElementType());

    for (int64_t i = 0; i < numResults; ++i) {
      resultTypes[i] = resultType;
    }

    auto multiPad =
        enzymexla::MultiPadOp::create(rewriter, op.getLoc(), resultTypes, input,
                                      paddingValue, dimension, amount);

    if (commonSharding.has_value() && commonSharding.value()) {
      auto shardings = commonSharding.value().getShardings();
      if (!shardings.empty()) {
        sdy::TensorShardingAttr singleShard = shardings[0];
        SmallVector<sdy::TensorShardingAttr> newShardings(numResults,
                                                          singleShard);
        sdy::setShardings(multiPad, sdy::TensorShardingPerValueAttr::get(
                                        op.getContext(), newShardings));
      }
    }

    for (auto pad : pads) {
      int64_t idx = pad.getEdgePaddingHigh()[dimension];
      rewriter.replaceOp(pad, multiPad.getResult(idx));
    }

    return success();
  }
};

struct ScatterOpCanon final
    : CheckedOpRewritePattern<stablehlo::ScatterOp, ScatterOpCanon> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp op,
                                    PatternRewriter &rewriter) {
    if (trySimplifyTrivialScatter(op, rewriter).succeeded()) {
      return success();
    }

    if (trySimplifyWithIotaIndexing(op, rewriter).succeeded()) {
      return success();
    }

    return failure();
  }

private:
  // if the scatter is essentially `return arg0`, then
  // we can remove the scatter op altogether
  LogicalResult trySimplifyTrivialScatter(stablehlo::ScatterOp op,
                                          PatternRewriter &rewriter) {
    auto &block = op.getUpdateComputation().front();
    if (!isSetindexBlock(&block, [&](stablehlo::ReturnOp retOp) {
          return retOp.getOperand(0) == block.getArgument(0);
        })) {
      return failure();
    }

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

  LogicalResult
  scatterRegionToFunction(stablehlo::ScatterOp op, func::FuncOp &func,
                          Type &resElemType,
                          llvm::SetVector<Value> &additionalInputs,
                          PatternRewriter &rewriter) const {
    auto modOp = op->getParentOfType<ModuleOp>();
    if (!modOp) {
      return failure();
    }

    auto &block = op.getUpdateComputation().front();

    SmallVector<Type> returnTypes;
    ExtractBlockIntoFunction(&block, modOp, func, additionalInputs, returnTypes,
                             rewriter);

    if (returnTypes.size() != 1) {
      return failure();
    }

    resElemType = cast<RankedTensorType>(returnTypes[0]).getElementType();
    return success();
  }

  /// Helper function to create batched scatter operations.
  /// This encapsulates the common logic for:
  /// 1. Calling scatterRegionToFunction to extract the region into a function
  /// 2. Broadcasting additional inputs to match batch dimensions
  /// 3. Creating BatchOps and DynamicUpdateSliceOps for each input
  /// 4. Inlining the batch operations
  LogicalResult createBatchedScatterOps(
      stablehlo::ScatterOp op, ValueRange inputs, ArrayRef<Value> slices,
      ArrayRef<Value> updates, Value startIndex, int64_t count,
      PatternRewriter &rewriter, SmallVectorImpl<Value> &results) const {
    // Use batching to compute the result of the region
    func::FuncOp func;
    Type resElemType;
    llvm::SetVector<Value> additionalInputsOrig;
    if (failed(scatterRegionToFunction(op, func, resElemType,
                                       additionalInputsOrig, rewriter))) {
      return failure();
    }

    SmallVector<Value> bcastedExtraInputs;
    for (auto addArg : additionalInputsOrig) {
      auto origType = cast<RankedTensorType>(addArg.getType());
      auto outShape = llvm::to_vector(origType.getShape());
      outShape.insert(outShape.begin(), count);

      if (count == 1) {
        // For count=1, we can use reshape instead of broadcast
        bcastedExtraInputs.push_back(stablehlo::ReshapeOpCreate(
            rewriter, op.getLoc(), addArg, outShape));
      } else {
        // For count>1, we need to broadcast
        SmallVector<int64_t> bcastDims(origType.getRank());
        std::iota(bcastDims.begin(), bcastDims.end(), 1);
        bcastedExtraInputs.push_back(stablehlo::BroadcastInDimOp::create(
            rewriter, op.getLoc(),
            RankedTensorType::get(outShape, origType.getElementType()), addArg,
            rewriter.getDenseI64ArrayAttr(bcastDims)));
      }
    }

    SmallVector<enzyme::BatchOp> batchOps;
    for (auto [input, slice, update] :
         llvm::zip_equal(inputs, slices, updates)) {
      SmallVector<Value> allInputs = {slice, update};
      allInputs.append(bcastedExtraInputs.begin(), bcastedExtraInputs.end());

      auto newBatchOp = enzyme::BatchOp::create(
          rewriter, op.getLoc(), RankedTensorType::get({count}, resElemType),
          mlir::FlatSymbolRefAttr::get(op.getContext(), func.getName()),
          allInputs, rewriter.getDenseI64ArrayAttr({count}));

      auto newDusOp = stablehlo::DynamicUpdateSliceOp::create(
          rewriter, op.getLoc(), input, newBatchOp.getResult(0),
          ValueRange(startIndex));
      results.push_back(newDusOp.getResult());
      batchOps.push_back(newBatchOp);
    }

    auto funcOpInterface = cast<FunctionOpInterface>(func.getOperation());
    for (auto batchOp : batchOps) {
      enzyme::batchutils::batchOperationInline(rewriter, batchOp,
                                               funcOpInterface);
    }

    return success();
  }

  LogicalResult trySimplifyWithIotaIndexing(stablehlo::ScatterOp op,
                                            PatternRewriter &rewriter) const {
    auto inputs = op.getInputs();
    auto inputTy = cast<RankedTensorType>(inputs[0].getType());
    // TODO: check if this optimization is possible for higher dimenional
    //       tensors?
    if (inputTy.getRank() != 1) {
      return failure();
    }

    auto indices = op.getScatterIndices();

    // size 1 index can be trivially simplified to a DUS
    if (indices.getType().getNumElements() == 1) {
      auto scalarIndex =
          stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), indices, {});

      // Extract all of the values being updated
      SmallVector<Value> slices, updates;
      for (auto [input, update] : llvm::zip_equal(inputs, op.getUpdates())) {
        slices.push_back(stablehlo::DynamicSliceOpCreate(
            rewriter, op.getLoc(), input, {scalarIndex}, {1}));
        updates.push_back(
            stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), update, {1}));
      }

      SmallVector<Value> results;
      if (failed(createBatchedScatterOps(op, inputs, slices, updates,
                                         scalarIndex, /*count=*/1, rewriter,
                                         results))) {
        return failure();
      }

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

    auto iotaLike = detectIotaLikeTensor(indices);
    if (!iotaLike) {
      return failure();
    }

    auto iota = *iotaLike;
    auto dimNumbers = op.getScatterDimensionNumbers();

    if (dimNumbers.getScatterDimsToOperandDims().size() <= iota.dimension ||
        dimNumbers.getScatterDimsToOperandDims()[iota.dimension] != 0) {
      return failure();
    }

    auto indicesTy = cast<RankedTensorType>(indices.getType());
    int64_t indexVectorDim = dimNumbers.getIndexVectorDim();
    if (indexVectorDim < indicesTy.getRank() &&
        indicesTy.getDimSize(indexVectorDim) != 1) {
      return failure();
    }

    int64_t start = cast<IntegerAttr>(iota.start).getValue().getSExtValue();
    int64_t count = indicesTy.getDimSize(iota.dimension);
    int64_t stride = cast<IntegerAttr>(iota.scale).getValue().getSExtValue();

    if (std::abs(stride) != 1) {
      return failure(); // DUS doesn't support strided store.
    }

    auto updateTy = cast<RankedTensorType>(op.getUpdates()[0].getType());
    if (updateTy.getNumElements() != count) {
      LLVM_DEBUG(op->emitError("expected num elements of result to match"));
      return failure();
    }

    int64_t limit;
    bool needsReverse = false;
    if (stride > 0) {
      limit = start + count * stride;
    } else {
      needsReverse = true;
      // For negative stride, the iota accesses indices in descending order:
      //   start, start+stride, start+2*stride, ... (where stride < 0)
      // The smallest index is start + (count-1)*stride, the largest is start.
      // We slice in the forward direction with abs(stride), then reverse.
      limit = start + 1;                    // one past the largest index
      start = start + (count - 1) * stride; // smallest index
      stride = -stride;
    }

    if (limit > inputTy.getDimSize(0)) { // gather clamps indices
      return failure();
    }

    // Extract all of the values being updated
    SmallVector<Value> slices, updates;
    for (auto [input, update] : llvm::zip_equal(inputs, op.getUpdates())) {
      auto slice = stablehlo::SliceOpCreate(rewriter, op.getLoc(), input,
                                            {start}, {limit}, {stride});
      Value newUpdate =
          stablehlo::ReshapeOpCreate(rewriter, op.getLoc(), update, {count});
      if (needsReverse) {
        newUpdate =
            stablehlo::ReverseOp::create(rewriter, op.getLoc(), newUpdate,
                                         rewriter.getDenseI64ArrayAttr({0}));
      }
      slices.push_back(slice);
      updates.push_back(newUpdate);
    }

    auto constStartVal = stablehlo::ConstantOp::create(
        rewriter, op.getLoc(),
        cast<ElementsAttr>(
            makeAttr(RankedTensorType::get({}, rewriter.getI32Type()), start)));

    SmallVector<Value> results;
    if (failed(createBatchedScatterOps(op, inputs, slices, updates,
                                       constStartVal, count, rewriter,
                                       results))) {
      return failure();
    }

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

// Helper function to build scatter update computation region based on the
// scatter operation kind. For constant-based operations, the constant value
// must be provided as the real or imaginary part (caller extracts from
// complex).
static void buildScatterUpdateRegion(PatternRewriter &rewriter, Location loc,
                                     stablehlo::ScatterOp scatterOp,
                                     stablehlo::ScatterOpKind kind,
                                     Type elemType,
                                     Attribute constantValue = nullptr) {
  OpBuilder::InsertionGuard guard(rewriter);
  auto &updateRegion = scatterOp.getUpdateComputation();
  auto *block = rewriter.createBlock(&updateRegion);
  auto argType = RankedTensorType::get({}, elemType);
  block->addArgument(argType, loc);
  block->addArgument(argType, loc);
  rewriter.setInsertionPointToStart(block);

  Value result;
  switch (kind) {
  case stablehlo::ScatterOpKind::Setindex:
    // Setindex: return the update value (arg1)
    result = block->getArgument(1);
    break;
  case stablehlo::ScatterOpKind::ConstantSetindex: {
    // ConstantSetindex: return the constant value
    assert(constantValue && "ConstantSetindex requires a constant value");
    auto constAttr = DenseElementsAttr::get(argType, constantValue);
    result = stablehlo::ConstantOp::create(rewriter, loc, constAttr);
    break;
  }
  case stablehlo::ScatterOpKind::Add:
    // Add: return arg0 + arg1
    result = stablehlo::AddOp::create(rewriter, loc, block->getArgument(0),
                                      block->getArgument(1));
    break;
  case stablehlo::ScatterOpKind::AddConstantUpdate: {
    // AddConstantUpdate: return arg0 + constant (constant replaces update/arg1)
    assert(constantValue && "AddConstantUpdate requires a constant value");
    auto constAttr = DenseElementsAttr::get(argType, constantValue);
    auto constOp = stablehlo::ConstantOp::create(rewriter, loc, constAttr);
    result =
        stablehlo::AddOp::create(rewriter, loc, block->getArgument(0), constOp);
    break;
  }
  case stablehlo::ScatterOpKind::AddConstantInput: {
    // AddConstantInput: return constant + arg1 (constant replaces input/arg0)
    assert(constantValue && "AddConstantInput requires a constant value");
    auto constAttr = DenseElementsAttr::get(argType, constantValue);
    auto constOp = stablehlo::ConstantOp::create(rewriter, loc, constAttr);
    result =
        stablehlo::AddOp::create(rewriter, loc, constOp, block->getArgument(1));
    break;
  }
  case stablehlo::ScatterOpKind::Sub:
    // Sub: return arg0 - arg1
    result = stablehlo::SubtractOp::create(rewriter, loc, block->getArgument(0),
                                           block->getArgument(1));
    break;
  default:
    llvm_unreachable("Unsupported scatter operation kind");
  }

  stablehlo::ReturnOp::create(rewriter, loc, result);
}

static Value concatRealAndImag(PatternRewriter &rewriter, Location loc,
                               Value realPart, Value imagPart) {
  auto tensorType = cast<RankedTensorType>(realPart.getType());
  auto shape = tensorType.getShape();

  // Reshape each to add a leading dimension of size 1:
  //   [d0, ..., dn] -> [1, d0, ..., dn]
  SmallVector<int64_t> reshapedShape;
  reshapedShape.push_back(1);
  reshapedShape.append(shape.begin(), shape.end());

  auto realReshaped =
      stablehlo::ReshapeOpCreate(rewriter, loc, realPart, reshapedShape);
  auto imagReshaped =
      stablehlo::ReshapeOpCreate(rewriter, loc, imagPart, reshapedShape);

  // Concatenate along the new leading dimension (dim = 0):
  //   [1, d0, ..., dn] + [1, d0, ..., dn] -> [2, d0, ..., dn]
  return stablehlo::ConcatenateOp::create(
             rewriter, loc, ValueRange{realReshaped, imagReshaped},
             rewriter.getI64IntegerAttr(0))
      .getResult();
}

// Helper: Split a complex tensor into real+imag and concatenate along a new
// trailing dimension. Given a complex tensor of shape [d0, ..., dn],
// returns a real-valued tensor of shape [d0, ..., dn, 2] where [..., 0] =
// real and [..., 1] = imag.
static Value splitComplexAndConcat(PatternRewriter &rewriter, Location loc,
                                   Value complexTensor) {
  auto realPart = stablehlo::RealOp::create(rewriter, loc, complexTensor);
  auto imagPart = stablehlo::ImagOp::create(rewriter, loc, complexTensor);
  return concatRealAndImag(rewriter, loc, realPart, imagPart);
}

// Helper: Combine real and imaginary parts from a trailing dimension into a
// complex tensor. Given a real-valued tensor of shape [d0, ..., dn, 2],
// returns a complex tensor of shape [d0, ..., dn] where real comes from
// [..., 0] and imag from [..., 1].
static Value sliceAndCombineComplex(PatternRewriter &rewriter, Location loc,
                                    Value splitTensor) {
  auto tensorType = cast<RankedTensorType>(splitTensor.getType());
  auto shape = tensorType.getShape();
  int64_t rank = tensorType.getRank();

  // Result shape: [d0, ..., dn]
  SmallVector<int64_t> resultShape(shape.begin() + 1, shape.end());

  SmallVector<int64_t> sliceStart(rank, 0);
  SmallVector<int64_t> sliceLimit(shape.begin(), shape.end());
  SmallVector<int64_t> sliceStride(rank, 1);

  // Real part: slice [0:1, ...]
  sliceLimit[0] = 1;

  auto realSlice = stablehlo::SliceOpCreate(
      rewriter, loc, splitTensor, sliceStart, sliceLimit, sliceStride);

  // Imag part: slice [1:2, ...]
  sliceStart[0] = 1;
  sliceLimit[0] = 2;

  auto imagSlice = stablehlo::SliceOpCreate(
      rewriter, loc, splitTensor, sliceStart, sliceLimit, sliceStride);

  // Reshape to remove leading dim of 1
  auto realResult =
      stablehlo::ReshapeOpCreate(rewriter, loc, realSlice, resultShape);
  auto imagResult =
      stablehlo::ReshapeOpCreate(rewriter, loc, imagSlice, resultShape);

  // Combine
  return stablehlo::ComplexOp::create(rewriter, loc, realResult, imagResult);
}

// See discussion in
// https://github.com/EnzymeAD/Reactant.jl/pull/2219#issuecomment-3801834108
struct SplitComplexScatter final
    : CheckedOpRewritePattern<stablehlo::ScatterOp, SplitComplexScatter> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ScatterOp op,
                                    PatternRewriter &rewriter) const {
    auto checkCommonScatterOp = mlir::stablehlo::CheckCommonScatterOp(op);

    // Support Setindex, ConstantSetindex, Add, AddConstantUpdate,
    // AddConstantInput, and Sub scatter operations
    if (checkCommonScatterOp.kind != stablehlo::ScatterOpKind::Setindex &&
        checkCommonScatterOp.kind !=
            stablehlo::ScatterOpKind::ConstantSetindex &&
        checkCommonScatterOp.kind != stablehlo::ScatterOpKind::Add &&
        checkCommonScatterOp.kind !=
            stablehlo::ScatterOpKind::AddConstantUpdate &&
        checkCommonScatterOp.kind !=
            stablehlo::ScatterOpKind::AddConstantInput &&
        checkCommonScatterOp.kind != stablehlo::ScatterOpKind::Sub) {
      return failure();
    }

    auto elemType = cast<RankedTensorType>(op.getType(0)).getElementType();
    if (!isa<ComplexType>(elemType)) {
      return failure();
    }

    // Get the inner element type (e.g., f32 from complex<f32>)
    auto complexType = cast<ComplexType>(elemType);
    auto innerElemType = complexType.getElementType();
    auto loc = op.getLoc();

    // Extract real and imaginary parts of the constant if present
    Attribute realConstant = nullptr;
    Attribute imagConstant = nullptr;
    if (checkCommonScatterOp.constant) {
      auto complexVal =
          checkCommonScatterOp.constant.getSplatValue<mlir::Complex<APFloat>>();
      realConstant = FloatAttr::get(innerElemType, complexVal.real());
      imagConstant = FloatAttr::get(innerElemType, complexVal.imag());
    }

    // Determine the effective scatter kind for the float scatter.
    // For constant-based operations, we materialize the constants into
    // update tensors and convert to simpler operations.
    stablehlo::ScatterOpKind effectiveKind;
    switch (checkCommonScatterOp.kind) {
    case stablehlo::ScatterOpKind::Setindex:
    case stablehlo::ScatterOpKind::Add:
    case stablehlo::ScatterOpKind::Sub:
      effectiveKind = checkCommonScatterOp.kind;
      break;
    case stablehlo::ScatterOpKind::ConstantSetindex:
      // Materialize constant as update, use Setindex
      effectiveKind = stablehlo::ScatterOpKind::Setindex;
      break;
    case stablehlo::ScatterOpKind::AddConstantUpdate:
      // Materialize constant as update, use Add
      effectiveKind = stablehlo::ScatterOpKind::Add;
      break;
    case stablehlo::ScatterOpKind::AddConstantInput:
      // Precompute constant + update, use Setindex
      effectiveKind = stablehlo::ScatterOpKind::Setindex;
      break;
    default:
      return failure();
    }

    // Step 1: For each input/update pair, split into real/imag, reshape,
    // and concatenate along a new leading dimension of size 2.
    SmallVector<Value> concatInputs, concatUpdates;
    for (auto [input, update] :
         llvm::zip_equal(op.getInputs(), op.getUpdates())) {
      // Process input: real/imag -> reshape -> concat
      concatInputs.push_back(splitComplexAndConcat(rewriter, loc, input));

      // Process update: depends on scatter kind
      auto curUpdateType = cast<RankedTensorType>(update.getType());
      auto updateShape = curUpdateType.getShape();
      auto realUpdateType = RankedTensorType::get(updateShape, innerElemType);

      Value realUpdate, imagUpdate;
      if (checkCommonScatterOp.kind ==
          stablehlo::ScatterOpKind::ConstantSetindex) {
        // Create constant update tensors to replace the ignored updates
        realUpdate = stablehlo::ConstantOp::create(
            rewriter, loc,
            DenseElementsAttr::get(realUpdateType, realConstant));
        imagUpdate = stablehlo::ConstantOp::create(
            rewriter, loc,
            DenseElementsAttr::get(realUpdateType, imagConstant));
      } else if (checkCommonScatterOp.kind ==
                 stablehlo::ScatterOpKind::AddConstantUpdate) {
        // Materialize constants as updates (arg0 + constant in region)
        realUpdate = stablehlo::ConstantOp::create(
            rewriter, loc,
            DenseElementsAttr::get(realUpdateType, realConstant));
        imagUpdate = stablehlo::ConstantOp::create(
            rewriter, loc,
            DenseElementsAttr::get(realUpdateType, imagConstant));
      } else if (checkCommonScatterOp.kind ==
                 stablehlo::ScatterOpKind::AddConstantInput) {
        // Precompute constant + update, then use Setindex
        auto realUp = stablehlo::RealOp::create(rewriter, loc, update);
        auto imagUp = stablehlo::ImagOp::create(rewriter, loc, update);
        auto realConstOp = stablehlo::ConstantOp::create(
            rewriter, loc,
            DenseElementsAttr::get(realUpdateType, realConstant));
        auto imagConstOp = stablehlo::ConstantOp::create(
            rewriter, loc,
            DenseElementsAttr::get(realUpdateType, imagConstant));
        realUpdate =
            stablehlo::AddOp::create(rewriter, loc, realConstOp, realUp);
        imagUpdate =
            stablehlo::AddOp::create(rewriter, loc, imagConstOp, imagUp);
      } else {
        // Setindex, Add, Sub: just extract real/imag of update
        realUpdate = stablehlo::RealOp::create(rewriter, loc, update);
        imagUpdate = stablehlo::ImagOp::create(rewriter, loc, update);
      }

      auto concatUpdate =
          concatRealAndImag(rewriter, loc, realUpdate, imagUpdate);
      concatUpdates.push_back(concatUpdate);
    }

    // Step 2: Adjust scatter dimension numbers for the extra leading
    // dimension. The new leading input dim maps to the new leading update
    // dim via update_window_dims.
    auto dnums = op.getScatterDimensionNumbers();

    SmallVector<int64_t> newUpdateWindowDims;
    newUpdateWindowDims.push_back(0);
    for (auto d : dnums.getUpdateWindowDims()) {
      newUpdateWindowDims.push_back(d + 1);
    }

    SmallVector<int64_t> newInsertedWindowDims;
    for (auto d : dnums.getInsertedWindowDims()) {
      newInsertedWindowDims.push_back(d + 1);
    }

    SmallVector<int64_t> newInputBatchingDims;
    for (auto d : dnums.getInputBatchingDims()) {
      newInputBatchingDims.push_back(d + 1);
    }

    SmallVector<int64_t> newScatterDimsToOperandDims;
    for (auto d : dnums.getScatterDimsToOperandDims()) {
      newScatterDimsToOperandDims.push_back(d + 1);
    }

    auto newDnumsAttr = stablehlo::ScatterDimensionNumbersAttr::get(
        op.getContext(), newUpdateWindowDims, newInsertedWindowDims,
        newInputBatchingDims, dnums.getScatterIndicesBatchingDims(),
        newScatterDimsToOperandDims, dnums.getIndexVectorDim());

    // Step 3: Create a single scatter op on the concatenated tensors
    auto scatterOp = stablehlo::ScatterOp::create(
        rewriter, loc, concatInputs, op.getScatterIndices(), concatUpdates,
        newDnumsAttr, op.getIndicesAreSortedAttr(), op.getUniqueIndicesAttr());

    // Build update computation region with the effective kind
    buildScatterUpdateRegion(rewriter, loc, scatterOp, effectiveKind,
                             innerElemType);

    // Step 4: Slice results back into real and imag parts, reshape, and
    // recombine with ComplexOp.
    SmallVector<Value> results;
    for (auto result : scatterOp.getResults()) {
      results.push_back(sliceAndCombineComplex(rewriter, loc, result));
    }

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

struct SplitComplexGather final
    : CheckedOpRewritePattern<stablehlo::GatherOp, SplitComplexGather> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::GatherOp op,
                                    PatternRewriter &rewriter) const {
    auto resultType = cast<RankedTensorType>(op.getType());
    auto elemType = resultType.getElementType();

    // Only match if the result is a complex type
    if (!isa<ComplexType>(elemType)) {
      return failure();
    }

    auto complexType = cast<ComplexType>(elemType);
    auto innerElemType = complexType.getElementType();

    // Step 1: Extract real and imaginary parts, reshape each to add a
    // leading dimension of size 1, then concatenate along the leading
    // dimension to form a float tensor with leading dim 2.
    // e.g., tensor<4xcomplex<f32>> -> real: tensor<4xf32> -> tensor<1x4xf32>
    //                                 imag: tensor<4xf32> -> tensor<1x4xf32>
    //                                 concat: tensor<2x4xf32>
    auto concatOperand =
        splitComplexAndConcat(rewriter, op.getLoc(), op.getOperand());

    // Step 2: Adjust the gather dimension numbers and slice sizes to account
    // for the extra leading dimension.
    auto dnums = op.getDimensionNumbers();

    // Add the new leading dimension to offset_dims
    SmallVector<int64_t> newOffsetDims;
    newOffsetDims.push_back(0);
    for (auto d : dnums.getOffsetDims()) {
      newOffsetDims.push_back(d + 1);
    }
    // The new offset dim corresponds to the first dimension of the new result

    SmallVector<int64_t> newCollapsedSliceDims;
    for (auto d : dnums.getCollapsedSliceDims()) {
      newCollapsedSliceDims.push_back(d + 1);
    }

    SmallVector<int64_t> newOperandBatchingDims;
    for (auto d : dnums.getOperandBatchingDims()) {
      newOperandBatchingDims.push_back(d + 1);
    }

    SmallVector<int64_t> newStartIndexMap;
    for (auto d : dnums.getStartIndexMap()) {
      newStartIndexMap.push_back(d + 1);
    }

    auto newDnums = stablehlo::GatherDimensionNumbersAttr::get(
        op.getContext(), newOffsetDims, newCollapsedSliceDims,
        newOperandBatchingDims, dnums.getStartIndicesBatchingDims(),
        newStartIndexMap, dnums.getIndexVectorDim());

    // Prepend 2 to slice_sizes for the new leading dimension
    SmallVector<int64_t> newSliceSizes;
    newSliceSizes.push_back(2);
    newSliceSizes.append(op.getSliceSizes().begin(), op.getSliceSizes().end());
    auto newSliceSizesAttr =
        DenseI64ArrayAttr::get(op.getContext(), newSliceSizes);

    // Step 3: Perform the gather on the concatenated operand
    // Result type: trailing dim of size 2 + original result shape
    auto resultShape = resultType.getShape();
    SmallVector<int64_t> newResultShape;
    newResultShape.push_back(2);
    newResultShape.append(resultShape.begin(), resultShape.end());
    auto newResultType = RankedTensorType::get(newResultShape, innerElemType);

    auto gatherOp = stablehlo::GatherOp::create(
        rewriter, op.getLoc(), newResultType, concatOperand,
        op.getStartIndices(), newDnums, newSliceSizesAttr,
        op.getIndicesAreSortedAttr());

    // Step 4: Slice the result along the leading dimension to extract
    // real and imag parts, reshape, and recombine.
    auto complexResult =
        sliceAndCombineComplex(rewriter, op.getLoc(), gatherOp);

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

// binop(complex, complex) -> complex(binop, binop)
template <typename OpTy>
struct BinaryOpComplexSimplifyBase
    : public CheckedOpRewritePattern<OpTy, BinaryOpComplexSimplifyBase<OpTy>> {
  using CheckedOpRewritePattern<
      OpTy, BinaryOpComplexSimplifyBase<OpTy>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(OpTy op, PatternRewriter &rewriter) const {
    auto lhs = op.getLhs();
    auto rhs = op.getRhs();

    auto lhsComplexOp = lhs.template getDefiningOp<stablehlo::ComplexOp>();
    auto rhsComplexOp = rhs.template getDefiningOp<stablehlo::ComplexOp>();

    if (!lhsComplexOp || !rhsComplexOp ||
        !isOnlyUsedInOperation(lhsComplexOp, op) ||
        !isOnlyUsedInOperation(rhsComplexOp, op)) {
      return failure();
    }

    // If any of the components are zero, we can potentially do more aggressive
    // constant folding
    bool lhsRealIsZero =
        guaranteedPurelyImagResult(lhsComplexOp.getResult(), rewriter);
    bool lhsImagIsZero =
        guaranteedPurelyRealResult(lhsComplexOp.getResult(), rewriter);
    bool rhsRealIsZero =
        guaranteedPurelyImagResult(rhsComplexOp.getResult(), rewriter);
    bool rhsImagIsZero =
        guaranteedPurelyRealResult(rhsComplexOp.getResult(), rewriter);

    if (!(lhsRealIsZero || lhsImagIsZero || rhsRealIsZero || rhsImagIsZero)) {
      return failure();
    }

    auto createOrFold = [&](Value lhs, Value rhs, bool lhsIsZero,
                            bool rhsIsZero) -> Value {
      if constexpr (std::is_same_v<OpTy, stablehlo::AddOp>) {
        if (lhsIsZero) {
          return rhs;
        }
        if (rhsIsZero) {
          return lhs;
        }
      }
      if constexpr (std::is_same_v<OpTy, stablehlo::SubtractOp>) {
        if (rhsIsZero) {
          return lhs;
        }
        if (lhsIsZero) {
          return stablehlo::NegOp::create(rewriter, op.getLoc(), rhs);
        }
      }
      return OpTy::create(rewriter, op.getLoc(), lhs, rhs);
    };

    auto newReal = createOrFold(lhsComplexOp.getLhs(), rhsComplexOp.getLhs(),
                                lhsRealIsZero, rhsRealIsZero);
    auto newImag = createOrFold(lhsComplexOp.getRhs(), rhsComplexOp.getRhs(),
                                lhsImagIsZero, rhsImagIsZero);
    rewriter.replaceOpWithNewOp<stablehlo::ComplexOp>(op, newReal, newImag);
    return success();
  }
};

// Pattern: convert[wide_float->narrow_float](mul(convert[int->wide_float](x),
// val))
// -> mul(convert[int->narrow_float](x), convert[wide_float->narrow_float](val))
//
// This eliminates an unnecessary intermediate wide-float precision step when
// the result is immediately narrowed. The convert of a constant 'val' will be
// folded to a narrower-type constant by subsequent constant-folding passes.
//
// Example from iota-based coordinate calculations:
//   convert[f64->f32](mul(convert[i64->f64](iota+c), f64_const))
//   -> mul(convert[i64->f32](iota+c), f32_const)
struct ConvertMulConvert final
    : CheckedOpRewritePattern<stablehlo::ConvertOp, ConvertMulConvert> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    auto dstType = op.getType();
    auto dstElemTy = dstType.getElementType();
    // Destination must be float
    if (!isa<FloatType>(dstElemTy))
      return failure();

    // Source must be a wider float (narrowing conversion)
    auto srcElemTy =
        cast<RankedTensorType>(op.getOperand().getType()).getElementType();
    if (!isa<FloatType>(srcElemTy))
      return failure();
    if (srcElemTy.getIntOrFloatBitWidth() <= dstElemTy.getIntOrFloatBitWidth())
      return failure();

    // Operand must be a multiply in the wide float type
    auto mul = op.getOperand().getDefiningOp<stablehlo::MulOp>();
    if (!mul)
      return failure();

    // Find which operand of the multiply was converted from an integer type
    Value intVal;   // original integer value (before widening convert)
    Value otherVal; // the other mul operand (to be converted to narrow float)

    auto lhsConv = mul.getLhs().getDefiningOp<stablehlo::ConvertOp>();
    if (lhsConv &&
        isa<IntegerType>(cast<RankedTensorType>(lhsConv.getOperand().getType())
                             .getElementType())) {
      intVal = lhsConv.getOperand();
      otherVal = mul.getRhs();
    } else {
      auto rhsConv = mul.getRhs().getDefiningOp<stablehlo::ConvertOp>();
      if (!rhsConv || !isa<IntegerType>(
                          cast<RankedTensorType>(rhsConv.getOperand().getType())
                              .getElementType()))
        return failure();
      intVal = rhsConv.getOperand();
      otherVal = mul.getLhs();
    }

    // Build convert[int->dstFloat](intVal)
    auto intValType = cast<RankedTensorType>(intVal.getType());
    auto newIntConv = stablehlo::ConvertOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(intValType.getShape(), dstElemTy), intVal);

    // Build convert[wideFloat->dstFloat](otherVal)
    // If otherVal is a constant, this will be folded to a narrower constant.
    auto otherValType = cast<RankedTensorType>(otherVal.getType());
    auto newOtherConv = stablehlo::ConvertOp::create(
        rewriter, op.getLoc(),
        RankedTensorType::get(otherValType.getShape(), dstElemTy), otherVal);

    // Replace the convert(mul(...)) with mul in the narrow float type
    rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, newIntConv, newOtherConv);
    return success();
  }
};

// Pattern: convert(binop(convert(x), convert(y))) -> binop(x, y)
// when x and y have the same type as the outer convert result.
// Valid for operations like min/max where the binop result is preserved
// through matching convert-unconvert pairs.
template <typename BinOp>
struct ConvertBinopConvert final
    : CheckedOpRewritePattern<stablehlo::ConvertOp,
                              ConvertBinopConvert<BinOp>> {
  using CheckedOpRewritePattern<
      stablehlo::ConvertOp,
      ConvertBinopConvert<BinOp>>::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::ConvertOp op,
                                    PatternRewriter &rewriter) const {
    auto resultType = op.getType();

    if (!isa<FloatType>(resultType.getElementType()))
      return failure();

    auto binop = op.getOperand().template getDefiningOp<BinOp>();
    if (!binop)
      return failure();

    auto lhsConvert =
        binop.getLhs().template getDefiningOp<stablehlo::ConvertOp>();
    auto rhsConvert =
        binop.getRhs().template getDefiningOp<stablehlo::ConvertOp>();
    if (!lhsConvert || !rhsConvert)
      return failure();

    auto lhs = lhsConvert.getOperand();
    auto rhs = rhsConvert.getOperand();

    if (lhs.getType() != resultType || rhs.getType() != resultType)
      return failure();

    rewriter.replaceOpWithNewOp<BinOp>(op, lhs, rhs);
    return success();
  }
};

// Pattern: negate(reduce_window(x, 0, \acc,elem -> acc - elem))
//       -> reduce_window(x, 0, \acc,elem -> acc + elem)
//
// When the reduce_window body subtracts each element from the running
// accumulator starting at zero, the result is the negated window sum:
//   reduce_window(x, 0, sub) = 0 - x_0 - x_1 - ... = -sum(window(x))
// Negating that recovers the plain sum, which equals a reduce_window with add.
struct NegateReduceWindowSub final
    : CheckedOpRewritePattern<stablehlo::NegOp, NegateReduceWindowSub> {
  using CheckedOpRewritePattern::CheckedOpRewritePattern;

  LogicalResult matchAndRewriteImpl(stablehlo::NegOp op,
                                    PatternRewriter &rewriter) const {
    auto rw = op.getOperand().getDefiningOp<stablehlo::ReduceWindowOp>();
    if (!rw || !rw->hasOneUse())
      return failure();

    // Only handle single input + single init value
    if (rw.getInputs().size() != 1 || rw.getInitValues().size() != 1)
      return failure();

    // Inspect the reduce_window body
    Region &body = rw.getBody();
    if (!body.hasOneBlock())
      return failure();
    Block &block = body.front();
    if (block.getNumArguments() != 2)
      return failure();

    // Body must be: %res = subtract(%acc, %elem); return %res
    auto returnOp = dyn_cast<stablehlo::ReturnOp>(block.getTerminator());
    if (!returnOp || returnOp.getNumOperands() != 1)
      return failure();

    auto subOp = returnOp.getOperand(0).getDefiningOp<stablehlo::SubtractOp>();
    if (!subOp)
      return failure();

    // Must be acc - elem (arg0 - arg1)
    if (subOp.getLhs() != block.getArgument(0) ||
        subOp.getRhs() != block.getArgument(1))
      return failure();

    // Init value must be zero
    if (!matchPattern(rw.getInitValues()[0], m_AnyZeroFloat()))
      return failure();

    // Create a new reduce_window identical to the original except the body
    // uses addition instead of subtraction.
    Value inputs[1] = {rw.getInputs()[0]};
    auto newRw = stablehlo::ReduceWindowOp::create(
        rewriter, rw.getLoc(), rw.getResultTypes(), inputs, rw.getInitValues(),
        rw.getWindowDimensionsAttr(), rw.getWindowStridesAttr(),
        rw.getBaseDilationsAttr(), rw.getWindowDilationsAttr(),
        rw.getPaddingAttr());

    // Build the add body using the same block argument types as the original
    {
      OpBuilder::InsertionGuard guard(rewriter);
      auto argTypes = block.getArgumentTypes();
      SmallVector<Location> argLocs(argTypes.size(), rw.getLoc());
      auto *newBlock =
          rewriter.createBlock(&newRw.getBody(), {}, argTypes, argLocs);
      rewriter.setInsertionPointToStart(newBlock);
      auto addOp = stablehlo::AddOp::create(rewriter, rw.getLoc(),
                                            newBlock->getArgument(0),
                                            newBlock->getArgument(1));
      stablehlo::ReturnOp::create(rewriter, rw.getLoc(), addOp.getResult());
    }

    // Replace negate(reduce_window(..., sub)) with reduce_window(..., add)
    rewriter.replaceOp(op, newRw.getResults());
    rewriter.eraseOp(rw);
    return success();
  }
};

///////////////  End Imported from stablehlo

namespace mlir {
namespace enzyme {

// runs for all floats and integers only when a % b == 0 (i.e. the division is
// exact).
static bool divMulReassociable(mlir::Value res, mlir::Attribute aAttr,
                               mlir::Attribute bAttr) {
  auto st = mlir::dyn_cast<mlir::ShapedType>(res.getType());
  if (!st)
    return false;
  if (mlir::isa<mlir::FloatType>(st.getElementType()))
    return true;
  auto a = mlir::dyn_cast<mlir::DenseIntElementsAttr>(aAttr);
  auto b = mlir::dyn_cast<mlir::DenseIntElementsAttr>(bAttr);
  if (!a || !b)
    return false;
  auto exact = [](const llvm::APInt &av, const llvm::APInt &bv) {
    return !bv.isZero() && av.srem(bv).isZero();
  };
  bool aSplat = a.isSplat(), bSplat = b.isSplat();
  if (aSplat && bSplat)
    return exact(a.getSplatValue<llvm::APInt>(),
                 b.getSplatValue<llvm::APInt>());
  if (aSplat) {
    auto av = a.getSplatValue<llvm::APInt>();
    return llvm::all_of(b.getValues<llvm::APInt>(),
                        [&](const llvm::APInt &bv) { return exact(av, bv); });
  }
  if (bSplat) {
    auto bv = b.getSplatValue<llvm::APInt>();
    return llvm::all_of(a.getValues<llvm::APInt>(),
                        [&](const llvm::APInt &av) { return exact(av, bv); });
  }
  if (a.getNumElements() != b.getNumElements())
    return false;
  for (const auto [av, bv] : llvm::zip_equal(a.getValues<llvm::APInt>(),
                                             b.getValues<llvm::APInt>())) {
    if (!exact(av, bv))
      return false;
  }
  return true;
}

// clang-format off
#include "src/enzyme_ad/jax/Passes/StablehloOptPatterns.cpp.inc"

}; // namespace enzyme
}; // namespace mlir

#include "src/enzyme_ad/jax/Passes/AutoBatching.h"
#include "src/enzyme_ad/jax/Passes/EnzymeHLOUnroll.h"

#include "src/enzyme_ad/jax/Passes/EnzymeHLOPatterns.cpp.inc"
// clang-format on

void mlir::transform::addEnzymeHLOUnroll(RewritePatternSet &patterns,
                                         int64_t maxNumIterations,
                                         MLIRContext &context,
                                         PatternBenefit benefit) {
  patterns.insert<WhileUnroll>(maxNumIterations, 128, &context, benefit);
}

void mlir::transform::addPadDotGeneral(RewritePatternSet &patterns,
                                       bool postPad, MLIRContext &context,
                                       PatternBenefit benefit) {
  patterns.insert<PadDotGeneral>(postPad, &context, benefit);
}

void mlir::transform::addIotaSimplify(RewritePatternSet &patterns,
                                      int64_t maxConstantExpansion,
                                      MLIRContext &context,
                                      PatternBenefit benefit) {
  patterns.insert<IotaSimplify>(maxConstantExpansion, &context, benefit);
}

void mlir::transform::addRecognizeFromConstant(RewritePatternSet &patterns,
                                               int64_t minFoldSize,
                                               MLIRContext &context,
                                               PatternBenefit benefit) {
  patterns.insert<RecognizeFromConstant>(minFoldSize, &context, benefit);
}

void mlir::transform::addConcatConstProp(RewritePatternSet &patterns,
                                         int64_t maxConstantExpansion,
                                         MLIRContext &context,
                                         PatternBenefit benefit) {
  patterns.insert<ConcatConstProp>(maxConstantExpansion, &context, benefit);
}

void mlir::transform::addScatterConstFold(RewritePatternSet &patterns,
                                          int64_t maxConstantExpansion,
                                          MLIRContext &context,
                                          PatternBenefit benefit) {
  patterns.insert<ScatterConstFold>(maxConstantExpansion, &context, benefit);
}

void mlir::transform::addPadSimplify(RewritePatternSet &patterns,
                                     int64_t maxConstantExpansion,
                                     MLIRContext &context,
                                     PatternBenefit benefit) {
  patterns.insert<PadSimplify>(maxConstantExpansion, &context, benefit);
}

void mlir::transform::addDynamicUpdateSliceConstProp(
    RewritePatternSet &patterns, int64_t maxConstantExpansion,
    MLIRContext &context, PatternBenefit benefit) {
  patterns.insert<DynamicUpdateSliceConstProp>(maxConstantExpansion, &context,
                                               benefit);
}

void mlir::transform::addWhileSimplify(RewritePatternSet &patterns,
                                       bool hoistAll, MLIRContext &context,
                                       PatternBenefit benefit) {
  patterns.insert<WhileSimplify>(hoistAll, &context, benefit);
}

void mlir::transform::addWhileLICM(RewritePatternSet &patterns, bool hoistAll,
                                   MLIRContext &context,
                                   PatternBenefit benefit) {
  patterns.insert<WhileLICM>(hoistAll, &context, benefit);
}

void mlir::transform::addLoopUnswitch(RewritePatternSet &patterns,
                                      int64_t maxOpsOutside,
                                      MLIRContext &context,
                                      PatternBenefit benefit) {
  patterns.insert<LoopUnswitch>(maxOpsOutside, &context, benefit);
}

void mlir::transform::addSliceLICM(RewritePatternSet &patterns,
                                   bool single_user, MLIRContext &context,
                                   PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::SliceOp>>(single_user, &context, benefit);
}

void mlir::transform::addDotGeneralLICM(RewritePatternSet &patterns,
                                        bool single_user, MLIRContext &context,
                                        PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::DotGeneralOp>>(single_user, &context,
                                                 benefit);
}

void mlir::transform::addReverseLICM(RewritePatternSet &patterns,
                                     bool single_user, MLIRContext &context,
                                     PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::ReverseOp>>(single_user, &context, benefit);
}

void mlir::transform::addReduceLICM(RewritePatternSet &patterns,
                                    bool single_user, MLIRContext &context,
                                    PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::ReduceOp>>(single_user, &context, benefit);
}

void mlir::transform::addReduceWindowLICM(RewritePatternSet &patterns,
                                          bool single_user,
                                          MLIRContext &context,
                                          PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::ReduceWindowOp>>(single_user, &context,
                                                   benefit);
}

void mlir::transform::addDUSLICM(RewritePatternSet &patterns, bool single_user,
                                 MLIRContext &context, PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::DynamicUpdateSliceOp>>(single_user, &context,
                                                         benefit);
}

void mlir::transform::addSumToConv(RewritePatternSet &patterns,
                                   bool collapseDims, MLIRContext &context,
                                   PatternBenefit benefit) {
  patterns
      .insert<SumToConv<stablehlo::AddOp>, SumToConv<stablehlo::SubtractOp>>(
          collapseDims, &context, benefit);
}

void mlir::transform::addPadLICM(RewritePatternSet &patterns, bool single_user,
                                 MLIRContext &context, PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::PadOp>>(single_user, &context, benefit);
}

void mlir::transform::addRotateLICM(RewritePatternSet &patterns,
                                    bool single_user, MLIRContext &context,
                                    PatternBenefit benefit) {
  patterns.insert<LICM<enzymexla::RotateOp>>(single_user, &context, benefit);
}

void mlir::transform::addWrapLICM(RewritePatternSet &patterns, bool single_user,
                                  MLIRContext &context,
                                  PatternBenefit benefit) {
  patterns.insert<LICM<enzymexla::WrapOp>>(single_user, &context, benefit);
}

void mlir::transform::addExtendLICM(RewritePatternSet &patterns,
                                    bool single_user, MLIRContext &context,
                                    PatternBenefit benefit) {
  patterns.insert<LICM<enzymexla::ExtendOp>>(single_user, &context, benefit);
}

void mlir::transform::addMultiSliceLICM(RewritePatternSet &patterns,
                                        bool single_user, MLIRContext &context,
                                        PatternBenefit benefit) {
  patterns.insert<LICM<enzymexla::MultiSliceOp>>(single_user, &context,
                                                 benefit);
}

void mlir::transform::addMultiRotateLICM(RewritePatternSet &patterns,
                                         bool single_user, MLIRContext &context,
                                         PatternBenefit benefit) {
  patterns.insert<LICM<enzymexla::MultiRotateOp>>(single_user, &context,
                                                  benefit);
}

void mlir::transform::addElementwiseLICM(RewritePatternSet &patterns,
                                         bool single_user, MLIRContext &context,
                                         PatternBenefit benefit) {
  patterns.insert<LICMElementwise>(single_user, &context, benefit);
}

void mlir::transform::addConcatenateLICM(RewritePatternSet &patterns,
                                         bool single_user, MLIRContext &context,
                                         PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::ConcatenateOp>>(single_user, &context,
                                                  benefit);
}

void mlir::transform::addBroadcastInDimLICM(RewritePatternSet &patterns,
                                            bool single_user,
                                            MLIRContext &context,
                                            PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::BroadcastInDimOp>>(single_user, &context,
                                                     benefit);
}

void mlir::transform::addReshapeLICM(RewritePatternSet &patterns,
                                     bool single_user, MLIRContext &context,
                                     PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::ReshapeOp>>(single_user, &context, benefit);
}

void mlir::transform::addTransposeLICM(RewritePatternSet &patterns,
                                       bool single_user, MLIRContext &context,
                                       PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::TransposeOp>>(single_user, &context, benefit);
}

void mlir::transform::addConvolutionLICM(RewritePatternSet &patterns,
                                         bool single_user, MLIRContext &context,
                                         PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::ConvolutionOp>>(single_user, &context,
                                                  benefit);
}

void mlir::transform::addDynamicSliceLICM(RewritePatternSet &patterns,
                                          bool single_user,
                                          MLIRContext &context,
                                          PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::DynamicSliceOp>>(single_user, &context,
                                                   benefit);
}

void mlir::transform::addScatterLICM(RewritePatternSet &patterns,
                                     bool single_user, MLIRContext &context,
                                     PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::ScatterOp>>(single_user, &context, benefit);
}

void mlir::transform::addGatherLICM(RewritePatternSet &patterns,
                                    bool single_user, MLIRContext &context,
                                    PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::GatherOp>>(single_user, &context, benefit);
}

void mlir::transform::addIotaLICM(RewritePatternSet &patterns, bool single_user,
                                  MLIRContext &context,
                                  PatternBenefit benefit) {
  patterns.insert<LICM<stablehlo::IotaOp>>(single_user, &context, benefit);
}

void mlir::transform::addNoNanAddSubSimplify(RewritePatternSet &patterns,
                                             bool allowOnFloatingPointMath,
                                             MLIRContext &context,
                                             PatternBenefit benefit) {
  patterns.insert<NoNanAddSubSimplify>(allowOnFloatingPointMath, &context,
                                       benefit);
}

void mlir::transform::addNoNanCompareSimplify(RewritePatternSet &patterns,
                                              bool allowOnFloatingPointMath,
                                              MLIRContext &context,
                                              PatternBenefit benefit) {
  patterns.insert<NoNanCompareSimplify>(allowOnFloatingPointMath, &context,
                                        benefit);
}

void mlir::transform::addNoNanSelfSubSimplify(RewritePatternSet &patterns,
                                              bool allowOnFloatingPointMath,
                                              MLIRContext &context,
                                              PatternBenefit benefit) {
  patterns.insert<NoNanSelfSubSimplify>(allowOnFloatingPointMath, &context,
                                        benefit);
}

void mlir::transform::addNoNanMulSimplify(RewritePatternSet &patterns,
                                          bool allowOnFloatingPointMath,
                                          MLIRContext &context,
                                          PatternBenefit benefit) {
  patterns.insert<NoNanMulSimplify>(allowOnFloatingPointMath, &context,
                                    benefit);
}

void mlir::transform::addNoNanDivSimplify(RewritePatternSet &patterns,
                                          bool allowOnFloatingPointMath,
                                          MLIRContext &context,
                                          PatternBenefit benefit) {
  patterns.insert<NoNanDivSimplify>(allowOnFloatingPointMath, &context,
                                    benefit);
}

void mlir::transform::addNoNanZeroBasePowSimplify(RewritePatternSet &patterns,
                                                  bool allowOnFloatingPointMath,
                                                  MLIRContext &context,
                                                  PatternBenefit benefit) {
  patterns.insert<NoNanZeroBasePowSimplify>(allowOnFloatingPointMath, &context,
                                            benefit);
}

void mlir::transform::addSelfSubtractToConvolutionLike(
    RewritePatternSet &patterns, bool allowEmitConvolution,
    MLIRContext &context, PatternBenefit benefit) {
  patterns.insert<SelfElementwiseToConvolutionLike<stablehlo::SubtractOp>>(
      allowEmitConvolution, &context, benefit);
}

void mlir::transform::addSelfAddToConvolutionLike(RewritePatternSet &patterns,
                                                  bool allowEmitConvolution,
                                                  MLIRContext &context,
                                                  PatternBenefit benefit) {
  patterns.insert<SelfElementwiseToConvolutionLike<stablehlo::AddOp>>(
      allowEmitConvolution, &context, benefit);
}

void mlir::transform::addSelfMulToConvolutionLike(RewritePatternSet &patterns,
                                                  bool allowEmitConvolution,
                                                  MLIRContext &context,
                                                  PatternBenefit benefit) {
  patterns.insert<SelfElementwiseToConvolutionLike<stablehlo::MulOp>>(
      allowEmitConvolution, &context, benefit);
}

void mlir::transform::addBroadcastInDimSimplify(RewritePatternSet &patterns,
                                                int64_t maxConstantExpansion,
                                                MLIRContext &context,
                                                PatternBenefit benefit) {
  patterns.insert<BroadcastInDimSimplify>(maxConstantExpansion, &context,
                                          benefit);
}

void mlir::transform::addSelectOpCanon(RewritePatternSet &patterns,
                                       int64_t maxConstantExpansion,
                                       MLIRContext &context,
                                       PatternBenefit benefit) {
  patterns.insert<SelectOpCanon>(maxConstantExpansion, &context, benefit);
}

void mlir::transform::addConcatenateOpCanon(RewritePatternSet &patterns,
                                            int64_t maxConstantExpansion,
                                            MLIRContext &context,
                                            PatternBenefit benefit) {
  patterns.insert<ConcatenateOpCanon>(maxConstantExpansion, &context, benefit);
}

void mlir::transform::addTransposeElementwise(RewritePatternSet &patterns,
                                              bool onlySingleUser,
                                              MLIRContext &context,
                                              PatternBenefit benefit) {
  patterns.insert<TransposeElementwise>(onlySingleUser, &context, benefit);
}

void mlir::transform::addTransposeLikeBroadcastElementwise(
    RewritePatternSet &patterns, bool onlySingleUser, MLIRContext &context,
    PatternBenefit benefit) {
  patterns.insert<TransposeLikeBroadcastElementwise>(onlySingleUser, &context,
                                                     benefit);
}

void mlir::transform::addReshapeElementwise(RewritePatternSet &patterns,
                                            bool onlySingleUser,
                                            MLIRContext &context,
                                            PatternBenefit benefit) {
  patterns.insert<ReshapeElementwise>(onlySingleUser, false, &context, benefit);
}

void mlir::transform::addReshapeElementwiseOnlyFusible(
    RewritePatternSet &patterns, bool onlySingleUser, MLIRContext &context,
    PatternBenefit benefit) {
  patterns.insert<ReshapeElementwise>(onlySingleUser, true, &context, benefit);
}

void mlir::transform::addReshapeSlice(RewritePatternSet &patterns,
                                      bool onlySingleUser, MLIRContext &context,
                                      PatternBenefit benefit) {
  patterns.insert<ReshapeSlice>(onlySingleUser, &context, benefit);
}

void mlir::transform::addReshapeDynamicSlice(RewritePatternSet &patterns,
                                             bool onlySingleUser,
                                             MLIRContext &context,
                                             PatternBenefit benefit) {
  patterns.insert<ReshapeDynamicSlice>(onlySingleUser, &context, benefit);
}

void mlir::transform::addExtendUnaryElementwise(RewritePatternSet &patterns,
                                                bool onlySingleUser,
                                                MLIRContext &context,
                                                PatternBenefit benefit) {
  patterns.insert<ExtendUnaryElementwise>(onlySingleUser, &context, benefit);
}

void mlir::transform::addWrapUnaryElementwise(RewritePatternSet &patterns,
                                              bool onlySingleUser,
                                              MLIRContext &context,
                                              PatternBenefit benefit) {
  patterns.insert<WrapUnaryElementwise>(onlySingleUser, &context, benefit);
}

namespace {

struct EnzymeHLOOptPass
    : public enzyme::impl::EnzymeHLOOptPassBase<EnzymeHLOOptPass> {
  using EnzymeHLOOptPassBase::EnzymeHLOOptPassBase;

  void runOnOperation() override {
    auto context = getOperation()->getContext();

    RewritePatternSet patterns(context);
    mlir::enzyme::populateWithGenerated(patterns);

    patterns.add<SliceExtend>(context);
    patterns.add<SliceRotate>(context);
    patterns.add<SliceWrap>(context);
    patterns.add<ReshapeWrap>(context);
    patterns.add<ReshapeExtend>(context);
    patterns.add<ReshapeRotate>(context);
    patterns.add<TransposeWrap>(context);
    patterns.add<TransposeExtend>(context);
    patterns.add<TransposeRotate>(context);
    patterns.add<SelectPad>(context);
    patterns.add<BitcastConvertCancellation>(context);

    patterns.add<
        AddSimplify, SubSimplify, AndSimplify, MaxSimplify, MinSimplify,
        OrSimplify, XorSimplify, MulSimplify, DivSimplify, RemSimplify,
        PowSimplify, NoopSlice, NoopReverse, SliceReverse, SliceSlice,
        DynamicSliceDynamicSlice, DynamicSliceSlice, SliceDynamicSlice,
        LogSimplify, ShiftRightLogicalSimplify, NegativePadToSlice,
        SliceSimplify, ConvertSimplify, TransposeSimplify, DotGeneralSimplify,
        DotGeneralReshape, DiagonalTensorDotGeneralRewrite,
        DynamicSliceToStatic, DynamicUpdateSliceElim, ReduceToReshape,
        BroadcastToReshape, ReshapeEmptyBroadcast, ReshapeBroadcast,
        BroadcastReshape, ConstPropThroughBarrier, ReplaceNegAddWithSubtract,
        ReplaceSubtractNegWithAdd, SignAbsSimplify, AbsPositiveSimplify,
        SimplifyBoundary<enzymexla::ExtendOp>,
        SimplifyBoundary<enzymexla::WrapOp>,
        SimplifyBoundary<enzymexla::RotateOp>, TransposeReshapeToBroadcast,
        ReshapeTransposeToBroadcast, SelectBroadcastInDim, PowerMultiplyToPower,
        NegMulConstSimplify, NegDivConstSimplify,
        NegatedConstantMulFactoring<stablehlo::AddOp>,
        NegatedConstantMulFactoring<stablehlo::SubtractOp>,
        ReshapeDeletionsBroadcastInDimSimplify,
        ReshapeInsertionsBroadcastInDimSimplify, CompareIotaConstSimplify,
        ConvertIotaSimplify, MinMaxIotaConstSimplify<stablehlo::MaxOp>,
        MinMaxIotaConstSimplify<stablehlo::MinOp>, ClampIotaConstSimplify,
        CompareAbs, CompareMul, CompareConvert, AddSelects,
        CompareNegateConstSimplify, CompareSubtractConstSimplify,
        SelectSimplify, DynamicSliceReshapeDynamicSlice,
        DynamicSliceReshapeSlice, SliceReshapeDynamicSlice, SliceReshapeSlice,
        ExponentialMinusOneFuse, ExponentialMinusOneAddFuse>(
        context, PatternBenefit(65000));

    patterns.add<IotaSimplify, BroadcastInDimSimplify, ConcatConstProp,
                 DynamicUpdateSliceConstProp, PadSimplify, ScatterConstFold,
                 RecognizeFromConstant>(max_constant_expansion, context,
                                        PatternBenefit(65000));

    patterns.add<
        ConvertConcat, DynamicUpdateToConcat, SliceOfDynamicUpdate,
        SliceOfUpdateWithoutCorners, SliceElementwise, SliceReshapeElementwise,
        DynamicSliceElementwise, SlicePad, SliceReshapePad, ReshapeSliceReshape,
        DotReshapeDot, ChloInfConstProp, CHLOLGammaConstProp, GammaConstProp,
        TGammaConstProp, LGammaConstProp, BinomialProgressConstProp, ConcatFuse,
        ConcatToBroadcast, ConcatSlicesToReverse, PadPad, PadReshapePad,
        ConcatPushBinop<stablehlo::AddOp>, ConcatPushBinop<stablehlo::MulOp>,
        ScatterToDynamicUpdateSlice, ReduceConcat, ConcatSlice, ConcatMultiPad,
        ConcatWrap, WidenWrap, WidenExtend, ConcatConcatAxisSwap, SliceConcat,
        SliceIf, SliceReshapeConcat, BinBroadcastSplat<stablehlo::AddOp>,
        BinBroadcastSplat<stablehlo::SubtractOp>,
        BinBroadcastSplat<stablehlo::DivOp>,
        BinBroadcastSplat<stablehlo::MulOp>, RotatePad, ConjReal,
        ConvertMulConvert, ConvertBinopConvert<stablehlo::MinOp>,
        ConvertBinopConvert<stablehlo::MaxOp>, NegateReduceWindowSub>(context);

    // Unary constant propagation patterns
    patterns
        .add<UnaryConstProp<stablehlo::NotOp, stablehlo::notOp>,
             UnaryConstProp<stablehlo::IsFiniteOp, stablehlo::isFiniteOp>,
             UnaryConstProp<stablehlo::LogOp, stablehlo::logOp>,
             UnaryConstProp<stablehlo::Log1pOp, stablehlo::log1pOp>,
             UnaryConstProp<stablehlo::AbsOp, stablehlo::absOp>,
             UnaryConstProp<stablehlo::NegOp, stablehlo::negOp>,
             UnaryConstProp<stablehlo::SqrtOp, stablehlo::sqrtOp>,
             UnaryConstProp<stablehlo::RsqrtOp, stablehlo::rsqrtOp>,
             UnaryConstProp<stablehlo::CosineOp, stablehlo::cosineOp>,
             UnaryConstProp<stablehlo::SineOp, stablehlo::sineOp>,
             UnaryConstProp<stablehlo::ExpOp, stablehlo::exponentialOp>,
             UnaryConstProp<stablehlo::Expm1Op, stablehlo::expm1Op>,
             UnaryConstProp<stablehlo::TanhOp, stablehlo::tanhOp>,
             UnaryConstProp<stablehlo::LogisticOp, stablehlo::logisticOp>,
             UnaryConstProp<chlo::ConjOp, conjOp>,
             UnaryConstProp<stablehlo::CeilOp, stablehlo::ceilOp>,
             UnaryConstProp<stablehlo::CbrtOp, stablehlo::cbrtOp>,
             UnaryConstProp<stablehlo::RealOp, stablehlo::realOp>,
             UnaryConstProp<stablehlo::ImagOp, stablehlo::imagOp>,
             UnaryConstProp<stablehlo::RoundOp, stablehlo::roundOp>,
             UnaryConstProp<stablehlo::RoundNearestEvenOp,
                            stablehlo::roundNearestEvenOp>,
             UnaryConstProp<stablehlo::SignOp, stablehlo::signOp>,
             UnaryConstProp<stablehlo::FloorOp, stablehlo::floorOp>,
             UnaryConstProp<stablehlo::TanOp, stablehlo::tanOp>,
             UnaryConstProp<enzymexla::ReluOp, reluOp>,
             UnaryConstProp<enzymexla::SoftplusOp, softplusOp>, GeluConstProp>(
            context);

    // binary constant propagation patterns
    patterns.add<BinaryConstProp<stablehlo::AddOp, stablehlo::addOp>,
                 BinaryConstProp<stablehlo::AndOp, stablehlo::andOp>,
                 BinaryConstProp<stablehlo::Atan2Op, stablehlo::atan2Op>,
                 BinaryConstProp<stablehlo::ComplexOp, stablehlo::complexOp>,
                 BinaryConstProp<stablehlo::DivOp, stablehlo::divideOp>,
                 BinaryConstProp<stablehlo::MaxOp, stablehlo::maxOp>,
                 BinaryConstProp<stablehlo::MinOp, stablehlo::minOp>,
                 BinaryConstProp<stablehlo::MulOp, stablehlo::multiplyOp>,
                 BinaryConstProp<stablehlo::OrOp, stablehlo::orOp>,
                 BinaryConstProp<stablehlo::PowOp, stablehlo::powerOp>,
                 BinaryConstProp<stablehlo::RemOp, stablehlo::remOp>,
                 BinaryConstProp<stablehlo::SubtractOp, stablehlo::subtractOp>,
                 BinaryConstProp<stablehlo::XorOp, stablehlo::xorOp>>(context);

    patterns.add<GatherConstProp, ClampConstProp>(context);

    patterns.add<ElementwiseAllTransposeOperandsSimplify,
                 TransposeElementwiseTransposeSimplify,
                 AssociativeBinaryOpReordering,
                 CommonAssociativeCommutativeOpReorder>(context);

    patterns.add<BinopPadToConcat<stablehlo::AddOp>,
                 BinopPadToConcat<stablehlo::MulOp>, ConcatPad,
                 PadConcatToConcatPad, SliceSelect, PadReduceWindow,
                 ConvolutionPad>(context);

    if (passses & 512) {
      patterns.add<TransposeDotReorder, DotTranspose, ConvolutionTranspose,
                   TransposeConvolution, EinsumTranspose, TransposeEinsum,
                   ConvertConvertFloat, ConvertConvertInt, ConcatToPad,
                   ConcatAppendingReshape, ReshapeIota, DUSDUS, DUSDUSConcat,
                   DUSConcat, DUSPad, DUSDUSSubsuming, SliceDUSToConcat,
                   ConcatConcatToDUS>(context);
      patterns.add<LICM<stablehlo::DynamicUpdateSliceOp>,
                   LICM<stablehlo::DynamicSliceOp>>(false, context);
    }

    if (passses & 1024)
      patterns.add<FullReduceReshapeOrTranspose>(context);

    if (passses & 1)
      patterns.add<SliceTransposeBase<stablehlo::SliceOp>,
                   SliceTransposeBase<stablehlo::DynamicSliceOp>,
                   SliceReshapeTranspose, SliceBroadcast, SliceReduceWindow>(
          context);

    if (passses & 2)
      patterns.add<ReducePad, BroadcastPad>(context);
    if (passses & 4)
      patterns.add<MulZeroPad, DivZeroPad, ZeroProductReshapePad>(context);
    if (passses & 8)
      patterns.add<BinopConstReshapePad, BinopConstPad<stablehlo::AddOp>,
                   BinopConstPad<stablehlo::SubtractOp>,
                   BinopConstPad<stablehlo::MulOp>,
                   BinopConstPad<stablehlo::DivOp>>(context);

    if (passses & 16)
      patterns.add<
          BinopBinopPadPad<stablehlo::AddOp>, AddPadPadToConcat,
          BinopBinopPadPad<stablehlo::MulOp>, BinopPadPad<stablehlo::AddOp>,
          BinopPadPad<stablehlo::SubtractOp>, BinopPadPad<stablehlo::MulOp>,
          BinopPadPad<stablehlo::DivOp>, BinopPadPad<stablehlo::MinOp>,
          BinopPadPad<stablehlo::MaxOp>>(context);

    if (passses & 32)
      patterns
          .add<UnaryPadPush<stablehlo::ConvertOp>,
               UnaryPadPush<stablehlo::TanhOp>, UnaryPadPush<stablehlo::ExpOp>>(
              context);

    if (passses & 64)
      patterns.add<TransposePad>(context);

    if (passses & 128)
      patterns.add<ReshapePad>(context);

    if (cse) {
      patterns.add<
          CSE<stablehlo::BroadcastInDimOp>, CSE<stablehlo::SliceOp>,
          CSE<stablehlo::TransposeOp>, CSE<stablehlo::ConvertOp>,
          CSE<stablehlo::PadOp>, CSE<stablehlo::DotGeneralOp>,
          CSE<stablehlo::ReshapeOp>, CSE<stablehlo::MulOp>,
          CSE<stablehlo::DivOp>, CSE<stablehlo::AddOp>,
          CSE<stablehlo::SubtractOp>, CSE<stablehlo::PowOp>,
          CSE<stablehlo::MinOp>, CSE<stablehlo::ConcatenateOp>,
          CSE<stablehlo::MaxOp>, CSE<stablehlo::NegOp>, CSE<stablehlo::AbsOp>,
          CSE<enzymexla::RotateOp>, CSE<enzymexla::WrapOp>,
          CSE<enzymexla::ExtendOp>, CSEIota, CSE<stablehlo::CompareOp>,
          CSE<stablehlo::GatherOp>, CSE<stablehlo::ScatterOp>,
          CSE<stablehlo::SelectOp>, CSE<stablehlo::RealOp>, CSE<chlo::ConjOp>,
          CSE<stablehlo::ImagOp>, CSE<stablehlo::BatchNormTrainingOp>,
          CSE<stablehlo::BatchNormInferenceOp>, CSE<stablehlo::BatchNormGradOp>,
          CSE<stablehlo::ExpOp>, CSE<stablehlo::Expm1Op>, CSE<stablehlo::LogOp>,
          CSE<stablehlo::Log1pOp>, CSE<stablehlo::TanhOp>,
          CSE<stablehlo::LogisticOp>, CSE<stablehlo::SqrtOp>,
          CSE<stablehlo::RsqrtOp>, CSE<stablehlo::CbrtOp>,
          CSE<stablehlo::SineOp>, CSE<stablehlo::CosineOp>,
          CSE<stablehlo::Atan2Op>, CSE<stablehlo::FloorOp>,
          CSE<stablehlo::CeilOp>, CSE<stablehlo::RoundOp>,
          CSE<stablehlo::SignOp>, CSE<stablehlo::IsFiniteOp>,
          CSE<stablehlo::ClampOp>, CSE<stablehlo::RemOp>, CSE<stablehlo::AndOp>,
          CSE<stablehlo::OrOp>, CSE<stablehlo::XorOp>,
          CSE<stablehlo::ConvolutionOp>, CSE<stablehlo::FftOp>,
          CSE<stablehlo::DynamicSliceOp>, CSE<stablehlo::DynamicUpdateSliceOp>,
          CSE<stablehlo::ReverseOp>>(context, PatternBenefit(65000));
    }

    if (passses & 256)
      patterns.add<TransposeConvert>(context);

    if (passses & 2048)
      patterns.add<TransposeTranspose>(context);

    if (passses & (2048 * 2))
      patterns.add<BroadcastReduce, SliceDotGeneral, SliceReshapeDotGeneral>(
          context);

    if (passses & (2048 * 4)) {
      patterns.add<PadDotGeneral>(false, context);
      patterns.add<DotReshapePad>(context);
    }
    if (passses & (2048 * 8))
      patterns.add<SliceReshape>(context);

    if (passses & (2048 * 16)) {
      patterns.add<PadDotGeneral>(true, context);
      patterns.add<DotReshapePad>(context);
    }

    if (passses & (2048 * 32)) {
      patterns.add<TransposeWhile, TransposeSliceBase<stablehlo::SliceOp>,
                   TransposeLikeBroadcastSliceBase<stablehlo::SliceOp>,
                   TransposeConcat, TransposeDUS, TransposeIota,
                   TransposeReduceWindow, TransposeReduce, TransposeSelect,
                   TransposeSliceBase<stablehlo::DynamicSliceOp>,
                   TransposeLikeBroadcastSliceBase<stablehlo::DynamicSliceOp>,
                   TransposeReverse, TransposeBatchNormTraining,
                   TransposeBatchNormInference, TransposeBatchNormGrad,
                   TransposeIf, TransposeFFT, TransposeReshape>(context);
      patterns.add<TransposeElementwise>(true, context);
    }

    if (passses & (2048 * 64)) {
      // add reshape push up cases here
      patterns.add<ReshapeSlice, ReshapeDynamicSlice>(true, context);
      patterns.add<ReshapeElementwise>(true, false, context);
      patterns.add<ReshapeOfConcatToConcatOfReshape, ReshapeDUS, ReshapePad,
                   ReshapeReduceWindow, ReshapeSelect, DeleteDimsBroadcast>(
          context);
    }

    if (passses & (2048 * 128)) {
      // Conflicts with ConcatPad
      patterns.add<ConstPadConcatToConcat>(context);
    }

    if (passses & (2048 * 256)) {
      patterns.add<RecognizeRotate, RecognizeWrap, RecognizeExtend,
                   RecognizeUpdateWithoutCorners>(context);
    }

    if (passses & (2048 * 512)) {
      patterns.add<LowerRotate, LowerWrap, LowerExtend,
                   LowerUpdateWithoutCorners, LowerPiecewiseSelect>(context);
    }

    if (passses & (2048 * 1024)) {
      patterns.add<ConcatToOneDimDUS>(context);
    }

    if (passses & (2048 * 2048)) {
      // push reshapes down
      patterns.add<ElementwiseReshapeLike>(context);
    }

    if (licm_passes) {
      // LICM passes
      patterns.add<LICM<stablehlo::ScatterOp>>(false, context);
      patterns.add<LICM<stablehlo::GatherOp>>(false, context);
      patterns.add<LICM<stablehlo::SliceOp>>(false, context);
      patterns.add<LICM<stablehlo::DotGeneralOp>>(false, context);
      patterns.add<LICM<stablehlo::ReverseOp>>(false, context);
      patterns.add<LICM<stablehlo::ReduceOp>>(false, context);
      patterns.add<LICM<stablehlo::ReduceWindowOp>>(false, context);
      patterns.add<LICM<stablehlo::DynamicUpdateSliceOp>>(false, context);
      patterns.add<LICM<stablehlo::PadOp>>(false, context);
      patterns.add<LICMElementwise>(false, context);
      patterns.add<LICM<stablehlo::ConcatenateOp>>(false, context);
      patterns.add<LICM<stablehlo::BroadcastInDimOp>>(false, context);
      patterns.add<LICM<stablehlo::ReshapeOp>>(false, context);
      patterns.add<LICM<stablehlo::TransposeOp>>(false, context);
      patterns.add<LICM<stablehlo::ConvolutionOp>>(false, context);
      patterns.add<WhileLICM>(false, context);
      patterns.add<LICM<stablehlo::IotaOp>>(false, context);
      patterns.add<LICM<enzymexla::ExtendOp>>(false, context);
      patterns.add<LICM<enzymexla::WrapOp>>(false, context);
      patterns.add<LICM<enzymexla::RotateOp>>(false, context);
    }

    if (all_finite)
      patterns.add<AllFiniteIsFinite, AllFiniteIsInf, AllFiniteIsPosInf,
                   AllFiniteIsNegInf>(context);

    patterns.add<NoNanCompareSimplify, NoNanSelfSubSimplify,
                 NoNanAddSubSimplify, NoNanMulSimplify, NoNanDivSimplify>(
        (no_nan || all_finite), context);

    patterns.add<TransposeSymmetricSimplify, TransposePartialSymmetrySimplify>(
        context);
    patterns.add<FactorScalarsInDotGeneral>(context);

    // clang-format off
    // structured tensor optimization patterns
    patterns.add<
        TransposeSyrkToSyrk,
        FuseMulIntoSyrk,
        FuseAddIntoSyrk,
        FuseAddIntoSymm,
        FuseMulIntoSymm,
        SyrkSimplifyOutputUplo
      >(context);
    // clang-format on

    if (structured_tensors_detection) {
      patterns.add<DotGeneralToSymm, DotGeneralToSyrk>(context);
    }

    // clang-format off
    patterns.add<
        WhileRepeatedInductionReduction,
        WhileOpInductionReplacement,
        WhilePadInductionReduction,
        WhileIdempotentDUS,
        ConcatBroadcastToPad,
        SliceBroadcast2,
        BroadcastInDimOpCanon,
        ChainedDynamicBroadcastInDimCanonicalization,
        CompareOpCanon,
        CompareExt,
        ConjComplexNegate,
        NegateImagConj,
        ConvertOpCanon,
        DivideSqrtToMultiplyRsqrt,
        DynamicBroadcastInDimAllDimsNonExpanding,
        DynamicBroadcastInDimOpNotActuallyDynamic,
        DynamicGatherOpIsNotDynamic,
        DynamicReshapeOpCanon,
        EmptyReduceOpCanon,
        GatherOpCanon,
        ScatterOpCanon,
        GetDimensionSizeOpCanon,
        GetTupleElementOpCanon,
        IfRemoveUnused,
        IfInline,
        IfToSelect,
        IfPredPropagation,
        ImagOpCanon,
        MergeConsecutiveReshapes,
        NoopReduceOpCanon,
        RealOpCanon,
        ReorderElementwiseAndShapeOp,
        ReshapeOpCanon,
        SelectCompIotaConstToDUS,
        SelectCompIotaConstSimplify,
        SelectPadToDUS,
        SelectSelectSameCond,
        SelectSelectNegCond,
        SelectToLogical,
        AndPadPad,
        SelectOpUsedWithinIf,
        TransposeBroadcastInDimToBroadcastInDim,
        BroadcastInDimTransposeToBroadcastInDim,
        TransposeIsReshape,
        BroadcastInDimIsReshape,
        ReshuffleAndsCompares,
        WhileDeadResults,
        ZeroExtentTensorCanon,
        CompareSelectSimplify,
        NotSelectSimplify,
        CommonCompareExpressionRewrite,
        ScatterUpdateComputationConstProp,
        ScatterIndicesAreUnique,
        ReduceTransposeSimplify,
        WhileScatterAccumulatorNoAdd,
        BroadcastIotaSimplify,
        BroadcastIota,
        BroadcastCompare,
        NotCompare,
        SliceInternal,
        SquareAbsSimplify,
        DivideDivideSimplify,
        ConcatReshapeSlice,
        ConcatBroadcastSlice,
        ConcatReshapeReduce,
        ConcatElementwise,
        ConcatReshapeElementwise,
        TransposeAllUsersSlice,
        ReduceReduce,
        IfOpLiftCommonOps,
        InvolutionSimplify<stablehlo::NegOp>,
        InvolutionSimplify<stablehlo::NotOp>,
        InvolutionSimplify<chlo::ConjOp>,
        RealConjSimplify,
        RealConvertSimplify,
        ConjComplexSimplify,
        ElementwiseComplexSimplify,
        SplitConvolutionIntoReverseConvolution,
        ScatterMultiplySimplify,
        ScatterDivSimplify,
        ScatterAddSimplify,
        ScatterSubSimplify,
        ScatterOfScatterSimplify,
        UnaryElementwiseScatterSimplify,
        GatherElementwise,
        ElementwiseGather,
        ElementwisePad,
        ConcatenateSubtractToSubtractPad,
        ConcatenateAddToAddPad,
        ConcatenateBroadcastInDim,
        ElementwiseRotate,
        ElementwiseWrap,
        ElementwiseExtend,
        SubtractMultiplyConstToAddMulConst,
        ReduceMulToDotGeneral,
        SplitReduceAddMulToAddDotGeneral,
        DotGeneralDistributiveSimplify<stablehlo::AddOp>,
        DotGeneralDistributiveSimplify<stablehlo::SubtractOp>,
        TrivialReduceWindowToReduceOp,
        NoopReduceWindowOpCanon,
        AddReduceSliceFusion,
        MulReduceSliceFusion,
        MinReduceSliceFusion,
        MaxReduceSliceFusion,
        AndReduceSliceFusion,
        OrReduceSliceFusion,
        XorReduceSliceFusion,
        CaseToIf,
        DUSToDynamicPad,
        DynamicPadToPad,
        RemoveNoOpsFromWhileLoop,
        SplitVariadicScatterOp,
        DynamicSliceSimplify,
        DotGeneralOnlyDiagonalAccess,
        BinaryNegatedOperandsSimplify<stablehlo::MulOp>,
        BinaryNegatedOperandsSimplify<stablehlo::DivOp>,
        DotGeneralBroadcastInDim,
        DotGeneralBroadcastInDimSortDims,
        DotGeneralRemoveBatchDimensions,
        DUSDynamicSliceSimplify,
        WhileDUSDSSimplify,
        WhileDUSDUSSimplify,
        WhileDUS,
        DeleteDimsReduce,
        ReduceDeleteDims,
        DotGeneralInsertDimContractionSimplification,
        FuseReshapeCollapseOrExpandDimsIntoReduce,
        GatherOfScatterSimplify,
        ReduceWindowWrapSimplify,
        SplitComplexScatter,
        SplitComplexGather,
        ReduceMaxMinMulPositiveScalar,
        BinaryOpComplexSimplifyBase<stablehlo::AddOp>,
        BinaryOpComplexSimplifyBase<stablehlo::SubtractOp>
      >(context);

    patterns.add<ReshapeElementwise>(true, true, context);

    patterns.add<
        SelfElementwiseToConvolutionLike<stablehlo::SubtractOp>,
        SelfElementwiseToConvolutionLike<stablehlo::AddOp>,
        SelfElementwiseToConvolutionLike<stablehlo::MulOp>
      >(enable_convert_to_convolution, context);

    patterns.add<SumToReduceWindow<stablehlo::AddOp>,
    SumToReduceWindow<stablehlo::SubtractOp>>(context);

    patterns.add<WhileSimplify>(false, context);

    patterns.add<WhileLICM>(false, context);

    patterns.add<LoopUnswitch>(/*maxOpsOutside=*/10, context);

    // clang-format on
    patterns.add<SelectOpCanon>(max_constant_expansion, context,
                                PatternBenefit(65000));
    patterns.add<ConcatenateOpCanon>(max_constant_expansion, context,
                                     PatternBenefit(65000));

    if (enable_auto_batching_passes) {
      mlir::enzyme::AutoBatchingPassPipelineOptions options{
          true, true, "greedy", true, true, true};
      mlir::enzyme::populateAutoBatchingPassPatterns(patterns, context,
                                                     options);
    }

    GreedyRewriteConfig config;
    config.setMaxIterations(max_iterations);
    config.setUseTopDownTraversal(top_down);
    config.enableFolding();
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns),
                                     config))) {
      signalPassFailure();
    }
  }
};

} // end anonymous namespace
