[mlir][memref] Add pass to expand realloc operations, simplify lowering to LLVM

[maerhart]

Summary

There are two motivations for this change:

1. It considerably simplifies adding support for the realloc operation to the new buffer deallocation pass by lowering the realloc such that no deallocation operation is inserted and the deallocation pass itself can insert that dealloc
2. The lowering is expressed on a higher level and thus easier to understand, and the lowerings of the memref operations it is composed of don't have to be duplicated in the MemRefToLLVM lowering (also see discussion in https://reviews.llvm.org/D133424)

[maerhart]

fix cmake build

[springerm]

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

209

memref.realloc

210

a

239

I would call it emitDeallocs with a default value of true.

mlir/lib/Dialect/MemRef/Transforms/ExpandRealloc.cpp

1

There is usually a small description here

59

I would put an assert(op.getType().getRank() == 1), same for the operand.

61

If it's not already there, let's add a check to the verifier that the operand memref or the result memref type have an identity layout map. And then let's also put an assert here, just to be safe.

131–140

You can use mlir::getValueOrCreateConstantIndexOp.

143

can be const

[maerhart]

address comments

[maerhart]

Forgot to hit 'submit' yesterday for the attached comment. Thanks for reviewing!

mlir/lib/Dialect/MemRef/Transforms/ExpandRealloc.cpp

61

Both the source and result are already verified to have identity layout, but I've added assertions.

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 class AffineDialect;
 class ModuleOp;
 
 namespace func {
+namespace arith {
+class ArithDialect;
+} // namespace arith
 class FuncDialect;
 } // namespace func
+namespace scf {
+class SCFDialect;
+} // namespace scf
 namespace tensor {
 class TensorDialect;
 } // namespace tensor
 namespace vector {
 class VectorDialect;
 } // namespace vector
 
 namespace memref {

 /// `InferShapedTypeOpInterface` or the `ReifyRankedShapeTypeShapeOpInterface`,
 /// in terms of shapes of its input operands.
 std::unique_ptr<Pass> createResolveShapedTypeResultDimsPass();
 
 /// Creates an operation pass to expand some memref operation into
 /// easier to reason about operations.
 std::unique_ptr<Pass> createExpandStridedMetadataPass();
 
+/// Creates an operation pass to expand `memref.realloc` operations into their
+/// components.
+std::unique_ptr<Pass> createExpandReallocPass(bool doNotDealloc = false);
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 
 } // namespace memref
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

let description = [{
     operations that model the effect of this operation on the different metadata.
     This pass uses affine constructs to materialize these effects.
   }];
   let constructor = "mlir::memref::createExpandStridedMetadataPass()";
   let dependentDialects = [
       "affine::AffineDialect", "memref::MemRefDialect"
   ];
 }
+
+def ExpandRealloc : Pass<"expand-realloc"> {
+  let summary = "Expand memref.realloc operations into its components";
+  let description = [{
+    The `realloc` operation performs a conditional allocation and copy to
+    increase the size of a buffer if necessary. This pass converts the `realloc`
+    operation into this sequence of simpler operations such that other passes
+    at a later stage in the compilation pipeline do not have to consider the
+    `realloc` operation anymore (e.g., the buffer deallocation pass and the
+    conversion pass to LLVM).
+
+    Example of an expansion:
+    ```mlir
+    %realloc = memref.realloc %alloc (%size) : memref<?xf32> to memref<?xf32>
+    ```
+    is expanded to
+    ```mlir
+    %c0 = arith.constant 0 : index
+    %dim = memref.dim %alloc, %c0 : memref<?xf32>
+    %is_old_smaller = arith.cmpi ult, %dim, %arg1
+    %realloc = scf.if %is_old_smaller -> (memref<?xf32>) {
+      %new_alloc = memref.alloc(%size) : memref<?xf32>
+      %subview = memref.subview %new_alloc[0] [%dim] [1]
+      memref.copy %alloc, %subview
+      memref.dealloc %alloc
+      scf.yield %alloc_0 : memref<?xf32>
+    } else {
+      %reinterpret_cast = memref.reinterpret_cast %alloc to
+        offset: [0], sizes: [%size], strides: [1]
+      scf.yield %reinterpret_cast : memref<?xf32>
+    }
+    ```
+  }];
+  let options = [
+    Option<"doNotDeallocate", "do-not-deallocate", "bool",
+           /*default=*/"false", "Do not deallocate the original MemRef">,
+  ];
+  let constructor = "mlir::memref::createExpandReallocPass()";
+  let dependentDialects = [
+      "arith::ArithDialect", "scf::SCFDialect", "memref::MemRefDialect"
+  ];
+}
+
 #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES
 

mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h

 /// Appends patterns for expanding memref operations that modify the metadata
 /// (sizes, offset, strides) of a memref into easier to analyze constructs.
 void populateExpandStridedMetadataPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for resolving `memref.extract_strided_metadata` into
 /// `memref.extract_strided_metadata` of its source.
 void populateResolveExtractStridedMetadataPatterns(RewritePatternSet &patterns);
 
+/// Appends patterns for expanding `memref.realloc` operations.
+void populateExpandReallocPatterns(RewritePatternSet &patterns,
+                                   bool doNotDealloc = false);
+
 /// Appends patterns for emulating wide integer memref operations with ops over
 /// narrower integer types.
 void populateMemRefWideIntEmulationPatterns(
     arith::WideIntEmulationConverter &typeConverter,
     RewritePatternSet &patterns);
 
 /// Appends type conversions for emulating wide integer memref operations with
 /// ops over narrowe integer types.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
     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(const 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, reallocOp), 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.
-      Type indexType = ConvertToLLVMPattern::getIndexType();
-      Value numElements =
-          type.isDynamicDim(0)
-              ? getDynamicSize()
-              : createIndexAttrConstant(rewriter, loc, indexType,
-                                        type.getDimSize(0));
-      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(const 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(const 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);
patterns.add<
       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);
+    patterns.add<AlignedAllocOpLowering, DeallocOpLowering>(converter);
   else if (allocLowering == LowerToLLVMOptions::AllocLowering::Malloc)
-    patterns.add<AllocOpLowering, ReallocOpLowering, DeallocOpLowering>(
-        converter);
+    patterns.add<AllocOpLowering, DeallocOpLowering>(converter);
 }
 
 namespace {
 struct FinalizeMemRefToLLVMConversionPass
     : public impl::FinalizeMemRefToLLVMConversionPassBase<
           FinalizeMemRefToLLVMConversionPass> {
   using FinalizeMemRefToLLVMConversionPassBase::
       FinalizeMemRefToLLVMConversionPassBase;

mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefTransforms
   BufferizableOpInterfaceImpl.cpp
   ComposeSubView.cpp
   ExpandOps.cpp
+  ExpandRealloc.cpp
   ExpandStridedMetadata.cpp
   EmulateWideInt.cpp
   EmulateNarrowType.cpp
   ExtractAddressComputations.cpp
   FoldMemRefAliasOps.cpp
   IndependenceTransforms.cpp
   MultiBuffer.cpp
   NormalizeMemRefs.cpp

mlir/lib/Dialect/MemRef/Transforms/ExpandRealloc.cpp

+//===- ExpandRealloc.cpp --------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace memref {
+#define GEN_PASS_DEF_EXPANDREALLOC
+#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
+} // namespace memref
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+/// The `realloc` operation performs a conditional allocation and copy to
+/// increase the size of a buffer if necessary. This pattern converts the
+/// `realloc` operation into this sequence of simpler operations.
+
+/// Example of an expansion:
+/// ```mlir
+/// %realloc = memref.realloc %alloc (%size) : memref<?xf32> to memref<?xf32>
+/// ```
+/// is expanded to
+/// ```mlir
+/// %c0 = arith.constant 0 : index
+/// %dim = memref.dim %alloc, %c0 : memref<?xf32>
+/// %is_old_smaller = arith.cmpi ult, %dim, %arg1
+/// %realloc = scf.if %is_old_smaller -> (memref<?xf32>) {
+///   %new_alloc = memref.alloc(%size) : memref<?xf32>
+///   %subview = memref.subview %new_alloc[0] [%dim] [1]
+///   memref.copy %alloc, %subview
+///   memref.dealloc %alloc
+///   scf.yield %alloc_0 : memref<?xf32>
+/// } else {
+///   %reinterpret_cast = memref.reinterpret_cast %alloc to
+///     offset: [0], sizes: [%size], strides: [1]
+///   scf.yield %reinterpret_cast : memref<?xf32>
+/// }
+/// ```
+struct ExpandReallocOpPattern : public OpRewritePattern<memref::ReallocOp> {
+  ExpandReallocOpPattern(MLIRContext *ctx, bool doNotDealloc)
+      : OpRewritePattern(ctx), doNotDealloc(doNotDealloc) {}
+
+  LogicalResult matchAndRewrite(memref::ReallocOp op,
+                                PatternRewriter &rewriter) const final {
+    Location loc = op.getLoc();
+
+    // Get the size of the original buffer.
+    int64_t inputSize =
+        op.getSource().getType().cast<BaseMemRefType>().getDimSize(0);
+    OpFoldResult currSize = rewriter.getIndexAttr(inputSize);
+    if (ShapedType::isDynamic(inputSize)) {
+      Value dimZero = materialize(rewriter.getIndexAttr(0), rewriter, loc);
+      currSize = rewriter.create<memref::DimOp>(loc, op.getSource(), dimZero)
+                     .getResult();
+    }
+
+    // Get the requested size that the new buffer should have.
+    int64_t outputSize =
+        op.getResult().getType().cast<BaseMemRefType>().getDimSize(0);
+    OpFoldResult targetSize = ShapedType::isDynamic(outputSize)
+                                  ? OpFoldResult{op.getDynamicResultSize()}
+                                  : rewriter.getIndexAttr(outputSize);
+
+    // Only allocate a new buffer and copy over the values in the old buffer if
+    // the old buffer is smaller than the requested size.
+    Value lhs = materialize(currSize, rewriter, loc);
+    Value rhs = materialize(targetSize, rewriter, loc);
+    Value cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
+                                                lhs, rhs);
+    auto ifOp = rewriter.create<scf::IfOp>(
+        loc, cond,
+        [&](OpBuilder &builder, Location loc) {
+          // Allocate the new buffer. If it is a dynamic memref we need to pass
+          // an additional operand for the size at runtime, otherwise the static
+          // size is encoded in the result type.
+          SmallVector<Value> dynamicSizeOperands;
+          if (op.getDynamicResultSize())
+            dynamicSizeOperands.push_back(op.getDynamicResultSize());
+
+          Value newAlloc = builder.create<memref::AllocOp>(
+              loc, op.getResult().getType(), dynamicSizeOperands,
+              op.getAlignmentAttr());
+
+          // Take a subview of the new (bigger) buffer such that we can copy the
+          // old values over (the copy operation requires both operands to have
+          // the same shape).
+          Value subview = builder.create<memref::SubViewOp>(
+              loc, newAlloc, ArrayRef<OpFoldResult>{rewriter.getIndexAttr(0)},
+              ArrayRef<OpFoldResult>{currSize},
+              ArrayRef<OpFoldResult>{rewriter.getIndexAttr(1)});
+          builder.create<memref::CopyOp>(loc, op.getSource(), subview);
+
+          // Insert the deallocation of the old buffer only if requested
+          // (enabled by default).
+          if (!doNotDealloc)
+            builder.create<memref::DeallocOp>(loc, op.getSource());
+
+          builder.create<scf::YieldOp>(loc, newAlloc);
+        },
+        [&](OpBuilder &builder, Location loc) {
+          // We need to reinterpret-cast here because either the input or output
+          // type might be static, which means we need to cast from static to
+          // dynamic or vice-versa. If both are static and the original buffer
+          // is already bigger than the requested size, the cast represents a
+          // subview operation.
+          Value casted = builder.create<memref::ReinterpretCastOp>(
+              loc, op.getResult().getType().cast<MemRefType>(), op.getSource(),
+              rewriter.getIndexAttr(0), ArrayRef<OpFoldResult>{targetSize},
+              ArrayRef<OpFoldResult>{rewriter.getIndexAttr(1)});
+          builder.create<scf::YieldOp>(loc, casted);
+        });
+
+    rewriter.replaceOp(op, ifOp.getResult(0));
+    return success();
+  }
+
+  Value materialize(OpFoldResult value, PatternRewriter &rewriter,
+                    Location loc) const {
+    if (value.is<Value>())
+      return value.get<Value>();
+
+    Attribute attr = value.get<Attribute>();
+    assert(isa<IntegerAttr>(attr) &&
+           "Only materializing integer attributes is supported.");
+    return rewriter.create<arith::ConstantOp>(loc, cast<IntegerAttr>(attr));
+  }
+
+private:
+  bool doNotDealloc;
+};
+
+struct ExpandReallocPass
+    : public memref::impl::ExpandReallocBase<ExpandReallocPass> {
+  ExpandReallocPass(bool doNotDealloc)
+      : memref::impl::ExpandReallocBase<ExpandReallocPass>() {
+    doNotDeallocate.setValue(doNotDealloc);
+  }
+  void runOnOperation() override {
+    MLIRContext &ctx = getContext();
+
+    RewritePatternSet patterns(&ctx);
+    memref::populateExpandReallocPatterns(patterns, doNotDeallocate.getValue());
+    ConversionTarget target(ctx);
+
+    target.addLegalDialect<arith::ArithDialect, scf::SCFDialect,
+                           memref::MemRefDialect>();
+    target.addIllegalOp<memref::ReallocOp>();
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+void mlir::memref::populateExpandReallocPatterns(RewritePatternSet &patterns,
+                                                 bool doNotDealloc) {
+  patterns.add<ExpandReallocOpPattern>(patterns.getContext(), doNotDealloc);
+}
+
+std::unique_ptr<Pass> mlir::memref::createExpandReallocPass(bool doNotDealloc) {
+  return std::make_unique<ExpandReallocPass>(doNotDealloc);
+}

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

if (gpuCodegen) {
     pm.addNestedPass<gpu::GPUModuleOp>(createConvertSCFToCFPass());
     pm.addNestedPass<gpu::GPUModuleOp>(createLowerGpuOpsToNVVMOpsPass());
   }
 
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add
   // it to this pipeline.
   pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
   pm.addNestedPass<func::FuncOp>(createConvertVectorToSCFPass());
+  pm.addNestedPass<func::FuncOp>(memref::createExpandReallocPass());
   pm.addNestedPass<func::FuncOp>(createConvertSCFToCFPass());
   pm.addPass(memref::createExpandStridedMetadataPass());
   pm.addPass(createLowerAffinePass());
   pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
   pm.addPass(createFinalizeMemRefToLLVMConversionPass());
   pm.addNestedPass<func::FuncOp>(createConvertComplexToStandardPass());
   pm.addNestedPass<func::FuncOp>(arith::createArithExpandOpsPass());
   pm.addNestedPass<func::FuncOp>(createConvertMathToLLVMPass());

mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir

func.func @ranked_unranked() {
   // Static alignment should be computed as ceilPowerOf2(sizeof(index) +
   // sizeof(pointer)) = 16.
   // ALIGNED-ALLOC: llvm.mlir.constant(16 : index)
   // ALIGNED-ALLOC: llvm.call @aligned_alloc
   %0 = memref.alloc() : memref<1 x memref<* x f32>>
   memref.cast %0 : memref<1 x memref<* x f32>> to memref<* x memref<* x f32>>
   return
 }
-
-// -----
-
-// CHECK-LABEL:     func.func @realloc_dynamic(
-// CHECK-SAME:      %[[arg0:.*]]: memref<?xf32>,
-// CHECK-SAME:      %[[arg1:.*]]: index) -> memref<?xf32> {
-func.func @realloc_dynamic(%in: memref<?xf32>, %d: index) -> memref<?xf32>{
-// CHECK:           %[[descriptor:.*]] = builtin.unrealized_conversion_cast %[[arg0]]
-// CHECK:           %[[src_dim:.*]] = llvm.extractvalue %[[descriptor]][3, 0]
-// CHECK:           %[[dst_dim:.*]] = builtin.unrealized_conversion_cast %[[arg1]] : index to i64
-// CHECK:           %[[cond:.*]] = llvm.icmp "ugt" %[[dst_dim]], %[[src_dim]] : i64
-// CHECK:           llvm.cond_br %[[cond]], ^bb1, ^bb2(%[[descriptor]]
-// CHECK:           ^bb1:
-// CHECK:           %[[dst_null:.*]] = llvm.mlir.null : !llvm.ptr
-// CHECK:           %[[dst_gep:.*]] = llvm.getelementptr %[[dst_null]][1]
-// CHECK:           %[[dst_es:.*]] = llvm.ptrtoint %[[dst_gep]] : !llvm.ptr to i64
-// CHECK:           %[[dst_size:.*]] = llvm.mul %[[dst_dim]], %[[dst_es]]
-// CHECK:           %[[src_size:.*]] = llvm.mul %[[src_dim]], %[[dst_es]]
-// CHECK:           %[[new_buffer_raw:.*]] = llvm.call @malloc(%[[dst_size]])
-// CHECK:           %[[old_buffer_aligned:.*]] = llvm.extractvalue %[[descriptor]][1]
-// CHECK:           "llvm.intr.memcpy"(%[[new_buffer_raw]], %[[old_buffer_aligned]], %[[src_size]]) <{isVolatile = false}>
-// CHECK:           %[[old_buffer_unaligned:.*]] = llvm.extractvalue %[[descriptor]][0]
-// CHECK:           llvm.call @free(%[[old_buffer_unaligned]])
-// CHECK:           %[[descriptor_update1:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor]][0]
-// CHECK:           %[[descriptor_update2:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor_update1]][1]
-// CHECK:           llvm.br ^bb2(%[[descriptor_update2]]
-// CHECK:           ^bb2(%[[descriptor_update3:.*]]: !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>):
-// CHECK:           %[[descriptor_update4:.*]] = llvm.insertvalue %[[dst_dim]], %[[descriptor_update3]][3, 0]
-// CHECK:           %[[descriptor_update5:.*]] = builtin.unrealized_conversion_cast %[[descriptor_update4]]
-// CHECK:           return %[[descriptor_update5]] : memref<?xf32>
-
-  %out = memref.realloc %in(%d) : memref<?xf32> to memref<?xf32>
-  return %out : memref<?xf32>
-}
-
-// -----
-
-// CHECK-LABEL:     func.func @realloc_dynamic_alignment(
-// CHECK-SAME:      %[[arg0:.*]]: memref<?xf32>,
-// CHECK-SAME:      %[[arg1:.*]]: index) -> memref<?xf32> {
-// ALIGNED-ALLOC-LABEL:   func.func @realloc_dynamic_alignment(
-// ALIGNED-ALLOC-SAME:    %[[arg0:.*]]: memref<?xf32>,
-// ALIGNED-ALLOC-SAME:    %[[arg1:.*]]: index) -> memref<?xf32> {
-func.func @realloc_dynamic_alignment(%in: memref<?xf32>, %d: index) -> memref<?xf32>{
-// CHECK:           %[[descriptor:.*]] = builtin.unrealized_conversion_cast %[[arg0]]
-// CHECK:           %[[drc_dim:.*]] = llvm.extractvalue %[[descriptor]][3, 0]
-// CHECK:           %[[dst_dim:.*]] = builtin.unrealized_conversion_cast %[[arg1]] : index to i64
-// CHECK:           %[[cond:.*]] = llvm.icmp "ugt" %[[dst_dim]], %[[drc_dim]] : i64
-// CHECK:           llvm.cond_br %[[cond]], ^bb1, ^bb2(%[[descriptor]]
-// CHECK:           ^bb1:
-// CHECK:           %[[dst_null:.*]] = llvm.mlir.null : !llvm.ptr
-// CHECK:           %[[dst_gep:.*]] = llvm.getelementptr %[[dst_null]][1]
-// CHECK:           %[[dst_es:.*]] = llvm.ptrtoint %[[dst_gep]] : !llvm.ptr to i64
-// CHECK:           %[[dst_size:.*]] = llvm.mul %[[dst_dim]], %[[dst_es]]
-// CHECK:           %[[src_size:.*]] = llvm.mul %[[drc_dim]], %[[dst_es]]
-// CHECK:           %[[alignment:.*]] = llvm.mlir.constant(8 : index) : i64
-// CHECK:           %[[adjust_dst_size:.*]] = llvm.add %[[dst_size]], %[[alignment]]
-// CHECK:           %[[new_buffer_raw:.*]] = llvm.call @malloc(%[[adjust_dst_size]])
-// CHECK:           %[[new_buffer_int:.*]] = llvm.ptrtoint %[[new_buffer_raw]] : !llvm.ptr
-// CHECK:           %[[const_1:.*]] = llvm.mlir.constant(1 : index) : i64
-// CHECK:           %[[alignment_m1:.*]] = llvm.sub %[[alignment]], %[[const_1]]
-// CHECK:           %[[ptr_alignment_m1:.*]] = llvm.add %[[new_buffer_int]], %[[alignment_m1]]
-// CHECK:           %[[padding:.*]] = llvm.urem %[[ptr_alignment_m1]], %[[alignment]]
-// CHECK:           %[[new_buffer_aligned_int:.*]] = llvm.sub %[[ptr_alignment_m1]], %[[padding]]
-// CHECK:           %[[new_buffer_aligned:.*]] = llvm.inttoptr %[[new_buffer_aligned_int]] : i64 to !llvm.ptr
-// CHECK:           %[[old_buffer_aligned:.*]] = llvm.extractvalue %[[descriptor]][1]
-// CHECK:           "llvm.intr.memcpy"(%[[new_buffer_aligned]], %[[old_buffer_aligned]], %[[src_size]]) <{isVolatile = false}>
-// CHECK:           %[[old_buffer_unaligned:.*]] = llvm.extractvalue %[[descriptor]][0]
-// CHECK:           llvm.call @free(%[[old_buffer_unaligned]])
-// CHECK:           %[[descriptor_update1:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor]][0]
-// CHECK:           %[[descriptor_update2:.*]] = llvm.insertvalue %[[new_buffer_aligned]], %[[descriptor_update1]][1]
-// CHECK:           llvm.br ^bb2(%[[descriptor_update2]]
-// CHECK:           ^bb2(%[[descriptor_update3:.*]]: !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>):
-// CHECK:           %[[descriptor_update4:.*]] = llvm.insertvalue %[[dst_dim]], %[[descriptor_update3]][3, 0]
-// CHECK:           %[[descriptor_update5:.*]] = builtin.unrealized_conversion_cast %[[descriptor_update4]]
-// CHECK:           return %[[descriptor_update5]] : memref<?xf32>
-
-// ALIGNED-ALLOC:           %[[descriptor:.*]] = builtin.unrealized_conversion_cast %[[arg0]]
-// ALIGNED-ALLOC:           %[[drc_dim:.*]] = llvm.extractvalue %[[descriptor]][3, 0]
-// ALIGNED-ALLOC:           %[[dst_dim:.*]] = builtin.unrealized_conversion_cast %[[arg1]] : index to i64
-// ALIGNED-ALLOC:           %[[cond:.*]] = llvm.icmp "ugt" %[[dst_dim]], %[[drc_dim]] : i64
-// ALIGNED-ALLOC:           llvm.cond_br %[[cond]], ^bb1, ^bb2(%[[descriptor]]
-// ALIGNED-ALLOC:           ^bb1:
-// ALIGNED-ALLOC:           %[[dst_null:.*]] = llvm.mlir.null : !llvm.ptr
-// ALIGNED-ALLOC:           %[[dst_gep:.*]] = llvm.getelementptr %[[dst_null]][1]
-// ALIGNED-ALLOC:           %[[dst_es:.*]] = llvm.ptrtoint %[[dst_gep]] : !llvm.ptr to i64
-// ALIGNED-ALLOC:           %[[dst_size:.*]] = llvm.mul %[[dst_dim]], %[[dst_es]]
-// ALIGNED-ALLOC:           %[[src_size:.*]] = llvm.mul %[[drc_dim]], %[[dst_es]]
-// ALIGNED-ALLOC-DAG:       %[[alignment:.*]] = llvm.mlir.constant(8 : index) : i64
-// ALIGNED-ALLOC-DAG:       %[[const_1:.*]] = llvm.mlir.constant(1 : index) : i64
-// ALIGNED-ALLOC:           %[[alignment_m1:.*]] = llvm.sub %[[alignment]], %[[const_1]]
-// ALIGNED-ALLOC:           %[[size_alignment_m1:.*]] = llvm.add %[[dst_size]], %[[alignment_m1]]
-// ALIGNED-ALLOC:           %[[padding:.*]] = llvm.urem %[[size_alignment_m1]], %[[alignment]]
-// ALIGNED-ALLOC:           %[[adjust_dst_size:.*]] = llvm.sub %[[size_alignment_m1]], %[[padding]]
-// ALIGNED-ALLOC:           %[[new_buffer_raw:.*]] = llvm.call @aligned_alloc(%[[alignment]], %[[adjust_dst_size]])
-// ALIGNED-ALLOC:           %[[old_buffer_aligned:.*]] = llvm.extractvalue %[[descriptor]][1]
-// ALIGNED-ALLOC:           "llvm.intr.memcpy"(%[[new_buffer_raw]], %[[old_buffer_aligned]], %[[src_size]]) <{isVolatile = false}>
-// ALIGNED-ALLOC:           %[[old_buffer_unaligned:.*]] = llvm.extractvalue %[[descriptor]][0]
-// ALIGNED-ALLOC:           llvm.call @free(%[[old_buffer_unaligned]])
-// ALIGNED-ALLOC:           %[[descriptor_update1:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor]][0]
-// ALIGNED-ALLOC:           %[[descriptor_update2:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor_update1]][1]
-// ALIGNED-ALLOC:           llvm.br ^bb2(%[[descriptor_update2]] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>)
-// ALIGNED-ALLOC:           ^bb2(%[[descriptor_update3:.*]]: !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>):
-// ALIGNED-ALLOC:           %[[descriptor_update4:.*]] = llvm.insertvalue %[[dst_dim]], %[[descriptor_update3]][3, 0]
-// ALIGNED-ALLOC:           %[[descriptor_update5:.*]] = builtin.unrealized_conversion_cast %[[descriptor_update4]]
-// ALIGNED-ALLOC:           return %[[descriptor_update5]] : memref<?xf32>
-
-  %out = memref.realloc %in(%d)  {alignment = 8} : memref<?xf32> to memref<?xf32>
-  return %out : memref<?xf32>
-}
-

mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir

func.func @memref.reshape_index(%arg0: memref<?x?xi32>, %shape: memref<1xindex>) ->  memref<?xi32> {
 
   // CHECK: %[[result_cast:.*]] = builtin.unrealized_conversion_cast %[[insert4:.*]] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi32>
   // CHECK: return %[[result_cast:.*]] : memref<?xi32>
 
   %1 = memref.reshape %arg0(%shape) : (memref<?x?xi32>, memref<1xindex>) -> memref<?xi32>
   return %1 : memref<?xi32>
 }
 
-// -----
-
-// CHECK-LABEL:     func.func @realloc_static(
-// CHECK-SAME:      %[[arg0:.*]]: memref<2xi32>) -> memref<4xi32> {
-func.func @realloc_static(%in: memref<2xi32>) -> memref<4xi32>{
-// CHECK:           %[[descriptor:.*]] = builtin.unrealized_conversion_cast %[[arg0]] : memref<2xi32> to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
-// CHECK:           %[[src_dim:.*]] = llvm.mlir.constant(2 : index) : i64
-// CHECK:           %[[dst_dim:.*]] = llvm.mlir.constant(4 : index) : i64
-// CHECK:           %[[cond:.*]] = llvm.icmp "ugt" %[[dst_dim]], %[[src_dim]]
-// CHECK:           llvm.cond_br %[[cond]], ^bb1, ^bb2(%[[descriptor]]
-// CHECK:           ^bb1:
-// CHECK:           %[[dst_null:.*]] = llvm.mlir.null : !llvm.ptr
-// CHECK:           %[[dst_gep:.*]] = llvm.getelementptr %[[dst_null]][1]
-// CHECK:           %[[dst_es:.*]] = llvm.ptrtoint %[[dst_gep]] : !llvm.ptr to i64
-// CHECK:           %[[dst_size:.*]] = llvm.mul %[[dst_dim]], %[[dst_es]]
-// CHECK:           %[[src_size:.*]] = llvm.mul %[[src_dim]], %[[dst_es]]
-// CHECK:           %[[new_buffer_raw:.*]] = llvm.call @malloc(%[[dst_size]])
-// CHECK:           %[[old_buffer_aligned:.*]] = llvm.extractvalue %[[descriptor]][1]
-// CHECK:           "llvm.intr.memcpy"(%[[new_buffer_raw]], %[[old_buffer_aligned]], %[[src_size]]) <{isVolatile = false}>
-// CHECK:           %[[old_buffer_unaligned:.*]] = llvm.extractvalue %[[descriptor]][0]
-// CHECK:           llvm.call @free(%[[old_buffer_unaligned]])
-// CHECK:           %[[descriptor_update1:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor]][0]
-// CHECK:           %[[descriptor_update2:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor_update1]][1]
-// CHECK:           llvm.br ^bb2(%[[descriptor_update2]] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>)
-// CHECK:           ^bb2(%[[descriptor_update3:.*]]: !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>):
-// CHECK:           %[[descriptor_update4:.*]] = llvm.insertvalue %[[dst_dim]], %[[descriptor_update3]][3, 0]
-// CHECK:           %[[descriptor_update5:.*]] = builtin.unrealized_conversion_cast %[[descriptor_update4]]
-// CHECK:           return %[[descriptor_update5]] : memref<4xi32>
-
-  %out = memref.realloc %in : memref<2xi32> to memref<4xi32>
-  return %out : memref<4xi32>
-}
-
-// -----
-
-// CHECK-LABEL:     func.func @realloc_static_alignment(
-// CHECK-SAME:      %[[arg0:.*]]: memref<2xf32>) -> memref<4xf32> {
-func.func @realloc_static_alignment(%in: memref<2xf32>) -> memref<4xf32>{
-// CHECK:           %[[descriptor:.*]] = builtin.unrealized_conversion_cast %[[arg0]] : memref<2xf32> to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
-// CHECK:           %[[src_dim:.*]] = llvm.mlir.constant(2 : index) : i64
-// CHECK:           %[[dst_dim:.*]] = llvm.mlir.constant(4 : index) : i64
-// CHECK:           %[[cond:.*]] = llvm.icmp "ugt" %[[dst_dim]], %[[src_dim]] : i64
-// CHECK:           llvm.cond_br %[[cond]], ^bb1, ^bb2(%[[descriptor]]
-// CHECK:           ^bb1:
-// CHECK:           %[[dst_null:.*]] = llvm.mlir.null : !llvm.ptr
-// CHECK:           %[[dst_gep:.*]] = llvm.getelementptr %[[dst_null]][1]
-// CHECK:           %[[dst_es:.*]] = llvm.ptrtoint %[[dst_gep]] : !llvm.ptr to i64
-// CHECK:           %[[dst_size:.*]] = llvm.mul %[[dst_dim]], %[[dst_es]]
-// CHECK:           %[[src_size:.*]] = llvm.mul %[[src_dim]], %[[dst_es]]
-// CHECK:           %[[alignment:.*]] = llvm.mlir.constant(8 : index) : i64
-// CHECK:           %[[adjust_dst_size:.*]] = llvm.add %[[dst_size]], %[[alignment]]
-// CHECK:           %[[new_buffer_raw:.*]] = llvm.call @malloc(%[[adjust_dst_size]])
-// CHECK:           %[[new_buffer_int:.*]] = llvm.ptrtoint %[[new_buffer_raw]] : !llvm.ptr
-// CHECK:           %[[const_1:.*]] = llvm.mlir.constant(1 : index) : i64
-// CHECK:           %[[alignment_m1:.*]] = llvm.sub %[[alignment]], %[[const_1]]
-// CHECK:           %[[ptr_alignment_m1:.*]] = llvm.add %[[new_buffer_int]], %[[alignment_m1]]
-// CHECK:           %[[padding:.*]] = llvm.urem %[[ptr_alignment_m1]], %[[alignment]]
-// CHECK:           %[[new_buffer_aligned_int:.*]] = llvm.sub %[[ptr_alignment_m1]], %[[padding]]
-// CHECK:           %[[new_buffer_aligned:.*]] = llvm.inttoptr %[[new_buffer_aligned_int]] : i64 to !llvm.ptr
-// CHECK:           %[[old_buffer_aligned:.*]] = llvm.extractvalue %[[descriptor]][1]
-// CHECK:           "llvm.intr.memcpy"(%[[new_buffer_aligned]], %[[old_buffer_aligned]], %[[src_size]]) <{isVolatile = false}>
-// CHECK:           %[[old_buffer_unaligned:.*]] = llvm.extractvalue %[[descriptor]][0]
-// CHECK:           llvm.call @free(%[[old_buffer_unaligned]])
-// CHECK:           %[[descriptor_update1:.*]] = llvm.insertvalue %[[new_buffer_raw]], %[[descriptor]][0]
-// CHECK:           %[[descriptor_update2:.*]] = llvm.insertvalue %[[new_buffer_aligned]], %[[descriptor_update1]][1]
-// CHECK:           llvm.br ^bb2(%[[descriptor_update2]]
-// CHECK:           ^bb2(%[[descriptor_update3:.*]]: !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>):
-// CHECK:           %[[descriptor_update4:.*]] = llvm.insertvalue %[[dst_dim]], %[[descriptor_update3]][3, 0]
-// CHECK:           %[[descriptor_update5:.*]] = builtin.unrealized_conversion_cast %[[descriptor_update4]]
-// CHECK:           return %[[descriptor_update5]] : memref<4xf32>
-
-
-  %out = memref.realloc %in {alignment = 8} : memref<2xf32> to memref<4xf32>
-  return %out : memref<4xf32>
-}
-
-// -----
-
 // CHECK-LABEL: @memref_memory_space_cast
 func.func @memref_memory_space_cast(%input : memref<?xf32>) -> memref<?xf32, 1> {
   %cast = memref.memory_space_cast %input : memref<?xf32> to memref<?xf32, 1>
   return %cast : memref<?xf32, 1>
 }
 // CHECK: [[INPUT:%.*]] = builtin.unrealized_conversion_cast %{{.*}}
 //  CHECK: [[ALLOC:%.*]] = llvm.extractvalue [[INPUT]][0]
 //  CHECK: [[ALIGN:%.*]] = llvm.extractvalue [[INPUT]][1]

mlir/test/Conversion/MemRefToLLVM/typed-pointers.mlir

 // CHECK:   [[TARGET_STRIDE_PTR:%.*]] = llvm.getelementptr [[STRIDES_PTR]]{{\[}}[[DIM]]]
 // CHECK:   llvm.store [[CUR_STRIDE]], [[TARGET_STRIDE_PTR]] : !llvm.ptr<i64>
 // CHECK:   [[UPDATE_STRIDE:%.*]] = llvm.mul [[CUR_STRIDE]], [[SIZE]] : i64
 // CHECK:   [[STRIDE_COND:%.*]] = llvm.sub [[DIM]], [[C1_]] : i64
 // CHECK:   llvm.br ^bb1([[STRIDE_COND]], [[UPDATE_STRIDE]] : i64, i64)
 
 // CHECK: ^bb3:
 // CHECK:   return
-
-// -----
-
-// CHECK-LABEL:     func.func @realloc_dynamic(
-// CHECK-SAME:      %[[arg0:.*]]: memref<?xf32>,
-// CHECK-SAME:      %[[arg1:.*]]: index) -> memref<?xf32> {
-func.func @realloc_dynamic(%in: memref<?xf32>, %d: index) -> memref<?xf32>{
-// CHECK:           %[[descriptor:.*]] = builtin.unrealized_conversion_cast %[[arg0]]
-// CHECK:           %[[src_dim:.*]] = llvm.extractvalue %[[descriptor]][3, 0]
-// CHECK:           %[[dst_dim:.*]] = builtin.unrealized_conversion_cast %[[arg1]] : index to i64
-// CHECK:           %[[cond:.*]] = llvm.icmp "ugt" %[[dst_dim]], %[[src_dim]] : i64
-// CHECK:           llvm.cond_br %[[cond]], ^bb1, ^bb2(%[[descriptor]]
-// CHECK:           ^bb1:
-// CHECK:           %[[dst_null:.*]] = llvm.mlir.null : !llvm.ptr<f32>
-// CHECK:           %[[dst_gep:.*]] = llvm.getelementptr %[[dst_null]][1]
-// CHECK:           %[[dst_es:.*]] = llvm.ptrtoint %[[dst_gep]] : !llvm.ptr<f32> to i64
-// CHECK:           %[[dst_size:.*]] = llvm.mul %[[dst_dim]], %[[dst_es]]
-// CHECK:           %[[src_size:.*]] = llvm.mul %[[src_dim]], %[[dst_es]]
-// CHECK:           %[[new_buffer_raw:.*]] = llvm.call @malloc(%[[dst_size]])
-// CHECK:           %[[new_buffer:.*]] = llvm.bitcast %[[new_buffer_raw]] : !llvm.ptr<i8> to !llvm.ptr<f32>
-// CHECK:           %[[old_buffer_aligned:.*]] = llvm.extractvalue %[[descriptor]][1]
-// CHECK-DAG:       %[[new_buffer_void:.*]] = llvm.bitcast %[[new_buffer]] : !llvm.ptr<f32> to !llvm.ptr<i8>
-// CHECK-DAG:       %[[old_buffer_void:.*]] = llvm.bitcast %[[old_buffer_aligned]] : !llvm.ptr<f32> to !llvm.ptr<i8>
-// CHECK:           "llvm.intr.memcpy"(%[[new_buffer_void]], %[[old_buffer_void]], %[[src_size]]) <{isVolatile = false}>
-// CHECK:           %[[old_buffer_unaligned:.*]] = llvm.extractvalue %[[descriptor]][0]
-// CHECK:           %[[old_buffer_unaligned_void:.*]] = llvm.bitcast %[[old_buffer_unaligned]] : !llvm.ptr<f32> to !llvm.ptr<i8>
-// CHECK:           llvm.call @free(%[[old_buffer_unaligned_void]])
-// CHECK:           %[[descriptor_update1:.*]] = llvm.insertvalue %[[new_buffer]], %[[descriptor]][0]
-// CHECK:           %[[descriptor_update2:.*]] = llvm.insertvalue %[[new_buffer]], %[[descriptor_update1]][1]
-// CHECK:           llvm.br ^bb2(%[[descriptor_update2]]
-// CHECK:           ^bb2(%[[descriptor_update3:.*]]: !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)>):
-// CHECK:           %[[descriptor_update4:.*]] = llvm.insertvalue %[[dst_dim]], %[[descriptor_update3]][3, 0]
-// CHECK:           %[[descriptor_update5:.*]] = builtin.unrealized_conversion_cast %[[descriptor_update4]]
-// CHECK:           return %[[descriptor_update5]] : memref<?xf32>
-
-  %out = memref.realloc %in(%d) : memref<?xf32> to memref<?xf32>
-  return %out : memref<?xf32>
-}
-

mlir/test/Dialect/MemRef/expand-realloc.mlir

+// RUN: mlir-opt --expand-realloc %s --split-input-file --verify-diagnostics | FileCheck %s
+// RUN: mlir-opt --expand-realloc=do-not-deallocate %s --split-input-file --verify-diagnostics | FileCheck %s --check-prefix=NODEALLOC
+
+func.func @reallow_lowering_example(%init_size: index, %new_size: index) -> memref<?xf32> {
+  %alloc = memref.alloc(%init_size) : memref<?xf32>
+  %realloc = memref.realloc %alloc (%new_size) {alignment = 8}: memref<?xf32> to memref<?xf32>
+  return %realloc : memref<?xf32>
+}
+
+// CHECK-LABEL: func @reallow_lowering_example
+//  CHECK-SAME: ([[INIT_SIZE:%.+]]: index, [[NEW_SIZE:%.+]]: index)
+//  CHECK-NEXT:   [[OLD_ALLOC:%.+]] = memref.alloc([[INIT_SIZE]]) : memref<?xf32>
+//  CHECK-NEXT:   [[C0:%.+]] = arith.constant 0
+//  CHECK-NEXT:   [[CURR_SIZE:%.+]] = memref.dim [[OLD_ALLOC]], [[C0]]
+//  CHECK-NEXT:   [[COND:%.+]] = arith.cmpi ult, [[CURR_SIZE]], [[NEW_SIZE]]
+//  CHECK-NEXT:   [[REALLOC:%.+]] = scf.if [[COND]]
+//  CHECK-NEXT:     [[NEW_ALLOC:%.+]] = memref.alloc([[NEW_SIZE]]) {alignment = 8 : i64} : memref<?xf32>
+//  CHECK-NEXT:     [[SUBVIEW:%.+]] = memref.subview [[NEW_ALLOC]][0] [[[CURR_SIZE]]] [1]
+//  CHECK-NEXT:     memref.copy [[OLD_ALLOC]], [[SUBVIEW]]
+//  CHECK-NEXT:     memref.dealloc [[OLD_ALLOC]]
+//  CHECK-NEXT:     scf.yield [[NEW_ALLOC]]
+//       CHECK:     [[CAST:%.+]] = memref.reinterpret_cast [[OLD_ALLOC]] to offset: [0], sizes: [[[NEW_SIZE]]], strides: [1] : memref<?xf32> to memref<?xf32>
+//  CHECK-NEXT:     scf.yield [[CAST]]
+//       CHECK:   return [[REALLOC]]
+
+// NODEALLOC-LABEL: func @reallow_lowering_example
+// NODEALLOC-NOT: dealloc
+
+// -----
+
+func.func @reallow_lowering_example() -> memref<4xf32> {
+  %alloc = memref.alloc() : memref<2xf32>
+  %realloc = memref.realloc %alloc {alignment = 8}: memref<2xf32> to memref<4xf32>
+  return %realloc : memref<4xf32>
+}
+
+// CHECK-LABEL: func @reallow_lowering_example
+//  CHECK-NEXT:   [[OLD_ALLOC:%.+]] = memref.alloc() : memref<2xf32>
+//  CHECK-NEXT:   [[CURR_SIZE:%.+]] = arith.constant 2
+//  CHECK-NEXT:   [[NEW_SIZE:%.+]] = arith.constant 4
+//  CHECK-NEXT:   [[COND:%.+]] = arith.cmpi ult, [[CURR_SIZE]], [[NEW_SIZE]]
+//  CHECK-NEXT:   [[REALLOC:%.+]] = scf.if [[COND]]
+//  CHECK-NEXT:     [[NEW_ALLOC:%.+]] = memref.alloc() {alignment = 8 : i64} : memref<4xf32>
+//  CHECK-NEXT:     [[SUBVIEW:%.+]] = memref.subview [[NEW_ALLOC]][0] [2] [1]
+//  CHECK-NEXT:     memref.copy [[OLD_ALLOC]], [[SUBVIEW]]
+//  CHECK-NEXT:     memref.dealloc [[OLD_ALLOC]]
+//  CHECK-NEXT:     scf.yield [[NEW_ALLOC]]
+//       CHECK:     [[CAST:%.+]] = memref.reinterpret_cast [[OLD_ALLOC]] to offset: [0], sizes: [4], strides: [1] : memref<2xf32> to memref<4xf32>
+//  CHECK-NEXT:     scf.yield [[CAST]]
+//       CHECK:   return [[REALLOC]]
+
+// NODEALLOC-LABEL: func @reallow_lowering_example
+// NODEALLOC-NOT: dealloc
+
+// -----
+
+func.func @reallow_lowering_example(%init_size: index) -> memref<4xf32> {
+  %alloc = memref.alloc(%init_size) : memref<?xf32>
+  %realloc = memref.realloc %alloc : memref<?xf32> to memref<4xf32>
+  return %realloc : memref<4xf32>
+}
+
+// CHECK-LABEL: func @reallow_lowering_example
+//  CHECK-SAME: ([[INIT_SIZE:%.+]]: index)
+//  CHECK-NEXT:   [[OLD_ALLOC:%.+]] = memref.alloc([[INIT_SIZE]]) : memref<?xf32>
+//  CHECK-NEXT:   [[C0:%.+]] = arith.constant 0
+//  CHECK-NEXT:   [[CURR_SIZE:%.+]] = memref.dim [[OLD_ALLOC]], [[C0]]
+//  CHECK-NEXT:   [[NEW_SIZE:%.+]] = arith.constant 4
+//  CHECK-NEXT:   [[COND:%.+]] = arith.cmpi ult, [[CURR_SIZE]], [[NEW_SIZE]]
+//  CHECK-NEXT:   [[REALLOC:%.+]] = scf.if [[COND]]
+//  CHECK-NEXT:     [[NEW_ALLOC:%.+]] = memref.alloc() : memref<4xf32>
+//  CHECK-NEXT:     [[SUBVIEW:%.+]] = memref.subview [[NEW_ALLOC]][0] [[[CURR_SIZE]]] [1]
+//  CHECK-NEXT:     memref.copy [[OLD_ALLOC]], [[SUBVIEW]]
+//  CHECK-NEXT:     memref.dealloc [[OLD_ALLOC]]
+//  CHECK-NEXT:     scf.yield [[NEW_ALLOC]]
+//       CHECK:     [[CAST:%.+]] = memref.reinterpret_cast [[OLD_ALLOC]] to offset: [0], sizes: [4], strides: [1] : memref<?xf32> to memref<4xf32>
+//  CHECK-NEXT:     scf.yield [[CAST]]
+//       CHECK:   return [[REALLOC]]
+
+// NODEALLOC-LABEL: func @reallow_lowering_example
+// NODEALLOC-NOT: dealloc
+
+// -----
+
+func.func @reallow_lowering_example(%new_size: index) -> memref<?xf32> {
+  %alloc = memref.alloc() : memref<2xf32>
+  %realloc = memref.realloc %alloc(%new_size) : memref<2xf32> to memref<?xf32>
+  return %realloc : memref<?xf32>
+}
+
+// CHECK-LABEL: func @reallow_lowering_example
+//  CHECK-SAME: ([[NEW_SIZE:%.+]]: index)
+//  CHECK-NEXT:   [[OLD_ALLOC:%.+]] = memref.alloc() : memref<2xf32>
+//  CHECK-NEXT:   [[CURR_SIZE:%.+]] = arith.constant 2
+//  CHECK-NEXT:   [[COND:%.+]] = arith.cmpi ult, [[CURR_SIZE]], [[NEW_SIZE]]
+//  CHECK-NEXT:   [[REALLOC:%.+]] = scf.if [[COND]]
+//  CHECK-NEXT:     [[NEW_ALLOC:%.+]] = memref.alloc([[NEW_SIZE]]) : memref<?xf32>
+//  CHECK-NEXT:     [[SUBVIEW:%.+]] = memref.subview [[NEW_ALLOC]][0] [2] [1]
+//  CHECK-NEXT:     memref.copy [[OLD_ALLOC]], [[SUBVIEW]]
+//  CHECK-NEXT:     memref.dealloc [[OLD_ALLOC]]
+//  CHECK-NEXT:     scf.yield [[NEW_ALLOC]]
+//       CHECK:     [[CAST:%.+]] = memref.reinterpret_cast [[OLD_ALLOC]] to offset: [0], sizes: [[[NEW_SIZE]]], strides: [1] : memref<2xf32> to memref<?xf32>
+//  CHECK-NEXT:     scf.yield [[CAST]]
+//       CHECK:   return [[REALLOC]]
+
+// NODEALLOC-LABEL: func @reallow_lowering_example
+// NODEALLOC-NOT: dealloc

mlir/test/Integration/Dialect/Vector/CPU/test-realloc.mlir

-// RUN: mlir-opt %s -convert-vector-to-scf -convert-scf-to-cf -convert-vector-to-llvm -finalize-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts |\
+// RUN: mlir-opt %s -convert-vector-to-scf -expand-realloc -expand-strided-metadata -convert-scf-to-cf -convert-vector-to-llvm -finalize-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts |\
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_c_runner_utils
-// RUN: mlir-opt %s -convert-vector-to-scf -convert-scf-to-cf -convert-vector-to-llvm -finalize-memref-to-llvm='use-aligned-alloc=1' -convert-func-to-llvm -arith-expand -reconcile-unrealized-casts |\
+// RUN: mlir-opt %s -convert-vector-to-scf -expand-realloc -expand-strided-metadata -convert-scf-to-cf -convert-vector-to-llvm -finalize-memref-to-llvm='use-aligned-alloc=1' -convert-func-to-llvm -arith-expand -reconcile-unrealized-casts |\
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_c_runner_utils | FileCheck %s
 
 // FIXME: Windows does not have aligned_alloc
 // UNSUPPORTED: system-windows
 
 func.func @entry() {
   // Set up memory.

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":InferTypeOpInterface",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",
         ":MemRefUtils",
         ":NVGPUDialect",
         ":Pass",
         ":RuntimeVerifiableOpInterface",
+        ":SCFDialect",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":ValueBoundsOpInterface",
         ":VectorDialect",
         "//llvm:Support",
     ],
 )