[mlir][memref] Fix num elements in lowering of memref.alloca op to LLVM

[fmorac]

Summary

Fixes a mistake in the lowering of memref.alloca to llvm.alloca, as llvm.alloca uses the number of elements to allocate in the stack and not the size in bytes.

Reference:
LLVM IR: https://llvm.org/docs/LangRef.html#alloca-instruction
LLVM MLIR: https://mlir.llvm.org/docs/Dialects/LLVM/#llvmalloca-mlirllvmallocaop

[fmorac]

Fixed the tests.

[fmorac]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

108

The bug happens here, it shouldn't be sizeBytes, instead it should be sizeBytes / sizeof(elementType).

[ftynse]

We should be already computing the number of elements to get the size in bytes. I don't think we want to first multiply that with the byte size only to divide it back. Consider refactoring allocateBuffer so it accepts the number of elements instead.

[fmorac]

In D150705#4349828, @ftynse wrote:

We should be already computing the number of elements to get the size in bytes. I don't think we want to first multiply that with the byte size only to divide it back. Consider refactoring allocateBuffer so it accepts the number of elements instead.

I agree with you, this was the quick fix. A proper solution requires modifying more parts.

We first need a function to compute the numElements in general:

As getMemRefDescriptorSizes always return the size in bytes, and getNumElements doesn't automatically handle static shapes.
https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/LLVMCommon/Pattern.cpp#L121-L137

I would introduce a new getNumElements with a similar signature as getMemRefDescriptorSizes to compute the number of elements in general.

Since AllocaOpLowering inherits from AllocLikeOpLLVMLowering, either we modify AllocLikeOpLLVMLowering or stop inheriting from it. I prefer stop inheriting from it, as it's alloca is not exactly alloc like.
And the quick fix for modifying AllocLikeOpLLVMLowering is introducing a conditional for testing AllocaOp, which I'm not to thrilled.
https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp#L159-L187

My proposal is then: introducing a new getNumElements, stop inheriting from AllocLikeOpLLVMLowering.

[fmorac]

Refactors allocateBuffer to use numElements. For achieving the refactor the function ConvertToLLVMPattern::getMemRefDescriptorSizes was modified to accept an optional parameter specifying to return numElements instead of sizeInBytes.

[ftynse]

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

171 ↗

(On Diff #524318)

Nit: please use braces symmetrically.

201–212 ↗

(On Diff #524318)

This will no longer handle 0D memrefs correctly, which was special-cased before.

203 ↗

(On Diff #524318)

Nit: please expand auto unless the type is obvious from context (e.g., there is a cast on the RHS).

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

177 ↗

(On Diff #524318)

Nit: leftover commented-out code.

[fmorac]

Fixed comments.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
     // 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 sizeOfElement = getSizeInBytes(loc, elementType, rewriter);
+    auto numElements =
+        rewriter.create<LLVM::UDivOp>(loc, sizeBytes, sizeOfElement);
 
     auto allocatedElementPtr = rewriter.create<LLVM::AllocaOp>(
-        loc, elementPtrType, elementType, sizeBytes,
+        loc, elementPtrType, elementType, numElements,
         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.