ArithToLLVM.cpp

Go to the documentation of this file.

//===- ArithToLLVM.cpp - Arithmetic to LLVM dialect conversion -------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 
#include "mlir/Analysis/DataLayoutAnalysis.h"
#include "mlir/Conversion/ArithAttrToLLVM/AttrToLLVMConverter.h"
#include "mlir/Conversion/ArithCommon/CeilFloorDivExpansion.h"
#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/VectorPattern.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/DialectResourceBlobManager.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Pass/Pass.h"
#include <type_traits>
 
namespace mlir {
#define GEN_PASS_DEF_ARITHTOLLVMCONVERSIONPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
 
using namespace mlir;
 
namespace {
 
/// Lowering pattern that matches only when the source op's rounding mode
/// presence agrees with `HasRoundingMode`. This allows registering two
/// instances of the same pattern for one source op: one that handles the
/// unconstrained case (no rounding mode, lowering to a regular LLVM op) and
/// one that handles the constrained case (rounding mode present, lowering to
/// a constrained LLVM intrinsic).
///
/// * `HasRoundingMode`: the pattern matches if and only if the source op has
///   a rounding mode attribute.
/// * `AttrConvert`: attribute converter to translate source attributes to
///   target attributes.
/// * `FailOnUnsupportedFP`: whether to fail if the source op has unsupported
///   floating point types.
template <typename SourceOp, typename TargetOp, bool HasRoundingMode,
          template <typename, typename> typename AttrConvert =
              AttrConvertPassThrough,
          bool FailOnUnsupportedFP = false>
struct ConstrainedVectorConvertToLLVMPattern
    : public VectorConvertToLLVMPattern<SourceOp, TargetOp, AttrConvert,
                                        FailOnUnsupportedFP> {
  using VectorConvertToLLVMPattern<
      SourceOp, TargetOp, AttrConvert,
      FailOnUnsupportedFP>::VectorConvertToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (HasRoundingMode != static_cast<bool>(op.getRoundingModeAttr()))
      return failure();
    return VectorConvertToLLVMPattern<
        SourceOp, TargetOp, AttrConvert,
        FailOnUnsupportedFP>::matchAndRewrite(op, adaptor, rewriter);
  }
};
 
/// No-op bitcast. Propagate type input arg if converted source and dest types
/// are the same.
struct IdentityBitcastLowering final
    : public OpConversionPattern<arith::BitcastOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(arith::BitcastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    Value src = adaptor.getIn();
    Type resultType = getTypeConverter()->convertType(op.getType());
    if (src.getType() != resultType)
      return rewriter.notifyMatchFailure(op, "Types are different");
 
    rewriter.replaceOp(op, src);
    return success();
  }
};
 
//===----------------------------------------------------------------------===//
// Straightforward Op Lowerings
//===----------------------------------------------------------------------===//
 
using AddFOpLowering =
    ConstrainedVectorConvertToLLVMPattern<arith::AddFOp, LLVM::FAddOp,
                                          /*HasRoundingMode=*/false,
                                          arith::AttrConvertFastMathToLLVM,
                                          /*FailOnUnsupportedFP=*/true>;
using ConstrainedAddFOpLowering = ConstrainedVectorConvertToLLVMPattern<
    arith::AddFOp, LLVM::ConstrainedFAddIntr, /*HasRoundingMode=*/true,
    arith::AttrConverterConstrainedFPToLLVM, /*FailOnUnsupportedFP=*/true>;
using AddIOpLowering =
    VectorConvertToLLVMPattern<arith::AddIOp, LLVM::AddOp,
                               arith::AttrConvertOverflowToLLVM>;
using AndIOpLowering = VectorConvertToLLVMPattern<arith::AndIOp, LLVM::AndOp>;
using BitcastOpLowering =
    VectorConvertToLLVMPattern<arith::BitcastOp, LLVM::BitcastOp>;
using DivFOpLowering =
    ConstrainedVectorConvertToLLVMPattern<arith::DivFOp, LLVM::FDivOp,
                                          /*HasRoundingMode=*/false,
                                          arith::AttrConvertFastMathToLLVM,
                                          /*FailOnUnsupportedFP=*/true>;
using ConstrainedDivFOpLowering = ConstrainedVectorConvertToLLVMPattern<
    arith::DivFOp, LLVM::ConstrainedFDivIntr, /*HasRoundingMode=*/true,
    arith::AttrConverterConstrainedFPToLLVM, /*FailOnUnsupportedFP=*/true>;
using DivSIOpLowering =
    VectorConvertToLLVMPattern<arith::DivSIOp, LLVM::SDivOp>;
using DivUIOpLowering =
    VectorConvertToLLVMPattern<arith::DivUIOp, LLVM::UDivOp>;
using ExtFOpLowering =
    VectorConvertToLLVMPattern<arith::ExtFOp, LLVM::FPExtOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using ExtSIOpLowering =
    VectorConvertToLLVMPattern<arith::ExtSIOp, LLVM::SExtOp>;
using ExtUIOpLowering =
    VectorConvertToLLVMPattern<arith::ExtUIOp, LLVM::ZExtOp,
                               arith::AttrConvertNonNegToLLVM>;
using FPToSIOpLowering =
    VectorConvertToLLVMPattern<arith::FPToSIOp, LLVM::FPToSIOp,
                               AttrConvertPassThrough,
                               /*FailOnUnsupportedFP=*/true>;
using FPToUIOpLowering =
    VectorConvertToLLVMPattern<arith::FPToUIOp, LLVM::FPToUIOp,
                               AttrConvertPassThrough,
                               /*FailOnUnsupportedFP=*/true>;
using MaximumFOpLowering =
    VectorConvertToLLVMPattern<arith::MaximumFOp, LLVM::MaximumOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using MaxNumFOpLowering =
    VectorConvertToLLVMPattern<arith::MaxNumFOp, LLVM::MaxNumOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using MaximumNumFOpLowering =
    VectorConvertToLLVMPattern<arith::MaximumNumFOp, LLVM::MaximumNumOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using MaxSIOpLowering =
    VectorConvertToLLVMPattern<arith::MaxSIOp, LLVM::SMaxOp>;
using MaxUIOpLowering =
    VectorConvertToLLVMPattern<arith::MaxUIOp, LLVM::UMaxOp>;
using MinimumFOpLowering =
    VectorConvertToLLVMPattern<arith::MinimumFOp, LLVM::MinimumOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using MinNumFOpLowering =
    VectorConvertToLLVMPattern<arith::MinNumFOp, LLVM::MinNumOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using MinimumNumFOpLowering =
    VectorConvertToLLVMPattern<arith::MinimumNumFOp, LLVM::MinimumNumOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using MinSIOpLowering =
    VectorConvertToLLVMPattern<arith::MinSIOp, LLVM::SMinOp>;
using MinUIOpLowering =
    VectorConvertToLLVMPattern<arith::MinUIOp, LLVM::UMinOp>;
using MulFOpLowering =
    ConstrainedVectorConvertToLLVMPattern<arith::MulFOp, LLVM::FMulOp,
                                          /*HasRoundingMode=*/false,
                                          arith::AttrConvertFastMathToLLVM,
                                          /*FailOnUnsupportedFP=*/true>;
using ConstrainedMulFOpLowering = ConstrainedVectorConvertToLLVMPattern<
    arith::MulFOp, LLVM::ConstrainedFMulIntr, /*HasRoundingMode=*/true,
    arith::AttrConverterConstrainedFPToLLVM, /*FailOnUnsupportedFP=*/true>;
using MulIOpLowering =
    VectorConvertToLLVMPattern<arith::MulIOp, LLVM::MulOp,
                               arith::AttrConvertOverflowToLLVM>;
using NegFOpLowering =
    VectorConvertToLLVMPattern<arith::NegFOp, LLVM::FNegOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using OrIOpLowering = VectorConvertToLLVMPattern<arith::OrIOp, LLVM::OrOp>;
using RemFOpLowering =
    VectorConvertToLLVMPattern<arith::RemFOp, LLVM::FRemOp,
                               arith::AttrConvertFastMathToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using RemSIOpLowering =
    VectorConvertToLLVMPattern<arith::RemSIOp, LLVM::SRemOp>;
using RemUIOpLowering =
    VectorConvertToLLVMPattern<arith::RemUIOp, LLVM::URemOp>;
using SelectOpLowering =
    VectorConvertToLLVMPattern<arith::SelectOp, LLVM::SelectOp>;
using ShLIOpLowering =
    VectorConvertToLLVMPattern<arith::ShLIOp, LLVM::ShlOp,
                               arith::AttrConvertOverflowToLLVM>;
using ShRSIOpLowering =
    VectorConvertToLLVMPattern<arith::ShRSIOp, LLVM::AShrOp>;
using ShRUIOpLowering =
    VectorConvertToLLVMPattern<arith::ShRUIOp, LLVM::LShrOp>;
using SIToFPOpLowering =
    VectorConvertToLLVMPattern<arith::SIToFPOp, LLVM::SIToFPOp,
                               AttrConvertPassThrough,
                               /*FailOnUnsupportedFP=*/true>;
using SubFOpLowering =
    ConstrainedVectorConvertToLLVMPattern<arith::SubFOp, LLVM::FSubOp,
                                          /*HasRoundingMode=*/false,
                                          arith::AttrConvertFastMathToLLVM,
                                          /*FailOnUnsupportedFP=*/true>;
using ConstrainedSubFOpLowering = ConstrainedVectorConvertToLLVMPattern<
    arith::SubFOp, LLVM::ConstrainedFSubIntr, /*HasRoundingMode=*/true,
    arith::AttrConverterConstrainedFPToLLVM, /*FailOnUnsupportedFP=*/true>;
using SubIOpLowering =
    VectorConvertToLLVMPattern<arith::SubIOp, LLVM::SubOp,
                               arith::AttrConvertOverflowToLLVM>;
using TruncFOpLowering =
    ConstrainedVectorConvertToLLVMPattern<arith::TruncFOp, LLVM::FPTruncOp,
                                          /*HasRoundingMode=*/false,
                                          arith::AttrConvertFastMathToLLVM,
                                          /*FailOnUnsupportedFP=*/true>;
using ConstrainedTruncFOpLowering = ConstrainedVectorConvertToLLVMPattern<
    arith::TruncFOp, LLVM::ConstrainedFPTruncIntr, /*HasRoundingMode=*/true,
    arith::AttrConverterConstrainedFPToLLVM, /*FailOnUnsupportedFP=*/true>;
using TruncIOpLowering =
    VectorConvertToLLVMPattern<arith::TruncIOp, LLVM::TruncOp,
                               arith::AttrConvertOverflowToLLVM>;
using UIToFPOpLowering =
    VectorConvertToLLVMPattern<arith::UIToFPOp, LLVM::UIToFPOp,
                               arith::AttrConvertNonNegToLLVM,
                               /*FailOnUnsupportedFP=*/true>;
using XOrIOpLowering = VectorConvertToLLVMPattern<arith::XOrIOp, LLVM::XOrOp>;
 
//===----------------------------------------------------------------------===//
// Op Lowering Patterns
//===----------------------------------------------------------------------===//
 
/// Directly lower to LLVM op.
struct ConstantOpLowering : public ConvertOpToLLVMPattern<arith::ConstantOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// The lowering of index_cast becomes an integer conversion since index
/// becomes an integer.  If the bit width of the source and target integer
/// types is the same, just erase the cast.  If the target type is wider,
/// sign-extend the value, otherwise truncate it.
template <typename OpTy, typename ExtCastTy>
struct IndexCastOpLowering : public ConvertOpToLLVMPattern<OpTy> {
  using ConvertOpToLLVMPattern<OpTy>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
using IndexCastOpSILowering =
    IndexCastOpLowering<arith::IndexCastOp, LLVM::SExtOp>;
using IndexCastOpUILowering =
    IndexCastOpLowering<arith::IndexCastUIOp, LLVM::ZExtOp>;
 
struct AddUIExtendedOpLowering
    : public ConvertOpToLLVMPattern<arith::AddUIExtendedOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(arith::AddUIExtendedOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
struct SubUIExtendedOpLowering
    : public ConvertOpToLLVMPattern<arith::SubUIExtendedOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(arith::SubUIExtendedOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
template <typename ArithMulOp, bool IsSigned>
struct MulIExtendedOpLowering : public ConvertOpToLLVMPattern<ArithMulOp> {
  using ConvertOpToLLVMPattern<ArithMulOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(ArithMulOp op, typename ArithMulOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
using MulSIExtendedOpLowering =
    MulIExtendedOpLowering<arith::MulSIExtendedOp, true>;
using MulUIExtendedOpLowering =
    MulIExtendedOpLowering<arith::MulUIExtendedOp, false>;
 
struct CmpIOpLowering : public ConvertOpToLLVMPattern<arith::CmpIOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(arith::CmpIOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
struct CmpFOpLowering : public ConvertOpToLLVMPattern<arith::CmpFOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(arith::CmpFOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// Lower arith.convertf (same-bitwidth FP cast) to LLVM.
///
/// Extends to f32 via llvm.fpext, then truncates to the target type via
/// llvm.fptrunc. This handles bf16 <-> f16, which is the only same-bitwidth
/// pair of LLVM-supported FP types.
struct ConvertFOpLowering : public ConvertOpToLLVMPattern<arith::ConvertFOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(arith::ConvertFOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (LLVM::detail::opHasUnsupportedFloatingPointTypes(op,
                                                         *getTypeConverter()))
      return rewriter.notifyMatchFailure(op, "unsupported floating point type");
 
    // Only bf16 <-> f16 conversions are supported. There is currently no other
    // pair of FP types that are valid LLVM types.
    [[maybe_unused]] auto srcType = getElementTypeOrSelf(op.getIn().getType());
    [[maybe_unused]] auto dstType = getElementTypeOrSelf(op.getType());
    assert(((srcType.isBF16() && dstType.isF16()) ||
            (srcType.isF16() && dstType.isBF16())) &&
           "only bf16 <-> f16 conversions are supported");
 
    Type convertedType = getTypeConverter()->convertType(op.getType());
    if (!convertedType)
      return rewriter.notifyMatchFailure(op, "failed to convert result type");
 
    Value input = adaptor.getIn();
    Location loc = op.getLoc();
 
    if (!isa<LLVM::LLVMArrayType>(input.getType())) {
      rewriter.replaceOp(op,
                         emitConversion(rewriter, loc, input, convertedType));
      return success();
    }
 
    if (!isa<VectorType>(op.getType()))
      return rewriter.notifyMatchFailure(op, "expected vector result type");
 
    return LLVM::detail::handleMultidimensionalVectors(
        op.getOperation(), adaptor.getOperands(), *getTypeConverter(),
        [&](Type llvm1DVectorTy, ValueRange operands) -> Value {
          return emitConversion(rewriter, loc, operands.front(),
                                llvm1DVectorTy);
        },
        rewriter);
  }
 
private:
  static Value emitConversion(ConversionPatternRewriter &rewriter, Location loc,
                              Value input, Type targetType) {
    Type f32Scalar = Float32Type::get(rewriter.getContext());
    Type f32Ty = f32Scalar;
    if (auto vecTy = dyn_cast<VectorType>(targetType))
      f32Ty = VectorType::get(vecTy.getShape(), f32Scalar);
 
    Value ext = LLVM::FPExtOp::create(rewriter, loc, f32Ty, input);
    return LLVM::FPTruncOp::create(rewriter, loc, targetType, ext);
  }
};
 
struct SelectOpOneToNLowering : public ConvertOpToLLVMPattern<arith::SelectOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
  using Adaptor = ConvertOpToLLVMPattern<arith::SelectOp>::OneToNOpAdaptor;
 
  LogicalResult
  matchAndRewrite(arith::SelectOp op, Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
} // namespace
 
//===----------------------------------------------------------------------===//
// ConstantOpLowering
//===----------------------------------------------------------------------===//
 
/// Retypes `attr` for a `llvm.mlir.constant` of `resultType`. `arith.constant`
/// requires the value attribute and the result to have the same type, but the
/// type converter may map the element type to a different one, e.g. `index` to
/// `i32` or `i64` depending on the configured index bitwidth. Build the
/// attribute from the converted type so that the two agree wherever that is
/// representable. Returns a null attribute if `attr` cannot be used for
/// `resultType`.
static TypedAttr convertConstantValue(TypedAttr attr, Type resultType) {
  // Compare the element types, but explicitly ignore the non-scalar portions of
  // types. This relaxation is required to support multi-dimensional vector
  // constant that result in nested LLVM array values.
  Type sourceElementType = LLVM::getConstantElementType(attr.getType());
  Type targetElementType = LLVM::getConstantElementType(resultType);
  if (sourceElementType == targetElementType)
    return attr;
 
  auto targetIntType = dyn_cast<IntegerType>(targetElementType);
  if (!targetIntType)
    return {};
 
  // The converter maps the low-precision float types that have no LLVM
  // equivalent to an integer of the same width. The attribute stays a float
  // attribute in that case.
  if (auto sourceFloatType = dyn_cast<FloatType>(sourceElementType)) {
    if (sourceFloatType.getWidth() != targetIntType.getWidth())
      return {};
    return attr;
  }
 
  // Apart from those floats, `index` is the only element type the converter
  // rewrites, so anything else is a malformed `arith.constant`. Bail out rather
  // than reinterpret its value, which would silently change the constant (e.g.
  // sign-extending an `i1` `true` to -1).
  if (!isa<IndexType>(sourceElementType))
    return {};
  // `index` is signless but holds signed values, so narrowing to a smaller
  // index bitwidth truncates and widening sign-extends.
  unsigned width = targetIntType.getWidth();
 
  if (auto intAttr = dyn_cast<IntegerAttr>(attr))
    return IntegerAttr::get(targetIntType,
                            intAttr.getValue().sextOrTrunc(width));
 
  auto retypeValues = [&](DenseIntElementsAttr values) {
    return values.mapValues(targetIntType, [&](const APInt &value) {
      return value.sextOrTrunc(width);
    });
  };
 
  if (auto denseAttr = dyn_cast<DenseIntElementsAttr>(attr))
    return retypeValues(denseAttr);
 
  if (auto sparseAttr = dyn_cast<SparseElementsAttr>(attr))
    return SparseElementsAttr::get(
        cast<ShapedType>(attr.getType()).clone(targetIntType),
        sparseAttr.getIndices(),
        retypeValues(cast<DenseIntElementsAttr>(sparseAttr.getValues())));
 
  // A resource-backed elements attribute refers to a blob laid out for its own
  // element type. The blob cannot be rewritten here, only reinterpreted, which
  // is correct exactly when the target type has the same width as the storage
  // `index` uses in a blob.
  if (auto resourceAttr = dyn_cast<DenseResourceElementsAttr>(attr)) {
    if (width != IndexType::kInternalStorageBitWidth)
      return {};
    return DenseResourceElementsAttr::get(
        cast<ShapedType>(attr.getType()).clone(targetIntType),
        resourceAttr.getRawHandle());
  }
 
  return {};
}
 
LogicalResult
ConstantOpLowering::matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor,
                                    ConversionPatternRewriter &rewriter) const {
  Type resultType = getTypeConverter()->convertType(op.getType());
  if (!resultType)
    return rewriter.notifyMatchFailure(op, "failed to convert result type");
 
  TypedAttr value = convertConstantValue(op.getValue(), resultType);
  if (!value)
    return rewriter.notifyMatchFailure(
        op, "failed to convert value attribute to the converted result type");
 
  // `arith.constant` has no operands and a single result, so there is nothing
  // for `oneToOneRewrite` to do here beyond converting `resultType` a second
  // time.
  DictionaryAttr discardableAttrs = op->getDiscardableAttrDictionary();
  auto constantOp =
      LLVM::ConstantOp::create(rewriter, op.getLoc(), resultType, value);
  constantOp->setDiscardableAttrs(discardableAttrs);
  rewriter.replaceOp(op, constantOp);
  return success();
}
 
//===----------------------------------------------------------------------===//
// IndexCastOpLowering
//===----------------------------------------------------------------------===//
 
template <typename OpTy, typename ExtCastTy>
LogicalResult IndexCastOpLowering<OpTy, ExtCastTy>::matchAndRewrite(
    OpTy op, typename OpTy::Adaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Type resultType = op.getResult().getType();
  Type targetElementType =
      this->typeConverter->convertType(getElementTypeOrSelf(resultType));
  Type sourceElementType =
      this->typeConverter->convertType(getElementTypeOrSelf(op.getIn()));
  unsigned targetBits = targetElementType.getIntOrFloatBitWidth();
  unsigned sourceBits = sourceElementType.getIntOrFloatBitWidth();
 
  if (targetBits == sourceBits) {
    rewriter.replaceOp(op, adaptor.getIn());
    return success();
  }
 
  // Memref index_cast is a no-op at the LLVM level since LLVM uses opaque
  // pointers and memrefs of different integer/index element types all convert
  // to the same LLVM struct type.
  if (isa<MemRefType>(op.getIn().getType())) {
    rewriter.replaceOp(op, adaptor.getIn());
    return success();
  }
 
  bool isNonNeg = false;
  if constexpr (std::is_same_v<ExtCastTy, LLVM::ZExtOp>)
    isNonNeg = op.getNonNeg();
 
  // Handle the scalar and 1D vector cases.
  Type operandType = adaptor.getIn().getType();
  if (!isa<LLVM::LLVMArrayType>(operandType)) {
    Type targetType = this->typeConverter->convertType(resultType);
    if (targetBits < sourceBits) {
      rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, targetType,
                                                 adaptor.getIn());
    } else {
      auto extOp = rewriter.replaceOpWithNewOp<ExtCastTy>(op, targetType,
                                                          adaptor.getIn());
      if constexpr (std::is_same_v<ExtCastTy, LLVM::ZExtOp>)
        extOp.setNonNeg(isNonNeg);
    }
    return success();
  }
 
  if (!isa<VectorType>(resultType))
    return rewriter.notifyMatchFailure(op, "expected vector result type");
 
  return LLVM::detail::handleMultidimensionalVectors(
      op.getOperation(), adaptor.getOperands(), *(this->getTypeConverter()),
      [&](Type llvm1DVectorTy, ValueRange operands) -> Value {
        typename OpTy::Adaptor adaptor(operands);
        if (targetBits < sourceBits) {
          return LLVM::TruncOp::create(rewriter, op.getLoc(), llvm1DVectorTy,
                                       adaptor.getIn());
        }
        auto extOp = ExtCastTy::create(rewriter, op.getLoc(), llvm1DVectorTy,
                                       adaptor.getIn());
        if constexpr (std::is_same_v<ExtCastTy, LLVM::ZExtOp>) {
          if (isNonNeg)
            extOp.setNonNeg(true);
        }
        return extOp;
      },
      rewriter);
}
 
//===----------------------------------------------------------------------===//
// AddUIExtendedOpLowering
//===----------------------------------------------------------------------===//
 
LogicalResult AddUIExtendedOpLowering::matchAndRewrite(
    arith::AddUIExtendedOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Type operandType = adaptor.getLhs().getType();
  Type sumResultType = op.getSum().getType();
  Type overflowResultType = op.getOverflow().getType();
 
  if (!LLVM::isCompatibleType(operandType))
    return failure();
 
  MLIRContext *ctx = rewriter.getContext();
  Location loc = op.getLoc();
 
  // Handle the scalar and 1D vector cases.
  if (!isa<LLVM::LLVMArrayType>(operandType)) {
    Type newSumType = typeConverter->convertType(sumResultType);
    Type newOverflowType = typeConverter->convertType(overflowResultType);
    Type structType =
        LLVM::LLVMStructType::getLiteral(ctx, {newSumType, newOverflowType});
    Value addOverflow = LLVM::UAddWithOverflowOp::create(
        rewriter, loc, structType, adaptor.getLhs(), adaptor.getRhs());
    Value sumExtracted =
        LLVM::ExtractValueOp::create(rewriter, loc, addOverflow, 0);
    Value overflowExtracted =
        LLVM::ExtractValueOp::create(rewriter, loc, addOverflow, 1);
    rewriter.replaceOp(op, {sumExtracted, overflowExtracted});
    return success();
  }
 
  if (!isa<VectorType>(sumResultType))
    return rewriter.notifyMatchFailure(loc, "expected vector result types");
 
  return rewriter.notifyMatchFailure(loc,
                                     "ND vector types are not supported yet");
}
 
//===----------------------------------------------------------------------===//
// SubUIExtendedOpLowering
//===----------------------------------------------------------------------===//
 
LogicalResult SubUIExtendedOpLowering::matchAndRewrite(
    arith::SubUIExtendedOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Type operandType = adaptor.getLhs().getType();
  Type diffResultType = op.getDiff().getType();
  Type borrowResultType = op.getBorrow().getType();
 
  if (!LLVM::isCompatibleType(operandType))
    return failure();
 
  MLIRContext *ctx = rewriter.getContext();
  Location loc = op.getLoc();
 
  // Handle the scalar and 1D vector cases.
  if (!isa<LLVM::LLVMArrayType>(operandType)) {
    Type newBorrowType = typeConverter->convertType(borrowResultType);
    Type structType =
        LLVM::LLVMStructType::getLiteral(ctx, {diffResultType, newBorrowType});
    Value subOverflow = LLVM::USubWithOverflowOp::create(
        rewriter, loc, structType, adaptor.getLhs(), adaptor.getRhs());
    Value diffExtracted =
        LLVM::ExtractValueOp::create(rewriter, loc, subOverflow, 0);
    Value borrowExtracted =
        LLVM::ExtractValueOp::create(rewriter, loc, subOverflow, 1);
    rewriter.replaceOp(op, {diffExtracted, borrowExtracted});
    return success();
  }
 
  if (!isa<VectorType>(diffResultType))
    return rewriter.notifyMatchFailure(loc, "expected vector result types");
 
  return rewriter.notifyMatchFailure(loc,
                                     "ND vector types are not supported yet");
}
 
//===----------------------------------------------------------------------===//
// MulIExtendedOpLowering
//===----------------------------------------------------------------------===//
 
template <typename ArithMulOp, bool IsSigned>
LogicalResult MulIExtendedOpLowering<ArithMulOp, IsSigned>::matchAndRewrite(
    ArithMulOp op, typename ArithMulOp::Adaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Type resultType = adaptor.getLhs().getType();
 
  if (!LLVM::isCompatibleType(resultType))
    return failure();
 
  Location loc = op.getLoc();
 
  // Handle the scalar and 1D vector cases. Because LLVM does not have a
  // matching extended multiplication intrinsic, perform regular multiplication
  // on operands zero-extended to i(2*N) bits, and truncate the results back to
  // iN types.
  if (!isa<LLVM::LLVMArrayType>(resultType)) {
    // Shift amount necessary to extract the high bits from widened result.
    TypedAttr shiftValAttr;
 
    if (auto intTy = dyn_cast<IntegerType>(resultType)) {
      unsigned resultBitwidth = intTy.getWidth();
      auto attrTy = rewriter.getIntegerType(resultBitwidth * 2);
      shiftValAttr = rewriter.getIntegerAttr(attrTy, resultBitwidth);
    } else {
      auto vecTy = cast<VectorType>(resultType);
      unsigned resultBitwidth = vecTy.getElementTypeBitWidth();
      auto attrTy = VectorType::get(
          vecTy.getShape(), rewriter.getIntegerType(resultBitwidth * 2));
      shiftValAttr = SplatElementsAttr::get(
          attrTy, APInt(resultBitwidth * 2, resultBitwidth));
    }
    Type wideType = shiftValAttr.getType();
    assert(LLVM::isCompatibleType(wideType) &&
           "LLVM dialect should support all signless integer types");
 
    using LLVMExtOp = std::conditional_t<IsSigned, LLVM::SExtOp, LLVM::ZExtOp>;
    Value lhsExt = LLVMExtOp::create(rewriter, loc, wideType, adaptor.getLhs());
    Value rhsExt = LLVMExtOp::create(rewriter, loc, wideType, adaptor.getRhs());
    Value mulExt = LLVM::MulOp::create(rewriter, loc, wideType, lhsExt, rhsExt);
 
    // Split the 2*N-bit wide result into two N-bit values.
    Value low = LLVM::TruncOp::create(rewriter, loc, resultType, mulExt);
    Value shiftVal = LLVM::ConstantOp::create(rewriter, loc, shiftValAttr);
    Value highExt = LLVM::LShrOp::create(rewriter, loc, mulExt, shiftVal);
    Value high = LLVM::TruncOp::create(rewriter, loc, resultType, highExt);
 
    rewriter.replaceOp(op, {low, high});
    return success();
  }
 
  if (!isa<VectorType>(resultType))
    return rewriter.notifyMatchFailure(op, "expected vector result type");
 
  return rewriter.notifyMatchFailure(op,
                                     "ND vector types are not supported yet");
}
 
//===----------------------------------------------------------------------===//
// CmpIOpLowering
//===----------------------------------------------------------------------===//
 
// Convert arith.cmp predicate into the LLVM dialect CmpPredicate. The two enums
// share numerical values so just cast.
template <typename LLVMPredType, typename PredType>
static LLVMPredType convertCmpPredicate(PredType pred) {
  return static_cast<LLVMPredType>(pred);
}
 
LogicalResult
CmpIOpLowering::matchAndRewrite(arith::CmpIOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
  Type operandType = adaptor.getLhs().getType();
  Type resultType = op.getResult().getType();
 
  // Handle the scalar and 1D vector cases.
  if (!isa<LLVM::LLVMArrayType>(operandType)) {
    rewriter.replaceOpWithNewOp<LLVM::ICmpOp>(
        op, typeConverter->convertType(resultType),
        convertCmpPredicate<LLVM::ICmpPredicate>(op.getPredicate()),
        adaptor.getLhs(), adaptor.getRhs());
    return success();
  }
 
  if (!isa<VectorType>(resultType))
    return rewriter.notifyMatchFailure(op, "expected vector result type");
 
  return LLVM::detail::handleMultidimensionalVectors(
      op.getOperation(), adaptor.getOperands(), *getTypeConverter(),
      [&](Type llvm1DVectorTy, ValueRange operands) {
        OpAdaptor adaptor(operands);
        return LLVM::ICmpOp::create(
            rewriter, op.getLoc(), llvm1DVectorTy,
            convertCmpPredicate<LLVM::ICmpPredicate>(op.getPredicate()),
            adaptor.getLhs(), adaptor.getRhs());
      },
      rewriter);
}
 
//===----------------------------------------------------------------------===//
// CmpFOpLowering
//===----------------------------------------------------------------------===//
 
LogicalResult
CmpFOpLowering::matchAndRewrite(arith::CmpFOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
  if (LLVM::detail::isUnsupportedFloatingPointType(*this->getTypeConverter(),
                                                   op.getLhs().getType()))
    return rewriter.notifyMatchFailure(op, "unsupported floating point type");
 
  Type operandType = adaptor.getLhs().getType();
  Type resultType = op.getResult().getType();
  LLVM::FastmathFlags fmf =
      arith::convertArithFastMathFlagsToLLVM(op.getFastmath());
 
  // Handle the scalar and 1D vector cases.
  if (!isa<LLVM::LLVMArrayType>(operandType)) {
    rewriter.replaceOpWithNewOp<LLVM::FCmpOp>(
        op, typeConverter->convertType(resultType),
        convertCmpPredicate<LLVM::FCmpPredicate>(op.getPredicate()),
        adaptor.getLhs(), adaptor.getRhs(), fmf);
    return success();
  }
 
  if (!isa<VectorType>(resultType))
    return rewriter.notifyMatchFailure(op, "expected vector result type");
 
  return LLVM::detail::handleMultidimensionalVectors(
      op.getOperation(), adaptor.getOperands(), *getTypeConverter(),
      [&](Type llvm1DVectorTy, ValueRange operands) {
        OpAdaptor adaptor(operands);
        return LLVM::FCmpOp::create(
            rewriter, op.getLoc(), llvm1DVectorTy,
            convertCmpPredicate<LLVM::FCmpPredicate>(op.getPredicate()),
            adaptor.getLhs(), adaptor.getRhs(), fmf);
      },
      rewriter);
}
 
//===----------------------------------------------------------------------===//
// SelectOpOneToNLowering
//===----------------------------------------------------------------------===//
 
/// Pattern for arith.select where the true/false values lower to multiple
/// SSA values (1:N conversion). This pattern generates multiple arith.select
/// than can be lowered by the 1:1 arith.select pattern.
LogicalResult SelectOpOneToNLowering::matchAndRewrite(
    arith::SelectOp op, Adaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  // In case of a 1:1 conversion, the 1:1 pattern will match.
  if (llvm::hasSingleElement(adaptor.getTrueValue()))
    return rewriter.notifyMatchFailure(
        op, "not a 1:N conversion, 1:1 pattern will match");
  if (!op.getCondition().getType().isInteger(1))
    return rewriter.notifyMatchFailure(op,
                                       "non-i1 conditions are not supported");
  SmallVector<Value> results;
  for (auto [trueValue, falseValue] :
       llvm::zip_equal(adaptor.getTrueValue(), adaptor.getFalseValue()))
    results.push_back(arith::SelectOp::create(
        rewriter, op.getLoc(), op.getCondition(), trueValue, falseValue));
  rewriter.replaceOpWithMultiple(op, {results});
  return success();
}
 
//===----------------------------------------------------------------------===//
// Pass Definition
//===----------------------------------------------------------------------===//
 
namespace {
struct ArithToLLVMConversionPass
    : public impl::ArithToLLVMConversionPassBase<ArithToLLVMConversionPass> {
  using Base::Base;
 
  void runOnOperation() override {
    LLVMConversionTarget target(getContext());
    RewritePatternSet patterns(&getContext());
 
    const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
    LowerToLLVMOptions options(&getContext(),
                               dataLayoutAnalysis.getAtOrAbove(getOperation()));
    if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
      options.overrideIndexBitwidth(indexBitwidth);
 
    LLVMTypeConverter converter(&getContext(), options, &dataLayoutAnalysis);
    arith::populateCeilFloorDivExpandOpsPatterns(patterns);
    arith::populateArithToLLVMConversionPatterns(converter, patterns);
 
    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns))))
      signalPassFailure();
  }
};
} // namespace
 
//===----------------------------------------------------------------------===//
// ConvertToLLVMPatternInterface implementation
//===----------------------------------------------------------------------===//
 
namespace {
/// Implement the interface to convert MemRef to LLVM.
struct ArithToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
  ArithToLLVMDialectInterface(Dialect *dialect)
      : ConvertToLLVMPatternInterface(dialect) {}
 
  void loadDependentDialects(MLIRContext *context) const final {
    context->loadDialect<LLVM::LLVMDialect>();
  }
 
  /// Hook for derived dialect interface to provide conversion patterns
  /// and mark dialect legal for the conversion target.
  void populateConvertToLLVMConversionPatterns(
      ConversionTarget &target, LLVMTypeConverter &typeConverter,
      RewritePatternSet &patterns) const final {
    arith::populateCeilFloorDivExpandOpsPatterns(patterns);
    arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);
  }
};
} // namespace
 
void mlir::arith::registerConvertArithToLLVMInterface(
    DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
    dialect->addInterfaces<ArithToLLVMDialectInterface>();
  });
}
 
//===----------------------------------------------------------------------===//
// Pattern Population
//===----------------------------------------------------------------------===//
 
void mlir::arith::populateArithToLLVMConversionPatterns(
    const LLVMTypeConverter &converter, RewritePatternSet &patterns) {
 
  // Set a higher pattern benefit for IdentityBitcastLowering so it will run
  // before BitcastOpLowering.
  patterns.add<IdentityBitcastLowering>(converter, patterns.getContext(),
                                        /*patternBenefit*/ 10);
 
  // clang-format off
  patterns.add<
    AddFOpLowering,
    ConstrainedAddFOpLowering,
    AddIOpLowering,
    AndIOpLowering,
    AddUIExtendedOpLowering,
    SubUIExtendedOpLowering,
    BitcastOpLowering,
    ConstantOpLowering,
    CmpFOpLowering,
    CmpIOpLowering,
    DivFOpLowering,
    ConstrainedDivFOpLowering,
    DivSIOpLowering,
    DivUIOpLowering,
    ExtFOpLowering,
    ExtSIOpLowering,
    ExtUIOpLowering,
    ConvertFOpLowering,
    FPToSIOpLowering,
    FPToUIOpLowering,
    IndexCastOpSILowering,
    IndexCastOpUILowering,
    MaximumFOpLowering,
    MaxNumFOpLowering,
    MaximumNumFOpLowering,
    MaxSIOpLowering,
    MaxUIOpLowering,
    MinimumFOpLowering,
    MinNumFOpLowering,
    MinimumNumFOpLowering,
    MinSIOpLowering,
    MinUIOpLowering,
    MulFOpLowering,
    ConstrainedMulFOpLowering,
    MulIOpLowering,
    MulSIExtendedOpLowering,
    MulUIExtendedOpLowering,
    NegFOpLowering,
    OrIOpLowering,
    RemFOpLowering,
    RemSIOpLowering,
    RemUIOpLowering,
    SelectOpLowering,
    SelectOpOneToNLowering,
    ShLIOpLowering,
    ShRSIOpLowering,
    ShRUIOpLowering,
    SIToFPOpLowering,
    SubFOpLowering,
    ConstrainedSubFOpLowering,
    SubIOpLowering,
    TruncFOpLowering,
    ConstrainedTruncFOpLowering,
    TruncIOpLowering,
    UIToFPOpLowering,
    XOrIOpLowering
  >(converter);
  // clang-format on
}