//===- MemRefToLLVM.cpp - MemRef 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/MemRefToLLVM/MemRefToLLVM.h"

#include "mlir/Analysis/DataLayoutAnalysis.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Conversion/MemRefToLLVM/AllocLikeConversion.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/MathExtras.h"
#include "llvm/ADT/SmallBitVector.h"
#include <optional>

namespace mlir {
#define GEN_PASS_DEF_FINALIZEMEMREFTOLLVMCONVERSIONPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;

namespace {

bool isStaticStrideOrOffset(int64_t strideOrOffset) {
  return !ShapedType::isDynamic(strideOrOffset);
}

LLVM::LLVMFuncOp getFreeFn(LLVMTypeConverter *typeConverter, ModuleOp module) {
  bool useGenericFn = typeConverter->getOptions().useGenericFunctions;

  if (useGenericFn)
    return LLVM::lookupOrCreateGenericFreeFn(
        module, typeConverter->useOpaquePointers());

  return LLVM::lookupOrCreateFreeFn(module, typeConverter->useOpaquePointers());
}

struct AllocOpLowering : public AllocLikeOpLLVMLowering {
  AllocOpLowering(LLVMTypeConverter &converter)
      : AllocLikeOpLLVMLowering(memref::AllocOp::getOperationName(),
                                converter) {}
  std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                          Location loc, Value sizeBytes,
                                          Operation *op) const override {
    return allocateBufferManuallyAlign(
        rewriter, loc, sizeBytes, op,
        getAlignment(rewriter, loc, cast<memref::AllocOp>(op)));
  }
};

struct AlignedAllocOpLowering : public AllocLikeOpLLVMLowering {
  AlignedAllocOpLowering(LLVMTypeConverter &converter)
      : AllocLikeOpLLVMLowering(memref::AllocOp::getOperationName(),
                                converter) {}
  std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                          Location loc, Value sizeBytes,
                                          Operation *op) const override {
    Value ptr = allocateBufferAutoAlign(
        rewriter, loc, sizeBytes, op, &defaultLayout,
        alignedAllocationGetAlignment(rewriter, loc, cast<memref::AllocOp>(op),
                                      &defaultLayout));
    return std::make_tuple(ptr, ptr);
  }

private:
  /// Default layout to use in absence of the corresponding analysis.
  DataLayout defaultLayout;
};

struct AllocaOpLowering : public AllocLikeOpLLVMLowering {
  AllocaOpLowering(LLVMTypeConverter &converter)
      : AllocLikeOpLLVMLowering(memref::AllocaOp::getOperationName(),
                                converter) {
    setRequiresNumElements();
  }

  /// Allocates the underlying buffer using the right call. `allocatedBytePtr`
  /// is set to null for stack allocations. `accessAlignment` is set if
  /// alignment is needed post allocation (for eg. in conjunction with malloc).
  std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                          Location loc, Value size,
                                          Operation *op) const override {

    // With alloca, one gets a pointer to the element type right away.
    // For stack allocations.
    auto allocaOp = cast<memref::AllocaOp>(op);
    auto elementType =
        typeConverter->convertType(allocaOp.getType().getElementType());
    unsigned addrSpace =
        *getTypeConverter()->getMemRefAddressSpace(allocaOp.getType());
    auto elementPtrType =
        getTypeConverter()->getPointerType(elementType, addrSpace);

    auto allocatedElementPtr =
        rewriter.create<LLVM::AllocaOp>(loc, elementPtrType, elementType, size,
                                        allocaOp.getAlignment().value_or(0));

    return std::make_tuple(allocatedElementPtr, allocatedElementPtr);
  }
};

/// The base class for lowering realloc op, to support the implementation of
/// realloc via allocation methods that may or may not support alignment.
/// A derived class should provide an implementation of allocateBuffer using
/// the underline allocation methods.
struct ReallocOpLoweringBase : public AllocationOpLLVMLowering {
  using OpAdaptor = typename memref::ReallocOp::Adaptor;

  ReallocOpLoweringBase(LLVMTypeConverter &converter)
      : AllocationOpLLVMLowering(memref::ReallocOp::getOperationName(),
                                 converter) {}

  /// Allocates the new buffer. Returns the allocated pointer and the
  /// aligned pointer.
  virtual std::tuple<Value, Value>
  allocateBuffer(ConversionPatternRewriter &rewriter, Location loc,
                 Value sizeBytes, memref::ReallocOp op) const = 0;

  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const final {
    auto reallocOp = cast<memref::ReallocOp>(op);
    return matchAndRewrite(reallocOp,
                           OpAdaptor(operands,
                                     op->getDiscardableAttrDictionary(),
                                     reallocOp.getProperties()),
                           rewriter);
  }

  // A `realloc` is converted as follows:
  //   If new_size > old_size
  //     1. allocates a new buffer
  //     2. copies the content of the old buffer to the new buffer
  //     3. release the old buffer
  //     3. updates the buffer pointers in the memref descriptor
  //   Update the size in the memref descriptor
  // Alignment request is handled by allocating `alignment` more bytes than
  // requested and shifting the aligned pointer relative to the allocated
  // memory.
  LogicalResult matchAndRewrite(memref::ReallocOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
    OpBuilder::InsertionGuard guard(rewriter);
    Location loc = op.getLoc();

    auto computeNumElements =
        [&](MemRefType type, function_ref<Value()> getDynamicSize) -> Value {
      // Compute number of elements.
      Value numElements =
          type.isDynamicDim(0)
              ? getDynamicSize()
              : createIndexConstant(rewriter, loc, type.getDimSize(0));
      Type indexType = getIndexType();
      if (numElements.getType() != indexType)
        numElements = typeConverter->materializeTargetConversion(
            rewriter, loc, indexType, numElements);
      return numElements;
    };

    MemRefDescriptor desc(adaptor.getSource());
    Value oldDesc = desc;

    // Split the block right before the current op into two blocks.
    Block *currentBlock = rewriter.getInsertionBlock();
    Block *block =
        rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
    // Add a block argument by creating an empty block with the argument type
    // and then merging the block into the empty block.
    Block *endBlock = rewriter.createBlock(
        block->getParent(), Region::iterator(block), oldDesc.getType(), loc);
    rewriter.mergeBlocks(block, endBlock, {});
    // Add a new block for the true branch of the conditional statement we will
    // add.
    Block *trueBlock = rewriter.createBlock(
        currentBlock->getParent(), std::next(Region::iterator(currentBlock)));

    rewriter.setInsertionPointToEnd(currentBlock);
    Value src = op.getSource();
    auto srcType = dyn_cast<MemRefType>(src.getType());
    Value srcNumElements = computeNumElements(
        srcType, [&]() -> Value { return desc.size(rewriter, loc, 0); });
    auto dstType = cast<MemRefType>(op.getType());
    Value dstNumElements = computeNumElements(
        dstType, [&]() -> Value { return op.getDynamicResultSize(); });
    Value cond = rewriter.create<LLVM::ICmpOp>(
        loc, IntegerType::get(rewriter.getContext(), 1),
        LLVM::ICmpPredicate::ugt, dstNumElements, srcNumElements);
    rewriter.create<LLVM::CondBrOp>(loc, cond, trueBlock, ArrayRef<Value>(),
                                    endBlock, ValueRange{oldDesc});

    rewriter.setInsertionPointToStart(trueBlock);
    Value sizeInBytes = getSizeInBytes(loc, dstType.getElementType(), rewriter);
    // Compute total byte size.
    auto dstByteSize =
        rewriter.create<LLVM::MulOp>(loc, dstNumElements, sizeInBytes);
    // Since the src and dst memref are guarantee to have the same
    // element type by the verifier, it is safe here to reuse the
    // type size computed from dst memref.
    auto srcByteSize =
        rewriter.create<LLVM::MulOp>(loc, srcNumElements, sizeInBytes);
    // Allocate a new buffer.
    auto [dstRawPtr, dstAlignedPtr] =
        allocateBuffer(rewriter, loc, dstByteSize, op);
    // Copy the data from the old buffer to the new buffer.
    Value srcAlignedPtr = desc.alignedPtr(rewriter, loc);
    auto toVoidPtr = [&](Value ptr) -> Value {
      if (getTypeConverter()->useOpaquePointers())
        return ptr;
      return rewriter.create<LLVM::BitcastOp>(loc, getVoidPtrType(), ptr);
    };
    rewriter.create<LLVM::MemcpyOp>(loc, toVoidPtr(dstAlignedPtr),
                                    toVoidPtr(srcAlignedPtr), srcByteSize,
                                    /*isVolatile=*/false);
    // Deallocate the old buffer.
    LLVM::LLVMFuncOp freeFunc =
        getFreeFn(getTypeConverter(), op->getParentOfType<ModuleOp>());
    rewriter.create<LLVM::CallOp>(loc, freeFunc,
                                  toVoidPtr(desc.allocatedPtr(rewriter, loc)));
    // Replace the old buffer addresses in the MemRefDescriptor with the new
    // buffer addresses.
    desc.setAllocatedPtr(rewriter, loc, dstRawPtr);
    desc.setAlignedPtr(rewriter, loc, dstAlignedPtr);
    rewriter.create<LLVM::BrOp>(loc, Value(desc), endBlock);

    rewriter.setInsertionPoint(op);
    // Update the memref size.
    MemRefDescriptor newDesc(endBlock->getArgument(0));
    newDesc.setSize(rewriter, loc, 0, dstNumElements);
    rewriter.replaceOp(op, {newDesc});
    return success();
  }

private:
  using ConvertToLLVMPattern::matchAndRewrite;
};

struct ReallocOpLowering : public ReallocOpLoweringBase {
  ReallocOpLowering(LLVMTypeConverter &converter)
      : ReallocOpLoweringBase(converter) {}
  std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                          Location loc, Value sizeBytes,
                                          memref::ReallocOp op) const override {
    return allocateBufferManuallyAlign(rewriter, loc, sizeBytes, op,
                                       getAlignment(rewriter, loc, op));
  }
};

struct AlignedReallocOpLowering : public ReallocOpLoweringBase {
  AlignedReallocOpLowering(LLVMTypeConverter &converter)
      : ReallocOpLoweringBase(converter) {}
  std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                          Location loc, Value sizeBytes,
                                          memref::ReallocOp op) const override {
    Value ptr = allocateBufferAutoAlign(
        rewriter, loc, sizeBytes, op, &defaultLayout,
        alignedAllocationGetAlignment(rewriter, loc, op, &defaultLayout));
    return std::make_tuple(ptr, ptr);
  }

private:
  /// Default layout to use in absence of the corresponding analysis.
  DataLayout defaultLayout;
};

struct AllocaScopeOpLowering
    : public ConvertOpToLLVMPattern<memref::AllocaScopeOp> {
  using ConvertOpToLLVMPattern<memref::AllocaScopeOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::AllocaScopeOp allocaScopeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    OpBuilder::InsertionGuard guard(rewriter);
    Location loc = allocaScopeOp.getLoc();

    // Split the current block before the AllocaScopeOp to create the inlining
    // point.
    auto *currentBlock = rewriter.getInsertionBlock();
    auto *remainingOpsBlock =
        rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
    Block *continueBlock;
    if (allocaScopeOp.getNumResults() == 0) {
      continueBlock = remainingOpsBlock;
    } else {
      continueBlock = rewriter.createBlock(
          remainingOpsBlock, allocaScopeOp.getResultTypes(),
          SmallVector<Location>(allocaScopeOp->getNumResults(),
                                allocaScopeOp.getLoc()));
      rewriter.create<LLVM::BrOp>(loc, ValueRange(), remainingOpsBlock);
    }

    // Inline body region.
    Block *beforeBody = &allocaScopeOp.getBodyRegion().front();
    Block *afterBody = &allocaScopeOp.getBodyRegion().back();
    rewriter.inlineRegionBefore(allocaScopeOp.getBodyRegion(), continueBlock);

    // Save stack and then branch into the body of the region.
    rewriter.setInsertionPointToEnd(currentBlock);
    auto stackSaveOp =
        rewriter.create<LLVM::StackSaveOp>(loc, getVoidPtrType());
    rewriter.create<LLVM::BrOp>(loc, ValueRange(), beforeBody);

    // Replace the alloca_scope return with a branch that jumps out of the body.
    // Stack restore before leaving the body region.
    rewriter.setInsertionPointToEnd(afterBody);
    auto returnOp =
        cast<memref::AllocaScopeReturnOp>(afterBody->getTerminator());
    auto branchOp = rewriter.replaceOpWithNewOp<LLVM::BrOp>(
        returnOp, returnOp.getResults(), continueBlock);

    // Insert stack restore before jumping out the body of the region.
    rewriter.setInsertionPoint(branchOp);
    rewriter.create<LLVM::StackRestoreOp>(loc, stackSaveOp);

    // Replace the op with values return from the body region.
    rewriter.replaceOp(allocaScopeOp, continueBlock->getArguments());

    return success();
  }
};

struct AssumeAlignmentOpLowering
    : public ConvertOpToLLVMPattern<memref::AssumeAlignmentOp> {
  using ConvertOpToLLVMPattern<
      memref::AssumeAlignmentOp>::ConvertOpToLLVMPattern;
  explicit AssumeAlignmentOpLowering(LLVMTypeConverter &converter)
      : ConvertOpToLLVMPattern<memref::AssumeAlignmentOp>(converter) {}

  LogicalResult
  matchAndRewrite(memref::AssumeAlignmentOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value memref = adaptor.getMemref();
    unsigned alignment = op.getAlignment();
    auto loc = op.getLoc();

    auto srcMemRefType = cast<MemRefType>(op.getMemref().getType());
    Value ptr = getStridedElementPtr(loc, srcMemRefType, memref, /*indices=*/{},
                                     rewriter);

    // Emit llvm.assume(memref & (alignment - 1) == 0).
    //
    // This relies on LLVM's CSE optimization (potentially after SROA), since
    // after CSE all memref instances should get de-duplicated into the same
    // pointer SSA value.
    MemRefDescriptor memRefDescriptor(memref);
    auto intPtrType =
        getIntPtrType(memRefDescriptor.getElementPtrType().getAddressSpace());
    Value zero = createIndexAttrConstant(rewriter, loc, intPtrType, 0);
    Value mask =
        createIndexAttrConstant(rewriter, loc, intPtrType, alignment - 1);
    Value ptrValue = rewriter.create<LLVM::PtrToIntOp>(loc, intPtrType, ptr);
    rewriter.create<LLVM::AssumeOp>(
        loc, rewriter.create<LLVM::ICmpOp>(
                 loc, LLVM::ICmpPredicate::eq,
                 rewriter.create<LLVM::AndOp>(loc, ptrValue, mask), zero));

    rewriter.eraseOp(op);
    return success();
  }
};

// A `dealloc` is converted into a call to `free` on the underlying data buffer.
// The memref descriptor being an SSA value, there is no need to clean it up
// in any way.
struct DeallocOpLowering : public ConvertOpToLLVMPattern<memref::DeallocOp> {
  using ConvertOpToLLVMPattern<memref::DeallocOp>::ConvertOpToLLVMPattern;

  explicit DeallocOpLowering(LLVMTypeConverter &converter)
      : ConvertOpToLLVMPattern<memref::DeallocOp>(converter) {}

  LogicalResult
  matchAndRewrite(memref::DeallocOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Insert the `free` declaration if it is not already present.
    LLVM::LLVMFuncOp freeFunc =
        getFreeFn(getTypeConverter(), op->getParentOfType<ModuleOp>());
    Value allocatedPtr;
    if (auto unrankedTy =
            llvm::dyn_cast<UnrankedMemRefType>(op.getMemref().getType())) {
      Type elementType = unrankedTy.getElementType();
      Type llvmElementTy = getTypeConverter()->convertType(elementType);
      LLVM::LLVMPointerType elementPtrTy = getTypeConverter()->getPointerType(
          llvmElementTy, unrankedTy.getMemorySpaceAsInt());
      allocatedPtr = UnrankedMemRefDescriptor::allocatedPtr(
          rewriter, op.getLoc(),
          UnrankedMemRefDescriptor(adaptor.getMemref())
              .memRefDescPtr(rewriter, op.getLoc()),
          elementPtrTy);
    } else {
      allocatedPtr = MemRefDescriptor(adaptor.getMemref())
                         .allocatedPtr(rewriter, op.getLoc());
    }
    if (!getTypeConverter()->useOpaquePointers())
      allocatedPtr = rewriter.create<LLVM::BitcastOp>(
          op.getLoc(), getVoidPtrType(), allocatedPtr);

    rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, freeFunc, allocatedPtr);
    return success();
  }
};

// A `dim` is converted to a constant for static sizes and to an access to the
// size stored in the memref descriptor for dynamic sizes.
struct DimOpLowering : public ConvertOpToLLVMPattern<memref::DimOp> {
  using ConvertOpToLLVMPattern<memref::DimOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::DimOp dimOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Type operandType = dimOp.getSource().getType();
    if (isa<UnrankedMemRefType>(operandType)) {
      FailureOr<Value> extractedSize = extractSizeOfUnrankedMemRef(
          operandType, dimOp, adaptor.getOperands(), rewriter);
      if (failed(extractedSize))
        return failure();
      rewriter.replaceOp(dimOp, {*extractedSize});
      return success();
    }
    if (isa<MemRefType>(operandType)) {
      rewriter.replaceOp(
          dimOp, {extractSizeOfRankedMemRef(operandType, dimOp,
                                            adaptor.getOperands(), rewriter)});
      return success();
    }
    llvm_unreachable("expected MemRefType or UnrankedMemRefType");
  }

private:
  FailureOr<Value>
  extractSizeOfUnrankedMemRef(Type operandType, memref::DimOp dimOp,
                              OpAdaptor adaptor,
                              ConversionPatternRewriter &rewriter) const {
    Location loc = dimOp.getLoc();

    auto unrankedMemRefType = cast<UnrankedMemRefType>(operandType);
    auto scalarMemRefType =
        MemRefType::get({}, unrankedMemRefType.getElementType());
    FailureOr<unsigned> maybeAddressSpace =
        getTypeConverter()->getMemRefAddressSpace(unrankedMemRefType);
    if (failed(maybeAddressSpace)) {
      dimOp.emitOpError("memref memory space must be convertible to an integer "
                        "address space");
      return failure();
    }
    unsigned addressSpace = *maybeAddressSpace;

    // Extract pointer to the underlying ranked descriptor and bitcast it to a
    // memref<element_type> descriptor pointer to minimize the number of GEP
    // operations.
    UnrankedMemRefDescriptor unrankedDesc(adaptor.getSource());
    Value underlyingRankedDesc = unrankedDesc.memRefDescPtr(rewriter, loc);

    Type elementType = typeConverter->convertType(scalarMemRefType);
    Value scalarMemRefDescPtr;
    if (getTypeConverter()->useOpaquePointers())
      scalarMemRefDescPtr = underlyingRankedDesc;
    else
      scalarMemRefDescPtr = rewriter.create<LLVM::BitcastOp>(
          loc, LLVM::LLVMPointerType::get(elementType, addressSpace),
          underlyingRankedDesc);

    // Get pointer to offset field of memref<element_type> descriptor.
    Type indexPtrTy = getTypeConverter()->getPointerType(
        getTypeConverter()->getIndexType(), addressSpace);
    Value offsetPtr = rewriter.create<LLVM::GEPOp>(
        loc, indexPtrTy, elementType, scalarMemRefDescPtr,
        ArrayRef<LLVM::GEPArg>{0, 2});

    // The size value that we have to extract can be obtained using GEPop with
    // `dimOp.index() + 1` index argument.
    Value idxPlusOne = rewriter.create<LLVM::AddOp>(
        loc, createIndexConstant(rewriter, loc, 1), adaptor.getIndex());
    Value sizePtr = rewriter.create<LLVM::GEPOp>(
        loc, indexPtrTy, getTypeConverter()->getIndexType(), offsetPtr,
        idxPlusOne);
    return rewriter
        .create<LLVM::LoadOp>(loc, getTypeConverter()->getIndexType(), sizePtr)
        .getResult();
  }

  std::optional<int64_t> getConstantDimIndex(memref::DimOp dimOp) const {
    if (auto idx = dimOp.getConstantIndex())
      return idx;

    if (auto constantOp = dimOp.getIndex().getDefiningOp<LLVM::ConstantOp>())
      return cast<IntegerAttr>(constantOp.getValue()).getValue().getSExtValue();

    return std::nullopt;
  }

  Value extractSizeOfRankedMemRef(Type operandType, memref::DimOp dimOp,
                                  OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
    Location loc = dimOp.getLoc();

    // Take advantage if index is constant.
    MemRefType memRefType = cast<MemRefType>(operandType);
    if (std::optional<int64_t> index = getConstantDimIndex(dimOp)) {
      int64_t i = *index;
      if (i >= 0 && i < memRefType.getRank()) {
        if (memRefType.isDynamicDim(i)) {
          // extract dynamic size from the memref descriptor.
          MemRefDescriptor descriptor(adaptor.getSource());
          return descriptor.size(rewriter, loc, i);
        }
        // Use constant for static size.
        int64_t dimSize = memRefType.getDimSize(i);
        return createIndexConstant(rewriter, loc, dimSize);
      }
    }
    Value index = adaptor.getIndex();
    int64_t rank = memRefType.getRank();
    MemRefDescriptor memrefDescriptor(adaptor.getSource());
    return memrefDescriptor.size(rewriter, loc, index, rank);
  }
};

/// Common base for load and store operations on MemRefs. Restricts the match
/// to supported MemRef types. Provides functionality to emit code accessing a
/// specific element of the underlying data buffer.
template <typename Derived>
struct LoadStoreOpLowering : public ConvertOpToLLVMPattern<Derived> {
  using ConvertOpToLLVMPattern<Derived>::ConvertOpToLLVMPattern;
  using ConvertOpToLLVMPattern<Derived>::isConvertibleAndHasIdentityMaps;
  using Base = LoadStoreOpLowering<Derived>;

  LogicalResult match(Derived op) const override {
    MemRefType type = op.getMemRefType();
    return isConvertibleAndHasIdentityMaps(type) ? success() : failure();
  }
};

/// Wrap a llvm.cmpxchg operation in a while loop so that the operation can be
/// retried until it succeeds in atomically storing a new value into memory.
///
///      +---------------------------------+
///      |   <code before the AtomicRMWOp> |
///      |   <compute initial %loaded>     |
///      |   cf.br loop(%loaded)              |
///      +---------------------------------+
///             |
///  -------|   |
///  |      v   v
///  |   +--------------------------------+
///  |   | loop(%loaded):                 |
///  |   |   <body contents>              |
///  |   |   %pair = cmpxchg              |
///  |   |   %ok = %pair[0]               |
///  |   |   %new = %pair[1]              |
///  |   |   cf.cond_br %ok, end, loop(%new) |
///  |   +--------------------------------+
///  |          |        |
///  |-----------        |
///                      v
///      +--------------------------------+
///      | end:                           |
///      |   <code after the AtomicRMWOp> |
///      +--------------------------------+
///
struct GenericAtomicRMWOpLowering
    : public LoadStoreOpLowering<memref::GenericAtomicRMWOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(memref::GenericAtomicRMWOp atomicOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = atomicOp.getLoc();
    Type valueType = typeConverter->convertType(atomicOp.getResult().getType());

    // Split the block into initial, loop, and ending parts.
    auto *initBlock = rewriter.getInsertionBlock();
    auto *loopBlock = rewriter.splitBlock(initBlock, Block::iterator(atomicOp));
    loopBlock->addArgument(valueType, loc);

    auto *endBlock =
        rewriter.splitBlock(loopBlock, Block::iterator(atomicOp)++);

    // Compute the loaded value and branch to the loop block.
    rewriter.setInsertionPointToEnd(initBlock);
    auto memRefType = cast<MemRefType>(atomicOp.getMemref().getType());
    auto dataPtr = getStridedElementPtr(loc, memRefType, adaptor.getMemref(),
                                        adaptor.getIndices(), rewriter);
    Value init = rewriter.create<LLVM::LoadOp>(
        loc, typeConverter->convertType(memRefType.getElementType()), dataPtr);
    rewriter.create<LLVM::BrOp>(loc, init, loopBlock);

    // Prepare the body of the loop block.
    rewriter.setInsertionPointToStart(loopBlock);

    // Clone the GenericAtomicRMWOp region and extract the result.
    auto loopArgument = loopBlock->getArgument(0);
    IRMapping mapping;
    mapping.map(atomicOp.getCurrentValue(), loopArgument);
    Block &entryBlock = atomicOp.body().front();
    for (auto &nestedOp : entryBlock.without_terminator()) {
      Operation *clone = rewriter.clone(nestedOp, mapping);
      mapping.map(nestedOp.getResults(), clone->getResults());
    }
    Value result = mapping.lookup(entryBlock.getTerminator()->getOperand(0));

    // Prepare the epilog of the loop block.
    // Append the cmpxchg op to the end of the loop block.
    auto successOrdering = LLVM::AtomicOrdering::acq_rel;
    auto failureOrdering = LLVM::AtomicOrdering::monotonic;
    auto cmpxchg = rewriter.create<LLVM::AtomicCmpXchgOp>(
        loc, dataPtr, loopArgument, result, successOrdering, failureOrdering);
    // Extract the %new_loaded and %ok values from the pair.
    Value newLoaded = rewriter.create<LLVM::ExtractValueOp>(loc, cmpxchg, 0);
    Value ok = rewriter.create<LLVM::ExtractValueOp>(loc, cmpxchg, 1);

    // Conditionally branch to the end or back to the loop depending on %ok.
    rewriter.create<LLVM::CondBrOp>(loc, ok, endBlock, ArrayRef<Value>(),
                                    loopBlock, newLoaded);

    rewriter.setInsertionPointToEnd(endBlock);

    // The 'result' of the atomic_rmw op is the newly loaded value.
    rewriter.replaceOp(atomicOp, {newLoaded});

    return success();
  }
};

/// Returns the LLVM type of the global variable given the memref type `type`.
static Type convertGlobalMemrefTypeToLLVM(MemRefType type,
                                          LLVMTypeConverter &typeConverter) {
  // LLVM type for a global memref will be a multi-dimension array. For
  // declarations or uninitialized global memrefs, we can potentially flatten
  // this to a 1D array. However, for memref.global's with an initial value,
  // we do not intend to flatten the ElementsAttribute when going from std ->
  // LLVM dialect, so the LLVM type needs to me a multi-dimension array.
  Type elementType = typeConverter.convertType(type.getElementType());
  Type arrayTy = elementType;
  // Shape has the outermost dim at index 0, so need to walk it backwards
  for (int64_t dim : llvm::reverse(type.getShape()))
    arrayTy = LLVM::LLVMArrayType::get(arrayTy, dim);
  return arrayTy;
}

/// GlobalMemrefOp is lowered to a LLVM Global Variable.
struct GlobalMemrefOpLowering
    : public ConvertOpToLLVMPattern<memref::GlobalOp> {
  using ConvertOpToLLVMPattern<memref::GlobalOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::GlobalOp global, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    MemRefType type = global.getType();
    if (!isConvertibleAndHasIdentityMaps(type))
      return failure();

    Type arrayTy = convertGlobalMemrefTypeToLLVM(type, *getTypeConverter());

    LLVM::Linkage linkage =
        global.isPublic() ? LLVM::Linkage::External : LLVM::Linkage::Private;

    Attribute initialValue = nullptr;
    if (!global.isExternal() && !global.isUninitialized()) {
      auto elementsAttr = llvm::cast<ElementsAttr>(*global.getInitialValue());
      initialValue = elementsAttr;

      // For scalar memrefs, the global variable created is of the element type,
      // so unpack the elements attribute to extract the value.
      if (type.getRank() == 0)
        initialValue = elementsAttr.getSplatValue<Attribute>();
    }

    uint64_t alignment = global.getAlignment().value_or(0);
    FailureOr<unsigned> addressSpace =
        getTypeConverter()->getMemRefAddressSpace(type);
    if (failed(addressSpace))
      return global.emitOpError(
          "memory space cannot be converted to an integer address space");
    auto newGlobal = rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
        global, arrayTy, global.getConstant(), linkage, global.getSymName(),
        initialValue, alignment, *addressSpace);
    if (!global.isExternal() && global.isUninitialized()) {
      Block *blk = new Block();
      newGlobal.getInitializerRegion().push_back(blk);
      rewriter.setInsertionPointToStart(blk);
      Value undef[] = {
          rewriter.create<LLVM::UndefOp>(global.getLoc(), arrayTy)};
      rewriter.create<LLVM::ReturnOp>(global.getLoc(), undef);
    }
    return success();
  }
};

/// GetGlobalMemrefOp is lowered into a Memref descriptor with the pointer to
/// the first element stashed into the descriptor. This reuses
/// `AllocLikeOpLowering` to reuse the Memref descriptor construction.
struct GetGlobalMemrefOpLowering : public AllocLikeOpLLVMLowering {
  GetGlobalMemrefOpLowering(LLVMTypeConverter &converter)
      : AllocLikeOpLLVMLowering(memref::GetGlobalOp::getOperationName(),
                                converter) {}

  /// Buffer "allocation" for memref.get_global op is getting the address of
  /// the global variable referenced.
  std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                          Location loc, Value sizeBytes,
                                          Operation *op) const override {
    auto getGlobalOp = cast<memref::GetGlobalOp>(op);
    MemRefType type = cast<MemRefType>(getGlobalOp.getResult().getType());

    // This is called after a type conversion, which would have failed if this
    // call fails.
    unsigned memSpace = *getTypeConverter()->getMemRefAddressSpace(type);

    Type arrayTy = convertGlobalMemrefTypeToLLVM(type, *getTypeConverter());
    Type resTy = getTypeConverter()->getPointerType(arrayTy, memSpace);
    auto addressOf =
        rewriter.create<LLVM::AddressOfOp>(loc, resTy, getGlobalOp.getName());

    // Get the address of the first element in the array by creating a GEP with
    // the address of the GV as the base, and (rank + 1) number of 0 indices.
    Type elementType = typeConverter->convertType(type.getElementType());
    Type elementPtrType =
        getTypeConverter()->getPointerType(elementType, memSpace);

    auto gep = rewriter.create<LLVM::GEPOp>(
        loc, elementPtrType, arrayTy, addressOf,
        SmallVector<LLVM::GEPArg>(type.getRank() + 1, 0));

    // We do not expect the memref obtained using `memref.get_global` to be
    // ever deallocated. Set the allocated pointer to be known bad value to
    // help debug if that ever happens.
    auto intPtrType = getIntPtrType(memSpace);
    Value deadBeefConst =
        createIndexAttrConstant(rewriter, op->getLoc(), intPtrType, 0xdeadbeef);
    auto deadBeefPtr =
        rewriter.create<LLVM::IntToPtrOp>(loc, elementPtrType, deadBeefConst);

    // Both allocated and aligned pointers are same. We could potentially stash
    // a nullptr for the allocated pointer since we do not expect any dealloc.
    return std::make_tuple(deadBeefPtr, gep);
  }
};

// Load operation is lowered to obtaining a pointer to the indexed element
// and loading it.
struct LoadOpLowering : public LoadStoreOpLowering<memref::LoadOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto type = loadOp.getMemRefType();

    Value dataPtr =
        getStridedElementPtr(loadOp.getLoc(), type, adaptor.getMemref(),
                             adaptor.getIndices(), rewriter);
    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(
        loadOp, typeConverter->convertType(type.getElementType()), dataPtr, 0,
        false, loadOp.getNontemporal());
    return success();
  }
};

// Store operation is lowered to obtaining a pointer to the indexed element,
// and storing the given value to it.
struct StoreOpLowering : public LoadStoreOpLowering<memref::StoreOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(memref::StoreOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto type = op.getMemRefType();

    Value dataPtr = getStridedElementPtr(op.getLoc(), type, adaptor.getMemref(),
                                         adaptor.getIndices(), rewriter);
    rewriter.replaceOpWithNewOp<LLVM::StoreOp>(op, adaptor.getValue(), dataPtr,
                                               0, false, op.getNontemporal());
    return success();
  }
};

// The prefetch operation is lowered in a way similar to the load operation
// except that the llvm.prefetch operation is used for replacement.
struct PrefetchOpLowering : public LoadStoreOpLowering<memref::PrefetchOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(memref::PrefetchOp prefetchOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto type = prefetchOp.getMemRefType();
    auto loc = prefetchOp.getLoc();

    Value dataPtr = getStridedElementPtr(loc, type, adaptor.getMemref(),
                                         adaptor.getIndices(), rewriter);

    // Replace with llvm.prefetch.
    IntegerAttr isWrite = rewriter.getI32IntegerAttr(prefetchOp.getIsWrite());
    IntegerAttr localityHint = prefetchOp.getLocalityHintAttr();
    IntegerAttr isData =
        rewriter.getI32IntegerAttr(prefetchOp.getIsDataCache());
    rewriter.replaceOpWithNewOp<LLVM::Prefetch>(prefetchOp, dataPtr, isWrite,
                                                localityHint, isData);
    return success();
  }
};

struct RankOpLowering : public ConvertOpToLLVMPattern<memref::RankOp> {
  using ConvertOpToLLVMPattern<memref::RankOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::RankOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    Type operandType = op.getMemref().getType();
    if (auto unrankedMemRefType = dyn_cast<UnrankedMemRefType>(operandType)) {
      UnrankedMemRefDescriptor desc(adaptor.getMemref());
      rewriter.replaceOp(op, {desc.rank(rewriter, loc)});
      return success();
    }
    if (auto rankedMemRefType = dyn_cast<MemRefType>(operandType)) {
      rewriter.replaceOp(
          op, {createIndexConstant(rewriter, loc, rankedMemRefType.getRank())});
      return success();
    }
    return failure();
  }
};

struct MemRefCastOpLowering : public ConvertOpToLLVMPattern<memref::CastOp> {
  using ConvertOpToLLVMPattern<memref::CastOp>::ConvertOpToLLVMPattern;

  LogicalResult match(memref::CastOp memRefCastOp) const override {
    Type srcType = memRefCastOp.getOperand().getType();
    Type dstType = memRefCastOp.getType();

    // memref::CastOp reduce to bitcast in the ranked MemRef case and can be
    // used for type erasure. For now they must preserve underlying element type
    // and require source and result type to have the same rank. Therefore,
    // perform a sanity check that the underlying structs are the same. Once op
    // semantics are relaxed we can revisit.
    if (isa<MemRefType>(srcType) && isa<MemRefType>(dstType))
      return success(typeConverter->convertType(srcType) ==
                     typeConverter->convertType(dstType));

    // At least one of the operands is unranked type
    assert(isa<UnrankedMemRefType>(srcType) ||
           isa<UnrankedMemRefType>(dstType));

    // Unranked to unranked cast is disallowed
    return !(isa<UnrankedMemRefType>(srcType) &&
             isa<UnrankedMemRefType>(dstType))
               ? success()
               : failure();
  }

  void rewrite(memref::CastOp memRefCastOp, OpAdaptor adaptor,
               ConversionPatternRewriter &rewriter) const override {
    auto srcType = memRefCastOp.getOperand().getType();
    auto dstType = memRefCastOp.getType();
    auto targetStructType = typeConverter->convertType(memRefCastOp.getType());
    auto loc = memRefCastOp.getLoc();

    // For ranked/ranked case, just keep the original descriptor.
    if (isa<MemRefType>(srcType) && isa<MemRefType>(dstType))
      return rewriter.replaceOp(memRefCastOp, {adaptor.getSource()});

    if (isa<MemRefType>(srcType) && isa<UnrankedMemRefType>(dstType)) {
      // Casting ranked to unranked memref type
      // Set the rank in the destination from the memref type
      // Allocate space on the stack and copy the src memref descriptor
      // Set the ptr in the destination to the stack space
      auto srcMemRefType = cast<MemRefType>(srcType);
      int64_t rank = srcMemRefType.getRank();
      // ptr = AllocaOp sizeof(MemRefDescriptor)
      auto ptr = getTypeConverter()->promoteOneMemRefDescriptor(
          loc, adaptor.getSource(), rewriter);

      // voidptr = BitCastOp srcType* to void*
      Value voidPtr;
      if (getTypeConverter()->useOpaquePointers())
        voidPtr = ptr;
      else
        voidPtr = rewriter.create<LLVM::BitcastOp>(loc, getVoidPtrType(), ptr);

      // rank = ConstantOp srcRank
      auto rankVal = rewriter.create<LLVM::ConstantOp>(
          loc, getIndexType(), rewriter.getIndexAttr(rank));
      // undef = UndefOp
      UnrankedMemRefDescriptor memRefDesc =
          UnrankedMemRefDescriptor::undef(rewriter, loc, targetStructType);
      // d1 = InsertValueOp undef, rank, 0
      memRefDesc.setRank(rewriter, loc, rankVal);
      // d2 = InsertValueOp d1, voidptr, 1
      memRefDesc.setMemRefDescPtr(rewriter, loc, voidPtr);
      rewriter.replaceOp(memRefCastOp, (Value)memRefDesc);

    } else if (isa<UnrankedMemRefType>(srcType) && isa<MemRefType>(dstType)) {
      // Casting from unranked type to ranked.
      // The operation is assumed to be doing a correct cast. If the destination
      // type mismatches the unranked the type, it is undefined behavior.
      UnrankedMemRefDescriptor memRefDesc(adaptor.getSource());
      // ptr = ExtractValueOp src, 1
      auto ptr = memRefDesc.memRefDescPtr(rewriter, loc);
      // castPtr = BitCastOp i8* to structTy*
      Value castPtr;
      if (getTypeConverter()->useOpaquePointers())
        castPtr = ptr;
      else
        castPtr = rewriter.create<LLVM::BitcastOp>(
            loc, LLVM::LLVMPointerType::get(targetStructType), ptr);

      // struct = LoadOp castPtr
      auto loadOp =
          rewriter.create<LLVM::LoadOp>(loc, targetStructType, castPtr);
      rewriter.replaceOp(memRefCastOp, loadOp.getResult());
    } else {
      llvm_unreachable("Unsupported unranked memref to unranked memref cast");
    }
  }
};

/// Pattern to lower a `memref.copy` to llvm.
///
/// For memrefs with identity layouts, the copy is lowered to the llvm
/// `memcpy` intrinsic. For non-identity layouts, the copy is lowered to a call
/// to the generic `MemrefCopyFn`.
struct MemRefCopyOpLowering : public ConvertOpToLLVMPattern<memref::CopyOp> {
  using ConvertOpToLLVMPattern<memref::CopyOp>::ConvertOpToLLVMPattern;

  LogicalResult
  lowerToMemCopyIntrinsic(memref::CopyOp op, OpAdaptor adaptor,
                          ConversionPatternRewriter &rewriter) const {
    auto loc = op.getLoc();
    auto srcType = dyn_cast<MemRefType>(op.getSource().getType());

    MemRefDescriptor srcDesc(adaptor.getSource());

    // Compute number of elements.
    Value numElements = rewriter.create<LLVM::ConstantOp>(
        loc, getIndexType(), rewriter.getIndexAttr(1));
    for (int pos = 0; pos < srcType.getRank(); ++pos) {
      auto size = srcDesc.size(rewriter, loc, pos);
      numElements = rewriter.create<LLVM::MulOp>(loc, numElements, size);
    }

    // Get element size.
    auto sizeInBytes = getSizeInBytes(loc, srcType.getElementType(), rewriter);
    // Compute total.
    Value totalSize =
        rewriter.create<LLVM::MulOp>(loc, numElements, sizeInBytes);

    Type elementType = typeConverter->convertType(srcType.getElementType());

    Value srcBasePtr = srcDesc.alignedPtr(rewriter, loc);
    Value srcOffset = srcDesc.offset(rewriter, loc);
    Value srcPtr = rewriter.create<LLVM::GEPOp>(
        loc, srcBasePtr.getType(), elementType, srcBasePtr, srcOffset);
    MemRefDescriptor targetDesc(adaptor.getTarget());
    Value targetBasePtr = targetDesc.alignedPtr(rewriter, loc);
    Value targetOffset = targetDesc.offset(rewriter, loc);
    Value targetPtr = rewriter.create<LLVM::GEPOp>(
        loc, targetBasePtr.getType(), elementType, targetBasePtr, targetOffset);
    rewriter.create<LLVM::MemcpyOp>(loc, targetPtr, srcPtr, totalSize,
                                    /*isVolatile=*/false);
    rewriter.eraseOp(op);

    return success();
  }

  LogicalResult
  lowerToMemCopyFunctionCall(memref::CopyOp op, OpAdaptor adaptor,
                             ConversionPatternRewriter &rewriter) const {
    auto loc = op.getLoc();
    auto srcType = cast<BaseMemRefType>(op.getSource().getType());
    auto targetType = cast<BaseMemRefType>(op.getTarget().getType());

    // First make sure we have an unranked memref descriptor representation.
    auto makeUnranked = [&, this](Value ranked, MemRefType type) {
      auto rank = rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                    type.getRank());
      auto *typeConverter = getTypeConverter();
      auto ptr =
          typeConverter->promoteOneMemRefDescriptor(loc, ranked, rewriter);

      Value voidPtr;
      if (getTypeConverter()->useOpaquePointers())
        voidPtr = ptr;
      else
        voidPtr = rewriter.create<LLVM::BitcastOp>(loc, getVoidPtrType(), ptr);

      auto unrankedType =
          UnrankedMemRefType::get(type.getElementType(), type.getMemorySpace());
      return UnrankedMemRefDescriptor::pack(rewriter, loc, *typeConverter,
                                            unrankedType,
                                            ValueRange{rank, voidPtr});
    };

    // Save stack position before promoting descriptors
    auto stackSaveOp =
        rewriter.create<LLVM::StackSaveOp>(loc, getVoidPtrType());

    auto srcMemRefType = dyn_cast<MemRefType>(srcType);
    Value unrankedSource =
        srcMemRefType ? makeUnranked(adaptor.getSource(), srcMemRefType)
                      : adaptor.getSource();
    auto targetMemRefType = dyn_cast<MemRefType>(targetType);
    Value unrankedTarget =
        targetMemRefType ? makeUnranked(adaptor.getTarget(), targetMemRefType)
                         : adaptor.getTarget();

    // Now promote the unranked descriptors to the stack.
    auto one = rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                 rewriter.getIndexAttr(1));
    auto promote = [&](Value desc) {
      Type ptrType = getTypeConverter()->getPointerType(desc.getType());
      auto allocated =
          rewriter.create<LLVM::AllocaOp>(loc, ptrType, desc.getType(), one);
      rewriter.create<LLVM::StoreOp>(loc, desc, allocated);
      return allocated;
    };

    auto sourcePtr = promote(unrankedSource);
    auto targetPtr = promote(unrankedTarget);

    unsigned typeSize =
        mlir::DataLayout::closest(op).getTypeSize(srcType.getElementType());
    auto elemSize = rewriter.create<LLVM::ConstantOp>(
        loc, getIndexType(), rewriter.getIndexAttr(typeSize));
    auto copyFn = LLVM::lookupOrCreateMemRefCopyFn(
        op->getParentOfType<ModuleOp>(), getIndexType(), sourcePtr.getType());
    rewriter.create<LLVM::CallOp>(loc, copyFn,
                                  ValueRange{elemSize, sourcePtr, targetPtr});

    // Restore stack used for descriptors
    rewriter.create<LLVM::StackRestoreOp>(loc, stackSaveOp);

    rewriter.eraseOp(op);

    return success();
  }

  LogicalResult
  matchAndRewrite(memref::CopyOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto srcType = cast<BaseMemRefType>(op.getSource().getType());
    auto targetType = cast<BaseMemRefType>(op.getTarget().getType());

    auto isContiguousMemrefType = [&](BaseMemRefType type) {
      auto memrefType = dyn_cast<mlir::MemRefType>(type);
      // We can use memcpy for memrefs if they have an identity layout or are
      // contiguous with an arbitrary offset. Ignore empty memrefs, which is a
      // special case handled by memrefCopy.
      return memrefType &&
             (memrefType.getLayout().isIdentity() ||
              (memrefType.hasStaticShape() && memrefType.getNumElements() > 0 &&
               memref::isStaticShapeAndContiguousRowMajor(memrefType)));
    };

    if (isContiguousMemrefType(srcType) && isContiguousMemrefType(targetType))
      return lowerToMemCopyIntrinsic(op, adaptor, rewriter);

    return lowerToMemCopyFunctionCall(op, adaptor, rewriter);
  }
};

struct MemorySpaceCastOpLowering
    : public ConvertOpToLLVMPattern<memref::MemorySpaceCastOp> {
  using ConvertOpToLLVMPattern<
      memref::MemorySpaceCastOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::MemorySpaceCastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();

    Type resultType = op.getDest().getType();
    if (auto resultTypeR = dyn_cast<MemRefType>(resultType)) {
      auto resultDescType =
          cast<LLVM::LLVMStructType>(typeConverter->convertType(resultTypeR));
      Type newPtrType = resultDescType.getBody()[0];

      SmallVector<Value> descVals;
      MemRefDescriptor::unpack(rewriter, loc, adaptor.getSource(), resultTypeR,
                               descVals);
      descVals[0] =
          rewriter.create<LLVM::AddrSpaceCastOp>(loc, newPtrType, descVals[0]);
      descVals[1] =
          rewriter.create<LLVM::AddrSpaceCastOp>(loc, newPtrType, descVals[1]);
      Value result = MemRefDescriptor::pack(rewriter, loc, *getTypeConverter(),
                                            resultTypeR, descVals);
      rewriter.replaceOp(op, result);
      return success();
    }
    if (auto resultTypeU = dyn_cast<UnrankedMemRefType>(resultType)) {
      // Since the type converter won't be doing this for us, get the address
      // space.
      auto sourceType = cast<UnrankedMemRefType>(op.getSource().getType());
      FailureOr<unsigned> maybeSourceAddrSpace =
          getTypeConverter()->getMemRefAddressSpace(sourceType);
      if (failed(maybeSourceAddrSpace))
        return rewriter.notifyMatchFailure(loc,
                                           "non-integer source address space");
      unsigned sourceAddrSpace = *maybeSourceAddrSpace;
      FailureOr<unsigned> maybeResultAddrSpace =
          getTypeConverter()->getMemRefAddressSpace(resultTypeU);
      if (failed(maybeResultAddrSpace))
        return rewriter.notifyMatchFailure(loc,
                                           "non-integer result address space");
      unsigned resultAddrSpace = *maybeResultAddrSpace;

      UnrankedMemRefDescriptor sourceDesc(adaptor.getSource());
      Value rank = sourceDesc.rank(rewriter, loc);
      Value sourceUnderlyingDesc = sourceDesc.memRefDescPtr(rewriter, loc);

      // Create and allocate storage for new memref descriptor.
      auto result = UnrankedMemRefDescriptor::undef(
          rewriter, loc, typeConverter->convertType(resultTypeU));
      result.setRank(rewriter, loc, rank);
      SmallVector<Value, 1> sizes;
      UnrankedMemRefDescriptor::computeSizes(rewriter, loc, *getTypeConverter(),
                                             result, resultAddrSpace, sizes);
      Value resultUnderlyingSize = sizes.front();
      Value resultUnderlyingDesc = rewriter.create<LLVM::AllocaOp>(
          loc, getVoidPtrType(), rewriter.getI8Type(), resultUnderlyingSize);
      result.setMemRefDescPtr(rewriter, loc, resultUnderlyingDesc);

      // Copy pointers, performing address space casts.
      Type llvmElementType =
          typeConverter->convertType(sourceType.getElementType());
      LLVM::LLVMPointerType sourceElemPtrType =
          getTypeConverter()->getPointerType(llvmElementType, sourceAddrSpace);
      auto resultElemPtrType =
          getTypeConverter()->getPointerType(llvmElementType, resultAddrSpace);

      Value allocatedPtr = sourceDesc.allocatedPtr(
          rewriter, loc, sourceUnderlyingDesc, sourceElemPtrType);
      Value alignedPtr =
          sourceDesc.alignedPtr(rewriter, loc, *getTypeConverter(),
                                sourceUnderlyingDesc, sourceElemPtrType);
      allocatedPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
          loc, resultElemPtrType, allocatedPtr);
      alignedPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
          loc, resultElemPtrType, alignedPtr);

      result.setAllocatedPtr(rewriter, loc, resultUnderlyingDesc,
                             resultElemPtrType, allocatedPtr);
      result.setAlignedPtr(rewriter, loc, *getTypeConverter(),
                           resultUnderlyingDesc, resultElemPtrType, alignedPtr);

      // Copy all the index-valued operands.
      Value sourceIndexVals =
          sourceDesc.offsetBasePtr(rewriter, loc, *getTypeConverter(),
                                   sourceUnderlyingDesc, sourceElemPtrType);
      Value resultIndexVals =
          result.offsetBasePtr(rewriter, loc, *getTypeConverter(),
                               resultUnderlyingDesc, resultElemPtrType);

      int64_t bytesToSkip =
          2 *
          ceilDiv(getTypeConverter()->getPointerBitwidth(resultAddrSpace), 8);
      Value bytesToSkipConst = rewriter.create<LLVM::ConstantOp>(
          loc, getIndexType(), rewriter.getIndexAttr(bytesToSkip));
      Value copySize = rewriter.create<LLVM::SubOp>(
          loc, getIndexType(), resultUnderlyingSize, bytesToSkipConst);
      rewriter.create<LLVM::MemcpyOp>(loc, resultIndexVals, sourceIndexVals,
                                      copySize, /*isVolatile=*/false);

      rewriter.replaceOp(op, ValueRange{result});
      return success();
    }
    return rewriter.notifyMatchFailure(loc, "unexpected memref type");
  }
};

/// Extracts allocated, aligned pointers and offset from a ranked or unranked
/// memref type. In unranked case, the fields are extracted from the underlying
/// ranked descriptor.
static void extractPointersAndOffset(Location loc,
                                     ConversionPatternRewriter &rewriter,
                                     LLVMTypeConverter &typeConverter,
                                     Value originalOperand,
                                     Value convertedOperand,
                                     Value *allocatedPtr, Value *alignedPtr,
                                     Value *offset = nullptr) {
  Type operandType = originalOperand.getType();
  if (isa<MemRefType>(operandType)) {
    MemRefDescriptor desc(convertedOperand);
    *allocatedPtr = desc.allocatedPtr(rewriter, loc);
    *alignedPtr = desc.alignedPtr(rewriter, loc);
    if (offset != nullptr)
      *offset = desc.offset(rewriter, loc);
    return;
  }

  // These will all cause assert()s on unconvertible types.
  unsigned memorySpace = *typeConverter.getMemRefAddressSpace(
      cast<UnrankedMemRefType>(operandType));
  Type elementType = cast<UnrankedMemRefType>(operandType).getElementType();
  Type llvmElementType = typeConverter.convertType(elementType);
  LLVM::LLVMPointerType elementPtrType =
      typeConverter.getPointerType(llvmElementType, memorySpace);

  // Extract pointer to the underlying ranked memref descriptor and cast it to
  // ElemType**.
  UnrankedMemRefDescriptor unrankedDesc(convertedOperand);
  Value underlyingDescPtr = unrankedDesc.memRefDescPtr(rewriter, loc);

  *allocatedPtr = UnrankedMemRefDescriptor::allocatedPtr(
      rewriter, loc, underlyingDescPtr, elementPtrType);
  *alignedPtr = UnrankedMemRefDescriptor::alignedPtr(
      rewriter, loc, typeConverter, underlyingDescPtr, elementPtrType);
  if (offset != nullptr) {
    *offset = UnrankedMemRefDescriptor::offset(
        rewriter, loc, typeConverter, underlyingDescPtr, elementPtrType);
  }
}

struct MemRefReinterpretCastOpLowering
    : public ConvertOpToLLVMPattern<memref::ReinterpretCastOp> {
  using ConvertOpToLLVMPattern<
      memref::ReinterpretCastOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::ReinterpretCastOp castOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Type srcType = castOp.getSource().getType();

    Value descriptor;
    if (failed(convertSourceMemRefToDescriptor(rewriter, srcType, castOp,
                                               adaptor, &descriptor)))
      return failure();
    rewriter.replaceOp(castOp, {descriptor});
    return success();
  }

private:
  LogicalResult convertSourceMemRefToDescriptor(
      ConversionPatternRewriter &rewriter, Type srcType,
      memref::ReinterpretCastOp castOp,
      memref::ReinterpretCastOp::Adaptor adaptor, Value *descriptor) const {
    MemRefType targetMemRefType =
        cast<MemRefType>(castOp.getResult().getType());
    auto llvmTargetDescriptorTy = dyn_cast_or_null<LLVM::LLVMStructType>(
        typeConverter->convertType(targetMemRefType));
    if (!llvmTargetDescriptorTy)
      return failure();

    // Create descriptor.
    Location loc = castOp.getLoc();
    auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);

    // Set allocated and aligned pointers.
    Value allocatedPtr, alignedPtr;
    extractPointersAndOffset(loc, rewriter, *getTypeConverter(),
                             castOp.getSource(), adaptor.getSource(),
                             &allocatedPtr, &alignedPtr);
    desc.setAllocatedPtr(rewriter, loc, allocatedPtr);
    desc.setAlignedPtr(rewriter, loc, alignedPtr);

    // Set offset.
    if (castOp.isDynamicOffset(0))
      desc.setOffset(rewriter, loc, adaptor.getOffsets()[0]);
    else
      desc.setConstantOffset(rewriter, loc, castOp.getStaticOffset(0));

    // Set sizes and strides.
    unsigned dynSizeId = 0;
    unsigned dynStrideId = 0;
    for (unsigned i = 0, e = targetMemRefType.getRank(); i < e; ++i) {
      if (castOp.isDynamicSize(i))
        desc.setSize(rewriter, loc, i, adaptor.getSizes()[dynSizeId++]);
      else
        desc.setConstantSize(rewriter, loc, i, castOp.getStaticSize(i));

      if (castOp.isDynamicStride(i))
        desc.setStride(rewriter, loc, i, adaptor.getStrides()[dynStrideId++]);
      else
        desc.setConstantStride(rewriter, loc, i, castOp.getStaticStride(i));
    }
    *descriptor = desc;
    return success();
  }
};

struct MemRefReshapeOpLowering
    : public ConvertOpToLLVMPattern<memref::ReshapeOp> {
  using ConvertOpToLLVMPattern<memref::ReshapeOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::ReshapeOp reshapeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Type srcType = reshapeOp.getSource().getType();

    Value descriptor;
    if (failed(convertSourceMemRefToDescriptor(rewriter, srcType, reshapeOp,
                                               adaptor, &descriptor)))
      return failure();
    rewriter.replaceOp(reshapeOp, {descriptor});
    return success();
  }

private:
  LogicalResult
  convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
                                  Type srcType, memref::ReshapeOp reshapeOp,
                                  memref::ReshapeOp::Adaptor adaptor,
                                  Value *descriptor) const {
    auto shapeMemRefType = cast<MemRefType>(reshapeOp.getShape().getType());
    if (shapeMemRefType.hasStaticShape()) {
      MemRefType targetMemRefType =
          cast<MemRefType>(reshapeOp.getResult().getType());
      auto llvmTargetDescriptorTy = dyn_cast_or_null<LLVM::LLVMStructType>(
          typeConverter->convertType(targetMemRefType));
      if (!llvmTargetDescriptorTy)
        return failure();

      // Create descriptor.
      Location loc = reshapeOp.getLoc();
      auto desc =
          MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);

      // Set allocated and aligned pointers.
      Value allocatedPtr, alignedPtr;
      extractPointersAndOffset(loc, rewriter, *getTypeConverter(),
                               reshapeOp.getSource(), adaptor.getSource(),
                               &allocatedPtr, &alignedPtr);
      desc.setAllocatedPtr(rewriter, loc, allocatedPtr);
      desc.setAlignedPtr(rewriter, loc, alignedPtr);

      // Extract the offset and strides from the type.
      int64_t offset;
      SmallVector<int64_t> strides;
      if (failed(getStridesAndOffset(targetMemRefType, strides, offset)))
        return rewriter.notifyMatchFailure(
            reshapeOp, "failed to get stride and offset exprs");

      if (!isStaticStrideOrOffset(offset))
        return rewriter.notifyMatchFailure(reshapeOp,
                                           "dynamic offset is unsupported");

      desc.setConstantOffset(rewriter, loc, offset);

      assert(targetMemRefType.getLayout().isIdentity() &&
             "Identity layout map is a precondition of a valid reshape op");

      Value stride = nullptr;
      int64_t targetRank = targetMemRefType.getRank();
      for (auto i : llvm::reverse(llvm::seq<int64_t>(0, targetRank))) {
        if (!ShapedType::isDynamic(strides[i])) {
          // If the stride for this dimension is dynamic, then use the product
          // of the sizes of the inner dimensions.
          stride = createIndexConstant(rewriter, loc, strides[i]);
        } else if (!stride) {
          // `stride` is null only in the first iteration of the loop.  However,
          // since the target memref has an identity layout, we can safely set
          // the innermost stride to 1.
          stride = createIndexConstant(rewriter, loc, 1);
        }

        Value dimSize;
        // If the size of this dimension is dynamic, then load it at runtime
        // from the shape operand.
        if (!targetMemRefType.isDynamicDim(i)) {
          dimSize = createIndexConstant(rewriter, loc,
                                        targetMemRefType.getDimSize(i));
        } else {
          Value shapeOp = reshapeOp.getShape();
          Value index = createIndexConstant(rewriter, loc, i);
          dimSize = rewriter.create<memref::LoadOp>(loc, shapeOp, index);
          Type indexType = getIndexType();
          if (dimSize.getType() != indexType)
            dimSize = typeConverter->materializeTargetConversion(
                rewriter, loc, indexType, dimSize);
          assert(dimSize && "Invalid memref element type");
        }

        desc.setSize(rewriter, loc, i, dimSize);
        desc.setStride(rewriter, loc, i, stride);

        // Prepare the stride value for the next dimension.
        stride = rewriter.create<LLVM::MulOp>(loc, stride, dimSize);
      }

      *descriptor = desc;
      return success();
    }

    // The shape is a rank-1 tensor with unknown length.
    Location loc = reshapeOp.getLoc();
    MemRefDescriptor shapeDesc(adaptor.getShape());
    Value resultRank = shapeDesc.size(rewriter, loc, 0);

    // Extract address space and element type.
    auto targetType = cast<UnrankedMemRefType>(reshapeOp.getResult().getType());
    unsigned addressSpace =
        *getTypeConverter()->getMemRefAddressSpace(targetType);
    Type elementType = targetType.getElementType();

    // Create the unranked memref descriptor that holds the ranked one. The
    // inner descriptor is allocated on stack.
    auto targetDesc = UnrankedMemRefDescriptor::undef(
        rewriter, loc, typeConverter->convertType(targetType));
    targetDesc.setRank(rewriter, loc, resultRank);
    SmallVector<Value, 4> sizes;
    UnrankedMemRefDescriptor::computeSizes(rewriter, loc, *getTypeConverter(),
                                           targetDesc, addressSpace, sizes);
    Value underlyingDescPtr = rewriter.create<LLVM::AllocaOp>(
        loc, getVoidPtrType(), IntegerType::get(getContext(), 8),
        sizes.front());
    targetDesc.setMemRefDescPtr(rewriter, loc, underlyingDescPtr);

    // Extract pointers and offset from the source memref.
    Value allocatedPtr, alignedPtr, offset;
    extractPointersAndOffset(loc, rewriter, *getTypeConverter(),
                             reshapeOp.getSource(), adaptor.getSource(),
                             &allocatedPtr, &alignedPtr, &offset);

    // Set pointers and offset.
    Type llvmElementType = typeConverter->convertType(elementType);
    LLVM::LLVMPointerType elementPtrType =
        getTypeConverter()->getPointerType(llvmElementType, addressSpace);

    UnrankedMemRefDescriptor::setAllocatedPtr(rewriter, loc, underlyingDescPtr,
                                              elementPtrType, allocatedPtr);
    UnrankedMemRefDescriptor::setAlignedPtr(rewriter, loc, *getTypeConverter(),
                                            underlyingDescPtr, elementPtrType,
                                            alignedPtr);
    UnrankedMemRefDescriptor::setOffset(rewriter, loc, *getTypeConverter(),
                                        underlyingDescPtr, elementPtrType,
                                        offset);

    // Use the offset pointer as base for further addressing. Copy over the new
    // shape and compute strides. For this, we create a loop from rank-1 to 0.
    Value targetSizesBase = UnrankedMemRefDescriptor::sizeBasePtr(
        rewriter, loc, *getTypeConverter(), underlyingDescPtr, elementPtrType);
    Value targetStridesBase = UnrankedMemRefDescriptor::strideBasePtr(
        rewriter, loc, *getTypeConverter(), targetSizesBase, resultRank);
    Value shapeOperandPtr = shapeDesc.alignedPtr(rewriter, loc);
    Value oneIndex = createIndexConstant(rewriter, loc, 1);
    Value resultRankMinusOne =
        rewriter.create<LLVM::SubOp>(loc, resultRank, oneIndex);

    Block *initBlock = rewriter.getInsertionBlock();
    Type indexType = getTypeConverter()->getIndexType();
    Block::iterator remainingOpsIt = std::next(rewriter.getInsertionPoint());

    Block *condBlock = rewriter.createBlock(initBlock->getParent(), {},
                                            {indexType, indexType}, {loc, loc});

    // Move the remaining initBlock ops to condBlock.
    Block *remainingBlock = rewriter.splitBlock(initBlock, remainingOpsIt);
    rewriter.mergeBlocks(remainingBlock, condBlock, ValueRange());

    rewriter.setInsertionPointToEnd(initBlock);
    rewriter.create<LLVM::BrOp>(loc, ValueRange({resultRankMinusOne, oneIndex}),
                                condBlock);
    rewriter.setInsertionPointToStart(condBlock);
    Value indexArg = condBlock->getArgument(0);
    Value strideArg = condBlock->getArgument(1);

    Value zeroIndex = createIndexConstant(rewriter, loc, 0);
    Value pred = rewriter.create<LLVM::ICmpOp>(
        loc, IntegerType::get(rewriter.getContext(), 1),
        LLVM::ICmpPredicate::sge, indexArg, zeroIndex);

    Block *bodyBlock =
        rewriter.splitBlock(condBlock, rewriter.getInsertionPoint());
    rewriter.setInsertionPointToStart(bodyBlock);

    // Copy size from shape to descriptor.
    Type llvmIndexPtrType = getTypeConverter()->getPointerType(indexType);
    Value sizeLoadGep = rewriter.create<LLVM::GEPOp>(
        loc, llvmIndexPtrType,
        typeConverter->convertType(shapeMemRefType.getElementType()),
        shapeOperandPtr, indexArg);
    Value size = rewriter.create<LLVM::LoadOp>(loc, indexType, sizeLoadGep);
    UnrankedMemRefDescriptor::setSize(rewriter, loc, *getTypeConverter(),
                                      targetSizesBase, indexArg, size);

    // Write stride value and compute next one.
    UnrankedMemRefDescriptor::setStride(rewriter, loc, *getTypeConverter(),
                                        targetStridesBase, indexArg, strideArg);
    Value nextStride = rewriter.create<LLVM::MulOp>(loc, strideArg, size);

    // Decrement loop counter and branch back.
    Value decrement = rewriter.create<LLVM::SubOp>(loc, indexArg, oneIndex);
    rewriter.create<LLVM::BrOp>(loc, ValueRange({decrement, nextStride}),
                                condBlock);

    Block *remainder =
        rewriter.splitBlock(bodyBlock, rewriter.getInsertionPoint());

    // Hook up the cond exit to the remainder.
    rewriter.setInsertionPointToEnd(condBlock);
    rewriter.create<LLVM::CondBrOp>(loc, pred, bodyBlock, std::nullopt,
                                    remainder, std::nullopt);

    // Reset position to beginning of new remainder block.
    rewriter.setInsertionPointToStart(remainder);

    *descriptor = targetDesc;
    return success();
  }
};

/// RessociatingReshapeOp must be expanded before we reach this stage.
/// Report that information.
template <typename ReshapeOp>
class ReassociatingReshapeOpConversion
    : public ConvertOpToLLVMPattern<ReshapeOp> {
public:
  using ConvertOpToLLVMPattern<ReshapeOp>::ConvertOpToLLVMPattern;
  using ReshapeOpAdaptor = typename ReshapeOp::Adaptor;

  LogicalResult
  matchAndRewrite(ReshapeOp reshapeOp, typename ReshapeOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    return rewriter.notifyMatchFailure(
        reshapeOp,
        "reassociation operations should have been expanded beforehand");
  }
};

/// Subviews must be expanded before we reach this stage.
/// Report that information.
struct SubViewOpLowering : public ConvertOpToLLVMPattern<memref::SubViewOp> {
  using ConvertOpToLLVMPattern<memref::SubViewOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::SubViewOp subViewOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    return rewriter.notifyMatchFailure(
        subViewOp, "subview operations should have been expanded beforehand");
  }
};

/// Conversion pattern that transforms a transpose op into:
///   1. A function entry `alloca` operation to allocate a ViewDescriptor.
///   2. A load of the ViewDescriptor from the pointer allocated in 1.
///   3. Updates to the ViewDescriptor to introduce the data ptr, offset, size
///      and stride. Size and stride are permutations of the original values.
///   4. A store of the resulting ViewDescriptor to the alloca'ed pointer.
/// The transpose op is replaced by the alloca'ed pointer.
class TransposeOpLowering : public ConvertOpToLLVMPattern<memref::TransposeOp> {
public:
  using ConvertOpToLLVMPattern<memref::TransposeOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::TransposeOp transposeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = transposeOp.getLoc();
    MemRefDescriptor viewMemRef(adaptor.getIn());

    // No permutation, early exit.
    if (transposeOp.getPermutation().isIdentity())
      return rewriter.replaceOp(transposeOp, {viewMemRef}), success();

    auto targetMemRef = MemRefDescriptor::undef(
        rewriter, loc,
        typeConverter->convertType(transposeOp.getIn().getType()));

    // Copy the base and aligned pointers from the old descriptor to the new
    // one.
    targetMemRef.setAllocatedPtr(rewriter, loc,
                                 viewMemRef.allocatedPtr(rewriter, loc));
    targetMemRef.setAlignedPtr(rewriter, loc,
                               viewMemRef.alignedPtr(rewriter, loc));

    // Copy the offset pointer from the old descriptor to the new one.
    targetMemRef.setOffset(rewriter, loc, viewMemRef.offset(rewriter, loc));

    // Iterate over the dimensions and apply size/stride permutation.
    for (const auto &en :
         llvm::enumerate(transposeOp.getPermutation().getResults())) {
      int sourcePos = en.index();
      int targetPos = en.value().cast<AffineDimExpr>().getPosition();
      targetMemRef.setSize(rewriter, loc, targetPos,
                           viewMemRef.size(rewriter, loc, sourcePos));
      targetMemRef.setStride(rewriter, loc, targetPos,
                             viewMemRef.stride(rewriter, loc, sourcePos));
    }

    rewriter.replaceOp(transposeOp, {targetMemRef});
    return success();
  }
};

/// Conversion pattern that transforms an op into:
///   1. An `llvm.mlir.undef` operation to create a memref descriptor
///   2. Updates to the descriptor to introduce the data ptr, offset, size
///      and stride.
/// The view op is replaced by the descriptor.
struct ViewOpLowering : public ConvertOpToLLVMPattern<memref::ViewOp> {
  using ConvertOpToLLVMPattern<memref::ViewOp>::ConvertOpToLLVMPattern;

  // Build and return the value for the idx^th shape dimension, either by
  // returning the constant shape dimension or counting the proper dynamic size.
  Value getSize(ConversionPatternRewriter &rewriter, Location loc,
                ArrayRef<int64_t> shape, ValueRange dynamicSizes,
                unsigned idx) const {
    assert(idx < shape.size());
    if (!ShapedType::isDynamic(shape[idx]))
      return createIndexConstant(rewriter, loc, shape[idx]);
    // Count the number of dynamic dims in range [0, idx]
    unsigned nDynamic =
        llvm::count_if(shape.take_front(idx), ShapedType::isDynamic);
    return dynamicSizes[nDynamic];
  }

  // Build and return the idx^th stride, either by returning the constant stride
  // or by computing the dynamic stride from the current `runningStride` and
  // `nextSize`. The caller should keep a running stride and update it with the
  // result returned by this function.
  Value getStride(ConversionPatternRewriter &rewriter, Location loc,
                  ArrayRef<int64_t> strides, Value nextSize,
                  Value runningStride, unsigned idx) const {
    assert(idx < strides.size());
    if (!ShapedType::isDynamic(strides[idx]))
      return createIndexConstant(rewriter, loc, strides[idx]);
    if (nextSize)
      return runningStride
                 ? rewriter.create<LLVM::MulOp>(loc, runningStride, nextSize)
                 : nextSize;
    assert(!runningStride);
    return createIndexConstant(rewriter, loc, 1);
  }

  LogicalResult
  matchAndRewrite(memref::ViewOp viewOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = viewOp.getLoc();

    auto viewMemRefType = viewOp.getType();
    auto targetElementTy =
        typeConverter->convertType(viewMemRefType.getElementType());
    auto targetDescTy = typeConverter->convertType(viewMemRefType);
    if (!targetDescTy || !targetElementTy ||
        !LLVM::isCompatibleType(targetElementTy) ||
        !LLVM::isCompatibleType(targetDescTy))
      return viewOp.emitWarning("Target descriptor type not converted to LLVM"),
             failure();

    int64_t offset;
    SmallVector<int64_t, 4> strides;
    auto successStrides = getStridesAndOffset(viewMemRefType, strides, offset);
    if (failed(successStrides))
      return viewOp.emitWarning("cannot cast to non-strided shape"), failure();
    assert(offset == 0 && "expected offset to be 0");

    // Target memref must be contiguous in memory (innermost stride is 1), or
    // empty (special case when at least one of the memref dimensions is 0).
    if (!strides.empty() && (strides.back() != 1 && strides.back() != 0))
      return viewOp.emitWarning("cannot cast to non-contiguous shape"),
             failure();

    // Create the descriptor.
    MemRefDescriptor sourceMemRef(adaptor.getSource());
    auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);

    // Field 1: Copy the allocated pointer, used for malloc/free.
    Value allocatedPtr = sourceMemRef.allocatedPtr(rewriter, loc);
    auto srcMemRefType = cast<MemRefType>(viewOp.getSource().getType());
    unsigned sourceMemorySpace =
        *getTypeConverter()->getMemRefAddressSpace(srcMemRefType);
    Value bitcastPtr;
    if (getTypeConverter()->useOpaquePointers())
      bitcastPtr = allocatedPtr;
    else
      bitcastPtr = rewriter.create<LLVM::BitcastOp>(
          loc, LLVM::LLVMPointerType::get(targetElementTy, sourceMemorySpace),
          allocatedPtr);

    targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);

    // Field 2: Copy the actual aligned pointer to payload.
    Value alignedPtr = sourceMemRef.alignedPtr(rewriter, loc);
    alignedPtr = rewriter.create<LLVM::GEPOp>(
        loc, alignedPtr.getType(),
        typeConverter->convertType(srcMemRefType.getElementType()), alignedPtr,
        adaptor.getByteShift());

    if (getTypeConverter()->useOpaquePointers()) {
      bitcastPtr = alignedPtr;
    } else {
      bitcastPtr = rewriter.create<LLVM::BitcastOp>(
          loc, LLVM::LLVMPointerType::get(targetElementTy, sourceMemorySpace),
          alignedPtr);
    }

    targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);

    // Field 3: The offset in the resulting type must be 0. This is because of
    // the type change: an offset on srcType* may not be expressible as an
    // offset on dstType*.
    targetMemRef.setOffset(rewriter, loc,
                           createIndexConstant(rewriter, loc, offset));

    // Early exit for 0-D corner case.
    if (viewMemRefType.getRank() == 0)
      return rewriter.replaceOp(viewOp, {targetMemRef}), success();

    // Fields 4 and 5: Update sizes and strides.
    Value stride = nullptr, nextSize = nullptr;
    for (int i = viewMemRefType.getRank() - 1; i >= 0; --i) {
      // Update size.
      Value size = getSize(rewriter, loc, viewMemRefType.getShape(),
                           adaptor.getSizes(), i);
      targetMemRef.setSize(rewriter, loc, i, size);
      // Update stride.
      stride = getStride(rewriter, loc, strides, nextSize, stride, i);
      targetMemRef.setStride(rewriter, loc, i, stride);
      nextSize = size;
    }

    rewriter.replaceOp(viewOp, {targetMemRef});
    return success();
  }
};

//===----------------------------------------------------------------------===//
// AtomicRMWOpLowering
//===----------------------------------------------------------------------===//

/// Try to match the kind of a memref.atomic_rmw to determine whether to use a
/// lowering to llvm.atomicrmw or fallback to llvm.cmpxchg.
static std::optional<LLVM::AtomicBinOp>
matchSimpleAtomicOp(memref::AtomicRMWOp atomicOp) {
  switch (atomicOp.getKind()) {
  case arith::AtomicRMWKind::addf:
    return LLVM::AtomicBinOp::fadd;
  case arith::AtomicRMWKind::addi:
    return LLVM::AtomicBinOp::add;
  case arith::AtomicRMWKind::assign:
    return LLVM::AtomicBinOp::xchg;
  case arith::AtomicRMWKind::maxs:
    return LLVM::AtomicBinOp::max;
  case arith::AtomicRMWKind::maxu:
    return LLVM::AtomicBinOp::umax;
  case arith::AtomicRMWKind::mins:
    return LLVM::AtomicBinOp::min;
  case arith::AtomicRMWKind::minu:
    return LLVM::AtomicBinOp::umin;
  case arith::AtomicRMWKind::ori:
    return LLVM::AtomicBinOp::_or;
  case arith::AtomicRMWKind::andi:
    return LLVM::AtomicBinOp::_and;
  default:
    return std::nullopt;
  }
  llvm_unreachable("Invalid AtomicRMWKind");
}

struct AtomicRMWOpLowering : public LoadStoreOpLowering<memref::AtomicRMWOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(memref::AtomicRMWOp atomicOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (failed(match(atomicOp)))
      return failure();
    auto maybeKind = matchSimpleAtomicOp(atomicOp);
    if (!maybeKind)
      return failure();
    auto memRefType = atomicOp.getMemRefType();
    auto dataPtr =
        getStridedElementPtr(atomicOp.getLoc(), memRefType, adaptor.getMemref(),
                             adaptor.getIndices(), rewriter);
    rewriter.replaceOpWithNewOp<LLVM::AtomicRMWOp>(
        atomicOp, *maybeKind, dataPtr, adaptor.getValue(),
        LLVM::AtomicOrdering::acq_rel);
    return success();
  }
};

/// Unpack the pointer returned by a memref.extract_aligned_pointer_as_index.
class ConvertExtractAlignedPointerAsIndex
    : public ConvertOpToLLVMPattern<memref::ExtractAlignedPointerAsIndexOp> {
public:
  using ConvertOpToLLVMPattern<
      memref::ExtractAlignedPointerAsIndexOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::ExtractAlignedPointerAsIndexOp extractOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    MemRefDescriptor desc(adaptor.getSource());
    rewriter.replaceOpWithNewOp<LLVM::PtrToIntOp>(
        extractOp, getTypeConverter()->getIndexType(),
        desc.alignedPtr(rewriter, extractOp->getLoc()));
    return success();
  }
};

/// Materialize the MemRef descriptor represented by the results of
/// ExtractStridedMetadataOp.
class ExtractStridedMetadataOpLowering
    : public ConvertOpToLLVMPattern<memref::ExtractStridedMetadataOp> {
public:
  using ConvertOpToLLVMPattern<
      memref::ExtractStridedMetadataOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(memref::ExtractStridedMetadataOp extractStridedMetadataOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    if (!LLVM::isCompatibleType(adaptor.getOperands().front().getType()))
      return failure();

    // Create the descriptor.
    MemRefDescriptor sourceMemRef(adaptor.getSource());
    Location loc = extractStridedMetadataOp.getLoc();
    Value source = extractStridedMetadataOp.getSource();

    auto sourceMemRefType = cast<MemRefType>(source.getType());
    int64_t rank = sourceMemRefType.getRank();
    SmallVector<Value> results;
    results.reserve(2 + rank * 2);

    // Base buffer.
    Value baseBuffer = sourceMemRef.allocatedPtr(rewriter, loc);
    Value alignedBuffer = sourceMemRef.alignedPtr(rewriter, loc);
    MemRefDescriptor dstMemRef = MemRefDescriptor::fromStaticShape(
        rewriter, loc, *getTypeConverter(),
        cast<MemRefType>(extractStridedMetadataOp.getBaseBuffer().getType()),
        baseBuffer, alignedBuffer);
    results.push_back((Value)dstMemRef);

    // Offset.
    results.push_back(sourceMemRef.offset(rewriter, loc));

    // Sizes.
    for (unsigned i = 0; i < rank; ++i)
      results.push_back(sourceMemRef.size(rewriter, loc, i));
    // Strides.
    for (unsigned i = 0; i < rank; ++i)
      results.push_back(sourceMemRef.stride(rewriter, loc, i));

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

} // namespace

void mlir::populateFinalizeMemRefToLLVMConversionPatterns(
    LLVMTypeConverter &converter, RewritePatternSet &patterns) {
  // clang-format off
  patterns.add<
      AllocaOpLowering,
      AllocaScopeOpLowering,
      AtomicRMWOpLowering,
      AssumeAlignmentOpLowering,
      ConvertExtractAlignedPointerAsIndex,
      DimOpLowering,
      ExtractStridedMetadataOpLowering,
      GenericAtomicRMWOpLowering,
      GlobalMemrefOpLowering,
      GetGlobalMemrefOpLowering,
      LoadOpLowering,
      MemRefCastOpLowering,
      MemRefCopyOpLowering,
      MemorySpaceCastOpLowering,
      MemRefReinterpretCastOpLowering,
      MemRefReshapeOpLowering,
      PrefetchOpLowering,
      RankOpLowering,
      ReassociatingReshapeOpConversion<memref::ExpandShapeOp>,
      ReassociatingReshapeOpConversion<memref::CollapseShapeOp>,
      StoreOpLowering,
      SubViewOpLowering,
      TransposeOpLowering,
      ViewOpLowering>(converter);
  // clang-format on
  auto allocLowering = converter.getOptions().allocLowering;
  if (allocLowering == LowerToLLVMOptions::AllocLowering::AlignedAlloc)
    patterns.add<AlignedAllocOpLowering, AlignedReallocOpLowering,
                 DeallocOpLowering>(converter);
  else if (allocLowering == LowerToLLVMOptions::AllocLowering::Malloc)
    patterns.add<AllocOpLowering, ReallocOpLowering, DeallocOpLowering>(
        converter);
}

namespace {
struct FinalizeMemRefToLLVMConversionPass
    : public impl::FinalizeMemRefToLLVMConversionPassBase<
          FinalizeMemRefToLLVMConversionPass> {
  using FinalizeMemRefToLLVMConversionPassBase::
      FinalizeMemRefToLLVMConversionPassBase;

  void runOnOperation() override {
    Operation *op = getOperation();
    const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
    LowerToLLVMOptions options(&getContext(),
                               dataLayoutAnalysis.getAtOrAbove(op));
    options.allocLowering =
        (useAlignedAlloc ? LowerToLLVMOptions::AllocLowering::AlignedAlloc
                         : LowerToLLVMOptions::AllocLowering::Malloc);

    options.useGenericFunctions = useGenericFunctions;
    options.useOpaquePointers = useOpaquePointers;

    if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
      options.overrideIndexBitwidth(indexBitwidth);

    LLVMTypeConverter typeConverter(&getContext(), options,
                                    &dataLayoutAnalysis);
    RewritePatternSet patterns(&getContext());
    populateFinalizeMemRefToLLVMConversionPatterns(typeConverter, patterns);
    LLVMConversionTarget target(getContext());
    target.addLegalOp<func::FuncOp>();
    if (failed(applyPartialConversion(op, target, std::move(patterns))))
      signalPassFailure();
  }
};
} // namespace