[mlir] Fix LLVM lowering of memory space in unranked memrefs

[dcaballe]

Summary

This patch lowers the memory space of an unraked memref to the
proper address space in LLVM. We were always lowering it to the
default memory space.

[ftynse]

Hmm, this may not be as straightforward as it seems. The pointer in the unranked memref descriptor is to the ranked descriptor struct rather than to the data. Since we don't have two memory spaces in the type, we need to at least define a convention. One can be that the descriptor of the memref is allocated in the same memory space as the memref itself, which requires changes to memref_cast lowering and thinking about what it means for descriptors that live in registers. Another convention can be that the descriptor itself lives in the default memory space, and then this change isn't necessary, only making sure that memref_cast produces the llvm.bitcast to a descriptor struct with correct memory spaces inside.

[dcaballe]

Sorry, I misunderstood what an unranked memref models. This indeed doesn't make sense as is. Thanks.

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

public:
   /// Utilities used to generate integer types.
   static LLVMType getIntNTy(MLIRContext *context, unsigned numBits);
   static LLVMType getInt1Ty(MLIRContext *context) {
     return getIntNTy(context, /*numBits=*/1);
   }
   static LLVMType getInt8Ty(MLIRContext *context) {
     return getIntNTy(context, /*numBits=*/8);
   }
-  static LLVMType getInt8PtrTy(MLIRContext *context) {
-    return getInt8Ty(context).getPointerTo();
+  static LLVMType getInt8PtrTy(MLIRContext *context, unsigned addrSpace = 0) {
+    return getInt8Ty(context).getPointerTo(addrSpace);
   }
   static LLVMType getInt16Ty(MLIRContext *context) {
     return getIntNTy(context, /*numBits=*/16);
   }
   static LLVMType getInt32Ty(MLIRContext *context) {
     return getIntNTy(context, /*numBits=*/32);
   }
   static LLVMType getInt64Ty(MLIRContext *context) {

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 //    stack allocated (alloca) copy of a MemRef descriptor that got casted to
 //    be unranked.
 
 static constexpr unsigned kRankInUnrankedMemRefDescriptor = 0;
 static constexpr unsigned kPtrInUnrankedMemRefDescriptor = 1;
 
 Type LLVMTypeConverter::convertUnrankedMemRefType(UnrankedMemRefType type) {
   auto rankTy = getIndexType();
-  auto ptrTy = LLVM::LLVMType::getInt8PtrTy(&getContext());
+  auto ptrTy =
+      LLVM::LLVMType::getInt8PtrTy(&getContext(), type.getMemorySpace());
   return LLVM::LLVMType::getStructTy(rankTy, ptrTy);
 }
 
 // Convert an n-D vector type to an LLVM vector type via (n-1)-D array type when
 // n > 1.
 // For example, `vector<4 x f32>` converts to `!llvm.type<"<4 x float>">` and
 // `vector<4 x 8 x 16 f32>` converts to `!llvm<"[4 x [8 x <16 x float>]]">`.
 Type LLVMTypeConverter::convertVectorType(VectorType type) {

mlir/test/Conversion/StandardToLLVM/convert-static-memref-ops.mlir

// BAREPTR:  llvm.mlir.constant(13 : index) : !llvm.i64
   %c3 = constant 3 : index
   %3 = dim %static, %c3 : memref<42x32x15x13x27xf32>
 // CHECK:  llvm.mlir.constant(27 : index) : !llvm.i64
 // BAREPTR:  llvm.mlir.constant(27 : index) : !llvm.i64
   %c4 = constant 4 : index
   %4 = dim %static, %c4 : memref<42x32x15x13x27xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @check_unranked_addrspace
+func @check_unranked_addrspace(%in : memref<*xf32, 2>) {
+  // CHECK: llvm.mlir.undef : !llvm.struct<(i64, ptr<i8, 2>)>
+  return
+}
+