Repository: llvm/llvm-project
File: mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (C++)

Code before:
  matchAndRewrite(memref::AllocOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    MemRefType memRefType = op.getType();
    if (!isConvertibleAndHasIdentityMaps(memRefType))
      return rewriter.notifyMatchFailure(op, "incompatible memref type");

    // Get or insert alloc function into the module.
    FailureOr<LLVM::LLVMFuncOp> allocFuncOp =
        getNotalignedAllocFn(rewriter, getTypeConverter(),
                             op->getParentWithTrait<OpTrait::SymbolTable>(),
                             getIndexType(), symbolTables);
    if (failed(allocFuncOp))
      return failure();

    // Get actual sizes of the memref as values: static sizes are constant
    // values and dynamic sizes are passed to 'alloc' as operands.  In case of
    // zero-dimensional memref, assume a scalar (size 1).
    SmallVector<Value, 4> sizes;
    SmallVector<Value, 4> strides;
    Value sizeBytes;

    if (failed(this->getMemRefDescriptorSizes(loc, memRefType,
                                              adaptor.getOperands(), rewriter,
                                              sizes, strides, sizeBytes, true)))
      return rewriter.notifyMatchFailure(
          op, "unable to compute memref descriptor sizes");

    Value alignment = getAlignment(rewriter, loc, op);
    if (alignment) {
      // Adjust the allocation size to consider alignment.
      sizeBytes = LLVM::AddOp::create(rewriter, loc, sizeBytes, alignment);
    }

    // Allocate the underlying buffer.
    Type elementPtrType = this->getElementPtrType(memRefType);
    assert(elementPtrType && "could not compute element ptr type");
    auto results =
        LLVM::CallOp::create(rewriter, loc, allocFuncOp.value(), sizeBytes);

    Value allocatedPtr =
        castAllocFuncResult(rewriter, loc, results.getResult(), memRefType,
                            elementPtrType, *getTypeConverter());
    Value alignedPtr = allocatedPtr;
    if (alignment) {
      // Compute the aligned pointer.
      Value allocatedInt =
          LLVM::PtrToIntOp::create(rewriter, loc, getIndexType(), allocatedPtr);
      Value alignmentInt =
          createAligned(rewriter, loc, allocatedInt, alignment);
      alignedPtr =

Review comment (bug):
Here and in the other patterns: `matchAndRewrite` [must not return "failure" when the IR has already been modified](https://mlir.llvm.org/docs/PatternRewriter/#restrictions). This "if" check must be before `getNotalignedAllocFn`.

However, `getNotalignedAllocFn` can also fail, so it must be before `getMemRefDescriptorSizes`.

The size check must be in a separate function that returns `LogicalResult`. `getMemRefDescriptorSizes` should keep returning `void`.


Code after:
  matchAndRewrite(memref::AllocOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    MemRefType memRefType = op.getType();
    if (!isConvertibleAndHasIdentityMaps(memRefType))
      return rewriter.notifyMatchFailure(op, "incompatible memref type");

    // Get or insert alloc function into the module.
    FailureOr<LLVM::LLVMFuncOp> allocFuncOp =
        getNotalignedAllocFn(rewriter, getTypeConverter(),
                             op->getParentWithTrait<OpTrait::SymbolTable>(),
                             getIndexType(), symbolTables);
    if (failed(allocFuncOp))
      return failure();

    // Get actual sizes of the memref as values: static sizes are constant
    // values and dynamic sizes are passed to 'alloc' as operands.  In case of
    // zero-dimensional memref, assume a scalar (size 1).
    SmallVector<Value, 4> sizes;
    SmallVector<Value, 4> strides;
    Value sizeBytes;

    this->getMemRefDescriptorSizes(loc, memRefType, adaptor.getOperands(),
                                   rewriter, sizes, strides, sizeBytes, true);

    Value alignment = getAlignment(rewriter, loc, op);
    if (alignment) {
      // Adjust the allocation size to consider alignment.
      sizeBytes = LLVM::AddOp::create(rewriter, loc, sizeBytes, alignment);
    }

    // Allocate the underlying buffer.
    Type elementPtrType = this->getElementPtrType(memRefType);
    assert(elementPtrType && "could not compute element ptr type");
    auto results =
        LLVM::CallOp::create(rewriter, loc, allocFuncOp.value(), sizeBytes);

    Value allocatedPtr =
        castAllocFuncResult(rewriter, loc, results.getResult(), memRefType,
                            elementPtrType, *getTypeConverter());
    Value alignedPtr = allocatedPtr;
    if (alignment) {
      // Compute the aligned pointer.
      Value allocatedInt =
          LLVM::PtrToIntOp::create(rewriter, loc, getIndexType(), allocatedPtr);
      Value alignmentInt =
          createAligned(rewriter, loc, allocatedInt, alignment);
      alignedPtr =
          LLVM::IntToPtrOp::create(rewriter, loc, elementPtrType, alignmentInt);
    }