[mlir] Handle different pointer sizes in unranked memref descriptors

[krzysz00]

Summary

The code for unranked memref descriptors assumed that
sizeof(!llvm.ptr) == lizeof(!llvm.ptr<N>) for all address spaces N.
This is not always true (ex. the AMDGPU compiler backend has
sizeof(!llvm.ptr) = 64 bits but sizeof(!llvm.ptr<5>) = 32 bits, where
address space 5 is used for stack allocations). While this is merely
an overallocation in the case where a non-0 address space has pointers
smaller than the default, the existing code could cause OOB memory
accesses when sizeof(!llvm.ptr<N>) > sizeof(!llvm.ptr).

So, add an address spaces parameter to computeSizes in order to
partially resolve this class of bugs. Note that the LLVM data layout
in the conversion passes is currently set to "" and not constructed
from the MLIR data layout or some other source, but this could change
in the future.

Depends on D142159

[ftynse]

I'm not sure this will in fact work for the GPU case that is used as a motivation. Could you add a test for that?

mlir/lib/Conversion/LLVMCommon/MemRefBuilder.cpp

361

Nit: we can use c++17 structured bindings now.

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

235–238

Nit: let's drop the no-longer necessary number of stack elements since we are here.

240

This wouldn't work with non-integer memory spaces such as the ones we now use for GPU, which seems to be the motivation for this patch.

[krzysz00]

I'll note that the non-integer GPU address spaces are handled not in the type converter but by running a pass before the *ToLLVM calls get their hands no the types, so the assumpiton that getMemorySpaceAsInt() works is currently valid ... but it's probably not a bad idea to account for the more general case now just in case.

[krzysz00]

Address style comments, add a few consts on getters

[krzysz00]

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

240

To rephrase, the type converter for MemRefType already assumes that getMemorySpaceAsInt() works. Getting that cleaned up would probably require a broader refactoring of the [LLVM] type converter to keep or allow creating an address space -> int mapping.

[ftynse]

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

240

I understand that the converter already assumes the int address space (it shouldn't, we should handle that more gracefully than asserting there, but it's irrelevant for this patch). Here, we should have access to the converted pointer type though, so we could rather extract the integer address space from the LLVM pointer type and avoid propagating the assumption about it being an integer in all memrefs here.

[krzysz00]

Swap getMemorySpaceAsInt() for the more general mechanism from the previous patch.

[krzysz00]

Update new test, fix -Wunused-variable from previous patch

mlir/include/mlir/Conversion/LLVMCommon/MemRefBuilder.h

 public:
   /// Construct a helper for the given descriptor value.
   explicit UnrankedMemRefDescriptor(Value descriptor);
   /// Builds IR creating an `undef` value of the descriptor type.
   static UnrankedMemRefDescriptor undef(OpBuilder &builder, Location loc,
                                         Type descriptorType);
 
   /// Builds IR extracting the rank from the descriptor
-  Value rank(OpBuilder &builder, Location loc);
+  Value rank(OpBuilder &builder, Location loc) const;
   /// Builds IR setting the rank in the descriptor
   void setRank(OpBuilder &builder, Location loc, Value value);
   /// Builds IR extracting ranked memref descriptor ptr
-  Value memRefDescPtr(OpBuilder &builder, Location loc);
+  Value memRefDescPtr(OpBuilder &builder, Location loc) const;
   /// Builds IR setting ranked memref descriptor ptr
   void setMemRefDescPtr(OpBuilder &builder, Location loc, Value value);
 
   /// Builds IR populating an unranked MemRef descriptor structure from a list
   /// of individual constituent values in the following order:
   /// - rank of the memref;
   /// - pointer to the memref descriptor.
   static Value pack(OpBuilder &builder, Location loc,
                     LLVMTypeConverter &converter, UnrankedMemRefType type,
                     ValueRange values);
 
   /// Builds IR extracting individual elements that compose an unranked memref
   /// descriptor and returns them as `results` list.
   static void unpack(OpBuilder &builder, Location loc, Value packed,
                      SmallVectorImpl<Value> &results);
 
   /// Returns the number of non-aggregate values that would be produced by
   /// `unpack`.
   static unsigned getNumUnpackedValues() { return 2; }
 
   /// Builds IR computing the sizes in bytes (suitable for opaque allocation)
-  /// and appends the corresponding values into `sizes`.
+  /// and appends the corresponding values into `sizes`. `addressSpaces`
+  /// which must have the same length as `values`, is needed to handle layouts
+  /// where sizeof(ptr addrspace(N)) != sizeof(ptr addrspace(0)).
   static void computeSizes(OpBuilder &builder, Location loc,
                            LLVMTypeConverter &typeConverter,
                            ArrayRef<UnrankedMemRefDescriptor> values,
+                           ArrayRef<unsigned> addressSpaces,
                            SmallVectorImpl<Value> &sizes);
 
   /// TODO: The following accessors don't take alignment rules between elements
   /// of the descriptor struct into account. For some architectures, it might be
   /// necessary to extend them and to use `llvm::DataLayout` contained in
   /// `LLVMTypeConverter`.
 
   /// Builds IR extracting the allocated pointer from the descriptor.

mlir/include/mlir/Conversion/LLVMCommon/StructBuilder.h

 protected:
   // LLVM value
   Value value;
   // Cached struct type.
   Type structType;
 
 protected:
   /// Builds IR to extract a value from the struct at position pos
-  Value extractPtr(OpBuilder &builder, Location loc, unsigned pos);
+  Value extractPtr(OpBuilder &builder, Location loc, unsigned pos) const;
   /// Builds IR to set a value in the struct at position pos
   void setPtr(OpBuilder &builder, Location loc, unsigned pos, Value ptr);
 };
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_LLVMCOMMON_STRUCTBUILDER_H

mlir/lib/Conversion/LLVMCommon/MemRefBuilder.cpp

 
 /// Builds IR creating an `undef` value of the descriptor type.
 UnrankedMemRefDescriptor UnrankedMemRefDescriptor::undef(OpBuilder &builder,
                                                          Location loc,
                                                          Type descriptorType) {
   Value descriptor = builder.create<LLVM::UndefOp>(loc, descriptorType);
   return UnrankedMemRefDescriptor(descriptor);
 }
-Value UnrankedMemRefDescriptor::rank(OpBuilder &builder, Location loc) {
+Value UnrankedMemRefDescriptor::rank(OpBuilder &builder, Location loc) const {
   return extractPtr(builder, loc, kRankInUnrankedMemRefDescriptor);
 }
 void UnrankedMemRefDescriptor::setRank(OpBuilder &builder, Location loc,
                                        Value v) {
   setPtr(builder, loc, kRankInUnrankedMemRefDescriptor, v);
 }
 Value UnrankedMemRefDescriptor::memRefDescPtr(OpBuilder &builder,
-                                              Location loc) {
+                                              Location loc) const {
   return extractPtr(builder, loc, kPtrInUnrankedMemRefDescriptor);
 }
 void UnrankedMemRefDescriptor::setMemRefDescPtr(OpBuilder &builder,
                                                 Location loc, Value v) {
   setPtr(builder, loc, kPtrInUnrankedMemRefDescriptor, v);
 }
 
 /// Builds IR populating an unranked MemRef descriptor structure from a list
void UnrankedMemRefDescriptor::unpack(OpBuilder &builder, Location loc,
   UnrankedMemRefDescriptor d(packed);
   results.reserve(results.size() + 2);
   results.push_back(d.rank(builder, loc));
   results.push_back(d.memRefDescPtr(builder, loc));
 }
 
 void UnrankedMemRefDescriptor::computeSizes(
     OpBuilder &builder, Location loc, LLVMTypeConverter &typeConverter,
-    ArrayRef<UnrankedMemRefDescriptor> values, SmallVectorImpl<Value> &sizes) {
+    ArrayRef<UnrankedMemRefDescriptor> values, ArrayRef<unsigned> addressSpaces,
+    SmallVectorImpl<Value> &sizes) {
   if (values.empty())
     return;
-
+  assert(values.size() == addressSpaces.size() &&
+         "must provide address space for each descriptor");
   // Cache the index type.
   Type indexType = typeConverter.getIndexType();
 
   // Initialize shared constants.
   Value one = createIndexAttrConstant(builder, loc, indexType, 1);
   Value two = createIndexAttrConstant(builder, loc, indexType, 2);
-  Value pointerSize = createIndexAttrConstant(
-      builder, loc, indexType, ceilDiv(typeConverter.getPointerBitwidth(), 8));
   Value indexSize =
       createIndexAttrConstant(builder, loc, indexType,
                               ceilDiv(typeConverter.getIndexTypeBitwidth(), 8));
 
   sizes.reserve(sizes.size() + values.size());
-  for (UnrankedMemRefDescriptor desc : values) {
+  for (auto [desc, addressSpace] : llvm::zip(values, addressSpaces)) {
     // Emit IR computing the memory necessary to store the descriptor. This
     // assumes the descriptor to be
     //   { type*, type*, index, index[rank], index[rank] }
     // and densely packed, so the total size is
     //   2 * sizeof(pointer) + (1 + 2 * rank) * sizeof(index).
     // TODO: consider including the actual size (including eventual padding due
     // to data layout) into the unranked descriptor.
+    Value pointerSize = createIndexAttrConstant(
+        builder, loc, indexType,
+        ceilDiv(typeConverter.getPointerBitwidth(addressSpace), 8));
     Value doublePointerSize =
         builder.create<LLVM::MulOp>(loc, indexType, two, pointerSize);
 
     // (1 + 2 * rank) * sizeof(index)
     Value rank = desc.rank(builder, loc);
     Value doubleRank = builder.create<LLVM::MulOp>(loc, indexType, two, rank);
     Value doubleRankIncremented =
         builder.create<LLVM::AddOp>(loc, indexType, doubleRank, one);

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

 
 LogicalResult ConvertToLLVMPattern::copyUnrankedDescriptors(
     OpBuilder &builder, Location loc, TypeRange origTypes,
     SmallVectorImpl<Value> &operands, bool toDynamic) const {
   assert(origTypes.size() == operands.size() &&
          "expected as may original types as operands");
 
   // Find operands of unranked memref type and store them.
-  SmallVector<UnrankedMemRefDescriptor, 4> unrankedMemrefs;
-  for (unsigned i = 0, e = operands.size(); i < e; ++i)
-    if (origTypes[i].isa<UnrankedMemRefType>())
+  SmallVector<UnrankedMemRefDescriptor> unrankedMemrefs;
+  SmallVector<unsigned> unrankedAddressSpaces;
+  for (unsigned i = 0, e = operands.size(); i < e; ++i) {
+    if (auto memRefType = origTypes[i].dyn_cast<UnrankedMemRefType>()) {
       unrankedMemrefs.emplace_back(operands[i]);
+      FailureOr<unsigned> addressSpace =
+          getTypeConverter()->getMemRefAddressSpace(memRefType);
+      if (failed(addressSpace))
+        return failure();
+      unrankedAddressSpaces.emplace_back(*addressSpace);
+    }
+  }
 
   if (unrankedMemrefs.empty())
     return success();
 
   // Compute allocation sizes.
-  SmallVector<Value, 4> sizes;
+  SmallVector<Value> sizes;
   UnrankedMemRefDescriptor::computeSizes(builder, loc, *getTypeConverter(),
-                                         unrankedMemrefs, sizes);
+                                         unrankedMemrefs, unrankedAddressSpaces,
+                                         sizes);
 
   // Get frequently used types.
   MLIRContext *context = builder.getContext();
   auto i1Type = IntegerType::get(context, 1);
   Type indexType = getTypeConverter()->getIndexType();
 
   // Find the malloc and free, or declare them if necessary.
   auto module = builder.getInsertionPoint()->getParentOfType<ModuleOp>();

mlir/lib/Conversion/LLVMCommon/StructBuilder.cpp

 //===----------------------------------------------------------------------===//
 
 StructBuilder::StructBuilder(Value v) : value(v), structType(v.getType()) {
   assert(value != nullptr && "value cannot be null");
   assert(LLVM::isCompatibleType(structType) && "expected llvm type");
 }
 
 Value StructBuilder::extractPtr(OpBuilder &builder, Location loc,
-                                unsigned pos) {
+                                unsigned pos) const {
   return builder.create<LLVM::ExtractValueOp>(loc, value, pos);
 }
 
 void StructBuilder::setPtr(OpBuilder &builder, Location loc, unsigned pos,
                            Value ptr) {
   value = builder.create<LLVM::InsertValueOp>(loc, value, ptr, pos);
 }

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
 
     // Create the unranked memref descriptor that holds the ranked one. The
     // inner descriptor is allocated on stack.
     auto targetDesc = UnrankedMemRefDescriptor::undef(
         rewriter, loc, typeConverter->convertType(targetType));
     targetDesc.setRank(rewriter, loc, resultRank);
     SmallVector<Value, 4> sizes;
     UnrankedMemRefDescriptor::computeSizes(rewriter, loc, *getTypeConverter(),
-                                           targetDesc, sizes);
+                                           targetDesc, addressSpace, sizes);
     Value underlyingDescPtr = rewriter.create<LLVM::AllocaOp>(
         loc, getVoidPtrType(), IntegerType::get(getContext(), 8),
         sizes.front());
     targetDesc.setMemRefDescPtr(rewriter, loc, underlyingDescPtr);
 
     // Extract pointers and offset from the source memref.
     Value allocatedPtr, alignedPtr, offset;
     extractPointersAndOffset(loc, rewriter, *getTypeConverter(),

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

convertTransferReadToLoads(vector::TransferReadOp op,
                            llvm::DenseMap<Value, Value> &valueMapping) {
   OpBuilder b(op);
 
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo))
     return failure();
 
-  Attribute memorySpace =
-      op.getSource().getType().cast<MemRefType>().getMemorySpace();
   bool isLdMatrixCompatible =
       isSharedMemory(op.getSource().getType().cast<MemRefType>()) &&
       nvgpu::inferTileWidthInBits(*warpMatrixInfo) == 128;
 
   VectorType vecTy = op.getVectorType();
   int64_t bitWidth = vecTy.getElementType().getIntOrFloatBitWidth();
 
   // When we are transposing the B operand, ldmatrix will only work if we have

mlir/test/Conversion/FuncToLLVM/calling-convention.mlir

 // CHECK-LABEL: llvm.func @return_var_memref_caller
 func.func @return_var_memref_caller(%arg0: memref<4x3xf32>) {
   // CHECK: %[[CALL_RES:.*]] = llvm.call @return_var_memref
   %0 = call @return_var_memref(%arg0) : (memref<4x3xf32>) -> memref<*xf32>
 
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : index)
   // CHECK: %[[TWO:.*]] = llvm.mlir.constant(2 : index)
   // These sizes may depend on the data layout, not matching specific values.
-  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[IDX_SIZE:.*]] = llvm.mlir.constant
 
+  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[DOUBLE_PTR_SIZE:.*]] = llvm.mul %[[TWO]], %[[PTR_SIZE]]
   // CHECK: %[[RANK:.*]] = llvm.extractvalue %[[CALL_RES]][0] : !llvm.struct<(i64, ptr)>
   // CHECK: %[[DOUBLE_RANK:.*]] = llvm.mul %[[TWO]], %[[RANK]]
   // CHECK: %[[DOUBLE_RANK_INC:.*]] = llvm.add %[[DOUBLE_RANK]], %[[ONE]]
   // CHECK: %[[TABLES_SIZE:.*]] = llvm.mul %[[DOUBLE_RANK_INC]], %[[IDX_SIZE]]
   // CHECK: %[[ALLOC_SIZE:.*]] = llvm.add %[[DOUBLE_PTR_SIZE]], %[[TABLES_SIZE]]
   // CHECK: %[[FALSE:.*]] = llvm.mlir.constant(false)
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[ALLOC_SIZE]] x i8
func.func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> attributes { llvm.emit_c_interface } {
   // Match the construction of the unranked descriptor.
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca
   // CHECK: %[[RANK:.*]] = llvm.mlir.constant(2 : index)
   // CHECK: %[[DESC_0:.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr)>
   // CHECK: %[[DESC_1:.*]] = llvm.insertvalue %[[RANK]], %[[DESC_0]][0]
   // CHECK: %[[DESC_2:.*]] = llvm.insertvalue %[[ALLOCA]], %[[DESC_1]][1]
   %0 = memref.cast %arg0: memref<4x3xf32> to memref<*xf32>
 
-
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : index)
   // CHECK: %[[TWO:.*]] = llvm.mlir.constant(2 : index)
   // These sizes may depend on the data layout, not matching specific values.
-  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[IDX_SIZE:.*]] = llvm.mlir.constant
 
+  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[DOUBLE_PTR_SIZE:.*]] = llvm.mul %[[TWO]], %[[PTR_SIZE]]
   // CHECK: %[[DOUBLE_RANK:.*]] = llvm.mul %[[TWO]], %[[RANK]]
   // CHECK: %[[DOUBLE_RANK_INC:.*]] = llvm.add %[[DOUBLE_RANK]], %[[ONE]]
   // CHECK: %[[TABLES_SIZE:.*]] = llvm.mul %[[DOUBLE_RANK_INC]], %[[IDX_SIZE]]
   // CHECK: %[[ALLOC_SIZE:.*]] = llvm.add %[[DOUBLE_PTR_SIZE]], %[[TABLES_SIZE]]
   // CHECK: %[[FALSE:.*]] = llvm.mlir.constant(false)
   // CHECK: %[[ALLOCATED:.*]] = llvm.call @malloc(%[[ALLOC_SIZE]])
   // CHECK: "llvm.intr.memcpy"(%[[ALLOCATED]], %[[ALLOCA]], %[[ALLOC_SIZE]], %[[FALSE]])

mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir

 // CHECK: [[SHAPE:%.*]] = builtin.unrealized_conversion_cast %{{.*}} to [[SHAPE_TY:!.*]]
 // CHECK: [[RANK:%.*]] = llvm.extractvalue [[SHAPE]][3, 0] : [[SHAPE_TY]]
 // CHECK: [[UNRANKED_OUT_O:%.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr)>
 // CHECK: [[UNRANKED_OUT_1:%.*]] = llvm.insertvalue [[RANK]], [[UNRANKED_OUT_O]][0] : !llvm.struct<(i64, ptr)>
 
 // Compute size in bytes to allocate result ranked descriptor
 // CHECK: [[C1:%.*]] = llvm.mlir.constant(1 : index) : i64
 // CHECK: [[C2:%.*]] = llvm.mlir.constant(2 : index) : i64
-// CHECK: [[PTR_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
 // CHECK: [[INDEX_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
+// CHECK: [[PTR_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
 // CHECK: [[DOUBLE_PTR_SIZE:%.*]] = llvm.mul [[C2]], [[PTR_SIZE]] : i64
 // CHECK: [[DESC_ALLOC_SIZE:%.*]] = llvm.add [[DOUBLE_PTR_SIZE]], %{{.*}}
 // CHECK: [[UNDERLYING_DESC:%.*]] = llvm.alloca [[DESC_ALLOC_SIZE]] x i8
 // CHECK: llvm.insertvalue [[UNDERLYING_DESC]], [[UNRANKED_OUT_1]][1]
 
 // Set allocated, aligned pointers and offset.
 // CHECK: [[ALLOC_PTR:%.*]] = llvm.extractvalue [[INPUT]][0] : [[INPUT_TY]]
 // CHECK: [[ALIGN_PTR:%.*]] = llvm.extractvalue [[INPUT]][1] : [[INPUT_TY]]

mlir/test/Conversion/MemRefToLLVM/typed-pointers.mlir

 // CHECK: [[SHAPE:%.*]] = builtin.unrealized_conversion_cast %{{.*}} to [[SHAPE_TY:!.*]]
 // CHECK: [[RANK:%.*]] = llvm.extractvalue [[SHAPE]][3, 0] : [[SHAPE_TY]]
 // CHECK: [[UNRANKED_OUT_O:%.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
 // CHECK: [[UNRANKED_OUT_1:%.*]] = llvm.insertvalue [[RANK]], [[UNRANKED_OUT_O]][0] : !llvm.struct<(i64, ptr<i8>)>
 
 // Compute size in bytes to allocate result ranked descriptor
 // CHECK: [[C1:%.*]] = llvm.mlir.constant(1 : index) : i64
 // CHECK: [[C2:%.*]] = llvm.mlir.constant(2 : index) : i64
-// CHECK: [[PTR_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
 // CHECK: [[INDEX_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
+// CHECK: [[PTR_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
 // CHECK: [[DOUBLE_PTR_SIZE:%.*]] = llvm.mul [[C2]], [[PTR_SIZE]] : i64
 // CHECK: [[DESC_ALLOC_SIZE:%.*]] = llvm.add [[DOUBLE_PTR_SIZE]], %{{.*}}
 // CHECK: [[UNDERLYING_DESC:%.*]] = llvm.alloca [[DESC_ALLOC_SIZE]] x i8
 // CHECK: llvm.insertvalue [[UNDERLYING_DESC]], [[UNRANKED_OUT_1]][1]
 
 // Set allocated, aligned pointers and offset.
 // CHECK: [[ALLOC_PTR:%.*]] = llvm.extractvalue [[INPUT]][0] : [[INPUT_TY]]
 // CHECK: [[ALIGN_PTR:%.*]] = llvm.extractvalue [[INPUT]][1] : [[INPUT_TY]]