[mlir][memref] Add realloc op.

Add memref.realloc and canonicalization of the op. Add conversion patterns for
lowering the op to LLVM using unaligned alloc or aligned alloc based on the
conversion option.

Add filecheck tests for parsing and converting the op. Add an integration test.

Reviewed By: ftynse

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

Author: bixia1

mlir/include/mlir/Conversion/MemRefToLLVM/AllocLikeConversion.h

@@ -13,34 +13,107 @@
 
 namespace mlir {
 
-/// Lowering for AllocOp and AllocaOp.
-struct AllocLikeOpLLVMLowering : public ConvertToLLVMPattern {
+/// Lowering for memory allocation ops.
+struct AllocationOpLLVMLowering : public ConvertToLLVMPattern {
   using ConvertToLLVMPattern::createIndexConstant;
   using ConvertToLLVMPattern::getIndexType;
   using ConvertToLLVMPattern::getVoidPtrType;
 
-  explicit AllocLikeOpLLVMLowering(StringRef opName,
-                                   LLVMTypeConverter &converter)
+  explicit AllocationOpLLVMLowering(StringRef opName,
+                                    LLVMTypeConverter &converter)
       : ConvertToLLVMPattern(opName, &converter.getContext(), converter) {}
 
 protected:
-  // Returns 'input' aligned up to 'alignment'. Computes
-  // bumped = input + alignement - 1
-  // aligned = bumped - bumped % alignment
+  /// Computes the aligned value for 'input' as follows:
+  ///   bumped = input + alignement - 1
+  ///   aligned = bumped - bumped % alignment
   static Value createAligned(ConversionPatternRewriter &rewriter, Location loc,
                              Value input, Value alignment);
 
+  static MemRefType getMemRefResultType(Operation *op) {
+    return op->getResult(0).getType().cast<MemRefType>();
+  }
+
+  /// Computes the alignment for the given memory allocation op.
+  template <typename OpType>
+  Value getAlignment(ConversionPatternRewriter &rewriter, Location loc,
+                     OpType op) const {
+    MemRefType memRefType = op.getType();
+    Value alignment;
+    if (auto alignmentAttr = op.getAlignment()) {
+      alignment = createIndexConstant(rewriter, loc, *alignmentAttr);
+    } else if (!memRefType.getElementType().isSignlessIntOrIndexOrFloat()) {
+      // In the case where no alignment is specified, we may want to override
+      // `malloc's` behavior. `malloc` typically aligns at the size of the
+      // biggest scalar on a target HW. For non-scalars, use the natural
+      // alignment of the LLVM type given by the LLVM DataLayout.
+      alignment = getSizeInBytes(loc, memRefType.getElementType(), rewriter);
+    }
+    return alignment;
+  }
+
+  /// Computes the alignment for aligned_alloc used to allocate the buffer for
+  /// the memory allocation op.
+  ///
+  /// Aligned_alloc requires the allocation size to be a power of two, and the
+  /// allocation size to be a multiple of the alignment.
+  template <typename OpType>
+  int64_t alignedAllocationGetAlignment(ConversionPatternRewriter &rewriter,
+                                        Location loc, OpType op,
+                                        const DataLayout *defaultLayout) const {
+    if (Optional<uint64_t> alignment = op.getAlignment())
+      return *alignment;
+
+    // Whenever we don't have alignment set, we will use an alignment
+    // consistent with the element type; since the allocation size has to be a
+    // power of two, we will bump to the next power of two if it isn't.
+    unsigned eltSizeBytes =
+        getMemRefEltSizeInBytes(op.getType(), op, defaultLayout);
+    return std::max(kMinAlignedAllocAlignment,
+                    llvm::PowerOf2Ceil(eltSizeBytes));
+  }
+
+  /// Allocates a memory buffer using an allocation method that doesn't
+  /// guarantee alignment. Returns the pointer and its aligned value.
+  std::tuple<Value, Value>
+  allocateBufferManuallyAlign(ConversionPatternRewriter &rewriter, Location loc,
+                              Value sizeBytes, Operation *op,
+                              Value alignment) const;
+
+  /// Allocates a memory buffer using an aligned allocation method.
+  Value allocateBufferAutoAlign(ConversionPatternRewriter &rewriter,
+                                Location loc, Value sizeBytes, Operation *op,
+                                const DataLayout *defaultLayout,
+                                int64_t alignment) const;
+
+private:
+  /// Computes the byte size for the MemRef element type.
+  unsigned getMemRefEltSizeInBytes(MemRefType memRefType, Operation *op,
+                                   const DataLayout *defaultLayout) const;
+
+  /// Returns true if the memref size in bytes is known to be a multiple of
+  /// factor.
+  bool isMemRefSizeMultipleOf(MemRefType type, uint64_t factor, Operation *op,
+                              const DataLayout *defaultLayout) const;
+
+  /// The minimum alignment to use with aligned_alloc (has to be a power of 2).
+  static constexpr uint64_t kMinAlignedAllocAlignment = 16UL;
+};
+
+/// Lowering for AllocOp and AllocaOp.
+struct AllocLikeOpLLVMLowering : public AllocationOpLLVMLowering {
+  explicit AllocLikeOpLLVMLowering(StringRef opName,
+                                   LLVMTypeConverter &converter)
+      : AllocationOpLLVMLowering(opName, converter) {}
+
+protected:
   /// Allocates the underlying buffer. Returns the allocated pointer and the
   /// aligned pointer.
   virtual std::tuple<Value, Value>
   allocateBuffer(ConversionPatternRewriter &rewriter, Location loc,
                  Value sizeBytes, Operation *op) const = 0;
 
 private:
-  static MemRefType getMemRefResultType(Operation *op) {
-    return op->getResult(0).getType().cast<MemRefType>();
-  }
-
   // An `alloc` is converted into a definition of a memref descriptor value and
   // a call to `malloc` to allocate the underlying data buffer.  The memref
   // descriptor is of the LLVM structure type where:

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

@@ -178,6 +178,99 @@ def MemRef_AllocOp : AllocLikeOp<"alloc", DefaultResource, []> {
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// ReallocOp
+//===----------------------------------------------------------------------===//
+
+
+def MemRef_ReallocOp : MemRef_Op<"realloc"> {
+  let summary = "memory reallocation operation";
+  let description = [{
+    The `realloc` operation changes the size of a memory region. The memory
+    region is specified by a 1D source memref and the size of the new memory
+    region is specified by a 1D result memref type and an optional dynamic Value
+    of `Index` type. The source and the result memref must be in the same memory
+    space and have the same element type.
+
+    The operation may move the memory region to a new location. In this case,
+    the content of the memory block is preserved up to the lesser of the new
+    and old sizes. If the new size if larger, the value of the extended memory
+    is undefined. This is consistent with the ISO C realloc.
+
+    The operation returns an SSA value for the memref.
+
+    Example:
+
+    ```mlir
+    %0 = memref.realloc %src : memref<64xf32> to memref<124xf32>
+    ```
+
+    The source memref may have a dynamic shape, in which case, the compiler will
+    generate code to extract its size from the runtime data structure for the
+    memref.
+
+    ```mlir
+    %1 = memref.realloc %src : memref<?xf32> to memref<124xf32>
+    ```
+
+    If the result memref has a dynamic shape, a result dimension operand is
+    needed to spefify its dynamic dimension. In the example below, the ssa value
+    '%d' specifies the unknown dimension of the result memref.
+
+    ```mlir
+    %2 = memref.realloc %src(%d) : memref<?xf32> to memref<?xf32>
+    ```
+
+    An optional `alignment` attribute may be specified to ensure that the
+    region of memory that will be indexed is aligned at the specified byte
+    boundary.  This is consistent with the fact that memref.alloc supports such
+    an optional alignment attribute. Note that in ISO C standard, neither alloc
+    nor realloc supports alignment, though there is aligned_alloc but not
+    aligned_realloc.
+
+    ```mlir
+    %3 = memref.ralloc %src {alignment = 8} : memref<64xf32> to memref<124xf32>
+    ```
+
+    Referencing the memref through the old SSA value after realloc is undefined
+    behavior.
+
+    ```mlir
+    %new = memref.realloc %old : memref<64xf32> to memref<124xf32>
+    %4 = memref.load %new[%index]   // ok
+    %5 = memref.load %old[%index]   // undefined behavior
+    ```
+  }];
+
+  let arguments = (ins MemRefRankOf<[AnyType], [1]>:$source,
+                   Optional<Index>:$dynamicResultSize,
+                   ConfinedAttr<OptionalAttr<I64Attr>,
+                                [IntMinValue<0>]>:$alignment);
+
+  let results = (outs MemRefRankOf<[AnyType], [1]>);
+
+  let builders = [
+    OpBuilder<(ins "MemRefType":$resultType,
+                  "Value":$source,
+                  CArg<"Value", "Value()">:$dynamicResultSize), [{
+      return build($_builder, $_state, resultType, source, dynamicResultSize,
+                   IntegerAttr());
+    }]>];
+
+    let extraClassDeclaration = [{
+    /// The result of a realloc is always a memref.
+    MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
+  }];
+
+  let assemblyFormat = [{
+    $source (`(` $dynamicResultSize^ `)`)? attr-dict
+    `:` type($source) `to` type(results)
+  }];
+
+  let hasCanonicalizer = 1;
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // AllocaOp
 //===----------------------------------------------------------------------===//

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

@@ -7,11 +7,40 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/MemRefToLLVM/AllocLikeConversion.h"
+#include "mlir/Analysis/DataLayoutAnalysis.h"
+#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 
 using namespace mlir;
 
-Value AllocLikeOpLLVMLowering::createAligned(
+namespace {
+// TODO: Fix the LLVM utilities for looking up functions to take Operation*
+// with SymbolTable trait instead of ModuleOp and make similar change here. This
+// allows call sites to use getParentWithTrait<OpTrait::SymbolTable> instead
+// of getParentOfType<ModuleOp> to pass down the operation.
+LLVM::LLVMFuncOp getNotalignedAllocFn(LLVMTypeConverter *typeConverter,
+                                      ModuleOp module, Type indexType) {
+  bool useGenericFn = typeConverter->getOptions().useGenericFunctions;
+
+  if (useGenericFn)
+    return LLVM::lookupOrCreateGenericAllocFn(module, indexType);
+
+  return LLVM::lookupOrCreateMallocFn(module, indexType);
+}
+
+LLVM::LLVMFuncOp getAlignedAllocFn(LLVMTypeConverter *typeConverter,
+                                   ModuleOp module, Type indexType) {
+  bool useGenericFn = typeConverter->getOptions().useGenericFunctions;
+
+  if (useGenericFn)
+    return LLVM::lookupOrCreateGenericAlignedAllocFn(module, indexType);
+
+  return LLVM::lookupOrCreateAlignedAllocFn(module, indexType);
+}
+
+} // end namespace
+
+Value AllocationOpLLVMLowering::createAligned(
     ConversionPatternRewriter &rewriter, Location loc, Value input,
     Value alignment) {
   Value one = createIndexAttrConstant(rewriter, loc, alignment.getType(), 1);
@@ -21,6 +50,88 @@ Value AllocLikeOpLLVMLowering::createAligned(
   return rewriter.create<LLVM::SubOp>(loc, bumped, mod);
 }
 
+std::tuple<Value, Value> AllocationOpLLVMLowering::allocateBufferManuallyAlign(
+    ConversionPatternRewriter &rewriter, Location loc, Value sizeBytes,
+    Operation *op, Value alignment) const {
+  if (alignment) {
+    // Adjust the allocation size to consider alignment.
+    sizeBytes = rewriter.create<LLVM::AddOp>(loc, sizeBytes, alignment);
+  }
+
+  MemRefType memRefType = getMemRefResultType(op);
+  // Allocate the underlying buffer.
+  Type elementPtrType = this->getElementPtrType(memRefType);
+  LLVM::LLVMFuncOp allocFuncOp = getNotalignedAllocFn(
+      getTypeConverter(), op->getParentOfType<ModuleOp>(), getIndexType());
+  auto results = rewriter.create<LLVM::CallOp>(loc, allocFuncOp, sizeBytes);
+  Value allocatedPtr = rewriter.create<LLVM::BitcastOp>(loc, elementPtrType,
+                                                        results.getResult());
+
+  Value alignedPtr = allocatedPtr;
+  if (alignment) {
+    // Compute the aligned pointer.
+    Value allocatedInt =
+        rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), allocatedPtr);
+    Value alignmentInt = createAligned(rewriter, loc, allocatedInt, alignment);
+    alignedPtr =
+        rewriter.create<LLVM::IntToPtrOp>(loc, elementPtrType, alignmentInt);
+  }
+
+  return std::make_tuple(allocatedPtr, alignedPtr);
+}
+
+unsigned AllocationOpLLVMLowering::getMemRefEltSizeInBytes(
+    MemRefType memRefType, Operation *op,
+    const DataLayout *defaultLayout) const {
+  const DataLayout *layout = defaultLayout;
+  if (const DataLayoutAnalysis *analysis =
+          getTypeConverter()->getDataLayoutAnalysis()) {
+    layout = &analysis->getAbove(op);
+  }
+  Type elementType = memRefType.getElementType();
+  if (auto memRefElementType = elementType.dyn_cast<MemRefType>())
+    return getTypeConverter()->getMemRefDescriptorSize(memRefElementType,
+                                                       *layout);
+  if (auto memRefElementType = elementType.dyn_cast<UnrankedMemRefType>())
+    return getTypeConverter()->getUnrankedMemRefDescriptorSize(
+        memRefElementType, *layout);
+  return layout->getTypeSize(elementType);
+}
+
+bool AllocationOpLLVMLowering::isMemRefSizeMultipleOf(
+    MemRefType type, uint64_t factor, Operation *op,
+    const DataLayout *defaultLayout) const {
+  uint64_t sizeDivisor = getMemRefEltSizeInBytes(type, op, defaultLayout);
+  for (unsigned i = 0, e = type.getRank(); i < e; i++) {
+    if (ShapedType::isDynamic(type.getDimSize(i)))
+      continue;
+    sizeDivisor = sizeDivisor * type.getDimSize(i);
+  }
+  return sizeDivisor % factor == 0;
+}
+
+Value AllocationOpLLVMLowering::allocateBufferAutoAlign(
+    ConversionPatternRewriter &rewriter, Location loc, Value sizeBytes,
+    Operation *op, const DataLayout *defaultLayout, int64_t alignment) const {
+  Value allocAlignment = createIndexConstant(rewriter, loc, alignment);
+
+  MemRefType memRefType = getMemRefResultType(op);
+  // Function aligned_alloc requires size to be a multiple of alignment; we pad
+  // the size to the next multiple if necessary.
+  if (!isMemRefSizeMultipleOf(memRefType, alignment, op, defaultLayout))
+    sizeBytes = createAligned(rewriter, loc, sizeBytes, allocAlignment);
+
+  Type elementPtrType = this->getElementPtrType(memRefType);
+  LLVM::LLVMFuncOp allocFuncOp = getAlignedAllocFn(
+      getTypeConverter(), op->getParentOfType<ModuleOp>(), getIndexType());
+  auto results = rewriter.create<LLVM::CallOp>(
+      loc, allocFuncOp, ValueRange({allocAlignment, sizeBytes}));
+  Value allocatedPtr = rewriter.create<LLVM::BitcastOp>(loc, elementPtrType,
+                                                        results.getResult());
+
+  return allocatedPtr;
+}
+
 LogicalResult AllocLikeOpLLVMLowering::matchAndRewrite(
     Operation *op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {