Basic.cpp

//===----------------------------------------------------------------------===//
//
// 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
// Also available under a BSD-style license. See LICENSE.
//
//===----------------------------------------------------------------------===//

#include "torch-mlir/Conversion/TorchToStablehlo/TorchToStablehlo.h"

#include "PopulatePatterns.h"
#include "Utils.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"
#include "torch-mlir/Conversion/TorchToStablehlo/StablehloLegalizeUtils.h"
#include "torch-mlir/Conversion/Utils/Utils.h"
#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
#include "torch-mlir/Dialect/Torch/IR/TorchTypes.h"
#include "torch-mlir/Dialect/Torch/Utils/TorchUpstream.h"
#include "torch-mlir/Dialect/Torch/Utils/Utils.h"
#include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
#include <cmath>
#include <numeric>
#include <type_traits>

using namespace mlir;
using namespace mlir::torch;
using namespace mlir::torch::Torch;
using namespace mlir::torch::torch_to_stablehlo;

LogicalResult broadcastRanks(PatternRewriter &rewriter, Operation *op,
                             mlir::Value &self, mlir::Value &other) {
  auto selfTy = dyn_cast<RankedTensorType>(self.getType());
  auto otherTy = dyn_cast<RankedTensorType>(other.getType());
  auto selfRank = selfTy.getRank();
  auto otherRank = otherTy.getRank();
  if (selfRank == 0 || otherRank == 0)
    return success();
  if (selfRank > otherRank) {
    auto unsqueezeDims =
        llvm::to_vector<4>(llvm::seq<int64_t>(0, selfRank - otherRank));
    auto unsqueezeInfo =
        hlo::unsqueezeTensor(rewriter, op, other, unsqueezeDims);
    if (failed(unsqueezeInfo))
      return failure();
    other = *unsqueezeInfo;
  } else if (otherRank > selfRank) {
    auto unsqueezeDims =
        llvm::to_vector<4>(llvm::seq<int64_t>(0, otherRank - selfRank));
    auto unsqueezeInfo =
        hlo::unsqueezeTensor(rewriter, op, self, unsqueezeDims);
    if (failed(unsqueezeInfo))
      return failure();
    self = *unsqueezeInfo;
  }
  return success();
}

bool skipMultiplyAlpha(Value alphaValue) {
  double doubleValue;
  auto isFloat = matchPattern(alphaValue, m_TorchConstantFloat(&doubleValue));

  int64_t intValue;
  auto isInt = matchPattern(alphaValue, m_TorchConstantInt(&intValue));

  return ((isFloat && doubleValue == 1.0) || (isInt && intValue == 1.0));
}

static FailureOr<Value> getMaxValueOfDtype(Operation *op, Type elementType,
                                           PatternRewriter &rewriter,
                                           bool allowNonFinites) {
  auto constType = RankedTensorType::get({}, elementType);
  if (isa<mlir::FloatType>(elementType)) {
    auto constAttr = SplatElementsAttr::get(
        constType, getFloatInf(cast<mlir::FloatType>(elementType),
                               /*negative=*/false, allowNonFinites));
    return stablehlo::ConstantOp::create(rewriter, op->getLoc(), constType,
                                         constAttr)
        .getResult();
  }
  if (isa<mlir::IntegerType>(elementType)) {
    auto integerType = cast<mlir::IntegerType>(elementType);
    DenseElementsAttr constAttr;
    if (integerType.isUnsigned()) {
      constAttr = SplatElementsAttr::get(
          constType, APInt::getMaxValue(integerType.getWidth()));
    } else {
      constAttr = SplatElementsAttr::get(
          constType, APInt::getSignedMaxValue(integerType.getWidth()));
    }
    return stablehlo::ConstantOp::create(rewriter, op->getLoc(), constType,
                                         constAttr)
        .getResult();
  }
  return failure();
}

static FailureOr<Value> getMinValueOfDtype(Operation *op, Type elementType,
                                           PatternRewriter &rewriter,
                                           bool allowNonFinites) {
  auto constType = RankedTensorType::get({}, elementType);
  if (isa<mlir::FloatType>(elementType)) {
    auto constAttr = SplatElementsAttr::get(
        constType, getFloatInf(cast<mlir::FloatType>(elementType),
                               /*negative=*/true, allowNonFinites));
    return stablehlo::ConstantOp::create(rewriter, op->getLoc(), constType,
                                         constAttr)
        .getResult();
  }
  if (isa<mlir::IntegerType>(elementType)) {
    auto integerType = cast<mlir::IntegerType>(elementType);
    DenseElementsAttr constAttr;
    if (integerType.isUnsigned()) {
      constAttr = SplatElementsAttr::get(
          constType, APInt::getMinValue(integerType.getWidth()));
    } else {
      constAttr = SplatElementsAttr::get(
          constType, APInt::getSignedMinValue(integerType.getWidth()));
    }
    return stablehlo::ConstantOp::create(rewriter, op->getLoc(), constType,
                                         constAttr)
        .getResult();
  }
  return failure();
}

// These legalizations are for unary ops.
namespace {
template <typename AtenOpT, typename StablehloOpT>
class ConvertAtenUnaryOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value self = adaptor.getSelf();
    auto selfType = cast<TensorType>(self.getType());
    if (!selfType) {
      return op.emitError("only Tensor types supported in StableHLO");
    }
    auto outType = cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));
    self =
        hlo::promoteType(rewriter, op.getLoc(), self, outType.getElementType());
    rewriter.replaceOpWithNewOp<StablehloOpT>(op, outType, self);
    return success();
  }
};
} // namespace

// These legalizations are for unary ops with only for floating point datatypes.
// There is no supported quantized integer mode for these.
namespace {
template <typename AtenOpT, typename StablehloOpT>
class ConvertAtenUnaryFPOnlyOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value self = adaptor.getSelf();
    auto selfTy = cast<TensorType>(self.getType());

    if (!selfTy)
      return op.emitError("only Tensor types supported in StableHLO");

    if (isa<mlir::FloatType>(selfTy.getElementType())) {
      rewriter.replaceOpWithNewOp<StablehloOpT>(
          op,
          OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
              op.getType()),
          self);
      return success();
    } else {
      return op.emitError(
          "only floating-point datatype legalization supported");
    }
  }
};
} // namespace

// These legalizations are for unary ops with promoting to floating point
// datatypes.
namespace {
template <typename AtenOpT, typename StablehloOpT>
class ConvertAtenUnaryPromoteToFPOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value self = adaptor.getSelf();
    auto selfTy = cast<TensorType>(self.getType());
    if (!selfTy)
      return op.emitError("only Tensor types supported in StableHLO");
    auto resultTy = cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));

    if (isa<mlir::FloatType>(resultTy.getElementType())) {
      Value src = hlo::promoteType(rewriter, op.getLoc(), self,
                                   resultTy.getElementType());
      rewriter.replaceOpWithNewOp<StablehloOpT>(op, resultTy, src);
      return success();
    } else {
      return op.emitError(
          "only result to be floating-point datatype legalization supported");
    }
  }
};
} // namespace

// aten.ones & aten.zeros
// Ref: Error checking based on the Torch to TOSA lowering
namespace {
template <typename AtenOpT, int fillVal>
class ConvertAtenConstPatternOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    auto outType = dyn_cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));

    if (!outType)
      return op.emitError("only Tensor types supported in StableHLO");

    Type outElemTy = outType.getElementType();
    if (!outElemTy.isIntOrFloat())
      return op.emitError(
          "only floating-point or integer datatype legalization supported");

    SmallVector<int64_t> shape;
    if (!matchPattern(op.getSize(), m_TorchListOfConstantInts(shape))) {
      return op.emitError("shape must be a list of Scalar constants");
    }

    int64_t size = 1;
    for (auto s : shape)
      size *= s;

    SmallVector<int32_t> values(size, fillVal);
    auto constOp =
        hlo::getConstTensor<int32_t>(rewriter, op, values, shape).value();

    rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(op, outType, constOp);
    return success();
  }
};

} // namespace

namespace {
// Casts a tensor of exactly one element to an elemental type.
// Many codes borrowed from
// `lib/Conversion/TorchToLinalg/TensorScalarInterop.cpp`
template <typename AtenOpT>
class ConvertAtenTensorToScalarLikeOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto inputType = dyn_cast<RankedTensorType>(adaptor.getA().getType());
    if (!inputType)
      op.emitError("only Tensor types supported in StableHLO");

    Location loc = op.getLoc();
    Value input = adaptor.getA();
    SmallVector<Value> inputSizes = getTensorSizes(rewriter, loc, input);
    int64_t inputRank = inputSizes.size();
    Type inputDtype = cast<BaseTensorType>(op.getA().getType()).getDtype();

    Value constantOne =
        arith::ConstantOp::create(rewriter, loc, rewriter.getI64IntegerAttr(1));
    for (int64_t i = 0; i < inputRank; i++)
      checkDimEqualHelper(rewriter, loc, inputSizes[i], constantOne);

    // handle unsigned interger
    if (inputType.getElementType().isUnsignedInteger()) {
      input = stablehlo::ConvertOp::create(
          rewriter, loc, input,
          rewriter.getIntegerType(
              inputType.getElementType().getIntOrFloatBitWidth()));
    }

    Value constantZero =
        arith::ConstantOp::create(rewriter, loc, rewriter.getIndexAttr(0));
    SmallVector<Value> indices(inputRank, constantZero);
    Value result = tensor::ExtractOp::create(rewriter, loc, input, indices);
    Type resultType =
        this->getTypeConverter()->convertType(op->getResult(0).getType());
    rewriter.replaceOp(
        op,
        convertScalarToDtype(rewriter, loc, result, resultType, inputDtype,
                             /*srcOriginalDtype=*/inputType.getElementType()));
    return success();
  }
};
} // namespace

// The binary broadcast patterns
namespace {
template <typename AtenOpT, typename ChloOpT>
class ConvertAtenBinaryBroadcastOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value lhs = adaptor.getSelf();
    auto lhsTy = cast<TensorType>(lhs.getType());
    Value rhs = adaptor.getOther();
    auto rhsTy = cast<TensorType>(rhs.getType());

    if (!lhsTy || !rhsTy)
      return op.emitError("only Tensor types supported");

    auto outTy = cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));

    lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outTy.getElementType());
    rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outTy.getElementType());

    rewriter.replaceOpWithNewOp<ChloOpT>(op, outTy, lhs, rhs,
                                         /*broadcast_attr*/ nullptr);
    return success();
  }
};
} // namespace

// These binary op legalizations are specific to add/sub which have an
// alpha multiplier.
namespace {
template <typename AtenOpT, typename ChloOpT>
class ConvertAtenAddSubOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value lhs = adaptor.getSelf();
    RankedTensorType lhsType = dyn_cast<RankedTensorType>(lhs.getType());
    Value rhs = adaptor.getOther();
    RankedTensorType rhsType = dyn_cast<RankedTensorType>(rhs.getType());

    if (!lhsType)
      return op.emitError("only Tensor types supported in StableHLO");

    TensorType outType = cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));

    Type outElemTy = outType.getElementType();
    if (!outElemTy.isIntOrFloat()) {
      return op.emitError(
          "only floating-point or integer datatype legalization supported");
    }

    if (!rhsType) {
      rhs = hlo::scalarToStablehloTensor(rewriter, op, adaptor.getOther(),
                                         outElemTy);
      if (isa<AtenRsubScalarOp>(op)) {
        std::swap(lhs, rhs);
      }
    }

    lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outElemTy);
    rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outElemTy);

    if (!skipMultiplyAlpha(op.getAlpha())) {
      Value alpha = hlo::scalarToStablehloTensor(rewriter, op,
                                                 adaptor.getAlpha(), outElemTy);
      DenseI64ArrayAttr bcastDimensions;
      rhs = chlo::BroadcastMulOp::create(rewriter, op->getLoc(), rhs, alpha,
                                         bcastDimensions);
    }

    DenseI64ArrayAttr bcastDimensions;
    rewriter.replaceOpWithNewOp<ChloOpT>(op, outType, lhs, rhs,
                                         bcastDimensions);
    return success();
  }
};
} // namespace

// Binary op legalizations for Mul/Div variants.
namespace {
template <typename AtenOpT, typename ChloOpT>
class ConvertAtenMulDivOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value lhs = adaptor.getSelf();
    auto lhsType = dyn_cast<TensorType>(lhs.getType());
    Value rhs = adaptor.getOther();
    TensorType rhsType = dyn_cast<TensorType>(rhs.getType());

    if (!lhsType)
      return op.emitError("only Tensor types supported in StableHLO");

    auto outType = cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));

    Type outElemTy = outType.getElementType();
    if (!outElemTy.isIntOrFloat()) {
      return op.emitError(
          "only floating-point or integer datatype legalization supported");
    }

    if constexpr (std::is_same<AtenOpT, AtenSquareOp>()) {
      rhs = lhs;
    } else {
      if (!rhsType) {
        rhs = hlo::scalarToStablehloTensor(rewriter, op, adaptor.getOther(),
                                           outElemTy);
      }
    }
    DenseI64ArrayAttr bcastDimensions;
    lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outElemTy);
    rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outElemTy);
    auto loc = op.getLoc();
    Value result =
        ChloOpT::create(rewriter, loc, outType, lhs, rhs, bcastDimensions);

    if constexpr (!std::is_same<AtenDivTensorModeOp, AtenOpT>() &&
                  !std::is_same<AtenDivScalarModeOp, AtenOpT>()) {
      rewriter.replaceOp(op, result);
      return success();
    }

    auto tensorOp = dyn_cast<AtenDivTensorModeOp>(op.getOperation());
    auto opRoundingMode =
        tensorOp
            ? tensorOp.getRoundingMode()
            : cast<AtenDivScalarModeOp>(op.getOperation()).getRoundingMode();

    std::string roundingMode;
    if (!matchPattern(opRoundingMode, m_TorchConstantStr(roundingMode))) {
      return rewriter.notifyMatchFailure(
          op, "only support constant str rounding mode");
    }

    // if trunc and int, do nothing
    if (roundingMode == "trunc" && isa<mlir::FloatType>(outElemTy)) {
      // "trunc" - rounds the results of the division towards zero. Equivalent
      // to C-style integer division.
      auto sign = stablehlo::SignOp::create(rewriter, loc, result);
      auto abs = stablehlo::AbsOp::create(rewriter, loc, result);
      auto floor = stablehlo::FloorOp::create(rewriter, loc, abs);
      result = stablehlo::MulOp::create(rewriter, loc, sign, floor).getResult();
    }
    if (roundingMode == "floor") {
      // "floor" - rounds the results of the division down. Equivalent to
      // floor division in Python (the // operator)
      if (isa<mlir::FloatType>(outElemTy))
        result = stablehlo::FloorOp::create(rewriter, loc, result).getResult();
      else if (!outElemTy.isUnsignedInteger()) {
        Type defaultIntToFloatType = rewriter.getF64Type();
        lhs =
            hlo::promoteType(rewriter, op.getLoc(), lhs, defaultIntToFloatType);
        rhs =
            hlo::promoteType(rewriter, op.getLoc(), rhs, defaultIntToFloatType);
        result = ChloOpT::create(
            rewriter, loc,
            outType.cloneWith(outType.getShape(), defaultIntToFloatType), lhs,
            rhs, bcastDimensions);
        result = stablehlo::FloorOp::create(rewriter, loc, result).getResult();
        result = hlo::promoteType(rewriter, op.getLoc(), result,
                                  outType.getElementType());
      }
    }
    rewriter.replaceOp(op, result);
    return success();
  }
};
} // namespace

// Binary op legalizations for comparator ops.
namespace {
template <typename AtenOpT>
class ConvertAtenCompareOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;

  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value lhs = adaptor.getSelf();
    Value rhs = adaptor.getOther();
    RankedTensorType lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
    RankedTensorType rhsTy = dyn_cast<RankedTensorType>(rhs.getType());

    if (!lhsTy) {
      return op.emitError("only Tensor types supported in StableHLO");
    }
    bool isRhsScalar = false;
    if (!rhsTy) {
      rhs = hlo::scalarToStablehloTensor(rewriter, op, adaptor.getOther(),
                                         rhs.getType());
      rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
      isRhsScalar = true;
    }

    auto outType = cast<RankedTensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));

    Type lhsElemTy = lhsTy.getElementType();
    Type rhsElemTy = rhsTy.getElementType();
    if (!lhsElemTy.isIntOrFloat() || !rhsElemTy.isIntOrFloat()) {
      return op.emitError(
          "only floating-point or integer datatype legalization supported");
    }

    if (isa<mlir::IntegerType>(lhsElemTy) && isa<mlir::FloatType>(rhsElemTy)) {
      // torch.lt(x_int, 1.1) use fp32 as compute type
      // torch.lt(x_int, y_float) use y's float type as compute type
      Type promoteTo = isRhsScalar ? rewriter.getF32Type() : rhsElemTy;
      lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, promoteTo);
      rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, promoteTo);
    } else if (isa<mlir::FloatType>(lhsElemTy) &&
               isa<mlir::IntegerType>(rhsElemTy)) {
      // always use lhs's float type as compute type
      rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, lhsElemTy);
    } else {
      if (isRhsScalar) {
        // torch.lt(x_float, 1.1) use x's float type as compute type
        // torch.lt(x_int, 1) use x's int type as compute type
        rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, lhsElemTy);
      } else {
        // torch.lt(x_float, y_float) use higher bitwidth as compute type
        Type promoteTo = lhsElemTy.getIntOrFloatBitWidth() >
                                 rhsElemTy.getIntOrFloatBitWidth()
                             ? lhsElemTy
                             : rhsElemTy;
        lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, promoteTo);
        rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, promoteTo);
      }
    }
    lhsElemTy = dyn_cast<RankedTensorType>(lhs.getType()).getElementType();

    chlo::ComparisonTypeAttr compareTypeAttr;
    chlo::ComparisonDirectionAttr compareDirectionAttr;

    if (isa<mlir::FloatType>(lhsElemTy)) {
      compareTypeAttr = chlo::ComparisonTypeAttr::get(
          op->getContext(), chlo::ComparisonType::FLOAT);
    } else if (isa<mlir::IntegerType>(lhsElemTy)) {
      compareTypeAttr = chlo::ComparisonTypeAttr::get(
          op->getContext(), chlo::ComparisonType::SIGNED);
    }

    if constexpr (std::is_same<AtenOpT, AtenLtTensorOp>() ||
                  std::is_same<AtenOpT, AtenLtScalarOp>()) {
      compareDirectionAttr = chlo::ComparisonDirectionAttr::get(
          op->getContext(), chlo::ComparisonDirection::LT);
    } else if constexpr (std::is_same<AtenOpT, AtenGtTensorOp>() ||
                         std::is_same<AtenOpT, AtenGtScalarOp>()) {
      compareDirectionAttr = chlo::ComparisonDirectionAttr::get(
          op->getContext(), chlo::ComparisonDirection::GT);
    } else if constexpr (std::is_same<AtenOpT, AtenGeTensorOp>() ||
                         std::is_same<AtenOpT, AtenGeScalarOp>()) {
      compareDirectionAttr = chlo::ComparisonDirectionAttr::get(
          op->getContext(), chlo::ComparisonDirection::GE);
    } else if constexpr (std::is_same<AtenOpT, AtenEqTensorOp>() ||
                         std::is_same<AtenOpT, AtenEqScalarOp>()) {
      compareDirectionAttr = chlo::ComparisonDirectionAttr::get(
          op->getContext(), chlo::ComparisonDirection::EQ);
    } else if constexpr (std::is_same<AtenOpT, AtenNeTensorOp>() ||
                         std::is_same<AtenOpT, AtenNeScalarOp>()) {
      compareDirectionAttr = chlo::ComparisonDirectionAttr::get(
          op->getContext(), chlo::ComparisonDirection::NE);
    } else if constexpr (std::is_same<AtenOpT, AtenLtTensorOp>() ||
                         std::is_same<AtenOpT, AtenLtScalarOp>()) {
      compareDirectionAttr = chlo::ComparisonDirectionAttr::get(
          op->getContext(), chlo::ComparisonDirection::LT);
    } else if constexpr (std::is_same<AtenOpT, AtenLeTensorOp>() ||
                         std::is_same<AtenOpT, AtenLeScalarOp>()) {
      compareDirectionAttr = chlo::ComparisonDirectionAttr::get(
          op->getContext(), chlo::ComparisonDirection::LE);
    } else {
      return op.emitError("operator haven't been supported");
    }
    DenseI64ArrayAttr bcastDimensions;
    rewriter.replaceOpWithNewOp<chlo::BroadcastCompareOp>(
        op, outType, lhs, rhs, bcastDimensions, compareDirectionAttr,
        compareTypeAttr);
    return success();
  }
};
} // namespace

// Binary op legalizations for Logical And/Or/Xor.
namespace {
template <typename AtenOpT, typename ChloOpT>
class ConvertAtenLogicalBinaryOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;

  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value lhs = adaptor.getSelf();
    Value rhs = adaptor.getOther();

    RankedTensorType lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
    RankedTensorType rhsTy = dyn_cast<RankedTensorType>(rhs.getType());

    if (!lhsTy)
      return op.emitError("lhs must be a ranked tensor type");

    TensorType outType = cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));
    Type outElemTy = outType.getElementType();
    lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outElemTy);
    if (!rhsTy) {
      rhs = hlo::scalarToStablehloTensor(rewriter, op, rhs, outElemTy);
    }
    rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outElemTy);

    DenseI64ArrayAttr bcastDimensions;
    rewriter.replaceOpWithNewOp<ChloOpT>(op, outType, lhs, rhs,
                                         bcastDimensions);
    return success();
  }
};
} // namespace

// AtenTransposeIntOp
namespace {
class ConvertAtenTransposeIntOp
    : public OpConversionPattern<AtenTransposeIntOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(AtenTransposeIntOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value self = adaptor.getSelf();
    int64_t dim0;
    if (!matchPattern(op.getDim0(), m_TorchConstantInt(&dim0))) {
      return rewriter.notifyMatchFailure(op, "dim0 must be constant");
    }
    int64_t dim1;
    if (!matchPattern(op.getDim1(), m_TorchConstantInt(&dim1))) {
      return rewriter.notifyMatchFailure(op, "dim1 must be constant");
    }

    auto inType = cast<RankedTensorType>(self.getType());
    auto inputRank = inType.getRank();
    auto outType = cast<RankedTensorType>(
        getTypeConverter()->convertType(op->getResult(0).getType()));

    dim0 = toPositiveDim(dim0, inputRank);
    if (!isValidDim(dim0, inputRank)) {
      return rewriter.notifyMatchFailure(op, "dim0 out of range");
    }
    dim1 = toPositiveDim(dim1, inputRank);
    if (!isValidDim(dim1, inputRank)) {
      return rewriter.notifyMatchFailure(op, "dim1 out of range");
    }

    SmallVector<int64_t> permValues(inputRank);
    std::iota(std::begin(permValues), std::end(permValues), 0);
    std::swap(permValues[dim0], permValues[dim1]);
    rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(op, outType, self,
                                                        permValues);
    return success();
  }
};
} // namespace

// AtenToDtypeOp
template <>
LogicalResult ConvertAtenOp<AtenToDtypeOp>::matchAndRewrite(
    AtenToDtypeOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value self = adaptor.getSelf();
  auto outType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(op, outType, self);
  return success();
}

template <>
LogicalResult ConvertAtenOp<AtenSizeIntOp>::matchAndRewrite(
    AtenSizeIntOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  // Not a tensor type.
  auto selfType = dyn_cast<TensorType>(adaptor.getSelf().getType());
  if (!selfType)
    return op.emitError("only tensor types are currently supported");

  Value dim;
  int64_t dimInt;
  if (matchPattern(op.getDim(), m_TorchConstantInt(&dimInt))) {
    dimInt = toPositiveDim(dimInt, selfType.getRank());
    if (!isValidDim(dimInt, selfType.getRank()))
      return rewriter.notifyMatchFailure(op, "dim is statically invalid");
    dim = arith::ConstantIndexOp::create(rewriter, op.getLoc(), dimInt);
  } else {
    Value inputRank = arith::ConstantOp::create(
        rewriter, op.getLoc(), rewriter.getI64IntegerAttr(selfType.getRank()));
    dim = toPositiveDimDynamic(rewriter, op.getLoc(), adaptor.getDim(),
                               inputRank);
    dim = arith::IndexCastOp::create(rewriter, op.getLoc(),
                                     rewriter.getIndexType(), dim);
  }

  auto dimSize = tensor::DimOp::create(
      rewriter, op.getLoc(), rewriter.getIndexType(), adaptor.getSelf(), dim);

  rewriter.replaceOpWithNewOp<arith::IndexCastOp>(
      op, getTypeConverter()->convertType(op.getType()), dimSize);

  return success();
}

template <>
LogicalResult ConvertAtenOp<AtenWhereSelfOp>::matchAndRewrite(
    AtenWhereSelfOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value self = adaptor.getSelf();
  Value cond = adaptor.getCondition();
  Value other = adaptor.getOther();

  auto outType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  // promote self and other types
  self =
      hlo::promoteType(rewriter, op.getLoc(), self, outType.getElementType());
  other =
      hlo::promoteType(rewriter, op.getLoc(), other, outType.getElementType());

  if (failed(broadcastRanks(rewriter, op, self, cond)))
    return op.emitError("failed broadcast self and condition ranks");

  if (failed(broadcastRanks(rewriter, op, other, cond)))
    return op.emitError("failed broadcast other and condition ranks");

  rewriter.replaceOpWithNewOp<chlo::BroadcastSelectOp>(
      op, getTypeConverter()->convertType(op.getType()),
      ArrayRef<Value>{cond, self, other});
  return success();
}

// AtenBroadcastToOp
template <>
LogicalResult ConvertAtenOp<AtenBroadcastToOp>::matchAndRewrite(
    AtenBroadcastToOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value self = adaptor.getSelf();
  auto selfTy = cast<RankedTensorType>(self.getType());
  auto outType = cast<RankedTensorType>(
      getTypeConverter()->convertType(op->getResult(0).getType()));

  if (options.enableStaticShape && selfTy.hasStaticShape()) {
    Value bcastOp =
        hlo::promoteAndBroadcast(rewriter, self, outType, std::nullopt);
    rewriter.replaceOp(op, bcastOp);
    return success();
  }

  SmallVector<Value> shape;
  if (!(getListConstructElements(adaptor.getSize(), shape))) {
    return op->emitError("desired shape must be a list of scalar");
  }
  SmallVector<Value> bcastShapeVec;
  int64_t totalRank = shape.size();
  int64_t selfRank = selfTy.getRank();
  int64_t leadingRank = totalRank - selfRank;

  for (int64_t i = 0; i < totalRank; ++i) {
    Value dValue = shape[i];
    Value newD;
    int64_t dInt;
    if (i >= leadingRank && matchPattern(dValue, m_TorchConstantInt(&dInt)) &&
        dInt == -1) {
      newD = mlir::tensor::DimOp::create(rewriter, op->getLoc(), self,
                                         i - leadingRank);
    } else {
      dValue = torch::TorchConversion::ToI64Op::create(rewriter, op->getLoc(),
                                                       dValue);
      newD = mlir::arith::IndexCastOp::create(rewriter, op->getLoc(),
                                              rewriter.getIndexType(), dValue);
    }
    bcastShapeVec.push_back(newD);
  }

  if (options.dimSizeIndexBits == 32) {
    for (auto &dsize : bcastShapeVec) {
      auto dsizeI64 = mlir::arith::IndexCastOp::create(
          rewriter, op->getLoc(), rewriter.getI64Type(), dsize);
      dsize = arith::TruncIOp::create(rewriter, op->getLoc(),
                                      rewriter.getI32Type(), dsizeI64);
    }
  }

  if (bcastShapeVec.size() == 0) {
    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, outType, self);
  } else {
    Value bcastShapeTensor = mlir::tensor::FromElementsOp::create(
        rewriter, op->getLoc(), ValueRange{bcastShapeVec});
    auto dimensionNumbers =
        llvm::to_vector<4>(llvm::seq<int64_t>(leadingRank, totalRank));
    rewriter.replaceOpWithNewOp<stablehlo::DynamicBroadcastInDimOp>(
        op, outType, self, bcastShapeTensor,
        rewriter.getDenseI64ArrayAttr(dimensionNumbers));
  }
  return success();
}

// AtenPermuteOp
template <>
LogicalResult ConvertAtenOp<AtenPermuteOp>::matchAndRewrite(
    AtenPermuteOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value self = adaptor.getSelf();
  // Not a ranked tensor type
  auto inType = dyn_cast<RankedTensorType>(self.getType());
  auto outType = cast<RankedTensorType>(
      getTypeConverter()->convertType(op->getResult(0).getType()));
  if (!inType)
    return op.emitError("only ranked tensor types with static shapes are "
                        "currently supported");

  SmallVector<int64_t> permValues;
  if (!matchPattern(adaptor.getDims(), m_TorchListOfConstantInts(permValues)))
    return rewriter.notifyMatchFailure(
        op, "only constant dimensions are currently supported");

  int64_t inRank = inType.getRank();
  for (auto &d : permValues) {
    d = toPositiveDim(d, inRank);
    if (!isValidDim(d, inRank))
      return op.emitError("not all dims are valid");
  }

  rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(op, outType, self,
                                                      permValues);
  return success();
}

// ValueTensorLiteralOp
template <>
LogicalResult ConvertAtenOp<ValueTensorLiteralOp>::matchAndRewrite(
    ValueTensorLiteralOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  RankedTensorType resultType = cast<RankedTensorType>(
      getTypeConverter()->convertType(op->getResult(0).getType()));

  // Tensors with integer types need to be converted to signless integer
  // element type. All tensors with element types other than integer can reuse
  // existing elements attribute.
  if (auto elements = dyn_cast<DenseIntElementsAttr>(op.getValueAttr())) {
    Type builtinTensorElemTy = resultType.getElementType();
    unsigned bitWidth = builtinTensorElemTy.getIntOrFloatBitWidth();
    bool isUnsigned =
        cast<IntegerType>(builtinTensorElemTy).isUnsignedInteger();

    DenseElementsAttr valueAttr =
        elements.mapValues(builtinTensorElemTy, [&](const APInt &v) {
          APInt intValue =
              isUnsigned ? v.zextOrTrunc(bitWidth) : v.sextOrTrunc(bitWidth);
          return intValue;
        });
    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, resultType,
                                                       valueAttr);
    return success();
  }

  rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, resultType,
                                                     adaptor.getValue());
  return success();
}

// AtenTensorIntOp
template <>
LogicalResult ConvertAtenOp<AtenTensorIntOp>::matchAndRewrite(
    AtenTensorIntOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  RankedTensorType resultType = cast<RankedTensorType>(
      getTypeConverter()->convertType(op->getResult(0).getType()));
  Type outElementType = resultType.getElementType();
  Value innerValue = adaptor.getT();
  Value stablehloTensor =
      hlo::scalarToStablehloTensor(rewriter, op, innerValue, outElementType);
  rewriter.replaceOp(op, stablehloTensor);
  return success();
}

// AtenReciprocalOp
// Reciprocal(x) = Div(1, x)
template <>
LogicalResult ConvertAtenOp<AtenReciprocalOp>::matchAndRewrite(
    AtenReciprocalOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getSelf();
  auto inputTy = cast<RankedTensorType>(input.getType());
  auto outTy =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  if (!isa<mlir::FloatType>(inputTy.getElementType())) {
    return op.emitError("only floating-point datatype legalization supported "
                        "for AtenReciprocalOp");
  }

  Value oneTensor =
      hlo::getConstantLike<int64_t>(rewriter, op->getLoc(), 1, input);
  rewriter.replaceOpWithNewOp<stablehlo::DivOp>(op, outTy, oneTensor, input);
  return success();
}

namespace {
template <typename AtenOpT>
class ConvertAtenPowOp : public OpConversionPattern<AtenOpT> {
public:
  using OpConversionPattern<AtenOpT>::OpConversionPattern;
  using OpAdaptor = typename AtenOpT::Adaptor;
  LogicalResult
  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto outType = cast<TensorType>(
        OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
            op.getType()));

    Type outElemTy = outType.getElementType();
    if (!outElemTy.isIntOrFloat()) {
      return op.emitError(
          "only floating-point or integer datatype legalization supported");
    }

    Value lhs = adaptor.getSelf();
    auto lhsType = dyn_cast<TensorType>(lhs.getType());
    Value rhs = adaptor.getExponent();
    auto rhsType = dyn_cast<TensorType>(rhs.getType());

    if (!lhsType && !rhsType) {
      return op.emitError("only Tensor types supported in StableHLO");
    }
    if (!lhsType) {
      lhs = hlo::scalarToStablehloTensor(rewriter, op, lhs, outElemTy);
    }
    if (!rhsType) {
      rhs = hlo::scalarToStablehloTensor(rewriter, op, rhs, outElemTy);
    }

    lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outElemTy);
    rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outElemTy);
    DenseI64ArrayAttr bcastDimensions;
    rewriter.replaceOpWithNewOp<chlo::BroadcastPowOp>(op, outType, lhs, rhs,
                                                      bcastDimensions);
    return success();
  }
};
} // namespace

// PrimNumToTensorScalarOp
template <>
LogicalResult ConvertAtenOp<PrimNumToTensorScalarOp>::matchAndRewrite(
    PrimNumToTensorScalarOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  RankedTensorType outputType = cast<RankedTensorType>(
      getTypeConverter()->convertType(op->getResult(0).getType()));
  auto outputElemType = outputType.getElementType();
  Value stablehloTensor = hlo::scalarToStablehloTensor(
      rewriter, op, adaptor.getA(), outputElemType);
  rewriter.replaceOp(op, stablehloTensor);
  return success();
}

// AtenScalarImplicitOp
template <>
LogicalResult ConvertAtenOp<AtenScalarImplicitOp>::matchAndRewrite(
    AtenScalarImplicitOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op.getLoc();
  Type inputDtype = cast<BaseTensorType>(op.getA().getType()).getDtype();
  Type resultType =
      this->getTypeConverter()->convertType(op->getResult(0).getType());
  auto result = tensor::ExtractOp::create(rewriter, loc, adaptor.getA());

  rewriter.replaceOp(
      op, convertScalarToDtype(rewriter, loc, result, resultType, inputDtype));
  return success();
}

// AtenContiguousOp
// Ref: TosaToTosa.cpp for implementation details
template <>
LogicalResult ConvertAtenOp<AtenContiguousOp>::matchAndRewrite(
    AtenContiguousOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {

  // Not a tensor type.
  auto selfType = dyn_cast<TensorType>(adaptor.getSelf().getType());
  if (!selfType)
    return op.emitError("only tensor types are currently supported");

  // FIXME: memory_format is not handled.

  rewriter.replaceOp(op, adaptor.getSelf());

  return success();
}

// AtenReluOp
// Relu(x) = Max(0, x)
template <>
LogicalResult ConvertAtenOp<AtenReluOp>::matchAndRewrite(
    AtenReluOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value lhs = adaptor.getSelf();
  auto lhsTy = cast<RankedTensorType>(lhs.getType());
  auto lhsElemTy = lhsTy.getElementType();

  if (!isa<mlir::FloatType>(lhsElemTy)) {
    return op->emitError("only float tensor in relu op is supported");
  }

  Value zeroTensor =
      hlo::getConstantLike<int64_t>(rewriter, op->getLoc(), 0, lhs);
  rewriter.replaceOpWithNewOp<stablehlo::MaxOp>(op, lhs, zeroTensor);
  return success();
}

// Convert a Aten::GELU to HLO
// Gelu(x, "none") = x * 0.5 * (1 + erf(x/(sqrt(2))))
// Gelu(x, "tanh") = x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
template <>
LogicalResult ConvertAtenOp<AtenGeluOp>::matchAndRewrite(
    AtenGeluOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op.getLoc();
  Value input = adaptor.getSelf();
  auto inputTy = dyn_cast<RankedTensorType>(input.getType());
  if (!inputTy) {
    return op.emitError("only ranked tensor type is supported.");
  }

  std::string approximate;
  if (!matchPattern(op.getApproximate(), m_TorchConstantStr(approximate))) {
    return op.emitError("approximate must be constant string");
  }
  if (approximate != "none" && approximate != "tanh") {
    return op.emitError("unsupported approximate: ") << approximate;
  }

  Value one = hlo::getConstantLike(rewriter, loc, 1.0, input);
  Value two = hlo::getConstantLike(rewriter, loc, 2.0, input);
  Value three = hlo::getConstantLike(rewriter, loc, 3.0, input);
  Value half = hlo::getConstantLike(rewriter, loc, 0.5, input);
  // 2/pi
  Value twoDivPi = hlo::getConstantLike(rewriter, loc, M_2_PI, input);
  Value t = hlo::getConstantLike(rewriter, loc, 0.044715, input);

  // x * 0.5
  auto inputMulHalf = stablehlo::MulOp::create(rewriter, loc, input, half);
  if (approximate == "none") {
    auto rsqrtTwo = stablehlo::RsqrtOp::create(rewriter, loc, two);
    auto erfElement = stablehlo::MulOp::create(rewriter, loc, input, rsqrtTwo);
    auto erf = chlo::ErfOp::create(rewriter, loc, erfElement);
    auto erfAdd = stablehlo::AddOp::create(rewriter, loc, erf, one);
    rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, erfAdd, inputMulHalf);
    return success();
  } else {
    auto sqrtTwoPi = stablehlo::SqrtOp::create(rewriter, loc, twoDivPi);
    // x^3
    auto powThree = stablehlo::PowOp::create(rewriter, loc, input, three);
    // x + 0.044715 * x^3
    auto add = stablehlo::AddOp::create(
        rewriter, loc, input,
        stablehlo::MulOp::create(rewriter, loc, t, powThree));
    auto tanh = stablehlo::TanhOp::create(
        rewriter, loc, stablehlo::MulOp::create(rewriter, loc, sqrtTwoPi, add));
    auto tanhAdd = stablehlo::AddOp::create(rewriter, loc, tanh, one);
    rewriter.replaceOpWithNewOp<stablehlo::MulOp>(op, tanhAdd, inputMulHalf);
    return success();
  }
}

// AtenLog2Op
template <>
LogicalResult ConvertAtenOp<AtenLog2Op>::matchAndRewrite(
    AtenLog2Op op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getSelf();
  auto inputTy = dyn_cast<RankedTensorType>(input.getType());
  if (!inputTy) {
    return op.emitError("only ranked tensor type is supported.");
  }
  auto outTy = cast<TensorType>(getTypeConverter()->convertType(op.getType()));
  input =
      hlo::promoteType(rewriter, op.getLoc(), input, outTy.getElementType());

  auto two = hlo::getConstantLike(rewriter, op.getLoc(), 2.0, input);
  auto log2Op = stablehlo::LogOp::create(rewriter, op.getLoc(), two);
  auto logInputOp = stablehlo::LogOp::create(rewriter, op.getLoc(), input);

  rewriter.replaceOpWithNewOp<stablehlo::DivOp>(op, outTy, logInputOp, log2Op);
  return success();
}

// AtenLog10Op
template <>
LogicalResult ConvertAtenOp<AtenLog10Op>::matchAndRewrite(
    AtenLog10Op op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getSelf();
  auto inputTy = dyn_cast<RankedTensorType>(input.getType());
  if (!inputTy) {
    return op.emitError("only ranked tensor type is supported.");
  }

  auto outTy = cast<TensorType>(getTypeConverter()->convertType(op.getType()));
  input =
      hlo::promoteType(rewriter, op.getLoc(), input, outTy.getElementType());

  auto ten = hlo::getConstantLike(rewriter, op.getLoc(), 10.0, input);
  auto log10Op = stablehlo::LogOp::create(rewriter, op.getLoc(), ten);
  auto logInputOp = stablehlo::LogOp::create(rewriter, op.getLoc(), input);

  rewriter.replaceOpWithNewOp<stablehlo::DivOp>(op, outTy, logInputOp, log10Op);
  return success();
}

// AtenLogitOp
template <>
LogicalResult ConvertAtenOp<AtenLogitOp>::matchAndRewrite(
    AtenLogitOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto loc = op.getLoc();

  Value self = adaptor.getSelf();
  auto selfTy = dyn_cast<RankedTensorType>(self.getType());
  if (!selfTy) {
    return op.emitError("only ranked tensor type is supported.");
  }

  auto outTy = cast<TensorType>(getTypeConverter()->convertType(op.getType()));
  self = hlo::promoteType(rewriter, op.getLoc(), self, outTy.getElementType());

  selfTy = dyn_cast<RankedTensorType>(self.getType());

  Value eps = adaptor.getEps();
  auto epsTy = eps.getType();
  Value newSelf;
  if (!isa<Torch::NoneType>(epsTy)) {
    auto epsTensor = hlo::scalarToStablehloTensor(rewriter, op, eps,
                                                  selfTy.getElementType());
    Value oneEpsTensor = hlo::getConstantLike(rewriter, loc, 1.0, epsTensor);
    auto max =
        stablehlo::SubtractOp::create(rewriter, loc, oneEpsTensor, epsTensor);
    newSelf = stablehlo::ClampOp::create(rewriter, loc, epsTensor, self, max);
  } else {
    newSelf = self;
  }

  Value one = hlo::getConstantLike(rewriter, loc, 1.0, self);
  Value zi1 = stablehlo::SubtractOp::create(rewriter, loc, one, newSelf);
  Value newZi = stablehlo::DivOp::create(rewriter, loc, newSelf, zi1);

  Value log = stablehlo::LogOp::create(rewriter, loc, outTy, newZi);

  rewriter.replaceOp(op, log);

  return success();
}

// AtenErfOp
template <>
LogicalResult ConvertAtenOp<AtenErfOp>::matchAndRewrite(
    AtenErfOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getSelf();
  auto inputType = cast<TensorType>(input.getType());
  if (!isa<mlir::FloatType>(inputType.getElementType())) {
    return rewriter.notifyMatchFailure(op, "only float tensor is supported");
  }
  rewriter.replaceOpWithNewOp<chlo::ErfOp>(
      op, getTypeConverter()->convertType(op.getType()), input);
  return success();
}

// AtenBatchNormOp
template <>
LogicalResult ConvertAtenOp<AtenBatchNormOp>::matchAndRewrite(
    AtenBatchNormOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getInput();
  auto inputTy = cast<RankedTensorType>(input.getType());
  Value weight = adaptor.getWeight();
  Value bias = adaptor.getBias();
  Value runningMean = adaptor.getRunningMean();
  Value runningVar = adaptor.getRunningVar();
  // momentum is ignored
  Value momentum = adaptor.getMomentum();
  (void)momentum;

  // handle feature index, see torch's BatchNorm1d, BatchNorm2d, BatchNorm3d,
  // all of NC, NCL, NCHW, NCDHW's feature index is 1.
  int64_t feature_index = 1;

  if (!isa<mlir::FloatType>(inputTy.getElementType())) {
    return op.emitError("only input tensor of float type is supported");
  }
  auto inputElemTy = cast<mlir::FloatType>(inputTy.getElementType());

  Value channelDim =
      tensor::DimOp::create(rewriter, op->getLoc(), input, feature_index);

  if (options.dimSizeIndexBits == 32) {
    auto channelDimI64 = mlir::arith::IndexCastOp::create(
        rewriter, op->getLoc(), rewriter.getI64Type(), channelDim);
    channelDim = arith::TruncIOp::create(rewriter, op->getLoc(),
                                         rewriter.getI32Type(), channelDimI64);
  }

  Value channelShape = tensor::FromElementsOp::create(rewriter, op->getLoc(),
                                                      ValueRange{channelDim});
  if (failed(checkNotNone(rewriter, op, weight))) {
    weight = hlo::getConstantOfShape(
        rewriter, op->getLoc(), APFloat(inputElemTy.getFloatSemantics(), 1),
        channelShape,
        RankedTensorType::get({inputTy.getShape()[1]},
                              inputTy.getElementType()));
  }
  if (failed(checkNotNone(rewriter, op, bias))) {
    bias = hlo::getConstantOfShape(
        rewriter, op->getLoc(), APFloat(inputElemTy.getFloatSemantics(), 0),
        channelShape,
        RankedTensorType::get({inputTy.getShape()[1]},
                              inputTy.getElementType()));
  }
  if (failed(checkNotNone(rewriter, op, runningVar))) {
    runningVar = hlo::getConstantOfShape(
        rewriter, op->getLoc(), APFloat(inputElemTy.getFloatSemantics(), 1),
        channelShape,
        RankedTensorType::get({inputTy.getShape()[1]},
                              inputTy.getElementType()));
  }
  if (failed(checkNotNone(rewriter, op, runningMean))) {
    runningMean = hlo::getConstantOfShape(
        rewriter, op->getLoc(), APFloat(inputElemTy.getFloatSemantics(), 0),
        channelShape,
        RankedTensorType::get({inputTy.getShape()[1]},
                              inputTy.getElementType()));
  }

  auto weightTy = cast<RankedTensorType>(weight.getType());
  auto biasTy = cast<RankedTensorType>(bias.getType());
  auto runningMeanTy = cast<RankedTensorType>(runningMean.getType());
  auto runningVarTy = cast<RankedTensorType>(runningVar.getType());

  if (weightTy.getRank() != 1 || biasTy.getRank() != 1 ||
      runningMeanTy.getRank() != 1 || runningVarTy.getRank() != 1) {
    return rewriter.notifyMatchFailure(
        op, "expect weight, bias, running_mean and running_var to be rank 1");
  }
  if (!isa<mlir::FloatType>(weightTy.getElementType()) ||
      !isa<mlir::FloatType>(biasTy.getElementType()) ||
      !isa<mlir::FloatType>(runningMeanTy.getElementType()) ||
      !isa<mlir::FloatType>(runningVarTy.getElementType())) {
    return op.emitError("only float weight/bias/runningMean/runningVar tensor "
                        "of float type is supported");
  }

  double eps = 0.0;
  if (!matchPattern(op.getEps(), m_TorchConstantFloat(&eps))) {
    return rewriter.notifyMatchFailure(op, "non-float(double) eps unsupported");
  }
  bool training = false;
  if (!matchPattern(op.getTraining(), m_TorchConstantBool(&training))) {
    return rewriter.notifyMatchFailure(op, "non-bool training unsupported");
  }
  // TODO: handle cudnnEnabled parameter. Here, we just ignore it!
  bool cudnnEnabled = false;
  if (!matchPattern(op.getCudnnEnabled(), m_TorchConstantBool(&cudnnEnabled))) {
    return rewriter.notifyMatchFailure(op,
                                       "non-bool cudnn_enabled unsupported");
  }
  if (training) {
    TensorType outputTy =
        cast<TensorType>(getTypeConverter()->convertType(op.getType()));

    Value output;
    // supported mixed types, like input type is fp16 and weight type is fp32.
    if (inputTy.getElementType() != weightTy.getElementType()) {
      Type computeType = inputTy.getElementType();
      if (weightTy.getElementType().getIntOrFloatBitWidth() >
          inputTy.getElementType().getIntOrFloatBitWidth()) {
        computeType = weightTy.getElementType();
      }
      input = hlo::promoteType(rewriter, op.getLoc(), input, computeType);
      weight = hlo::promoteType(rewriter, op.getLoc(), weight, computeType);
      bias = hlo::promoteType(rewriter, op.getLoc(), bias, computeType);
      auto batchNormTrainingResult = stablehlo::BatchNormTrainingOp::create(
          rewriter, op.getLoc(),
          RankedTensorType::get(inputTy.getShape(), computeType),
          RankedTensorType::get(weightTy.getShape(), computeType),
          RankedTensorType::get(weightTy.getShape(), computeType), input,
          weight, bias, rewriter.getF32FloatAttr(eps),
          rewriter.getI64IntegerAttr(feature_index));
      output = hlo::promoteType(rewriter, op.getLoc(),
                                batchNormTrainingResult.getResult(0),
                                outputTy.getElementType());
    } else {
      auto batchNormTrainingResult = stablehlo::BatchNormTrainingOp::create(
          rewriter, op.getLoc(), outputTy, weightTy, weightTy, input, weight,
          bias, rewriter.getF32FloatAttr(eps),
          rewriter.getI64IntegerAttr(feature_index));
      output = batchNormTrainingResult.getResult(0);
    }
    rewriter.replaceOp(op, output);
    return success();
  } else {
    TensorType outputTy =
        cast<TensorType>(getTypeConverter()->convertType(op.getType()));
    SmallVector<int64_t, 4> castShape{inputTy.getShape().begin(),
                                      inputTy.getShape().end()};
    castShape[1] = weightTy.getShape()[0];
    auto castTy = RankedTensorType::get(castShape, inputTy.getElementType());
    // Feature counts must match among operands of
    // stablehlo::BatchNormInferenceOp.
    Value inputCasted =
        tensor::CastOp::create(rewriter, op.getLoc(), castTy, input);

    Value output;
    // supported mixed types, like input type is fp16 and weight type is fp32.
    if (inputTy.getElementType() != weightTy.getElementType()) {
      Type computeType = inputTy.getElementType();
      if (weightTy.getElementType().getIntOrFloatBitWidth() >
          inputTy.getElementType().getIntOrFloatBitWidth()) {
        computeType = weightTy.getElementType();
      }
      input = hlo::promoteType(rewriter, op.getLoc(), inputCasted, computeType);
      weight = hlo::promoteType(rewriter, op.getLoc(), weight, computeType);
      bias = hlo::promoteType(rewriter, op.getLoc(), bias, computeType);
      runningMean =
          hlo::promoteType(rewriter, op.getLoc(), runningMean, computeType);
      runningVar =
          hlo::promoteType(rewriter, op.getLoc(), runningVar, computeType);
      Value bnResult = stablehlo::BatchNormInferenceOp::create(
          rewriter, op.getLoc(),
          RankedTensorType::get(inputTy.getShape(), computeType), input, weight,
          bias, runningMean, runningVar, rewriter.getF32FloatAttr(eps),
          rewriter.getI64IntegerAttr(feature_index));
      output = hlo::promoteType(rewriter, op.getLoc(), bnResult,
                                outputTy.getElementType());
    } else {
      output = stablehlo::BatchNormInferenceOp::create(
          rewriter, op.getLoc(), inputCasted.getType(), inputCasted, weight,
          bias, runningMean, runningVar,
          // 'epsilon' must satisfy constraint: 32-bit float attribute.
          rewriter.getF32FloatAttr(eps),
          rewriter.getI64IntegerAttr(feature_index));
    }
    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, outputTy, output);
    return success();
  }
}

// AtenNativeLayerNormOp
template <>
LogicalResult ConvertAtenOp<AtenNativeLayerNormOp>::matchAndRewrite(
    AtenNativeLayerNormOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getInput();
  auto inputTy = cast<RankedTensorType>(input.getType());
  auto inputShape = inputTy.getShape();
  auto inputRank = inputTy.getRank();
  Value weight = adaptor.getWeight();
  Value bias = adaptor.getBias();

  if (!inputTy.hasStaticShape()) {
    return op->emitError("dynamic shaped input is not supported");
  }

  SmallVector<int64_t> normalizedShape;
  if (!matchPattern(op.getNormalizedShape(),
                    m_TorchListOfConstantInts(normalizedShape))) {
    return rewriter.notifyMatchFailure(
        op, "normalized_shape must be a list of const int");
  }
  double eps = 0;
  if (!matchPattern(op.getEps(), m_TorchConstantFloat(&eps))) {
    return rewriter.notifyMatchFailure(op,
                                       "non const float eps is unsupported");
  }
  if (failed(checkNotNone(rewriter, op, weight)) ||
      failed(checkNotNone(rewriter, op, bias))) {
    return op->emitError("none weight or bias is unsupported");
  }
  auto weightTy = cast<RankedTensorType>(weight.getType());
  auto biasTy = cast<RankedTensorType>(bias.getType());

  if (!isa<mlir::FloatType>(inputTy.getElementType()) ||
      !isa<mlir::FloatType>(biasTy.getElementType()) ||
      !isa<mlir::FloatType>(weightTy.getElementType())) {
    return op->emitError("currently only float data type are supported");
  }
  int64_t normalizedShapeRank = normalizedShape.size();
  if (weightTy.getRank() != normalizedShapeRank ||
      biasTy.getRank() != normalizedShapeRank ||
      inputRank < normalizedShapeRank || normalizedShapeRank < 1) {
    return rewriter.notifyMatchFailure(op, "input or weight or bias shape or"
                                           "normalized shape not compatible");
  }
  for (int64_t i = 1; i <= normalizedShapeRank; i++) {
    if (inputShape[inputRank - i] != normalizedShape[normalizedShapeRank - i] ||
        weightTy.getShape()[normalizedShapeRank - i] !=
            normalizedShape[normalizedShapeRank - i] ||
        biasTy.getShape()[normalizedShapeRank - i] !=
            normalizedShape[normalizedShapeRank - i]) {
      return op.emitError("mismatching contracting dimension");
    }
  }

  // Flatten dims to fit batch_norm operation.
  int64_t numFeatureDimSize = 1;
  int64_t numEmbeddingDimSize = 1;
  for (int64_t i = 0; i < inputRank - normalizedShapeRank; i++) {
    numFeatureDimSize *= inputShape[i];
  }
  for (int64_t i = 0; i < normalizedShapeRank; i++) {
    numEmbeddingDimSize *= normalizedShape[i];
  }
  SmallVector<int64_t> inputFlattenShape{1, numFeatureDimSize,
                                         numEmbeddingDimSize};
  SmallVector<int64_t> meanOrVarStablehloOutShape{numFeatureDimSize};

  auto stablehloBatchNormOutTy =
      RankedTensorType::get(inputFlattenShape, inputTy.getElementType());
  auto stablehloBathNormOutMeanOrVarTy = RankedTensorType::get(
      meanOrVarStablehloOutShape, inputTy.getElementType());

  // Reshape input
  auto stablehloInput = stablehlo::DynamicReshapeOp::create(
      rewriter, op->getLoc(), stablehloBatchNormOutTy, input,
      hlo::getConstTensor(rewriter, op, llvm::ArrayRef(inputFlattenShape),
                          {static_cast<int64_t>(inputFlattenShape.size())})
          .value());

  // Generate "scale" and "offset" Value for stablehlo.BatchNormTrainingOp.
  SmallVector<APFloat> zeroConstVec(
      numFeatureDimSize,
      APFloat::getZero(
          cast<mlir::FloatType>(inputTy.getElementType()).getFloatSemantics()));
  SmallVector<APFloat> oneConstVec(
      numFeatureDimSize,
      APFloat(
          cast<mlir::FloatType>(inputTy.getElementType()).getFloatSemantics(),
          1));
  auto oneOrZeroConstType =
      RankedTensorType::get({numFeatureDimSize}, inputTy.getElementType());

  Value scale = stablehlo::ConstantOp::create(
      rewriter, op->getLoc(), oneOrZeroConstType,
      DenseElementsAttr::get(oneOrZeroConstType, oneConstVec));
  Value offset = stablehlo::ConstantOp::create(
      rewriter, op->getLoc(), oneOrZeroConstType,
      DenseElementsAttr::get(oneOrZeroConstType, zeroConstVec));
  auto batchNormTrainingResult = stablehlo::BatchNormTrainingOp::create(
      rewriter, op->getLoc(), stablehloBatchNormOutTy,
      stablehloBathNormOutMeanOrVarTy, stablehloBathNormOutMeanOrVarTy,
      stablehloInput, scale, offset, rewriter.getF32FloatAttr(eps),
      rewriter.getI64IntegerAttr(1));

  // Reshape back
  auto outputTy =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType(0)));
  auto outputMeanOrVarTy =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType(1)));

  auto output = stablehlo::DynamicReshapeOp::create(
      rewriter, op->getLoc(), outputTy, batchNormTrainingResult.getResult(0),
      hlo::getConstTensor(rewriter, op, outputTy.getShape(),
                          {static_cast<int64_t>(outputTy.getShape().size())})
          .value());
  auto mean = stablehlo::DynamicReshapeOp::create(
      rewriter, op->getLoc(), outputMeanOrVarTy,
      batchNormTrainingResult.getResult(1),
      hlo::getConstTensor(
          rewriter, op, outputMeanOrVarTy.getShape(),
          {static_cast<int64_t>(outputMeanOrVarTy.getShape().size())})
          .value());
  auto var = stablehlo::DynamicReshapeOp::create(
      rewriter, op->getLoc(), outputMeanOrVarTy,
      batchNormTrainingResult.getResult(2),
      hlo::getConstTensor(
          rewriter, op, outputMeanOrVarTy.getShape(),
          {static_cast<int64_t>(outputMeanOrVarTy.getShape().size())})
          .value());

  // Apply affine transform: output x weight + bias [element-wise]
  auto bcastedWeight =
      hlo::promoteAndBroadcast(rewriter, weight, outputTy, std::nullopt);
  auto bcastedBias =
      hlo::promoteAndBroadcast(rewriter, bias, outputTy, std::nullopt);
  auto outputMulWeight =
      stablehlo::MulOp::create(rewriter, op->getLoc(), output, bcastedWeight);
  auto finalOuput = stablehlo::AddOp::create(rewriter, op->getLoc(),
                                             outputMulWeight, bcastedBias);
  rewriter.replaceOp(op, {finalOuput, mean, var});
  return success();
}

// AtenCatOp
template <>
LogicalResult ConvertAtenOp<AtenCatOp>::matchAndRewrite(
    AtenCatOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto outType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  int64_t dim;
  if (!matchPattern(op.getDim(), m_TorchConstantInt(&dim))) {
    return rewriter.notifyMatchFailure(op,
                                       "only constant dim param is supported");
  }
  dim = toPositiveDim(dim, outType.getRank());
  if (!isValidDim(dim, outType.getRank()))
    return rewriter.notifyMatchFailure(op, "dim is statically invalid");

  SmallVector<Value> torchTensors;
  if (!getListConstructElements(op.getTensors(), torchTensors)) {
    return rewriter.notifyMatchFailure(
        op, "input should comes from a PrimListConstructOp");
  }
  SmallVector<Value> builtinTensors = getTypeConvertedValues(
      rewriter, op->getLoc(), getTypeConverter(), torchTensors);

  // Promote type
  for (auto &v : builtinTensors) {
    v = hlo::promoteType(rewriter, op->getLoc(), v, outType.getElementType());
  }

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

// AtenNumelOp
template <>
LogicalResult ConvertAtenOp<AtenNumelOp>::matchAndRewrite(
    AtenNumelOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto self = adaptor.getSelf();
  auto selfTy = dyn_cast<RankedTensorType>(self.getType());
  size_t rank = selfTy.getRank();

  Type intType = rewriter.getIntegerType(options.dimSizeIndexBits);
  auto loc = op->getLoc();
  Value numel = arith::ConstantOp::create(rewriter, loc,
                                          rewriter.getIntegerAttr(intType, 1));
  for (size_t d = 0; d < rank; ++d) {
    Value dimSize = arith::IndexCastOp::create(
        rewriter, loc, intType, tensor::DimOp::create(rewriter, loc, self, d));
    numel = arith::MulIOp::create(rewriter, loc, numel, dimSize);
  }

  auto outTy = getTypeConverter()->convertType(op.getType());
  if (outTy != numel.getType()) {
    rewriter.replaceOpWithNewOp<arith::ExtSIOp>(op, outTy, numel);
  } else {
    rewriter.replaceOp(op, numel);
  }
  return success();
}

// AtenClampOp
template <>
LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewrite(
    AtenClampOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getSelf();
  auto inputType = cast<RankedTensorType>(input.getType());
  auto inputElemType = inputType.getElementType();
  Value minValue = adaptor.getMin();
  Value maxValue = adaptor.getMax();
  if (failed(checkNotNone(rewriter, op, minValue)) &&
      failed(checkNotNone(rewriter, op, maxValue))) {
    return rewriter.notifyMatchFailure(
        op, "this op should be folded as its `min` and `max` both are none");
  } else if (failed(checkNotNone(rewriter, op, minValue))) {
    maxValue =
        hlo::scalarToStablehloTensor(rewriter, op, maxValue, inputElemType);
    auto minInfo = getMinValueOfDtype(op, inputElemType, rewriter,
                                      options.allowNonFinites);
    if (failed(minInfo)) {
      return rewriter.notifyMatchFailure(
          op, "failed to generate min value of dtype");
    }
    minValue = *minInfo;
  } else if (failed(checkNotNone(rewriter, op, maxValue))) {
    minValue =
        hlo::scalarToStablehloTensor(rewriter, op, minValue, inputElemType);
    auto maxInfo = getMaxValueOfDtype(op, inputElemType, rewriter,
                                      options.allowNonFinites);
    if (failed(maxInfo)) {
      return rewriter.notifyMatchFailure(
          op, "failed to generate max value of dtype");
    }
    maxValue = *maxInfo;
  } else {
    minValue =
        hlo::scalarToStablehloTensor(rewriter, op, minValue, inputElemType);
    maxValue =
        hlo::scalarToStablehloTensor(rewriter, op, maxValue, inputElemType);
  }
  rewriter.replaceOpWithNewOp<stablehlo::ClampOp>(op, minValue, input,
                                                  maxValue);
  return success();
}

// AtenClampTensorOp
template <>
LogicalResult ConvertAtenOp<AtenClampTensorOp>::matchAndRewrite(
    AtenClampTensorOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value input = adaptor.getSelf();
  auto inputType = cast<RankedTensorType>(input.getType());
  auto inputElemType = inputType.getElementType();
  Value minValue = adaptor.getMin();
  Value maxValue = adaptor.getMax();
  auto minIsNotNone = checkNotNone(rewriter, op, minValue);
  auto maxIsNotNone = checkNotNone(rewriter, op, maxValue);
  if (failed(minIsNotNone) && failed(maxIsNotNone)) {
    return rewriter.notifyMatchFailure(
        op, "this op should be folded as its `min` and `max` both are none");
  } else if (failed(minIsNotNone)) {
    auto minInfo = getMinValueOfDtype(op, inputElemType, rewriter,
                                      options.allowNonFinites);
    if (failed(minInfo)) {
      return rewriter.notifyMatchFailure(
          op, "failed to generate min value of dtype");
    }
    minValue = *minInfo;
  } else if (failed(maxIsNotNone)) {
    auto maxInfo = getMaxValueOfDtype(op, inputElemType, rewriter,
                                      options.allowNonFinites);
    if (failed(maxInfo)) {
      return rewriter.notifyMatchFailure(
          op, "failed to generate max value of dtype");
    }
    maxValue = *maxInfo;
  }
  if (inputType.hasStaticShape()) {
    minValue =
        hlo::promoteAndBroadcast(rewriter, minValue, inputType, std::nullopt);
    maxValue =
        hlo::promoteAndBroadcast(rewriter, maxValue, inputType, std::nullopt);
  }
  rewriter.replaceOpWithNewOp<stablehlo::ClampOp>(op, minValue, input,
                                                  maxValue);
  return success();
}

// AtenArangeStartStepOp
// aten.arange.start_step = range(ceil((end-start)/step)) * step + start.
template <>
LogicalResult ConvertAtenOp<AtenArangeStartStepOp>::matchAndRewrite(
    AtenArangeStartStepOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op->getLoc();

  // Get element type of resultType as dtype
  auto outType = cast<RankedTensorType>(
      this->getTypeConverter()->convertType(op.getType()));
  auto dtype = outType.getElementType();
  if (!isa<mlir::IntegerType>(dtype) && !isa<mlir::FloatType>(dtype)) {
    return rewriter.notifyMatchFailure(
        op, "unimplemented: only int or float dtype supported");
  }

  Value start =
      hlo::scalarToStablehloTensor(rewriter, op, adaptor.getStart(), dtype);
  Value end =
      hlo::scalarToStablehloTensor(rewriter, op, adaptor.getEnd(), dtype);
  Value step =
      hlo::scalarToStablehloTensor(rewriter, op, adaptor.getStep(), dtype);

  // Get length of the 1-d output tensor
  Value subOut = stablehlo::SubtractOp::create(rewriter, loc, end, start);
  // promote div to f64
  Type divType = RankedTensorType::get({}, rewriter.getF64Type());
  Value divOut = stablehlo::DivOp::create(
      rewriter, loc,
      stablehlo::ConvertOp::create(rewriter, loc, divType, subOut),
      stablehlo::ConvertOp::create(rewriter, loc, divType, step));
  // ceil to i64
  Value resultLength = stablehlo::ConvertOp::create(
      rewriter, loc, RankedTensorType::get({}, rewriter.getI64Type()),
      stablehlo::CeilOp::create(rewriter, loc, divOut));
  resultLength = stablehlo::ReshapeOp::create(
      rewriter, loc, RankedTensorType::get({1}, rewriter.getI64Type()),
      resultLength);

  Value window =
      stablehlo::DynamicIotaOp::create(rewriter, loc, outType, resultLength, 0);
  DenseI64ArrayAttr broadcastDimensions;
  Value mulOut = chlo::BroadcastMulOp::create(rewriter, loc, window, step,
                                              broadcastDimensions);
  rewriter.replaceOpWithNewOp<chlo::BroadcastAddOp>(op, mulOut, start,
                                                    broadcastDimensions);
  return success();
}

template <>
LogicalResult ConvertAtenOp<AtenConstantPadNdOp>::matchAndRewrite(
    AtenConstantPadNdOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value self = adaptor.getSelf();
  auto selfTy = cast<RankedTensorType>(self.getType());
  auto selfElemTy = selfTy.getElementType();
  int64_t rank = selfTy.getRank();

  SmallVector<int64_t> padInts;
  if (!matchPattern(op.getPad(), m_TorchListOfConstantInts(padInts)))
    return rewriter.notifyMatchFailure(op,
                                       "only support constant int pad ranges");
  uint64_t padRank = padInts.size() / 2;
  if (padRank * 2 != padInts.size())
    return rewriter.notifyMatchFailure(op, "pad range size is not even");
  if (rank < 0 || padRank > (uint64_t)rank)
    return rewriter.notifyMatchFailure(op, "padding exceeds tensor rank");

  // Initialize low/high paddings with 0 for all the dims.
  SmallVector<int64_t> lowPadding(/*Size=*/rank, /*Value=*/0);
  SmallVector<int64_t> highPadding(/*Size=*/rank, /*Value=*/0);
  // Add the requested padding - note op.pad() is highest dim first ordered
  // pairs of low,high.
  // Add the requested padding - note op.pad() is highest dim first ordered
  // pairs of low,high.
  for (uint64_t i = 0; i < padRank; ++i) {
    lowPadding[rank - i - 1] = padInts[i * 2];
    highPadding[rank - i - 1] = padInts[i * 2 + 1];
  }

  Value constantValue = hlo::scalarToStablehloTensor(
      rewriter, op, adaptor.getValue(), selfElemTy);

  SmallVector<int64_t> interiorPadding(rank, 0);
  rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
      op, self, constantValue, lowPadding, highPadding, interiorPadding);
  return success();
}

template <>
LogicalResult ConvertAtenOp<AtenReflectionPad1dOp>::matchAndRewrite(
    AtenReflectionPad1dOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op.getLoc();
  Value self = adaptor.getSelf();
  auto selfTy = cast<RankedTensorType>(self.getType());
  if (!selfTy.hasStaticShape()) {
    return rewriter.notifyMatchFailure(op, "only support static shape");
  }
  int64_t rank = selfTy.getRank();
  int64_t dim = rank - 1;

  SmallVector<int64_t> padInts;
  if (!matchPattern(op.getPadding(), m_TorchListOfConstantInts(padInts))) {
    return rewriter.notifyMatchFailure(op,
                                       "only support constant int pad ranges");
  }
  if (padInts.size() != 2) {
    return rewriter.notifyMatchFailure(op, "pad size must be 2");
  }
  if (padInts[0] >= selfTy.getDimSize(dim) ||
      padInts[1] >= selfTy.getDimSize(dim)) {
    return rewriter.notifyMatchFailure(op,
                                       "pad size must be less than dim size");
  }

  Value left;
  {
    SmallVector<int64_t> startIndices(rank, 0);
    SmallVector<int64_t> limitIndices(selfTy.getShape().begin(),
                                      selfTy.getShape().end());
    SmallVector<int64_t> strides(rank, 1);
    startIndices[dim] = 1;
    limitIndices[dim] = padInts[0] + 1;
    left = stablehlo::SliceOp::create(rewriter, loc, self, startIndices,
                                      limitIndices, strides);
    left = stablehlo::ReverseOp::create(rewriter, loc, left,
                                        ArrayRef<int64_t>({dim}));
  }
  Value right;
  {
    SmallVector<int64_t> startIndices(rank, 0);
    SmallVector<int64_t> limitIndices(selfTy.getShape().begin(),
                                      selfTy.getShape().end());
    SmallVector<int64_t> strides(rank, 1);
    startIndices[dim] = selfTy.getDimSize(dim) - 1 - padInts[1];
    limitIndices[dim] = selfTy.getDimSize(dim) - 1;
    right = stablehlo::SliceOp::create(rewriter, loc, self, startIndices,
                                       limitIndices, strides);
    right = stablehlo::ReverseOp::create(rewriter, loc, right,
                                         ArrayRef<int64_t>({dim}));
  }
  Value result = stablehlo::ConcatenateOp::create(
      rewriter, loc, ValueRange{left, self, right}, dim);
  rewriter.replaceOp(op, result);
  return success();
}

template <>
LogicalResult ConvertAtenOp<AtenGeluBackwardOp>::matchAndRewrite(
    AtenGeluBackwardOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op.getLoc();
  Value input = adaptor.getSelf();
  auto outType =
      cast<TensorType>(this->getTypeConverter()->convertType(op.getType()));
  if (!outType) {
    return op.emitError("only tensor type is supported");
  }
  // TODO: Handle approximate.
  std::string approximate;
  if (!matchPattern(op.getApproximate(), m_TorchConstantStr(approximate)) ||
      approximate != "none") {
    return rewriter.notifyMatchFailure(op, "Unsupported value of approximate");
  }
  // Create constant value
  Value kAlpha =
      hlo::getConstantLike(rewriter, loc, 0.70710678118654752440, input);
  Value cstAlpha0 =
      hlo::getConstantLike(rewriter, loc, 1.12837916709551257390, input);
  Value half = hlo::getConstantLike(rewriter, loc, .5, input);
  Value one = hlo::getConstantLike(rewriter, loc, 1.0, input);
  Value negHalf = hlo::getConstantLike(rewriter, loc, -0.5, input);

  // Compute
  Value kBeta0 =
      stablehlo::MulOp::create(rewriter, loc, outType, kAlpha, cstAlpha0);
  Value kBeta = stablehlo::MulOp::create(rewriter, loc, outType, kBeta0, half);
  Value erfArg = stablehlo::MulOp::create(rewriter, loc, outType, kAlpha,
                                          adaptor.getSelf());
  Value erf = mlir::chlo::ErfOp::create(rewriter, loc, outType, erfArg);
  Value erfAdd = stablehlo::AddOp::create(rewriter, loc, outType, erf, one);
  Value cdf = stablehlo::MulOp::create(rewriter, loc, outType, erfAdd, half);
  Value inputSquared = stablehlo::MulOp::create(
      rewriter, loc, outType, adaptor.getSelf(), adaptor.getSelf());
  Value negHalfInputSquared =
      stablehlo::MulOp::create(rewriter, loc, outType, inputSquared, negHalf);
  Value expRes =
      stablehlo::ExpOp::create(rewriter, loc, outType, negHalfInputSquared);
  Value pdf = stablehlo::MulOp::create(rewriter, loc, outType, kBeta, expRes);
  Value pdfTimesInput =
      stablehlo::MulOp::create(rewriter, loc, outType, pdf, adaptor.getSelf());
  Value pdfTimesInputAddCdf =
      stablehlo::AddOp::create(rewriter, loc, outType, pdfTimesInput, cdf);
  rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
      op, outType, adaptor.getGradOutput(), pdfTimesInputAddCdf);
  return success();
}

// Converts `aten.empty.memory_format` to `tensor.empty` op.
template <>
LogicalResult ConvertAtenOp<AtenEmptyMemoryFormatOp>::matchAndRewrite(
    AtenEmptyMemoryFormatOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {

  // TODO: Add support pin_memory and memory_format features.
  // At this point all tensors should have value semantics, and hence the
  // `layout` check can be ignored.

  // The pin_memory should be either `False` or `none`.
  bool pinMemory;
  if (!isa<Torch::NoneType>(op.getPinMemory().getType()) &&
      (!matchPattern(op.getPinMemory(), m_TorchConstantBool(&pinMemory)) ||
       pinMemory))
    return rewriter.notifyMatchFailure(
        op, "unimplemented: pin_memory must be either None or false");

  // Only `none`, `contiguous` and `preserve` memory_format is supported.
  if (!isa<Torch::NoneType>(op.getMemoryFormat().getType())) {
    int64_t memoryFormat;
    if (!matchPattern(op.getMemoryFormat(), m_TorchConstantInt(&memoryFormat)))
      return rewriter.notifyMatchFailure(
          op, "unimplemented: the memory format should be specified in "
              "an integer constant");
    if (memoryFormat != torch_upstream::MemoryFormat::Contiguous &&
        memoryFormat != torch_upstream::MemoryFormat::Preserve)
      return rewriter.notifyMatchFailure(
          op, "unimplemented: only none, contiguous and preserve "
              "memory_format is supported");
  }

  if (!isa<Torch::NoneType>(op.getDevice().getType())) {
    std::string device;
    if (!matchPattern(op.getDevice(), m_TorchConstantDevice(device)))
      return rewriter.notifyMatchFailure(
          op, "unimplemented: device must be a constant str");
  }

  // TODO: Add support for non-strided layout.
  // torch.layout is by default strided i.e. 0.
  if (!isa<Torch::NoneType>(op.getLayout().getType())) {
    int64_t tensorLayout;
    if (!matchPattern(op.getLayout(), m_TorchConstantInt(&tensorLayout)))
      return rewriter.notifyMatchFailure(
          op, "unimplemented: layout must be a constant");
    else if (tensorLayout != torch_upstream::Layout::Strided)
      return rewriter.notifyMatchFailure(
          op, "unimplemented: layout is expected to be strided");
  }

  Location loc = op.getLoc();
  const TypeConverter *typeConverter = this->getTypeConverter();
  SmallVector<Value> resultSizeTorchInt, resultSize, resultSizeIndex;
  if (!getListConstructElements(op.getSize(), resultSizeTorchInt)) {
    return rewriter.notifyMatchFailure(
        op, "unimplemented: size must be constructed using ListConstruct");
  }
  resultSize =
      getTypeConvertedValues(rewriter, loc, typeConverter, resultSizeTorchInt);
  for (auto size : resultSize)
    resultSizeIndex.push_back(castIntToIndex(rewriter, loc, size));

  auto resultType =
      cast<RankedTensorType>(typeConverter->convertType(op.getType()));
  Type resultElementType;
  if (isa<Torch::NoneType>(op.getDtype().getType())) {
    resultElementType = resultType.getElementType();
  } else {
    int64_t dtypeInt;
    if (!matchPattern(op.getDtype(), m_TorchConstantInt(&dtypeInt)))
      return rewriter.notifyMatchFailure(
          op, "unimplemented: dtype must be a constant integer or none");
    FailureOr<Type> maybeResultElementType =
        torch_to_stablehlo::getBackendTypeForScalarType(
            op->getContext(), (torch_upstream::ScalarType)dtypeInt);
    if (failed(maybeResultElementType)) {
      return rewriter.notifyMatchFailure(
          op, "unable to convert `dtypeInt` to builtin type");
    }
    resultElementType = *maybeResultElementType;
  }

  // Create an uninitialized tensor of `resultSize` shape.
  Value initTensor = tensor::EmptyOp::create(
      rewriter, loc, getAsOpFoldResult(resultSizeIndex), resultElementType);
  rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, initTensor);
  return success();
}

// RuntimeAssertOp
namespace {
class ConvertRuntimeAssertOp : public OpConversionPattern<RuntimeAssertOp> {
public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(RuntimeAssertOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    bool condition;
    if (!matchPattern(op.getCondition(), m_TorchConstantBool(&condition))) {
      return rewriter.notifyMatchFailure(
          op, "unimplemented: condition must be a constant");
    }
    if (!condition) {
      return op->emitError("condition must be true");
    }
    rewriter.eraseOp(op);
    return success();
  }
};
} // namespace

// AtenFillScalarOp
template <>
LogicalResult ConvertAtenOp<AtenFillScalarOp>::matchAndRewrite(
    AtenFillScalarOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto outType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  auto dtype = outType.getElementType();
  Value scalarTensor =
      hlo::scalarToStablehloTensor(rewriter, op, adaptor.getValue(), dtype);
  Value shapeTensor =
      shape::ShapeOfOp::create(rewriter, op->getLoc(), adaptor.getSelf());
  Value bcastScalar = stablehlo::DynamicBroadcastInDimOp::create(
      rewriter, op->getLoc(), outType, scalarTensor, shapeTensor,
      rewriter.getDenseI64ArrayAttr({}));
  rewriter.replaceOp(op, bcastScalar);
  return success();
}

// AtenFlipOp
template <>
LogicalResult ConvertAtenOp<AtenFlipOp>::matchAndRewrite(
    AtenFlipOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value self = adaptor.getSelf();
  auto outType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));

  SmallVector<int64_t> dims;
  if (!matchPattern(op.getDims(), m_TorchListOfConstantInts(dims))) {
    return rewriter.notifyMatchFailure(op, "dims must be a list of const int");
  }
  for (unsigned i = 0, e = dims.size(); i < e; i++) {
    dims[i] = toPositiveDim(dims[i], outType.getRank());
    if (!isValidDim(dims[i], outType.getRank())) {
      return rewriter.notifyMatchFailure(op, "dim is statically invalid");
    }
  }

  rewriter.replaceOpWithNewOp<stablehlo::ReverseOp>(op, outType, self, dims);
  return success();
}

// AtenFmodTensorOp
// torch.fmod(a, b) == a - a.div(b, rounding_mode="trunc") * b
template <>
LogicalResult ConvertAtenOp<AtenFmodTensorOp>::matchAndRewrite(
    AtenFmodTensorOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto loc = op->getLoc();
  Value lhs = adaptor.getSelf();
  Value rhs = adaptor.getOther();

  auto resultType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  lhs = hlo::promoteType(rewriter, op->getLoc(), lhs,
                         resultType.getElementType());
  rhs = hlo::promoteType(rewriter, op->getLoc(), rhs,
                         resultType.getElementType());

  stablehlo::MulOp mul;
  auto div = stablehlo::DivOp::create(rewriter, loc, lhs, rhs);
  if (isa<mlir::FloatType>(resultType.getElementType())) {
    // rounding mode is trunc
    auto sign = stablehlo::SignOp::create(rewriter, loc, div);
    auto abs = stablehlo::AbsOp::create(rewriter, loc, div);
    auto floor = stablehlo::FloorOp::create(rewriter, loc, abs);
    auto trunc = stablehlo::MulOp::create(rewriter, loc, sign, floor);
    mul = stablehlo::MulOp::create(rewriter, loc, trunc, rhs);
  } else {
    mul = stablehlo::MulOp::create(rewriter, loc, div, rhs);
  }
  rewriter.replaceOpWithNewOp<stablehlo::SubtractOp>(op, lhs, mul);
  return success();
}

// AtenBitwiseLeftShiftTensorOp
template <>
LogicalResult ConvertAtenOp<AtenBitwiseLeftShiftTensorOp>::matchAndRewrite(
    AtenBitwiseLeftShiftTensorOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value lhs = adaptor.getSelf();
  Value rhs = adaptor.getOther();

  auto resultType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  rhs = hlo::promoteAndBroadcast(rewriter, rhs, resultType, std::nullopt);
  rewriter.replaceOpWithNewOp<stablehlo::ShiftLeftOp>(op, lhs, rhs);
  return success();
}

// AtenBitwiseRightShiftTensorOp
template <>
LogicalResult ConvertAtenOp<AtenBitwiseRightShiftTensorOp>::matchAndRewrite(
    AtenBitwiseRightShiftTensorOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value lhs = adaptor.getSelf();
  Value rhs = adaptor.getOther();

  auto resultType =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  rhs = hlo::promoteAndBroadcast(rewriter, rhs, resultType, std::nullopt);
  rewriter.replaceOpWithNewOp<stablehlo::ShiftRightArithmeticOp>(op, lhs, rhs);
  return success();
}

template <>
LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewrite(
    AtenTrilOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {

  Location loc = op.getLoc();

  Value self = adaptor.getSelf();

  auto selfTy = cast<RankedTensorType>(self.getType());
  if (!selfTy.hasStaticShape()) {
    return op->emitError("dynamic shaped input is not supported");
  }

  ArrayRef<int64_t> selfShape = selfTy.getShape();
  int64_t selfRank = selfTy.getRank();
  auto iotaElementTy = mlir::IntegerType::get(op.getContext(), 64);
  auto iotaTy = RankedTensorType::get(
      {selfShape[selfRank - 2], selfShape[selfRank - 1]}, iotaElementTy);
  Value colIdxTensor =
      stablehlo::IotaOp::create(rewriter, loc, iotaTy, 1).getResult();
  Value rowIdxTensor =
      stablehlo::IotaOp::create(rewriter, loc, iotaTy, 0).getResult();

  Value diagonal = adaptor.getDiagonal();
  Value diagonalTensor =
      tensor::FromElementsOp::create(rewriter, loc, diagonal).getResult();

  auto bcastDimensions = rewriter.getDenseI64ArrayAttr({1});
  Value shiftedRowIdxTensor = chlo::BroadcastAddOp::create(
      rewriter, loc, rowIdxTensor, diagonalTensor, bcastDimensions);

  auto cmpDirectionAttr = stablehlo::ComparisonDirectionAttr::get(
      rewriter.getContext(), stablehlo::ComparisonDirection::LE);
  auto cmpTypeAttr = stablehlo::ComparisonTypeAttr::get(
      rewriter.getContext(), stablehlo::ComparisonType::SIGNED);
  auto cmpTy = iotaTy.clone(rewriter.getI1Type());
  Value cmpRes = stablehlo::CompareOp::create(rewriter, loc, cmpTy,
                                              colIdxTensor, shiftedRowIdxTensor,
                                              cmpDirectionAttr, cmpTypeAttr);

  auto resTy =
      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));

  auto bcastTy = resTy.clone(rewriter.getI1Type());
  auto bcastAttr = rewriter.getDenseI64ArrayAttr({selfRank - 2, selfRank - 1});
  Value bcastedCmpRes = stablehlo::BroadcastInDimOp::create(
      rewriter, loc, bcastTy, cmpRes, bcastAttr);

  auto resElemTy = resTy.getElementType();
  Value zeroTensor;
  if (isa<mlir::FloatType>(resElemTy)) {
    auto constAttr = SplatElementsAttr::get(
        resTy, llvm::APFloat::getZero(
                   cast<FloatType>(resElemTy).getFloatSemantics(), false));
    zeroTensor = stablehlo::ConstantOp::create(rewriter, loc, resTy, constAttr);
  } else if (isa<mlir::IntegerType>(resElemTy)) {
    auto constAttr = SplatElementsAttr::get(
        resTy,
        llvm::APInt::getZero(cast<mlir::IntegerType>(resElemTy).getWidth()));
    zeroTensor = stablehlo::ConstantOp::create(rewriter, loc, resTy, constAttr);
  } else {
    return op.emitError("element type is not float or integer");
  }

  rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(
      op.getOperation(), resTy, bcastedCmpRes, self, zeroTensor);

  return success();
}

template <>
LogicalResult ConvertAtenOp<AtenIsfiniteOp>::matchAndRewrite(
    AtenIsfiniteOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Value self = adaptor.getSelf();
  auto selfTy = cast<RankedTensorType>(self.getType());
  if (!selfTy)
    return rewriter.notifyMatchFailure(
        op, "Only Tensor types are currently supported");

  auto outType =
      dyn_cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
  Type outElemTy = outType.getElementType();
  if (!outElemTy.isInteger(1)) {
    return rewriter.notifyMatchFailure(
        op, "Only i1 output element type is supported");
  }

  rewriter.replaceOpWithNewOp<stablehlo::IsFiniteOp>(op.getOperation(), outType,
                                                     self);

  return success();
}

void mlir::torch::torch_to_stablehlo::populateBasicOpPatternsAndLegality(
    TypeConverter &typeConverter, RewritePatternSet &patterns,
    ConversionTarget &target, const TorchToStablehloOptions &options) {
  MLIRContext *context = patterns.getContext();

  target.addIllegalOp<AtenTransposeIntOp>();
  patterns.add<ConvertAtenTransposeIntOp>(typeConverter, context);
  target.addIllegalOp<RuntimeAssertOp>();
  patterns.add<ConvertRuntimeAssertOp>(typeConverter, context);

#define INSERT_UNARY_PATTERN(AtenOp, StablehloOp)                              \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenUnaryOp<AtenOp, StablehloOp>>(typeConverter, context)
  INSERT_UNARY_PATTERN(AtenNegOp, stablehlo::NegOp);
  INSERT_UNARY_PATTERN(AtenLogicalNotOp, stablehlo::NotOp);
  INSERT_UNARY_PATTERN(AtenBitwiseNotOp, stablehlo::NotOp);
  INSERT_UNARY_PATTERN(AtenAbsOp, stablehlo::AbsOp);
  INSERT_UNARY_PATTERN(AtenExpm1Op, stablehlo::Expm1Op);
#undef INSERT_UNARY_PATTERN

#define INSERT_UNARY_FPONLY_PATTERN(AtenOp, StablehloOp)                       \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenUnaryFPOnlyOp<AtenOp, StablehloOp>>(typeConverter,   \
                                                              context)
  INSERT_UNARY_FPONLY_PATTERN(AtenCeilOp, stablehlo::CeilOp);
  INSERT_UNARY_FPONLY_PATTERN(AtenFloorOp, stablehlo::FloorOp);
  INSERT_UNARY_FPONLY_PATTERN(AtenRoundOp, stablehlo::RoundNearestEvenOp);
#undef INSERT_UNARY_FPONLY_PATTERN

#define INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenOp, StablehloOp)                \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenUnaryPromoteToFPOp<AtenOp, StablehloOp>>(            \
      typeConverter, context)
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenLogOp, stablehlo::LogOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenLog1pOp, stablehlo::Log1pOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenExpOp, stablehlo::ExpOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenSqrtOp, stablehlo::SqrtOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenRsqrtOp, stablehlo::RsqrtOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenSigmoidOp, stablehlo::LogisticOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenTanhOp, stablehlo::TanhOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenSinOp, stablehlo::SineOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenCosOp, stablehlo::CosineOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenTanOp, chlo::TanOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenAsinOp, chlo::AsinOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenSinhOp, chlo::SinhOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenAcosOp, chlo::AcosOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenCoshOp, chlo::CoshOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenAtanOp, chlo::AtanOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenAsinhOp, chlo::AsinhOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenAcoshOp, chlo::AcoshOp);
  INSERT_UNARY_PROMOTE_TO_FP_PATTERN(AtenAtanhOp, chlo::AtanhOp);
#undef INSERT_UNARY_PROMOTE_TO_FP_PATTERN

#define INSERT_CONSTANT_FILL_PATTERN(AtenOp, fillVal)                          \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenConstPatternOp<AtenOp, fillVal>>(typeConverter,      \
                                                           context)
  INSERT_CONSTANT_FILL_PATTERN(AtenOnesOp, 1);
  INSERT_CONSTANT_FILL_PATTERN(AtenZerosOp, 0);
#undef INSERT_CONSTANT_FILL_PATTERN

#define INSERT_TENSOR_TO_SCALAR_PATTERN(AtenOp)                                \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenTensorToScalarLikeOp<AtenOp>>(typeConverter, context)

  INSERT_TENSOR_TO_SCALAR_PATTERN(AtenIntTensorOp);
  INSERT_TENSOR_TO_SCALAR_PATTERN(AtenFloatTensorOp);
  INSERT_TENSOR_TO_SCALAR_PATTERN(AtenBoolTensorOp);
#undef INSERT_TENSOR_TO_SCALAR_PATTERN

#define INSERT_BINARY_ADDSUB_PATTERN(AtenOp, ChloOp)                           \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenAddSubOp<AtenOp, ChloOp>>(typeConverter, context)
  INSERT_BINARY_ADDSUB_PATTERN(AtenAddTensorOp, chlo::BroadcastAddOp);
  INSERT_BINARY_ADDSUB_PATTERN(AtenAddScalarOp, chlo::BroadcastAddOp);
  INSERT_BINARY_ADDSUB_PATTERN(AtenSubTensorOp, chlo::BroadcastSubOp);
  INSERT_BINARY_ADDSUB_PATTERN(AtenSubScalarOp, chlo::BroadcastSubOp);
  INSERT_BINARY_ADDSUB_PATTERN(AtenRsubScalarOp, chlo::BroadcastSubOp);
#undef INSERT_BINARY_ADDSUB_PATTERN

#define INSERT_BINARY_MULDIV_PATTERN(AtenOp, ChloOp)                           \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenMulDivOp<AtenOp, ChloOp>>(typeConverter, context)
  INSERT_BINARY_MULDIV_PATTERN(AtenMulTensorOp, chlo::BroadcastMulOp);
  INSERT_BINARY_MULDIV_PATTERN(AtenMulScalarOp, chlo::BroadcastMulOp);
  INSERT_BINARY_MULDIV_PATTERN(AtenDivTensorOp, chlo::BroadcastDivOp);
  INSERT_BINARY_MULDIV_PATTERN(AtenDivTensorModeOp, chlo::BroadcastDivOp);
  INSERT_BINARY_MULDIV_PATTERN(AtenDivScalarOp, chlo::BroadcastDivOp);
  INSERT_BINARY_MULDIV_PATTERN(AtenDivScalarModeOp, chlo::BroadcastDivOp);
  INSERT_BINARY_MULDIV_PATTERN(AtenRemainderTensorOp, chlo::BroadcastRemOp);
  INSERT_BINARY_MULDIV_PATTERN(AtenRemainderScalarOp, chlo::BroadcastRemOp);
#undef INSERT_BINARY_MULDIV_PATTERN

#define INSERT_BINARY_COMPARE_PATTERN(AtenOp)                                  \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenCompareOp<AtenOp>>(typeConverter, context)

  INSERT_BINARY_COMPARE_PATTERN(AtenGtTensorOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenGtScalarOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenGeTensorOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenGeScalarOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenLtTensorOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenLtScalarOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenLeTensorOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenLeScalarOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenEqTensorOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenEqScalarOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenNeTensorOp);
  INSERT_BINARY_COMPARE_PATTERN(AtenNeScalarOp);
#undef INSERT_BINARY_COMPARE_PATTERN

#define INSERT_BINARY_LOGICAL_PATTERN(AtenOp, ChloOp)                          \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenLogicalBinaryOp<AtenOp, ChloOp>>(typeConverter,      \
                                                           context)

  INSERT_BINARY_LOGICAL_PATTERN(AtenLogicalOrOp, chlo::BroadcastOrOp);
  INSERT_BINARY_LOGICAL_PATTERN(AtenLogicalAndOp, chlo::BroadcastAndOp);
  INSERT_BINARY_LOGICAL_PATTERN(AtenLogicalXorOp, chlo::BroadcastXorOp);
  INSERT_BINARY_LOGICAL_PATTERN(AtenBitwiseAndScalarOp, chlo::BroadcastAndOp);
  INSERT_BINARY_LOGICAL_PATTERN(AtenBitwiseXorScalarOp, chlo::BroadcastXorOp);
  INSERT_BINARY_LOGICAL_PATTERN(AtenBitwiseOrScalarOp, chlo::BroadcastOrOp);

#undef INSERT_BINARY_LOGICAL_PATTERN

#define INSERT_BINARY_POW_PATTERN(AtenOp)                                      \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenPowOp<AtenOp>>(typeConverter, context)
  INSERT_BINARY_POW_PATTERN(AtenPowTensorScalarOp);
  INSERT_BINARY_POW_PATTERN(AtenPowTensorTensorOp);
  INSERT_BINARY_POW_PATTERN(AtenPowScalarOp);
#undef INSERT_BINARY_ADDSUB_PATTERN

#define INSERT_ATENOP_PATTERN(AtenOp)                                          \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenOp<AtenOp>>(typeConverter, context, options)

  INSERT_ATENOP_PATTERN(AtenBroadcastToOp);
  INSERT_ATENOP_PATTERN(AtenPermuteOp);

  INSERT_ATENOP_PATTERN(ValueTensorLiteralOp);
  INSERT_ATENOP_PATTERN(AtenTensorIntOp);
  INSERT_ATENOP_PATTERN(AtenReciprocalOp);
  INSERT_ATENOP_PATTERN(PrimNumToTensorScalarOp);
  INSERT_ATENOP_PATTERN(AtenScalarImplicitOp);
  INSERT_ATENOP_PATTERN(AtenContiguousOp);
  INSERT_ATENOP_PATTERN(AtenConstantPadNdOp);
  INSERT_ATENOP_PATTERN(AtenReflectionPad1dOp);

  INSERT_ATENOP_PATTERN(AtenReluOp);
  INSERT_ATENOP_PATTERN(AtenGeluOp);
  INSERT_ATENOP_PATTERN(AtenLog2Op);
  INSERT_ATENOP_PATTERN(AtenLog10Op);
  INSERT_ATENOP_PATTERN(AtenLogitOp);
  INSERT_ATENOP_PATTERN(AtenErfOp);
  INSERT_ATENOP_PATTERN(AtenGeluBackwardOp);

  INSERT_ATENOP_PATTERN(AtenCatOp);
  INSERT_ATENOP_PATTERN(AtenClampOp);
  INSERT_ATENOP_PATTERN(AtenClampTensorOp);
  INSERT_ATENOP_PATTERN(AtenArangeStartStepOp);

  INSERT_ATENOP_PATTERN(AtenBatchNormOp);
  INSERT_ATENOP_PATTERN(AtenNativeLayerNormOp);
  INSERT_ATENOP_PATTERN(AtenNumelOp);
  INSERT_ATENOP_PATTERN(AtenSizeIntOp);
  INSERT_ATENOP_PATTERN(AtenToDtypeOp);
  INSERT_ATENOP_PATTERN(AtenWhereSelfOp);

  INSERT_ATENOP_PATTERN(AtenEmptyMemoryFormatOp);
  INSERT_ATENOP_PATTERN(AtenFillScalarOp);
  INSERT_ATENOP_PATTERN(AtenFlipOp);
  INSERT_ATENOP_PATTERN(AtenFmodTensorOp);
  INSERT_ATENOP_PATTERN(AtenBitwiseLeftShiftTensorOp);
  INSERT_ATENOP_PATTERN(AtenBitwiseRightShiftTensorOp);

  INSERT_ATENOP_PATTERN(AtenTrilOp);
  INSERT_ATENOP_PATTERN(AtenIsfiniteOp);
#undef INSERT_ATENOP_PATTERN

#define INSERT_BINARY_BROADCAST_PATTERN(AtenOp, StablehloOp)                   \
  target.addIllegalOp<AtenOp>();                                               \
  patterns.add<ConvertAtenBinaryBroadcastOp<AtenOp, StablehloOp>>(             \
      typeConverter, context)
  INSERT_BINARY_BROADCAST_PATTERN(AtenMaximumOp, chlo::BroadcastMaxOp);
  INSERT_BINARY_BROADCAST_PATTERN(AtenMinimumOp, chlo::BroadcastMinOp);
  INSERT_BINARY_BROADCAST_PATTERN(Aten__And__TensorOp, chlo::BroadcastAndOp);
  INSERT_BINARY_BROADCAST_PATTERN(AtenBitwiseAndTensorOp, chlo::BroadcastAndOp);
  INSERT_BINARY_BROADCAST_PATTERN(AtenBitwiseOrTensorOp, chlo::BroadcastOrOp);
  INSERT_BINARY_BROADCAST_PATTERN(AtenBitwiseXorTensorOp, chlo::BroadcastXorOp);
  INSERT_BINARY_BROADCAST_PATTERN(AtenAtan2Op, chlo::BroadcastAtan2Op);
#undef INSERT_BINARY_BROADCAST_PATTERN
}