[mlir][MemRef] Add required address space cast when lowering alloc to LLVM

[zero9178]

Summary

alloc uses either malloc or a plugable allocation function for allocating the required memory. Both of these functions always return a llvm.ptr<i8>, aka a pointer in the default address space. When allocating for a memref in a different memory space however, no address space cast is created, leading to invalid LLVM IR being generated.

This is currently not caught by the verifier since the pointer to the memory is always bitcast which currently lacks a verifier disallowing address space casts. Translating to actual LLVM IR would cause the verifier to go off, since bitcast cannot translate from one address space to another: https://godbolt.org/z/3a1z97rc9

This patch fixes that issue by generating an address space cast if the address space of the allocation function does not match the address space of the resulting memref.

Not sure whether this is actually a real life problem. I found this issue while converting the pass to using opaque pointers which gets rid of all the bitcasts and hence caused type errors without the address space cast.

[qcolombet]

Nice catch!

LGTM

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

Value AllocationOpLLVMLowering::createAligned(
     Value alignment) {
   Value one = createIndexAttrConstant(rewriter, loc, alignment.getType(), 1);
   Value bump = rewriter.create<LLVM::SubOp>(loc, alignment, one);
   Value bumped = rewriter.create<LLVM::AddOp>(loc, input, bump);
   Value mod = rewriter.create<LLVM::URemOp>(loc, bumped, alignment);
   return rewriter.create<LLVM::SubOp>(loc, bumped, mod);
 }
 
+static Value castAllocFuncResult(ConversionPatternRewriter &rewriter,
+                                 Location loc, Value allocatedPtr,
+                                 MemRefType memRefType, Type elementPtrType,
+                                 LLVMTypeConverter &typeConverter) {
+  auto allocatedPtrTy = allocatedPtr.getType().cast<LLVM::LLVMPointerType>();
+  if (allocatedPtrTy.getAddressSpace() != memRefType.getMemorySpaceAsInt())
+    allocatedPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
+        loc,
+        LLVM::LLVMPointerType::get(allocatedPtrTy.getElementType(),
+                                   memRefType.getMemorySpaceAsInt()),
+        allocatedPtr);
+
+  allocatedPtr =
+      rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, allocatedPtr);
+  return allocatedPtr;
+}
+
 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 allocatedPtr =
+      castAllocFuncResult(rewriter, loc, results.getResult(), memRefType,
+                          elementPtrType, *getTypeConverter());
 
   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 =
Value AllocationOpLLVMLowering::allocateBufferAutoAlign(
   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;
+  return castAllocFuncResult(rewriter, loc, results.getResult(), memRefType,
+                             elementPtrType, *getTypeConverter());
 }
 
 LogicalResult AllocLikeOpLLVMLowering::matchAndRewrite(
     Operation *op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   MemRefType memRefType = getMemRefResultType(op);
   if (!isConvertibleAndHasIdentityMaps(memRefType))
     return rewriter.notifyMatchFailure(op, "incompatible memref type");

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

 // CHECK: %[[SIZE:.*]] = llvm.load %[[SIZE_PTR]] : !llvm.ptr<i64>
 
 // CHECK32: %[[SIZE:.*]] = llvm.load %{{.*}} : !llvm.ptr<i32>
 
 // -----
 
 // CHECK-LABEL: func @address_space(
 func.func @address_space(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>, 7>) {
+  // CHECK: %[[MEMORY:.*]] = llvm.call @malloc(%{{.*}})
+  // CHECK: %[[CAST:.*]] = llvm.addrspacecast %[[MEMORY]] : !llvm.ptr<i8> to !llvm.ptr<i8, 5>
+  // CHECK: %[[BCAST:.*]] = llvm.bitcast %[[CAST]]
+  // CHECK: llvm.insertvalue %[[BCAST]], %{{[[:alnum:]]+}}[0]
+  // CHECK: llvm.insertvalue %[[BCAST]], %{{[[:alnum:]]+}}[1]
   %0 = memref.alloc() : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   %1 = arith.constant 7 : index
   // CHECK: llvm.load %{{.*}} : !llvm.ptr<f32, 5>
   %2 = memref.load %0[%1] : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   func.return
 }
 
 // -----