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

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)

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D159430

Author: maerhart

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

@@ -21,8 +21,14 @@ 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
@@ -67,6 +73,10 @@ std::unique_ptr<Pass> createResolveShapedTypeResultDimsPass();
 /// 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 emitDeallocs = true);
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//

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

@@ -202,5 +202,48 @@ def ExpandStridedMetadata : Pass<"expand-strided-metadata"> {
       "affine::AffineDialect", "memref::MemRefDialect"
   ];
 }
+
+def ExpandRealloc : Pass<"expand-realloc"> {
+  let summary = "Expand memref.realloc operations into its components";
+  let description = [{
+    The `memref.realloc` operation performs a conditional allocation and copy to
+    increase the size of a buffer if necessary. This pass converts a `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<"emitDeallocs", "emit-deallocs", "bool", /*default=*/"true",
+           "Emit deallocation operations for 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

@@ -63,6 +63,10 @@ void populateExpandStridedMetadataPatterns(RewritePatternSet &patterns);
 /// `memref.extract_strided_metadata` of its source.
 void populateResolveExtractStridedMetadataPatterns(RewritePatternSet &patterns);
 
+/// Appends patterns for expanding `memref.realloc` operations.
+void populateExpandReallocPatterns(RewritePatternSet &patterns,
+                                   bool emitDeallocs = true);
+
 /// Appends patterns for emulating wide integer memref operations with ops over
 /// narrower integer types.
 void populateMemRefWideIntEmulationPatterns(

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

@@ -116,165 +116,6 @@ struct AllocaOpLowering : public AllocLikeOpLLVMLowering {
   }
 };
 
-/// 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;
@@ -1899,11 +1740,9 @@ void mlir::populateFinalizeMemRefToLLVMConversionPatterns(
   // 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 {

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

@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRMemRefTransforms
   BufferizableOpInterfaceImpl.cpp
   ComposeSubView.cpp
   ExpandOps.cpp
+  ExpandRealloc.cpp
   ExpandStridedMetadata.cpp
   EmulateWideInt.cpp
   EmulateNarrowType.cpp